oneflow.nn¶
These are the basic building blocks for graphs:
oneflow.nn
Containers¶
Base class for all neural network modules. |
|
A sequential container. |
|
Holds submodules in a list. |
|
Holds submodules in a dictionary. |
|
Holds parameters in a list. |
|
Holds parameters in a dictionary. |
nn.Module¶
Adds a child module to the current module. |
|
Applies |
|
Returns an iterator over module buffers. |
|
Returns an iterator over immediate children modules. |
|
Moves all model parameters and buffers to the CPU. |
|
Moves all model parameters and buffers to the GPU. |
|
Casts all floating point parameters and buffers to |
|
Sets the module in training mode. |
|
Sets the module in evaluation mode. |
|
Set the extra representation of the module |
|
Casts all floating point parameters and buffers to |
|
Copies parameters and buffers from |
|
Returns an iterator over all modules in the network. |
|
Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself. |
|
Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself. |
|
Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself. |
|
Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself. |
|
Returns an iterator over module parameters. |
|
Adds a buffer to the module. |
|
Registers a forward hook on the module. |
|
Registers a forward pre-hook on the module. |
|
Registers a backward hook on the module. |
|
Registers a backward hook on the module. |
|
These hooks will be called with arguments: |
|
Adds a parameter to the module. |
|
Change if autograd should record operations on parameters in this module. |
|
Returns a dictionary containing a whole state of the module. |
|
Moves and/or casts the parameters and buffers. |
|
Sets gradients of all model parameters to zero. |
Containers
Convolution Layers¶
Applies a 1D convolution over an input signal composed of several input planes. |
|
Applies a 2D convolution over an input signal composed of several input planes. |
|
Applies a 3D convolution over an input signal composed of several input planes. |
|
Applies a 1D transposed convolution operator over an input image composed of several input planes. |
|
Applies a 2D transposed convolution operator over an input image composed of several input planes. |
|
Applies a 3D transposed convolution operator over an input image composed of several input planes. |
|
The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.nn.Unfold.html. |
|
The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.nn.Fold.html. |
Pooling Layers¶
Applies a 1D max pooling over an input signal composed of several input planes. |
|
Applies a 2D max pooling over an input signal composed of several input planes. |
|
Applies a 3D max pooling over an input signal composed of several input planes. |
|
Computes a partial inverse of |
|
Computes a partial inverse of |
|
Computes a partial inverse of |
|
Applies a 1D adaptive average pooling over an input signal composed of several input planes. |
|
Applies a 2D adaptive average pooling over an input signal composed of several input planes. |
|
Applies a 3D adaptive average pooling over an input signal composed of several input planes. |
|
Applies a 1D adaptive max pooling over an input signal composed of several input planes. |
|
Applies a 2D adaptive max pooling over an input signal composed of several input planes. |
|
Applies a 3D adaptive max pooling over an input signal composed of several input planes. |
|
Applies a 1D average pooling over an input signal composed of several input planes. |
|
Performs the 2d-average pooling on the input. |
|
Applies a 3D average pooling over an input signal composed of several input planes. |
Padding Layers¶
Pads the input tensor boundaries with a constant value. |
|
This operator pads the input with constant value that user specifies. |
|
Pads the input tensor boundaries with a constant value. |
|
This operator pads the input tensor using the reflection of the input boundary. |
|
This operator pads the input tensor using the reflection of the input boundary. |
|
Pads the input tensor using replication of the input boundary. |
|
Pads the input tensor using the replication of the input boundary. |
|
Pads the input tensor boundaries with zero. |
Non-linear Activations (weighted sum, nonlinearity)¶
Applies the element-wise function |
|
The Hardshrink activation. |
|
Applies the element-wise function: |
|
Applies the hardswish function, element-wise, as described in the paper Searching for MobileNetV3. |
|
Applies the HardTanh function element-wise |
|
Applies the element-wise function: |
|
Applies the element-wise function: |
|
Applies the element-wise function: |
|
Applies the rectified linear unit function element-wise: |
|
Applies the element-wise function: |
|
Applies the element-wise function: |
|
Applies the element-wise function: |
|
The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.nn.GELU.html. |
|
Applies GELU approximation that is fast but somewhat inaccurate. |
|
SiLU(Swish) activation: |
|
Applies the element-wise function: |
|
Applies the element-wise function: |
|
Applies the element-wise function: |
|
The Softshrink activation. |
|
The SoftSign activation. |
|
This operator computes the hyperbolic tangent value of Tensor. |
|
The Threshold Activation. |
|
The GLU activation. |
Non-linear Activations (other)¶
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. |
|
Applies the LogSoftmax function to an n-dimensional input Tensor. |
Normalization Layers¶
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 . |
|
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 . |
|
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 . |
|
Applies Batch Normalization over a N-Dimensional input (a mini-batch of [N-2]D inputs with additional channel dimension) as described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift . |
|
Applies Fused Batch Normalization over a 2D or 3D input, the formula is: |
|
Applies Fused Batch Normalization over a 4D input, the formula is: |
|
Applies Fused Batch Normalization over a 5D input, the formula is: |
|
Applies Group Normalization over a mini-batch of inputs as described in the paper Group Normalization |
|
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. |
|
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. |
|
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. |
|
Applies Layer Normalization over a mini-batch of inputs as described in the paper Layer Normalization |
|
Construct a layernorm module in the T5 style. |
|
Applies Root Mean Square Layer Normalization over a mini-batch of inputs as described in the paper Root Mean Square Layer Normalization |
Recurrent Layers¶
Applies a multi-layer Elman RNN with tanhtanh or text{ReLU}ReLU non-linearity to an input sequence. |
|
Applies a multi-layer long short-term memory (LSTM) RNN to an input sequence. |
|
Applies a multi-layer gated recurrent unit (GRU) RNN to an input sequence. |
|
An Elman RNN cell with tanh or ReLU non-linearity. |
|
A long short-term memory (LSTM) cell. |
|
A gated recurrent unit (GRU) cell |
Linear Layers¶
A placeholder identity operator that is argument-insensitive. |
|
Applies a linear transformation to the incoming data: \(y = xA^T + b\) |
Dropout Layers¶
During training, randomly zeroes some of the elements of the input tensor with probability |
|
Randomly zero out entire channels (a channel is a 1D feature map, e.g., the \(j\)-th channel of the \(i\)-th sample in the batched input is a 1D tensor :math:` ext{input}[i, j]`). |
|
Randomly zero out entire channels (a channel is a 2D feature map, e.g., the \(j\)-th channel of the \(i\)-th sample in the batched input is a 2D tensor :math:` ext{input}[i, j]`). |
|
Randomly zero out entire channels (a channel is a 3D feature map, e.g., the \(j\)-th channel of the \(i\)-th sample in the batched input is a 3D tensor :math:` ext{input}[i, j]`). |
Sparse Layers¶
A simple lookup table that stores embeddings of a fixed dictionary and size. |
Distance Functions¶
Returns cosine similarity between \(x_1\) and \(x_2\), computed along dim. |
|
Computes the pairwise distance between vectors \(v_1\), \(v_2\) using the p-norm: |
Loss Functions¶
This operator computes the binary cross entropy loss. |
|
This operator combines the Sigmoid and BCELoss together. |
|
The Connectionist Temporal Classification loss. |
|
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. |
|
The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.nn.CrossEntropyLoss.html. |
|
The Kullback-Leibler divergence loss measure |
|
This operator computes the L1 Loss between each element in input and target. |
|
Creates a criterion that measures the mean squared error (squared L2 norm) between each element in the input \(x\) and target \(y\). |
|
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). |
|
The negative log likelihood loss. |
|
Creates a criterion that uses a squared term if the absolute element-wise error falls below beta and an L1 term otherwise. |
|
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¶
alias of |
|
Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D (volumetric) data. |
|
Applies a 2D bilinear upsampling to an input signal composed of several input channels. |
|
Applies a 2D nearest neighbor upsampling to an input signal composed of several input channels. |
Data loading and preprocessing Layers¶
Generates random boolean values following a bernoulli distribution. |
|
Performs fused cropping, normalization, format conversion (NHWC to NCHW) if desired, and type casting. |
|
This operator reads an tensor as bytes. |
|
Quantization Aware Training¶
Compute the quantization parameters of the input tensor. |
|
Compute the quantization parameters based on the moving average of the input tensor’s min and max values. |
|
Simulate the quantize and dequantize operations in training time. |
|
A Conv1d module attached with nn.MinMaxObserver, nn.MovingAverageMinMaxObserver and nn.FakeQuantization modules for weight and input, used for quantization aware training. |
|
A Conv2d module attached with nn.MinMaxObserver, nn.MovingAverageMinMaxObserver and nn.FakeQuantization modules for weight and input, used for quantization aware training. |
|
A Conv3d module attached with nn.MinMaxObserver, nn.MovingAverageMinMaxObserver and nn.FakeQuantization modules for weight and input, used for quantization aware training. |
Utilities¶
From the oneflow.nn.utils
module
Clips gradient norm of an iterable of parameters. |
|
Clips gradient of an iterable of parameters at specified value. |
|
Applies weight normalization to a parameter in the given module. |
|
Removes the weight normalization reparameterization from a module. |
Utility functions in other modules
The interface is consistent with PyTorch. |
|
The interface is consistent with PyTorch. |
|
The interface is consistent with PyTorch. |
|
The interface is consistent with PyTorch. |
|
Packs a list of variable length Tensors |
Flattens a contiguous range of dims into a tensor. |
Quantized Functions¶
Quantization refers to techniques for performing computations and storing tensors at lower bitwidths than floating point precision.
Simulate the quantize and dequantize operations in training time. |
|
Compute the quantization parameters of the input tensor. |
|
Compute the quantization parameters based on the moving average of the input tensor’s min and max values. |
|
Simulate the quantize operation in inference time. |