GitHub

MLPMixer-like models

This is the API reference for the MLPMixer-like models supported by Metalhead.jl.

The higher-level model constructors

Metalhead.MLPMixer — Type
MLPMixer(config::Symbol; patch_size::Dims{2} = (16, 16), imsize::Dims{2} = (224, 224),
         inchannels::Integer = 3, nclasses::Integer = 1000)

Creates a model with the MLPMixer architecture. (reference).

Arguments

  • config: the size of the model - one of :small, :base, :large or :huge
  • patch_size: the size of the patches
  • imsize: the size of the input image
  • stochastic_depth_prob: Stochastic depth probability
  • inchannels: the number of input channels
  • nclasses: number of output classes

See also Metalhead.mlpmixer.

source
Metalhead.ResMLP — Type
ResMLP(config::Symbol; patch_size::Dims{2} = (16, 16), imsize::Dims{2} = (224, 224),
       inchannels::Integer = 3, nclasses::Integer = 1000)

Creates a model with the ResMLP architecture. (reference).

Arguments

  • config: the size of the model - one of :small, :base, :large or :huge
  • patch_size: the size of the patches
  • imsize: the size of the input image
  • inchannels: the number of input channels
  • nclasses: number of output classes

See also Metalhead.mlpmixer.

source
Metalhead.gMLP — Type
gMLP(config::Symbol; patch_size::Dims{2} = (16, 16), imsize::Dims{2} = (224, 224),
     inchannels::Integer = 3, nclasses::Integer = 1000)

Creates a model with the gMLP architecture. (reference).

Arguments

  • config: the size of the model - one of :small, :base, :large or :huge
  • patch_size: the size of the patches
  • imsize: the size of the input image
  • inchannels: the number of input channels
  • nclasses: number of output classes

See also Metalhead.mlpmixer.

source

The core MLPMixer function

Metalhead.mlpmixer — Function
mlpmixer(block, imsize::Dims{2} = (224, 224); inchannels::Integer = 3, norm_layer = LayerNorm,
         patch_size::Dims{2} = (16, 16), embedplanes = 512, stochastic_depth_prob = 0.,
         depth::Integer = 12, nclasses::Integer = 1000, kwargs...)

Creates a model with the MLPMixer architecture. (reference).

Arguments

  • block: the type of mixer block to use in the model - architecture dependent (a constructor of the form block(embedplanes, npatches; stochastic_depth_prob, kwargs...))
  • imsize: the size of the input image
  • inchannels: the number of input channels
  • norm_layer: the normalization layer to use in the model
  • patch_size: the size of the patches
  • embedplanes: the number of channels after the patch embedding (denotes the hidden dimension)
  • stochastic_depth_prob: Stochastic depth probability
  • depth: the number of blocks in the model
  • nclasses: number of output classes
  • kwargs: additional arguments (if any) to pass to the mixer block. Will use the defaults if not specified.
source

The block functions

Metalhead.mixerblock — Function
mixerblock(planes::Integer, npatches::Integer; mlp_layer = mlp_block,
           mlp_ratio = (0.5, 4.0), dropout_prob = 0.0, stochastic_depth_prob = 0.0,
           activation = gelu)

Creates a feedforward block for the MLPMixer architecture. (reference)

Arguments

  • planes: the number of planes in the block
  • npatches: the number of patches of the input
  • mlp_ratio: number(s) that determine(s) the number of hidden channels in the token mixing MLP and/or the channel mixing MLP as a ratio to the number of planes in the block.
  • mlp_layer: the MLP layer to use in the block
  • dropout_prob: the dropout probability to use in the MLP blocks
  • stochastic_depth_prob: Stochastic depth probability
  • activation: the activation function to use in the MLP blocks
source
Metalhead.resmixerblock — Function
resmixerblock(planes, npatches; dropout_prob = 0., stochastic_depth_prob = 0., mlp_ratio = 4.0,
              activation = gelu, layerscale_init = 1e-4)

Creates a block for the ResMixer architecture. (reference).

Arguments

  • planes: the number of planes in the block
  • npatches: the number of patches of the input
  • mlp_ratio: ratio of the number of hidden channels in the channel mixing MLP to the number of planes in the block
  • mlp_layer: the MLP block to use
  • dropout_prob: the dropout probability to use in the MLP blocks
  • stochastic_depth_prob: Stochastic depth probability
  • activation: the activation function to use in the MLP blocks
  • layerscale_init: initialisation constant for the LayerScale
source
Metalhead.SpatialGatingUnit — Type
SpatialGatingUnit(planes::Integer, npatches::Integer; norm_layer = LayerNorm)

Creates a spatial gating unit as described in the gMLP paper. (reference)

Arguments

  • planes: the number of planes in the block
  • npatches: the number of patches of the input
  • norm_layer: the normalisation layer to use
source
Metalhead.spatialgatingblock — Function
spatialgatingblock(planes::Integer, npatches::Integer; mlp_ratio = 4.0,
                   norm_layer = LayerNorm, mlp_layer = gated_mlp_block,
                   dropout_prob = 0.0, stochastic_depth_prob = 0.0,
                   activation = gelu)

Creates a feedforward block based on the gMLP model architecture described in the paper. (reference)

Arguments

  • planes: the number of planes in the block
  • npatches: the number of patches of the input
  • mlp_ratio: ratio of the number of hidden channels in the channel mixing MLP to the number of planes in the block
  • norm_layer: the normalisation layer to use
  • dropout_prob: the dropout probability to use in the MLP blocks
  • stochastic_depth_prob: Stochastic depth probability
  • activation: the activation function to use in the MLP blocks
source