Base class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:
Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.
training (bool) – Boolean represents whether this module is in training or evaluation mode.
Adds a child module to the current module.
The module can be accessed as an attribute using the given name.
Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also torch.nn.init).
fn (Module -> None) – function to be applied to each submodule
self
Module
Example:
Casts all floating point parameters and buffers to bfloat16 datatype.
self
Module
Returns an iterator over module buffers.
recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.
torch.Tensor – module buffer
Example:
Returns an iterator over immediate children modules.
Module – a child module
Moves all model parameters and buffers to the CPU.
self
Module
Moves all model parameters and buffers to the GPU.
This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.
device (int, optional) – if specified, all parameters will be copied to that device
self
Module
Casts all floating point parameters and buffers to double datatype.
self
Module
This allows better BC support for load_state_dict(). In state_dict(), the version number will be saved as in the attribute _metadata of the returned state dict, and thus pickled. _metadata is a dictionary with keys that follow the naming convention of state dict. See _load_from_state_dict on how to use this information in loading.
If new parameters/buffers are added/removed from a module, this number shall be bumped, and the module’s _load_from_state_dict method can compare the version number and do appropriate changes if the state dict is from before the change.
Sets the module in evaluation mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.
This is equivalent with self.train(False).
self
Module
Set the extra representation of the module
To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.
Casts all floating point parameters and buffers to float datatype.
self
Module
Defines the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
Casts all floating point parameters and buffers to half datatype.
self
Module
Copies parameters and buffers from state_dict into this module and its descendants. If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.
NamedTuple with missing_keys and unexpected_keys fields
Returns an iterator over all modules in the network.
Module – a module in the network
Note
Duplicate modules are returned only once. In the following example, l will be returned only once.
Example:
Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
(string, torch.Tensor) – Tuple containing the name and buffer
Example:
Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
(string, Module) – Tuple containing a name and child module
Example:
Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
(string, Module) – Tuple of name and module
Note
Duplicate modules are returned only once. In the following example, l will be returned only once.
Example:
Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
(string, Parameter) – Tuple containing the name and parameter
Example:
Returns an iterator over module parameters.
This is typically passed to an optimizer.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Parameter – module parameter
Example:
Registers a backward hook on the module.
This function is deprecated in favor of nn.Module.register_full_backward_hook() and the behavior of this function will change in future versions.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle
Adds a buffer to the module.
This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.
Buffers can be accessed as attributes using given names.
Example:
Registers a forward hook on the module.
The hook will be called every time after forward() has computed an output. It should have the following signature:
The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle
Registers a forward pre-hook on the module.
The hook will be called every time before forward() is invoked. It should have the following signature:
The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle
Registers a backward hook on the module.
The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature:
The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.
Warning
Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle
Adds a parameter to the module.
The parameter can be accessed as an attribute using given name.
Change if autograd should record operations on parameters in this module.
This method sets the parameters’ requires_grad attributes in-place.
This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).
requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.
self
Module
Returns a dictionary containing a whole state of the module.
Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names.
a dictionary containing a whole state of the module
dict
Example:
Moves and/or casts the parameters and buffers.
This can be called as
Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtype`s. In addition, this method will
only cast the floating point or complex parameters and buffers to :attr:`dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.
See below for examples.
Note
This method modifies the module in-place.
self
Module
Examples:
Sets the module in training mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.
mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.
self
Module
Casts all parameters and buffers to dst_type.
dst_type (type or string) – the desired type
self
Module
Moves all model parameters and buffers to the XPU.
This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.
device (int, optional) – if specified, all parameters will be copied to that device
self
Module
Sets gradients of all model parameters to zero. See similar function under torch.optim.Optimizer for more context.
set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.


Adds a child module to the current module.
The module can be accessed as an attribute using the given name.


Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also torch.nn.init).
fn (Module -> None) – function to be applied to each submodule
self
Module
Example:


Casts all floating point parameters and buffers to bfloat16 datatype.
self
Module


Returns an iterator over module buffers.
recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.
torch.Tensor – module buffer
Example:


Returns an iterator over immediate children modules.
Module – a child module


Moves all model parameters and buffers to the CPU.
self
Module


Moves all model parameters and buffers to the GPU.
This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.
device (int, optional) – if specified, all parameters will be copied to that device
self
Module


Casts all floating point parameters and buffers to double datatype.
self
Module


This allows better BC support for load_state_dict(). In state_dict(), the version number will be saved as in the attribute _metadata of the returned state dict, and thus pickled. _metadata is a dictionary with keys that follow the naming convention of state dict. See _load_from_state_dict on how to use this information in loading.
If new parameters/buffers are added/removed from a module, this number shall be bumped, and the module’s _load_from_state_dict method can compare the version number and do appropriate changes if the state dict is from before the change.


Sets the module in evaluation mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.
This is equivalent with self.train(False).
self
Module


Set the extra representation of the module
To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.


Casts all floating point parameters and buffers to float datatype.
self
Module


Defines the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.


Casts all floating point parameters and buffers to half datatype.
self
Module


Copies parameters and buffers from state_dict into this module and its descendants. If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.
NamedTuple with missing_keys and unexpected_keys fields


Returns an iterator over all modules in the network.
Module – a module in the network
Note
Duplicate modules are returned only once. In the following example, l will be returned only once.
Example:


Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
(string, torch.Tensor) – Tuple containing the name and buffer
Example:


Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
(string, Module) – Tuple containing a name and child module
Example:


Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
(string, Module) – Tuple of name and module
Note
Duplicate modules are returned only once. In the following example, l will be returned only once.
Example:


Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
(string, Parameter) – Tuple containing the name and parameter
Example:


Returns an iterator over module parameters.
This is typically passed to an optimizer.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Parameter – module parameter
Example:


Registers a backward hook on the module.
This function is deprecated in favor of nn.Module.register_full_backward_hook() and the behavior of this function will change in future versions.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle


Adds a buffer to the module.
This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.
Buffers can be accessed as attributes using given names.
Example:


Registers a forward hook on the module.
The hook will be called every time after forward() has computed an output. It should have the following signature:
The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle


Registers a forward pre-hook on the module.
The hook will be called every time before forward() is invoked. It should have the following signature:
The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle


Registers a backward hook on the module.
The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature:
The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.
Warning
Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle


Adds a parameter to the module.
The parameter can be accessed as an attribute using given name.


Change if autograd should record operations on parameters in this module.
This method sets the parameters’ requires_grad attributes in-place.
This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).
requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.
self
Module


Returns a dictionary containing a whole state of the module.
Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names.
a dictionary containing a whole state of the module
dict
Example:


Moves and/or casts the parameters and buffers.
This can be called as
Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtype`s. In addition, this method will
only cast the floating point or complex parameters and buffers to :attr:`dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.
See below for examples.
Note
This method modifies the module in-place.
self
Module
Examples:


Sets the module in training mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.
mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.
self
Module


Casts all parameters and buffers to dst_type.
dst_type (type or string) – the desired type
self
Module


Moves all model parameters and buffers to the XPU.
This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.
device (int, optional) – if specified, all parameters will be copied to that device
self
Module


Sets gradients of all model parameters to zero. See similar function under torch.optim.Optimizer for more context.
set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.


Holds submodules in a dictionary.
ModuleDict can be indexed like a regular Python dictionary, but modules it contains are properly registered, and will be visible by all Module methods.
ModuleDict is an ordered dictionary that respects
Note that update() with other unordered mapping types (e.g., Python’s plain dict before Python version 3.6) does not preserve the order of the merged mapping.
modules (iterable, optional) – a mapping (dictionary) of (string: module) or an iterable of key-value pairs of type (string, module)
Example:
Remove all items from the ModuleDict.
Return an iterable of the ModuleDict key/value pairs.
Return an iterable of the ModuleDict keys.
Remove key from the ModuleDict and return its module.
key (string) – key to pop from the ModuleDict
Update the ModuleDict with the key-value pairs from a mapping or an iterable, overwriting existing keys.
Note
If modules is an OrderedDict, a ModuleDict, or an iterable of key-value pairs, the order of new elements in it is preserved.
modules (iterable) – a mapping (dictionary) from string to Module, or an iterable of key-value pairs of type (string, Module)
Return an iterable of the ModuleDict values.


