mlpmixer

Documentation

Overview

Package mlpmixer implements the MLP-Mixer (Tolstikhin et al., 2021).

Index

• type Config
• type FeedForward
  • func (m *FeedForward) Forward(xs ...ag.Node) []ag.Node
• type MixerBlock
  • func New[T float.DType](config Config) *MixerBlock
  • func (m *MixerBlock) Forward(xs ...ag.Node) []ag.Node

Types

type Config

type Config struct {
	InputSize               int
	HiddenSizeTokenMixer    int
	HiddenSizeChannelMixer  int
	Channels                int
	ActFunctionTokenMixer   activation.Name
	ActFunctionChannelMixer activation.Name
	Eps                     float64
}

Config provides configuration settings for a MixerBlock.

type FeedForward

type FeedForward struct {
	nn.Module
	Layers []nn.StandardModel
}

FeedForward is the model for feed-forward operations of a MixerBlock.

func (*FeedForward) Forward

func (m *FeedForward) Forward(xs ...ag.Node) []ag.Node

type MixerBlock

type MixerBlock struct {
	nn.Module
	Config
	TokenLayerNorm   *layernorm.Model
	TokenMixerFF     *FeedForward
	ChannelLayerNorm *layernorm.Model
	ChannelMixerFF   *FeedForward
}

MixerBlock contains the serializable parameters.

func New

func New[T float.DType](config Config) *MixerBlock

New returns a new model with parameters initialized to zeros.

func (*MixerBlock) Forward

func (m *MixerBlock) Forward(xs ...ag.Node) []ag.Node

Forward performs the forward step for each input node and returns the result.