oneflow.nn

These are the basic building blocks for graphs:

Containers

Module

Base class for all neural network modules.

Sequential

A sequential container.

ModuleList

Holds submodules in a list.

ModuleDict

Holds submodules in a dictionary.

ParameterList

Holds parameters in a list.

ParameterDict

Holds parameters in a dictionary.

Convolution Layers

nn.Conv1d

Applies a 1D convolution over an input signal composed of several input planes.

nn.Conv2d

Applies a 2D convolution over an input signal composed of several input planes.

nn.Conv3d

Applies a 3D convolution over an input signal composed of several input planes.

nn.ConvTranspose1d

Applies a 1D transposed convolution operator over an input image composed of several input planes.

nn.ConvTranspose2d

Applies a 2D transposed convolution operator over an input image composed of several input planes.

nn.ConvTranspose3d

Applies a 3D transposed convolution operator over an input image composed of several input planes.

nn.Unfold

This op extracts elements in a local window from input tensor, it also called img2col.

nn.Fold

Combines an array of sliding local blocks into a large containing tensor, it also called col2img.

Pooling Layers

nn.MaxPool1d

Applies a 1D max pooling over an input signal composed of several input planes.

nn.MaxPool2d

Applies a 2D max pooling over an input signal composed of several input planes.

nn.MaxPool3d

Applies a 3D max pooling over an input signal composed of several input planes.

nn.AdaptiveAvgPool1d

Applies a 1D adaptive average pooling over an input signal composed of several input planes.

nn.AdaptiveAvgPool2d

Applies a 2D adaptive average pooling over an input signal composed of several input planes.

nn.AdaptiveAvgPool3d

Applies a 3D adaptive average pooling over an input signal composed of several input planes.

nn.AvgPool1d

Applies a 1D average pooling over an input signal composed of several input planes.

nn.AvgPool2d

Performs the 2d-average pooling on the input.

nn.AvgPool3d

Applies a 3D average pooling over an input signal composed of several input planes.

Padding Layers

nn.ConstantPad1d

Pads the input tensor boundaries with a constant value.

nn.ConstantPad2d

This operator pads the input with constant value that user specifies.

nn.ConstantPad3d

Pads the input tensor boundaries with a constant value.

nn.ReflectionPad2d

This operator pads the input tensor using the reflection of the input boundary.

nn.ReplicationPad2d

Pads the input tensor using the replication of the input boundary.

nn.ZeroPad2d

Pads the input tensor boundaries with zero.

Non-linear Activations (weighted sum, nonlinearity)

nn.ELU

Applies the element-wise function

nn.Hardshrink

The Hardshrink activation.

nn.Hardsigmoid

Applies the element-wise function:

nn.Hardswish

Applies the hardswish function, element-wise, as described in the paper Searching for MobileNetV3.

nn.Hardtanh

Applies the HardTanh function element-wise

nn.LeakyReLU

Applies the element-wise function:

nn.LogSigmoid

Applies the element-wise function:

nn.PReLU

Applies the element-wise function:

nn.ReLU

Applies the rectified linear unit function element-wise:

nn.ReLU6

Applies the element-wise function:

nn.SELU

Applies the element-wise function:

nn.CELU

Applies the element-wise function:

nn.GELU

Gelu activation operator.

nn.SiLU

SiLU(Swish) activation:

nn.Sigmoid

Applies the element-wise function:

nn.Mish

Applies the element-wise function:

nn.Softplus

Applies the element-wise function:

nn.Softshrink

The Softshrink activation.

nn.Softsign

The SoftSign activation.

nn.Tanh

This operator computes the hyperbolic tangent value of Tensor.

nn.Threshold

The Threshold Activation.

nn.GLU

The GLU activation.

Non-linear Activations (other)

nn.Softmax

Applies the Softmax function to an n-dimensional input Tensor rescaling them so that the elements of the n-dimensional output Tensor lie in the range [0,1] and sum to 1.

nn.LogSoftmax

Applies the LogSoftmax function to an n-dimensional input Tensor.

Normalization Layers

nn.BatchNorm1d

Applies Batch Normalization over a 2D or 3D input (a mini-batch of 1D inputs with optional additional channel dimension) as described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .

nn.BatchNorm2d

Applies Batch Normalization over a 4D input (a mini-batch of 2D inputs with additional channel dimension) as described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .

nn.BatchNorm3d

Applies Batch Normalization over a 5D input (a mini-batch of 3D inputs with additional channel dimension) as described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .

nn.FusedBatchNorm1d

Applies Fused Batch Normalization over a 2D or 3D input, the formula is:

nn.FusedBatchNorm2d

Applies Fused Batch Normalization over a 4D input, the formula is:

nn.FusedBatchNorm3d