Return an iterable of the ModuleDict key/value pairs.


Return an iterable of the ModuleDict keys.


Remove key from the ModuleDict and return its module.
key (string) – key to pop from the ModuleDict


Update the ModuleDict with the key-value pairs from a mapping or an iterable, overwriting existing keys.
Note
If modules is an OrderedDict, a ModuleDict, or an iterable of key-value pairs, the order of new elements in it is preserved.
modules (iterable) – a mapping (dictionary) from string to Module, or an iterable of key-value pairs of type (string, Module)


Return an iterable of the ModuleDict values.


Holds submodules in a list.
ModuleList can be indexed like a regular Python list, but modules it contains are properly registered, and will be visible by all Module methods.
modules (iterable, optional) – an iterable of modules to add
Example:
Appends a given module to the end of the list.
module (nn.Module) – module to append
Appends modules from a Python iterable to the end of the list.
modules (iterable) – iterable of modules to append
Insert a given module before a given index in the list.


Appends a given module to the end of the list.
module (nn.Module) – module to append


Appends modules from a Python iterable to the end of the list.
modules (iterable) – iterable of modules to append


Insert a given module before a given index in the list.


A mixin for modules that lazily initialize parameters, also known as “lazy modules.”
Modules that lazily initialize parameters, or “lazy modules”, derive the shapes of their parameters from the first input(s) to their forward method. Until that first forward they contain torch.nn.UninitializedParameter`s that should not be accessed
or used, and afterward they contain regular :class:`torch.nn.Parameter`s.
Lazy modules are convenient since they don't require computing some
module arguments, like the `in_features argument of a typical torch.nn.Linear.
After construction, networks with lazy modules should first be converted to the desired dtype and placed on the desired device. The lazy modules should then be initialized with one or more “dry runs”. These “dry runs” send inputs of the correct size, dtype, and device through the network and to each one of its lazy modules. After this the network can be used as usual.
A final caveat when using lazy modules is that the order of initialization of a network’s parameters may change, since the lazy modules are always initialized after other modules. This can cause the parameters of a network using lazy modules to be initialized differently than the parameters of a network without lazy modules. For example, if the LazyMLP class defined above had a torch.nn.LazyLinear module first and then a regular torch.nn.Linear second, the second module would be initialized on construction and the first module would be initialized during the first dry run.
Lazy modules can be serialized with a state dict like other modules. For example:
Lazy modules can also load regular torch.nn.Parameter s, which replace their torch.nn.UninitializedParameter s:
Note, however, that lazy modules cannot validate that the shape of parameters they load is correct.
Check if a module has parameters that are not initialized
Initialize parameters according to the input batch properties. This adds an interface to isolate parameter initialization from the forward pass when doing parameter shape inference.


Check if a module has parameters that are not initialized


Initialize parameters according to the input batch properties. This adds an interface to isolate parameter initialization from the forward pass when doing parameter shape inference.


Registers a backward hook common to all the modules.
This function is deprecated in favor of nn.module.register_module_full_backward_hook() and the behavior of this function will change in future versions.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle


Registers a global forward hook for all the modules
Warning
This adds global state to the nn.module module and it is only intended for debugging/profiling purposes.
The hook will be called every time after forward() has computed an output. It should have the following signature:
The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle
This hook will be executed before specific module hooks registered with register_forward_hook.


Registers a forward pre-hook common to all modules.
Warning
This adds global state to the nn.module module and it is only intended for debugging/profiling purposes.
The hook will be called every time before forward() is invoked. It should have the following signature:
The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).
This hook has precedence over the specific module hooks registered with register_forward_pre_hook.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle


Creates a criterion that measures the mean squared error (squared L2 norm) between each element in the input xx  and target yy .
The unreduced (i.e. with reduction set to 'none') loss can be described as:
where NN  is the batch size. If reduction is not 'none' (default 'mean'), then:
xx  and yy  are tensors of arbitrary shapes with a total of nn  elements each.
The mean operation still operates over all the elements, and divides by nn .
The division by nn  can be avoided if one sets reduction = 'sum'.
Examples:


Allows the model to jointly attend to information from different representation subspaces. See Attention Is All You Need
where headi=Attention(QWiQ,KWiK,VWiV)head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) .
Note that if kdim and vdim are None, they will be set to embed_dim such that query, key, and value have the same number of features.
Examples:
attn_mask: if a 2D mask: (L,S)(L, S)  where L is the target sequence length, S is the source sequence length.
If a 3D mask: (N⋅num_heads,L,S)(N\cdot\text{num\_heads}, L, S)  where N is the batch size, L is the target sequence length, S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend while the zero positions will be unchanged. If a BoolTensor is provided, positions with True is not allowed to attend while False values will be unchanged. If a FloatTensor is provided, it will be added to the attention weight.


attn_mask: if a 2D mask: (L,S)(L, S)  where L is the target sequence length, S is the source sequence length.
If a 3D mask: (N⋅num_heads,L,S)(N\cdot\text{num\_heads}, L, S)  where N is the batch size, L is the target sequence length, S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend while the zero positions will be unchanged. If a BoolTensor is provided, positions with True is not allowed to attend while False values will be unchanged. If a FloatTensor is provided, it will be added to the attention weight.


Creates a criterion that optimizes a multi-class multi-classification hinge loss (margin-based loss) between input xx  (a 2D mini-batch Tensor) and output yy  (which is a 2D Tensor of target class indices). For each sample in the mini-batch:
where x∈{0,⋯,x.size(0)−1}x \in \left\{0, \; \cdots , \; \text{x.size}(0) - 1\right\} , y∈{0,⋯,y.size(0)−1}y \in \left\{0, \; \cdots , \; \text{y.size}(0) - 1\right\} , 0≤y[j]≤x.size(0)−10 \leq y[j] \leq \text{x.size}(0)-1 , and i≠y[j]i \neq y[j]  for all ii  and jj .
yy  and xx  must have the same size.
The criterion only considers a contiguous block of non-negative targets that starts at the front.
This allows for different samples to have variable amounts of target classes.
Examples:


Creates a criterion that optimizes a multi-label one-versus-all loss based on max-entropy, between input xx  and target yy  of size (N,C)(N, C) . For each sample in the minibatch:
where i∈{0,⋯,x.nElement()−1}i \in \left\{0, \; \cdots , \; \text{x.nElement}() - 1\right\} , y[i]∈{0,1}y[i] \in \left\{0, \; 1\right\} .


Creates a criterion that optimizes a multi-class classification hinge loss (margin-based loss) between input xx  (a 2D mini-batch Tensor) and output yy  (which is a 1D tensor of target class indices, 0≤y≤x.size(1)−10 \leq y \leq \text{x.size}(1)-1 ):
For each mini-batch sample, the loss in terms of the 1D input xx  and scalar output yy  is:
where x∈{0,⋯,x.size(0)−1}x \in \left\{0, \; \cdots , \; \text{x.size}(0) - 1\right\}  and i≠yi \neq y .
Optionally, you can give non-equal weighting on the classes by passing a 1D weight tensor into the constructor.
The loss function then becomes:


The negative log likelihood loss. It is useful to train a classification problem with C classes.
If provided, the optional argument weight should be a 1D Tensor assigning weight to each of the classes. This is particularly useful when you have an unbalanced training set.
The input given through a forward call is expected to contain log-probabilities of each class. input has to be a Tensor of size either (minibatch,C)(minibatch, C)  or (minibatch,C,d1,d2,...,dK)(minibatch, C, d_1, d_2, ..., d_K)  with K≥1K \geq 1  for the K-dimensional case (described later).
Obtaining log-probabilities in a neural network is easily achieved by adding a LogSoftmax layer in the last layer of your network. You may use CrossEntropyLoss instead, if you prefer not to add an extra layer.
The target that this loss expects should be a class index in the range [0,C−1][0, C-1]  where C = number of classes; if ignore_index is specified, this loss also accepts this class index (this index may not necessarily be in the class range).
The unreduced (i.e. with reduction set to 'none') loss can be described as:
where xx  is the input, yy  is the target, ww  is the weight, and NN  is the batch size. If reduction is not 'none' (default 'mean'), then
Can also be used for higher dimension inputs, such as 2D images, by providing an input of size (minibatch,C,d1,d2,...,dK)(minibatch, C, d_1, d_2, ..., d_K)  with K≥1K \geq 1 , where KK  is the number of dimensions, and a target of appropriate shape (see below). In the case of images, it computes NLL loss per-pixel.
Examples:


