Built-In Layers
Index
Lux.AdaptiveMaxPoolLux.AdaptiveMeanPoolLux.AlphaDropoutLux.BatchNormLux.BilinearLux.BranchLayerLux.ChainLux.ConvLux.ConvTransposeLux.CrossCorLux.DenseLux.DropoutLux.EmbeddingLux.FlattenLayerLux.GRUCellLux.GlobalMaxPoolLux.GlobalMeanPoolLux.GroupNormLux.InstanceNormLux.LSTMCellLux.LayerNormLux.MaxPoolLux.MaxoutLux.MeanPoolLux.NoOpLayerLux.PairwiseFusionLux.ParallelLux.RNNCellLux.RecurrenceLux.RepeatedLayerLux.ReshapeLayerLux.ScaleLux.SelectDimLux.SkipConnectionLux.StatefulRecurrentCellLux.UpsampleLux.VariationalHiddenDropoutLux.WeightNormLux.WrappedFunctionLux.PixelShuffle
Containers
BranchLayer(layers...)
BranchLayer(; name=nothing, layers...)Takes an input x and passes it through all the layers and returns a tuple of the outputs.
Arguments
- Layers can be specified in two formats:
A list of
NLux layersSpecified as
Nkeyword arguments.
Keyword Arguments
name: Name of the layer (optional)
Inputs
x: Will be directly passed to each of thelayers
Returns
Tuple:
(layer_1(x), layer_2(x), ..., layer_N(x))(naming changes if using the kwargs API)Updated state of the
layers
Parameters
- Parameters of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
States
- States of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
Comparison with Parallel
This is slightly different from Parallel(nothing, layers...)
If the input is a tuple,
Parallelwill pass each element individually to each layer.BranchLayeressentially assumes 1 input comes in and is branched out intoNoutputs.
Example
An easy way to replicate an input to an NTuple is to do
l = BranchLayer(NoOpLayer(), NoOpLayer(), NoOpLayer())Chain(layers...; name=nothing, disable_optimizations::Bool = false)
Chain(; layers..., name=nothing, disable_optimizations::Bool = false)Collects multiple layers / functions to be called in sequence on a given input.
Arguments
- Layers can be specified in two formats:
A list of
NLux layersSpecified as
Nkeyword arguments.
Keyword Arguments
disable_optimizations: Prevents any structural optimizationname: Name of the layer (optional)
Inputs
Input x is passed sequentially to each layer, and must conform to the input requirements of the internal layers.
Returns
Output after sequentially applying all the layers to
xUpdated model states
Parameters
- Parameters of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
States
- States of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
Optimizations
Performs a few optimizations to generate reasonable architectures. Can be disabled using keyword argument disable_optimizations.
All sublayers are recursively optimized.
If a function
fis passed as a layer and it doesn't take 3 inputs, it is converted to aWrappedFunction(f) which takes only one input.If the layer is a Chain, it is flattened.
NoOpLayers are removed.If there is only 1 layer (left after optimizations), then it is returned without the
Chainwrapper.If there are no layers (left after optimizations), a
NoOpLayeris returned.
Miscellaneous Properties
- Allows indexing. We can access the
ith layer usingm[i]. We can also index using ranges or arrays.
Example
c = Chain(Dense(2, 3, relu), BatchNorm(3), Dense(3, 2))PairwiseFusion(connection, layers...; name=nothing)
PairwiseFusion(connection; name=nothing, layers...)x1 → layer1 → y1 ↘
connection → layer2 → y2 ↘
x2 ↗ connection → y3
x3 ↗Arguments
connection: Takes 2 inputs and combines themlayers:AbstractExplicitLayers. Layers can be specified in two formats:A list of
NLux layersSpecified as
Nkeyword arguments.
Keyword Arguments
name: Name of the layer (optional)
Inputs
Layer behaves differently based on input type:
- If the input
xis a tuple of lengthN + 1, then thelayersmust be a tuple of lengthN. The computation is as follows
y = x[1]
for i in 1:N
y = connection(x[i + 1], layers[i](y))
end- Any other kind of input
y = x
for i in 1:N
y = connection(x, layers[i](y))
endReturns
See Inputs section for how the return value is computed
Updated model state for all the contained layers
Parameters
- Parameters of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
States
- States of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
Parallel(connection, layers...; name=nothing)
Parallel(connection; name=nothing, layers...)Create a layer which passes an input to each path in layers, before reducing the output with connection.
Arguments
connection: AnN-argument function that is called after passing the input through each layer. Ifconnection = nothing, we return a tupleParallel(nothing, f, g)(x, y) = (f(x), g(y))Layers can be specified in two formats:
A list of
NLux layersSpecified as
Nkeyword arguments.
Keyword Arguments
name: Name of the layer (optional)
Inputs
x: Ifxis not a tuple, then return is computed asconnection([l(x) for l in layers]...). Else one is passed to each layer, thusParallel(+, f, g)(x, y) = f(x) + g(y).
Returns
See the Inputs section for how the output is computed
Updated state of the
layers
Parameters
- Parameters of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
States
- States of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
See also SkipConnection which is Parallel with one identity.
SkipConnection(layer, connection; name=nothing)Create a skip connection which consists of a layer or Chain of consecutive layers and a shortcut connection linking the block's input to the output through a user-supplied 2-argument callable. The first argument to the callable will be propagated through the given layer while the second is the unchanged, "skipped" input.
The simplest "ResNet"-type connection is just SkipConnection(layer, +).
Arguments
layer: Layer orChainof layers to be applied to the inputconnection:A 2-argument function that takes
layer(input)and the input ORAn AbstractExplicitLayer that takes
(layer(input), input)as input
Keyword Arguments
name: Name of the layer (optional)
Inputs
x: Will be passed directly tolayer
Returns
Output of
connection(layer(input), input)Updated state of
layer
Parameters
Parameters of
layerORIf
connectionis an AbstractExplicitLayer, then NamedTuple with fields:layersand:connection
States
States of
layerORIf
connectionis an AbstractExplicitLayer, then NamedTuple with fields:layersand:connection
See Parallel for a more general implementation.
RepeatedLayer(model; repeats::Val = Val(10), input_injection::Val = Val(false))Iteratively applies model for repeats number of times. The initial input is passed into the model repeatedly if input_injection = Val(true). This layer unrolls the computation, however, semantically this is same as:
input_injection = Val(false)
res = x
for i in 1:repeats
res, st = model(res, ps, st)
endinput_injection = Val(true)
res = x
for i in 1:repeats
res, st = model((res, x), ps, st)
endIt is expected that repeats will be a reasonable number below 20, beyond that compile times for gradients might be unreasonably high.
Arguments
modelmust be anAbstractExplicitLayer
Keyword Arguments
repeats: Number of times to apply the modelinput_injection: Iftrue, then the input is passed to the model along with the output
Inputs
x: Input as described above
Returns
Output is computed by as described above
Updated state of the
model
Parameters
- Parameters of
model
States
- State of
model
Convolutional Layers
Conv(k::NTuple{N,Integer}, (in_chs => out_chs)::Pair{<:Integer,<:Integer},
activation=identity; init_weight=glorot_uniform, init_bias=zeros32, stride=1,
pad=0, dilation=1, groups=1, use_bias=true)Standard convolutional layer.
Image data should be stored in WHCN order (width, height, channels, batch). In other words, a 100 x 100 RGB image would be a 100 x 100 x 3 x 1 array, and a batch of 50 would be a 100 x 100 x 3 x 50 array. This has N = 2 spatial dimensions, and needs a kernel size like (5, 5), a 2-tuple of integers. To take convolutions along N feature dimensions, this layer expects as input an array with ndims(x) == N + 2, where size(x, N + 1) == in_chs is the number of input channels, and size(x, ndims(x)) is the number of observations in a batch.
Warning
Frameworks like Pytorch perform cross-correlation in their convolution layers
Arguments
k: Tuple of integers specifying the size of the convolutional kernel. Eg, for 2D convolutionslength(k) == 2in_chs: Number of input channelsout_chs: Number of input and output channelsactivation: Activation Function
Keyword Arguments
init_weight: Controls the initialization of the weight parameterinit_bias: Controls the initialization of the bias parameterstride: Should each be either single integer, or a tuple withNintegersdilation: Should each be either single integer, or a tuple withNintegerspad: Specifies the number of elements added to the borders of the data array. It can bea single integer for equal padding all around,
a tuple of
Nintegers, to apply the same padding at begin/end of each spatial dimension,a tuple of
2*Nintegers, for asymmetric padding, orthe singleton
SamePad(), to calculate padding such thatsize(output,d) == size(x,d) / stride(possibly rounded) for each spatial dimension.Periodic padding can achieved by pre-empting the layer with a
WrappedFunction(x -> NNlib.circular_pad(x, N_pad; dims=pad_dims))
groups: Expected to be anInt. It specifies the number of groups to divide a convolution into (setgroups = in_chsfor Depthwise Convolutions).in_chsandout_chsmust be divisible bygroups.use_bias: Trainable bias can be disabled entirely by setting this tofalse.allow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)
Inputs
x: Data satisfyingndims(x) == N + 2 && size(x, N - 1) == in_chs, i.e.size(x) = (I_N, ..., I_1, C_in, N)
Returns
- Output of the convolution
yof size(O_N, ..., O_1, C_out, N)where
- Empty
NamedTuple()
Parameters
weight: Convolution kernelbias: Bias (present ifuse_bias=true)
ConvTranspose(k::NTuple{N,Integer}, (in_chs => out_chs)::Pair{<:Integer,<:Integer},
activation=identity; init_weight=glorot_uniform, init_bias=zeros32,
stride=1, pad=0, dilation=1, groups=1, use_bias=true)Standard convolutional transpose layer.
Arguments
k: Tuple of integers specifying the size of the convolutional kernel. Eg, for 2D convolutionslength(k) == 2in_chs: Number of input channelsout_chs: Number of input and output channelsactivation: Activation Function
Keyword Arguments
init_weight: Controls the initialization of the weight parameterinit_bias: Controls the initialization of the bias parameterstride: Should each be either single integer, or a tuple withNintegersdilation: Should each be either single integer, or a tuple withNintegerspad: Specifies the number of elements added to the borders of the data array. It can bea single integer for equal padding all around,
a tuple of
Nintegers, to apply the same padding at begin/end of each spatial dimension,a tuple of
2*Nintegers, for asymmetric padding, orthe singleton
SamePad(), to calculate padding such thatsize(output,d) == size(x,d) * stride(possibly rounded) for each spatial dimension.
groups: Expected to be anInt. It specifies the number of groups to divide a convolution into (setgroups = in_chsfor Depthwise Convolutions).in_chsandout_chsmust be divisible bygroups.use_bias: Trainable bias can be disabled entirely by setting this tofalse.allow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)
Inputs
x: Data satisfyingndims(x) == N + 2 && size(x, N - 1) == in_chs, i.e.size(x) = (I_N, ..., I_1, C_in, N)
Returns
Output of the convolution transpose
yof size(O_N, ..., O_1, C_out, N)whereEmpty
NamedTuple()
Parameters
weight: Convolution Transpose kernelbias: Bias (present ifuse_bias=true)
CrossCor(k::NTuple{N,Integer}, (in_chs => out_chs)::Pair{<:Integer,<:Integer},
activation=identity; init_weight=glorot_uniform, init_bias=zeros32, stride=1,
pad=0, dilation=1, use_bias=true)Cross Correlation layer.
Image data should be stored in WHCN order (width, height, channels, batch). In other words, a 100 x 100 RGB image would be a 100 x 100 x 3 x 1 array, and a batch of 50 would be a 100 x 100 x 3 x 50 array. This has N = 2 spatial dimensions, and needs a kernel size like (5, 5), a 2-tuple of integers. To take convolutions along N feature dimensions, this layer expects as input an array with ndims(x) == N + 2, where size(x, N + 1) == in_chs is the number of input channels, and size(x, ndims(x)) is the number of observations in a batch.
Arguments
k: Tuple of integers specifying the size of the convolutional kernel. Eg, for 2D convolutionslength(k) == 2in_chs: Number of input channelsout_chs: Number of input and output channelsactivation: Activation Function
Keyword Arguments
init_weight: Controls the initialization of the weight parameterinit_bias: Controls the initialization of the bias parameterstride: Should each be either single integer, or a tuple withNintegersdilation: Should each be either single integer, or a tuple withNintegerspad: Specifies the number of elements added to the borders of the data array. It can bea single integer for equal padding all around,
a tuple of
Nintegers, to apply the same padding at begin/end of each spatial dimension,a tuple of
2*Nintegers, for asymmetric padding, orthe singleton
SamePad(), to calculate padding such thatsize(output,d) == size(x,d) / stride(possibly rounded) for each spatial dimension.
use_bias: Trainable bias can be disabled entirely by setting this tofalse.allow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)
Inputs
x: Data satisfyingndims(x) == N + 2 && size(x, N - 1) == in_chs, i.e.size(x) = (I_N, ..., I_1, C_in, N)
Returns
- Output of the convolution
yof size(O_N, ..., O_1, C_out, N)where
- Empty
NamedTuple()
Parameters
weight: Convolution kernelbias: Bias (present ifuse_bias=true)
Dropout Layers
AlphaDropout(p::Real)AlphaDropout layer.
Arguments
p: Probability of Dropoutif
p = 0thenNoOpLayeris returned.if
p = 1thenWrappedLayer(Base.Fix1(broadcast, zero))is returned.
Inputs
x: Must be an AbstractArray
Returns
xwith dropout mask applied iftraining=Val(true)else justxState with updated
rng
States
rng: Pseudo Random Number Generatortraining: Used to check if training/inference mode
Call Lux.testmode to switch to test mode.
See also Dropout, VariationalHiddenDropout
Dropout(p; dims=:)Dropout layer.
Arguments
p: Probability of Dropout (ifp = 0thenNoOpLayeris returned)
Keyword Arguments
- To apply dropout along certain dimension(s), specify the
dimskeyword. e.g.Dropout(p; dims = 3)will randomly zero out entire channels on WHCN input (also called 2D dropout).
Inputs
x: Must be an AbstractArray
Returns
xwith dropout mask applied iftraining=Val(true)else justxState with updated
rng
States
rng: Pseudo Random Number Generatortraining: Used to check if training/inference mode
Call Lux.testmode to switch to test mode.
See also AlphaDropout, VariationalHiddenDropout
VariationalHiddenDropout(p; dims=:)VariationalHiddenDropout layer. The only difference from Dropout is that the mask is retained until Lux.update_state(l, :update_mask, Val(true)) is called.
Arguments
p: Probability of Dropout (ifp = 0thenNoOpLayeris returned)
Keyword Arguments
- To apply dropout along certain dimension(s), specify the
dimskeyword. e.g.VariationalHiddenDropout(p; dims = 3)will randomly zero out entire channels on WHCN input (also called 2D dropout).
Inputs
x: Must be an AbstractArray
Returns
xwith dropout mask applied iftraining=Val(true)else justxState with updated
rng
States
rng: Pseudo Random Number Generatortraining: Used to check if training/inference modemask: Dropout mask. Initilly set to nothing. After every run, contains the mask applied in that callupdate_mask: Stores whether new mask needs to be generated in the current call
Call Lux.testmode to switch to test mode.
See also AlphaDropout, Dropout
Pooling Layers
AdaptiveMaxPool(out::NTuple)Adaptive Max Pooling layer. Calculates the necessary window size such that its output has size(y)[1:N] == out.
Arguments
out: Size of the firstNdimensions for the output
Inputs
x: Expects as input an array withndims(x) == N+2, i.e. channel and batch dimensions, after theNfeature dimensions, whereN = length(out).
Returns
Output of size
(out..., C, N)Empty
NamedTuple()
See also MaxPool, AdaptiveMeanPool.
AdaptiveMeanPool(out::NTuple)Adaptive Mean Pooling layer. Calculates the necessary window size such that its output has size(y)[1:N] == out.
Arguments
out: Size of the firstNdimensions for the output
Inputs
x: Expects as input an array withndims(x) == N+2, i.e. channel and batch dimensions, after theNfeature dimensions, whereN = length(out).
Returns
Output of size
(out..., C, N)Empty
NamedTuple()
See also MeanPool, AdaptiveMaxPool.
GlobalMaxPool()Global Max Pooling layer. Transforms (w,h,c,b)-shaped input into (1,1,c,b)-shaped output, by performing max pooling on the complete (w,h)-shaped feature maps.
Inputs
x: Data satisfyingndims(x) > 2, i.e.size(x) = (I_N, ..., I_1, C, N)
Returns
Output of the pooling
yof size(1, ..., 1, C, N)Empty
NamedTuple()
See also MaxPool, AdaptiveMaxPool, GlobalMeanPool
GlobalMeanPool()Global Mean Pooling layer. Transforms (w,h,c,b)-shaped input into (1,1,c,b)-shaped output, by performing mean pooling on the complete (w,h)-shaped feature maps.
Inputs
x: Data satisfyingndims(x) > 2, i.e.size(x) = (I_N, ..., I_1, C, N)
Returns
Output of the pooling
yof size(1, ..., 1, C, N)Empty
NamedTuple()
See also MeanPool, AdaptiveMeanPool, GlobalMaxPool
MaxPool(window::NTuple; pad=0, stride=window)Max pooling layer, which replaces all pixels in a block of size window with the maximum value.
Arguments
window: Tuple of integers specifying the size of the window. Eg, for 2D poolinglength(window) == 2
Keyword Arguments
stride: Should each be either single integer, or a tuple withNintegerspad: Specifies the number of elements added to the borders of the data array. It can bea single integer for equal padding all around,
a tuple of
Nintegers, to apply the same padding at begin/end of each spatial dimension,a tuple of
2*Nintegers, for asymmetric padding, orthe singleton
SamePad(), to calculate padding such thatsize(output,d) == size(x,d) / stride(possibly rounded) for each spatial dimension.
Inputs
x: Data satisfyingndims(x) == N + 2, i.e.size(x) = (I_N, ..., I_1, C, N)
Returns
- Output of the pooling
yof size(O_N, ..., O_1, C, N)where
- Empty
NamedTuple()
See also Conv, MeanPool, GlobalMaxPool, AdaptiveMaxPool
MeanPool(window::NTuple; pad=0, stride=window)Mean pooling layer, which replaces all pixels in a block of size window with the mean value.
Arguments
window: Tuple of integers specifying the size of the window. Eg, for 2D poolinglength(window) == 2
Keyword Arguments
stride: Should each be either single integer, or a tuple withNintegerspad: Specifies the number of elements added to the borders of the data array. It can bea single integer for equal padding all around,
a tuple of
Nintegers, to apply the same padding at begin/end of each spatial dimension,a tuple of
2*Nintegers, for asymmetric padding, orthe singleton
SamePad(), to calculate padding such thatsize(output,d) == size(x,d) / stride(possibly rounded) for each spatial dimension.
Inputs
x: Data satisfyingndims(x) == N + 2, i.e.size(x) = (I_N, ..., I_1, C, N)
Returns
- Output of the pooling
yof size(O_N, ..., O_1, C, N)where
- Empty
NamedTuple()
See also Conv, MaxPool, GlobalMeanPool, AdaptiveMeanPool
Recurrent Layers
GRUCell((in_dims, out_dims)::Pair{<:Int,<:Int}; use_bias=true, train_state::Bool=false,
init_weight::Tuple{Function,Function,Function}=(glorot_uniform, glorot_uniform,
glorot_uniform),
init_bias::Tuple{Function,Function,Function}=(zeros32, zeros32, zeros32),
init_state::Function=zeros32)Gated Recurrent Unit (GRU) Cell
Arguments
in_dims: Input Dimensionout_dims: Output (Hidden State) Dimensionuse_bias: Set to false to deactivate biastrain_state: Trainable initial hidden state can be activated by setting this totrueinit_bias: Initializer for bias. Must be a tuple containing 3 functionsinit_weight: Initializer for weight. Must be a tuple containing 3 functionsinit_state: Initializer for hidden state
Inputs
Case 1a: Only a single input
xof shape(in_dims, batch_size),train_stateis set tofalse- Creates a hidden state usinginit_stateand proceeds to Case 2.Case 1b: Only a single input
xof shape(in_dims, batch_size),train_stateis set totrue- Repeatshidden_statefrom parameters to match the shape ofxand proceeds to Case 2.Case 2: Tuple
(x, (h, ))is provided, then the output and a tuple containing the updated hidden state is returned.
Returns
Tuple containing
Output
of shape (out_dims, batch_size)Tuple containing new hidden state
Updated model state
Parameters
weight_i: Concatenated Weights to map from input space. weight_h: Concatenated Weights to map from hidden space. bias_i: Bias vector (; not present if use_bias=false).bias_h: Concatenated Bias vector for the hidden space(not present if use_bias=false).hidden_state: Initial hidden state vector (not present iftrain_state=false).
States
rng: Controls the randomness (if any) in the initial state generation
LSTMCell(in_dims => out_dims; use_bias::Bool=true, train_state::Bool=false,
train_memory::Bool=false,
init_weight=(glorot_uniform, glorot_uniform, glorot_uniform, glorot_uniform),
init_bias=(zeros32, zeros32, ones32, zeros32), init_state=zeros32,
init_memory=zeros32)Long Short-Term (LSTM) Cell
Arguments
in_dims: Input Dimensionout_dims: Output (Hidden State & Memory) Dimensionuse_bias: Set to false to deactivate biastrain_state: Trainable initial hidden state can be activated by setting this totruetrain_memory: Trainable initial memory can be activated by setting this totrueinit_bias: Initializer for bias. Must be a tuple containing 4 functionsinit_weight: Initializer for weight. Must be a tuple containing 4 functionsinit_state: Initializer for hidden stateinit_memory: Initializer for memory
Inputs
Case 1a: Only a single input
xof shape(in_dims, batch_size),train_stateis set tofalse,train_memoryis set tofalse- Creates a hidden state usinginit_state, hidden memory usinginit_memoryand proceeds to Case 2.Case 1b: Only a single input
xof shape(in_dims, batch_size),train_stateis set totrue,train_memoryis set tofalse- Repeatshidden_statevector from the parameters to match the shape ofx, creates hidden memory usinginit_memoryand proceeds to Case 2.Case 1c: Only a single input
xof shape(in_dims, batch_size),train_stateis set tofalse,train_memoryis set totrue- Creates a hidden state usinginit_state, repeats the memory vector from parameters to match the shape ofxand proceeds to Case 2.Case 1d: Only a single input
xof shape(in_dims, batch_size),train_stateis set totrue,train_memoryis set totrue- Repeats the hidden state and memory vectors from the parameters to match the shape ofxand proceeds to Case 2.Case 2: Tuple
(x, (h, c))is provided, then the output and a tuple containing the updated hidden state and memory is returned.
Returns
Tuple Containing
Output
of shape (out_dims, batch_size)Tuple containing new hidden state
and new memory
Updated model state
Parameters
weight_i: Concatenated Weights to map from input space. weight_h: Concatenated Weights to map from hidden spacebias: Bias vector (not present ifuse_bias=false)hidden_state: Initial hidden state vector (not present iftrain_state=false)memory: Initial memory vector (not present iftrain_memory=false)
States
rng: Controls the randomness (if any) in the initial state generation
RNNCell(in_dims => out_dims, activation=tanh; bias::Bool=true,
train_state::Bool=false, init_bias=zeros32, init_weight=glorot_uniform,
init_state=ones32)An Elman RNNCell cell with activation (typically set to tanh or relu).
Arguments
in_dims: Input Dimensionout_dims: Output (Hidden State) Dimensionactivation: Activation functionbias: Set to false to deactivate biastrain_state: Trainable initial hidden state can be activated by setting this totrueinit_bias: Initializer for biasinit_weight: Initializer for weightinit_state: Initializer for hidden state
Inputs
Case 1a: Only a single input
xof shape(in_dims, batch_size),train_stateis set tofalse- Creates a hidden state usinginit_stateand proceeds to Case 2.Case 1b: Only a single input
xof shape(in_dims, batch_size),train_stateis set totrue- Repeatshidden_statefrom parameters to match the shape ofxand proceeds to Case 2.Case 2: Tuple
(x, (h, ))is provided, then the output and a tuple containing the updated hidden state is returned.
Returns
Tuple containing
Output
of shape (out_dims, batch_size)Tuple containing new hidden state
Updated model state
Parameters
weight_ih: Maps the input to the hidden state.weight_hh: Maps the hidden state to the hidden state.bias: Bias vector (not present ifuse_bias=false)hidden_state: Initial hidden state vector (not present iftrain_state=false)
States
rng: Controls the randomness (if any) in the initial state generation
Recurrence(cell;
ordering::AbstractTimeSeriesDataBatchOrdering=BatchLastIndex(),
return_sequence::Bool=false)Wraps a recurrent cell (like RNNCell, LSTMCell, GRUCell) to automatically operate over a sequence of inputs.
Warning
This is completely distinct from Flux.Recur. It doesn't make the cell stateful, rather allows operating on an entire sequence of inputs at once. See StatefulRecurrentCell for functionality similar to Flux.Recur.
Arguments
cell: A recurrent cell. SeeRNNCell,LSTMCell,GRUCell, for how the inputs/outputs of a recurrent cell must be structured.
Keyword Arguments
return_sequence: Iftruereturns the entire sequence of outputs, else returns only the last output. Defaults tofalse.ordering: The ordering of the batch and time dimensions in the input. Defaults toBatchLastIndex(). Alternatively can be set toTimeLastIndex().
Inputs
- If
xis aTuple or Vector: Each element is fed to the
cellsequentially.Array (except a Vector): It is spliced along the penultimate dimension and each slice is fed to the
cellsequentially.
Returns
Output of the
cellfor the entire sequence.Update state of the
cell.
Parameters
- Same as
cell.
States
- Same as
cell.
Tip
Frameworks like Tensorflow have special implementation of MultiRNNCell to handle sequentially composed RNN Cells. In Lux, one can simple stack multiple Recurrence blocks in a Chain to achieve the same.
Chain(
Recurrence(RNNCell(inputsize => latentsize); return_sequence=true),
Recurrence(RNNCell(latentsize => latentsize); return_sequence=true),
:
x -> stack(x; dims=2)
)For some discussion on this topic, see https://github.com/LuxDL/Lux.jl/issues/472.
StatefulRecurrentCell(cell)Wraps a recurrent cell (like RNNCell, LSTMCell, GRUCell) and makes it stateful.
Tip
This is very similar to Flux.Recur
To avoid undefined behavior, once the processing of a single sequence of data is complete, update the state with Lux.update_state(st, :carry, nothing).
Arguments
cell: A recurrent cell. SeeRNNCell,LSTMCell,GRUCell, for how the inputs/outputs of a recurrent cell must be structured.
Inputs
- Input to the
cell.
Returns
Output of the
cellfor the entire sequence.Update state of the
celland updatedcarry.
Parameters
- Same as
cell.
States
- NamedTuple containing:
cell: Same ascell.carry: The carry state of thecell.
Linear Layers
Bilinear((in1_dims, in2_dims) => out, activation=identity; init_weight=glorot_uniform,
init_bias=zeros32, use_bias::Bool=true, allow_fast_activation::Bool=true)
Bilinear(in12_dims => out, activation=identity; init_weight=glorot_uniform,
init_bias=zeros32, use_bias::Bool=true, allow_fast_activation::Bool=true)Create a fully connected layer between two inputs and an output, and otherwise similar to Dense. Its output, given vectors x & y, is another vector z with, for all i in 1:out:
z[i] = activation(x' * W[i, :, :] * y + bias[i])
If x and y are matrices, then each column of the output z = B(x, y) is of this form, with B the Bilinear layer.
Arguments
in1_dims: number of input dimensions ofxin2_dims: number of input dimensions ofyin12_dims: If specified, thenin1_dims = in2_dims = in12_dimsout: number of output dimensionsactivation: activation function
Keyword Arguments
init_weight: initializer for the weight matrix (weight = init_weight(rng, out_dims, in1_dims, in2_dims))init_bias: initializer for the bias vector (ignored ifuse_bias=false)use_bias: Trainable bias can be disabled entirely by setting this tofalseallow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)
Input
A 2-Tuple containing
xmust be an AbstractArray withsize(x, 1) == in1_dimsymust be an AbstractArray withsize(y, 1) == in2_dims
If the input is an AbstractArray, then
x = y
Returns
AbstractArray with dimensions
(out_dims, size(x, 2))Empty
NamedTuple()
Parameters
weight: Weight Matrix of size(out_dims, in1_dims, in2_dims)bias: Bias of size(out_dims, 1)(present ifuse_bias=true)
Dense(in_dims => out_dims, activation=identity; init_weight=glorot_uniform,
init_bias=zeros32, use_bias::Bool=true, allow_fast_activation::Bool=true)Create a traditional fully connected layer, whose forward pass is given by: y = activation.(weight * x .+ bias)
Arguments
in_dims: number of input dimensionsout_dims: number of output dimensionsactivation: activation function
Keyword Arguments
init_weight: initializer for the weight matrix (weight = init_weight(rng, out_dims, in_dims))init_bias: initializer for the bias vector (ignored ifuse_bias=false)use_bias: Trainable bias can be disabled entirely by setting this tofalseallow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)
Input
xmust be an AbstractArray withsize(x, 1) == in_dims
Returns
AbstractArray with dimensions
(out_dims, ...)where...are the dimensions ofxEmpty
NamedTuple()
Parameters
weight: Weight Matrix of size(out_dims, in_dims)bias: Bias of size(out_dims, 1)(present ifuse_bias=true)
Embedding(in_dims => out_dims; init_weight=randn32)A lookup table that stores embeddings of dimension out_dims for a vocabulary of size in_dims.
This layer is often used to store word embeddings and retrieve them using indices.
Warning
Unlike Flux.Embedding, this layer does not support using OneHotArray as an input.
Arguments
in_dims: number of input dimensionsout_dims: number of output dimensions
Keyword Arguments
init_weight: initializer for the weight matrix (weight = init_weight(rng, out_dims, in_dims))
Input
Integer OR
Abstract Vector of Integers OR
Abstract Array of Integers
Returns
Returns the embedding corresponding to each index in the input. For an N dimensional input, an N + 1 dimensional output is returned.
Empty
NamedTuple()
Scale(dims, activation=identity; init_weight=ones32, init_bias=zeros32, bias::Bool=true)Create a Sparsely Connected Layer with a very specific structure (only Diagonal Elements are non-zero). The forward pass is given by: y = activation.(weight .* x .+ bias)
Arguments
dims: size of the learnable scale and bias parameters.activation: activation function
Keyword Arguments
init_weight: initializer for the weight matrix (weight = init_weight(rng, out_dims, in_dims))init_bias: initializer for the bias vector (ignored ifuse_bias=false)use_bias: Trainable bias can be disabled entirely by setting this tofalseallow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)
Input
xmust be an Array of size(dims..., B)or(dims...[0], ..., dims[k])fork ≤ size(dims)
Returns
Array of size
(dims..., B)or(dims...[0], ..., dims[k])fork ≤ size(dims)Empty
NamedTuple()
Parameters
weight: Weight Array of size(dims...)bias: Bias of size(dims...)
Misc. Helper Layers
FlattenLayer(N = nothing)Flattens the passed array into a matrix.
Arguments
N: Flatten the firstNdimensions of the input array. Ifnothing, then all dimensions (except) are flattened. Note that the batch dimension is never flattened.
Inputs
x: AbstractArray
Returns
AbstractMatrix of size
(:, size(x, ndims(x)))Empty
NamedTuple()
Maxout(layers...)
Maxout(; layers...)
Maxout(f::Function, n_alts::Int)This contains a number of internal layers, each of which receives the same input. Its output is the elementwise maximum of the the internal layers' outputs.
Maxout over linear dense layers satisfies the univeral approximation theorem. See [1].
See also Parallel to reduce with other operators.
Arguments
- Layers can be specified in three formats:
A list of
NLux layersSpecified as
Nkeyword arguments.A no argument function
fand an integern_altswhich specifies the number of layers.
Inputs
x: Input that is passed to each of the layers
Returns
Output is computed by taking elementwise
maxof the outputs of the individual layers.Updated state of the
layers
Parameters
- Parameters of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
States
- States of each
layerwrapped in a NamedTuple withfields = layer_1, layer_2, ..., layer_N(naming changes if using the kwargs API)
References
[1] Goodfellow, Warde-Farley, Mirza, Courville & Bengio "Maxout Networks" https://arxiv.org/abs/1302.4389
NoOpLayer()As the name suggests does nothing but allows pretty printing of layers. Whatever input is passed is returned.
ReshapeLayer(dims)Reshapes the passed array to have a size of (dims..., :)
Arguments
dims: The new dimensions of the array (excluding the last dimension).
Inputs
x: AbstractArray of any shape which can be reshaped in(dims..., size(x, ndims(x)))
Returns
AbstractArray of size
(dims..., size(x, ndims(x)))Empty
NamedTuple()
SelectDim(dim, i)Return a view of all the data of the input x where the index for dimension dim equals i. Equivalent to view(x,:,:,...,i,:,:,...) where i is in position d.
Arguments
dim: Dimension for indexingi: Index for dimensiondim
Inputs
x: AbstractArray that can be indexed withview(x,:,:,...,i,:,:,...)
Returns
view(x,:,:,...,i,:,:,...)whereiis in positiondEmpty
NamedTuple()
WrappedFunction(f)Wraps a stateless and parameter less function. Might be used when a function is added to Chain. For example, Chain(x -> relu.(x)) would not work and the right thing to do would be Chain((x, ps, st) -> (relu.(x), st)). An easier thing to do would be Chain(WrappedFunction(Base.Fix1(broadcast, relu)))
Arguments
f::Function: A stateless and parameterless function
Inputs
x: s.thasmethod(f, (typeof(x),))istrue
Returns
Output of
f(x)Empty
NamedTuple()
Normalization Layers
BatchNorm(chs::Integer, activation=identity; init_bias=zeros32, init_scale=ones32,
affine=true, track_stats=true, epsilon=1f-5, momentum=0.1f0,
allow_fast_activation::Bool=true)Batch Normalization layer.
BatchNorm computes the mean and variance for each
Arguments
chs: Size of the channel dimension in your data. Given an array withNdimensions, call theN-1th the channel dimension. For a batch of feature vectors this is just the data dimension, forWHCNimages it's the usual channel dimension.activation: After normalization, elementwise activationactivationis applied.
Keyword Arguments
If
track_stats=true, accumulates mean and variance statistics in training phase that will be used to renormalize the input in test phase.epsilon: a value added to the denominator for numerical stabilitymomentum: the value used for therunning_meanandrunning_varcomputationallow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)If
affine=true, it also applies a shift and a rescale to the input through to learnable per-channel bias and scale parameters.init_bias: Controls how thebiasis initiliazedinit_scale: Controls how thescaleis initiliazed
Inputs
x: Array wheresize(x, N - 1) = chsandndims(x) > 2
Returns
y: Normalized ArrayUpdate model state
Parameters
affine=truebias: Bias of shape(chs,)scale: Scale of shape(chs,)
affine=false- EmptyNamedTuple()
States
Statistics if
track_stats=truerunning_mean: Running mean of shape(chs,)running_var: Running variance of shape(chs,)
Statistics if
track_stats=falserunning_mean: nothingrunning_var: nothing
training: Used to check if training/inference mode
Use Lux.testmode during inference.
Example
m = Chain(Dense(784 => 64), BatchNorm(64, relu), Dense(64 => 10), BatchNorm(10))Warning
Passing a batch size of 1, during training will result in NaNs.
See also BatchNorm, InstanceNorm, LayerNorm, WeightNorm
GroupNorm(chs::Integer, groups::Integer, activation=identity; init_bias=zeros32,
init_scale=ones32, affine=true, epsilon=1f-5,
allow_fast_activation::Bool=true)Group Normalization layer.
Arguments
chs: Size of the channel dimension in your data. Given an array withNdimensions, call theN-1th the channel dimension. For a batch of feature vectors this is just the data dimension, forWHCNimages it's the usual channel dimension.groupsis the number of groups along which the statistics are computed. The number of channels must be an integer multiple of the number of groups.activation: After normalization, elementwise activationactivationis applied.
Keyword Arguments
epsilon: a value added to the denominator for numerical stabilityallow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)If
affine=true, it also applies a shift and a rescale to the input through to learnable per-channel bias and scale parameters.init_bias: Controls how thebiasis initiliazedinit_scale: Controls how thescaleis initiliazed
Inputs
x: Array wheresize(x, N - 1) = chsandndims(x) > 2
Returns
y: Normalized ArrayUpdate model state
Parameters
affine=truebias: Bias of shape(chs,)scale: Scale of shape(chs,)
affine=false- EmptyNamedTuple()
States
training: Used to check if training/inference mode
Use Lux.testmode during inference.
Example
m = Chain(Dense(784 => 64), GroupNorm(64, 4, relu), Dense(64 => 10), GroupNorm(10, 5))See also GroupNorm, InstanceNorm, LayerNorm, WeightNorm
InstanceNorm(chs::Integer, activation=identity; init_bias=zeros32, init_scale=ones32,
affine=true, epsilon=1f-5, allow_fast_activation::Bool=true)Instance Normalization. For details see [1].
Instance Normalization computes the mean and variance for each
Arguments
chs: Size of the channel dimension in your data. Given an array withNdimensions, call theN-1th the channel dimension. For a batch of feature vectors this is just the data dimension, forWHCNimages it's the usual channel dimension.activation: After normalization, elementwise activationactivationis applied.
Keyword Arguments
epsilon: a value added to the denominator for numerical stabilityallow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)If
affine=true, it also applies a shift and a rescale to the input through to learnable per-channel bias and scale parameters.init_bias: Controls how thebiasis initiliazedinit_scale: Controls how thescaleis initiliazed
Inputs
x: Array wheresize(x, N - 1) = chsandndims(x) > 2
Returns
y: Normalized ArrayUpdate model state
Parameters
affine=truebias: Bias of shape(chs,)scale: Scale of shape(chs,)
affine=false- EmptyNamedTuple()
States
training: Used to check if training/inference mode
Use Lux.testmode during inference.
Example
m = Chain(Dense(784 => 64), InstanceNorm(64, relu), Dense(64 => 10), InstanceNorm(10, 5))References
[1] Ulyanov, Dmitry, Andrea Vedaldi, and Victor Lempitsky. "Instance normalization: The missing ingredient for fast stylization." arXiv preprint arXiv:1607.08022 (2016).
See also BatchNorm, GroupNorm, LayerNorm, WeightNorm
LayerNorm(shape::NTuple{N, Int}, activation=identity; epsilon=1f-5, dims=Colon(),
affine::Bool=true, init_bias=zeros32, init_scale=ones32,)Computes mean and standard deviation over the whole input array, and uses these to normalize the whole array. Optionally applies an elementwise affine transformation afterwards.
Given an input array
where affine=true.
Warning
As of v0.5.0, the doc used to say affine::Bool=false, but the code actually had affine::Bool=true as the default. Now the doc reflects the code, so please check whether your assumptions about the default (if made) were invalid.
Arguments
shape: Broadcastable shape of input array excluding the batch dimension.activation: After normalization, elementwise activationactivationis applied.
Keyword Arguments
allow_fast_activation: Iftrue, then certain activations can be approximated with a faster version. The new activation function will be given byNNlib.fast_act(activation)epsilon: a value added to the denominator for numerical stability.dims: Dimensions to normalize the array over.If
affine=true, it also applies a shift and a rescale to the input through to learnable per-channel bias and scale parameters.init_bias: Controls how thebiasis initiliazedinit_scale: Controls how thescaleis initiliazed
Inputs
x: AbstractArray
Returns
y: Normalized ArrayEmpty NamedTuple()
Parameters
affine=false: EmptyNamedTuple()affine=truebias: Bias of shape(shape..., 1)scale: Scale of shape(shape..., 1)
WeightNorm(layer::AbstractExplicitLayer, which_params::NTuple{N,Symbol},
dims::Union{Tuple,Nothing}=nothing)Applies weight normalization to a parameter in the given layer.
Weight normalization is a reparameterization that decouples the magnitude of a weight tensor from its direction. This updates the parameters in which_params (e.g. weight) using two parameters: one specifying the magnitude (e.g. weight_g) and one specifying the direction (e.g. weight_v).
Arguments
layerwhose parameters are being reparameterizedwhich_params: parameter names for the parameters being reparameterizedBy default, a norm over the entire array is computed. Pass
dimsto modify the dimension.
Inputs
x: Should be of valid type for input tolayer
Returns
Output from
layerUpdated model state of
layer
Parameters
normalized: Parameters oflayerthat are being normalizedunnormalized: Parameters oflayerthat are not being normalized
States
- Same as that of
layer
Upsampling
PixelShuffle(r::Int)Pixel shuffling layer with upscale factor r. Usually used for generating higher resolution images while upscaling them.
See NNlib.pixel_shuffle for more details.
PixelShuffle is not a Layer, rather it returns a WrappedFunction with the function set to Base.Fix2(pixel_shuffle, r)
Arguments
r: Upscale factor
Inputs
x: For 4D-arrays representing N images, the operation converts inputsize(x) == (W, H, r² x C, N)to output of size(r x W, r x H, C, N). For D-dimensional data, it expectsndims(x) == D + 2with channel and batch dimensions, and divides the number of channels byrᴰ.
Returns
- Output of size
(r x W, r x H, C, N)for 4D-arrays, and(r x W, r x H, ..., C, N)for D-dimensional data, whereD = ndims(x) - 2
Upsample(mode = :nearest; [scale, size])
Upsample(scale, mode = :nearest)Upsampling Layer.
Layer Construction
Option 1
mode: Set to:nearest,:linear,:bilinearor:trilinear
Exactly one of two keywords must be specified:
If
scaleis a number, this applies to all but the last two dimensions (channel and batch) of the input. It may also be a tuple, to control dimensions individually.Alternatively, keyword
sizeaccepts a tuple, to directly specify the leading dimensions of the output.
Option 2
If
scaleis a number, this applies to all but the last two dimensions (channel and batch) of the input. It may also be a tuple, to control dimensions individually.mode: Set to:nearest,:bilinearor:trilinear
Currently supported upsampling modes and corresponding NNlib's methods are:
:nearest->NNlib.upsample_nearest:bilinear->NNlib.upsample_bilinear:trilinear->NNlib.upsample_trilinear
Inputs
x: For the input dimensions look into the documentation for the correspondingNNlibfunctionAs a rule of thumb,
:nearestshould work with arrays of arbitrary dimensions:bilinearworks with 4D Arrays:trilinearworks with 5D Arrays
Returns
Upsampled Input of size
sizeor of size(I_1 x scale[1], ..., I_N x scale[N], C, N)Empty
NamedTuple()