TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.API.Neural.Heads

Task heads for public neural-network models.

The definitions here package classifier, regression, and language-model heads that sit on top of the reusable block and Transformer APIs.

source
def TorchLean.nn.Internal.heads.classifier {s : Spec.Shape} (classes : ℕ) (seedW seedB : ℕ := 0) :
Sequential s (Spec.Shape.dim classes Spec.Shape.scalar)

Classification head: Flatten -> Linear.

Named head constructor built from nn.flattenLinear.

Instances For
    source
    def TorchLean.nn.Internal.heads.regressor {s : Spec.Shape} (outDim : ℕ := 1) (seedW seedB : ℕ := 0) :
    Sequential s (Spec.Shape.dim outDim Spec.Shape.scalar)

    Regression head: Flatten -> Linear with outDim outputs.

    Instances For
      source
      def TorchLean.nn.Internal.heads.classifierBatch {n : ℕ} {s : Spec.Shape} (classes : ℕ) (seedW seedB : ℕ := 0) :
      Sequential (Spec.Shape.dim n s) (Spec.Shape.dim n (Spec.Shape.dim classes Spec.Shape.scalar))

      Flatten(start_dim=1) -> Linear head for batched tensors.

      Input: N × σ Output: Mat N classes

      Instances For
        source
        def TorchLean.nn.Internal.heads.regressorBatch {n : ℕ} {s : Spec.Shape} (outDim : ℕ := 1) (seedW seedB : ℕ := 0) :
        Sequential (Spec.Shape.dim n s) (Spec.Shape.dim n (Spec.Shape.dim outDim Spec.Shape.scalar))

        Batched regression head: Flatten(start_dim=1) -> Linear(_, outDim) producing Mat N outDim.

        Instances For