Computes the batchwise pairwise distance between vectors v1v_1 , v2v_2  using the p-norm:


Evaluates module(input) in parallel across the GPUs given in device_ids.
This is the functional version of the DataParallel module.
a Tensor containing the result of module(input) located on output_device


Implements distributed data parallelism that is based on torch.distributed package at the module level.
This container parallelizes the application of the given module by splitting the input across the specified devices by chunking in the batch dimension. The module is replicated on each machine and each device, and each such replica handles a portion of the input. During the backwards pass, gradients from each node are averaged.
The batch size should be larger than the number of GPUs used locally.
See also: Basics and Use nn.parallel.DistributedDataParallel instead of multiprocessing or nn.DataParallel. The same constraints on input as in torch.nn.DataParallel apply.
Creation of this class requires that torch.distributed to be already initialized, by calling torch.distributed.init_process_group().
DistributedDataParallel is proven to be significantly faster than torch.nn.DataParallel for single-node multi-GPU data parallel training.
To use DistributedDataParallel on a host with N GPUs, you should spawn up N processes, ensuring that each process exclusively works on a single GPU from 0 to N-1. This can be done by either setting CUDA_VISIBLE_DEVICES for every process or by calling:
where i is from 0 to N-1. In each process, you should refer the following to construct this module:
In order to spawn up multiple processes per node, you can use either torch.distributed.launch or torch.multiprocessing.spawn.
Note
Please refer to PyTorch Distributed Overview for a brief introduction to all features related to distributed training.
Note
nccl backend is currently the fastest and highly recommended backend when using GPUs. This applies to both single-node and multi-node distributed training.
Note
This module also supports mixed-precision distributed training. This means that your model can have different types of parameters such as mixed types of fp16 and fp32, the gradient reduction on these mixed types of parameters will just work fine.
Note
If you use torch.save on one process to checkpoint the module, and torch.load on some other processes to recover it, make sure that map_location is configured properly for every process. Without map_location, torch.load would recover the module to devices where the module was saved from.
Note
When a model is trained on M nodes with batch=N, the gradient will be M times smaller when compared to the same model trained on a single node with batch=M*N if the loss is summed (NOT averaged as usual) across instances in a batch (because the gradients between different nodes are averaged). You should take this into consideration when you want to obtain a mathematically equivalent training process compared to the local training counterpart. But in most cases, you can just treat a DistributedDataParallel wrapped model, a DataParallel wrapped model and an ordinary model on a single GPU as the same (E.g. using the same learning rate for equivalent batch size).
Note
Parameters are never broadcast between processes. The module performs an all-reduce step on gradients and assumes that they will be modified by the optimizer in all processes in the same way. Buffers (e.g. BatchNorm stats) are broadcast from the module in process of rank 0, to all other replicas in the system in every iteration.
Note
If you are using DistributedDataParallel in conjunction with the Distributed RPC Framework, you should always use torch.distributed.autograd.backward() to compute gradients and torch.distributed.optim.DistributedOptimizer for optimizing parameters.
Example:
Warning
Constructor, forward method, and differentiation of the output (or a function of the output of this module) are distributed synchronization points. Take that into account in case different processes might be executing different code.
Warning
This module assumes all parameters are registered in the model by the time it is created. No parameters should be added nor removed later. Same applies to buffers.
Warning
This module assumes all parameters are registered in the model of each distributed processes are in the same order. The module itself will conduct gradient allreduce following the reverse order of the registered parameters of the model. In other words, it is users’ responsibility to ensure that each distributed process has the exact same model and thus the exact same parameter registration order.
Warning
This module allows parameters with non-rowmajor-contiguous strides. For example, your model may contain some parameters whose torch.memory_format is torch.contiguous_format and others whose format is torch.channels_last. However, corresponding parameters in different processes must have the same strides.
Warning
This module doesn’t work with torch.autograd.grad() (i.e. it will only work if gradients are to be accumulated in .grad attributes of parameters).
Warning
If you plan on using this module with a nccl backend or a gloo backend (that uses Infiniband), together with a DataLoader that uses multiple workers, please change the multiprocessing start method to forkserver (Python 3 only) or spawn. Unfortunately Gloo (that uses Infiniband) and NCCL2 are not fork safe, and you will likely experience deadlocks if you don’t change this setting.
Warning
Forward and backward hooks defined on module and its submodules won’t be invoked anymore, unless the hooks are initialized in the forward() method.
Warning
You should never try to change your model’s parameters after wrapping up your model with DistributedDataParallel. Because, when wrapping up your model with DistributedDataParallel, the constructor of DistributedDataParallel will register the additional gradient reduction functions on all the parameters of the model itself at the time of construction. If you change the model’s parameters afterwards, gradient redunction functions no longer match the correct set of parameters.
Warning
Using DistributedDataParallel in conjunction with the Distributed RPC Framework is experimental and subject to change.
Warning
The gradient_as_bucket_view mode does not yet work with Automatic Mixed Precision (AMP). AMP maintains stashed gradients that are used for unscaling gradients. With gradient_as_bucket_view=True, these stashed gradients will point to communication buckets in the first iteration. In the next iteration, the communication buckets are mutated and thus these stashed gradients will be unexpectedly mutated as well, which might lead to wrong results.
~DistributedDataParallel.module (Module) – the module to be parallelized.
Example:
A context manager to be used in conjunction with an instance of torch.nn.parallel.DistributedDataParallel to be able to train with uneven inputs across participating processes.
This context manager will keep track of already-joined DDP processes, and “shadow” the forward and backward passes by inserting collective communication operations to match with the ones created by non-joined DDP processes. This will ensure each collective call has a corresponding call by already-joined DDP processes, preventing hangs or errors that would otherwise happen when training with uneven inputs across processes.
Once all DDP processes have joined, the context manager will broadcast the model corresponding to the last joined process to all processes to ensure the model is the same across all processes (which is guaranteed by DDP).
To use this to enable training with uneven inputs across processes, simply wrap this context manager around your training loop. No further modifications to the model or data loading is required.
Warning
This module works only with the multi-process, single-device usage of torch.nn.parallel.DistributedDataParallel, which means that a single process works on a single GPU.
Warning
This module currently does not support custom distributed collective operations in the forward pass, such as SyncBatchNorm or other custom defined collectives in the model’s forward pass.
Example:
A context manager to disable gradient synchronizations across DDP processes. Within this context, gradients will be accumulated on module variables, which will later be synchronized in the first forward-backward pass exiting the context.
Example:
Registers a communication hook which is an enhancement that provides a flexible hook to users where they can specify how DDP aggregates gradients across multiple workers.
This hook would be very useful for researchers to try out new ideas. For example, this hook can be used to implement several algorithms like GossipGrad and gradient compression which involve different communication strategies for parameter syncs while running Distributed DataParallel training.
state (object) –
Passed to the hook to maintain any state information during the training process. Examples include error feedback in gradient compression, peers to communicate with next in GossipGrad, etc.
It is locally stored by each worker and shared by all the gradient tensors on the worker.
hook (callable) –
Averages gradient tensors across workers and defined as: hook(state: object, bucket: dist._GradBucket) -> torch.futures.Future:
This function is called once the bucket is ready. The hook can perform whatever processing is needed and return a Future indicating completion of any async work (ex: allreduce). If the hook doesn’t perform any communication, it can also just return a completed Future. The Future should hold the new value of grad bucket’s tensors. Once a bucket is ready, c10d reducer would call this hook and use the tensors returned by the Future and copy grads to individual parameters.
We also provide an API called get_future to retrieve a Future associated with the completion of c10d.ProcessGroup.work.
Warning
Grad bucket’s tensors will not be predivided by world_size. User is responsible to divide by the world_size in case of operations like allreduce.
Warning
DDP communication hook can only be registered once and should be registered before calling backward.
Warning
The Future object that hook returns should contain a result that has the same shape with the tensors inside grad bucket.
Warning
DDP communication hook does not support single-process multiple-device mode. Gradbucket tensors should consist of only a single tensor.
Warning
get_future API supports only NCCL backend and will return a torch._C.Future which is an internal type and should be used with caution. It can still be used by register_comm_hook API, but it is subject to some subtle differences compared to torch.futures.Future.
Warning
DDP communication hook is experimental and subject to change.
Below is an example of a noop hook that returns the same tensors.
Below is an example of a Parallel SGD algorithm where gradients are encoded before allreduce, and then decoded after allreduce.


