TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.API.Models.Vit

Vision Transformer #

Patch embedding is an arbitrary-dimensional convolution. The spatial output is flattened into a token axis before the Transformer block, so the construction applies equally to one-dimensional signals, images, volumes, and higher-dimensional grids.

source
structure TorchLean.nn.models.VitConfig (d : ℕ) :
Type

Configuration for a Transformer over patches from a d-dimensional spatial domain.

  • batch : ℕ

    Number of independent samples processed together.

  • inChannels : ℕ

    Number of channels in each input sample.

  • spatial : Vector ℕ d

    Extent of each spatial axis.

  • patch : Conv d

    Convolution that extracts and embeds patches.

  • outDim : ℕ

    Number of classifier outputs per sample.

  • numHeads : ℕ

    Number of attention heads.

  • headDim : ℕ

    Width of each attention head.

  • ffnHidden : ℕ

    Width of the feed-forward sublayer.

Instances For
    source
    def TorchLean.nn.models.VitConfig.patchSpatial {d : ℕ} (cfg : VitConfig d) :
    Vector ℕ d

    Spatial extent of the patch embedding.

    Instances For
      source
      def TorchLean.nn.models.VitConfig.seqLen {d : ℕ} (cfg : VitConfig d) :
      ℕ

      Number of patch tokens.

      Instances For
        source
        def TorchLean.nn.models.VitConfig.flatDim {d : ℕ} (cfg : VitConfig d) :
        ℕ

        Number of flattened features passed to the classifier.

        Instances For
          source
          def TorchLean.nn.models.vitInShape {d : ℕ} (cfg : VitConfig d) :
          Spec.Shape

          Input shape (batch, inChannels, spatial...).

          Instances For
            source
            def TorchLean.nn.models.vitConvOutShape {d : ℕ} (cfg : VitConfig d) :
            Spec.Shape

            Shape produced by patch embedding.

            Instances For
              source
              def TorchLean.nn.models.vitTokensShape {d : ℕ} (cfg : VitConfig d) :
              Spec.Shape

              Token shape (batch, sequence, embedding).

              Instances For
                source
                def TorchLean.nn.models.vitOutShape {d : ℕ} (cfg : VitConfig d) :
                Spec.Shape

                Classifier output shape (batch, outDim).

                Instances For
                  source
                  def TorchLean.nn.models.spatialToTokens {d : ℕ} (cfg : VitConfig d) :
                  LayerDef (vitConvOutShape cfg) (vitTokensShape cfg)

                  flatten the patch grid into a sequence and move channels to the final axis.

                  Instances For
                    source
                    def TorchLean.nn.models.vit {d : ℕ} (cfg : VitConfig d) (hInChannels : cfg.inChannels ≠ 0 := by decide) (hSeqLen : cfg.seqLen ≠ 0 := by decide) (hModel : cfg.patch.outChannels ≠ 0 := by decide) :
                    M (Sequential (vitInShape cfg) (vitOutShape cfg))

                    Build patch embedding, one Transformer encoder block, and a linear classifier.

                    Instances For