Applies Fused Batch Normalization over a 5D input, the formula is:

nn.GroupNorm

Applies Group Normalization over a mini-batch of inputs as described in the paper Group Normalization

nn.InstanceNorm1d

Applies Instance Normalization over a 3D input (a mini-batch of 1D inputs with optional additional channel dimension) as described in the paper Instance Normalization: The Missing Ingredient for Fast Stylization.

nn.InstanceNorm2d

Applies Instance Normalization over a 4D input (a mini-batch of 2D inputs with additional channel dimension) as described in the paper Instance Normalization: The Missing Ingredient for Fast Stylization.

nn.InstanceNorm3d

Applies Instance Normalization over a 5D input (a mini-batch of 3D inputs with additional channel dimension) as described in the paper Instance Normalization: The Missing Ingredient for Fast Stylization.

nn.LayerNorm

Applies Layer Normalization over a mini-batch of inputs as described in the paper Layer Normalization

Recurrent Layers

nn.RNN

Applies a multi-layer Elman RNN with tanhtanh or text{ReLU}ReLU non-linearity to an input sequence.

nn.LSTM

Applies a multi-layer long short-term memory (LSTM) RNN to an input sequence.

nn.GRU

Applies a multi-layer gated recurrent unit (GRU) RNN to an input sequence.

nn.RNNCell

An Elman RNN cell with tanh or ReLU non-linearity.

nn.LSTMCell

A long short-term memory (LSTM) cell.

nn.GRUCell

A gated recurrent unit (GRU) cell

Linear Layers

nn.Identity

A placeholder identity operator that is argument-insensitive.

nn.Linear

Applies a linear transformation to the incoming data: \(y = xA^T + b\)

Dropout Layers

nn.Dropout

During training, randomly zeroes some of the elements of the input tensor with probability p using samples from a Bernoulli distribution.

Sparse Layers

nn.Embedding

A simple lookup table that stores embeddings of a fixed dictionary and size.

Distance Functions

nn.CosineSimilarity

Returns cosine similarity between \(x_1\) and \(x_2\), computed along dim.

Loss Functions

nn.BCELoss

This operator computes the binary cross entropy loss.

nn.BCEWithLogitsLoss

This operator combines the Sigmoid and BCELoss together.

nn.CTCLoss

The Connectionist Temporal Classification loss.

nn.CombinedMarginLoss

The operation implements “margin_softmax” in InsightFace: https://github.com/deepinsight/insightface/blob/master/recognition/arcface_mxnet/train.py The implementation of margin_softmax in InsightFace is composed of multiple operators.

nn.CrossEntropyLoss

This criterion combines LogSoftmax and NLLLoss in one single class.

nn.KLDivLoss

The Kullback-Leibler divergence loss measure

nn.L1Loss

This operator computes the L1 Loss between each element in input and target.

nn.MSELoss

Creates a criterion that measures the mean squared error (squared L2 norm) between each element in the input \(x\) and target \(y\).

nn.MarginRankingLoss

Creates a criterion that measures the loss given inputs \(x1\), \(x2\), two 1D mini-batch Tensors, and a label 1D mini-batch tensor \(y\) (containing 1 or -1).

nn.NLLLoss

The negative log likelihood loss.

nn.SmoothL1Loss

Creates a criterion that uses a squared term if the absolute element-wise error falls below beta and an L1 term otherwise.

nn.TripletMarginLoss

Creates a criterion that measures the triplet loss given an input tensors \(x1\), \(x2\), \(x3\) and a margin with a value greater than \(0\).

Vision Layers

nn.PixelShuffle

alias of oneflow.nn.modules.pixelshuffle.PixelShufflev2

nn.Upsample

Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D (volumetric) data.

nn.UpsamplingBilinear2d

Applies a 2D bilinear upsampling to an input signal composed of several input channels.

nn.UpsamplingNearest2d

Applies a 2D nearest neighbor upsampling to an input signal composed of several input channels.

Utilities

From the oneflow.nn.utils module

clip_grad_norm_

Clips gradient norm of an iterable of parameters.

clip_grad_value_

Clips gradient of an iterable of parameters at specified value.

weight_norm

Applies weight normalization to a parameter in the given module.

remove_weight_norm

Removes the weight normalization reparameterization from a module.

Utility functions in other modules

nn.utils.rnn.PackedSequence

The interface is consistent with PyTorch.

nn.utils.rnn.pack_padded_sequence

The interface is consistent with PyTorch.

nn.utils.rnn.pad_packed_sequence

The interface is consistent with PyTorch.

nn.utils.rnn.pad_sequence

The interface is consistent with PyTorch.

nn.utils.rnn.pack_sequence

Packs a list of variable length Tensors

nn.Flatten

Flattens a contiguous range of dims into a tensor.