A context manager to be used in conjunction with an instance of torch.nn.parallel.DistributedDataParallel to be able to train with uneven inputs across participating processes.
This context manager will keep track of already-joined DDP processes, and “shadow” the forward and backward passes by inserting collective communication operations to match with the ones created by non-joined DDP processes. This will ensure each collective call has a corresponding call by already-joined DDP processes, preventing hangs or errors that would otherwise happen when training with uneven inputs across processes.
Once all DDP processes have joined, the context manager will broadcast the model corresponding to the last joined process to all processes to ensure the model is the same across all processes (which is guaranteed by DDP).
To use this to enable training with uneven inputs across processes, simply wrap this context manager around your training loop. No further modifications to the model or data loading is required.
Warning
This module works only with the multi-process, single-device usage of torch.nn.parallel.DistributedDataParallel, which means that a single process works on a single GPU.
Warning
This module currently does not support custom distributed collective operations in the forward pass, such as SyncBatchNorm or other custom defined collectives in the model’s forward pass.
Example:


A context manager to disable gradient synchronizations across DDP processes. Within this context, gradients will be accumulated on module variables, which will later be synchronized in the first forward-backward pass exiting the context.
Example:


Registers a communication hook which is an enhancement that provides a flexible hook to users where they can specify how DDP aggregates gradients across multiple workers.
This hook would be very useful for researchers to try out new ideas. For example, this hook can be used to implement several algorithms like GossipGrad and gradient compression which involve different communication strategies for parameter syncs while running Distributed DataParallel training.
state (object) –
Passed to the hook to maintain any state information during the training process. Examples include error feedback in gradient compression, peers to communicate with next in GossipGrad, etc.
It is locally stored by each worker and shared by all the gradient tensors on the worker.
hook (callable) –
Averages gradient tensors across workers and defined as: hook(state: object, bucket: dist._GradBucket) -> torch.futures.Future:
This function is called once the bucket is ready. The hook can perform whatever processing is needed and return a Future indicating completion of any async work (ex: allreduce). If the hook doesn’t perform any communication, it can also just return a completed Future. The Future should hold the new value of grad bucket’s tensors. Once a bucket is ready, c10d reducer would call this hook and use the tensors returned by the Future and copy grads to individual parameters.
We also provide an API called get_future to retrieve a Future associated with the completion of c10d.ProcessGroup.work.
Warning
Grad bucket’s tensors will not be predivided by world_size. User is responsible to divide by the world_size in case of operations like allreduce.
Warning
DDP communication hook can only be registered once and should be registered before calling backward.
Warning
The Future object that hook returns should contain a result that has the same shape with the tensors inside grad bucket.
Warning
DDP communication hook does not support single-process multiple-device mode. Gradbucket tensors should consist of only a single tensor.
Warning
get_future API supports only NCCL backend and will return a torch._C.Future which is an internal type and should be used with caution. It can still be used by register_comm_hook API, but it is subject to some subtle differences compared to torch.futures.Future.
Warning
DDP communication hook is experimental and subject to change.
Below is an example of a noop hook that returns the same tensors.
Below is an example of a Parallel SGD algorithm where gradients are encoded before allreduce, and then decoded after allreduce.


A kind of Tensor that is to be considered a module parameter.
Parameters are Tensor subclasses, that have a very special property when used with Module s - when they’re assigned as Module attributes they are automatically added to the list of its parameters, and will appear e.g. in parameters() iterator. Assigning a Tensor doesn’t have such effect. This is because one might want to cache some temporary state, like last hidden state of the RNN, in the model. If there was no such class as Parameter, these temporaries would get registered too.


A parameter that is not initialized.
Unitialized Parameters are a a special case of torch.nn.Parameter where the shape of the data is still unknown.
Unlikely a torch.nn.Parameter, uninitialized parameters hold no data and attempting to access some properties, like their shape, will throw a runtime error. The only operations that can be performed on a uninitialized parameter are changing its datatype, moving it to a different device and converting it to a regular torch.nn.Parameter.
Create a Parameter with the same properties of the uninitialized one. Given a shape, it materializes a parameter in the same device and with the same dtype as the current one or the specified ones in the arguments.


Create a Parameter with the same properties of the uninitialized one. Given a shape, it materializes a parameter in the same device and with the same dtype as the current one or the specified ones in the arguments.


Holds parameters in a dictionary.
ParameterDict can be indexed like a regular Python dictionary, but parameters it contains are properly registered, and will be visible by all Module methods.
ParameterDict is an ordered dictionary that respects
Note that update() with other unordered mapping types (e.g., Python’s plain dict) does not preserve the order of the merged mapping.
parameters (iterable, optional) – a mapping (dictionary) of (string : Parameter) or an iterable of key-value pairs of type (string, Parameter)
Example:
Remove all items from the ParameterDict.
Return an iterable of the ParameterDict key/value pairs.
Return an iterable of the ParameterDict keys.
Remove key from the ParameterDict and return its parameter.
key (string) – key to pop from the ParameterDict
Update the ParameterDict with the key-value pairs from a mapping or an iterable, overwriting existing keys.
Note
If parameters is an OrderedDict, a ParameterDict, or an iterable of key-value pairs, the order of new elements in it is preserved.
parameters (iterable) – a mapping (dictionary) from string to Parameter, or an iterable of key-value pairs of type (string, Parameter)
Return an iterable of the ParameterDict values.


Remove all items from the ParameterDict.


Return an iterable of the ParameterDict key/value pairs.


Return an iterable of the ParameterDict keys.


Remove key from the ParameterDict and return its parameter.
key (string) – key to pop from the ParameterDict


Update the ParameterDict with the key-value pairs from a mapping or an iterable, overwriting existing keys.
Note
If parameters is an OrderedDict, a ParameterDict, or an iterable of key-value pairs, the order of new elements in it is preserved.
parameters (iterable) – a mapping (dictionary) from string to Parameter, or an iterable of key-value pairs of type (string, Parameter)


Return an iterable of the ParameterDict values.


Holds parameters in a list.
ParameterList can be indexed like a regular Python list, but parameters it contains are properly registered, and will be visible by all Module methods.
parameters (iterable, optional) – an iterable of Parameter to add
Example:
Appends a given parameter at the end of the list.
parameter (nn.Parameter) – parameter to append
Appends parameters from a Python iterable to the end of the list.
parameters (iterable) – iterable of parameters to append


Appends a given parameter at the end of the list.
parameter (nn.Parameter) – parameter to append


Appends parameters from a Python iterable to the end of the list.
parameters (iterable) – iterable of parameters to append


Rearranges elements in a tensor of shape (∗,C×r2,H,W)(*, C \times r^2, H, W)  to a tensor of shape (∗,C,H×r,W×r)(*, C, H \times r, W \times r) , where r is an upscale factor.
This is useful for implementing efficient sub-pixel convolution with a stride of 1/r1/r .
See the paper: Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network by Shi et. al (2016) for more details.
upscale_factor (int) – factor to increase spatial resolution by
Examples:


Reverses the PixelShuffle operation by rearranging elements in a tensor of shape (∗,C,H×r,W×r)(*, C, H \times r, W \times r)  to a tensor of shape (∗,C×r2,H,W)(*, C \times r^2, H, W) , where r is a downscale factor.
See the paper: Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network by Shi et. al (2016) for more details.
downscale_factor (int) – factor to decrease spatial resolution by
Examples:


Negative log likelihood loss with Poisson distribution of target.
The loss can be described as:
The last term can be omitted or approximated with Stirling formula. The approximation is used for target values more than 1. For targets less or equal to 1 zeros are added to the loss.
full (bool, optional) –
whether to compute full loss, i. e. to add the Stirling approximation term
Examples:


Applies the element-wise function:
or
Here aa  is a learnable parameter. When called without arguments, nn.PReLU() uses a single parameter aa  across all input channels. If called with nn.PReLU(nChannels), a separate aa  is used for each input channel.
Note
weight decay should not be used when learning aa  for good performance.
Note
Channel dim is the 2nd dim of input. When input has dims < 2, then there is no channel dim and the number of channels = 1.
~PReLU.weight (Tensor) – the learnable weights of shape (num_parameters).
Examples:


A Conv2d module attached with FakeQuantize modules for weight, used for quantization aware training.
We adopt the same interface as torch.nn.Conv2d, please see https://pytorch.org/docs/stable/nn.html?highlight=conv2d#torch.nn.Conv2d for documentation.
Similar to torch.nn.Conv2d, with FakeQuantize modules initialized to default.
~Conv2d.weight_fake_quant – fake quant module for weight
Create a qat module from a float module or qparams_dict
Args: mod a float module, either produced by torch.quantization utilities or directly from user


Create a qat module from a float module or qparams_dict
Args: mod a float module, either produced by torch.quantization utilities or directly from user


A linear module attached with FakeQuantize modules for weight, used for quantization aware training.
We adopt the same interface as torch.nn.Linear, please see https://pytorch.org/docs/stable/nn.html#torch.nn.Linear for documentation.
Similar to torch.nn.Linear, with FakeQuantize modules initialized to default.
~Linear.weight – fake quant module for weight
Create a qat module from a float module or qparams_dict
Args: mod a float module, either produced by torch.quantization utilities or directly from user


This is the quantized version of BatchNorm2d.


This is the quantized version of BatchNorm3d.


Applies a 1D convolution over a quantized input signal composed of several quantized input planes.
For details on input arguments, parameters, and implementation see Conv1d.
Note
Only zeros is supported for the padding_mode argument.
Note
Only torch.quint8 is supported for the input data type.
See Conv1d for other attributes.
Examples:
Creates a quantized module from a float module or qparams_dict.
mod (Module) – a float module, either produced by torch.quantization utilities or provided by the user


Creates a quantized module from a float module or qparams_dict.
mod (Module) – a float module, either produced by torch.quantization utilities or provided by the user


Applies a 2D convolution over a quantized input signal composed of several quantized input planes.
For details on input arguments, parameters, and implementation see Conv2d.
Note
Only zeros is supported for the padding_mode argument.
Note
Only torch.quint8 is supported for the input data type.
See Conv2d for other attributes.
Examples:
Creates a quantized module from a float module or qparams_dict.
mod (Module) – a float module, either produced by torch.quantization utilities or provided by the user


Applies a 3D convolution over a quantized input signal composed of several quantized input planes.
For details on input arguments, parameters, and implementation see Conv3d.
Note
Only zeros is supported for the padding_mode argument.
Note
Only torch.quint8 is supported for the input data type.
See Conv3d for other attributes.
Examples:
Creates a quantized module from a float module or qparams_dict.
mod (Module) – a float module, either produced by torch.quantization utilities or provided by the user


A dynamic quantized linear module with floating point tensor as inputs and outputs. We adopt the same interface as torch.nn.Linear, please see https://pytorch.org/docs/stable/nn.html#torch.nn.Linear for documentation.
Similar to torch.nn.Linear, attributes will be randomly initialized at module creation time and will be overwritten later
Examples:
Create a dynamic quantized module from a float module or qparams_dict
mod (Module) – a float module, either produced by torch.quantization utilities or provided by the user
A dynamic quantized LSTM module with floating point tensor as inputs and outputs. We adopt the same interface as torch.nn.LSTM, please see https://pytorch.org/docs/stable/nn.html#torch.nn.LSTM for documentation.
Examples:
A long short-term memory (LSTM) cell.
A dynamic quantized LSTMCell module with floating point tensor as inputs and outputs. Weights are quantized to 8 bits. We adopt the same interface as torch.nn.LSTMCell, please see https://pytorch.org/docs/stable/nn.html#torch.nn.LSTMCell for documentation.
Examples:
A gated recurrent unit (GRU) cell
A dynamic quantized GRUCell module with floating point tensor as inputs and outputs. Weights are quantized to 8 bits. We adopt the same interface as torch.nn.GRUCell, please see https://pytorch.org/docs/stable/nn.html#torch.nn.GRUCell for documentation.
Examples:
An Elman RNN cell with tanh or ReLU non-linearity. A dynamic quantized RNNCell module with floating point tensor as inputs and outputs. Weights are quantized to 8 bits. We adopt the same interface as torch.nn.RNNCell, please see https://pytorch.org/docs/stable/nn.html#torch.nn.RNNCell for documentation.
Examples:


A gated recurrent unit (GRU) cell
A dynamic quantized GRUCell module with floating point tensor as inputs and outputs. Weights are quantized to 8 bits. We adopt the same interface as torch.nn.GRUCell, please see https://pytorch.org/docs/stable/nn.html#torch.nn.GRUCell for documentation.
Examples:


A dynamic quantized linear module with floating point tensor as inputs and outputs. We adopt the same interface as torch.nn.Linear, please see https://pytorch.org/docs/stable/nn.html#torch.nn.Linear for documentation.
Similar to torch.nn.Linear, attributes will be randomly initialized at module creation time and will be overwritten later
Examples:
Create a dynamic quantized module from a float module or qparams_dict
mod (Module) – a float module, either produced by torch.quantization utilities or provided by the user


Create a dynamic quantized module from a float module or qparams_dict
mod (Module) – a float module, either produced by torch.quantization utilities or provided by the user


A dynamic quantized LSTM module with floating point tensor as inputs and outputs. We adopt the same interface as torch.nn.LSTM, please see https://pytorch.org/docs/stable/nn.html#torch.nn.LSTM for documentation.
Examples:


A long short-term memory (LSTM) cell.
A dynamic quantized LSTMCell module with floating point tensor as inputs and outputs. Weights are quantized to 8 bits. We adopt the same interface as torch.nn.LSTMCell, please see https://pytorch.org/docs/stable/nn.html#torch.nn.LSTMCell for documentation.
Examples:


An Elman RNN cell with tanh or ReLU non-linearity. A dynamic quantized RNNCell module with floating point tensor as inputs and outputs. Weights are quantized to 8 bits. We adopt the same interface as torch.nn.RNNCell, please see https://pytorch.org/docs/stable/nn.html#torch.nn.RNNCell for documentation.
Examples:


This is the quantized equivalent of ELU.


A quantized Embedding module with quantized packed weights as inputs. We adopt the same interface as torch.nn.Embedding, please see https://pytorch.org/docs/stable/nn.html#torch.nn.Embedding for documentation.
Similar to Embedding, attributes will be randomly initialized at module creation time and will be overwritten later
~Embedding.weight (Tensor) – the non-learnable quantized weights of the module of shape (num_embeddings,embedding_dim)(\text{num\_embeddings}, \text{embedding\_dim}) .
Create a quantized embedding module from a float module
mod (Module) – a float module, either produced by torch.quantization utilities or provided by user


Create a quantized embedding module from a float module
mod (Module) – a float module, either produced by torch.quantization utilities or provided by user


A quantized EmbeddingBag module with quantized packed weights as inputs. We adopt the same interface as torch.nn.EmbeddingBag, please see https://pytorch.org/docs/stable/nn.html#torch.nn.EmbeddingBag for documentation.
Similar to EmbeddingBag, attributes will be randomly initialized at module creation time and will be overwritten later
~EmbeddingBag.weight (Tensor) – the non-learnable quantized weights of the module of shape (num_embeddings,embedding_dim)(\text{num\_embeddings}, \text{embedding\_dim}) .
Create a quantized embedding_bag module from a float module
mod (Module) – a float module, either produced by torch.quantization utilities or provided by user


Create a quantized embedding_bag module from a float module
mod (Module) – a float module, either produced by torch.quantization utilities or provided by user


State collector class for float operations.
The instance of this class can be used instead of the torch. prefix for some operations. See example usage below.
Note
This class does not provide a forward hook. Instead, you must use one of the underlying functions (e.g. add).
Examples:


Applies a 2D adaptive average pooling over a quantized input signal composed of several quantized input planes.
Note
The input quantization parameters propagate to the output.
See AdaptiveAvgPool2d for details and output shape.
output_size – the target output size (single integer or double-integer tuple)


Applies 2D average-pooling operation in kH×kWkH \times kW  regions by step size sH×sWsH \times sW  steps. The number of output features is equal to the number of input planes.
Note
The input quantization parameters propagate to the output.
See AvgPool2d for details and output shape.


Applies a 1D convolution over a quantized 1D input composed of several input planes.
See Conv1d for details and output shape.
Examples:


Applies a 2D convolution over a quantized 2D input composed of several input planes.
See Conv2d for details and output shape.
Examples:


Applies a 3D convolution over a quantized 3D input composed of several input planes.
See Conv3d for details and output shape.
Examples:


This is the quantized version of hardswish().


Down/up samples the input to either the given size or the given scale_factor
See torch.nn.functional.interpolate() for implementation details.
The input dimensions are interpreted in the form: mini-batch x channels x [optional depth] x [optional height] x width.
Note
The input quantization parameters propagate to the output.
Note
Only 2D/3D input is supported for quantized inputs
Note
Only the following modes are supported for the quantized inputs:


Applies a linear transformation to the incoming quantized data: y=xAT+by = xA^T + b . See Linear
Note
Current implementation packs weights on every call, which has penalty on performance. If you want to avoid the overhead, use Linear.


Applies a 2D max pooling over a quantized input signal composed of several quantized input planes.
Note
The input quantization parameters are propagated to the output.
See MaxPool2d for details.


Upsamples the input to either the given size or the given scale_factor
Warning
This function is deprecated in favor of torch.nn.quantized.functional.interpolate(). This is equivalent with nn.quantized.functional.interpolate(...).
See torch.nn.functional.interpolate() for implementation details.
The input dimensions are interpreted in the form: mini-batch x channels x [optional depth] x [optional height] x width.
Note
The input quantization parameters propagate to the output.
Note
Only 2D input is supported for quantized inputs
Note
Only the following modes are supported for the quantized inputs:
Warning
With align_corners = True, the linearly interpolating modes (bilinear) don’t proportionally align the output and input pixels, and thus the output values can depend on the input size. This was the default behavior for these modes up to version 0.3.1. Since then, the default behavior is align_corners = False. See Upsample for concrete examples on how this affects the outputs.


Upsamples the input, using bilinear upsampling.
Warning
This function is deprecated in favor of torch.nn.quantized.functional.interpolate(). This is equivalent with nn.quantized.functional.interpolate(..., mode='bilinear', align_corners=True).
Note
The input quantization parameters propagate to the output.
Note
Only 2D inputs are supported


Upsamples the input, using nearest neighbours’ pixel values.
Warning
This function is deprecated in favor of torch.nn.quantized.functional.interpolate(). This is equivalent with nn.quantized.functional.interpolate(..., mode='nearest').
Note
The input quantization parameters propagate to the output.
Note
Only 2D inputs are supported


This is the quantized version of GroupNorm.


This is the quantized version of Hardswish.


This is the quantized version of InstanceNorm1d.


This is the quantized version of InstanceNorm2d.


This is the quantized version of InstanceNorm3d.


This is the quantized version of LayerNorm.


A quantized linear module with quantized tensor as inputs and outputs. We adopt the same interface as torch.nn.Linear, please see https://pytorch.org/docs/stable/nn.html#torch.nn.Linear for documentation.
Similar to Linear, attributes will be randomly initialized at module creation time and will be overwritten later
Examples:
Create a quantized module from a float module or qparams_dict
mod (Module) – a float module, either produced by torch.quantization utilities or provided by the user


Create a quantized module from a float module or qparams_dict
mod (Module) – a float module, either produced by torch.quantization utilities or provided by the user


Wrapper class for quantized operations.
The instance of this class can be used instead of the torch.ops.quantized prefix. See example usage below.
Note
This class does not provide a forward hook. Instead, you must use one of the underlying functions (e.g. add).
Examples:


Quantizes an incoming tensor
zero_point, dtype (`scale`,) –


Applies the element-wise function:
ReLU6(x)=min(max(x0,x),q(6))\text{ReLU6}(x) = \min(\max(x_0, x), q(6)) , where x0x_0  is the zero_point, and q(6)q(6)  is the quantized representation of number 6.
inplace – can optionally do the operation in-place. Default: False
Examples:


Pads the input tensor using the reflection of the input boundary.
For N-dimensional padding, use torch.nn.functional.pad().
padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 2-tuple, uses (padding_left\text{padding\_left} , padding_right\text{padding\_right} )
Output: (N,C,Wout)(N, C, W_{out})  where
Wout=Win+padding_left+padding_rightW_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}
Examples:


Pads the input tensor using the reflection of the input boundary.
For N-dimensional padding, use torch.nn.functional.pad().
padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 4-tuple, uses (padding_left\text{padding\_left} , padding_right\text{padding\_right} , padding_top\text{padding\_top} , padding_bottom\text{padding\_bottom} )
Output: (N,C,Hout,Wout)(N, C, H_{out}, W_{out})  where
Hout=Hin+padding_top+padding_bottomH_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}
Wout=Win+padding_left+padding_rightW_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}
Examples:


Applies the rectified linear unit function element-wise:
ReLU(x)=(x)+=max(0,x)\text{ReLU}(x) = (x)^+ = \max(0, x)
inplace – can optionally do the operation in-place. Default: False
Examples:


Applies the element-wise function:
inplace – can optionally do the operation in-place. Default: False
Examples:


Pads the input tensor using replication of the input boundary.
For N-dimensional padding, use torch.nn.functional.pad().
padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 2-tuple, uses (padding_left\text{padding\_left} , padding_right\text{padding\_right} )
Output: (N,C,Wout)(N, C, W_{out})  where
Wout=Win+padding_left+padding_rightW_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}
Examples:


Pads the input tensor using replication of the input boundary.
For N-dimensional padding, use torch.nn.functional.pad().
padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 4-tuple, uses (padding_left\text{padding\_left} , padding_right\text{padding\_right} , padding_top\text{padding\_top} , padding_bottom\text{padding\_bottom} )
Output: (N,C,Hout,Wout)(N, C, H_{out}, W_{out})  where
Hout=Hin+padding_top+padding_bottomH_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}
Wout=Win+padding_left+padding_rightW_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}
Examples:


Pads the input tensor using replication of the input boundary.
For N-dimensional padding, use torch.nn.functional.pad().
padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 6-tuple, uses (padding_left\text{padding\_left} , padding_right\text{padding\_right} , padding_top\text{padding\_top} , padding_bottom\text{padding\_bottom} , padding_front\text{padding\_front} , padding_back\text{padding\_back} )
Output: (N,C,Dout,Hout,Wout)(N, C, D_{out}, H_{out}, W_{out})  where
Dout=Din+padding_front+padding_backD_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}
Hout=Hin+padding_top+padding_bottomH_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}
Wout=Win+padding_left+padding_rightW_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}
Examples:


Applies a multi-layer Elman RNN with tanh\tanh  or ReLU\text{ReLU}  non-linearity to an input sequence.
For each element in the input sequence, each layer computes the following function:
where hth_t  is the hidden state at time t, xtx_t  is the input at time t, and h(t−1)h_{(t-1)}  is the hidden state of the previous layer at time t-1 or the initial hidden state at time 0. If nonlinearity is 'relu', then ReLU\text{ReLU}  is used instead of tanh\tanh .
output of shape (seq_len, batch, num_directions * hidden_size): tensor containing the output features (h_t) from the last layer of the RNN, for each t. If a torch.nn.utils.rnn.PackedSequence has been given as the input, the output will also be a packed sequence.
For the unpacked case, the directions can be separated using output.view(seq_len, batch, num_directions, hidden_size), with forward and backward being direction 0 and 1 respectively. Similarly, the directions can be separated in the packed case.
h_n of shape (num_layers * num_directions, batch, hidden_size): tensor containing the hidden state for t = seq_len.
Like output, the layers can be separated using h_n.view(num_layers, num_directions, batch, hidden_size).
Note
All the weights and biases are initialized from U(−k,k)\mathcal{U}(-\sqrt{k}, \sqrt{k})  where k=1hidden_sizek = \frac{1}{\text{hidden\_size}}
Warning
There are known non-determinism issues for RNN functions on some versions of cuDNN and CUDA. You can enforce deterministic behavior by setting the following environment variables:
On CUDA 10.1, set environment variable CUDA_LAUNCH_BLOCKING=1. This may affect performance.
On CUDA 10.2 or later, set environment variable (note the leading colon symbol) CUBLAS_WORKSPACE_CONFIG=:16:8 or CUBLAS_WORKSPACE_CONFIG=:4096:2.
See the cuDNN 8 Release Notes for more information.
Note
If the following conditions are satisfied: 1) cudnn is enabled, 2) input data is on the GPU 3) input data has dtype torch.float16 4) V100 GPU is used, 5) input data is not in PackedSequence format persistent algorithm can be selected to improve performance.
Examples:


Resets parameter data pointer so that they can use faster code paths.
Right now, this works only if the module is on the GPU and cuDNN is enabled. Otherwise, it’s a no-op.


An Elman RNN cell with tanh or ReLU non-linearity.
If nonlinearity is ‘relu’, then ReLU is used in place of tanh.
Note
All the weights and biases are initialized from U(−k,k)\mathcal{U}(-\sqrt{k}, \sqrt{k})  where k=1hidden_sizek = \frac{1}{\text{hidden\_size}}
Examples:


Applies the randomized leaky rectified liner unit function, element-wise, as described in the paper:
Empirical Evaluation of Rectified Activations in Convolutional Network.
The function is defined as:
where aa  is randomly sampled from uniform distribution U(lower,upper)\mathcal{U}(\text{lower}, \text{upper}) .
See: https://arxiv.org/pdf/1505.00853.pdf
Examples:


Applied element-wise, as:
with α=1.6732632423543772848170429916717\alpha = 1.6732632423543772848170429916717  and scale=1.0507009873554804934193349852946\text{scale} = 1.0507009873554804934193349852946 .
More details can be found in the paper Self-Normalizing Neural Networks .
inplace (bool, optional) – can optionally do the operation in-place. Default: False
Examples:


A sequential container. Modules will be added to it in the order they are passed in the constructor. Alternatively, an ordered dict of modules can also be passed in.
To make it easier to understand, here is a small example:


Applies the element-wise function:
Examples:


Applies the silu function, element-wise.
Note
See Gaussian Error Linear Units (GELUs) where the SiLU (Sigmoid Linear Unit) was originally coined, and see Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning and Swish: a Self-Gated Activation Function where the SiLU was experimented with later.
Examples:


Creates a criterion that uses a squared term if the absolute element-wise error falls below beta and an L1 term otherwise. It is less sensitive to outliers than the torch.nn.MSELoss and in some cases prevents exploding gradients (e.g. see Fast R-CNN paper by Ross Girshick). Omitting a scaling factor of beta, this loss is also known as the Huber loss:
where ziz_{i}  is given by:
xx  and yy  arbitrary shapes with a total of nn  elements each the sum operation still operates over all the elements, and divides by nn .
beta is an optional parameter that defaults to 1.
Note: When beta is set to 0, this is equivalent to L1Loss. Passing a negative value in for beta will result in an exception.
The division by nn  can be avoided if sets reduction = 'sum'.


Creates a criterion that optimizes a two-class classification logistic loss between input tensor xx  and target tensor yy  (containing 1 or -1).


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.
Softmax is defined as:
When the input Tensor is a sparse tensor then the unspecifed values are treated as -inf.
a Tensor of the same dimension and shape as the input with values in the range [0, 1]
dim (int) – A dimension along which Softmax will be computed (so every slice along dim will sum to 1).
Note
This module doesn’t work directly with NLLLoss, which expects the Log to be computed between the Softmax and itself. Use LogSoftmax instead (it’s faster and has better numerical properties).
Examples:


Applies SoftMax over features to each spatial location.
When given an image of Channels x Height x Width, it will apply Softmax to each location (Channels,hi,wj)(Channels, h_i, w_j)
a Tensor of the same dimension and shape as the input with values in the range [0, 1]
Examples:


Applies the Softmin 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.
Softmin is defined as:
dim (int) – A dimension along which Softmin will be computed (so every slice along dim will sum to 1).
a Tensor of the same dimension and shape as the input, with values in the range [0, 1]
Examples:


Applies the element-wise function:
SoftPlus is a smooth approximation to the ReLU function and can be used to constrain the output of a machine to always be positive.
For numerical stability the implementation reverts to the linear function when input×β>thresholdinput \times \beta > threshold .
Examples:


Applies the soft shrinkage function elementwise:
lambd – the λ\lambda  (must be no less than zero) value for the Softshrink formulation. Default: 0.5
Examples:


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 .
The mean and standard-deviation are calculated per-dimension over all mini-batches of the same process groups. γ\gamma  and β\beta  are learnable parameter vectors of size C (where C is the input size). By default, the elements of γ\gamma  are sampled from U(0,1)\mathcal{U}(0, 1)  and the elements of β\beta  are set to 0. The standard-deviation is calculated via the biased estimator, equivalent to torch.var(input, unbiased=False).
Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default momentum of 0.1.
If track_running_stats is set to False, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well.
Note
This momentum argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is x^new=(1−momentum)×x^+momentum×xt\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t , where x^\hat{x}  is the estimated statistic and xtx_t  is the new observed value.
Because the Batch Normalization is done for each channel in the C dimension, computing statistics on (N, +) slices, it’s common terminology to call this Volumetric Batch Normalization or Spatio-temporal Batch Normalization.
Currently SyncBatchNorm only supports DistributedDataParallel (DDP) with single GPU per process. Use torch.nn.SyncBatchNorm.convert_sync_batchnorm() to convert BatchNorm*D layer to SyncBatchNorm before wrapping Network with DDP.
Examples:
Helper function to convert all BatchNorm*D layers in the model to torch.nn.SyncBatchNorm layers.
The original module with the converted torch.nn.SyncBatchNorm layers. If the original module is a BatchNorm*D layer, a new torch.nn.SyncBatchNorm layer object will be returned instead.
Example:


Helper function to convert all BatchNorm*D layers in the model to torch.nn.SyncBatchNorm layers.
The original module with the converted torch.nn.SyncBatchNorm layers. If the original module is a BatchNorm*D layer, a new torch.nn.SyncBatchNorm layer object will be returned instead.
Example:


Thresholds each element of the input Tensor.
Threshold is defined as:
Examples:


A transformer model. User is able to modify the attributes as needed. The architecture is based on the paper “Attention Is All You Need”. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems, pages 6000-6010. Users can build the BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters.
Note: A full example to apply nn.Transformer module for the word language model is available in https://github.com/pytorch/examples/tree/master/word_language_model
Take in and process masked source/target sequences.
Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend while the zero positions will be unchanged. If a BoolTensor is provided, positions with True are not allowed to attend while False values will be unchanged. If a FloatTensor is provided, it will be added to the attention weight. [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions will be unchanged. If a BoolTensor is provided, the positions with the value of True will be ignored while the position with the value of False will be unchanged.
Note: Due to the multi-head attention architecture in the transformer model, the output sequence length of a transformer is same as the input sequence (i.e. target) length of the decode.
where S is the source sequence length, T is the target sequence length, N is the batch size, E is the feature number
Generate a square mask for the sequence. The masked positions are filled with float(‘-inf’). Unmasked positions are filled with float(0.0).


Take in and process masked source/target sequences.
Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend while the zero positions will be unchanged. If a BoolTensor is provided, positions with True are not allowed to attend while False values will be unchanged. If a FloatTensor is provided, it will be added to the attention weight. [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions will be unchanged. If a BoolTensor is provided, the positions with the value of True will be ignored while the position with the value of False will be unchanged.
Note: Due to the multi-head attention architecture in the transformer model, the output sequence length of a transformer is same as the input sequence (i.e. target) length of the decode.
where S is the source sequence length, T is the target sequence length, N is the batch size, E is the feature number


Generate a square mask for the sequence. The masked positions are filled with float(‘-inf’). Unmasked positions are filled with float(0.0).


TransformerDecoder is a stack of N decoder layers
Pass the inputs (and mask) through the decoder layer in turn.
see the docs in Transformer class.


Pass the inputs (and mask) through the decoder layer in turn.
see the docs in Transformer class.


TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network. This standard decoder layer is based on the paper “Attention Is All You Need”. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems, pages 6000-6010. Users may modify or implement in a different way during application.
Pass the inputs (and mask) through the decoder layer.
see the docs in Transformer class.


Pass the inputs (and mask) through the decoder layer.
see the docs in Transformer class.


TransformerEncoder is a stack of N encoder layers
Pass the input through the encoder layers in turn.
see the docs in Transformer class.


Pass the input through the encoder layers in turn.
see the docs in Transformer class.


TransformerEncoderLayer is made up of self-attn and feedforward network. This standard encoder layer is based on the paper “Attention Is All You Need”. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems, pages 6000-6010. Users may modify or implement in a different way during application.
Pass the input through the encoder layer.
see the docs in Transformer class.


Pass the input through the encoder layer.
see the docs in Transformer class.


Creates a criterion that measures the triplet loss given an input tensors x1x1 , x2x2 , x3x3  and a margin with a value greater than 00 . This is used for measuring a relative similarity between samples. A triplet is composed by a, p and n (i.e., anchor, positive examples and negative examples respectively). The shapes of all input tensors should be (N,D)(N, D) .
The distance swap is described in detail in the paper Learning shallow convolutional feature descriptors with triplet losses by V. Balntas, E. Riba et al.
The loss function for each sample in the mini-batch is:
where
See also TripletMarginWithDistanceLoss, which computes the triplet margin loss for input tensors using a custom distance function.
otherwise.


Creates a criterion that measures the triplet loss given input tensors aa , pp , and nn  (representing anchor, positive, and negative examples, respectively), and a nonnegative, real-valued function (“distance function”) used to compute the relationship between the anchor and positive example (“positive distance”) and the anchor and negative example (“negative distance”).
The unreduced loss (i.e., with reduction set to 'none') can be described as:
where NN  is the batch size; dd  is a nonnegative, real-valued function quantifying the closeness of two tensors, referred to as the distance_function; and marginmargin  is a nonnegative margin representing the minimum difference between the positive and negative distances that is required for the loss to be 0. The input tensors have NN  elements each and can be of any shape that the distance function can handle.
If reduction is not 'none' (default 'mean'), then:
See also TripletMarginLoss, which computes the triplet loss for input tensors using the lpl_p  distance as the distance function.
Examples:
V. Balntas, et al.: Learning shallow convolutional feature descriptors with triplet losses: http://www.bmva.org/bmvc/2016/papers/paper119/index.html


Unflattens a tensor dim expanding it to a desired shape. For use with Sequential.
Adds a child module to the current module.
The module can be accessed as an attribute using the given name.
Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also torch.nn.init).
fn (Module -> None) – function to be applied to each submodule
self
Module
Example:
Casts all floating point parameters and buffers to bfloat16 datatype.
self
Module
Returns an iterator over module buffers.
recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.
torch.Tensor – module buffer
Example:
Returns an iterator over immediate children modules.
Module – a child module
Moves all model parameters and buffers to the CPU.
self
Module
Moves all model parameters and buffers to the GPU.
This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.
device (int, optional) – if specified, all parameters will be copied to that device
self
Module
Casts all floating point parameters and buffers to double datatype.
self
Module
Sets the module in evaluation mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.
This is equivalent with self.train(False).
self
Module
Casts all floating point parameters and buffers to float datatype.
self
Module
Casts all floating point parameters and buffers to half datatype.
self
Module
Copies parameters and buffers from state_dict into this module and its descendants. If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.
NamedTuple with missing_keys and unexpected_keys fields
Returns an iterator over all modules in the network.
Module – a module in the network
Note
Duplicate modules are returned only once. In the following example, l will be returned only once.
Example:
Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
(string, torch.Tensor) – Tuple containing the name and buffer
Example:
Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
(string, Module) – Tuple containing a name and child module
Example:
Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
(string, Module) – Tuple of name and module
Note
Duplicate modules are returned only once. In the following example, l will be returned only once.
Example:
Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
(string, Parameter) – Tuple containing the name and parameter
Example:
Returns an iterator over module parameters.
This is typically passed to an optimizer.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Parameter – module parameter
Example:
Registers a backward hook on the module.
This function is deprecated in favor of nn.Module.register_full_backward_hook() and the behavior of this function will change in future versions.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle
Adds a buffer to the module.
This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.
Buffers can be accessed as attributes using given names.
Example:
Registers a forward hook on the module.
The hook will be called every time after forward() has computed an output. It should have the following signature:
The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle
Registers a forward pre-hook on the module.
The hook will be called every time before forward() is invoked. It should have the following signature:
The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle
Registers a backward hook on the module.
The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature:
The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.
Warning
Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle
Adds a parameter to the module.
The parameter can be accessed as an attribute using given name.
Change if autograd should record operations on parameters in this module.
This method sets the parameters’ requires_grad attributes in-place.
This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).
requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.
self
Module
Returns a dictionary containing a whole state of the module.
Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names.
a dictionary containing a whole state of the module
dict
Example:
Moves and/or casts the parameters and buffers.
This can be called as
Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtype`s. In addition, this method will
only cast the floating point or complex parameters and buffers to :attr:`dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.
See below for examples.
Note
This method modifies the module in-place.
self
Module
Examples:
Sets the module in training mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.
mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.
self
Module
Casts all parameters and buffers to dst_type.
dst_type (type or string) – the desired type
self
Module
Moves all model parameters and buffers to the XPU.
This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.
device (int, optional) – if specified, all parameters will be copied to that device
self
Module
Sets gradients of all model parameters to zero. See similar function under torch.optim.Optimizer for more context.
set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.


Adds a child module to the current module.
The module can be accessed as an attribute using the given name.


Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also torch.nn.init).
fn (Module -> None) – function to be applied to each submodule
self
Module
Example:


Casts all floating point parameters and buffers to bfloat16 datatype.
self
Module


Returns an iterator over module buffers.
recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.
torch.Tensor – module buffer
Example:


Returns an iterator over immediate children modules.
Module – a child module


Moves all model parameters and buffers to the CPU.
self
Module


Moves all model parameters and buffers to the GPU.
This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.
device (int, optional) – if specified, all parameters will be copied to that device
self
Module


Casts all floating point parameters and buffers to double datatype.
self
Module


Sets the module in evaluation mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.
This is equivalent with self.train(False).
self
Module


Casts all floating point parameters and buffers to float datatype.
self
Module


Casts all floating point parameters and buffers to half datatype.
self
Module


Copies parameters and buffers from state_dict into this module and its descendants. If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.
NamedTuple with missing_keys and unexpected_keys fields


Returns an iterator over all modules in the network.
Module – a module in the network
Note
Duplicate modules are returned only once. In the following example, l will be returned only once.
Example:


Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
(string, torch.Tensor) – Tuple containing the name and buffer
Example:


Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
(string, Module) – Tuple containing a name and child module
Example:


Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
(string, Module) – Tuple of name and module
Note
Duplicate modules are returned only once. In the following example, l will be returned only once.
Example:


Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
(string, Parameter) – Tuple containing the name and parameter
Example:


Returns an iterator over module parameters.
This is typically passed to an optimizer.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Parameter – module parameter
Example:


Registers a backward hook on the module.
This function is deprecated in favor of nn.Module.register_full_backward_hook() and the behavior of this function will change in future versions.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle


Adds a buffer to the module.
This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.
Buffers can be accessed as attributes using given names.
Example:


Registers a forward hook on the module.
The hook will be called every time after forward() has computed an output. It should have the following signature:
The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called.
a handle that can be used to remove the added hook by calling handle.remove()
torch.utils.hooks.RemovableHandle