PyTorch documentation¶

PyTorch is an optimized tensor library for deep learning using GPUs and CPUs.

torch¶

The torch package contains data structures for multi-dimensional tensors and mathematical operations over these are defined. Additionally, it provides many utilities for efficient serializing of Tensors and arbitrary types, and other useful utilities.

It has a CUDA counterpart, that enables you to run your tensor computations on an NVIDIA GPU with compute capability >= 3.0.

Tensors¶

torch.is_tensor(obj)¶

Returns True if obj is a PyTorch tensor.

Parameters: obj (Object) – Object to test

torch.is_storage(obj)¶

Returns True if obj is a PyTorch storage object.

Parameters: obj (Object) – Object to test

torch.is_floating_point(tensor) -> (bool)¶

Returns True if the data type of tensor is a floating point data type i.e., one of torch.float64, torch.float32 and torch.float16.

Parameters: tensor (Tensor) – the PyTorch tensor to test

torch.set_default_dtype(d)¶

Sets the default floating point dtype to d. This type will be used as default floating point type for type inference in torch.tensor().

The default floating point dtype is initially torch.float32.

Parameters: d (torch.dtype) – the floating point dtype to make the default

Example:

>>> torch.tensor([1.2, 3]).dtype           # initial default for floating point is torch.float32
torch.float32
>>> torch.set_default_dtype(torch.float64)
>>> torch.tensor([1.2, 3]).dtype           # a new floating point tensor
torch.float64

torch.get_default_dtype() → torch.dtype¶

Get the current default floating point torch.dtype.

Example:

>>> torch.get_default_dtype()  # initial default for floating point is torch.float32
torch.float32
>>> torch.set_default_dtype(torch.float64)
>>> torch.get_default_dtype()  # default is now changed to torch.float64
torch.float64
>>> torch.set_default_tensor_type(torch.FloatTensor)  # setting tensor type also affects this
>>> torch.get_default_dtype()  # changed to torch.float32, the dtype for torch.FloatTensor
torch.float32

torch.set_default_tensor_type(t)¶

Sets the default torch.Tensor type to floating point tensor type t. This type will also be used as default floating point type for type inference in torch.tensor().

The default floating point tensor type is initially torch.FloatTensor.

Parameters: t (type or string) – the floating point tensor type or its name

Example:

>>> torch.tensor([1.2, 3]).dtype    # initial default for floating point is torch.float32
torch.float32
>>> torch.set_default_tensor_type(torch.DoubleTensor)
>>> torch.tensor([1.2, 3]).dtype    # a new floating point tensor
torch.float64

torch.numel(input) → int¶

Returns the total number of elements in the input tensor.

Parameters: input (Tensor) – the input tensor

Example:

>>> a = torch.randn(1, 2, 3, 4, 5)
>>> torch.numel(a)
120
>>> a = torch.zeros(4,4)
>>> torch.numel(a)
16

torch.set_printoptions(precision=None, threshold=None, edgeitems=None, linewidth=None, profile=None, sci_mode=None)¶

Set options for printing. Items shamelessly taken from NumPy

Parameters

precision – Number of digits of precision for floating point output (default = 4).
threshold – Total number of array elements which trigger summarization rather than full repr (default = 1000).
edgeitems – Number of array items in summary at beginning and end of each dimension (default = 3).
linewidth – The number of characters per line for the purpose of inserting line breaks (default = 80). Thresholded matrices will ignore this parameter.
profile – Sane defaults for pretty printing. Can override with any of the above options. (any one of default, short, full)
sci_mode – Enable (True) or disable (False) scientific notation. If None (default) is specified, the value is defined by _Formatter

torch.set_flush_denormal(mode) → bool¶

Disables denormal floating numbers on CPU.

Returns True if your system supports flushing denormal numbers and it successfully configures flush denormal mode. set_flush_denormal() is only supported on x86 architectures supporting SSE3.

Parameters: mode (bool) – Controls whether to enable flush denormal mode or not

Example:

>>> torch.set_flush_denormal(True)
True
>>> torch.tensor([1e-323], dtype=torch.float64)
tensor([ 0.], dtype=torch.float64)
>>> torch.set_flush_denormal(False)
True
>>> torch.tensor([1e-323], dtype=torch.float64)
tensor(9.88131e-324 *
       [ 1.0000], dtype=torch.float64)

Creation Ops¶

Note

Random sampling creation ops are listed under Random sampling and include: torch.rand() torch.rand_like() torch.randn() torch.randn_like() torch.randint() torch.randint_like() torch.randperm() You may also use torch.empty() with the In-place random sampling methods to create torch.Tensor s with values sampled from a broader range of distributions.

torch.tensor(data, dtype=None, device=None, requires_grad=False) → Tensor¶

Constructs a tensor with data.

Warning

torch.tensor() always copies data. If you have a Tensor data and want to avoid a copy, use torch.Tensor.requires_grad_() or torch.Tensor.detach(). If you have a NumPy ndarray and want to avoid a copy, use torch.as_tensor().

Warning

When data is a tensor x, torch.tensor() reads out ‘the data’ from whatever it is passed, and constructs a leaf variable. Therefore torch.tensor(x) is equivalent to x.clone().detach() and torch.tensor(x, requires_grad=True) is equivalent to x.clone().detach().requires_grad_(True). The equivalents using clone() and detach() are recommended.

Parameters

data (array_like) – Initial data for the tensor. Can be a list, tuple, NumPy ndarray, scalar, and other types.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, infers data type from data.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.tensor([[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]])
tensor([[ 0.1000,  1.2000],
        [ 2.2000,  3.1000],
        [ 4.9000,  5.2000]])

>>> torch.tensor([0, 1])  # Type inference on data
tensor([ 0,  1])

>>> torch.tensor([[0.11111, 0.222222, 0.3333333]],
                 dtype=torch.float64,
                 device=torch.device('cuda:0'))  # creates a torch.cuda.DoubleTensor
tensor([[ 0.1111,  0.2222,  0.3333]], dtype=torch.float64, device='cuda:0')

>>> torch.tensor(3.14159)  # Create a scalar (zero-dimensional tensor)
tensor(3.1416)

>>> torch.tensor([])  # Create an empty tensor (of size (0,))
tensor([])

torch.sparse_coo_tensor(indices, values, size=None, dtype=None, device=None, requires_grad=False) → Tensor¶

Constructs a sparse tensors in COO(rdinate) format with non-zero elements at the given indices with the given values. A sparse tensor can be uncoalesced, in that case, there are duplicate coordinates in the indices, and the value at that index is the sum of all duplicate value entries: torch.sparse.

Parameters

indices (array_like) – Initial data for the tensor. Can be a list, tuple, NumPy ndarray, scalar, and other types. Will be cast to a torch.LongTensor internally. The indices are the coordinates of the non-zero values in the matrix, and thus should be two-dimensional where the first dimension is the number of tensor dimensions and the second dimension is the number of non-zero values.
values (array_like) – Initial values for the tensor. Can be a list, tuple, NumPy ndarray, scalar, and other types.
size (list, tuple, or torch.Size, optional) – Size of the sparse tensor. If not provided the size will be inferred as the minimum size big enough to hold all non-zero elements.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, infers data type from values.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> i = torch.tensor([[0, 1, 1],
                      [2, 0, 2]])
>>> v = torch.tensor([3, 4, 5], dtype=torch.float32)
>>> torch.sparse_coo_tensor(i, v, [2, 4])
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3., 4., 5.]),
       size=(2, 4), nnz=3, layout=torch.sparse_coo)

>>> torch.sparse_coo_tensor(i, v)  # Shape inference
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3., 4., 5.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)

>>> torch.sparse_coo_tensor(i, v, [2, 4],
                            dtype=torch.float64,
                            device=torch.device('cuda:0'))
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3., 4., 5.]),
       device='cuda:0', size=(2, 4), nnz=3, dtype=torch.float64,
       layout=torch.sparse_coo)

# Create an empty sparse tensor with the following invariants:
#   1. sparse_dim + dense_dim = len(SparseTensor.shape)
#   2. SparseTensor._indices().shape = (sparse_dim, nnz)
#   3. SparseTensor._values().shape = (nnz, SparseTensor.shape[sparse_dim:])
#
# For instance, to create an empty sparse tensor with nnz = 0, dense_dim = 0 and
# sparse_dim = 1 (hence indices is a 2D tensor of shape = (1, 0))
>>> S = torch.sparse_coo_tensor(torch.empty([1, 0]), [], [1])
tensor(indices=tensor([], size=(1, 0)),
       values=tensor([], size=(0,)),
       size=(1,), nnz=0, layout=torch.sparse_coo)

# and to create an empty sparse tensor with nnz = 0, dense_dim = 1 and
# sparse_dim = 1
>>> S = torch.sparse_coo_tensor(torch.empty([1, 0]), torch.empty([0, 2]), [1, 2])
tensor(indices=tensor([], size=(1, 0)),
       values=tensor([], size=(0, 2)),
       size=(1, 2), nnz=0, layout=torch.sparse_coo)

torch.as_tensor(data, dtype=None, device=None) → Tensor¶

Convert the data into a torch.Tensor. If the data is already a Tensor with the same dtype and device, no copy will be performed, otherwise a new Tensor will be returned with computational graph retained if data Tensor has requires_grad=True. Similarly, if the data is an ndarray of the corresponding dtype and the device is the cpu, no copy will be performed.

Parameters

data (array_like) – Initial data for the tensor. Can be a list, tuple, NumPy ndarray, scalar, and other types.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, infers data type from data.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.

Example:

>>> a = numpy.array([1, 2, 3])
>>> t = torch.as_tensor(a)
>>> t
tensor([ 1,  2,  3])
>>> t[0] = -1
>>> a
array([-1,  2,  3])

>>> a = numpy.array([1, 2, 3])
>>> t = torch.as_tensor(a, device=torch.device('cuda'))
>>> t
tensor([ 1,  2,  3])
>>> t[0] = -1
>>> a
array([1,  2,  3])

torch.from_numpy(ndarray) → Tensor¶

Creates a Tensor from a numpy.ndarray.

The returned tensor and ndarray share the same memory. Modifications to the tensor will be reflected in the ndarray and vice versa. The returned tensor is not resizable.

Example:

>>> a = numpy.array([1, 2, 3])
>>> t = torch.from_numpy(a)
>>> t
tensor([ 1,  2,  3])
>>> t[0] = -1
>>> a
array([-1,  2,  3])

torch.zeros(*sizes, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a tensor filled with the scalar value 0, with the shape defined by the variable argument sizes.

Parameters

sizes (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.zeros(2, 3)
tensor([[ 0.,  0.,  0.],
        [ 0.,  0.,  0.]])

>>> torch.zeros(5)
tensor([ 0.,  0.,  0.,  0.,  0.])

torch.zeros_like(input, dtype=None, layout=None, device=None, requires_grad=False) → Tensor¶

Returns a tensor filled with the scalar value 0, with the same size as input. torch.zeros_like(input) is equivalent to torch.zeros(input.size(), dtype=input.dtype, layout=input.layout, device=input.device).

Warning

As of 0.4, this function does not support an out keyword. As an alternative, the old torch.zeros_like(input, out=output) is equivalent to torch.zeros(input.size(), out=output).

Parameters

input (Tensor) – the size of input will determine size of the output tensor
dtype (torch.dtype, optional) – the desired data type of returned Tensor. Default: if None, defaults to the dtype of input.
layout (torch.layout, optional) – the desired layout of returned tensor. Default: if None, defaults to the layout of input.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, defaults to the device of input.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> input = torch.empty(2, 3)
>>> torch.zeros_like(input)
tensor([[ 0.,  0.,  0.],
        [ 0.,  0.,  0.]])

torch.ones(*sizes, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a tensor filled with the scalar value 1, with the shape defined by the variable argument sizes.

Parameters

sizes (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.ones(2, 3)
tensor([[ 1.,  1.,  1.],
        [ 1.,  1.,  1.]])

>>> torch.ones(5)
tensor([ 1.,  1.,  1.,  1.,  1.])

torch.ones_like(input, dtype=None, layout=None, device=None, requires_grad=False) → Tensor¶

Returns a tensor filled with the scalar value 1, with the same size as input. torch.ones_like(input) is equivalent to torch.ones(input.size(), dtype=input.dtype, layout=input.layout, device=input.device).

Warning

As of 0.4, this function does not support an out keyword. As an alternative, the old torch.ones_like(input, out=output) is equivalent to torch.ones(input.size(), out=output).

Parameters

input (Tensor) – the size of input will determine size of the output tensor
dtype (torch.dtype, optional) – the desired data type of returned Tensor. Default: if None, defaults to the dtype of input.
layout (torch.layout, optional) – the desired layout of returned tensor. Default: if None, defaults to the layout of input.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, defaults to the device of input.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> input = torch.empty(2, 3)
>>> torch.ones_like(input)
tensor([[ 1.,  1.,  1.],
        [ 1.,  1.,  1.]])

torch.arange(start=0, end, step=1, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a 1-D tensor of size $\left\lfloor \frac{\text{end} - \text{start}}{\text{step}} \right\rfloor$ with values from the interval [start, end) taken with common difference step beginning from start.

Note that non-integer step is subject to floating point rounding errors when comparing against end; to avoid inconsistency, we advise adding a small epsilon to end in such cases.

\[\text{out}_{{i+1}} = \text{out}_{i} + \text{step} \]

Parameters

start (Number) – the starting value for the set of points. Default: 0.
end (Number) – the ending value for the set of points
step (Number) – the gap between each pair of adjacent points. Default: 1.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()). If dtype is not given, infer the data type from the other input arguments. If any of start, end, or stop are floating-point, the dtype is inferred to be the default dtype, see get_default_dtype(). Otherwise, the dtype is inferred to be torch.int64.
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.arange(5)
tensor([ 0,  1,  2,  3,  4])
>>> torch.arange(1, 4)
tensor([ 1,  2,  3])
>>> torch.arange(1, 2.5, 0.5)
tensor([ 1.0000,  1.5000,  2.0000])

torch.range(start=0, end, step=1, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a 1-D tensor of size $\left\lfloor \frac{\text{end} - \text{start}}{\text{step}} \right\rfloor + 1$ with values from start to end with step step. Step is the gap between two values in the tensor.

\[\text{out}_{i+1} = \text{out}_i + \text{step}. \]

Warning

This function is deprecated in favor of torch.arange().

Parameters

start (float) – the starting value for the set of points. Default: 0.
end (float) – the ending value for the set of points
step (float) – the gap between each pair of adjacent points. Default: 1.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.range(1, 4)
tensor([ 1.,  2.,  3.,  4.])
>>> torch.range(1, 4, 0.5)
tensor([ 1.0000,  1.5000,  2.0000,  2.5000,  3.0000,  3.5000,  4.0000])

torch.linspace(start, end, steps=100, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a one-dimensional tensor of steps equally spaced points between start and end.

The output tensor is 1-D of size steps.

Parameters

start (float) – the starting value for the set of points
end (float) – the ending value for the set of points
steps (int) – number of points to sample between start and end. Default: 100.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.linspace(3, 10, steps=5)
tensor([  3.0000,   4.7500,   6.5000,   8.2500,  10.0000])
>>> torch.linspace(-10, 10, steps=5)
tensor([-10.,  -5.,   0.,   5.,  10.])
>>> torch.linspace(start=-10, end=10, steps=5)
tensor([-10.,  -5.,   0.,   5.,  10.])
>>> torch.linspace(start=-10, end=10, steps=1)
tensor([-10.])

torch.logspace(start, end, steps=100, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a one-dimensional tensor of steps points logarithmically spaced between $10^{\text{start}}$ and $10^{\text{end}}$.

The output tensor is 1-D of size steps.

Parameters

start (float) – the starting value for the set of points
end (float) – the ending value for the set of points
steps (int) – number of points to sample between start and end. Default: 100.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.logspace(start=-10, end=10, steps=5)
tensor([ 1.0000e-10,  1.0000e-05,  1.0000e+00,  1.0000e+05,  1.0000e+10])
>>> torch.logspace(start=0.1, end=1.0, steps=5)
tensor([  1.2589,   2.1135,   3.5481,   5.9566,  10.0000])
>>> torch.logspace(start=0.1, end=1.0, steps=1)
tensor([1.2589])

torch.eye(n, m=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a 2-D tensor with ones on the diagonal and zeros elsewhere.

Parameters

n (int) – the number of rows
m (int, optional) – the number of columns with default being n
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Returns

A 2-D tensor with ones on the diagonal and zeros elsewhere

Return type

Tensor

Example:

>>> torch.eye(3)
tensor([[ 1.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  1.]])

torch.empty(*sizes, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a tensor filled with uninitialized data. The shape of the tensor is defined by the variable argument sizes.

Parameters

sizes (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.empty(2, 3)
tensor(1.00000e-08 *
       [[ 6.3984,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000]])

torch.empty_like(input, dtype=None, layout=None, device=None, requires_grad=False) → Tensor¶

Returns an uninitialized tensor with the same size as input. torch.empty_like(input) is equivalent to torch.empty(input.size(), dtype=input.dtype, layout=input.layout, device=input.device).

Parameters

input (Tensor) – the size of input will determine size of the output tensor
dtype (torch.dtype, optional) – the desired data type of returned Tensor. Default: if None, defaults to the dtype of input.
layout (torch.layout, optional) – the desired layout of returned tensor. Default: if None, defaults to the layout of input.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, defaults to the device of input.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.empty((2,3), dtype=torch.int64)
tensor([[ 9.4064e+13,  2.8000e+01,  9.3493e+13],
        [ 7.5751e+18,  7.1428e+18,  7.5955e+18]])

torch.full(size, fill_value, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a tensor of size size filled with fill_value.

Parameters

size (int...) – a list, tuple, or torch.Size of integers defining the shape of the output tensor.
fill_value – the number to fill the output tensor with.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.full((2, 3), 3.141592)
tensor([[ 3.1416,  3.1416,  3.1416],
        [ 3.1416,  3.1416,  3.1416]])

torch.full_like(input, fill_value, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a tensor with the same size as input filled with fill_value. torch.full_like(input, fill_value) is equivalent to torch.full(input.size(), fill_value, dtype=input.dtype, layout=input.layout, device=input.device).

Parameters

input (Tensor) – the size of input will determine size of the output tensor
fill_value – the number to fill the output tensor with.
dtype (torch.dtype, optional) – the desired data type of returned Tensor. Default: if None, defaults to the dtype of input.
layout (torch.layout, optional) – the desired layout of returned tensor. Default: if None, defaults to the layout of input.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, defaults to the device of input.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Indexing, Slicing, Joining, Mutating Ops¶

torch.cat(tensors, dim=0, out=None) → Tensor¶

Concatenates the given sequence of seq tensors in the given dimension. All tensors must either have the same shape (except in the concatenating dimension) or be empty.

torch.cat() can be seen as an inverse operation for torch.split() and torch.chunk().

torch.cat() can be best understood via examples.

Parameters

tensors (sequence of Tensors) – any python sequence of tensors of the same type. Non-empty tensors provided must have the same shape, except in the cat dimension.
dim (int, optional) – the dimension over which the tensors are concatenated
out (Tensor, optional) – the output tensor

Example:

>>> x = torch.randn(2, 3)
>>> x
tensor([[ 0.6580, -1.0969, -0.4614],
        [-0.1034, -0.5790,  0.1497]])
>>> torch.cat((x, x, x), 0)
tensor([[ 0.6580, -1.0969, -0.4614],
        [-0.1034, -0.5790,  0.1497],
        [ 0.6580, -1.0969, -0.4614],
        [-0.1034, -0.5790,  0.1497],
        [ 0.6580, -1.0969, -0.4614],
        [-0.1034, -0.5790,  0.1497]])
>>> torch.cat((x, x, x), 1)
tensor([[ 0.6580, -1.0969, -0.4614,  0.6580, -1.0969, -0.4614,  0.6580,
         -1.0969, -0.4614],
        [-0.1034, -0.5790,  0.1497, -0.1034, -0.5790,  0.1497, -0.1034,
         -0.5790,  0.1497]])

torch.chunk(tensor, chunks, dim=0) → List of Tensors¶

Splits a tensor into a specific number of chunks.

Last chunk will be smaller if the tensor size along the given dimension dim is not divisible by chunks.

Parameters

tensor (Tensor) – the tensor to split
chunks (int) – number of chunks to return
dim (int) – dimension along which to split the tensor

torch.gather(input, dim, index, out=None, sparse_grad=False) → Tensor¶

Gathers values along an axis specified by dim.

For a 3-D tensor the output is specified by:

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

If input is an n-dimensional tensor with size $(x_0, x_1..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})$ and dim = i, then index must be an $n$-dimensional tensor with size $(x_0, x_1, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})$ where $y \geq 1$ and out will have the same size as index.

Parameters

input (Tensor) – the source tensor
dim (int) – the axis along which to index
index (LongTensor) – the indices of elements to gather
out (Tensor, optional) – the destination tensor
sparse_grad (bool,optional) – If True, gradient w.r.t. input will be a sparse tensor.

Example:

>>> t = torch.tensor([[1,2],[3,4]])
>>> torch.gather(t, 1, torch.tensor([[0,0],[1,0]]))
tensor([[ 1,  1],
        [ 4,  3]])

torch.index_select(input, dim, index, out=None) → Tensor¶

Returns a new tensor which indexes the input tensor along dimension dim using the entries in index which is a LongTensor.

The returned tensor has the same number of dimensions as the original tensor (input). The dimth dimension has the same size as the length of index; other dimensions have the same size as in the original tensor.

Note

The returned tensor does not use the same storage as the original tensor. If out has a different shape than expected, we silently change it to the correct shape, reallocating the underlying storage if necessary.

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension in which we index
index (LongTensor) – the 1-D tensor containing the indices to index
out (Tensor, optional) – the output tensor

Example:

>>> x = torch.randn(3, 4)
>>> x
tensor([[ 0.1427,  0.0231, -0.5414, -1.0009],
        [-0.4664,  0.2647, -0.1228, -1.1068],
        [-1.1734, -0.6571,  0.7230, -0.6004]])
>>> indices = torch.tensor([0, 2])
>>> torch.index_select(x, 0, indices)
tensor([[ 0.1427,  0.0231, -0.5414, -1.0009],
        [-1.1734, -0.6571,  0.7230, -0.6004]])
>>> torch.index_select(x, 1, indices)
tensor([[ 0.1427, -0.5414],
        [-0.4664, -0.1228],
        [-1.1734,  0.7230]])

torch.masked_select(input, mask, out=None) → Tensor¶

Returns a new 1-D tensor which indexes the input tensor according to the binary mask mask which is a ByteTensor.

The shapes of the mask tensor and the input tensor don’t need to match, but they must be broadcastable.

Note

The returned tensor does not use the same storage as the original tensor

Parameters

input (Tensor) – the input data
mask (ByteTensor) – the tensor containing the binary mask to index with
out (Tensor, optional) – the output tensor

Example:

>>> x = torch.randn(3, 4)
>>> x
tensor([[ 0.3552, -2.3825, -0.8297,  0.3477],
        [-1.2035,  1.2252,  0.5002,  0.6248],
        [ 0.1307, -2.0608,  0.1244,  2.0139]])
>>> mask = x.ge(0.5)
>>> mask
tensor([[ 0,  0,  0,  0],
        [ 0,  1,  1,  1],
        [ 0,  0,  0,  1]], dtype=torch.uint8)
>>> torch.masked_select(x, mask)
tensor([ 1.2252,  0.5002,  0.6248,  2.0139])

torch.narrow(input, dimension, start, length) → Tensor¶

Returns a new tensor that is a narrowed version of input tensor. The dimension dim is input from start to start + length. The returned tensor and input tensor share the same underlying storage.

Parameters

input (Tensor) – the tensor to narrow
dimension (int) – the dimension along which to narrow
start (int) – the starting dimension
length (int) – the distance to the ending dimension

Example:

>>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
>>> torch.narrow(x, 0, 0, 2)
tensor([[ 1,  2,  3],
        [ 4,  5,  6]])
>>> torch.narrow(x, 1, 1, 2)
tensor([[ 2,  3],
        [ 5,  6],
        [ 8,  9]])

torch.nonzero(input, out=None) → LongTensor¶

Returns a tensor containing the indices of all non-zero elements of input. Each row in the result contains the indices of a non-zero element in input.

If input has n dimensions, then the resulting indices tensor out is of size $(z \times n)$, where $z$ is the total number of non-zero elements in the input tensor.

Parameters

input (Tensor) – the input tensor
out (LongTensor, optional) – the output tensor containing indices

Example:

>>> torch.nonzero(torch.tensor([1, 1, 1, 0, 1]))
tensor([[ 0],
        [ 1],
        [ 2],
        [ 4]])
>>> torch.nonzero(torch.tensor([[0.6, 0.0, 0.0, 0.0],
                                [0.0, 0.4, 0.0, 0.0],
                                [0.0, 0.0, 1.2, 0.0],
                                [0.0, 0.0, 0.0,-0.4]]))
tensor([[ 0,  0],
        [ 1,  1],
        [ 2,  2],
        [ 3,  3]])

torch.reshape(input, shape) → Tensor¶

Returns a tensor with the same data and number of elements as input, but with the specified shape. When possible, the returned tensor will be a view of input. Otherwise, it will be a copy. Contiguous inputs and inputs with compatible strides can be reshaped without copying, but you should not depend on the copying vs. viewing behavior.

See torch.Tensor.view() on when it is possible to return a view.

A single dimension may be -1, in which case it’s inferred from the remaining dimensions and the number of elements in input.

Parameters

input (Tensor) – the tensor to be reshaped
shape (tuple of python:ints) – the new shape

Example:

>>> a = torch.arange(4.)
>>> torch.reshape(a, (2, 2))
tensor([[ 0.,  1.],
        [ 2.,  3.]])
>>> b = torch.tensor([[0, 1], [2, 3]])
>>> torch.reshape(b, (-1,))
tensor([ 0,  1,  2,  3])

torch.split(tensor, split_size_or_sections, dim=0)¶

Splits the tensor into chunks.

If split_size_or_sections is an integer type, then tensor will be split into equally sized chunks (if possible). Last chunk will be smaller if the tensor size along the given dimension dim is not divisible by split_size.

If split_size_or_sections is a list, then tensor will be split into len(split_size_or_sections) chunks with sizes in dim according to split_size_or_sections.

Parameters

tensor (Tensor) – tensor to split.
split_size_or_sections (int) or (list(int)) – size of a single chunk or list of sizes for each chunk
dim (int) – dimension along which to split the tensor.

torch.squeeze(input, dim=None, out=None) → Tensor¶

Returns a tensor with all the dimensions of input of size 1 removed.

For example, if input is of shape: $(A \times 1 \times B \times C \times 1 \times D)$ then the out tensor will be of shape: $(A \times B \times C \times D)$.

When dim is given, a squeeze operation is done only in the given dimension. If input is of shape: $(A \times 1 \times B)$, squeeze(input, 0) leaves the tensor unchanged, but squeeze(input, 1) will squeeze the tensor to the shape $(A \times B)$.

Note

The returned tensor shares the storage with the input tensor, so changing the contents of one will change the contents of the other.

Parameters

input (Tensor) – the input tensor
dim (int, optional) – if given, the input will be squeezed only in this dimension
out (Tensor, optional) – the output tensor

Example:

>>> x = torch.zeros(2, 1, 2, 1, 2)
>>> x.size()
torch.Size([2, 1, 2, 1, 2])
>>> y = torch.squeeze(x)
>>> y.size()
torch.Size([2, 2, 2])
>>> y = torch.squeeze(x, 0)
>>> y.size()
torch.Size([2, 1, 2, 1, 2])
>>> y = torch.squeeze(x, 1)
>>> y.size()
torch.Size([2, 2, 1, 2])

torch.stack(seq, dim=0, out=None) → Tensor¶

Concatenates sequence of tensors along a new dimension.

All tensors need to be of the same size.

Parameters

seq (sequence of Tensors) – sequence of tensors to concatenate
dim (int) – dimension to insert. Has to be between 0 and the number of dimensions of concatenated tensors (inclusive)
out (Tensor, optional) – the output tensor

torch.t(input) → Tensor¶

Expects input to be <= 2-D tensor and transposes dimensions 0 and 1.

0-D and 1-D tensors are returned as it is and 2-D tensor can be seen as a short-hand function for transpose(input, 0, 1).

Parameters: input (Tensor) – the input tensor

Example:

>>> x = torch.randn(())
>>> x
tensor(0.1995)
>>> torch.t(x)
tensor(0.1995)
>>> x = torch.randn(3)
>>> x
tensor([ 2.4320, -0.4608,  0.7702])
>>> torch.t(x)
tensor([.2.4320,.-0.4608,..0.7702])
>>> x = torch.randn(2, 3)
>>> x
tensor([[ 0.4875,  0.9158, -0.5872],
        [ 0.3938, -0.6929,  0.6932]])
>>> torch.t(x)
tensor([[ 0.4875,  0.3938],
        [ 0.9158, -0.6929],
        [-0.5872,  0.6932]])

torch.take(input, indices) → Tensor¶

Returns a new tensor with the elements of input at the given indices. The input tensor is treated as if it were viewed as a 1-D tensor. The result takes the same shape as the indices.

Parameters

input (Tensor) – the input tensor
indices (LongTensor) – the indices into tensor

Example:

>>> src = torch.tensor([[4, 3, 5],
                        [6, 7, 8]])
>>> torch.take(src, torch.tensor([0, 2, 5]))
tensor([ 4,  5,  8])

torch.transpose(input, dim0, dim1) → Tensor¶

Returns a tensor that is a transposed version of input. The given dimensions dim0 and dim1 are swapped.

The resulting out tensor shares it’s underlying storage with the input tensor, so changing the content of one would change the content of the other.

Parameters

input (Tensor) – the input tensor
dim0 (int) – the first dimension to be transposed
dim1 (int) – the second dimension to be transposed

Example:

>>> x = torch.randn(2, 3)
>>> x
tensor([[ 1.0028, -0.9893,  0.5809],
        [-0.1669,  0.7299,  0.4942]])
>>> torch.transpose(x, 0, 1)
tensor([[ 1.0028, -0.1669],
        [-0.9893,  0.7299],
        [ 0.5809,  0.4942]])

torch.unbind(tensor, dim=0) → seq¶

Removes a tensor dimension.

Returns a tuple of all slices along a given dimension, already without it.

Parameters

tensor (Tensor) – the tensor to unbind
dim (int) – dimension to remove

Example:

>>> torch.unbind(torch.tensor([[1, 2, 3],
>>>                            [4, 5, 6],
>>>                            [7, 8, 9]]))
(tensor([1, 2, 3]), tensor([4, 5, 6]), tensor([7, 8, 9]))

torch.unsqueeze(input, dim, out=None) → Tensor¶

Returns a new tensor with a dimension of size one inserted at the specified position.

The returned tensor shares the same underlying data with this tensor.

A dim value within the range [-input.dim() - 1, input.dim() + 1) can be used. Negative dim will correspond to unsqueeze() applied at dim = dim + input.dim() + 1.

Parameters

input (Tensor) – the input tensor
dim (int) – the index at which to insert the singleton dimension
out (Tensor, optional) – the output tensor

Example:

>>> x = torch.tensor([1, 2, 3, 4])
>>> torch.unsqueeze(x, 0)
tensor([[ 1,  2,  3,  4]])
>>> torch.unsqueeze(x, 1)
tensor([[ 1],
        [ 2],
        [ 3],
        [ 4]])

torch.where(condition, x, y) → Tensor¶

Return a tensor of elements selected from either x or y, depending on condition.

The operation is defined as:

\[out_i = \begin{cases} x_i & \text{if } \text{condition}_i \\ y_i & \text{otherwise} \\ \end{cases} \]

Note

The tensors condition, x, y must be broadcastable.

Parameters

condition (ByteTensor) – When True (nonzero), yield x, otherwise yield y
x (Tensor) – values selected at indices where condition is True
y (Tensor) – values selected at indices where condition is False

Returns

A tensor of shape equal to the broadcasted shape of condition, x, y

Return type

Tensor

Example:

>>> x = torch.randn(3, 2)
>>> y = torch.ones(3, 2)
>>> x
tensor([[-0.4620,  0.3139],
        [ 0.3898, -0.7197],
        [ 0.0478, -0.1657]])
>>> torch.where(x > 0, x, y)
tensor([[ 1.0000,  0.3139],
        [ 0.3898,  1.0000],
        [ 0.0478,  1.0000]])

Random sampling¶

torch.manual_seed(seed)¶

Sets the seed for generating random numbers. Returns a torch._C.Generator object.

Parameters: seed (int) – The desired seed.

torch.initial_seed()¶: Returns the initial seed for generating random numbers as a Python long.

torch.get_rng_state()¶: Returns the random number generator state as a torch.ByteTensor.

torch.set_rng_state(new_state)¶

Sets the random number generator state.

Parameters: new_state (torch.ByteTensor) – The desired state

torch.default_generator = <torch._C.Generator object>¶

torch.bernoulli(input, *, generator=None, out=None) → Tensor¶

Draws binary random numbers (0 or 1) from a Bernoulli distribution.

The input tensor should be a tensor containing probabilities to be used for drawing the binary random number. Hence, all values in input have to be in the range: $0 \leq \text{input}_i \leq 1$.

The $\text{i}^{th}$ element of the output tensor will draw a value $1$ according to the $\text{i}^{th}$ probability value given in input.

\[\text{out}_{i} \sim \mathrm{Bernoulli}(p = \text{input}_{i}) \]

The returned out tensor only has values 0 or 1 and is of the same shape as input.

out can have integral dtype, but input must have floating point dtype.

Parameters

input (Tensor) – the input tensor of probability values for the Bernoulli distribution
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.empty(3, 3).uniform_(0, 1)  # generate a uniform random matrix with range [0, 1]
>>> a
tensor([[ 0.1737,  0.0950,  0.3609],
        [ 0.7148,  0.0289,  0.2676],
        [ 0.9456,  0.8937,  0.7202]])
>>> torch.bernoulli(a)
tensor([[ 1.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 1.,  1.,  1.]])

>>> a = torch.ones(3, 3) # probability of drawing "1" is 1
>>> torch.bernoulli(a)
tensor([[ 1.,  1.,  1.],
        [ 1.,  1.,  1.],
        [ 1.,  1.,  1.]])
>>> a = torch.zeros(3, 3) # probability of drawing "1" is 0
>>> torch.bernoulli(a)
tensor([[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]])

torch.multinomial(input, num_samples, replacement=False, out=None) → LongTensor¶

Returns a tensor where each row contains num_samples indices sampled from the multinomial probability distribution located in the corresponding row of tensor input.

Note

The rows of input do not need to sum to one (in which case we use the values as weights), but must be non-negative, finite and have a non-zero sum.

Indices are ordered from left to right according to when each was sampled (first samples are placed in first column).

If input is a vector, out is a vector of size num_samples.

If input is a matrix with m rows, out is an matrix of shape $(m \times \text{num\_samples})$.

If replacement is True, samples are drawn with replacement.

If not, they are drawn without replacement, which means that when a sample index is drawn for a row, it cannot be drawn again for that row.

Note

When drawn without replacement, num_samples must be lower than number of non-zero elements in input (or the min number of non-zero elements in each row of input if it is a matrix).

Parameters

input (Tensor) – the input tensor containing probabilities
num_samples (int) – number of samples to draw
replacement (bool, optional) – whether to draw with replacement or not
out (Tensor, optional) – the output tensor

Example:

>>> weights = torch.tensor([0, 10, 3, 0], dtype=torch.float) # create a tensor of weights
>>> torch.multinomial(weights, 2)
tensor([1, 2])
>>> torch.multinomial(weights, 4) # ERROR!
RuntimeError: invalid argument 2: invalid multinomial distribution (with replacement=False,
not enough non-negative category to sample) at ../aten/src/TH/generic/THTensorRandom.cpp:320
>>> torch.multinomial(weights, 4, replacement=True)
tensor([ 2,  1,  1,  1])

torch.normal()¶

torch.normal(mean, std, out=None) → Tensor

Returns a tensor of random numbers drawn from separate normal distributions whose mean and standard deviation are given.

The mean is a tensor with the mean of each output element’s normal distribution

The std is a tensor with the standard deviation of each output element’s normal distribution

The shapes of mean and std don’t need to match, but the total number of elements in each tensor need to be the same.

Note

When the shapes do not match, the shape of mean is used as the shape for the returned output tensor

Parameters

mean (Tensor) – the tensor of per-element means
std (Tensor) – the tensor of per-element standard deviations
out (Tensor, optional) – the output tensor

Example:

>>> torch.normal(mean=torch.arange(1., 11.), std=torch.arange(1, 0, -0.1))
tensor([  1.0425,   3.5672,   2.7969,   4.2925,   4.7229,   6.2134,
          8.0505,   8.1408,   9.0563,  10.0566])

torch.normal(mean=0.0, std, out=None) → Tensor

Similar to the function above, but the means are shared among all drawn elements.

Parameters

mean (float, optional) – the mean for all distributions
std (Tensor) – the tensor of per-element standard deviations
out (Tensor, optional) – the output tensor

Example:

>>> torch.normal(mean=0.5, std=torch.arange(1., 6.))
tensor([-1.2793, -1.0732, -2.0687,  5.1177, -1.2303])

torch.normal(mean, std=1.0, out=None) → Tensor

Similar to the function above, but the standard-deviations are shared among all drawn elements.

Parameters

mean (Tensor) – the tensor of per-element means
std (float, optional) – the standard deviation for all distributions
out (Tensor, optional) – the output tensor

Example:

>>> torch.normal(mean=torch.arange(1., 6.))
tensor([ 1.1552,  2.6148,  2.6535,  5.8318,  4.2361])

torch.rand(*sizes, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a tensor filled with random numbers from a uniform distribution on the interval $[0, 1)$

The shape of the tensor is defined by the variable argument sizes.

Parameters

sizes (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.rand(4)
tensor([ 0.5204,  0.2503,  0.3525,  0.5673])
>>> torch.rand(2, 3)
tensor([[ 0.8237,  0.5781,  0.6879],
        [ 0.3816,  0.7249,  0.0998]])

torch.rand_like(input, dtype=None, layout=None, device=None, requires_grad=False) → Tensor¶

Returns a tensor with the same size as input that is filled with random numbers from a uniform distribution on the interval $[0, 1)$. torch.rand_like(input) is equivalent to torch.rand(input.size(), dtype=input.dtype, layout=input.layout, device=input.device).

Parameters

input (Tensor) – the size of input will determine size of the output tensor
dtype (torch.dtype, optional) – the desired data type of returned Tensor. Default: if None, defaults to the dtype of input.
layout (torch.layout, optional) – the desired layout of returned tensor. Default: if None, defaults to the layout of input.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, defaults to the device of input.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

torch.randint(low=0, high, size, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a tensor filled with random integers generated uniformly between low (inclusive) and high (exclusive).

The shape of the tensor is defined by the variable argument size.

Parameters

low (int, optional) – Lowest integer to be drawn from the distribution. Default: 0.
high (int) – One above the highest integer to be drawn from the distribution.
size (tuple) – a tuple defining the shape of the output tensor.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.randint(3, 5, (3,))
tensor([4, 3, 4])


>>> torch.randint(10, (2, 2))
tensor([[0, 2],
        [5, 5]])


>>> torch.randint(3, 10, (2, 2))
tensor([[4, 5],
        [6, 7]])

torch.randint_like(input, low=0, high, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a tensor with the same shape as Tensor input filled with random integers generated uniformly between low (inclusive) and high (exclusive).

Parameters

input (Tensor) – the size of input will determine size of the output tensor
low (int, optional) – Lowest integer to be drawn from the distribution. Default: 0.
high (int) – One above the highest integer to be drawn from the distribution.
dtype (torch.dtype, optional) – the desired data type of returned Tensor. Default: if None, defaults to the dtype of input.
layout (torch.layout, optional) – the desired layout of returned tensor. Default: if None, defaults to the layout of input.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, defaults to the device of input.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

torch.randn(*sizes, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Returns a tensor filled with random numbers from a normal distribution with mean 0 and variance 1 (also called the standard normal distribution).

\[\text{out}_{i} \sim \mathcal{N}(0, 1) \]

The shape of the tensor is defined by the variable argument sizes.

Parameters

sizes (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.randn(4)
tensor([-2.1436,  0.9966,  2.3426, -0.6366])
>>> torch.randn(2, 3)
tensor([[ 1.5954,  2.8929, -1.0923],
        [ 1.1719, -0.4709, -0.1996]])

torch.randn_like(input, dtype=None, layout=None, device=None, requires_grad=False) → Tensor¶

Returns a tensor with the same size as input that is filled with random numbers from a normal distribution with mean 0 and variance 1. torch.randn_like(input) is equivalent to torch.randn(input.size(), dtype=input.dtype, layout=input.layout, device=input.device).

Parameters

input (Tensor) – the size of input will determine size of the output tensor
dtype (torch.dtype, optional) – the desired data type of returned Tensor. Default: if None, defaults to the dtype of input.
layout (torch.layout, optional) – the desired layout of returned tensor. Default: if None, defaults to the layout of input.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, defaults to the device of input.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

torch.randperm(n, out=None, dtype=torch.int64, layout=torch.strided, device=None, requires_grad=False) → LongTensor¶

Returns a random permutation of integers from 0 to n - 1.

Parameters

n (int) – the upper bound (exclusive)
out (Tensor, optional) – the output tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: torch.int64.
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> torch.randperm(4)
tensor([2, 1, 0, 3])

In-place random sampling¶

There are a few more in-place random sampling functions defined on Tensors as well. Click through to refer to their documentation:

torch.Tensor.bernoulli_() - in-place version of torch.bernoulli()
torch.Tensor.cauchy_() - numbers drawn from the Cauchy distribution
torch.Tensor.exponential_() - numbers drawn from the exponential distribution
torch.Tensor.geometric_() - elements drawn from the geometric distribution
torch.Tensor.log_normal_() - samples from the log-normal distribution
torch.Tensor.normal_() - in-place version of torch.normal()
torch.Tensor.random_() - numbers sampled from the discrete uniform distribution
torch.Tensor.uniform_() - numbers sampled from the continuous uniform distribution

Serialization¶

torch.save(obj, f, pickle_module=<module 'pickle'>, pickle_protocol=2)¶

Saves an object to a disk file.

Parallelism¶

torch.get_num_threads() → int¶: Gets the number of OpenMP threads used for parallelizing CPU operations

torch.set_num_threads(int)¶: Sets the number of OpenMP threads used for parallelizing CPU operations

Locally disabling gradient computation¶

The context managers torch.no_grad(), torch.enable_grad(), and torch.set_grad_enabled() are helpful for locally disabling and enabling gradient computation. See Locally disabling gradient computation for more details on their usage.

Examples:

>>> x = torch.zeros(1, requires_grad=True)
>>> with torch.no_grad():
...     y = x * 2
>>> y.requires_grad
False

>>> is_train = False
>>> with torch.set_grad_enabled(is_train):
...     y = x * 2
>>> y.requires_grad
False

>>> torch.set_grad_enabled(True)  # this can also be used as a function
>>> y = x * 2
>>> y.requires_grad
True

>>> torch.set_grad_enabled(False)
>>> y = x * 2
>>> y.requires_grad
False

Math operations¶

Pointwise Ops¶

torch.abs(input, out=None) → Tensor¶

Computes the element-wise absolute value of the given input tensor.

\[\text{out}_{i} = |\text{input}_{i}| \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> torch.abs(torch.tensor([-1, -2, 3]))
tensor([ 1,  2,  3])

torch.acos(input, out=None) → Tensor¶

Returns a new tensor with the arccosine of the elements of input.

\[\text{out}_{i} = \cos^{-1}(\text{input}_{i}) \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 0.3348, -0.5889,  0.2005, -0.1584])
>>> torch.acos(a)
tensor([ 1.2294,  2.2004,  1.3690,  1.7298])

torch.add()¶

torch.add(input, value, out=None)

Adds the scalar value to each element of the input input and returns a new resulting tensor.

\[\text{out} = \text{input} + \text{value} \]

If input is of type FloatTensor or DoubleTensor, value must be a real number, otherwise it should be an integer.

Parameters

input (Tensor) – the input tensor
value (Number) – the number to be added to each element of input

Keyword Arguments

out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 0.0202,  1.0985,  1.3506, -0.6056])
>>> torch.add(a, 20)
tensor([ 20.0202,  21.0985,  21.3506,  19.3944])

torch.add(input, value=1, other, out=None)

Each element of the tensor other is multiplied by the scalar value and added to each element of the tensor input. The resulting tensor is returned.

The shapes of input and other must be broadcastable.

\[\text{out} = \text{input} + \text{value} \times \text{other} \]

If other is of type FloatTensor or DoubleTensor, value must be a real number, otherwise it should be an integer.

Parameters

input (Tensor) – the first input tensor
value (Number) – the scalar multiplier for other
other (Tensor) – the second input tensor

Keyword Arguments

out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([-0.9732, -0.3497,  0.6245,  0.4022])
>>> b = torch.randn(4, 1)
>>> b
tensor([[ 0.3743],
        [-1.7724],
        [-0.5811],
        [-0.8017]])
>>> torch.add(a, 10, b)
tensor([[  2.7695,   3.3930,   4.3672,   4.1450],
        [-18.6971, -18.0736, -17.0994, -17.3216],
        [ -6.7845,  -6.1610,  -5.1868,  -5.4090],
        [ -8.9902,  -8.3667,  -7.3925,  -7.6147]])

torch.addcdiv(tensor, value=1, tensor1, tensor2, out=None) → Tensor¶

Performs the element-wise division of tensor1 by tensor2, multiply the result by the scalar value and add it to tensor.

\[\text{out}_i = \text{tensor}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i} \]

The shapes of tensor, tensor1, and tensor2 must be broadcastable.

For inputs of type FloatTensor or DoubleTensor, value must be a real number, otherwise an integer.

Parameters

tensor (Tensor) – the tensor to be added
value (Number, optional) – multiplier for $\text{tensor1} / \text{tensor2}$
tensor1 (Tensor) – the numerator tensor
tensor2 (Tensor) – the denominator tensor
out (Tensor, optional) – the output tensor

Example:

>>> t = torch.randn(1, 3)
>>> t1 = torch.randn(3, 1)
>>> t2 = torch.randn(1, 3)
>>> torch.addcdiv(t, 0.1, t1, t2)
tensor([[-0.2312, -3.6496,  0.1312],
        [-1.0428,  3.4292, -0.1030],
        [-0.5369, -0.9829,  0.0430]])

torch.addcmul(tensor, value=1, tensor1, tensor2, out=None) → Tensor¶

Performs the element-wise multiplication of tensor1 by tensor2, multiply the result by the scalar value and add it to tensor.

\[\text{out}_i = \text{tensor}_i + \text{value} \times \text{tensor1}_i \times \text{tensor2}_i \]

The shapes of tensor, tensor1, and tensor2 must be broadcastable.

For inputs of type FloatTensor or DoubleTensor, value must be a real number, otherwise an integer.

Parameters

tensor (Tensor) – the tensor to be added
value (Number, optional) – multiplier for $tensor1 .* tensor2$
tensor1 (Tensor) – the tensor to be multiplied
tensor2 (Tensor) – the tensor to be multiplied
out (Tensor, optional) – the output tensor

Example:

>>> t = torch.randn(1, 3)
>>> t1 = torch.randn(3, 1)
>>> t2 = torch.randn(1, 3)
>>> torch.addcmul(t, 0.1, t1, t2)
tensor([[-0.8635, -0.6391,  1.6174],
        [-0.7617, -0.5879,  1.7388],
        [-0.8353, -0.6249,  1.6511]])

torch.asin(input, out=None) → Tensor¶

Returns a new tensor with the arcsine of the elements of input.

\[\text{out}_{i} = \sin^{-1}(\text{input}_{i}) \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([-0.5962,  1.4985, -0.4396,  1.4525])
>>> torch.asin(a)
tensor([-0.6387,     nan, -0.4552,     nan])

torch.atan(input, out=None) → Tensor¶

Returns a new tensor with the arctangent of the elements of input.

\[\text{out}_{i} = \tan^{-1}(\text{input}_{i}) \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 0.2341,  0.2539, -0.6256, -0.6448])
>>> torch.atan(a)
tensor([ 0.2299,  0.2487, -0.5591, -0.5727])

torch.atan2(input1, input2, out=None) → Tensor¶

Returns a new tensor with the arctangent of the elements of input1 and input2.

The shapes of input1 and input2 must be broadcastable.

Parameters

input1 (Tensor) – the first input tensor
input2 (Tensor) – the second input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 0.9041,  0.0196, -0.3108, -2.4423])
>>> torch.atan2(a, torch.randn(4))
tensor([ 0.9833,  0.0811, -1.9743, -1.4151])

torch.ceil(input, out=None) → Tensor¶

Returns a new tensor with the ceil of the elements of input, the smallest integer greater than or equal to each element.

\[\text{out}_{i} = \left\lceil \text{input}_{i} \right\rceil = \left\lfloor \text{input}_{i} \right\rfloor + 1 \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([-0.6341, -1.4208, -1.0900,  0.5826])
>>> torch.ceil(a)
tensor([-0., -1., -1.,  1.])

torch.clamp(input, min, max, out=None) → Tensor¶

Clamp all elements in input into the range [ min, max ] and return a resulting tensor:

\[y_i = \begin{cases} \text{min} & \text{if } x_i < \text{min} \\ x_i & \text{if } \text{min} \leq x_i \leq \text{max} \\ \text{max} & \text{if } x_i > \text{max} \end{cases} \]

If input is of type FloatTensor or DoubleTensor, args min and max must be real numbers, otherwise they should be integers.

Parameters

input (Tensor) – the input tensor
min (Number) – lower-bound of the range to be clamped to
max (Number) – upper-bound of the range to be clamped to
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([-1.7120,  0.1734, -0.0478, -0.0922])
>>> torch.clamp(a, min=-0.5, max=0.5)
tensor([-0.5000,  0.1734, -0.0478, -0.0922])

torch.clamp(input, *, min, out=None) → Tensor

Clamps all elements in input to be larger or equal min.

If input is of type FloatTensor or DoubleTensor, value should be a real number, otherwise it should be an integer.

Parameters

input (Tensor) – the input tensor
value (Number) – minimal value of each element in the output
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([-0.0299, -2.3184,  2.1593, -0.8883])
>>> torch.clamp(a, min=0.5)
tensor([ 0.5000,  0.5000,  2.1593,  0.5000])

torch.clamp(input, *, max, out=None) → Tensor

Clamps all elements in input to be smaller or equal max.

If input is of type FloatTensor or DoubleTensor, value should be a real number, otherwise it should be an integer.

Parameters

input (Tensor) – the input tensor
value (Number) – maximal value of each element in the output
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 0.7753, -0.4702, -0.4599,  1.1899])
>>> torch.clamp(a, max=0.5)
tensor([ 0.5000, -0.4702, -0.4599,  0.5000])

torch.cos(input, out=None) → Tensor¶

Returns a new tensor with the cosine of the elements of input.

\[\text{out}_{i} = \cos(\text{input}_{i}) \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 1.4309,  1.2706, -0.8562,  0.9796])
>>> torch.cos(a)
tensor([ 0.1395,  0.2957,  0.6553,  0.5574])

torch.cosh(input, out=None) → Tensor¶

Returns a new tensor with the hyperbolic cosine of the elements of input.

\[\text{out}_{i} = \cosh(\text{input}_{i}) \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 0.1632,  1.1835, -0.6979, -0.7325])
>>> torch.cosh(a)
tensor([ 1.0133,  1.7860,  1.2536,  1.2805])

torch.div()¶

torch.div(input, value, out=None) → Tensor

Divides each element of the input input with the scalar value and returns a new resulting tensor.

\[\text{out}_i = \frac{\text{input}_i}{\text{value}} \]

If input is of type FloatTensor or DoubleTensor, value should be a real number, otherwise it should be an integer

Parameters

input (Tensor) – the input tensor
value (Number) – the number to be divided to each element of input
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(5)
>>> a
tensor([ 0.3810,  1.2774, -0.2972, -0.3719,  0.4637])
>>> torch.div(a, 0.5)
tensor([ 0.7620,  2.5548, -0.5944, -0.7439,  0.9275])

torch.div(input, other, out=None) → Tensor

Each element of the tensor input is divided by each element of the tensor other. The resulting tensor is returned. The shapes of input and other must be broadcastable.

\[\text{out}_i = \frac{\text{input}_i}{\text{other}_i} \]

Parameters

input (Tensor) – the numerator tensor
other (Tensor) – the denominator tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[-0.3711, -1.9353, -0.4605, -0.2917],
        [ 0.1815, -1.0111,  0.9805, -1.5923],
        [ 0.1062,  1.4581,  0.7759, -1.2344],
        [-0.1830, -0.0313,  1.1908, -1.4757]])
>>> b = torch.randn(4)
>>> b
tensor([ 0.8032,  0.2930, -0.8113, -0.2308])
>>> torch.div(a, b)
tensor([[-0.4620, -6.6051,  0.5676,  1.2637],
        [ 0.2260, -3.4507, -1.2086,  6.8988],
        [ 0.1322,  4.9764, -0.9564,  5.3480],
        [-0.2278, -0.1068, -1.4678,  6.3936]])

torch.digamma(input, out=None) → Tensor¶

Computes the logarithmic derivative of the gamma function on input.

\[\psi(x) = \frac{d}{dx} \ln\left(\Gamma\left(x\right)\right) = \frac{\Gamma'(x)}{\Gamma(x)} \]

Parameters: input (Tensor) – the tensor to compute the digamma function on

Example:

>>> a = torch.tensor([1, 0.5])
>>> torch.digamma(a)
tensor([-0.5772, -1.9635])

torch.erf(tensor, out=None) → Tensor¶

Computes the error function of each element. The error function is defined as follows:

\[\mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_{0}^{x} e^{-t^2} dt \]

Parameters

tensor (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> torch.erf(torch.tensor([0, -1., 10.]))
tensor([ 0.0000, -0.8427,  1.0000])

torch.erfc(input, out=None) → Tensor¶

Computes the complementary error function of each element of input. The complementary error function is defined as follows:

\[\mathrm{erfc}(x) = 1 - \frac{2}{\sqrt{\pi}} \int_{0}^{x} e^{-t^2} dt \]

Parameters

tensor (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> torch.erfc(torch.tensor([0, -1., 10.]))
tensor([ 1.0000, 1.8427,  0.0000])

torch.erfinv(input, out=None) → Tensor¶

Computes the inverse error function of each element of input. The inverse error function is defined in the range $(-1, 1)$ as:

\[\mathrm{erfinv}(\mathrm{erf}(x)) = x \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> torch.erfinv(torch.tensor([0, 0.5, -1.]))
tensor([ 0.0000,  0.4769,    -inf])

torch.exp(input, out=None) → Tensor¶

Returns a new tensor with the exponential of the elements of the input tensor input.

\[y_{i} = e^{x_{i}} \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> torch.exp(torch.tensor([0, math.log(2.)]))
tensor([ 1.,  2.])

torch.expm1(input, out=None) → Tensor¶

Returns a new tensor with the exponential of the elements minus 1 of input.

\[y_{i} = e^{x_{i}} - 1 \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> torch.expm1(torch.tensor([0, math.log(2.)]))
tensor([ 0.,  1.])

torch.floor(input, out=None) → Tensor¶

Returns a new tensor with the floor of the elements of input, the largest integer less than or equal to each element.

\[\text{out}_{i} = \left\lfloor \text{input}_{i} \right\rfloor \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([-0.8166,  1.5308, -0.2530, -0.2091])
>>> torch.floor(a)
tensor([-1.,  1., -1., -1.])

torch.fmod(input, divisor, out=None) → Tensor¶

Computes the element-wise remainder of division.

The dividend and divisor may contain both for integer and floating point numbers. The remainder has the same sign as the dividend input.

When divisor is a tensor, the shapes of input and divisor must be broadcastable.

Parameters

input (Tensor) – the dividend
divisor (Tensor or float) – the divisor, which may be either a number or a tensor of the same shape as the dividend
out (Tensor, optional) – the output tensor

Example:

>>> torch.fmod(torch.tensor([-3., -2, -1, 1, 2, 3]), 2)
tensor([-1., -0., -1.,  1.,  0.,  1.])
>>> torch.fmod(torch.tensor([1., 2, 3, 4, 5]), 1.5)
tensor([ 1.0000,  0.5000,  0.0000,  1.0000,  0.5000])

torch.frac(input, out=None) → Tensor¶

Computes the fractional portion of each element in input.

\[\text{out}_{i} = \text{input}_{i} - \left\lfloor \text{input}_{i} \right\rfloor \]

Example:

>>> torch.frac(torch.tensor([1, 2.5, -3.2]))
tensor([ 0.0000,  0.5000, -0.2000])

torch.lerp(start, end, weight, out=None)¶

Does a linear interpolation of two tensors start and end based on a scalar or tensor weight and returns the resulting out tensor.

\[\text{out}_i = \text{start}_i + \text{weight}_i \times (\text{end}_i - \text{start}_i) \]

The shapes of start and end must be broadcastable. If weight is a tensor, then the shapes of start, end must be broadcastable.

Parameters

start (Tensor) – the tensor with the starting points
end (Tensor) – the tensor with the ending points
weight (float or tensor) – the weight for the interpolation formula
out (Tensor, optional) – the output tensor

Example:

>>> start = torch.arange(1., 5.)
>>> end = torch.empty(4).fill_(10)
>>> start
tensor([ 1.,  2.,  3.,  4.])
>>> end
tensor([ 10.,  10.,  10.,  10.])
>>> torch.lerp(start, end, 0.5)
tensor([ 5.5000,  6.0000,  6.5000,  7.0000])
>>> torch.lerp(start, end, torch.full_like(start, 0.5))
tensor([ 5.5000,  6.0000,  6.5000,  7.0000])

torch.log(input, out=None) → Tensor¶

Returns a new tensor with the natural logarithm of the elements of input.

\[y_{i} = \log_{e} (x_{i}) \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(5)
>>> a
tensor([-0.7168, -0.5471, -0.8933, -1.4428, -0.1190])
>>> torch.log(a)
tensor([ nan,  nan,  nan,  nan,  nan])

torch.log10(input, out=None) → Tensor¶

Returns a new tensor with the logarithm to the base 10 of the elements of input.

\[y_{i} = \log_{10} (x_{i}) \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.rand(5)
>>> a
tensor([ 0.5224,  0.9354,  0.7257,  0.1301,  0.2251])


>>> torch.log10(a)
tensor([-0.2820, -0.0290, -0.1392, -0.8857, -0.6476])

torch.log1p(input, out=None) → Tensor¶

Returns a new tensor with the natural logarithm of (1 + input).

\[y_i = \log_{e} (x_i + 1) \]

Note

This function is more accurate than torch.log() for small values of input

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(5)
>>> a
tensor([-1.0090, -0.9923,  1.0249, -0.5372,  0.2492])
>>> torch.log1p(a)
tensor([    nan, -4.8653,  0.7055, -0.7705,  0.2225])

torch.log2(input, out=None) → Tensor¶

Returns a new tensor with the logarithm to the base 2 of the elements of input.

\[y_{i} = \log_{2} (x_{i}) \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.rand(5)
>>> a
tensor([ 0.8419,  0.8003,  0.9971,  0.5287,  0.0490])


>>> torch.log2(a)
tensor([-0.2483, -0.3213, -0.0042, -0.9196, -4.3504])

torch.mul()¶

torch.mul(input, value, out=None)

Multiplies each element of the input input with the scalar value and returns a new resulting tensor.

\[\text{out}_i = \text{value} \times \text{input}_i \]

If input is of type FloatTensor or DoubleTensor, value should be a real number, otherwise it should be an integer

Parameters

input (Tensor) – the input tensor
value (Number) – the number to be multiplied to each element of input
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(3)
>>> a
tensor([ 0.2015, -0.4255,  2.6087])
>>> torch.mul(a, 100)
tensor([  20.1494,  -42.5491,  260.8663])

torch.mul(input, other, out=None)

Each element of the tensor input is multiplied by the corresponding element of the Tensor other. The resulting tensor is returned.

The shapes of input and other must be broadcastable.

\[\text{out}_i = \text{input}_i \times \text{other}_i \]

Parameters

input (Tensor) – the first multiplicand tensor
other (Tensor) – the second multiplicand tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4, 1)
>>> a
tensor([[ 1.1207],
        [-0.3137],
        [ 0.0700],
        [ 0.8378]])
>>> b = torch.randn(1, 4)
>>> b
tensor([[ 0.5146,  0.1216, -0.5244,  2.2382]])
>>> torch.mul(a, b)
tensor([[ 0.5767,  0.1363, -0.5877,  2.5083],
        [-0.1614, -0.0382,  0.1645, -0.7021],
        [ 0.0360,  0.0085, -0.0367,  0.1567],
        [ 0.4312,  0.1019, -0.4394,  1.8753]])

torch.mvlgamma(input, p) → Tensor¶

Computes the multivariate log-gamma function ([reference]) with dimension $p$ element-wise, given by

\[\log(\Gamma_{p}(a)) = C + \displaystyle \sum_{i=1}^{p} \log\left(\Gamma\left(a - \frac{i - 1}{2}\right)\right) \]

where $C = \log(\pi) \times \frac{p (p - 1)}{4}$ and $\Gamma(\cdot)$ is the Gamma function.

If any of the elements are less than or equal to $\frac{p - 1}{2}$, then an error is thrown.

Parameters

input (Tensor) – the tensor to compute the multivariate log-gamma function
p (int) – the number of dimensions

Example:

>>> a = torch.empty(2, 3).uniform_(1, 2)
>>> a
tensor([[1.6835, 1.8474, 1.1929],
        [1.0475, 1.7162, 1.4180]])
>>> torch.mvlgamma(a, 2)
tensor([[0.3928, 0.4007, 0.7586],
        [1.0311, 0.3901, 0.5049]])

torch.neg(input, out=None) → Tensor¶

Returns a new tensor with the negative of the elements of input.

\[\text{out} = -1 \times \text{input} \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(5)
>>> a
tensor([ 0.0090, -0.2262, -0.0682, -0.2866,  0.3940])
>>> torch.neg(a)
tensor([-0.0090,  0.2262,  0.0682,  0.2866, -0.3940])

torch.pow()¶

torch.pow(input, exponent, out=None) → Tensor

Takes the power of each element in input with exponent and returns a tensor with the result.

exponent can be either a single float number or a Tensor with the same number of elements as input.

When exponent is a scalar value, the operation applied is:

\[\text{out}_i = x_i ^ \text{exponent} \]

When exponent is a tensor, the operation applied is:

\[\text{out}_i = x_i ^ {\text{exponent}_i} \]

When exponent is a tensor, the shapes of input and exponent must be broadcastable.

Parameters

input (Tensor) – the input tensor
exponent (float or tensor) – the exponent value
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 0.4331,  1.2475,  0.6834, -0.2791])
>>> torch.pow(a, 2)
tensor([ 0.1875,  1.5561,  0.4670,  0.0779])
>>> exp = torch.arange(1., 5.)

>>> a = torch.arange(1., 5.)
>>> a
tensor([ 1.,  2.,  3.,  4.])
>>> exp
tensor([ 1.,  2.,  3.,  4.])
>>> torch.pow(a, exp)
tensor([   1.,    4.,   27.,  256.])

torch.pow(base, input, out=None) → Tensor

base is a scalar float value, and input is a tensor. The returned tensor out is of the same shape as input

The operation applied is:

\[out_i = base ^ {input_i} \]

Parameters

base (float) – the scalar base value for the power operation
input (Tensor) – the exponent tensor
out (Tensor, optional) – the output tensor

Example:

>>> exp = torch.arange(1., 5.)
>>> base = 2
>>> torch.pow(base, exp)
tensor([  2.,   4.,   8.,  16.])

torch.reciprocal(input, out=None) → Tensor¶

Returns a new tensor with the reciprocal of the elements of input

\[\text{out}_{i} = \frac{1}{\text{input}_{i}} \]

Parameters

input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([-0.4595, -2.1219, -1.4314,  0.7298])
>>> torch.reciprocal(a)
tensor([-2.1763, -0.4713, -0.6986,  1.3702])

torch.remainder(input, divisor, out=None) → Tensor¶

Computes the element-wise remainder of division.

The divisor and dividend may contain both for integer and floating point numbers. The remainder has the same sign as the divisor.

When divisor is a tensor, the shapes of input and divisor must be broadcastable.

Parameters

input (Tensor) – the dividend
divisor (Tensor or float) – the divisor that may be either a number or a Tensor of the same shape as the dividend
out (Tensor, optional) – the output tensor

Example:

>>> torch.remainder(torch.tensor([-3., -2, -1, 1, 2, 3]), 2)
tensor([ 1.,  0.,  1.,  1.,  0.,  1.])
>>> torch.remainder(torch.tensor([1., 2, 3, 4, 5]), 1.5)
tensor([ 1.0000,  0.5000,  0.0000,  1.0000,  0.5000])

Reduction Ops¶

torch.argmax(input, dim=None, keepdim=False)¶

Returns the indices of the maximum values of a tensor across a dimension.

This is the second value returned by torch.max(). See its documentation for the exact semantics of this method.

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension to reduce. If None, the argmax of the flattened input is returned.
keepdim (bool) – whether the output tensors have dim retained or not. Ignored if dim=None.

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 1.3398,  0.2663, -0.2686,  0.2450],
        [-0.7401, -0.8805, -0.3402, -1.1936],
        [ 0.4907, -1.3948, -1.0691, -0.3132],
        [-1.6092,  0.5419, -0.2993,  0.3195]])


>>> torch.argmax(a, dim=1)
tensor([ 0,  2,  0,  1])

torch.argmin(input, dim=None, keepdim=False)¶

Returns the indices of the minimum values of a tensor across a dimension.

This is the second value returned by torch.min(). See its documentation for the exact semantics of this method.

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension to reduce. If None, the argmin of the flattened input is returned.
keepdim (bool) – whether the output tensors have dim retained or not. Ignored if dim=None.

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 0.1139,  0.2254, -0.1381,  0.3687],
        [ 1.0100, -1.1975, -0.0102, -0.4732],
        [-0.9240,  0.1207, -0.7506, -1.0213],
        [ 1.7809, -1.2960,  0.9384,  0.1438]])


>>> torch.argmin(a, dim=1)
tensor([ 2,  1,  3,  1])

torch.cumprod(input, dim, dtype=None) → Tensor¶

Returns the cumulative product of elements of input in the dimension dim.

For example, if input is a vector of size N, the result will also be a vector of size N, with elements.

\[y_i = x_1 \times x_2\times x_3\times \dots \times x_i \]

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension to do the operation over
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

Example:

>>> a = torch.randn(10)
>>> a
tensor([ 0.6001,  0.2069, -0.1919,  0.9792,  0.6727,  1.0062,  0.4126,
        -0.2129, -0.4206,  0.1968])
>>> torch.cumprod(a, dim=0)
tensor([ 0.6001,  0.1241, -0.0238, -0.0233, -0.0157, -0.0158, -0.0065,
         0.0014, -0.0006, -0.0001])

>>> a[5] = 0.0
>>> torch.cumprod(a, dim=0)
tensor([ 0.6001,  0.1241, -0.0238, -0.0233, -0.0157, -0.0000, -0.0000,
         0.0000, -0.0000, -0.0000])

torch.cumsum(input, dim, out=None, dtype=None) → Tensor¶

Returns the cumulative sum of elements of input in the dimension dim.

For example, if input is a vector of size N, the result will also be a vector of size N, with elements.

\[y_i = x_1 + x_2 + x_3 + \dots + x_i \]

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension to do the operation over
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

Example:

>>> a = torch.randn(10)
>>> a
tensor([-0.8286, -0.4890,  0.5155,  0.8443,  0.1865, -0.1752, -2.0595,
         0.1850, -1.1571, -0.4243])
>>> torch.cumsum(a, dim=0)
tensor([-0.8286, -1.3175, -0.8020,  0.0423,  0.2289,  0.0537, -2.0058,
        -1.8209, -2.9780, -3.4022])

torch.dist(input, other, p=2) → Tensor¶

Returns the p-norm of (input - other)

The shapes of input and other must be broadcastable.

Parameters

input (Tensor) – the input tensor
other (Tensor) – the Right-hand-side input tensor
p (float, optional) – the norm to be computed

Example:

>>> x = torch.randn(4)
>>> x
tensor([-1.5393, -0.8675,  0.5916,  1.6321])
>>> y = torch.randn(4)
>>> y
tensor([ 0.0967, -1.0511,  0.6295,  0.8360])
>>> torch.dist(x, y, 3.5)
tensor(1.6727)
>>> torch.dist(x, y, 3)
tensor(1.6973)
>>> torch.dist(x, y, 0)
tensor(inf)
>>> torch.dist(x, y, 1)
tensor(2.6537)

torch.logsumexp(input, dim, keepdim=False, out=None)¶

Returns the log of summed exponentials of each row of the input tensor in the given dimension dim. The computation is numerically stabilized.

For summation index $j$ given by dim and other indices $i$, the result is

\[\text{logsumexp}(x)_{i} = \log \sum_j \exp(x_{ij}) \]

If keepdim is True, the output tensor is of the same size as input except in the dimension(s) dim where it is of size 1. Otherwise, dim is squeezed (see torch.squeeze()), resulting in the output tensor having 1 (or len(dim)) fewer dimension(s).

Parameters

input (Tensor) – the input tensor
dim (int or tuple of python:ints) – the dimension or dimensions to reduce
keepdim (bool) – whether the output tensor has dim retained or not
out (Tensor, optional) – the output tensor

Example::

>>> a = torch.randn(3, 3)
>>> torch.logsumexp(a, 1)
tensor([ 0.8442,  1.4322,  0.8711])

torch.mean()¶

torch.mean(input) → Tensor

Returns the mean value of all elements in the input tensor.

Parameters: input (Tensor) – the input tensor

Example:

>>> a = torch.randn(1, 3)
>>> a
tensor([[ 0.2294, -0.5481,  1.3288]])
>>> torch.mean(a)
tensor(0.3367)

torch.mean(input, dim, keepdim=False, out=None) → Tensor

Returns the mean value of each row of the input tensor in the given dimension dim. If dim is a list of dimensions, reduce over all of them.

Parameters

input (Tensor) – the input tensor
dim (int or tuple of python:ints) – the dimension or dimensions to reduce
keepdim (bool) – whether the output tensor has dim retained or not
out (Tensor) – the output tensor

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[-0.3841,  0.6320,  0.4254, -0.7384],
        [-0.9644,  1.0131, -0.6549, -1.4279],
        [-0.2951, -1.3350, -0.7694,  0.5600],
        [ 1.0842, -0.9580,  0.3623,  0.2343]])
>>> torch.mean(a, 1)
tensor([-0.0163, -0.5085, -0.4599,  0.1807])
>>> torch.mean(a, 1, True)
tensor([[-0.0163],
        [-0.5085],
        [-0.4599],
        [ 0.1807]])

torch.median()¶

torch.median(input) → Tensor

Returns the median value of all elements in the input tensor.

Parameters: input (Tensor) – the input tensor

Example:

>>> a = torch.randn(1, 3)
>>> a
tensor([[ 1.5219, -1.5212,  0.2202]])
>>> torch.median(a)
tensor(0.2202)

torch.median(input, dim=-1, keepdim=False, values=None, indices=None) -> (Tensor, LongTensor)

Returns a namedtuple (values, indices) where values is the median value of each row of the input tensor in the given dimension dim. And indices is the index location of each median value found.

By default, dim is the last dimension of the input tensor.

If keepdim is True, the output tensors are of the same size as input except in the dimension dim where they are of size 1. Otherwise, dim is squeezed (see torch.squeeze()), resulting in the outputs tensor having 1 fewer dimension than input.

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension to reduce
keepdim (bool) – whether the output tensors have dim retained or not
values (Tensor, optional) – the output tensor
indices (Tensor, optional) – the output index tensor

Example:

>>> a = torch.randn(4, 5)
>>> a
tensor([[ 0.2505, -0.3982, -0.9948,  0.3518, -1.3131],
        [ 0.3180, -0.6993,  1.0436,  0.0438,  0.2270],
        [-0.2751,  0.7303,  0.2192,  0.3321,  0.2488],
        [ 1.0778, -1.9510,  0.7048,  0.4742, -0.7125]])
>>> torch.median(a, 1)
torch.return_types.median(values=tensor([-0.3982,  0.2270,  0.2488,  0.4742]), indices=tensor([1, 4, 4, 3]))

torch.mode(input, dim=-1, keepdim=False, values=None, indices=None) -> (Tensor, LongTensor)¶

Returns a namedtuple (values, indices) where values is the mode value of each row of the input tensor in the given dimension dim, i.e. a value which appears most often in that row, and indices is the index location of each mode value found.

By default, dim is the last dimension of the input tensor.

Note

This function is not defined for torch.cuda.Tensor yet.

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension to reduce
keepdim (bool) – whether the output tensors have dim retained or not
values (Tensor, optional) – the output tensor
indices (Tensor, optional) – the output index tensor

Example:

>>> a = torch.randint(10, (5,))
>>> a
tensor([6, 5, 1, 0, 2])
>>> b = a + (torch.randn(50, 1) * 5).long()
>>> torch.mode(b, 0)
torch.return_types.mode(values=tensor([6, 5, 1, 0, 2]), indices=tensor([2, 2, 2, 2, 2]))

torch.norm(input, p='fro', dim=None, keepdim=False, out=None, dtype=None)¶

Returns the matrix norm or vector norm of a given tensor.

Parameters

input (Tensor) – the input tensor
p (int, float, inf, -inf, 'fro', 'nuc', optional) –
the order of norm. Default: 'fro' The following norms can be calculated:

ord

matrix norm

vector norm

None

Frobenius norm

2-norm

’fro’

Frobenius norm

–

‘nuc’

nuclear norm

–

Other

as vec norm when dim is None

sum(abs(x)**ord)**(1./ord)
dim (int, 2-tuple of python:ints, 2-list of python:ints, optional) – If it is an int, vector norm will be calculated, if it is 2-tuple of ints, matrix norm will be calculated. If the value is None, matrix norm will be calculated when the input tensor only has two dimensions, vector norm will be calculated when the input tensor only has one dimension. If the input tensor has more than two dimensions, the vector norm will be applied to last dimension.
keepdim (bool, optional) – whether the output tensors have dim retained or not. Ignored if dim = None and out = None. Default: False
out (Tensor, optional) – the output tensor. Ignored if dim = None and out = None.
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to :attr:’dtype’ while performing the operation. Default: None.

ord	matrix norm	vector norm
None	Frobenius norm	2-norm
’fro’	Frobenius norm	–
‘nuc’	nuclear norm	–
Other	as vec norm when dim is None	sum(abs(x)ord)(1./ord)

Example:

>>> import torch
>>> a = torch.arange(9, dtype= torch.float) - 4
>>> b = a.reshape((3, 3))
>>> torch.norm(a)
tensor(7.7460)
>>> torch.norm(b)
tensor(7.7460)
>>> torch.norm(a, float('inf'))
tensor(4.)
>>> torch.norm(b, float('inf'))
tensor(4.)
>>> c = torch.tensor([[ 1, 2, 3],[-1, 1, 4]] , dtype= torch.float)
>>> torch.norm(c, dim=0)
tensor([1.4142, 2.2361, 5.0000])
>>> torch.norm(c, dim=1)
tensor([3.7417, 4.2426])
>>> torch.norm(c, p=1, dim=1)
tensor([6., 6.])
>>> d = torch.arange(8, dtype= torch.float).reshape(2,2,2)
>>> torch.norm(d, dim=(1,2))
tensor([ 3.7417, 11.2250])
>>> torch.norm(d[0, :, :]), torch.norm(d[1, :, :])
(tensor(3.7417), tensor(11.2250))

torch.prod()¶

torch.prod(input, dtype=None) → Tensor

Returns the product of all elements in the input tensor.

Parameters

input (Tensor) – the input tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

Example:

>>> a = torch.randn(1, 3)
>>> a
tensor([[-0.8020,  0.5428, -1.5854]])
>>> torch.prod(a)
tensor(0.6902)

torch.prod(input, dim, keepdim=False, dtype=None) → Tensor

Returns the product of each row of the input tensor in the given dimension dim.

If keepdim is True, the output tensor is of the same size as input except in the dimension dim where it is of size 1. Otherwise, dim is squeezed (see torch.squeeze()), resulting in the output tensor having 1 fewer dimension than input.

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension to reduce
keepdim (bool) – whether the output tensor has dim retained or not
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

Example:

>>> a = torch.randn(4, 2)
>>> a
tensor([[ 0.5261, -0.3837],
        [ 1.1857, -0.2498],
        [-1.1646,  0.0705],
        [ 1.1131, -1.0629]])
>>> torch.prod(a, 1)
tensor([-0.2018, -0.2962, -0.0821, -1.1831])

torch.std()¶

torch.std(input, unbiased=True) → Tensor

Returns the standard-deviation of all elements in the input tensor.

If unbiased is False, then the standard-deviation will be calculated via the biased estimator. Otherwise, Bessel’s correction will be used.

Parameters

input (Tensor) – the input tensor
unbiased (bool) – whether to use the unbiased estimation or not

Example:

>>> a = torch.randn(1, 3)
>>> a
tensor([[-0.8166, -1.3802, -0.3560]])
>>> torch.std(a)
tensor(0.5130)

torch.std(input, dim, keepdim=False, unbiased=True, out=None) → Tensor

Returns the standard-deviation of each row of the input tensor in the dimension dim. If dim is a list of dimensions, reduce over all of them.

If unbiased is False, then the standard-deviation will be calculated via the biased estimator. Otherwise, Bessel’s correction will be used.

Parameters

input (Tensor) – the input tensor
dim (int or tuple of python:ints) – the dimension or dimensions to reduce
keepdim (bool) – whether the output tensor has dim retained or not
unbiased (bool) – whether to use the unbiased estimation or not
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 0.2035,  1.2959,  1.8101, -0.4644],
        [ 1.5027, -0.3270,  0.5905,  0.6538],
        [-1.5745,  1.3330, -0.5596, -0.6548],
        [ 0.1264, -0.5080,  1.6420,  0.1992]])
>>> torch.std(a, dim=1)
tensor([ 1.0311,  0.7477,  1.2204,  0.9087])

torch.sum()¶

torch.sum(input, dtype=None) → Tensor

Returns the sum of all elements in the input tensor.

Parameters

input (Tensor) – the input tensor
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

Example:

>>> a = torch.randn(1, 3)
>>> a
tensor([[ 0.1133, -0.9567,  0.2958]])
>>> torch.sum(a)
tensor(-0.5475)

torch.sum(input, dim, keepdim=False, dtype=None) → Tensor

Returns the sum of each row of the input tensor in the given dimension dim. If dim is a list of dimensions, reduce over all of them.

Parameters

input (Tensor) – the input tensor
dim (int or tuple of python:ints) – the dimension or dimensions to reduce
keepdim (bool) – whether the output tensor has dim retained or not
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 0.0569, -0.2475,  0.0737, -0.3429],
        [-0.2993,  0.9138,  0.9337, -1.6864],
        [ 0.1132,  0.7892, -0.1003,  0.5688],
        [ 0.3637, -0.9906, -0.4752, -1.5197]])
>>> torch.sum(a, 1)
tensor([-0.4598, -0.1381,  1.3708, -2.6217])
>>> b = torch.arange(4 * 5 * 6).view(4, 5, 6)
>>> torch.sum(b, (2, 1))
tensor([  435.,  1335.,  2235.,  3135.])

torch.unique(input, sorted=True, return_inverse=False, dim=None)¶

Returns the unique scalar elements of the input tensor as a 1-D tensor.

Parameters

input (Tensor) – the input tensor
sorted (bool) – Whether to sort the unique elements in ascending order before returning as output.
return_inverse (bool) – Whether to also return the indices for where elements in the original input ended up in the returned unique list.
dim (int) – the dimension to apply unique. If None, the unique of the flattened input is returned. default: None

Returns

A tensor or a tuple of tensors containing

output (Tensor): the output list of unique scalar elements.

inverse_indices (Tensor): (optional) if return_inverse is True, there will be a 2nd returned tensor (same shape as input) representing the indices for where elements in the original input map to in the output; otherwise, this function will only return a single tensor.

Return type

(Tensor, Tensor (optional))

Example:

>>> output = torch.unique(torch.tensor([1, 3, 2, 3], dtype=torch.long))
>>> output
tensor([ 2,  3,  1])

>>> output, inverse_indices = torch.unique(
        torch.tensor([1, 3, 2, 3], dtype=torch.long), sorted=True, return_inverse=True)
>>> output
tensor([ 1,  2,  3])
>>> inverse_indices
tensor([ 0,  2,  1,  2])

>>> output, inverse_indices = torch.unique(
        torch.tensor([[1, 3], [2, 3]], dtype=torch.long), sorted=True, return_inverse=True)
>>> output
tensor([ 1,  2,  3])
>>> inverse_indices
tensor([[ 0,  2],
        [ 1,  2]])

torch.var()¶

torch.var(input, unbiased=True) → Tensor

Returns the variance of all elements in the input tensor.

If unbiased is False, then the variance will be calculated via the biased estimator. Otherwise, Bessel’s correction will be used.

Parameters

input (Tensor) – the input tensor
unbiased (bool) – whether to use the unbiased estimation or not

Example:

>>> a = torch.randn(1, 3)
>>> a
tensor([[-0.3425, -1.2636, -0.4864]])
>>> torch.var(a)
tensor(0.2455)

torch.var(input, dim, keepdim=False, unbiased=True, out=None) → Tensor

Returns the variance of each row of the input tensor in the given dimension dim.

If unbiased is False, then the variance will be calculated via the biased estimator. Otherwise, Bessel’s correction will be used.

Parameters

input (Tensor) – the input tensor
dim (int or tuple of python:ints) – the dimension or dimensions to reduce
keepdim (bool) – whether the output tensor has dim retained or not
unbiased (bool) – whether to use the unbiased estimation or not
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[-0.3567,  1.7385, -1.3042,  0.7423],
        [ 1.3436, -0.1015, -0.9834, -0.8438],
        [ 0.6056,  0.1089, -0.3112, -1.4085],
        [-0.7700,  0.6074, -0.1469,  0.7777]])
>>> torch.var(a, 1)
tensor([ 1.7444,  1.1363,  0.7356,  0.5112])

Comparison Ops¶

torch.allclose(self, other, rtol=1e-05, atol=1e-08, equal_nan=False) → bool¶

This function checks if all self and other satisfy the condition:

\[\lvert \text{self} - \text{other} \rvert \leq \texttt{atol} + \texttt{rtol} \times \lvert \text{other} \rvert \]

elementwise, for all elements of self and other. The behaviour of this function is analogous to numpy.allclose

Parameters

self (Tensor) – first tensor to compare
other (Tensor) – second tensor to compare
atol (float, optional) – absolute tolerance. Default: 1e-08
rtol (float, optional) – relative tolerance. Default: 1e-05
equal_nan (float, optional) – if True, then two NaN s will be compared as equal. Default: False

Example:

>>> torch.allclose(torch.tensor([10000., 1e-07]), torch.tensor([10000.1, 1e-08]))
False
>>> torch.allclose(torch.tensor([10000., 1e-08]), torch.tensor([10000.1, 1e-09]))
True
>>> torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')]))
False
>>> torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')]), equal_nan=True)
True

torch.argsort(input, dim=-1, descending=False, out=None) → LongTensor¶

Returns the indices that sort a tensor along a given dimension in ascending order by value.

This is the second value returned by torch.sort(). See its documentation for the exact semantics of this method.

Parameters

input (Tensor) – the input tensor
dim (int, optional) – the dimension to sort along
descending (bool, optional) – controls the sorting order (ascending or descending)

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 0.0785,  1.5267, -0.8521,  0.4065],
        [ 0.1598,  0.0788, -0.0745, -1.2700],
        [ 1.2208,  1.0722, -0.7064,  1.2564],
        [ 0.0669, -0.2318, -0.8229, -0.9280]])


>>> torch.argsort(a, dim=1)
tensor([[2, 0, 3, 1],
        [3, 2, 1, 0],
        [2, 1, 0, 3],
        [3, 2, 1, 0]])

torch.eq(input, other, out=None) → Tensor¶

Computes element-wise equality

The second argument can be a number or a tensor whose shape is broadcastable with the first argument.

Parameters

input (Tensor) – the tensor to compare
other (Tensor or float) – the tensor or value to compare
out (Tensor, optional) – the output tensor. Must be a ByteTensor

Returns

A torch.ByteTensor containing a 1 at each location where comparison is true

Return type

Tensor

Example:

>>> torch.eq(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]]))
tensor([[ 1,  0],
        [ 0,  1]], dtype=torch.uint8)

torch.equal(tensor1, tensor2) → bool¶

True if two tensors have the same size and elements, False otherwise.

Example:

>>> torch.equal(torch.tensor([1, 2]), torch.tensor([1, 2]))
True

torch.ge(input, other, out=None) → Tensor¶

Computes $\text{input} \geq \text{other}$ element-wise.

The second argument can be a number or a tensor whose shape is broadcastable with the first argument.

Parameters

input (Tensor) – the tensor to compare
other (Tensor or float) – the tensor or value to compare
out (Tensor, optional) – the output tensor that must be a ByteTensor

Returns

A torch.ByteTensor containing a 1 at each location where comparison is true

Return type

Tensor

Example:

>>> torch.ge(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]]))
tensor([[ 1,  1],
        [ 0,  1]], dtype=torch.uint8)

torch.gt(input, other, out=None) → Tensor¶

Computes $\text{input} > \text{other}$ element-wise.

The second argument can be a number or a tensor whose shape is broadcastable with the first argument.

Parameters

input (Tensor) – the tensor to compare
other (Tensor or float) – the tensor or value to compare
out (Tensor, optional) – the output tensor that must be a ByteTensor

Returns

A torch.ByteTensor containing a 1 at each location where comparison is true

Return type

Tensor

Example:

>>> torch.gt(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]]))
tensor([[ 0,  1],
        [ 0,  0]], dtype=torch.uint8)

torch.isfinite(tensor)¶

Returns a new tensor with boolean elements representing if each element is Finite or not.

Parameters: tensor (Tensor) – A tensor to check
Returns: A torch.ByteTensor containing a 1 at each location of finite elements and 0 otherwise
Return type: Tensor

Example:

>>> torch.isfinite(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
tensor([ 1,  0,  1,  0,  0], dtype=torch.uint8)

torch.isinf(tensor)¶

Returns a new tensor with boolean elements representing if each element is +/-INF or not.

Parameters: tensor (Tensor) – A tensor to check
Returns: A torch.ByteTensor containing a 1 at each location of +/-INF elements and 0 otherwise
Return type: Tensor

Example:

>>> torch.isinf(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
tensor([ 0,  1,  0,  1,  0], dtype=torch.uint8)

torch.isnan()¶

Returns a new tensor with boolean elements representing if each element is NaN or not.

Parameters: tensor (Tensor) – A tensor to check
Returns: A torch.ByteTensor containing a 1 at each location of NaN elements.
Return type: Tensor

Example:

>>> torch.isnan(torch.tensor([1, float('nan'), 2]))
tensor([ 0,  1,  0], dtype=torch.uint8)

torch.kthvalue(input, k, dim=None, keepdim=False, out=None) -> (Tensor, LongTensor)¶

Returns a namedtuple (values, indices) where values is the k th smallest element of each row of the input tensor in the given dimension dim. And indices is the index location of each element found.

If dim is not given, the last dimension of the input is chosen.

If keepdim is True, both the values and indices tensors are the same size as input, except in the dimension dim where they are of size 1. Otherwise, dim is squeezed (see torch.squeeze()), resulting in both the values and indices tensors having 1 fewer dimension than the input tensor.

Parameters

input (Tensor) – the input tensor
k (int) – k for the k-th smallest element
dim (int, optional) – the dimension to find the kth value along
keepdim (bool) – whether the output tensors have dim retained or not
out (tuple, optional) – the output tuple of (Tensor, LongTensor) can be optionally given to be used as output buffers

Example:

>>> x = torch.arange(1., 6.)
>>> x
tensor([ 1.,  2.,  3.,  4.,  5.])
>>> torch.kthvalue(x, 4)
torch.return_types.kthvalue(values=tensor(4.), indices=tensor(3))

>>> x=torch.arange(1.,7.).resize_(2,3)
>>> x
tensor([[ 1.,  2.,  3.],
        [ 4.,  5.,  6.]])
>>> torch.kthvalue(x, 2, 0, True)
torch.return_types.kthvalue(values=tensor([[4., 5., 6.]]), indices=tensor([[1, 1, 1]]))

torch.le(input, other, out=None) → Tensor¶

Computes $\text{input} \leq \text{other}$ element-wise.

The second argument can be a number or a tensor whose shape is broadcastable with the first argument.

Parameters

input (Tensor) – the tensor to compare
other (Tensor or float) – the tensor or value to compare
out (Tensor, optional) – the output tensor that must be a ByteTensor

Returns

A torch.ByteTensor containing a 1 at each location where comparison is true

Return type

Tensor

Example:

>>> torch.le(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]]))
tensor([[ 1,  0],
        [ 1,  1]], dtype=torch.uint8)

torch.lt(input, other, out=None) → Tensor¶

Computes $\text{input} < \text{other}$ element-wise.

The second argument can be a number or a tensor whose shape is broadcastable with the first argument.

Parameters

input (Tensor) – the tensor to compare
other (Tensor or float) – the tensor or value to compare
out (Tensor, optional) – the output tensor that must be a ByteTensor

Returns

A torch.ByteTensor containing a 1 at each location where comparison is true

Return type

Tensor

Example:

>>> torch.lt(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]]))
tensor([[ 0,  0],
        [ 1,  0]], dtype=torch.uint8)

torch.max()¶

torch.max(input) → Tensor

Returns the maximum value of all elements in the input tensor.

Parameters: input (Tensor) – the input tensor

Example:

>>> a = torch.randn(1, 3)
>>> a
tensor([[ 0.6763,  0.7445, -2.2369]])
>>> torch.max(a)
tensor(0.7445)

torch.max(input, dim, keepdim=False, out=None) -> (Tensor, LongTensor)

Returns a namedtuple (values, indices) where values is the maximum value of each row of the input tensor in the given dimension dim. And indices is the index location of each maximum value found (argmax).

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension to reduce
keepdim (bool, optional) – whether the output tensors have dim retained or not. Default: False.
out (tuple, optional) – the result tuple of two output tensors (max, max_indices)

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[-1.2360, -0.2942, -0.1222,  0.8475],
        [ 1.1949, -1.1127, -2.2379, -0.6702],
        [ 1.5717, -0.9207,  0.1297, -1.8768],
        [-0.6172,  1.0036, -0.6060, -0.2432]])
>>> torch.max(a, 1)
torch.return_types.max(values=tensor([0.8475, 1.1949, 1.5717, 1.0036]), indices=tensor([3, 0, 0, 1]))

torch.max(input, other, out=None) → Tensor

Each element of the tensor input is compared with the corresponding element of the tensor other and an element-wise maximum is taken.

The shapes of input and other don’t need to match, but they must be broadcastable.

\[\text{out}_i = \max(\text{tensor}_i, \text{other}_i) \]

Note

When the shapes do not match, the shape of the returned output tensor follows the broadcasting rules.

Parameters

input (Tensor) – the input tensor
other (Tensor) – the second input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 0.2942, -0.7416,  0.2653, -0.1584])
>>> b = torch.randn(4)
>>> b
tensor([ 0.8722, -1.7421, -0.4141, -0.5055])
>>> torch.max(a, b)
tensor([ 0.8722, -0.7416,  0.2653, -0.1584])

torch.min()¶

torch.min(input) → Tensor

Returns the minimum value of all elements in the input tensor.

Parameters: input (Tensor) – the input tensor

Example:

>>> a = torch.randn(1, 3)
>>> a
tensor([[ 0.6750,  1.0857,  1.7197]])
>>> torch.min(a)
tensor(0.6750)

torch.min(input, dim, keepdim=False, out=None) -> (Tensor, LongTensor)

Returns a namedtuple (values, indices) where values is the minimum value of each row of the input tensor in the given dimension dim. And indices is the index location of each minimum value found (argmin).

Parameters

input (Tensor) – the input tensor
dim (int) – the dimension to reduce
keepdim (bool) – whether the output tensors have dim retained or not
out (tuple, optional) – the tuple of two output tensors (min, min_indices)

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[-0.6248,  1.1334, -1.1899, -0.2803],
        [-1.4644, -0.2635, -0.3651,  0.6134],
        [ 0.2457,  0.0384,  1.0128,  0.7015],
        [-0.1153,  2.9849,  2.1458,  0.5788]])
>>> torch.min(a, 1)
torch.return_types.min(values=tensor([-1.1899, -1.4644,  0.0384, -0.1153]), indices=tensor([2, 0, 1, 0]))

torch.min(input, other, out=None) → Tensor

Each element of the tensor input is compared with the corresponding element of the tensor other and an element-wise minimum is taken. The resulting tensor is returned.

The shapes of input and other don’t need to match, but they must be broadcastable.

\[\text{out}_i = \min(\text{tensor}_i, \text{other}_i) \]

Note

When the shapes do not match, the shape of the returned output tensor follows the broadcasting rules.

Parameters

input (Tensor) – the input tensor
other (Tensor) – the second input tensor
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4)
>>> a
tensor([ 0.8137, -1.1740, -0.6460,  0.6308])
>>> b = torch.randn(4)
>>> b
tensor([-0.1369,  0.1555,  0.4019, -0.1929])
>>> torch.min(a, b)
tensor([-0.1369, -1.1740, -0.6460, -0.1929])

torch.ne(input, other, out=None) → Tensor¶

Computes $input \neq other$ element-wise.

The second argument can be a number or a tensor whose shape is broadcastable with the first argument.

Parameters

input (Tensor) – the tensor to compare
other (Tensor or float) – the tensor or value to compare
out (Tensor, optional) – the output tensor that must be a ByteTensor

Returns

A torch.ByteTensor containing a 1 at each location where comparison is true.

Return type

Tensor

Example:

>>> torch.ne(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]]))
tensor([[ 0,  1],
        [ 1,  0]], dtype=torch.uint8)

torch.sort(input, dim=-1, descending=False, out=None) -> (Tensor, LongTensor)¶

Sorts the elements of the input tensor along a given dimension in ascending order by value.

If dim is not given, the last dimension of the input is chosen.

If descending is True then the elements are sorted in descending order by value.

A tuple of (sorted_tensor, sorted_indices) is returned, where the sorted_indices are the indices of the elements in the original input tensor.

Parameters

input (Tensor) – the input tensor
dim (int, optional) – the dimension to sort along
descending (bool, optional) – controls the sorting order (ascending or descending)
out (tuple, optional) – the output tuple of (Tensor, LongTensor) that can be optionally given to be used as output buffers

Example:

>>> x = torch.randn(3, 4)
>>> sorted, indices = torch.sort(x)
>>> sorted
tensor([[-0.2162,  0.0608,  0.6719,  2.3332],
        [-0.5793,  0.0061,  0.6058,  0.9497],
        [-0.5071,  0.3343,  0.9553,  1.0960]])
>>> indices
tensor([[ 1,  0,  2,  3],
        [ 3,  1,  0,  2],
        [ 0,  3,  1,  2]])

>>> sorted, indices = torch.sort(x, 0)
>>> sorted
tensor([[-0.5071, -0.2162,  0.6719, -0.5793],
        [ 0.0608,  0.0061,  0.9497,  0.3343],
        [ 0.6058,  0.9553,  1.0960,  2.3332]])
>>> indices
tensor([[ 2,  0,  0,  1],
        [ 0,  1,  1,  2],
        [ 1,  2,  2,  0]])

torch.topk(input, k, dim=None, largest=True, sorted=True, out=None) -> (Tensor, LongTensor)¶

Returns the k largest elements of the given input tensor along a given dimension.

If dim is not given, the last dimension of the input is chosen.

If largest is False then the k smallest elements are returned.

A tuple of (values, indices) is returned, where the indices are the indices of the elements in the original input tensor.

The boolean option sorted if True, will make sure that the returned k elements are themselves sorted

Parameters

input (Tensor) – the input tensor
k (int) – the k in “top-k”
dim (int, optional) – the dimension to sort along
largest (bool, optional) – controls whether to return largest or smallest elements
sorted (bool, optional) – controls whether to return the elements in sorted order
out (tuple, optional) – the output tuple of (Tensor, LongTensor) that can be optionally given to be used as output buffers

Example:

>>> x = torch.arange(1., 6.)
>>> x
tensor([ 1.,  2.,  3.,  4.,  5.])
>>> torch.topk(x, 3)
(tensor([ 5.,  4.,  3.]), tensor([ 4,  3,  2]))

Spectral Ops¶

torch.fft(input, signal_ndim, normalized=False) → Tensor¶

Complex-to-complex Discrete Fourier Transform

This method computes the complex-to-complex discrete Fourier transform. Ignoring the batch dimensions, it computes the following expression:

\[X[\omega_1, \dots, \omega_d] = \sum_{n_1=0}^{N_1-1} \dots \sum_{n_d=0}^{N_d-1} x[n_1, \dots, n_d] e^{-j\ 2 \pi \sum_{i=0}^d \frac{\omega_i n_i}{N_i}}, \]

where $d$ = signal_ndim is number of dimensions for the signal, and $N_i$ is the size of signal dimension $i$.

This method supports 1D, 2D and 3D complex-to-complex transforms, indicated by signal_ndim. input must be a tensor with last dimension of size 2, representing the real and imaginary components of complex numbers, and should have at least signal_ndim + 1 dimensions with optionally arbitrary number of leading batch dimensions. If normalized is set to True, this normalizes the result by dividing it with $\sqrt{\prod_{i=1}^K N_i}$ so that the operator is unitary.

Returns the real and the imaginary parts together as one tensor of the same shape of input.

The inverse of this function is ifft().

Note

For CUDA tensors, an LRU cache is used for cuFFT plans to speed up repeatedly running FFT methods on tensors of same geometry with same same configuration.

Changing torch.backends.cuda.cufft_plan_cache.max_size (default is 4096 on CUDA 10 and newer, and 1023 on older CUDA versions) controls the capacity of this cache. Some cuFFT plans may allocate GPU memory. You can use torch.backends.cuda.cufft_plan_cache.size to query the number of plans currently in cache, and torch.backends.cuda.cufft_plan_cache.clear() to clear the cache.

Warning

For CPU tensors, this method is currently only available with MKL. Use torch.backends.mkl.is_available() to check if MKL is installed.

Parameters

input (Tensor) – the input tensor of at least signal_ndim + 1 dimensions
signal_ndim (int) – the number of dimensions in each signal. signal_ndim can only be 1, 2 or 3
normalized (bool, optional) – controls whether to return normalized results. Default: False

Returns

A tensor containing the complex-to-complex Fourier transform result

Return type

Tensor

Example:

>>> # unbatched 2D FFT
>>> x = torch.randn(4, 3, 2)
>>> torch.fft(x, 2)
tensor([[[-0.0876,  1.7835],
         [-2.0399, -2.9754],
         [ 4.4773, -5.0119]],

        [[-1.5716,  2.7631],
         [-3.8846,  5.2652],
         [ 0.2046, -0.7088]],

        [[ 1.9938, -0.5901],
         [ 6.5637,  6.4556],
         [ 2.9865,  4.9318]],

        [[ 7.0193,  1.1742],
         [-1.3717, -2.1084],
         [ 2.0289,  2.9357]]])
>>> # batched 1D FFT
>>> torch.fft(x, 1)
tensor([[[ 1.8385,  1.2827],
         [-0.1831,  1.6593],
         [ 2.4243,  0.5367]],

        [[-0.9176, -1.5543],
         [-3.9943, -2.9860],
         [ 1.2838, -2.9420]],

        [[-0.8854, -0.6860],
         [ 2.4450,  0.0808],
         [ 1.3076, -0.5768]],

        [[-0.1231,  2.7411],
         [-0.3075, -1.7295],
         [-0.5384, -2.0299]]])
>>> # arbitrary number of batch dimensions, 2D FFT
>>> x = torch.randn(3, 3, 5, 5, 2)
>>> y = torch.fft(x, 2)
>>> y.shape
torch.Size([3, 3, 5, 5, 2])

torch.ifft(input, signal_ndim, normalized=False) → Tensor¶

Complex-to-complex Inverse Discrete Fourier Transform

This method computes the complex-to-complex inverse discrete Fourier transform. Ignoring the batch dimensions, it computes the following expression:

\[X[\omega_1, \dots, \omega_d] = \frac{1}{\prod_{i=1}^d N_i} \sum_{n_1=0}^{N_1-1} \dots \sum_{n_d=0}^{N_d-1} x[n_1, \dots, n_d] e^{\ j\ 2 \pi \sum_{i=0}^d \frac{\omega_i n_i}{N_i}}, \]

where $d$ = signal_ndim is number of dimensions for the signal, and $N_i$ is the size of signal dimension $i$.

The argument specifications are almost identical with fft(). However, if normalized is set to True, this instead returns the results multiplied by $\sqrt{\prod_{i=1}^d N_i}$, to become a unitary operator. Therefore, to invert a fft(), the normalized argument should be set identically for fft().

Returns the real and the imaginary parts together as one tensor of the same shape of input.

The inverse of this function is fft().

Note

For CUDA tensors, an LRU cache is used for cuFFT plans to speed up repeatedly running FFT methods on tensors of same geometry with same same configuration.

Warning

For CPU tensors, this method is currently only available with MKL. Use torch.backends.mkl.is_available() to check if MKL is installed.

Parameters

input (Tensor) – the input tensor of at least signal_ndim + 1 dimensions
signal_ndim (int) – the number of dimensions in each signal. signal_ndim can only be 1, 2 or 3
normalized (bool, optional) – controls whether to return normalized results. Default: False

Returns

A tensor containing the complex-to-complex inverse Fourier transform result

Return type

Tensor

Example:

>>> x = torch.randn(3, 3, 2)
>>> x
tensor([[[ 1.2766,  1.3680],
         [-0.8337,  2.0251],
         [ 0.9465, -1.4390]],

        [[-0.1890,  1.6010],
         [ 1.1034, -1.9230],
         [-0.9482,  1.0775]],

        [[-0.7708, -0.8176],
         [-0.1843, -0.2287],
         [-1.9034, -0.2196]]])
>>> y = torch.fft(x, 2)
>>> torch.ifft(y, 2)  # recover x
tensor([[[ 1.2766,  1.3680],
         [-0.8337,  2.0251],
         [ 0.9465, -1.4390]],

        [[-0.1890,  1.6010],
         [ 1.1034, -1.9230],
         [-0.9482,  1.0775]],

        [[-0.7708, -0.8176],
         [-0.1843, -0.2287],
         [-1.9034, -0.2196]]])

torch.rfft(input, signal_ndim, normalized=False, onesided=True) → Tensor¶

Real-to-complex Discrete Fourier Transform

This method computes the real-to-complex discrete Fourier transform. It is mathematically equivalent with fft() with differences only in formats of the input and output.

This method supports 1D, 2D and 3D real-to-complex transforms, indicated by signal_ndim. input must be a tensor with at least signal_ndim dimensions with optionally arbitrary number of leading batch dimensions. If normalized is set to True, this normalizes the result by dividing it with $\sqrt{\prod_{i=1}^K N_i}$ so that the operator is unitary, where $N_i$ is the size of signal dimension $i$.

The real-to-complex Fourier transform results follow conjugate symmetry:

\[X[\omega_1, \dots, \omega_d] = X^*[N_1 - \omega_1, \dots, N_d - \omega_d], \]

where the index arithmetic is computed modulus the size of the corresponding dimension, $\ ^*$ is the conjugate operator, and $d$ = signal_ndim. onesided flag controls whether to avoid redundancy in the output results. If set to True (default), the output will not be full complex result of shape $(*, 2)$, where $*$ is the shape of input, but instead the last dimension will be halfed as of size $\lfloor \frac{N_d}{2} \rfloor + 1$.

The inverse of this function is irfft().

Note

For CUDA tensors, an LRU cache is used for cuFFT plans to speed up repeatedly running FFT methods on tensors of same geometry with same same configuration.

Warning

For CPU tensors, this method is currently only available with MKL. Use torch.backends.mkl.is_available() to check if MKL is installed.

Parameters

input (Tensor) – the input tensor of at least signal_ndim dimensions
signal_ndim (int) – the number of dimensions in each signal. signal_ndim can only be 1, 2 or 3
normalized (bool, optional) – controls whether to return normalized results. Default: False
onesided (bool, optional) – controls whether to return half of results to avoid redundancy. Default: True

Returns

A tensor containing the real-to-complex Fourier transform result

Return type

Tensor

Example:

>>> x = torch.randn(5, 5)
>>> torch.rfft(x, 2).shape
torch.Size([5, 3, 2])
>>> torch.rfft(x, 2, onesided=False).shape
torch.Size([5, 5, 2])

torch.irfft(input, signal_ndim, normalized=False, onesided=True, signal_sizes=None) → Tensor¶

Complex-to-real Inverse Discrete Fourier Transform

This method computes the complex-to-real inverse discrete Fourier transform. It is mathematically equivalent with ifft() with differences only in formats of the input and output.

The argument specifications are almost identical with ifft(). Similar to ifft(), if normalized is set to True, this normalizes the result by multiplying it with $\sqrt{\prod_{i=1}^K N_i}$ so that the operator is unitary, where $N_i$ is the size of signal dimension $i$.

Due to the conjugate symmetry, input do not need to contain the full complex frequency values. Roughly half of the values will be sufficient, as is the case when input is given by rfft() with rfft(signal, onesided=True). In such case, set the onesided argument of this method to True. Moreover, the original signal shape information can sometimes be lost, optionally set signal_sizes to be the size of the original signal (without the batch dimensions if in batched mode) to recover it with correct shape.

Therefore, to invert an rfft(), the normalized and onesided arguments should be set identically for irfft(), and preferrably a signal_sizes is given to avoid size mismatch. See the example below for a case of size mismatch.

See rfft() for details on conjugate symmetry.

The inverse of this function is rfft().

Warning

Generally speaking, the input of this function should contain values following conjugate symmetry. Note that even if onesided is True, often symmetry on some part is still needed. When this requirement is not satisfied, the behavior of irfft() is undefined. Since torch.autograd.gradcheck() estimates numerical Jacobian with point perturbations, irfft() will almost certainly fail the check.

Note

For CUDA tensors, an LRU cache is used for cuFFT plans to speed up repeatedly running FFT methods on tensors of same geometry with same same configuration.

Warning

For CPU tensors, this method is currently only available with MKL. Use torch.backends.mkl.is_available() to check if MKL is installed.

Parameters

input (Tensor) – the input tensor of at least signal_ndim + 1 dimensions
signal_ndim (int) – the number of dimensions in each signal. signal_ndim can only be 1, 2 or 3
normalized (bool, optional) – controls whether to return normalized results. Default: False
onesided (bool, optional) – controls whether input was halfed to avoid redundancy, e.g., by rfft(). Default: True
signal_sizes (list or torch.Size, optional) – the size of the original signal (without batch dimension). Default: None

Returns

A tensor containing the complex-to-real inverse Fourier transform result

Return type

Tensor

Example:

>>> x = torch.randn(4, 4)
>>> torch.rfft(x, 2, onesided=True).shape
torch.Size([4, 3, 2])
>>>
>>> # notice that with onesided=True, output size does not determine the original signal size
>>> x = torch.randn(4, 5)

>>> torch.rfft(x, 2, onesided=True).shape
torch.Size([4, 3, 2])
>>>
>>> # now we use the original shape to recover x
>>> x
tensor([[-0.8992,  0.6117, -1.6091, -0.4155, -0.8346],
        [-2.1596, -0.0853,  0.7232,  0.1941, -0.0789],
        [-2.0329,  1.1031,  0.6869, -0.5042,  0.9895],
        [-0.1884,  0.2858, -1.5831,  0.9917, -0.8356]])
>>> y = torch.rfft(x, 2, onesided=True)
>>> torch.irfft(y, 2, onesided=True, signal_sizes=x.shape)  # recover x
tensor([[-0.8992,  0.6117, -1.6091, -0.4155, -0.8346],
        [-2.1596, -0.0853,  0.7232,  0.1941, -0.0789],
        [-2.0329,  1.1031,  0.6869, -0.5042,  0.9895],
        [-0.1884,  0.2858, -1.5831,  0.9917, -0.8356]])

torch.stft(input, n_fft, hop_length=None, win_length=None, window=None, center=True, pad_mode='reflect', normalized=False, onesided=True)¶

Short-time Fourier transform (STFT).

Ignoring the optional batch dimension, this method computes the following expression:

\[X[m, \omega] = \sum_{k = 0}^{\text{win\_length-1}}% \text{window}[k]\ \text{input}[m \times \text{hop\_length} + k]\ % \exp\left(- j \frac{2 \pi \cdot \omega k}{\text{win\_length}}\right), \]

where $m$ is the index of the sliding window, and $\omega$ is the frequency that $0 \leq \omega < \text{n\_fft}$. When onesided is the default value True,

input must be either a 1-D time sequence or a 2-D batch of time sequences.
If hop_length is None (default), it is treated as equal to floor(n_fft / 4).
If win_length is None (default), it is treated as equal to n_fft.
window can be a 1-D tensor of size win_length, e.g., from torch.hann_window(). If window is None (default), it is treated as if having $1$ everywhere in the window. If $\text{win\_length} < \text{n\_fft}$, window will be padded on both sides to length n_fft before being applied.
If center is True (default), input will be padded on both sides so that the $t$-th frame is centered at time $t \times \text{hop\_length}$. Otherwise, the $t$-th frame begins at time $t \times \text{hop\_length}$.
pad_mode determines the padding method used on input when center is True. See torch.nn.functional.pad() for all available options. Default is "reflect".
If onesided is True (default), only values for $\omega$ in $\left[0, 1, 2, \dots, \left\lfloor \frac{\text{n\_fft}}{2} \right\rfloor + 1\right]$ are returned because the real-to-complex Fourier transform satisfies the conjugate symmetry, i.e., $X[m, \omega] = X[m, \text{n\_fft} - \omega]^*$.
If normalized is True (default is False), the function returns the normalized STFT results, i.e., multiplied by $(\text{frame\_length})^{-0.5}$.

Returns the real and the imaginary parts together as one tensor of size $(* \times N \times T \times 2)$, where $*$ is the optional batch size of input, $N$ is the number of frequencies where STFT is applied, $T$ is the total number of frames used, and each pair in the last dimension represents a complex number as the real part and the imaginary part.

Warning

This function changed signature at version 0.4.1. Calling with the previous signature may cause error or return incorrect result.

Parameters

input (Tensor) – the input tensor
n_fft (int) – size of Fourier transform
hop_length (int, optional) – the distance between neighboring sliding window frames. Default: None (treated as equal to floor(n_fft / 4))
win_length (int, optional) – the size of window frame and STFT filter. Default: None (treated as equal to n_fft)
window (Tensor, optional) – the optional window function. Default: None (treated as window of all $1$ s)
center (bool, optional) – whether to pad input on both sides so that the $t$-th frame is centered at time $t \times \text{hop\_length}$. Default: True
pad_mode (string, optional) – controls the padding method used when center is True. Default: "reflect"
normalized (bool, optional) – controls whether to return the normalized STFT results Default: False
onesided (bool, optional) – controls whether to return half of results to avoid redundancy Default: True

Returns

A tensor containing the STFT result with shape described above

Return type

Tensor

torch.bartlett_window(window_length, periodic=True, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Bartlett window function.

\[w[n] = 1 - \left| \frac{2n}{N-1} - 1 \right| = \begin{cases} \frac{2n}{N - 1} & \text{if } 0 \leq n \leq \frac{N - 1}{2} \\ 2 - \frac{2n}{N - 1} & \text{if } \frac{N - 1}{2} < n < N \\ \end{cases}, \]

where $N$ is the full window size.

The input window_length is a positive integer controlling the returned window size. periodic flag determines whether the returned window trims off the last duplicate value from the symmetric window and is ready to be used as a periodic window with functions like torch.stft(). Therefore, if periodic is true, the $N$ in above formula is in fact $\text{window\_length} + 1$. Also, we always have torch.bartlett_window(L, periodic=True) equal to torch.bartlett_window(L + 1, periodic=False)[:-1]).

Note

If window_length $=1$, the returned window contains a single value 1.

Parameters

window_length (int) – the size of returned window
periodic (bool, optional) – If True, returns a window to be used as periodic function. If False, return a symmetric window.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()). Only floating point types are supported.
layout (torch.layout, optional) – the desired layout of returned window tensor. Only torch.strided (dense layout) is supported.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Returns

A 1-D tensor of size $(\text{window\_length},)$ containing the window

Return type

Tensor

torch.blackman_window(window_length, periodic=True, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Blackman window function.

\[w[n] = 0.42 - 0.5 \cos \left( \frac{2 \pi n}{N - 1} \right) + 0.08 \cos \left( \frac{4 \pi n}{N - 1} \right) \]

where $N$ is the full window size.

The input window_length is a positive integer controlling the returned window size. periodic flag determines whether the returned window trims off the last duplicate value from the symmetric window and is ready to be used as a periodic window with functions like torch.stft(). Therefore, if periodic is true, the $N$ in above formula is in fact $\text{window\_length} + 1$. Also, we always have torch.blackman_window(L, periodic=True) equal to torch.blackman_window(L + 1, periodic=False)[:-1]).

Note

If window_length $=1$, the returned window contains a single value 1.

Parameters

window_length (int) – the size of returned window
periodic (bool, optional) – If True, returns a window to be used as periodic function. If False, return a symmetric window.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()). Only floating point types are supported.
layout (torch.layout, optional) – the desired layout of returned window tensor. Only torch.strided (dense layout) is supported.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Returns

A 1-D tensor of size $(\text{window\_length},)$ containing the window

Return type

Tensor

torch.hamming_window(window_length, periodic=True, alpha=0.54, beta=0.46, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Hamming window function.

\[w[n] = \alpha - \beta\ \cos \left( \frac{2 \pi n}{N - 1} \right), \]

where $N$ is the full window size.

The input window_length is a positive integer controlling the returned window size. periodic flag determines whether the returned window trims off the last duplicate value from the symmetric window and is ready to be used as a periodic window with functions like torch.stft(). Therefore, if periodic is true, the $N$ in above formula is in fact $\text{window\_length} + 1$. Also, we always have torch.hamming_window(L, periodic=True) equal to torch.hamming_window(L + 1, periodic=False)[:-1]).

Note

If window_length $=1$, the returned window contains a single value 1.

Note

This is a generalized version of torch.hann_window().

Parameters

window_length (int) – the size of returned window
periodic (bool, optional) – If True, returns a window to be used as periodic function. If False, return a symmetric window.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()). Only floating point types are supported.
layout (torch.layout, optional) – the desired layout of returned window tensor. Only torch.strided (dense layout) is supported.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Returns

A 1-D tensor of size $(\text{window\_length},)$ containing the window

Return type

Tensor

torch.hann_window(window_length, periodic=True, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶

Hann window function.

\[w[n] = \frac{1}{2}\ \left[1 - \cos \left( \frac{2 \pi n}{N - 1} \right)\right] = \sin^2 \left( \frac{\pi n}{N - 1} \right), \]

where $N$ is the full window size.

The input window_length is a positive integer controlling the returned window size. periodic flag determines whether the returned window trims off the last duplicate value from the symmetric window and is ready to be used as a periodic window with functions like torch.stft(). Therefore, if periodic is true, the $N$ in above formula is in fact $\text{window\_length} + 1$. Also, we always have torch.hann_window(L, periodic=True) equal to torch.hann_window(L + 1, periodic=False)[:-1]).

Note

If window_length $=1$, the returned window contains a single value 1.

Parameters

window_length (int) – the size of returned window
periodic (bool, optional) – If True, returns a window to be used as periodic function. If False, return a symmetric window.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_tensor_type()). Only floating point types are supported.
layout (torch.layout, optional) – the desired layout of returned window tensor. Only torch.strided (dense layout) is supported.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_tensor_type()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Returns

A 1-D tensor of size $(\text{window\_length},)$ containing the window

Return type

Tensor

Other Operations¶

torch.bincount(self, weights=None, minlength=0) → Tensor¶

Count the frequency of each value in an array of non-negative ints.

The number of bins (size 1) is one larger than the largest value in input unless input is empty, in which case the result is a tensor of size 0. If minlength is specified, the number of bins is at least minlength and if input is empty, then the result is tensor of size minlength filled with zeros. If n is the value at position i, out[n] += weights[i] if weights is specified else out[n] += 1.

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour that is not easily switched off. Please see the notes on /notes/randomness for background.

Parameters

input (Tensor) – 1-d int tensor
weights (Tensor) – optional, weight for each value in the input tensor. Should be of same size as input tensor.
minlength (int) – optional, minimum number of bins. Should be non-negative.

Returns

a tensor of shape Size([max(input) + 1]) if input is non-empty, else Size(0)

Return type

output (Tensor)

Example:

>>> input = torch.randint(0, 8, (5,), dtype=torch.int64)
>>> weights = torch.linspace(0, 1, steps=5)
>>> input, weights
(tensor([4, 3, 6, 3, 4]),
 tensor([ 0.0000,  0.2500,  0.5000,  0.7500,  1.0000])

>>> torch.bincount(input)
tensor([0, 0, 0, 2, 2, 0, 1])

>>> input.bincount(weights)
tensor([0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 0.0000, 0.5000])

torch.broadcast_tensors(*tensors) → List of Tensors¶

Broadcasts the given tensors according to broadcasting-semantics.

Parameters: *tensors – any number of tensors of the same type

Warning

More than one element of a broadcasted tensor may refer to a single memory location. As a result, in-place operations (especially ones that are vectorized) may result in incorrect behavior. If you need to write to the tensors, please clone them first.

Example:

>>> x = torch.arange(3).view(1, 3)
>>> y = torch.arange(2).view(2, 1)
>>> a, b = torch.broadcast_tensors(x, y)
>>> a.size()
torch.Size([2, 3])
>>> a
tensor([[0, 1, 2],
        [0, 1, 2]])

torch.cartesian_prod(*tensors)¶

Do cartesian product of the given sequence of tensors. The behavior is similar to python’s itertools.product.

Parameters

*tensors – any number of 1 dimensional tensors.

Returns

A tensor equivalent to converting all the input tensors into lists,: do itertools.product on these lists, and finally convert the resulting list into tensor.

Return type

Tensor

Example:

>>> a = [1, 2, 3]
>>> b = [4, 5]
>>> list(itertools.product(a, b))
[(1, 4), (1, 5), (2, 4), (2, 5), (3, 4), (3, 5)]
>>> tensor_a = torch.tensor(a)
>>> tensor_b = torch.tensor(b)
>>> torch.cartesian_prod(tensor_a, tensor_b)
tensor([[1, 4],
        [1, 5],
        [2, 4],
        [2, 5],
        [3, 4],
        [3, 5]])

torch.combinations(tensor, r=2, with_replacement=False) → seq¶

Compute combinations of length $r$ of the given tensor. The behavior is similar to python’s itertools.combinations when with_replacement is set to False, and itertools.combinations_with_replacement when with_replacement is set to True.

Parameters

tensor (Tensor) – 1D vector.
r (int, optional) – number of elements to combine
with_replacement (boolean, optional) – whether to allow duplication in combination

Returns

A tensor equivalent to converting all the input tensors into lists, do itertools.combinations or itertools.combinations_with_replacement on these lists, and finally convert the resulting list into tensor.

Return type

Tensor

Example:

>>> a = [1, 2, 3]
>>> list(itertools.combinations(a, r=2))
[(1, 2), (1, 3), (2, 3)]
>>> list(itertools.combinations(a, r=3))
[(1, 2, 3)]
>>> list(itertools.combinations_with_replacement(a, r=2))
[(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]
>>> tensor_a = torch.tensor(a)
>>> torch.combinations(tensor_a)
tensor([[1, 2],
        [1, 3],
        [2, 3]])
>>> torch.combinations(tensor_a, r=3)
tensor([[1, 2, 3]])
>>> torch.combinations(tensor_a, with_replacement=True)
tensor([[1, 1],
        [1, 2],
        [1, 3],
        [2, 2],
        [2, 3],
        [3, 3]])

torch.cross(input, other, dim=-1, out=None) → Tensor¶

Returns the cross product of vectors in dimension dim of input and other.

input and other must have the same size, and the size of their dim dimension should be 3.

If dim is not given, it defaults to the first dimension found with the size 3.

Parameters

input (Tensor) – the input tensor
other (Tensor) – the second input tensor
dim (int, optional) – the dimension to take the cross-product in.
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4, 3)
>>> a
tensor([[-0.3956,  1.1455,  1.6895],
        [-0.5849,  1.3672,  0.3599],
        [-1.1626,  0.7180, -0.0521],
        [-0.1339,  0.9902, -2.0225]])
>>> b = torch.randn(4, 3)
>>> b
tensor([[-0.0257, -1.4725, -1.2251],
        [-1.1479, -0.7005, -1.9757],
        [-1.3904,  0.3726, -1.1836],
        [-0.9688, -0.7153,  0.2159]])
>>> torch.cross(a, b, dim=1)
tensor([[ 1.0844, -0.5281,  0.6120],
        [-2.4490, -1.5687,  1.9792],
        [-0.8304, -1.3037,  0.5650],
        [-1.2329,  1.9883,  1.0551]])
>>> torch.cross(a, b)
tensor([[ 1.0844, -0.5281,  0.6120],
        [-2.4490, -1.5687,  1.9792],
        [-0.8304, -1.3037,  0.5650],
        [-1.2329,  1.9883,  1.0551]])

torch.diag(input, diagonal=0, out=None) → Tensor¶

If input is a vector (1-D tensor), then returns a 2-D square tensor with the elements of input as the diagonal.
If input is a matrix (2-D tensor), then returns a 1-D tensor with the diagonal elements of input.

The argument diagonal controls which diagonal to consider:

If diagonal = 0, it is the main diagonal.
If diagonal > 0, it is above the main diagonal.
If diagonal < 0, it is below the main diagonal.

Parameters

input (Tensor) – the input tensor
diagonal (int, optional) – the diagonal to consider
out (Tensor, optional) – the output tensor

BLAS and LAPACK Operations¶

torch.addbmm(beta=1, mat, alpha=1, batch1, batch2, out=None) → Tensor¶

Performs a batch matrix-matrix product of matrices stored in batch1 and batch2, with a reduced add step (all matrix multiplications get accumulated along the first dimension). mat is added to the final result.

batch1 and batch2 must be 3-D tensors each containing the same number of matrices.

If batch1 is a $(b \times n \times m)$ tensor, batch2 is a $(b \times m \times p)$ tensor, mat must be broadcastable with a $(n \times p)$ tensor and out will be a $(n \times p)$ tensor.

\[out = \beta\ \text{mat} + \alpha\ (\sum_{i=0}^{b-1} \text{batch1}_i \mathbin{@} \text{batch2}_i) \]

For inputs of type FloatTensor or DoubleTensor, arguments beta and alpha must be real numbers, otherwise they should be integers.

Parameters

beta (Number, optional) – multiplier for mat ($\beta$)
mat (Tensor) – matrix to be added
alpha (Number, optional) – multiplier for batch1 @ batch2 ($\alpha$)
batch1 (Tensor) – the first batch of matrices to be multiplied
batch2 (Tensor) – the second batch of matrices to be multiplied
out (Tensor, optional) – the output tensor

Example:

>>> M = torch.randn(3, 5)
>>> batch1 = torch.randn(10, 3, 4)
>>> batch2 = torch.randn(10, 4, 5)
>>> torch.addbmm(M, batch1, batch2)
tensor([[  6.6311,   0.0503,   6.9768, -12.0362,  -2.1653],
        [ -4.8185,  -1.4255,  -6.6760,   8.9453,   2.5743],
        [ -3.8202,   4.3691,   1.0943,  -1.1109,   5.4730]])

torch.addmm(beta=1, mat, alpha=1, mat1, mat2, out=None) → Tensor¶

Performs a matrix multiplication of the matrices mat1 and mat2. The matrix mat is added to the final result.

If mat1 is a $(n \times m)$ tensor, mat2 is a $(m \times p)$ tensor, then mat must be broadcastable with a $(n \times p)$ tensor and out will be a $(n \times p)$ tensor.

alpha and beta are scaling factors on matrix-vector product between mat1 and mat2 and the added matrix mat respectively.

\[\text{out} = \beta\ \text{mat} + \alpha\ (\text{mat1}_i \mathbin{@} \text{mat2}_i) \]

For inputs of type FloatTensor or DoubleTensor, arguments beta and alpha must be real numbers, otherwise they should be integers.

Parameters

beta (Number, optional) – multiplier for mat ($\beta$)
mat (Tensor) – matrix to be added
alpha (Number, optional) – multiplier for $mat1 @ mat2$ ($\alpha$)
mat1 (Tensor) – the first matrix to be multiplied
mat2 (Tensor) – the second matrix to be multiplied
out (Tensor, optional) – the output tensor

Example:

>>> M = torch.randn(2, 3)
>>> mat1 = torch.randn(2, 3)
>>> mat2 = torch.randn(3, 3)
>>> torch.addmm(M, mat1, mat2)
tensor([[-4.8716,  1.4671, -1.3746],
        [ 0.7573, -3.9555, -2.8681]])

torch.addmv(beta=1, tensor, alpha=1, mat, vec, out=None) → Tensor¶

Performs a matrix-vector product of the matrix mat and the vector vec. The vector tensor is added to the final result.

If mat is a $(n \times m)$ tensor, vec is a 1-D tensor of size m, then tensor must be broadcastable with a 1-D tensor of size n and out will be 1-D tensor of size n.

alpha and beta are scaling factors on matrix-vector product between mat and vec and the added tensor tensor respectively.

\[\text{out} = \beta\ \text{tensor} + \alpha\ (\text{mat} \mathbin{@} \text{vec}) \]

For inputs of type FloatTensor or DoubleTensor, arguments beta and alpha must be real numbers, otherwise they should be integers

Parameters

beta (Number, optional) – multiplier for tensor ($\beta$)
tensor (Tensor) – vector to be added
alpha (Number, optional) – multiplier for $mat @ vec$ ($\alpha$)
mat (Tensor) – matrix to be multiplied
vec (Tensor) – vector to be multiplied
out (Tensor, optional) – the output tensor

Example:

>>> M = torch.randn(2)
>>> mat = torch.randn(2, 3)
>>> vec = torch.randn(3)
>>> torch.addmv(M, mat, vec)
tensor([-0.3768, -5.5565])

torch.addr(beta=1, mat, alpha=1, vec1, vec2, out=None) → Tensor¶

Performs the outer-product of vectors vec1 and vec2 and adds it to the matrix mat.

Optional values beta and alpha are scaling factors on the outer product between vec1 and vec2 and the added matrix mat respectively.

\[\text{out} = \beta\ \text{mat} + \alpha\ (\text{vec1} \otimes \text{vec2}) \]

If vec1 is a vector of size n and vec2 is a vector of size m, then mat must be broadcastable with a matrix of size $(n \times m)$ and out will be a matrix of size $(n \times m)$.

For inputs of type FloatTensor or DoubleTensor, arguments beta and alpha must be real numbers, otherwise they should be integers

Parameters

beta (Number, optional) – multiplier for mat ($\beta$)
mat (Tensor) – matrix to be added
alpha (Number, optional) – multiplier for $\text{vec1} \otimes \text{vec2}$ ($\alpha$)
vec1 (Tensor) – the first vector of the outer product
vec2 (Tensor) – the second vector of the outer product
out (Tensor, optional) – the output tensor

Example:

>>> vec1 = torch.arange(1., 4.)
>>> vec2 = torch.arange(1., 3.)
>>> M = torch.zeros(3, 2)
>>> torch.addr(M, vec1, vec2)
tensor([[ 1.,  2.],
        [ 2.,  4.],
        [ 3.,  6.]])

torch.baddbmm(beta=1, mat, alpha=1, batch1, batch2, out=None) → Tensor¶

Performs a batch matrix-matrix product of matrices in batch1 and batch2. mat is added to the final result.

batch1 and batch2 must be 3-D tensors each containing the same number of matrices.

If batch1 is a $(b \times n \times m)$ tensor, batch2 is a $(b \times m \times p)$ tensor, then mat must be broadcastable with a $(b \times n \times p)$ tensor and out will be a $(b \times n \times p)$ tensor. Both alpha and beta mean the same as the scaling factors used in torch.addbmm().

\[\text{out}_i = \beta\ \text{mat}_i + \alpha\ (\text{batch1}_i \mathbin{@} \text{batch2}_i) \]

For inputs of type FloatTensor or DoubleTensor, arguments beta and alpha must be real numbers, otherwise they should be integers.

Parameters

beta (Number, optional) – multiplier for mat ($\beta$)
mat (Tensor) – the tensor to be added
alpha (Number, optional) – multiplier for $\text{batch1} \mathbin{@} \text{batch2}$ ($\alpha$)
batch1 (Tensor) – the first batch of matrices to be multiplied
batch2 (Tensor) – the second batch of matrices to be multiplied
out (Tensor, optional) – the output tensor

Example:

>>> M = torch.randn(10, 3, 5)
>>> batch1 = torch.randn(10, 3, 4)
>>> batch2 = torch.randn(10, 4, 5)
>>> torch.baddbmm(M, batch1, batch2).size()
torch.Size([10, 3, 5])

torch.bmm(batch1, batch2, out=None) → Tensor¶

Performs a batch matrix-matrix product of matrices stored in batch1 and batch2.

batch1 and batch2 must be 3-D tensors each containing the same number of matrices.

If batch1 is a $(b \times n \times m)$ tensor, batch2 is a $(b \times m \times p)$ tensor, out will be a $(b \times n \times p)$ tensor.

\[\text{out}_i = \text{batch1}_i \mathbin{@} \text{batch2}_i \]

Note

This function does not broadcast. For broadcasting matrix products, see torch.matmul().

Parameters

batch1 (Tensor) – the first batch of matrices to be multiplied
batch2 (Tensor) – the second batch of matrices to be multiplied
out (Tensor, optional) – the output tensor

Example:

>>> batch1 = torch.randn(10, 3, 4)
>>> batch2 = torch.randn(10, 4, 5)
>>> res = torch.bmm(batch1, batch2)
>>> res.size()
torch.Size([10, 3, 5])

torch.btrifact(A, pivot=True) -> (Tensor, IntTensor)¶

Batch LU factorization.

Returns a tuple containing the LU factorization and pivots. Pivoting is done if pivot is set.

Note

LU factorization with pivot = True is not available for CPU, and attempting to do so will throw an error. However, LU factorization with pivot = True is available for CUDA.

Parameters

A (Tensor) – the tensor to factor
pivot (bool, optional) – controls whether pivoting is done

Returns

A tuple containing factorization and pivots.

Example:

>>> A = torch.randn(2, 3, 3)
>>> A_LU, pivots = torch.btrifact(A)
>>> A_LU
tensor([[[ 1.3506,  2.5558, -0.0816],
         [ 0.1684,  1.1551,  0.1940],
         [ 0.1193,  0.6189, -0.5497]],

        [[ 0.4526,  1.2526, -0.3285],
         [-0.7988,  0.7175, -0.9701],
         [ 0.2634, -0.9255, -0.3459]]])

>>> pivots
tensor([[ 3,  3,  3],
        [ 3,  3,  3]], dtype=torch.int32)

torch.btrifact_with_info(A, pivot=True) -> (Tensor, IntTensor, IntTensor)¶

Batch LU factorization with additional error information.

This is a version of torch.btrifact() that always creates an info IntTensor, and returns it as the third return value.

Parameters

A (Tensor) – the tensor to factor
pivot (bool, optional) – controls whether pivoting is done

Returns

A tuple containing factorization, pivots, and an IntTensor where non-zero values indicate whether factorization for each minibatch sample succeeds.

Example:

>>> A = torch.randn(2, 3, 3)
>>> A_LU, pivots, info = A.btrifact_with_info()
>>> if info.nonzero().size(0) == 0:
>>>   print('LU factorization succeeded for all samples!')
LU factorization succeeded for all samples!

torch.btrisolve(b, LU_data, LU_pivots) → Tensor¶

Batch LU solve.

Returns the LU solve of the linear system $Ax = b$.

Parameters

b (Tensor) – the RHS tensor
LU_data (Tensor) – the pivoted LU factorization of A from btrifact().
LU_pivots (IntTensor) – the pivots of the LU factorization

Example:

>>> A = torch.randn(2, 3, 3)
>>> b = torch.randn(2, 3)
>>> A_LU = torch.btrifact(A)
>>> x = torch.btrisolve(b, *A_LU)
>>> torch.norm(torch.bmm(A, x.unsqueeze(2)) - b.unsqueeze(2))
tensor(1.00000e-07 *
       2.8312)

torch.btriunpack(LU_data, LU_pivots, unpack_data=True, unpack_pivots=True)¶

Unpacks the data and pivots from a batched LU factorization (btrifact) of a tensor.

Returns a tuple of tensors as (the pivots, the L tensor, the U tensor).

Parameters

LU_data (Tensor) – the packed LU factorization data
LU_pivots (Tensor) – the packed LU factorization pivots
unpack_data (bool) – flag indicating if the data should be unpacked
unpack_pivots (bool) – flag indicating if the pivots should be unpacked

Example:

>>> A = torch.randn(2, 3, 3)
>>> A_LU, pivots = A.btrifact()
>>> P, A_L, A_U = torch.btriunpack(A_LU, pivots)
>>>
>>> # can recover A from factorization
>>> A_ = torch.bmm(P, torch.bmm(A_L, A_U))

torch.chain_matmul(*matrices)¶

Returns the matrix product of the $N$ 2-D tensors. This product is efficiently computed using the matrix chain order algorithm which selects the order in which incurs the lowest cost in terms of arithmetic operations ([CLRS]). Note that since this is a function to compute the product, $N$ needs to be greater than or equal to 2; if equal to 2 then a trivial matrix-matrix product is returned. If $N$ is 1, then this is a no-op - the original matrix is returned as is.

Parameters: matrices (Tensors...) – a sequence of 2 or more 2-D tensors whose product is to be determined.
Returns: if the $i^{th}$ tensor was of dimensions $p_{i} \times p_{i + 1}$, then the product would be of dimensions $p_{1} \times p_{N + 1}$.
Return type: Tensor

Example:

>>> a = torch.randn(3, 4)
>>> b = torch.randn(4, 5)
>>> c = torch.randn(5, 6)
>>> d = torch.randn(6, 7)
>>> torch.chain_matmul(a, b, c, d)
tensor([[ -2.3375,  -3.9790,  -4.1119,  -6.6577,   9.5609, -11.5095,  -3.2614],
        [ 21.4038,   3.3378,  -8.4982,  -5.2457, -10.2561,  -2.4684,   2.7163],
        [ -0.9647,  -5.8917,  -2.3213,  -5.2284,  12.8615, -12.2816,  -2.5095]])

torch.cholesky(A, upper=False, out=None) → Tensor¶

Computes the Cholesky decomposition of a symmetric positive-definite matrix $A$ or for batches of symmetric positive-definite matrices.

If upper is True, the returned matrix U is upper-triangular, and the decomposition has the form:

\[A = U^TU\]

If upper is False, the returned matrix L is lower-triangular, and the decomposition has the form:

\[A = LL^T\]

If upper is True, and A is a batch of symmetric positive-definite matrices, then the returned tensor will be composed of upper-triangular Cholesky factors of each of the individual matrices. Similarly, when upper is False, the returned tensor will be composed of lower-triangular Cholesky factors of each of the individual matrices.

Parameters

a (Tensor) – the input tensor of size (*, n, n) where * is zero or more batch dimensions consisting of symmetric positive-definite matrices.
upper (bool, optional) – flag that indicates whether to return a upper or lower triangular matrix. Default: False
out (Tensor, optional) – the output matrix

Example:

>>> a = torch.randn(3, 3)
>>> a = torch.mm(a, a.t()) # make symmetric positive-definite
>>> l = torch.cholesky(a)
>>> a
tensor([[ 2.4112, -0.7486,  1.4551],
        [-0.7486,  1.3544,  0.1294],
        [ 1.4551,  0.1294,  1.6724]])
>>> l
tensor([[ 1.5528,  0.0000,  0.0000],
        [-0.4821,  1.0592,  0.0000],
        [ 0.9371,  0.5487,  0.7023]])
>>> torch.mm(l, l.t())
tensor([[ 2.4112, -0.7486,  1.4551],
        [-0.7486,  1.3544,  0.1294],
        [ 1.4551,  0.1294,  1.6724]])
>>> a = torch.randn(3, 2, 2)
>>> a = torch.matmul(a, a.transpose(-1, -2)) + 1e-03 # make symmetric positive-definite
>>> l = torch.cholesky(a)
>>> z = torch.matmul(l, l.transpose(-1, -2))
>>> torch.max(torch.abs(z - a)) # Max non-zero
tensor(2.3842e-07)

torch.cholesky_solve(b, u, upper=False, out=None) → Tensor¶

Solves a linear system of equations with a positive semidefinite matrix to be inverted given its Cholesky factor matrix u.

If upper is False, u is and lower triangular and c is returned such that:

\[c = (u u^T)^{-1} b \]

If upper is True or not provided, u is upper triangular and c is returned such that:

\[c = (u^T u)^{-1} b \]

torch.cholesky_solve(b, u) can take in 2D inputs b, u or inputs that are batches of 2D matrices. If the inputs are batches, then returns batched outputs c

Note

The out keyword only supports 2D matrix inputs, that is, b, u must be 2D matrices.

Parameters

b (Tensor) – input matrix of size $(*, m, k)$, where $*$ is zero or more batch dimensions
u (Tensor) – input matrix of size $(*, m, m)$, where $*$ is zero of more batch dimensions composed of upper or lower triangular Cholesky factor
upper (bool, optional) – whether to consider the Cholesky factor as a lower or upper triangular matrix. Default: False.
out (Tensor, optional) – the output tensor for c

Example:

>>> a = torch.randn(3, 3)
>>> a = torch.mm(a, a.t()) # make symmetric positive definite
>>> u = torch.cholesky(a)
>>> a
tensor([[ 0.7747, -1.9549,  1.3086],
        [-1.9549,  6.7546, -5.4114],
        [ 1.3086, -5.4114,  4.8733]])
>>> b = torch.randn(3, 2)
>>> b
tensor([[-0.6355,  0.9891],
        [ 0.1974,  1.4706],
        [-0.4115, -0.6225]])
>>> torch.cholesky_solve(b, u)
tensor([[ -8.1625,  19.6097],
        [ -5.8398,  14.2387],
        [ -4.3771,  10.4173]])
>>> torch.mm(a.inverse(), b)
tensor([[ -8.1626,  19.6097],
        [ -5.8398,  14.2387],
        [ -4.3771,  10.4173]])

torch.dot(tensor1, tensor2) → Tensor¶

Computes the dot product (inner product) of two tensors.

Note

This function does not broadcast.

Example:

>>> torch.dot(torch.tensor([2, 3]), torch.tensor([2, 1]))
tensor(7)

torch.eig(a, eigenvectors=False, out=None) -> (Tensor, Tensor)¶

Computes the eigenvalues and eigenvectors of a real square matrix.

Note

Since eigenvalues and eigenvectors might be complex, backward pass is supported only

for torch.symeig()

Parameters

a (Tensor) – the square matrix of shape $(n \times n)$ for which the eigenvalues and eigenvectors will be computed
eigenvectors (bool) – True to compute both eigenvalues and eigenvectors; otherwise, only eigenvalues will be computed
out (tuple, optional) – the output tensors

Returns

A namedtuple (eigenvalues, eigenvectors) containing

eigenvalues (Tensor): Shape $(n \times 2)$. Each row is an eigenvalue of a, where the first element is the real part and the second element is the imaginary part. The eigenvalues are not necessarily ordered.

eigenvectors (Tensor): If eigenvectors=False, it’s an empty tensor. Otherwise, this tensor of shape $(n \times n)$ can be used to compute normalized (unit length) eigenvectors of corresponding eigenvalues as follows. If the corresponding eigenvalues[j] is a real number, column eigenvectors[:, j] is the eigenvector corresponding to eigenvalues[j]. If the corresponding eigenvalues[j] and eigenvalues[j + 1] form a complex conjugate pair, then the true eigenvectors can be computed as $\text{true eigenvector}[j] = eigenvectors[:, j] + i \times eigenvectors[:, j + 1]$, $\text{true eigenvector}[j + 1] = eigenvectors[:, j] - i \times eigenvectors[:, j + 1]$.

Return type

(Tensor, Tensor)

torch.gels(B, A, out=None) → Tensor¶

Computes the solution to the least squares and least norm problems for a full rank matrix $A$ of size $(m \times n)$ and a matrix $B$ of size $(m \times k)$.

If $m \geq n$, gels() solves the least-squares problem:

\[\begin{array}{ll} \min_X & \|AX-B\|_2. \end{array}\]

If $m < n$, gels() solves the least-norm problem:

\[\begin{array}{ll} \min_X & \|X\|_2 & \text{subject to} & AX = B. \end{array}\]

Returned tensor $X$ has shape $(\max(m, n) \times k)$. The first $n$ rows of $X$ contains the solution. If $m \geq n$, the residual sum of squares for the solution in each column is given by the sum of squares of elements in the remaining $m - n$ rows of that column.

Parameters

B (Tensor) – the matrix $B$
A (Tensor) – the $m$ by $n$ matrix $A$
out (tuple, optional) – the optional destination tensor

Returns

A tuple containing:

X (Tensor): the least squares solution

qr (Tensor): the details of the QR factorization

Return type

(Tensor, Tensor)

Note

The returned matrices will always be transposed, irrespective of the strides of the input matrices. That is, they will have stride (1, m) instead of (m, 1).

Example:

>>> A = torch.tensor([[1., 1, 1],
                      [2, 3, 4],
                      [3, 5, 2],
                      [4, 2, 5],
                      [5, 4, 3]])
>>> B = torch.tensor([[-10., -3],
                      [ 12, 14],
                      [ 14, 12],
                      [ 16, 16],
                      [ 18, 16]])
>>> X, _ = torch.gels(B, A)
>>> X
tensor([[  2.0000,   1.0000],
        [  1.0000,   1.0000],
        [  1.0000,   2.0000],
        [ 10.9635,   4.8501],
        [  8.9332,   5.2418]])

torch.geqrf(input, out=None) -> (Tensor, Tensor)¶

This is a low-level function for calling LAPACK directly. This function returns a namedtuple (a, tau) as defined in LAPACK documentation for geqrf .

You’ll generally want to use torch.qr() instead.

Computes a QR decomposition of input, but without constructing $Q$ and $R$ as explicit separate matrices.

Rather, this directly calls the underlying LAPACK function ?geqrf which produces a sequence of ‘elementary reflectors’.

See LAPACK documentation for geqrf for further details.

Parameters

input (Tensor) – the input matrix
out (tuple, optional) – the output tuple of (Tensor, Tensor)

torch.ger(vec1, vec2, out=None) → Tensor¶

Outer product of vec1 and vec2. If vec1 is a vector of size $n$ and vec2 is a vector of size $m$, then out must be a matrix of size $(n \times m)$.

Note

This function does not broadcast.

Parameters

vec1 (Tensor) – 1-D input vector
vec2 (Tensor) – 1-D input vector
out (Tensor, optional) – optional output matrix

Example:

>>> v1 = torch.arange(1., 5.)
>>> v2 = torch.arange(1., 4.)
>>> torch.ger(v1, v2)
tensor([[  1.,   2.,   3.],
        [  2.,   4.,   6.],
        [  3.,   6.,   9.],
        [  4.,   8.,  12.]])

torch.gesv(b, A, out=None)¶

This function returns the solution to the system of linear equations represented by $AX = B$ and the LU factorization of A, in order as a tuple X, LU.

For more information regarding torch.gesv(), please check torch.solve().

Warning

torch.gesv() is deprecated in favour of torch.solve() and will be removed in the next release. Please use torch.solve() instead.

torch.inverse(input, out=None) → Tensor¶

Takes the inverse of the square matrix input. input can be batches of 2D square tensors, in which case this function would return a tensor composed of individual inverses.

Note

Irrespective of the original strides, the returned tensors will be transposed, i.e. with strides like input.contiguous().transpose(-2, -1).strides()

Parameters

input (Tensor) – the input tensor of size (*, n, n) where * is zero or more batch dimensions
out (Tensor, optional) – the optional output tensor

Example:

>>> x = torch.rand(4, 4)
>>> y = torch.inverse(x)
>>> z = torch.mm(x, y)
>>> z
tensor([[ 1.0000, -0.0000, -0.0000,  0.0000],
        [ 0.0000,  1.0000,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  1.0000,  0.0000],
        [ 0.0000, -0.0000, -0.0000,  1.0000]])
>>> torch.max(torch.abs(z - torch.eye(4))) # Max non-zero
tensor(1.1921e-07)
>>> # Batched inverse example
>>> x = torch.randn(2, 3, 4, 4)
>>> y = torch.inverse(x)
>>> z = torch.matmul(x, y)
>>> torch.max(torch.abs(z - torch.eye(4).expand_as(x))) # Max non-zero
tensor(1.9073e-06)

torch.det(A) → Tensor¶

Calculates determinant of a 2D square tensor.

Note

Backward through det() internally uses SVD results when A is not invertible. In this case, double backward through det() will be unstable in when A doesn’t have distinct singular values. See svd() for details.

Parameters: A (Tensor) – The input 2D square tensor

Example:

>>> A = torch.randn(3, 3)
>>> torch.det(A)
tensor(3.7641)

torch.logdet(A) → Tensor¶

Calculates log determinant of a 2D square tensor.

Note

Result is -inf if A has zero log determinant, and is nan if A has negative determinant.

Note

Backward through logdet() internally uses SVD results when A is not invertible. In this case, double backward through logdet() will be unstable in when A doesn’t have distinct singular values. See svd() for details.

Parameters: A (Tensor) – The input 2D square tensor

Example:

>>> A = torch.randn(3, 3)
>>> torch.det(A)
tensor(0.2611)
>>> torch.logdet(A)
tensor(-1.3430)

torch.slogdet(A) -> (Tensor, Tensor)¶

Calculates the sign and log value of a 2D square tensor’s determinant.

Note

If A has zero determinant, this returns (0, -inf).

Note

Backward through slogdet() internally uses SVD results when A is not invertible. In this case, double backward through slogdet() will be unstable in when A doesn’t have distinct singular values. See svd() for details.

Parameters: A (Tensor) – The input 2D square tensor
Returns: A tuple containing the sign of the determinant, and the log value of the absolute determinant.

Example:

>>> A = torch.randn(3, 3)
>>> torch.det(A)
tensor(-4.8215)
>>> torch.logdet(A)
tensor(nan)
>>> torch.slogdet(A)
(tensor(-1.), tensor(1.5731))

torch.matmul(tensor1, tensor2, out=None) → Tensor¶

Matrix product of two tensors.

The behavior depends on the dimensionality of the tensors as follows:

If both tensors are 1-dimensional, the dot product (scalar) is returned.
If both arguments are 2-dimensional, the matrix-matrix product is returned.
If the first argument is 1-dimensional and the second argument is 2-dimensional, a 1 is prepended to its dimension for the purpose of the matrix multiply. After the matrix multiply, the prepended dimension is removed.
If the first argument is 2-dimensional and the second argument is 1-dimensional, the matrix-vector product is returned.
If both arguments are at least 1-dimensional and at least one argument is N-dimensional (where N > 2), then a batched matrix multiply is returned. If the first argument is 1-dimensional, a 1 is prepended to its dimension for the purpose of the batched matrix multiply and removed after. If the second argument is 1-dimensional, a 1 is appended to its dimension for the purpose of the batched matrix multiple and removed after. The non-matrix (i.e. batch) dimensions are broadcasted (and thus must be broadcastable). For example, if tensor1 is a $(j \times 1 \times n \times m)$ tensor and tensor2 is a $(k \times m \times p)$ tensor, out will be an $(j \times k \times n \times p)$ tensor.

Note

The 1-dimensional dot product version of this function does not support an out parameter.

Parameters

tensor1 (Tensor) – the first tensor to be multiplied
tensor2 (Tensor) – the second tensor to be multiplied
out (Tensor, optional) – the output tensor

Example:

>>> # vector x vector
>>> tensor1 = torch.randn(3)
>>> tensor2 = torch.randn(3)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([])
>>> # matrix x vector
>>> tensor1 = torch.randn(3, 4)
>>> tensor2 = torch.randn(4)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([3])
>>> # batched matrix x broadcasted vector
>>> tensor1 = torch.randn(10, 3, 4)
>>> tensor2 = torch.randn(4)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([10, 3])
>>> # batched matrix x batched matrix
>>> tensor1 = torch.randn(10, 3, 4)
>>> tensor2 = torch.randn(10, 4, 5)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([10, 3, 5])
>>> # batched matrix x broadcasted matrix
>>> tensor1 = torch.randn(10, 3, 4)
>>> tensor2 = torch.randn(4, 5)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([10, 3, 5])

torch.matrix_power(input, n) → Tensor¶

Returns the matrix raised to the power n for square matrices. For batch of matrices, each individual matrix is raised to the power n.

If n is negative, then the inverse of the matrix (if invertible) is raised to the power n. For a batch of matrices, the batched inverse (if invertible) is raised to the power n. If n is 0, then an identity matrix is returned.

Parameters

input (Tensor) – the input tensor
n (int) – the power to raise the matrix to

Example:

>>> a = torch.randn(2, 2, 2)
>>> a
tensor([[[-1.9975, -1.9610],
         [ 0.9592, -2.3364]],

        [[-1.2534, -1.3429],
         [ 0.4153, -1.4664]]])
>>> torch.matrix_power(a, 3)
tensor([[[  3.9392, -23.9916],
         [ 11.7357,  -0.2070]],

        [[  0.2468,  -6.7168],
         [  2.0774,  -0.8187]]])

torch.matrix_rank(input, tol=None, bool symmetric=False) → Tensor¶

Returns the numerical rank of a 2-D tensor. The method to compute the matrix rank is done using SVD by default. If symmetric is True, then input is assumed to be symmetric, and the computation of the rank is done by obtaining the eigenvalues.

tol is the threshold below which the singular values (or the eigenvalues when symmetric is True) are considered to be 0. If tol is not specified, tol is set to S.max() * max(S.size()) * eps where S is the singular values (or the eigenvalues when symmetric is True), and eps is the epsilon value for the datatype of input.

Parameters

input (Tensor) – the input 2-D tensor
tol (float, optional) – the tolerance value. Default: None
symmetric (bool, optional) – indicates whether input is symmetric. Default: False

Example:

>>> a = torch.eye(10)
>>> torch.matrix_rank(a)
tensor(10)
>>> b = torch.eye(10)
>>> b[0, 0] = 0
>>> torch.matrix_rank(b)
tensor(9)

torch.mm(mat1, mat2, out=None) → Tensor¶

Performs a matrix multiplication of the matrices mat1 and mat2.

If mat1 is a $(n \times m)$ tensor, mat2 is a $(m \times p)$ tensor, out will be a $(n \times p)$ tensor.

Note

This function does not broadcast. For broadcasting matrix products, see torch.matmul().

Parameters

mat1 (Tensor) – the first matrix to be multiplied
mat2 (Tensor) – the second matrix to be multiplied
out (Tensor, optional) – the output tensor

Example:

>>> mat1 = torch.randn(2, 3)
>>> mat2 = torch.randn(3, 3)
>>> torch.mm(mat1, mat2)
tensor([[ 0.4851,  0.5037, -0.3633],
        [-0.0760, -3.6705,  2.4784]])

torch.mv(mat, vec, out=None) → Tensor¶

Performs a matrix-vector product of the matrix mat and the vector vec.

If mat is a $(n \times m)$ tensor, vec is a 1-D tensor of size $m$, out will be 1-D of size $n$.

Note

This function does not broadcast.

Parameters

mat (Tensor) – matrix to be multiplied
vec (Tensor) – vector to be multiplied
out (Tensor, optional) – the output tensor

Example:

>>> mat = torch.randn(2, 3)
>>> vec = torch.randn(3)
>>> torch.mv(mat, vec)
tensor([ 1.0404, -0.6361])

torch.orgqr(a, tau) → Tensor¶

Computes the orthogonal matrix Q of a QR factorization, from the (a, tau) tuple returned by torch.geqrf().

This directly calls the underlying LAPACK function ?orgqr. See LAPACK documentation for orgqr for further details.

Parameters

a (Tensor) – the a from torch.geqrf().
tau (Tensor) – the tau from torch.geqrf().

torch.ormqr(a, tau, mat, left=True, transpose=False) → Tensor¶

Multiplies mat by the orthogonal Q matrix of the QR factorization formed by torch.geqrf() that is represented by (a, tau).

This directly calls the underlying LAPACK function ?ormqr. See LAPACK documentation for ormqr for further details.

Parameters

a (Tensor) – the a from torch.geqrf().
tau (Tensor) – the tau from torch.geqrf().
mat (Tensor) – the matrix to be multiplied.

torch.pinverse(input, rcond=1e-15) → Tensor¶

Calculates the pseudo-inverse (also known as the Moore-Penrose inverse) of a 2D tensor. Please look at Moore-Penrose inverse for more details

Note

This method is implemented using the Singular Value Decomposition.

Note

The pseudo-inverse is not necessarily a continuous function in the elements of the matrix [1]. Therefore, derivatives are not always existent, and exist for a constant rank only [2]. However, this method is backprop-able due to the implementation by using SVD results, and could be unstable. Double-backward will also be unstable due to the usage of SVD internally. See svd() for more details.

Parameters

input (Tensor) – The input 2D tensor of dimensions $m \times n$
rcond (float) – A floating point value to determine the cutoff for small singular values. Default: 1e-15

Returns

The pseudo-inverse of input of dimensions $n \times m$

Example:

>>> input = torch.randn(3, 5)
>>> input
tensor([[ 0.5495,  0.0979, -1.4092, -0.1128,  0.4132],
        [-1.1143, -0.3662,  0.3042,  1.6374, -0.9294],
        [-0.3269, -0.5745, -0.0382, -0.5922, -0.6759]])
>>> torch.pinverse(input)
tensor([[ 0.0600, -0.1933, -0.2090],
        [-0.0903, -0.0817, -0.4752],
        [-0.7124, -0.1631, -0.2272],
        [ 0.1356,  0.3933, -0.5023],
        [-0.0308, -0.1725, -0.5216]])

torch.potrf(a, upper=True, out=None)¶

Computes the Cholesky decomposition of a symmetric positive-definite matrix $A$.

For more information regarding torch.potrf(), please check torch.cholesky().

Warning

torch.potrf() is deprecated in favour of torch.cholesky() and will be removed in the next release. Please use torch.cholesky() instead and note that the upper argument in torch.cholesky() defaults to False.

torch.potri(u, upper=True, out=None) → Tensor¶

Computes the inverse of a positive semidefinite matrix given its Cholesky factor u: returns matrix inv

If upper is True or not provided, u is upper triangular such that the returned tensor is

\[inv = (u^T u)^{-1} \]

If upper is False, u is lower triangular such that the returned tensor is

\[inv = (uu^{T})^{-1} \]

Parameters

u (Tensor) – the input 2-D tensor, a upper or lower triangular Cholesky factor
upper (bool, optional) – whether to return a upper (default) or lower triangular matrix
out (Tensor, optional) – the output tensor for inv

Example:

>>> a = torch.randn(3, 3)
>>> a = torch.mm(a, a.t()) # make symmetric positive definite
>>> u = torch.cholesky(a)
>>> a
tensor([[  0.9935,  -0.6353,   1.5806],
        [ -0.6353,   0.8769,  -1.7183],
        [  1.5806,  -1.7183,  10.6618]])
>>> torch.potri(u)
tensor([[ 1.9314,  1.2251, -0.0889],
        [ 1.2251,  2.4439,  0.2122],
        [-0.0889,  0.2122,  0.1412]])
>>> a.inverse()
tensor([[ 1.9314,  1.2251, -0.0889],
        [ 1.2251,  2.4439,  0.2122],
        [-0.0889,  0.2122,  0.1412]])

torch.potrs(b, u, upper=True, out=None)¶

Solves a linear system of equations with a positive semidefinite matrix to be inverted given its Cholesky factor matrix u.

For more information regarding torch.potrs(), please check torch.cholesky_solve().

Warning

torch.potrs() is deprecated in favour of torch.cholesky_solve() and will be removed in the next release. Please use torch.cholesky_solve() instead and note that the upper argument in torch.cholesky_solve() defaults to False.

torch.pstrf(a, upper=True, out=None)¶

Computes the pivoted Cholesky decomposition of a symmetric positive-definite matrix a. returns a namedtuple (u, pivot) of matrice.

If upper is True or not provided, u is upper triangular such that $a = p^T u^T u p$, with p the permutation given by pivot.

If upper is False, u is lower triangular such that $a = p^T u u^T p$.

Warning

torch.pstrf() is deprecated in favour of torch.cholesky() and will be removed in the next release.

Parameters

a (Tensor) – the input 2-D tensor
upper (bool, optional) – whether to return a upper (default) or lower triangular matrix
out (tuple, optional) – namedtuple of u and pivot tensors

Example:

>>> a = torch.randn(3, 3)
>>> a = torch.mm(a, a.t()) # make symmetric positive definite
>>> a
tensor([[ 3.5405, -0.4577,  0.8342],
        [-0.4577,  1.8244, -0.1996],
        [ 0.8342, -0.1996,  3.7493]])
>>> u,piv = torch.pstrf(a)
>>> u
tensor([[ 1.9363,  0.4308, -0.1031],
        [ 0.0000,  1.8316, -0.2256],
        [ 0.0000,  0.0000,  1.3277]])
>>> piv
tensor([ 2,  0,  1], dtype=torch.int32)
>>> p = torch.eye(3).index_select(0,piv.long()).index_select(0,piv.long()).t() # make pivot permutation
>>> torch.mm(torch.mm(p.t(),torch.mm(u.t(),u)),p) # reconstruct
tensor([[ 3.5405, -0.4577,  0.8342],
        [-0.4577,  1.8244, -0.1996],
        [ 0.8342, -0.1996,  3.7493]])

torch.qr(input, out=None) -> (Tensor, Tensor)¶

Computes the QR decomposition of a matrix input, and returns a namedtuple (Q, R) of matrices such that $\text{input} = Q R$, with $Q$ being an orthogonal matrix and $R$ being an upper triangular matrix.

This returns the thin (reduced) QR factorization.

Note

precision may be lost if the magnitudes of the elements of input are large

Note

While it should always give you a valid decomposition, it may not give you the same one across platforms - it will depend on your LAPACK implementation.

Note

Irrespective of the original strides, the returned matrix $Q$ will be transposed, i.e. with strides (1, m) instead of (m, 1).

Parameters

input (Tensor) – the input 2-D tensor
out (tuple, optional) – tuple of Q and R tensors

Example:

>>> a = torch.tensor([[12., -51, 4], [6, 167, -68], [-4, 24, -41]])
>>> q, r = torch.qr(a)
>>> q
tensor([[-0.8571,  0.3943,  0.3314],
        [-0.4286, -0.9029, -0.0343],
        [ 0.2857, -0.1714,  0.9429]])
>>> r
tensor([[ -14.0000,  -21.0000,   14.0000],
        [   0.0000, -175.0000,   70.0000],
        [   0.0000,    0.0000,  -35.0000]])
>>> torch.mm(q, r).round()
tensor([[  12.,  -51.,    4.],
        [   6.,  167.,  -68.],
        [  -4.,   24.,  -41.]])
>>> torch.mm(q.t(), q).round()
tensor([[ 1.,  0.,  0.],
        [ 0.,  1., -0.],
        [ 0., -0.,  1.]])

torch.solve(B, A, out=None) -> (Tensor, Tensor)¶

This function returns the solution to the system of linear equations represented by $AX = B$ and the LU factorization of A, in order as a tuple X, LU.

LU contains L and U factors for LU factorization of A.

torch.solve(B, A) can take in 2D inputs B, A or inputs that are batches of 2D matrices. If the inputs are batches, then returns batched outputs X, LU.

Note

Irrespective of the original strides, the returned matrices X and LU will be transposed, i.e. with strides like B.contiguous().transpose(-1, -2).strides() and A.contiguous().transpose(-1, -2).strides() respectively.

Parameters

B (Tensor) – input matrix of size $(*, m, k)$ , where $*$ is zero or more batch dimensions.
A (Tensor) – input square matrix of size $(*, m, m)$, where $*$ is zero or more batch dimensions.
out ((Tensor, Tensor), optional) – optional output tuple.

Example:

>>> A = torch.tensor([[6.80, -2.11,  5.66,  5.97,  8.23],
                      [-6.05, -3.30,  5.36, -4.44,  1.08],
                      [-0.45,  2.58, -2.70,  0.27,  9.04],
                      [8.32,  2.71,  4.35,  -7.17,  2.14],
                      [-9.67, -5.14, -7.26,  6.08, -6.87]]).t()
>>> B = torch.tensor([[4.02,  6.19, -8.22, -7.57, -3.03],
                      [-1.56,  4.00, -8.67,  1.75,  2.86],
                      [9.81, -4.09, -4.57, -8.61,  8.99]]).t()
>>> X, LU = torch.solve(B, A)
>>> torch.dist(B, torch.mm(A, X))
tensor(1.00000e-06 *
       7.0977)

>>> # Batched solver example
>>> A = torch.randn(2, 3, 1, 4, 4)
>>> B = torch.randn(2, 3, 1, 4, 6)
>>> X, LU = torch.solve(B, A)
>>> torch.dist(B, A.matmul(X))
tensor(1.00000e-06 *
   3.6386)

torch.svd(input, some=True, compute_uv=True, out=None) -> (Tensor, Tensor, Tensor)¶

svd(A) returns a namedtuple (U, S, V) which the singular value decomposition of a input real matrix A of size (n x m) such that $A = USV^T$.

U is of shape $(n \times n)$.

S is a diagonal matrix of shape $(n \times m)$, represented as a vector of size $\min(n, m)$ containing the non-negative diagonal entries.

V is of shape $(m \times m)$.

If some is True (default), the returned U and V matrices will contain only $min(n, m)$ orthonormal columns.

If compute_uv is False, the returned U and V matrices will be zero matrices of shape $(n \times n)$ and $(m \times m)$ respectively. some will be ignored here.

Note

The implementation of SVD on CPU uses the LAPACK routine ?gesdd (a divide-and-conquer algorithm) instead of ?gesvd for speed. Analogously, the SVD on GPU uses the MAGMA routine gesdd as well.

Note

Irrespective of the original strides, the returned matrix U will be transposed, i.e. with strides (1, n) instead of (n, 1).

Note

Extra care needs to be taken when backward through U and V outputs. Such operation is really only stable when input is full rank with all distinct singular values. Otherwise, NaN can appear as the gradients are not properly defined. Also, notice that double backward will usually do an additional backward through U and V even if the original backward is only on S.

Note

When some = False, the gradients on U[:, min(n, m):] and V[:, min(n, m):] will be ignored in backward as those vectors can be arbitrary bases of the subspaces.

Note

When compute_uv = False, backward cannot be performed since U and V from the forward pass is required for the backward operation.

Parameters

input (Tensor) – the input 2-D tensor
some (bool, optional) – controls the shape of returned U and V
out (tuple, optional) – the output tuple of tensors

Example:

>>> a = torch.tensor([[8.79,  6.11, -9.15,  9.57, -3.49,  9.84],
                      [9.93,  6.91, -7.93,  1.64,  4.02,  0.15],
                      [9.83,  5.04,  4.86,  8.83,  9.80, -8.99],
                      [5.45, -0.27,  4.85,  0.74, 10.00, -6.02],
                      [3.16,  7.98,  3.01,  5.80,  4.27, -5.31]]).t()

>>> torch.svd(a).__class__
<class 'torch.return_types.svd'>
>>> u, s, v = torch.svd(a)
>>> u
tensor([[-0.5911,  0.2632,  0.3554,  0.3143,  0.2299],
        [-0.3976,  0.2438, -0.2224, -0.7535, -0.3636],
        [-0.0335, -0.6003, -0.4508,  0.2334, -0.3055],
        [-0.4297,  0.2362, -0.6859,  0.3319,  0.1649],
        [-0.4697, -0.3509,  0.3874,  0.1587, -0.5183],
        [ 0.2934,  0.5763, -0.0209,  0.3791, -0.6526]])
>>> s
tensor([ 27.4687,  22.6432,   8.5584,   5.9857,   2.0149])
>>> v
tensor([[-0.2514,  0.8148, -0.2606,  0.3967, -0.2180],
        [-0.3968,  0.3587,  0.7008, -0.4507,  0.1402],
        [-0.6922, -0.2489, -0.2208,  0.2513,  0.5891],
        [-0.3662, -0.3686,  0.3859,  0.4342, -0.6265],
        [-0.4076, -0.0980, -0.4933, -0.6227, -0.4396]])
>>> torch.dist(a, torch.mm(torch.mm(u, torch.diag(s)), v.t()))
tensor(1.00000e-06 *
       9.3738)

torch.symeig(input, eigenvectors=False, upper=True, out=None) -> (Tensor, Tensor)¶

This function returns eigenvalues and eigenvectors of a real symmetric matrix input, represented by a namedtuple (eigenvalues, eigenvectors).

input and $V$ are $(m \times m)$ matrices and $e$ is a $m$ dimensional vector.

This function calculates all eigenvalues (and vectors) of input such that $\text{input} = V \text{diag}(e) V^T$.

The boolean argument eigenvectors defines computation of eigenvectors or eigenvalues only.

If it is False, only eigenvalues are computed. If it is True, both eigenvalues and eigenvectors are computed.

Since the input matrix input is supposed to be symmetric, only the upper triangular portion is used by default.

If upper is False, then lower triangular portion is used.

Note

Irrespective of the original strides, the returned matrix V will be transposed, i.e. with strides (1, m) instead of (m, 1).

Note

Extra care needs to be taken when backward through outputs. Such operation is really only stable when all eigenvalues are distinct. Otherwise, NaN can appear as the gradients are not properly defined.

Parameters

input (Tensor) – the input symmetric matrix
eigenvectors (boolean, optional) – controls whether eigenvectors have to be computed
upper (boolean, optional) – controls whether to consider upper-triangular or lower-triangular region
out (tuple, optional) – the output tuple of (Tensor, Tensor)

Returns

A namedtuple (eigenvalues, eigenvectors) containing

eigenvalues (Tensor): Shape $(m)$. Each element is an eigenvalue of input, The eigenvalues are in ascending order.

eigenvectors (Tensor): Shape $(m \times m)$. If eigenvectors=False, it’s a tensor filled with zeros. Otherwise, this tensor contains the orthonormal eigenvectors of the input.

Return type

(Tensor, Tensor)

Examples:

>>> a = torch.tensor([[ 1.96,  0.00,  0.00,  0.00,  0.00],
                      [-6.49,  3.80,  0.00,  0.00,  0.00],
                      [-0.47, -6.39,  4.17,  0.00,  0.00],
                      [-7.20,  1.50, -1.51,  5.70,  0.00],
                      [-0.65, -6.34,  2.67,  1.80, -7.10]]).t()
>>> e, v = torch.symeig(a, eigenvectors=True)
>>> e
tensor([-11.0656,  -6.2287,   0.8640,   8.8655,  16.0948])
>>> v
tensor([[-0.2981, -0.6075,  0.4026, -0.3745,  0.4896],
        [-0.5078, -0.2880, -0.4066, -0.3572, -0.6053],
        [-0.0816, -0.3843, -0.6600,  0.5008,  0.3991],
        [-0.0036, -0.4467,  0.4553,  0.6204, -0.4564],
        [-0.8041,  0.4480,  0.1725,  0.3108,  0.1622]])

torch.trtrs(b, A, upper=True, transpose=False, unitriangular=False) -> (Tensor, Tensor)¶

Solves a system of equations with a triangular coefficient matrix $A$ and multiple right-hand sides b.

In particular, solves $AX = b$ and assumes $A$ is upper-triangular with the default keyword arguments.

Parameters

A (Tensor) – the input triangular coefficient matrix
b (Tensor) – multiple right-hand sides. Each column of $b$ is a right-hand side for the system of equations.
upper (bool, optional) – whether to solve the upper-triangular system of equations (default) or the lower-triangular system of equations. Default: True.
transpose (bool, optional) – whether $A$ should be transposed before being sent into the solver. Default: False.
unitriangular (bool, optional) – whether $A$ is unit triangular. If True, the diagonal elements of $A$ are assumed to be 1 and not referenced from $A$. Default: False.

Returns

A tuple $(X, M)$ where $M$ is a clone of $A$ and $X$ is the solution to $AX = b$ (or whatever variant of the system of equations, depending on the keyword arguments.)

Shape:

A: $(N, N)$
b: $(N, C)$
output[0]: $(N, C)$
output[1]: $(N, N)$

Examples:

>>> A = torch.randn(2, 2).triu()
>>> A
tensor([[ 1.1527, -1.0753],
        [ 0.0000,  0.7986]])
>>> b = torch.randn(2, 3)
>>> b
tensor([[-0.0210,  2.3513, -1.5492],
        [ 1.5429,  0.7403, -1.0243]])
>>> torch.trtrs(b, A)
(tensor([[ 1.7840,  2.9045, -2.5405],
        [ 1.9319,  0.9269, -1.2826]]), tensor([[ 1.1527, -1.0753],
        [ 0.0000,  0.7986]]))

Utilities¶

torch.compiled_with_cxx11_abi()¶: Returns whether PyTorch was built with _GLIBCXX_USE_CXX11_ABI=1

torch.Tensor¶

A torch.Tensor is a multi-dimensional matrix containing elements of a single data type.

Torch defines eight CPU tensor types and eight GPU tensor types:

Data type	dtype	CPU tensor	GPU tensor
32-bit floating point	`torch.float32` or `torch.float`	`torch.FloatTensor`	`torch.cuda.FloatTensor`
64-bit floating point	`torch.float64` or `torch.double`	`torch.DoubleTensor`	`torch.cuda.DoubleTensor`
16-bit floating point	`torch.float16` or `torch.half`	`torch.HalfTensor`	`torch.cuda.HalfTensor`
8-bit integer (unsigned)	`torch.uint8`	`torch.ByteTensor`	`torch.cuda.ByteTensor`
8-bit integer (signed)	`torch.int8`	`torch.CharTensor`	`torch.cuda.CharTensor`
16-bit integer (signed)	`torch.int16` or `torch.short`	`torch.ShortTensor`	`torch.cuda.ShortTensor`
32-bit integer (signed)	`torch.int32` or `torch.int`	`torch.IntTensor`	`torch.cuda.IntTensor`
64-bit integer (signed)	`torch.int64` or `torch.long`	`torch.LongTensor`	`torch.cuda.LongTensor`

torch.Tensor is an alias for the default tensor type (torch.FloatTensor).

A tensor can be constructed from a Python list or sequence using the torch.tensor() constructor:

>>> torch.tensor([[1., -1.], [1., -1.]])
tensor([[ 1.0000, -1.0000],
        [ 1.0000, -1.0000]])
>>> torch.tensor(np.array([[1, 2, 3], [4, 5, 6]]))
tensor([[ 1,  2,  3],
        [ 4,  5,  6]])

Warning

torch.tensor() always copies data. If you have a Tensor data and just want to change its requires_grad flag, use requires_grad_() or detach() to avoid a copy. If you have a numpy array and want to avoid a copy, use torch.as_tensor().

A tensor of specific data type can be constructed by passing a torch.dtype and/or a torch.device to a constructor or tensor creation op:

>>> torch.zeros([2, 4], dtype=torch.int32)
tensor([[ 0,  0,  0,  0],
        [ 0,  0,  0,  0]], dtype=torch.int32)
>>> cuda0 = torch.device('cuda:0')
>>> torch.ones([2, 4], dtype=torch.float64, device=cuda0)
tensor([[ 1.0000,  1.0000,  1.0000,  1.0000],
        [ 1.0000,  1.0000,  1.0000,  1.0000]], dtype=torch.float64, device='cuda:0')

The contents of a tensor can be accessed and modified using Python’s indexing and slicing notation:

>>> x = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> print(x[1][2])
tensor(6)
>>> x[0][1] = 8
>>> print(x)
tensor([[ 1,  8,  3],
        [ 4,  5,  6]])

Use torch.Tensor.item() to get a Python number from a tensor containing a single value:

>>> x = torch.tensor([[1]])
>>> x
tensor([[ 1]])
>>> x.item()
1
>>> x = torch.tensor(2.5)
>>> x
tensor(2.5000)
>>> x.item()
2.5

A tensor can be created with requires_grad=True so that torch.autograd records operations on them for automatic differentiation.

>>> x = torch.tensor([[1., -1.], [1., 1.]], requires_grad=True)
>>> out = x.pow(2).sum()
>>> out.backward()
>>> x.grad
tensor([[ 2.0000, -2.0000],
        [ 2.0000,  2.0000]])

Each tensor has an associated torch.Storage, which holds its data. The tensor class provides multi-dimensional, strided view of a storage and defines numeric operations on it.

Note

For more information on the torch.dtype, torch.device, and torch.layout attributes of a torch.Tensor, see Tensor Attributes.

Note

Methods which mutate a tensor are marked with an underscore suffix. For example, torch.FloatTensor.abs_() computes the absolute value in-place and returns the modified tensor, while torch.FloatTensor.abs() computes the result in a new tensor.

Note

To change an existing tensor’s torch.device and/or torch.dtype, consider using to() method on the tensor.

class torch.Tensor¶

There are a few main ways to create a tensor, depending on your use case.

To create a tensor with pre-existing data, use torch.tensor().
To create a tensor with specific size, use torch.* tensor creation ops (see Creation Ops).
To create a tensor with the same size (and similar types) as another tensor, use torch.*_like tensor creation ops (see Creation Ops).
To create a tensor with similar type but different size as another tensor, use tensor.new_* creation ops.

new_tensor(data, dtype=None, device=None, requires_grad=False) → Tensor¶

Returns a new Tensor with data as the tensor data. By default, the returned Tensor has the same torch.dtype and torch.device as this tensor.

Warning

new_tensor() always copies data. If you have a Tensor data and want to avoid a copy, use torch.Tensor.requires_grad_() or torch.Tensor.detach(). If you have a numpy array and want to avoid a copy, use torch.from_numpy().

Warning

When data is a tensor x, new_tensor() reads out ‘the data’ from whatever it is passed, and constructs a leaf variable. Therefore tensor.new_tensor(x) is equivalent to x.clone().detach() and tensor.new_tensor(x, requires_grad=True) is equivalent to x.clone().detach().requires_grad_(True). The equivalents using clone() and detach() are recommended.

Parameters

data (array_like) – The returned Tensor copies data.
dtype (torch.dtype, optional) – the desired type of returned tensor. Default: if None, same torch.dtype as this tensor.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, same torch.device as this tensor.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> tensor = torch.ones((2,), dtype=torch.int8)
>>> data = [[0, 1], [2, 3]]
>>> tensor.new_tensor(data)
tensor([[ 0,  1],
        [ 2,  3]], dtype=torch.int8)

new_full(size, fill_value, dtype=None, device=None, requires_grad=False) → Tensor¶

Returns a Tensor of size size filled with fill_value. By default, the returned Tensor has the same torch.dtype and torch.device as this tensor.

Parameters

fill_value (scalar) – the number to fill the output tensor with.
dtype (torch.dtype, optional) – the desired type of returned tensor. Default: if None, same torch.dtype as this tensor.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, same torch.device as this tensor.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> tensor = torch.ones((2,), dtype=torch.float64)
>>> tensor.new_full((3, 4), 3.141592)
tensor([[ 3.1416,  3.1416,  3.1416,  3.1416],
        [ 3.1416,  3.1416,  3.1416,  3.1416],
        [ 3.1416,  3.1416,  3.1416,  3.1416]], dtype=torch.float64)

new_empty(size, dtype=None, device=None, requires_grad=False) → Tensor¶

Returns a Tensor of size size filled with uninitialized data. By default, the returned Tensor has the same torch.dtype and torch.device as this tensor.

Parameters

dtype (torch.dtype, optional) – the desired type of returned tensor. Default: if None, same torch.dtype as this tensor.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, same torch.device as this tensor.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> tensor = torch.ones(())
>>> tensor.new_empty((2, 3))
tensor([[ 5.8182e-18,  4.5765e-41, -1.0545e+30],
        [ 3.0949e-41,  4.4842e-44,  0.0000e+00]])

new_ones(size, dtype=None, device=None, requires_grad=False) → Tensor¶

Returns a Tensor of size size filled with 1. By default, the returned Tensor has the same torch.dtype and torch.device as this tensor.

Parameters

size (int...) – a list, tuple, or torch.Size of integers defining the shape of the output tensor.
dtype (torch.dtype, optional) – the desired type of returned tensor. Default: if None, same torch.dtype as this tensor.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, same torch.device as this tensor.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> tensor = torch.tensor((), dtype=torch.int32)
>>> tensor.new_ones((2, 3))
tensor([[ 1,  1,  1],
        [ 1,  1,  1]], dtype=torch.int32)

new_zeros(size, dtype=None, device=None, requires_grad=False) → Tensor¶

Returns a Tensor of size size filled with 0. By default, the returned Tensor has the same torch.dtype and torch.device as this tensor.

Parameters

size (int...) – a list, tuple, or torch.Size of integers defining the shape of the output tensor.
dtype (torch.dtype, optional) – the desired type of returned tensor. Default: if None, same torch.dtype as this tensor.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, same torch.device as this tensor.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example:

>>> tensor = torch.tensor((), dtype=torch.float64)
>>> tensor.new_zeros((2, 3))
tensor([[ 0.,  0.,  0.],
        [ 0.,  0.,  0.]], dtype=torch.float64)

is_cuda¶: Is True if the Tensor is stored on the GPU, False otherwise.

device¶: Is the torch.device where this Tensor is.

abs() → Tensor¶: See torch.abs()

abs_() → Tensor¶: In-place version of abs()

acos() → Tensor¶: See torch.acos()

acos_() → Tensor¶: In-place version of acos()

add(value) → Tensor¶

add(value=1, other) -> Tensor

See torch.add()

add_(value) → Tensor¶

add_(value=1, other) -> Tensor

In-place version of add()

addbmm(beta=1, alpha=1, batch1, batch2) → Tensor¶: See torch.addbmm()

addbmm_(beta=1, alpha=1, batch1, batch2) → Tensor¶: In-place version of addbmm()

addcdiv(value=1, tensor1, tensor2) → Tensor¶: See torch.addcdiv()

addcdiv_(value=1, tensor1, tensor2) → Tensor¶: In-place version of addcdiv()

addcmul(value=1, tensor1, tensor2) → Tensor¶: See torch.addcmul()

addcmul_(value=1, tensor1, tensor2) → Tensor¶: In-place version of addcmul()

addmm(beta=1, alpha=1, mat1, mat2) → Tensor¶: See torch.addmm()

addmm_(beta=1, alpha=1, mat1, mat2) → Tensor¶: In-place version of addmm()

addmv(beta=1, alpha=1, mat, vec) → Tensor¶: See torch.addmv()

addmv_(beta=1, alpha=1, mat, vec) → Tensor¶: In-place version of addmv()

addr(beta=1, alpha=1, vec1, vec2) → Tensor¶: See torch.addr()

addr_(beta=1, alpha=1, vec1, vec2) → Tensor¶: In-place version of addr()

allclose(other, rtol=1e-05, atol=1e-08, equal_nan=False) → Tensor¶: See torch.allclose()

apply_(callable) → Tensor¶: Applies the function callable to each element in the tensor, replacing each element with the value returned by callable.

Note

This function only works with CPU tensors and should not be used in code sections that require high performance.

argmax(dim=None, keepdim=False)¶: See torch.argmax()

argmin(dim=None, keepdim=False)¶: See torch.argmin()

asin() → Tensor¶: See torch.asin()

asin_() → Tensor¶: In-place version of asin()

atan() → Tensor¶: See torch.atan()

atan2(other) → Tensor¶: See torch.atan2()

atan2_(other) → Tensor¶: In-place version of atan2()

atan_() → Tensor¶: In-place version of atan()

baddbmm(beta=1, alpha=1, batch1, batch2) → Tensor¶: See torch.baddbmm()

baddbmm_(beta=1, alpha=1, batch1, batch2) → Tensor¶: In-place version of baddbmm()

bernoulli(*, generator=None) → Tensor¶

Returns a result tensor where each $\texttt{result[i]}$ is independently sampled from $\text{Bernoulli}(\texttt{self[i]})$. self must have floating point dtype, and the result will have the same dtype.

See torch.bernoulli()

bernoulli_()¶

bernoulli_(p=0.5, *, generator=None) → Tensor: Fills each location of self with an independent sample from $\text{Bernoulli}(\texttt{p})$. self can have integral dtype.

bernoulli_(p_tensor, *, generator=None) → Tensor

p_tensor should be a tensor containing probabilities to be used for drawing the binary random number.

The $\text{i}^{th}$ element of self tensor will be set to a value sampled from $\text{Bernoulli}(\texttt{p\_tensor[i]})$.

self can have integral dtype, but p_tensor must have floating point dtype.

See also bernoulli() and torch.bernoulli()

bmm(batch2) → Tensor¶: See torch.bmm()

byte() → Tensor¶: self.byte() is equivalent to self.to(torch.uint8). See to().

btrifact(pivot=True) -> (Tensor, Tensor)¶: See torch.btrifact()

btrifact_with_info(pivot=True) -> (Tensor, Tensor, Tensor)¶: See torch.btrifact_with_info()

btrisolve(LU_data, LU_pivots) → Tensor¶: See torch.btrisolve()

cauchy_(median=0, sigma=1, *, generator=None) → Tensor¶: Fills the tensor with numbers drawn from the Cauchy distribution:

\[f(x) = \dfrac{1}{\pi} \dfrac{\sigma}{(x - \text{median})^2 + \sigma^2}\]

ceil() → Tensor¶: See torch.ceil()

ceil_() → Tensor¶: In-place version of ceil()

char() → Tensor¶: self.char() is equivalent to self.to(torch.int8). See to().

cholesky(upper=False) → Tensor¶: See torch.cholesky()

cholesky_solve(input2, upper=False) → Tensor¶: See torch.cholesky_solve()

chunk(chunks, dim=0) → List of Tensors¶: See torch.chunk()

clamp(min, max) → Tensor¶: See torch.clamp()

clamp_(min, max) → Tensor¶: In-place version of clamp()

clone() → Tensor¶: Returns a copy of the self tensor. The copy has the same size and data type as self.

Note

Unlike copy_(), this function is recorded in the computation graph. Gradients propagating to the cloned tensor will propagate to the original tensor.

contiguous() → Tensor¶: Returns a contiguous tensor containing the same data as self tensor. If self tensor is contiguous, this function returns the self tensor.

copy_(src, non_blocking=False) → Tensor¶

Copies the elements from src into self tensor and returns self.

The src tensor must be broadcastable with the self tensor. It may be of a different data type or reside on a different device.

Parameters

src (Tensor) – the source tensor to copy from
non_blocking (bool) – if True and this copy is between CPU and GPU, the copy may occur asynchronously with respect to the host. For other cases, this argument has no effect.

cos() → Tensor¶: See torch.cos()

cos_() → Tensor¶: In-place version of cos()

cosh() → Tensor¶: See torch.cosh()

cosh_() → Tensor¶: In-place version of cosh()

cpu() → Tensor¶

Returns a copy of this object in CPU memory.

If this object is already in CPU memory and on the correct device, then no copy is performed and the original object is returned.

cross(other, dim=-1) → Tensor¶: See torch.cross()

cuda(device=None, non_blocking=False) → Tensor¶

Returns a copy of this object in CUDA memory.

If this object is already in CUDA memory and on the correct device, then no copy is performed and the original object is returned.

Parameters

device (torch.device) – The destination GPU device. Defaults to the current CUDA device.
non_blocking (bool) – If True and the source is in pinned memory, the copy will be asynchronous with respect to the host. Otherwise, the argument has no effect. Default: False.

cumprod(dim, dtype=None) → Tensor¶: See torch.cumprod()

cumsum(dim, dtype=None) → Tensor¶: See torch.cumsum()

data_ptr() → int¶: Returns the address of the first element of self tensor.

det() → Tensor¶: See torch.det()

diag(diagonal=0) → Tensor¶: See torch.diag()

diag_embed(offset=0, dim1=-2, dim2=-1) → Tensor¶: See torch.diag_embed()

dim() → int¶: Returns the number of dimensions of self tensor.

dist(other, p=2) → Tensor¶: See torch.dist()

div(value) → Tensor¶: See torch.div()

div_(value) → Tensor¶: In-place version of div()

dot(tensor2) → Tensor¶: See torch.dot()

double() → Tensor¶: self.double() is equivalent to self.to(torch.float64). See to().

eig(eigenvectors=False) -> (Tensor, Tensor)¶: See torch.eig()

element_size() → int¶

Returns the size in bytes of an individual element.

Example:

>>> torch.tensor([]).element_size()
4
>>> torch.tensor([], dtype=torch.uint8).element_size()
1

eq(other) → Tensor¶: See torch.eq()

eq_(other) → Tensor¶: In-place version of eq()

equal(other) → bool¶: See torch.equal()

erf() → Tensor¶: See torch.erf()

erf_() → Tensor¶: In-place version of erf()

erfc() → Tensor¶: See torch.erfc()

erfc_() → Tensor¶: In-place version of erfc()

erfinv() → Tensor¶: See torch.erfinv()

erfinv_() → Tensor¶: In-place version of erfinv()

exp() → Tensor¶: See torch.exp()

exp_() → Tensor¶: In-place version of exp()

expm1() → Tensor¶: See torch.expm1()

expm1_() → Tensor¶: In-place version of expm1()

expand(*sizes) → Tensor¶

Returns a new view of the self tensor with singleton dimensions expanded to a larger size.

Passing -1 as the size for a dimension means not changing the size of that dimension.

Tensor can be also expanded to a larger number of dimensions, and the new ones will be appended at the front. For the new dimensions, the size cannot be set to -1.

Expanding a tensor does not allocate new memory, but only creates a new view on the existing tensor where a dimension of size one is expanded to a larger size by setting the stride to 0. Any dimension of size 1 can be expanded to an arbitrary value without allocating new memory.

Parameters: *sizes (torch.Size or int...) – the desired expanded size

Warning

More than one element of an expanded tensor may refer to a single memory location. As a result, in-place operations (especially ones that are vectorized) may result in incorrect behavior. If you need to write to the tensors, please clone them first.

Example:

>>> x = torch.tensor([[1], [2], [3]])
>>> x.size()
torch.Size([3, 1])
>>> x.expand(3, 4)
tensor([[ 1,  1,  1,  1],
        [ 2,  2,  2,  2],
        [ 3,  3,  3,  3]])
>>> x.expand(-1, 4)   # -1 means not changing the size of that dimension
tensor([[ 1,  1,  1,  1],
        [ 2,  2,  2,  2],
        [ 3,  3,  3,  3]])

expand_as(other) → Tensor¶

Expand this tensor to the same size as other. self.expand_as(other) is equivalent to self.expand(other.size()).

Please see expand() for more information about expand.

Parameters: other (torch.Tensor) – The result tensor has the same size as other.

exponential_(lambd=1, *, generator=None) → Tensor¶: Fills self tensor with elements drawn from the exponential distribution:

\[f(x) = \lambda e^{-\lambda x}\]

fill_(value) → Tensor¶: Fills self tensor with the specified value.

flatten(input, start_dim=0, end_dim=-1) → Tensor¶: see torch.flatten()

flip(dims) → Tensor¶: See torch.flip()

float() → Tensor¶: self.float() is equivalent to self.to(torch.float32). See to().

floor() → Tensor¶: See torch.floor()

floor_() → Tensor¶: In-place version of floor()

fmod(divisor) → Tensor¶: See torch.fmod()

fmod_(divisor) → Tensor¶: In-place version of fmod()

frac() → Tensor¶: See torch.frac()

frac_() → Tensor¶: In-place version of frac()

gather(dim, index) → Tensor¶: See torch.gather()

ge(other) → Tensor¶: See torch.ge()

ge_(other) → Tensor¶: In-place version of ge()

gels(A) → Tensor¶: See torch.gels()

geometric_(p, *, generator=None) → Tensor¶: Fills self tensor with elements drawn from the geometric distribution:

\[f(X=k) = (1 - p)^{k - 1} p\]

geqrf() -> (Tensor, Tensor)¶: See torch.geqrf()

ger(vec2) → Tensor¶: See torch.ger()

gesv(A)¶: See torch.solve()

get_device() -> Device ordinal (Integer)¶

For CUDA tensors, this function returns the device ordinal of the GPU on which the tensor resides. For CPU tensors, an error is thrown.

Example:

>>> x = torch.randn(3, 4, 5, device='cuda:0')
>>> x.get_device()
0
>>> x.cpu().get_device()  # RuntimeError: get_device is not implemented for type torch.FloatTensor

gt(other) → Tensor¶: See torch.gt()

gt_(other) → Tensor¶: In-place version of gt()

half() → Tensor¶: self.half() is equivalent to self.to(torch.float16). See to().

histc(bins=100, min=0, max=0) → Tensor¶: See torch.histc()

index_add_(dim, index, tensor) → Tensor¶

Accumulate the elements of tensor into the self tensor by adding to the indices in the order given in index. For example, if dim == 0 and index[i] == j, then the ith row of tensor is added to the jth row of self.

The dimth dimension of tensor must have the same size as the length of index (which must be a vector), and all other dimensions must match self, or an error will be raised.

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour that is not easily switched off. Please see the notes on /notes/randomness for background.

Parameters

dim (int) – dimension along which to index
index (LongTensor) – indices of tensor to select from
tensor (Tensor) – the tensor containing values to add

Example:

>>> x = torch.ones(5, 3)
>>> t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float)
>>> index = torch.tensor([0, 4, 2])
>>> x.index_add_(0, index, t)
tensor([[  2.,   3.,   4.],
        [  1.,   1.,   1.],
        [  8.,   9.,  10.],
        [  1.,   1.,   1.],
        [  5.,   6.,   7.]])

index_add(dim, index, tensor) → Tensor¶: Out-of-place version of torch.Tensor.index_add_()

index_copy_(dim, index, tensor) → Tensor¶

Copies the elements of tensor into the self tensor by selecting the indices in the order given in index. For example, if dim == 0 and index[i] == j, then the ith row of tensor is copied to the jth row of self.

The dimth dimension of tensor must have the same size as the length of index (which must be a vector), and all other dimensions must match self, or an error will be raised.

Parameters

dim (int) – dimension along which to index
index (LongTensor) – indices of tensor to select from
tensor (Tensor) – the tensor containing values to copy

Example:

>>> x = torch.zeros(5, 3)
>>> t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float)
>>> index = torch.tensor([0, 4, 2])
>>> x.index_copy_(0, index, t)
tensor([[ 1.,  2.,  3.],
        [ 0.,  0.,  0.],
        [ 7.,  8.,  9.],
        [ 0.,  0.,  0.],
        [ 4.,  5.,  6.]])

index_copy(dim, index, tensor) → Tensor¶: Out-of-place version of torch.Tensor.index_copy_()

index_fill_(dim, index, val) → Tensor¶

Fills the elements of the self tensor with value val by selecting the indices in the order given in index.

Parameters

dim (int) – dimension along which to index
index (LongTensor) – indices of self tensor to fill in
val (float) – the value to fill with

Example::

>>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float)
>>> index = torch.tensor([0, 2])
>>> x.index_fill_(1, index, -1)
tensor([[-1.,  2., -1.],
        [-1.,  5., -1.],
        [-1.,  8., -1.]])

index_fill(dim, index, value) → Tensor¶: Out-of-place version of torch.Tensor.index_fill_()

index_put_(indices, value, accumulate=False) → Tensor¶

Puts values from the tensor value into the tensor self using the indices specified in indices (which is a tuple of Tensors). The expression tensor.index_put_(indices, value) is equivalent to tensor[indices] = value. Returns self.

If accumulate is True, the elements in tensor are added to self. If accumulate is False, the behavior is undefined if indices contain duplicate elements.

Parameters

indices (tuple of LongTensor) – tensors used to index into self.
value (Tensor) – tensor of same dtype as self.
accumulate (bool) – whether to accumulate into self

index_put()¶

index_select(dim, index) → Tensor¶: See torch.index_select()

int() → Tensor¶: self.int() is equivalent to self.to(torch.int32). See to().

inverse() → Tensor¶: See torch.inverse()

is_contiguous() → bool¶: Returns True if self tensor is contiguous in memory in C order.

is_floating_point() → bool¶: Returns True if the data type of self is a floating point data type.

is_pinned()¶: Returns true if this tensor resides in pinned memory

is_set_to(tensor) → bool¶: Returns True if this object refers to the same THTensor object from the Torch C API as the given tensor.

is_signed()¶

item() → number¶

Returns the value of this tensor as a standard Python number. This only works for tensors with one element. For other cases, see tolist().

This operation is not differentiable.

Example:

>>> x = torch.tensor([1.0])
>>> x.item()
1.0

kthvalue(k, dim=None, keepdim=False) -> (Tensor, LongTensor)¶: See torch.kthvalue()

le(other) → Tensor¶: See torch.le()

le_(other) → Tensor¶: In-place version of le()

lerp(end, weight) → Tensor¶: See torch.lerp()

lerp_(end, weight) → Tensor¶: In-place version of lerp()

log() → Tensor¶: See torch.log()

log_() → Tensor¶: In-place version of log()

logdet() → Tensor¶: See torch.logdet()

log10() → Tensor¶: See torch.log10()

log10_() → Tensor¶: In-place version of log10()

log1p() → Tensor¶: See torch.log1p()

log1p_() → Tensor¶: In-place version of log1p()

log2() → Tensor¶: See torch.log2()

log2_() → Tensor¶: In-place version of log2()

log_normal_(mean=1, std=2, *, generator=None)¶: Fills self tensor with numbers samples from the log-normal distribution parameterized by the given mean $\mu$ and standard deviation $\sigma$. Note that mean and std are the mean and standard deviation of the underlying normal distribution, and not of the returned distribution:

\[f(x) = \dfrac{1}{x \sigma \sqrt{2\pi}}\ e^{-\frac{(\ln x - \mu)^2}{2\sigma^2}}\]

logsumexp(dim, keepdim=False) → Tensor¶: See torch.logsumexp()

long() → Tensor¶: self.long() is equivalent to self.to(torch.int64). See to().

lt(other) → Tensor¶: See torch.lt()

lt_(other) → Tensor¶: In-place version of lt()

map_(tensor, callable)¶

Applies callable for each element in self tensor and the given tensor and stores the results in self tensor. self tensor and the given tensor must be broadcastable.

The callable should have the signature:

def callable(a, b) -> number

masked_scatter_(mask, source)¶

Copies elements from source into self tensor at positions where the mask is one. The shape of mask must be broadcastable with the shape of the underlying tensor. The source should have at least as many elements as the number of ones in mask

Parameters

mask (ByteTensor) – the binary mask
source (Tensor) – the tensor to copy from

Note

The mask operates on the self tensor, not on the given source tensor.

masked_scatter(mask, tensor) → Tensor¶: Out-of-place version of torch.Tensor.masked_scatter_()

masked_fill_(mask, value)¶

Fills elements of self tensor with value where mask is one. The shape of mask must be broadcastable with the shape of the underlying tensor.

Parameters

mask (ByteTensor) – the binary mask
value (float) – the value to fill in with

masked_fill(mask, value) → Tensor¶: Out-of-place version of torch.Tensor.masked_fill_()

masked_select(mask) → Tensor¶: See torch.masked_select()

matmul(tensor2) → Tensor¶: See torch.matmul()

matrix_power(n) → Tensor¶: See torch.matrix_power()

max(dim=None, keepdim=False) -> Tensor or (Tensor, Tensor)¶: See torch.max()

mean(dim=None, keepdim=False) -> Tensor or (Tensor, Tensor)¶: See torch.mean()

median(dim=None, keepdim=False) -> (Tensor, LongTensor)¶: See torch.median()

min(dim=None, keepdim=False) -> Tensor or (Tensor, Tensor)¶: See torch.min()

mm(mat2) → Tensor¶: See torch.mm()

mode(dim=None, keepdim=False) -> (Tensor, LongTensor)¶: See torch.mode()

mul(value) → Tensor¶: See torch.mul()

mul_(value)¶: In-place version of mul()

multinomial(num_samples, replacement=False, *, generator=None) → Tensor¶: See torch.multinomial()

mv(vec) → Tensor¶: See torch.mv()

mvlgamma(p) → Tensor¶: See torch.mvlgamma()

mvlgamma_(p) → Tensor¶: In-place version of mvlgamma()

narrow(dimension, start, length) → Tensor¶

See torch.narrow()

Example:

>>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
>>> x.narrow(0, 0, 2)
tensor([[ 1,  2,  3],
        [ 4,  5,  6]])
>>> x.narrow(1, 1, 2)
tensor([[ 2,  3],
        [ 5,  6],
        [ 8,  9]])

ndimension() → int¶: Alias for dim()

ne(other) → Tensor¶: See torch.ne()

ne_(other) → Tensor¶: In-place version of ne()

neg() → Tensor¶: See torch.neg()

neg_() → Tensor¶: In-place version of neg()

nelement() → int¶: Alias for numel()

nonzero() → LongTensor¶: See torch.nonzero()

norm(p='fro', dim=None, keepdim=False, dtype=None)¶: See torch.norm()

normal_(mean=0, std=1, *, generator=None) → Tensor¶: Fills self tensor with elements samples from the normal distribution parameterized by mean and std.

numel() → int¶: See torch.numel()

numpy() → numpy.ndarray¶: Returns self tensor as a NumPy ndarray. This tensor and the returned ndarray share the same underlying storage. Changes to self tensor will be reflected in the ndarray and vice versa.

orgqr(input2) → Tensor¶: See torch.orgqr()

ormqr(input2, input3, left=True, transpose=False) → Tensor¶: See torch.ormqr()

permute(*dims) → Tensor¶

Permute the dimensions of this tensor.

Parameters: *dims (int...) – The desired ordering of dimensions

Example

>>> x = torch.randn(2, 3, 5)
>>> x.size()
torch.Size([2, 3, 5])
>>> x.permute(2, 0, 1).size()
torch.Size([5, 2, 3])

pin_memory()¶

pinverse() → Tensor¶: See torch.pinverse()

potrf(upper=True)¶: See torch.cholesky()

potri(upper=True) → Tensor¶: See torch.potri()

potrs(u, upper=True)¶: See torch.cholesky_solve()

pow(exponent) → Tensor¶: See torch.pow()

pow_(exponent) → Tensor¶: In-place version of pow()

prod(dim=None, keepdim=False, dtype=None) → Tensor¶: See torch.prod()

pstrf(upper=True)¶: See torch.pstrf()

put_(indices, tensor, accumulate=False) → Tensor¶

Copies the elements from tensor into the positions specified by indices. For the purpose of indexing, the self tensor is treated as if it were a 1-D tensor.

If accumulate is True, the elements in tensor are added to self. If accumulate is False, the behavior is undefined if indices contain duplicate elements.

Parameters

indices (LongTensor) – the indices into self
tensor (Tensor) – the tensor containing values to copy from
accumulate (bool) – whether to accumulate into self

Example:

>>> src = torch.tensor([[4, 3, 5],
                        [6, 7, 8]])
>>> src.put_(torch.tensor([1, 3]), torch.tensor([9, 10]))
tensor([[  4,   9,   5],
        [ 10,   7,   8]])

qr() -> (Tensor, Tensor)¶: See torch.qr()

random_(from=0, to=None, *, generator=None) → Tensor¶: Fills self tensor with numbers sampled from the discrete uniform distribution over [from, to - 1]. If not specified, the values are usually only bounded by self tensor’s data type. However, for floating point types, if unspecified, range will be [0, 2^mantissa] to ensure that every value is representable. For example, torch.tensor(1, dtype=torch.double).random_() will be uniform in [0, 2^53].

reciprocal() → Tensor¶: See torch.reciprocal()

reciprocal_() → Tensor¶: In-place version of reciprocal()

remainder(divisor) → Tensor¶: See torch.remainder()

remainder_(divisor) → Tensor¶: In-place version of remainder()

renorm(p, dim, maxnorm) → Tensor¶: See torch.renorm()

renorm_(p, dim, maxnorm) → Tensor¶: In-place version of renorm()

repeat(*sizes) → Tensor¶

Repeats this tensor along the specified dimensions.

Unlike expand(), this function copies the tensor’s data.

Warning

torch.repeat() behaves differently from numpy.repeat, but is more similar to numpy.tile.

Parameters: sizes (torch.Size or int...) – The number of times to repeat this tensor along each dimension

Example:

>>> x = torch.tensor([1, 2, 3])
>>> x.repeat(4, 2)
tensor([[ 1,  2,  3,  1,  2,  3],
        [ 1,  2,  3,  1,  2,  3],
        [ 1,  2,  3,  1,  2,  3],
        [ 1,  2,  3,  1,  2,  3]])
>>> x.repeat(4, 2, 1).size()
torch.Size([4, 2, 3])

requires_grad_(requires_grad=True) → Tensor¶

Change if autograd should record operations on this tensor: sets this tensor’s requires_grad attribute in-place. Returns this tensor.

require_grad_()’s main use case is to tell autograd to begin recording operations on a Tensor tensor. If tensor has requires_grad=False (because it was obtained through a DataLoader, or required preprocessing or initialization), tensor.requires_grad_() makes it so that autograd will begin to record operations on tensor.

Parameters: requires_grad (bool) – If autograd should record operations on this tensor. Default: True.

Example:

>>> # Let's say we want to preprocess some saved weights and use
>>> # the result as new weights.
>>> saved_weights = [0.1, 0.2, 0.3, 0.25]
>>> loaded_weights = torch.tensor(saved_weights)
>>> weights = preprocess(loaded_weights)  # some function
>>> weights
tensor([-0.5503,  0.4926, -2.1158, -0.8303])

>>> # Now, start to record operations done to weights
>>> weights.requires_grad_()
>>> out = weights.pow(2).sum()
>>> out.backward()
>>> weights.grad
tensor([-1.1007,  0.9853, -4.2316, -1.6606])

reshape(*shape) → Tensor¶

Returns a tensor with the same data and number of elements as self but with the specified shape. This method returns a view if shape is compatible with the current shape. See torch.Tensor.view() on when it is possible to return a view.

See torch.reshape()

Parameters: shape (tuple of python:ints or int...) – the desired shape

reshape_as(other) → Tensor¶

Returns this tensor as the same shape as other. self.reshape_as(other) is equivalent to self.reshape(other.sizes()). This method returns a view if other.sizes() is compatible with the current shape. See torch.Tensor.view() on when it is possible to return a view.

Please see reshape() for more information about reshape.

Parameters: other (torch.Tensor) – The result tensor has the same shape as other.

resize_(*sizes) → Tensor¶

Resizes self tensor to the specified size. If the number of elements is larger than the current storage size, then the underlying storage is resized to fit the new number of elements. If the number of elements is smaller, the underlying storage is not changed. Existing elements are preserved but any new memory is uninitialized.

Warning

This is a low-level method. The storage is reinterpreted as C-contiguous, ignoring the current strides (unless the target size equals the current size, in which case the tensor is left unchanged). For most purposes, you will instead want to use view(), which checks for contiguity, or reshape(), which copies data if needed. To change the size in-place with custom strides, see set_().

Parameters: sizes (torch.Size or int...) – the desired size

Example:

>>> x = torch.tensor([[1, 2], [3, 4], [5, 6]])
>>> x.resize_(2, 2)
tensor([[ 1,  2],
        [ 3,  4]])

resize_as_(tensor) → Tensor¶: Resizes the self tensor to be the same size as the specified tensor. This is equivalent to self.resize_(tensor.size()).

roll(shifts, dims) → Tensor¶: See torch.roll()

round() → Tensor¶: See torch.round()

round_() → Tensor¶: In-place version of round()

rsqrt() → Tensor¶: See torch.rsqrt()

rsqrt_() → Tensor¶: In-place version of rsqrt()

scatter_(dim, index, src) → Tensor¶

Writes all values from the tensor src into self at the indices specified in the index tensor. For each value in src, its output index is specified by its index in src for dimension != dim and by the corresponding value in index for dimension = dim.

For a 3-D tensor, self is updated as:

self[index[i][j][k]][j][k] = src[i][j][k]  # if dim == 0
self[i][index[i][j][k]][k] = src[i][j][k]  # if dim == 1
self[i][j][index[i][j][k]] = src[i][j][k]  # if dim == 2

This is the reverse operation of the manner described in gather().

self, index and src (if it is a Tensor) should have same number of dimensions. It is also required that index.size(d) <= src.size(d) for all dimensions d, and that index.size(d) <= self.size(d) for all dimensions d != dim.

Moreover, as for gather(), the values of index must be between 0 and self.size(dim) - 1 inclusive, and all values in a row along the specified dimension dim must be unique.

Parameters

dim (int) – the axis along which to index
index (LongTensor) – the indices of elements to scatter, can be either empty or the same size of src. When empty, the operation returns identity
src (Tensor) – the source element(s) to scatter, incase value is not specified
value (float) – the source element(s) to scatter, incase src is not specified

Example:

>>> x = torch.rand(2, 5)
>>> x
tensor([[ 0.3992,  0.2908,  0.9044,  0.4850,  0.6004],
        [ 0.5735,  0.9006,  0.6797,  0.4152,  0.1732]])
>>> torch.zeros(3, 5).scatter_(0, torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]), x)
tensor([[ 0.3992,  0.9006,  0.6797,  0.4850,  0.6004],
        [ 0.0000,  0.2908,  0.0000,  0.4152,  0.0000],
        [ 0.5735,  0.0000,  0.9044,  0.0000,  0.1732]])

>>> z = torch.zeros(2, 4).scatter_(1, torch.tensor([[2], [3]]), 1.23)
>>> z
tensor([[ 0.0000,  0.0000,  1.2300,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  1.2300]])

scatter(dim, index, source) → Tensor¶: Out-of-place version of torch.Tensor.scatter_()

scatter_add_(dim, index, other) → Tensor¶

Adds all values from the tensor other into self at the indices specified in the index tensor in a similar fashion as scatter_(). For each value in other, it is added to an index in self which is specified by its index in other for dimension != dim and by the corresponding value in index for dimension = dim.

For a 3-D tensor, self is updated as:

self[index[i][j][k]][j][k] += other[i][j][k]  # if dim == 0
self[i][index[i][j][k]][k] += other[i][j][k]  # if dim == 1
self[i][j][index[i][j][k]] += other[i][j][k]  # if dim == 2

self, index and other should have same number of dimensions. It is also required that index.size(d) <= other.size(d) for all dimensions d, and that index.size(d) <= self.size(d) for all dimensions d != dim.

Moreover, as for gather(), the values of index must be between 0 and self.size(dim) - 1 inclusive, and all values in a row along the specified dimension dim must be unique.

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour that is not easily switched off. Please see the notes on /notes/randomness for background.

Parameters

dim (int) – the axis along which to index
index (LongTensor) – the indices of elements to scatter and add, can be either empty or the same size of src. When empty, the operation returns identity.
other (Tensor) – the source elements to scatter and add

Example:

>>> x = torch.rand(2, 5)
>>> x
tensor([[0.7404, 0.0427, 0.6480, 0.3806, 0.8328],
        [0.7953, 0.2009, 0.9154, 0.6782, 0.9620]])
>>> torch.ones(3, 5).scatter_add_(0, torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]), x)
tensor([[1.7404, 1.2009, 1.9154, 1.3806, 1.8328],
        [1.0000, 1.0427, 1.0000, 1.6782, 1.0000],
        [1.7953, 1.0000, 1.6480, 1.0000, 1.9620]])

scatter_add(dim, index, source) → Tensor¶: Out-of-place version of torch.Tensor.scatter_add_()

select(dim, index) → Tensor¶

Slices the self tensor along the selected dimension at the given index. This function returns a tensor with the given dimension removed.

Parameters

dim (int) – the dimension to slice
index (int) – the index to select with

Note

select() is equivalent to slicing. For example, tensor.select(0, index) is equivalent to tensor[index] and tensor.select(2, index) is equivalent to tensor[:,:,index].

set_(source=None, storage_offset=0, size=None, stride=None) → Tensor¶

Sets the underlying storage, size, and strides. If source is a tensor, self tensor will share the same storage and have the same size and strides as source. Changes to elements in one tensor will be reflected in the other.

If source is a Storage, the method sets the underlying storage, offset, size, and stride.

Parameters

source (Tensor or Storage) – the tensor or storage to use
storage_offset (int, optional) – the offset in the storage
size (torch.Size, optional) – the desired size. Defaults to the size of the source.
stride (tuple, optional) – the desired stride. Defaults to C-contiguous strides.

share_memory_()¶

Moves the underlying storage to shared memory.

This is a no-op if the underlying storage is already in shared memory and for CUDA tensors. Tensors in shared memory cannot be resized.

short() → Tensor¶: self.short() is equivalent to self.to(torch.int16). See to().

sigmoid() → Tensor¶: See torch.sigmoid()

sigmoid_() → Tensor¶: In-place version of sigmoid()

sign() → Tensor¶: See torch.sign()

sign_() → Tensor¶: In-place version of sign()

sin() → Tensor¶: See torch.sin()

sin_() → Tensor¶: In-place version of sin()

sinh() → Tensor¶: See torch.sinh()

sinh_() → Tensor¶: In-place version of sinh()

size() → torch.Size¶

Returns the size of the self tensor. The returned value is a subclass of tuple.

Example:

>>> torch.empty(3, 4, 5).size()
torch.Size([3, 4, 5])

slogdet() -> (Tensor, Tensor)¶: See torch.slogdet()

solve(A) → Tensor, Tensor¶: See torch.solve()

sort(dim=-1, descending=False) -> (Tensor, LongTensor)¶: See torch.sort()

split(split_size, dim=0)¶: See torch.split()

sparse_mask(input, mask) → Tensor¶

Returns a new SparseTensor with values from Tensor input filtered by indices of mask and values are ignored. input and mask must have the same shape.

Parameters

input (Tensor) – an input Tensor
mask (SparseTensor) – a SparseTensor which we filter input based on its indices

Example:

>>> nnz = 5
>>> dims = [5, 5, 2, 2]
>>> I = torch.cat([torch.randint(0, dims[0], size=(nnz,)),
                   torch.randint(0, dims[1], size=(nnz,))], 0).reshape(2, nnz)
>>> V = torch.randn(nnz, dims[2], dims[3])
>>> size = torch.Size(dims)
>>> S = torch.sparse_coo_tensor(I, V, size).coalesce()
>>> D = torch.randn(dims)
>>> D.sparse_mask(S)
tensor(indices=tensor([[0, 0, 0, 2],
                       [0, 1, 4, 3]]),
       values=tensor([[[ 1.6550,  0.2397],
                       [-0.1611, -0.0779]],

                      [[ 0.2326, -1.0558],
                       [ 1.4711,  1.9678]],

                      [[-0.5138, -0.0411],
                       [ 1.9417,  0.5158]],

                      [[ 0.0793,  0.0036],
                       [-0.2569, -0.1055]]]),
       size=(5, 5, 2, 2), nnz=4, layout=torch.sparse_coo)

sqrt() → Tensor¶: See torch.sqrt()

sqrt_() → Tensor¶: In-place version of sqrt()

squeeze(dim=None) → Tensor¶: See torch.squeeze()

squeeze_(dim=None) → Tensor¶: In-place version of squeeze()

std(dim=None, unbiased=True, keepdim=False) → Tensor¶: See torch.std()

storage() → torch.Storage¶: Returns the underlying storage

storage_offset() → int¶

Returns self tensor’s offset in the underlying storage in terms of number of storage elements (not bytes).

Example:

>>> x = torch.tensor([1, 2, 3, 4, 5])
>>> x.storage_offset()
0
>>> x[3:].storage_offset()
3

storage_type()¶

stride(dim) → tuple or int¶

Returns the stride of self tensor.

Stride is the jump necessary to go from one element to the next one in the specified dimension dim. A tuple of all strides is returned when no argument is passed in. Otherwise, an integer value is returned as the stride in the particular dimension dim.

Parameters: dim (int, optional) – the desired dimension in which stride is required

Example:

>>> x = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
>>> x.stride()
(5, 1)
>>>x.stride(0)
5
>>> x.stride(-1)
1

sub(value, other) → Tensor¶

Subtracts a scalar or tensor from self tensor. If both value and other are specified, each element of other is scaled by value before being used.

When other is a tensor, the shape of other must be broadcastable with the shape of the underlying tensor.

sub_(x) → Tensor¶: In-place version of sub()

sum(dim=None, keepdim=False, dtype=None) → Tensor¶: See torch.sum()

svd(some=True, compute_uv=True) -> (Tensor, Tensor, Tensor)¶: See torch.svd()

symeig(eigenvectors=False, upper=True) -> (Tensor, Tensor)¶: See torch.symeig()

t() → Tensor¶: See torch.t()

t_() → Tensor¶: In-place version of t()

to(*args, **kwargs) → Tensor¶

Performs Tensor dtype and/or device conversion. A torch.dtype and torch.device are inferred from the arguments of self.to(*args, **kwargs).

Note

If the self Tensor already has the correct torch.dtype and torch.device, then self is returned. Otherwise, the returned tensor is a copy of self with the desired torch.dtype and torch.device.

Here are the ways to call to:

to(dtype, non_blocking=False, copy=False) → Tensor: Returns a Tensor with the specified dtype

to(device=None, dtype=None, non_blocking=False, copy=False) → Tensor: Returns a Tensor with the specified device and (optional) dtype. If dtype is None it is inferred to be self.dtype. When non_blocking, tries to convert asynchronously with respect to the host if possible, e.g., converting a CPU Tensor with pinned memory to a CUDA Tensor. When copy is set, a new Tensor is created even when the Tensor already matches the desired conversion.

to(other, non_blocking=False, copy=False) → Tensor: Returns a Tensor with same torch.dtype and torch.device as the Tensor other. When non_blocking, tries to convert asynchronously with respect to the host if possible, e.g., converting a CPU Tensor with pinned memory to a CUDA Tensor. When copy is set, a new Tensor is created even when the Tensor already matches the desired conversion.

Example:

>>> tensor = torch.randn(2, 2)  # Initially dtype=float32, device=cpu
>>> tensor.to(torch.float64)
tensor([[-0.5044,  0.0005],
        [ 0.3310, -0.0584]], dtype=torch.float64)

>>> cuda0 = torch.device('cuda:0')
>>> tensor.to(cuda0)
tensor([[-0.5044,  0.0005],
        [ 0.3310, -0.0584]], device='cuda:0')

>>> tensor.to(cuda0, dtype=torch.float64)
tensor([[-0.5044,  0.0005],
        [ 0.3310, -0.0584]], dtype=torch.float64, device='cuda:0')

>>> other = torch.randn((), dtype=torch.float64, device=cuda0)
>>> tensor.to(other, non_blocking=True)
tensor([[-0.5044,  0.0005],
        [ 0.3310, -0.0584]], dtype=torch.float64, device='cuda:0')

take(indices) → Tensor¶: See torch.take()

tan()¶

tan_() → Tensor¶: In-place version of tan()

tanh() → Tensor¶: See torch.tanh()

tanh_() → Tensor¶: In-place version of tanh()

tolist()¶

” tolist() -> list or number

Returns the tensor as a (nested) list. For scalars, a standard Python number is returned, just like with item(). Tensors are automatically moved to the CPU first if necessary.

This operation is not differentiable.

Examples:

>>> a = torch.randn(2, 2)
>>> a.tolist()
[[0.012766935862600803, 0.5415473580360413],
 [-0.08909505605697632, 0.7729271650314331]]
>>> a[0,0].tolist()
0.012766935862600803

topk(k, dim=None, largest=True, sorted=True) -> (Tensor, LongTensor)¶: See torch.topk()

to_sparse(sparseDims) → Tensor¶

Returns a sparse copy of the tensor. PyTorch supports sparse tensors in coordinate format.

Parameters: sparseDims (int, optional) – the number of sparse dimensions to include in the new sparse tensor

Example:

>>> d = torch.tensor([[0, 0, 0], [9, 0, 10], [0, 0, 0]])
>>> d
tensor([[ 0,  0,  0],
        [ 9,  0, 10],
        [ 0,  0,  0]])
>>> d.to_sparse()
tensor(indices=tensor([[1, 1],
                       [0, 2]]),
       values=tensor([ 9, 10]),
       size=(3, 3), nnz=2, layout=torch.sparse_coo)
>>> d.to_sparse(1)
tensor(indices=tensor([[1]]),
       values=tensor([[ 9,  0, 10]]),
       size=(3, 3), nnz=1, layout=torch.sparse_coo)

trace() → Tensor¶: See torch.trace()

transpose(dim0, dim1) → Tensor¶: See torch.transpose()

transpose_(dim0, dim1) → Tensor¶: In-place version of transpose()

tril(k=0) → Tensor¶: See torch.tril()

tril_(k=0) → Tensor¶: In-place version of tril()

triu(k=0) → Tensor¶: See torch.triu()

triu_(k=0) → Tensor¶: In-place version of triu()

trtrs(A, upper=True, transpose=False, unitriangular=False) -> (Tensor, Tensor)¶: See torch.trtrs()

trunc() → Tensor¶: See torch.trunc()

trunc_() → Tensor¶: In-place version of trunc()

type(dtype=None, non_blocking=False, **kwargs) → str or Tensor¶

Returns the type if dtype is not provided, else casts this object to the specified type.

If this is already of the correct type, no copy is performed and the original object is returned.

Parameters

dtype (type or string) – The desired type
non_blocking (bool) – If True, and the source is in pinned memory and destination is on the GPU or vice versa, the copy is performed asynchronously with respect to the host. Otherwise, the argument has no effect.
**kwargs – For compatibility, may contain the key async in place of the non_blocking argument. The async arg is deprecated.

type_as(tensor) → Tensor¶

Returns this tensor cast to the type of the given tensor.

This is a no-op if the tensor is already of the correct type. This is equivalent to self.type(tensor.type())

Parameters: tensor (Tensor) – the tensor which has the desired type

unfold(dim, size, step) → Tensor¶

Returns a tensor which contains all slices of size size from self tensor in the dimension dim.

Step between two slices is given by step.

If sizedim is the size of dimension dim for self, the size of dimension dim in the returned tensor will be (sizedim - size) / step + 1.

An additional dimension of size size is appended in the returned tensor.

Parameters

dim (int) – dimension in which unfolding happens
size (int) – the size of each slice that is unfolded
step (int) – the step between each slice

Example:

>>> x = torch.arange(1., 8)
>>> x
tensor([ 1.,  2.,  3.,  4.,  5.,  6.,  7.])
>>> x.unfold(0, 2, 1)
tensor([[ 1.,  2.],
        [ 2.,  3.],
        [ 3.,  4.],
        [ 4.,  5.],
        [ 5.,  6.],
        [ 6.,  7.]])
>>> x.unfold(0, 2, 2)
tensor([[ 1.,  2.],
        [ 3.,  4.],
        [ 5.,  6.]])

uniform_(from=0, to=1) → Tensor¶: Fills self tensor with numbers sampled from the continuous uniform distribution:

\[P(x) = \dfrac{1}{\text{to} - \text{from}} \]

unique(sorted=True, return_inverse=False, dim=None)¶

Returns the unique scalar elements of the tensor as a 1-D tensor.

See torch.unique()

unsqueeze(dim) → Tensor¶: See torch.unsqueeze()

unsqueeze_(dim) → Tensor¶: In-place version of unsqueeze()

var(dim=None, unbiased=True, keepdim=False) → Tensor¶: See torch.var()

view(*shape) → Tensor¶

Returns a new tensor with the same data as the self tensor but of a different shape.

The returned tensor shares the same data and must have the same number of elements, but may have a different size. For a tensor to be viewed, the new view size must be compatible with its original size and stride, i.e., each new view dimension must either be a subspace of an original dimension, or only span across original dimensions $d, d+1, \dots, d+k$ that satisfy the following contiguity-like condition that $\forall i = 0, \dots, k-1$,

\[\text{stride}[i] = \text{stride}[i+1] \times \text{size}[i+1]\]

Otherwise, contiguous() needs to be called before the tensor can be viewed. See also: reshape(), which returns a view if the shapes are compatible, and copies (equivalent to calling contiguous()) otherwise.

Parameters: shape (torch.Size or int...) – the desired size

Example:

>>> x = torch.randn(4, 4)
>>> x.size()
torch.Size([4, 4])
>>> y = x.view(16)
>>> y.size()
torch.Size([16])
>>> z = x.view(-1, 8)  # the size -1 is inferred from other dimensions
>>> z.size()
torch.Size([2, 8])

>>> a = torch.randn(1, 2, 3, 4)
>>> a.size()
torch.Size([1, 2, 3, 4])
>>> b = a.transpose(1, 2)  # Swaps 2nd and 3rd dimension
>>> b.size()
torch.Size([1, 3, 2, 4])
>>> c = a.view(1, 3, 2, 4)  # Does not change tensor layout in memory
>>> c.size()
torch.Size([1, 3, 2, 4])
>>> torch.equal(b, c)
False

view_as(other) → Tensor¶

View this tensor as the same size as other. self.view_as(other) is equivalent to self.view(other.size()).

Please see view() for more information about view.

Parameters: other (torch.Tensor) – The result tensor has the same size as other.

zero_() → Tensor¶: Fills self tensor with zeros.

class torch.ByteTensor¶

The following methods are unique to torch.ByteTensor.

all()¶

all() → bool

Returns True if all elements in the tensor are non-zero, False otherwise.

Example:

>>> a = torch.randn(1, 3).byte() % 2
>>> a
tensor([[1, 0, 0]], dtype=torch.uint8)
>>> a.all()
tensor(0, dtype=torch.uint8)

all(dim, keepdim=False, out=None) → Tensor

Returns True if all elements in each row of the tensor in the given dimension dim are non-zero, False otherwise.

Parameters

dim (int) – the dimension to reduce
keepdim (bool) – whether the output tensor has dim retained or not
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4, 2).byte() % 2
>>> a
tensor([[0, 0],
        [0, 0],
        [0, 1],
        [1, 1]], dtype=torch.uint8)
>>> a.all(dim=1)
tensor([0, 0, 0, 1], dtype=torch.uint8)

any()¶

any() → bool

Returns True if any elements in the tensor are non-zero, False otherwise.

Example:

>>> a = torch.randn(1, 3).byte() % 2
>>> a
tensor([[0, 0, 1]], dtype=torch.uint8)
>>> a.any()
tensor(1, dtype=torch.uint8)

any(dim, keepdim=False, out=None) → Tensor

Returns True if any elements in each row of the tensor in the given dimension dim are non-zero, False otherwise.

Parameters

dim (int) – the dimension to reduce
keepdim (bool) – whether the output tensor has dim retained or not
out (Tensor, optional) – the output tensor

Example:

>>> a = torch.randn(4, 2).byte() % 2
>>> a
tensor([[1, 0],
        [0, 0],
        [0, 1],
        [0, 0]], dtype=torch.uint8)
>>> a.any(dim=1)
tensor([1, 0, 1, 0], dtype=torch.uint8)

Tensor Attributes¶

Each torch.Tensor has a torch.dtype, torch.device, and torch.layout.

torch.dtype¶

class torch.dtype¶

A torch.dtype is an object that represents the data type of a torch.Tensor. PyTorch has eight different data types:

Data type	dtype	Tensor types
32-bit floating point	`torch.float32` or `torch.float`	`torch.*.FloatTensor`
64-bit floating point	`torch.float64` or `torch.double`	`torch.*.DoubleTensor`
16-bit floating point	`torch.float16` or `torch.half`	`torch.*.HalfTensor`
8-bit integer (unsigned)	`torch.uint8`	`torch.*.ByteTensor`
8-bit integer (signed)	`torch.int8`	`torch.*.CharTensor`
16-bit integer (signed)	`torch.int16` or `torch.short`	`torch.*.ShortTensor`
32-bit integer (signed)	`torch.int32` or `torch.int`	`torch.*.IntTensor`
64-bit integer (signed)	`torch.int64` or `torch.long`	`torch.*.LongTensor`

To find out if a torch.dtype is a floating point data type, the property is_floating_point can be used, which returns True if the data type is a floating point data type.

torch.device¶

class torch.device¶

A torch.device is an object representing the device on which a torch.Tensor is or will be allocated.

The torch.device contains a device type ('cpu' or 'cuda') and optional device ordinal for the device type. If the device ordinal is not present, this represents the current device for the device type; e.g. a torch.Tensor constructed with device 'cuda' is equivalent to 'cuda:X' where X is the result of torch.cuda.current_device().

A torch.Tensor’s device can be accessed via the Tensor.device property.

A torch.device can be constructed via a string or via a string and device ordinal

Via a string:

>>> torch.device('cuda:0')
device(type='cuda', index=0)

>>> torch.device('cpu')
device(type='cpu')

>>> torch.device('cuda')  # current cuda device
device(type='cuda')

Via a string and device ordinal:

>>> torch.device('cuda', 0)
device(type='cuda', index=0)

>>> torch.device('cpu', 0)
device(type='cpu', index=0)

Note

The torch.device argument in functions can generally be substituted with a string. This allows for fast prototyping of code.

>>> # Example of a function that takes in a torch.device
>>> cuda1 = torch.device('cuda:1')
>>> torch.randn((2,3), device=cuda1)

>>> # You can substitute the torch.device with a string
>>> torch.randn((2,3), device='cuda:1')

Note

For legacy reasons, a device can be constructed via a single device ordinal, which is treated as a cuda device. This matches Tensor.get_device(), which returns an ordinal for cuda tensors and is not supported for cpu tensors.

>>> torch.device(1)
device(type='cuda', index=1)

Note

Methods which take a device will generally accept a (properly formatted) string or (legacy) integer device ordinal, i.e. the following are all equivalent:

>>> torch.randn((2,3), device=torch.device('cuda:1'))
>>> torch.randn((2,3), device='cuda:1')
>>> torch.randn((2,3), device=1)  # legacy

torch.layout¶

class torch.layout¶

A torch.layout is an object that represents the memory layout of a torch.Tensor. Currently, we support torch.strided (dense Tensors) and have experimental support for torch.sparse_coo (sparse COO Tensors).

torch.strided represents dense Tensors and is the memory layout that is most commonly used. Each strided tensor has an associated torch.Storage, which holds its data. These tensors provide multi-dimensional, strided view of a storage. Strides are a list of integers: the k-th stride represents the jump in the memory necessary to go from one element to the next one in the k-th dimension of the Tensor. This concept makes it possible to perform many tensor operations efficiently.

Example:

>>> x = torch.Tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
>>> x.stride()
(5, 1)

>>> x.t().stride()
(1, 5)

For more information on torch.sparse_coo tensors, see torch.sparse.

Type Info¶

The numerical properties of a torch.dtype can be accessed through either the torch.finfo or the torch.iinfo.

torch.finfo¶

class torch.finfo¶

A torch.finfo is an object that represents the numerical properties of a floating point torch.dtype, (i.e. torch.float32, torch.float64, and torch.float16). This is similar to numpy.finfo.

A torch.finfo provides the following attributes:

Name	Type	Description
bits	int	The number of bits occupied by the type.
eps	float	The smallest representable number such that `1.0 + eps != 1.0`.
max	float	The largest representable number.
min	float	The smallest representable number (typically `-max`).
tiny	float	The smallest positive representable number.

Note

The constructor of torch.finfo can be called without argument, in which case the class is created for the pytorch default dtype (as returned by torch.get_default_dtype()).

torch.iinfo¶

class torch.iinfo¶

A torch.iinfo is an object that represents the numerical properties of a integer torch.dtype (i.e. torch.uint8, torch.int8, torch.int16, torch.int32, and torch.int64). This is similar to numpy.iinfo.

A torch.iinfo provides the following attributes:

Name	Type	Description
bits	int	The number of bits occupied by the type.
max	int	The largest representable number.
min	int	The smallest representable number.

torch.sparse¶

Warning

This API is currently experimental and may change in the near future.

Torch supports sparse tensors in COO(rdinate) format, which can efficiently store and process tensors for which the majority of elements are zeros.

A sparse tensor is represented as a pair of dense tensors: a tensor of values and a 2D tensor of indices. A sparse tensor can be constructed by providing these two tensors, as well as the size of the sparse tensor (which cannot be inferred from these tensors!) Suppose we want to define a sparse tensor with the entry 3 at location (0, 2), entry 4 at location (1, 0), and entry 5 at location (1, 2). We would then write:

>>> i = torch.LongTensor([[0, 1, 1],
                          [2, 0, 2]])
>>> v = torch.FloatTensor([3, 4, 5])
>>> torch.sparse.FloatTensor(i, v, torch.Size([2,3])).to_dense()
 0  0  3
 4  0  5
[torch.FloatTensor of size 2x3]

Note that the input to LongTensor is NOT a list of index tuples. If you want to write your indices this way, you should transpose before passing them to the sparse constructor:

>>> i = torch.LongTensor([[0, 2], [1, 0], [1, 2]])
>>> v = torch.FloatTensor([3,      4,      5    ])
>>> torch.sparse.FloatTensor(i.t(), v, torch.Size([2,3])).to_dense()
 0  0  3
 4  0  5
[torch.FloatTensor of size 2x3]

You can also construct hybrid sparse tensors, where only the first n dimensions are sparse, and the rest of the dimensions are dense.

>>> i = torch.LongTensor([[2, 4]])
>>> v = torch.FloatTensor([[1, 3], [5, 7]])
>>> torch.sparse.FloatTensor(i, v).to_dense()
 0  0
 0  0
 1  3
 0  0
 5  7
[torch.FloatTensor of size 5x2]

An empty sparse tensor can be constructed by specifying its size:

>>> torch.sparse.FloatTensor(2, 3)
SparseFloatTensor of size 2x3 with indices:
[torch.LongTensor with no dimension]
and values:
[torch.FloatTensor with no dimension]

SparseTensor has the following invariants:

sparse_dim + dense_dim = len(SparseTensor.shape)
SparseTensor._indices().shape = (sparse_dim, nnz)
SparseTensor._values().shape = (nnz, SparseTensor.shape[sparse_dim:])

Since SparseTensor._indices() is always a 2D tensor, the smallest sparse_dim = 1. Therefore, representation of a SparseTensor of sparse_dim = 0 is simply a dense tensor.

Note

Our sparse tensor format permits uncoalesced sparse tensors, where there may be duplicate coordinates in the indices; in this case, the interpretation is that the value at that index is the sum of all duplicate value entries. Uncoalesced tensors permit us to implement certain operators more efficiently.

For the most part, you shouldn’t have to care whether or not a sparse tensor is coalesced or not, as most operations will work identically given a coalesced or uncoalesced sparse tensor. However, there are two cases in which you may need to care.

First, if you repeatedly perform an operation that can produce duplicate entries (e.g., torch.sparse.FloatTensor.add()), you should occasionally coalesce your sparse tensors to prevent them from growing too large.

Second, some operators will produce different values depending on whether or not they are coalesced or not (e.g., torch.sparse.FloatTensor._values() and torch.sparse.FloatTensor._indices(), as well as torch.Tensor.sparse_mask()). These operators are prefixed by an underscore to indicate that they reveal internal implementation details and should be used with care, since code that works with coalesced sparse tensors may not work with uncoalesced sparse tensors; generally speaking, it is safest to explicitly coalesce before working with these operators.

For example, suppose that we wanted to implement an operator by operating directly on torch.sparse.FloatTensor._values(). Multiplication by a scalar can be implemented in the obvious way, as multiplication distributes over addition; however, square root cannot be implemented directly, since sqrt(a + b) != sqrt(a) + sqrt(b) (which is what would be computed if you were given an uncoalesced tensor.)

class torch.sparse.FloatTensor¶

add()¶

add_()¶

clone()¶

dim()¶

div()¶

div_()¶

get_device()¶

hspmm()¶

mm()¶

mul()¶

mul_()¶

narrow_copy()¶

resizeAs_()¶

size()¶

spadd()¶

spmm()¶

sspaddmm()¶

sspmm()¶

sub()¶

sub_()¶

t_()¶

toDense()¶

transpose()¶

transpose_()¶

zero_()¶

coalesce()¶

is_coalesced()¶

_indices()¶

_values()¶

_nnz()¶

Functions¶

torch.sparse.addmm(mat, mat1, mat2, beta=1, alpha=1)¶

This function does exact same thing as torch.addmm() in the forward, except that it supports backward for sparse matrix mat1. mat1 need to have sparse_dim = 2. Note that the gradients of mat1 is a coalesced sparse tensor.

Parameters

mat (Tensor) – a dense matrix to be added
mat1 (SparseTensor) – a sparse matrix to be multiplied
mat2 (Tensor) – a dense matrix be multiplied
beta (Number, optional) – multiplier for mat ($\beta$)
alpha (Number, optional) – multiplier for $mat1 @ mat2$ ($\alpha$)

torch.sparse.mm(mat1, mat2)¶

Performs a matrix multiplication of the sparse matrix mat1 and dense matrix mat2. Similar to torch.mm(), If mat1 is a $(n \times m)$ tensor, mat2 is a $(m \times p)$ tensor, out will be a $(n \times p)$ dense tensor. mat1 need to have sparse_dim = 2. This function also supports backward for both matrices. Note that the gradients of mat1 is a coalesced sparse tensor.

Parameters

mat1 (SparseTensor) – the first sparse matrix to be multiplied
mat2 (Tensor) – the second dense matrix to be multiplied

Example:

>>> a = torch.randn(2, 3).to_sparse().requires_grad_(True)
>>> a
tensor(indices=tensor([[0, 0, 0, 1, 1, 1],
                       [0, 1, 2, 0, 1, 2]]),
       values=tensor([ 1.5901,  0.0183, -0.6146,  1.8061, -0.0112,  0.6302]),
       size=(2, 3), nnz=6, layout=torch.sparse_coo, requires_grad=True)

>>> b = torch.randn(3, 2, requires_grad=True)
>>> b
tensor([[-0.6479,  0.7874],
        [-1.2056,  0.5641],
        [-1.1716, -0.9923]], requires_grad=True)

>>> y = torch.sparse.mm(a, b)
>>> y
tensor([[-0.3323,  1.8723],
        [-1.8951,  0.7904]], grad_fn=<SparseAddmmBackward>)
>>> y.sum().backward()
>>> a.grad
tensor(indices=tensor([[0, 0, 0, 1, 1, 1],
                       [0, 1, 2, 0, 1, 2]]),
       values=tensor([ 0.1394, -0.6415, -2.1639,  0.1394, -0.6415, -2.1639]),
       size=(2, 3), nnz=6, layout=torch.sparse_coo)

torch.sparse.sum(input, dim=None, dtype=None)¶

Returns the sum of each row of SparseTensor input in the given dimensions dim. If dim is a list of dimensions, reduce over all of them. When sum over all sparse_dim, this method returns a Tensor instead of SparseTensor.

All summed dim are squeezed (see torch.squeeze()), resulting an output tensor having dim fewer dimensions than input.

During backward, only gradients at nnz locations of input will propagate back. Note that the gradients of input is coalesced.

Parameters

input (Tensor) – the input SparseTensor
dim (int or tuple of python:ints) – a dimension or a list of dimensions to reduce. Default: reduce over all dims.
dtype (torch.dtype, optional) – the desired data type of returned Tensor. Default: dtype of input.

Example:

>>> nnz = 3
>>> dims = [5, 5, 2, 3]
>>> I = torch.cat([torch.randint(0, dims[0], size=(nnz,)),
                   torch.randint(0, dims[1], size=(nnz,))], 0).reshape(2, nnz)
>>> V = torch.randn(nnz, dims[2], dims[3])
>>> size = torch.Size(dims)
>>> S = torch.sparse_coo_tensor(I, V, size)
>>> S
tensor(indices=tensor([[2, 0, 3],
                       [2, 4, 1]]),
       values=tensor([[[-0.6438, -1.6467,  1.4004],
                       [ 0.3411,  0.0918, -0.2312]],

                      [[ 0.5348,  0.0634, -2.0494],
                       [-0.7125, -1.0646,  2.1844]],

                      [[ 0.1276,  0.1874, -0.6334],
                       [-1.9682, -0.5340,  0.7483]]]),
       size=(5, 5, 2, 3), nnz=3, layout=torch.sparse_coo)

# when sum over only part of sparse_dims, return a SparseTensor
>>> torch.sparse.sum(S, [1, 3])
tensor(indices=tensor([[0, 2, 3]]),
       values=tensor([[-1.4512,  0.4073],
                      [-0.8901,  0.2017],
                      [-0.3183, -1.7539]]),
       size=(5, 2), nnz=3, layout=torch.sparse_coo)

# when sum over all sparse dim, return a dense Tensor
# with summed dims squeezed
>>> torch.sparse.sum(S, [0, 1, 3])
tensor([-2.6596, -1.1450])

torch.cuda¶

This package adds support for CUDA tensor types, that implement the same function as CPU tensors, but they utilize GPUs for computation.

It is lazily initialized, so you can always import it, and use is_available() to determine if your system supports CUDA.

cuda-semantics has more details about working with CUDA.

torch.cuda.current_blas_handle()¶: Returns cublasHandle_t pointer to current cuBLAS handle

torch.cuda.current_device()¶: Returns the index of a currently selected device.

torch.cuda.current_stream(device=None)¶

Returns the currently selected Stream for a given device.

Parameters: device (torch.device or int, optional) – selected device. Returns the currently selected Stream for the current device, given by current_device(), if device is None (default).

torch.cuda.default_stream(device=None)¶

Returns the default Stream for a given device.

Parameters: device (torch.device or int, optional) – selected device. Returns the default Stream for the current device, given by current_device(), if device is None (default).

class torch.cuda.device(device)¶

Context-manager that changes the selected device.

Parameters: device (torch.device or int) – device index to select. It’s a no-op if this argument is a negative integer or None.

torch.cuda.device_count()¶: Returns the number of GPUs available.

class torch.cuda.device_of(obj)¶

Context-manager that changes the current device to that of given object.

You can use both tensors and storages as arguments. If a given object is not allocated on a GPU, this is a no-op.

Parameters: obj (Tensor or Storage) – object allocated on the selected device.

torch.cuda.empty_cache()¶: Releases all unoccupied cached memory currently held by the caching allocator so that those can be used in other GPU application and visible in nvidia-smi.

Note

empty_cache() doesn’t increase the amount of GPU memory available for PyTorch. See cuda-memory-management for more details about GPU memory management.

torch.cuda.get_device_capability(device=None)¶

Gets the cuda capability of a device.

Parameters: device (torch.device or int, optional) – device for which to return the device capability. This function is a no-op if this argument is a negative integer. Uses the current device, given by current_device(), if device is None (default).
Returns: the major and minor cuda capability of the device
Return type: tuple(int, int)

torch.cuda.get_device_name(device=None)¶

Gets the name of a device.

Parameters: device (torch.device or int, optional) – device for which to return the name. This function is a no-op if this argument is a negative integer. Uses the current device, given by current_device(), if device is None (default).

torch.cuda.init()¶

Initialize PyTorch’s CUDA state. You may need to call this explicitly if you are interacting with PyTorch via its C API, as Python bindings for CUDA functionality will not be until this initialization takes place. Ordinary users should not need this, as all of PyTorch’s CUDA methods automatically initialize CUDA state on-demand.

Does nothing if the CUDA state is already initialized.

torch.cuda.is_available()¶: Returns a bool indicating if CUDA is currently available.

torch.cuda.max_memory_allocated(device=None)¶

Returns the maximum GPU memory occupied by tensors in bytes for a given device.

By default, this returns the peak allocated memory since the beginning of this program. reset_max_memory_allocated() can be used to reset the starting point in tracking this metric. For example, these two functions can measure the peak allocated memory usage of each iteration in a training loop.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

torch.cuda.max_memory_cached(device=None)¶

Returns the maximum GPU memory managed by the caching allocator in bytes for a given device.

By default, this returns the peak cached memory since the beginning of this program. reset_max_memory_cached() can be used to reset the starting point in tracking this metric. For example, these two functions can measure the peak cached memory amount of each iteration in a training loop.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

torch.cuda.memory_allocated(device=None)¶

Returns the current GPU memory occupied by tensors in bytes for a given device.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

This is likely less than the amount shown in nvidia-smi since some unused memory can be held by the caching allocator and some context needs to be created on GPU. See cuda-memory-management for more details about GPU memory management.

torch.cuda.memory_cached(device=None)¶

Returns the current GPU memory managed by the caching allocator in bytes for a given device.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

torch.cuda.reset_max_memory_allocated(device=None)¶

Resets the starting point in tracking maximum GPU memory occupied by tensors for a given device.

See max_memory_allocated() for details.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

torch.cuda.reset_max_memory_cached(device=None)¶

Resets the starting point in tracking maximum GPU memory managed by the caching allocator for a given device.

See max_memory_cached() for details.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

torch.cuda.set_device(device)¶

Sets the current device.

Usage of this function is discouraged in favor of device. In most cases it’s better to use CUDA_VISIBLE_DEVICES environmental variable.

Parameters: device (torch.device or int) – selected device. This function is a no-op if this argument is negative.

torch.cuda.stream(stream)¶

Context-manager that selects a given stream.

All CUDA kernels queued within its context will be enqueued on a selected stream.

Parameters: stream (Stream) – selected stream. This manager is a no-op if it’s None.

Note

Streams are per-device. If the selected stream is not on the current device, this function will also change the current device to match the stream.

torch.cuda.synchronize()¶: Waits for all kernels in all streams on current device to complete.

Random Number Generator¶

torch.cuda.get_rng_state(device=device(type='cuda'))¶

Returns the random number generator state of the current GPU as a ByteTensor.

Parameters: device (torch.device or int, optional) – The device to return the RNG state of. Default: torch.device('cuda') (i.e., the current CUDA device).

Warning

This function eagerly initializes CUDA.

torch.cuda.get_rng_state_all()¶: Returns a tuple of ByteTensor representing the random number states of all devices.

torch.cuda.set_rng_state(new_state, device=device(type='cuda'))¶

Sets the random number generator state of the current GPU.

Parameters

new_state (torch.ByteTensor) – The desired state
device (torch.device or int, optional) – The device to set the RNG state. Default: torch.device('cuda') (i.e., the current CUDA device).

torch.cuda.set_rng_state_all(new_states)¶

Sets the random number generator state of all devices.

Parameters: new_state (tuple of torch.ByteTensor) – The desired state for each device

torch.cuda.manual_seed(seed)¶

Sets the seed for generating random numbers for the current GPU. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.

Parameters: seed (int) – The desired seed.

Warning

If you are working with a multi-GPU model, this function is insufficient to get determinism. To seed all GPUs, use manual_seed_all().

torch.cuda.manual_seed_all(seed)¶

Sets the seed for generating random numbers on all GPUs. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.

Parameters: seed (int) – The desired seed.

torch.cuda.seed()¶: Sets the seed for generating random numbers to a random number for the current GPU. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.

Warning

If you are working with a multi-GPU model, this function will only initialize the seed on one GPU. To initialize all GPUs, use seed_all().

torch.cuda.seed_all()¶: Sets the seed for generating random numbers to a random number on all GPUs. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.

torch.cuda.initial_seed()¶: Returns the current random seed of the current GPU.

Warning

This function eagerly initializes CUDA.

Communication collectives¶

torch.cuda.comm.broadcast(tensor, devices)¶

Broadcasts a tensor to a number of GPUs.

Parameters

tensor (Tensor) – tensor to broadcast.
devices (Iterable) – an iterable of devices among which to broadcast. Note that it should be like (src, dst1, dst2, …), the first element of which is the source device to broadcast from.

Returns

A tuple containing copies of the tensor, placed on devices corresponding to indices from devices.

torch.cuda.comm.broadcast_coalesced(tensors, devices, buffer_size=10485760)¶

Broadcasts a sequence tensors to the specified GPUs. Small tensors are first coalesced into a buffer to reduce the number of synchronizations.

Parameters

tensors (sequence) – tensors to broadcast.
devices (Iterable) – an iterable of devices among which to broadcast. Note that it should be like (src, dst1, dst2, …), the first element of which is the source device to broadcast from.
buffer_size (int) – maximum size of the buffer used for coalescing

Returns

A tuple containing copies of the tensor, placed on devices corresponding to indices from devices.

torch.cuda.comm.reduce_add(inputs, destination=None)¶

Sums tensors from multiple GPUs.

All inputs should have matching shapes.

Parameters

inputs (Iterable[Tensor]) – an iterable of tensors to add.
destination (int, optional) – a device on which the output will be placed (default: current device).

Returns

A tensor containing an elementwise sum of all inputs, placed on the destination device.

torch.cuda.comm.scatter(tensor, devices, chunk_sizes=None, dim=0, streams=None)¶

Scatters tensor across multiple GPUs.

Parameters

tensor (Tensor) – tensor to scatter.
devices (Iterable[int]) – iterable of ints, specifying among which devices the tensor should be scattered.
chunk_sizes (Iterable[int], optional) – sizes of chunks to be placed on each device. It should match devices in length and sum to tensor.size(dim). If not specified, the tensor will be divided into equal chunks.
dim (int, optional) – A dimension along which to chunk the tensor.

Returns

A tuple containing chunks of the tensor, spread across given devices.

torch.cuda.comm.gather(tensors, dim=0, destination=None)¶

Gathers tensors from multiple GPUs.

Tensor sizes in all dimension different than dim have to match.

Parameters

tensors (Iterable[Tensor]) – iterable of tensors to gather.
dim (int) – a dimension along which the tensors will be concatenated.
destination (int, optional) – output device (-1 means CPU, default: current device)

Returns

A tensor located on destination device, that is a result of concatenating tensors along dim.

Streams and events¶

class torch.cuda.Stream¶

Wrapper around a CUDA stream.

A CUDA stream is a linear sequence of execution that belongs to a specific device, independent from other streams. See cuda-semantics for details.

Parameters

device (torch.device or int, optional) – a device on which to allocate the stream. If device is None (default) or a negative integer, this will use the current device.
priority (int, optional) – priority of the stream. Lower numbers represent higher priorities.

query()¶

Checks if all the work submitted has been completed.

Returns: A boolean indicating if all kernels in this stream are completed.

record_event(event=None)¶

Records an event.

Parameters: event (Event, optional) – event to record. If not given, a new one will be allocated.
Returns: Recorded event.

synchronize()¶: Wait for all the kernels in this stream to complete.

Note

This is a wrapper around cudaStreamSynchronize(): see `CUDA documentation`_ for more info.

wait_event(event)¶

Makes all future work submitted to the stream wait for an event.

Parameters: event (Event) – an event to wait for.

Note

This is a wrapper around cudaStreamWaitEvent(): see `CUDA documentation`_ for more info.

This function returns without waiting for event: only future operations are affected.

wait_stream(stream)¶

Synchronizes with another stream.

All future work submitted to this stream will wait until all kernels submitted to a given stream at the time of call complete.

Parameters: stream (Stream) – a stream to synchronize.

Note

This function returns without waiting for currently enqueued kernels in stream: only future operations are affected.

class torch.cuda.Event¶

Wrapper around a CUDA event.

CUDA events are synchronization markers that can be used to monitor the device’s progress, to accurately measure timing, and to synchronize CUDA streams.

The underlying CUDA events are lazily initialized when the event is first recorded or exported to another process. After creation, only streams on the same device may record the event. However, streams on any device can wait on the event.

Parameters

enable_timing (bool, optional) – indicates if the event should measure time (default: False)
blocking (bool, optional) – if True, wait() will be blocking (default: False)
interprocess () – if True, the event can be shared between processes (default: False)

elapsed_time(end_event)¶: Returns the time elapsed in milliseconds after the event was recorded and before the end_event was recorded.

classmethod from_ipc_handle(device, handle)¶: Reconstruct an event from an IPC handle on the given device.

ipc_handle()¶: Returns an IPC handle of this event. If not recorded yet, the event will use the current device.

query()¶

Checks if all work currently captured by event has completed.

Returns: A boolean indicating if all work currently captured by event has completed.

record(stream=None)¶

Records the event in a given stream.

Uses torch.cuda.current_stream() if no stream is specified. The stream’s device must match the event’s device.

synchronize()¶

Waits for the event to complete.

Waits until the completion of all work currently captured in this event. This prevents the CPU thread from proceeding until the event completes.

Note

This is a wrapper around cudaEventSynchronize(): see `CUDA documentation`_ for more info.

wait(stream=None)¶

Makes all future work submitted to the given stream wait for this event.

Use torch.cuda.current_stream() if no stream is specified.

Memory management¶

torch.cuda.empty_cache(): Releases all unoccupied cached memory currently held by the caching allocator so that those can be used in other GPU application and visible in nvidia-smi.

Note

empty_cache() doesn’t increase the amount of GPU memory available for PyTorch. See cuda-memory-management for more details about GPU memory management.

torch.cuda.memory_allocated(device=None)

Returns the current GPU memory occupied by tensors in bytes for a given device.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

torch.cuda.max_memory_allocated(device=None)

Returns the maximum GPU memory occupied by tensors in bytes for a given device.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

torch.cuda.reset_max_memory_allocated(device=None)

Resets the starting point in tracking maximum GPU memory occupied by tensors for a given device.

See max_memory_allocated() for details.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

torch.cuda.memory_cached(device=None)

Returns the current GPU memory managed by the caching allocator in bytes for a given device.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

torch.cuda.max_memory_cached(device=None)

Returns the maximum GPU memory managed by the caching allocator in bytes for a given device.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

torch.cuda.reset_max_memory_cached(device=None)

Resets the starting point in tracking maximum GPU memory managed by the caching allocator for a given device.

See max_memory_cached() for details.

Parameters: device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

Note

See cuda-memory-management for more details about GPU memory management.

NVIDIA Tools Extension (NVTX)¶

torch.cuda.nvtx.mark(msg)¶

Describe an instantaneous event that occurred at some point.

Parameters: msg (string) – ASCII message to associate with the event.

torch.cuda.nvtx.range_push(msg)¶

Pushes a range onto a stack of nested range span. Returns zero-based depth of the range that is started.

Parameters: msg (string) – ASCII message to associate with range

torch.cuda.nvtx.range_pop()¶: Pops a range off of a stack of nested range spans. Returns the zero-based depth of the range that is ended.

torch.Storage¶

A torch.Storage is a contiguous, one-dimensional array of a single data type.

Every torch.Tensor has a corresponding storage of the same data type.

class torch.FloatStorage¶

bool()¶: Casts this storage to bool type

byte()¶: Casts this storage to byte type

char()¶: Casts this storage to char type

clone()¶: Returns a copy of this storage

copy_()¶

cpu()¶: Returns a CPU copy of this storage if it’s not already on the CPU

cuda(device=None, non_blocking=False, **kwargs)¶

Returns a copy of this object in CUDA memory.

If this object is already in CUDA memory and on the correct device, then no copy is performed and the original object is returned.

Parameters

device (int) – The destination GPU id. Defaults to the current device.
non_blocking (bool) – If True and the source is in pinned memory, the copy will be asynchronous with respect to the host. Otherwise, the argument has no effect.
**kwargs – For compatibility, may contain the key async in place of the non_blocking argument.

data_ptr()¶

double()¶: Casts this storage to double type

element_size()¶

fill_()¶

float()¶: Casts this storage to float type

static from_buffer()¶

static from_file(filename, shared=False, size=0) → Storage¶

If shared is True, then memory is shared between all processes. All changes are written to the file. If shared is False, then the changes on the storage do not affect the file.

size is the number of elements in the storage. If shared is False, then the file must contain at least size * sizeof(Type) bytes (Type is the type of storage). If shared is True the file will be created if needed.

Parameters

filename (str) – file name to map
shared (bool) – whether to share memory
size (int) – number of elements in the storage

half()¶: Casts this storage to half type

int()¶: Casts this storage to int type

is_cuda = False¶

is_pinned()¶

is_shared()¶

is_sparse = False¶

long()¶: Casts this storage to long type

new()¶

pin_memory()¶: Copies the storage to pinned memory, if it’s not already pinned.

resize_()¶

share_memory_()¶

Moves the storage to shared memory.

This is a no-op for storages already in shared memory and for CUDA storages, which do not need to be moved for sharing across processes. Storages in shared memory cannot be resized.

Returns: self

short()¶: Casts this storage to short type

size()¶

tolist()¶: Returns a list containing the elements of this storage

type(dtype=None, non_blocking=False, **kwargs)¶

Returns the type if dtype is not provided, else casts this object to the specified type.

If this is already of the correct type, no copy is performed and the original object is returned.

Parameters

dtype (type or string) – The desired type
non_blocking (bool) – If True, and the source is in pinned memory and destination is on the GPU or vice versa, the copy is performed asynchronously with respect to the host. Otherwise, the argument has no effect.
**kwargs – For compatibility, may contain the key async in place of the non_blocking argument. The async arg is deprecated.

torch.nn¶

Parameters¶

class torch.nn.Parameter¶

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.

Parameters

data (Tensor) – parameter tensor.
requires_grad (bool, optional) – if the parameter requires gradient. See excluding-subgraphs for more details. Default: True

Containers¶

Module¶

class torch.nn.Module¶

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:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
       x = F.relu(self.conv1(x))
       return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

add_module(name, module)¶

Adds a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters

name (string) – name of the child module. The child module can be accessed from this module using the given name
parameter (Module) – child module to be added to the module.

apply(fn)¶

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).

Parameters: fn (Module -> None) – function to be applied to each submodule
Returns: self
Return type: Module

Example:

>>> def init_weights(m):
        print(m)
        if type(m) == nn.Linear:
            m.weight.data.fill_(1.0)
            print(m.weight)

>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

buffers(recurse=True)¶

Returns an iterator over module buffers.

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

Example:

>>> for buf in model.buffers():
>>>     print(type(buf.data), buf.size())
<class 'torch.FloatTensor'> (20L,)
<class 'torch.FloatTensor'> (20L, 1L, 5L, 5L)

children()¶

Returns an iterator over immediate children modules.

Yields: Module – a child module

cpu()¶

Moves all model parameters and buffers to the CPU.

Returns: self
Return type: Module

cuda(device=None)¶

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.

Parameters: device (int, optional) – if specified, all parameters will be copied to that device
Returns: self
Return type: Module

double()¶

Casts all floating point parameters and buffers to double datatype.

Returns: self
Return type: Module

dump_patches = False¶

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.

eval()¶

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.

extra_repr()¶

Set the extra representation of the module

To print customized extra information, you should reimplement this method in your own modules. Both single-line and multi-line strings are acceptable.

float()¶

Casts all floating point parameters and buffers to float datatype.

Returns: self
Return type: Module

forward(*input)¶

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.

half()¶

Casts all floating point parameters and buffers to half datatype.

Returns: self
Return type: Module

load_state_dict(state_dict, strict=True)¶

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.

Parameters

state_dict (dict) – a dict containing parameters and persistent buffers.
strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

modules()¶

Returns an iterator over all modules in the network.

Yields: Module – a module in the network

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
        print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

named_buffers(prefix='', recurse=True)¶

Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters

prefix (str) – prefix to prepend to all buffer names.
recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields

(string, torch.Tensor) – Tuple containing the name and buffer

Example:

>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())

named_children()¶

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields: (string, Module) – Tuple containing a name and child module

Example:

>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo=None, prefix='')¶

Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Yields: (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:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True)¶

Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters

prefix (str) – prefix to prepend to all parameter names.
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields

(string, Parameter) – Tuple containing the name and parameter

Example:

>>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())

parameters(recurse=True)¶

Returns an iterator over module parameters.

This is typically passed to an optimizer.

Parameters: recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
Yields: Parameter – module parameter

Example:

>>> for param in model.parameters():
>>>     print(type(param.data), param.size())
<class 'torch.FloatTensor'> (20L,)
<class 'torch.FloatTensor'> (20L, 1L, 5L, 5L)

register_backward_hook(hook)¶

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:

hook(module, grad_input, grad_output) -> Tensor or None

The grad_input and grad_output may be tuples if the module has multiple inputs or outputs. The hook should not modify its arguments, but it can optionally return a new gradient with respect to input that will be used in place of grad_input in subsequent computations.

Returns: a handle that can be used to remove the added hook by calling handle.remove()
Return type: torch.utils.hooks.RemovableHandle

Warning

The current implementation will not have the presented behavior for complex Module that perform many operations. In some failure cases, grad_input and grad_output will only contain the gradients for a subset of the inputs and outputs. For such Module, you should use torch.Tensor.register_hook() directly on a specific input or output to get the required gradients.

register_buffer(name, tensor)¶

Adds a persistent 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 persistent state.

Buffers can be accessed as attributes using given names.

Parameters

name (string) – name of the buffer. The buffer can be accessed from this module using the given name
tensor (Tensor) – buffer to be registered.

Example:

>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_forward_hook(hook)¶

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:

hook(module, input, output) -> None

The hook should not modify the input or output.

Returns: a handle that can be used to remove the added hook by calling handle.remove()
Return type: torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook)¶

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:

hook(module, input) -> None

The hook should not modify the input.

Returns: a handle that can be used to remove the added hook by calling handle.remove()
Return type: torch.utils.hooks.RemovableHandle

register_parameter(name, param)¶

Adds a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters

name (string) – name of the parameter. The parameter can be accessed from this module using the given name
parameter (Parameter) – parameter to be added to the module.

state_dict(destination=None, prefix='', keep_vars=False)¶

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.

Returns: a dictionary containing a whole state of the module
Return type: dict

Example:

>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)¶

Moves and/or casts the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

Its signature is similar to torch.Tensor.to(), but only accepts floating point desired dtype s. In addition, this method will only cast the floating point parameters and buffers to 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.

Parameters

device (torch.device) – the desired device of the parameters and buffers in this module
dtype (torch.dtype) – the desired floating point type of the floating point parameters and buffers in this module
tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

Returns

self

Return type

Module

Example:

>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

train(mode=True)¶

Sets the module in training mode.

Returns: self
Return type: Module

type(dst_type)¶

Casts all parameters and buffers to dst_type.

Parameters: dst_type (type or string) – the desired type
Returns: self
Return type: Module

zero_grad()¶: Sets gradients of all model parameters to zero.

Sequential¶

class torch.nn.Sequential(*args)¶

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:

# Example of using Sequential
model = nn.Sequential(
          nn.Conv2d(1,20,5),
          nn.ReLU(),
          nn.Conv2d(20,64,5),
          nn.ReLU()
        )

# Example of using Sequential with OrderedDict
model = nn.Sequential(OrderedDict([
          ('conv1', nn.Conv2d(1,20,5)),
          ('relu1', nn.ReLU()),
          ('conv2', nn.Conv2d(20,64,5)),
          ('relu2', nn.ReLU())
        ]))

Shape:

Input: $(*)$ where * means, any number of additional dimensions
Output: $(*)$, same shape as the input

ModuleList¶

class torch.nn.ModuleList(modules=None)¶

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.

Parameters: modules (iterable, optional) – an iterable of modules to add

Example:

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])

    def forward(self, x):
        # ModuleList can act as an iterable, or be indexed using ints
        for i, l in enumerate(self.linears):
            x = self.linears[i // 2](x) + l(x)
        return x

append(module)¶

Appends a given module to the end of the list.

Parameters: module (nn.Module) – module to append

extend(modules)¶

Appends modules from a Python iterable to the end of the list.

Parameters: modules (iterable) – iterable of modules to append

insert(index, module)¶

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

Parameters

index (int) – index to insert.
module (nn.Module) – module to insert

ModuleDict¶

class torch.nn.ModuleDict(modules=None)¶

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

the order of insertion, and
in update(), the order of the merged OrderedDict or another ModuleDict (the argument to update()).

Note that update() with other unordered mapping types (e.g., Python’s plain dict) does not preserve the order of the merged mapping.

Parameters: modules (iterable, optional) – a mapping (dictionary) of (string: module) or an iterable of key-value pairs of type (string, module)

Example:

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.choices = nn.ModuleDict({
                'conv': nn.Conv2d(10, 10, 3),
                'pool': nn.MaxPool2d(3)
        })
        self.activations = nn.ModuleDict([
                ['lrelu', nn.LeakyReLU()],
                ['prelu', nn.PReLU()]
        ])

    def forward(self, x, choice, act):
        x = self.choices[choice](x)
        x = self.activations[act](x)
        return x

clear()¶: Remove all items from the ModuleDict.

items()¶: Return an iterable of the ModuleDict key/value pairs.

keys()¶: Return an iterable of the ModuleDict keys.

pop(key)¶

Remove key from the ModuleDict and return its module.

Parameters: key (string) – key to pop from the ModuleDict

update(modules)¶

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.

Parameters: modules (iterable) – a mapping (dictionary) from string to Module, or an iterable of key-value pairs of type (string, Module)

values()¶: Return an iterable of the ModuleDict values.

ParameterList¶

class torch.nn.ParameterList(parameters=None)¶

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: parameters (iterable, optional) – an iterable of Parameter to add

Example:

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.params = nn.ParameterList([nn.Parameter(torch.randn(10, 10)) for i in range(10)])

    def forward(self, x):
        # ParameterList can act as an iterable, or be indexed using ints
        for i, p in enumerate(self.params):
            x = self.params[i // 2].mm(x) + p.mm(x)
        return x

append(parameter)¶

Appends a given parameter at the end of the list.

Parameters: parameter (nn.Parameter) – parameter to append

extend(parameters)¶

Appends parameters from a Python iterable to the end of the list.

Parameters: parameters (iterable) – iterable of parameters to append

ParameterDict¶

class torch.nn.ParameterDict(parameters=None)¶

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

the order of insertion, and
in update(), the order of the merged OrderedDict or another ParameterDict (the argument to update()).

Note that update() with other unordered mapping types (e.g., Python’s plain dict) does not preserve the order of the merged mapping.

Parameters: parameters (iterable, optional) – a mapping (dictionary) of (string : Parameter) or an iterable of key-value pairs of type (string, Parameter)

Example:

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.params = nn.ParameterDict({
                'left': nn.Parameter(torch.randn(5, 10)),
                'right': nn.Parameter(torch.randn(5, 10))
        })

    def forward(self, x, choice):
        x = self.params[choice].mm(x)
        return x

clear()¶: Remove all items from the ParameterDict.

items()¶: Return an iterable of the ParameterDict key/value pairs.

keys()¶: Return an iterable of the ParameterDict keys.

pop(key)¶

Remove key from the ParameterDict and return its parameter.

Parameters: key (string) – key to pop from the ParameterDict

update(parameters)¶

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: parameters (iterable) – a mapping (dictionary) from string to Parameter, or an iterable of key-value pairs of type (string, Parameter)

values()¶: Return an iterable of the ParameterDict values.

Convolution layers¶

Conv1d¶

class torch.nn.Conv1d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')¶

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

In the simplest case, the output value of the layer with input size $(N, C_{\text{in}}, L)$ and output $(N, C_{\text{out}}, L_{\text{out}})$ can be precisely described as:

\[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k = 0}^{C_{in} - 1} \text{weight}(C_{\text{out}_j}, k) \star \text{input}(N_i, k) \]

where $\star$ is the valid cross-correlation operator, $N$ is a batch size, $C$ denotes a number of channels, $L$ is a length of signal sequence.

stride controls the stride for the cross-correlation, a single number or a one-element tuple.
padding controls the amount of implicit zero-paddings on both sides for padding number of points.
dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.
groups controls the connections between inputs and outputs. in_channels and out_channels must both be divisible by groups. For example,
- At groups=1, all inputs are convolved to all outputs.
- At groups=2, the operation becomes equivalent to having two conv layers side by side, each seeing half the input channels, and producing half the output channels, and both subsequently concatenated.
- At groups= in_channels, each input channel is convolved with its own set of filters, of size $\left\lfloor\frac{out\_channels}{in\_channels}\right\rfloor$.

Note

Depending of the size of your kernel, several (of the last) columns of the input might be lost, because it is a valid cross-correlation, and not a full cross-correlation. It is up to the user to add proper padding.

Note

When groups == in_channels and out_channels == K * in_channels, where K is a positive integer, this operation is also termed in literature as depthwise convolution.

In other words, for an input of size $(N, C_{in}, L_{in})$, a depthwise convolution with a depthwise multiplier K, can be constructed by arguments $(C_\text{in}=C_{in}, C_\text{out}=C_{in} \times K, ..., \text{groups}=C_{in})$.

Note

In some circumstances when using the CUDA backend with CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. Please see the notes on /notes/randomness for background.

Parameters

in_channels (int) – Number of channels in the input image
out_channels (int) – Number of channels produced by the convolution
kernel_size (int or tuple) – Size of the convolving kernel
stride (int or tuple, optional) – Stride of the convolution. Default: 1
padding (int or tuple, optional) – Zero-padding added to both sides of the input. Default: 0
padding_mode (string, optional) – zeros
dilation (int or tuple, optional) – Spacing between kernel elements. Default: 1
groups (int, optional) – Number of blocked connections from input channels to output channels. Default: 1
bias (bool, optional) – If True, adds a learnable bias to the output. Default: True

Shape:

Input: $(N, C_{in}, L_{in})$
Output: $(N, C_{out}, L_{out})$ where

\[L_{out} = \left\lfloor\frac{L_{in} + 2 \times \text{padding} - \text{dilation} \times (\text{kernel\_size} - 1) - 1}{\text{stride}} + 1\right\rfloor \]

Variables

~Conv1d.weight (Tensor) – the learnable weights of the module of shape $(\text{out\_channels}, \frac{\text{in\_channels}}{\text{groups}}, \text{kernel\_size})$. The values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \text{kernel\_size}}$
~Conv1d.bias (Tensor) – the learnable bias of the module of shape (out_channels). If bias is True, then the values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \text{kernel\_size}}$

Examples:

>>> m = nn.Conv1d(16, 33, 3, stride=2)
>>> input = torch.randn(20, 16, 50)
>>> output = m(input)

Conv2d¶

class torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')¶

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

In the simplest case, the output value of the layer with input size $(N, C_{\text{in}}, H, W)$ and output $(N, C_{\text{out}}, H_{\text{out}}, W_{\text{out}})$ can be precisely described as:

\[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k = 0}^{C_{\text{in}} - 1} \text{weight}(C_{\text{out}_j}, k) \star \text{input}(N_i, k) \]

where $\star$ is the valid 2D cross-correlation operator, $N$ is a batch size, $C$ denotes a number of channels, $H$ is a height of input planes in pixels, and $W$ is width in pixels.

stride controls the stride for the cross-correlation, a single number or a tuple.
padding controls the amount of implicit zero-paddings on both sides for padding number of points for each dimension.
dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.
groups controls the connections between inputs and outputs. in_channels and out_channels must both be divisible by groups. For example,
- At groups=1, all inputs are convolved to all outputs.
- At groups=2, the operation becomes equivalent to having two conv layers side by side, each seeing half the input channels, and producing half the output channels, and both subsequently concatenated.
- At groups= in_channels, each input channel is convolved with its own set of filters, of size: $\left\lfloor\frac{out\_channels}{in\_channels}\right\rfloor$.

The parameters kernel_size, stride, padding, dilation can either be:

a single int – in which case the same value is used for the height and width dimension

a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension

Note

When groups == in_channels and out_channels == K * in_channels, where K is a positive integer, this operation is also termed in literature as depthwise convolution.

In other words, for an input of size $(N, C_{in}, H_{in}, W_{in})$, a depthwise convolution with a depthwise multiplier K, can be constructed by arguments $(in\_channels=C_{in}, out\_channels=C_{in} \times K, ..., groups=C_{in})$.

Note

Parameters

in_channels (int) – Number of channels in the input image
out_channels (int) – Number of channels produced by the convolution
kernel_size (int or tuple) – Size of the convolving kernel
stride (int or tuple, optional) – Stride of the convolution. Default: 1
padding (int or tuple, optional) – Zero-padding added to both sides of the input. Default: 0
padding_mode (string, optional) – zeros
dilation (int or tuple, optional) – Spacing between kernel elements. Default: 1
groups (int, optional) – Number of blocked connections from input channels to output channels. Default: 1
bias (bool, optional) – If True, adds a learnable bias to the output. Default: True

Shape:

Input: $(N, C_{in}, H_{in}, W_{in})$
Output: $(N, C_{out}, H_{out}, W_{out})$ where

\[H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor \]

\[W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor \]

Variables

~Conv2d.weight (Tensor) –
the learnable weights of the module of shape :math:`(text{out_channels}, frac{text{in_channels}}{text{groups}},

text{kernel_size[0]}, text{kernel_size[1]})`.

The values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}$
~Conv2d.bias (Tensor) – the learnable bias of the module of shape (out_channels). If bias is True, then the values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}$

Examples:

>>> # With square kernels and equal stride
>>> m = nn.Conv2d(16, 33, 3, stride=2)
>>> # non-square kernels and unequal stride and with padding
>>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
>>> # non-square kernels and unequal stride and with padding and dilation
>>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(3, 1))
>>> input = torch.randn(20, 16, 50, 100)
>>> output = m(input)

Conv3d¶

class torch.nn.Conv3d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')¶

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

In the simplest case, the output value of the layer with input size $(N, C_{in}, D, H, W)$ and output $(N, C_{out}, D_{out}, H_{out}, W_{out})$ can be precisely described as:

\[out(N_i, C_{out_j}) = bias(C_{out_j}) + \sum_{k = 0}^{C_{in} - 1} weight(C_{out_j}, k) \star input(N_i, k) \]

where $\star$ is the valid 3D cross-correlation operator

stride controls the stride for the cross-correlation.
padding controls the amount of implicit zero-paddings on both sides for padding number of points for each dimension.
dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.
groups controls the connections between inputs and outputs. in_channels and out_channels must both be divisible by groups. For example,
- At groups=1, all inputs are convolved to all outputs.
- At groups=2, the operation becomes equivalent to having two conv layers side by side, each seeing half the input channels, and producing half the output channels, and both subsequently concatenated.
- At groups= in_channels, each input channel is convolved with its own set of filters, of size $\left\lfloor\frac{out\_channels}{in\_channels}\right\rfloor$.

The parameters kernel_size, stride, padding, dilation can either be:

a single int – in which case the same value is used for the depth, height and width dimension

a tuple of three ints – in which case, the first int is used for the depth dimension, the second int for the height dimension and the third int for the width dimension

Note

When groups == in_channels and out_channels == K * in_channels, where K is a positive integer, this operation is also termed in literature as depthwise convolution.

In other words, for an input of size $(N, C_{in}, D_{in}, H_{in}, W_{in})$, a depthwise convolution with a depthwise multiplier K, can be constructed by arguments $(in\_channels=C_{in}, out\_channels=C_{in} \times K, ..., groups=C_{in})$.

Note

Parameters

in_channels (int) – Number of channels in the input image
out_channels (int) – Number of channels produced by the convolution
kernel_size (int or tuple) – Size of the convolving kernel
stride (int or tuple, optional) – Stride of the convolution. Default: 1
padding (int or tuple, optional) – Zero-padding added to all three sides of the input. Default: 0
padding_mode (string, optional) – zeros
dilation (int or tuple, optional) – Spacing between kernel elements. Default: 1
groups (int, optional) – Number of blocked connections from input channels to output channels. Default: 1
bias (bool, optional) – If True, adds a learnable bias to the output. Default: True

Shape:

Input: $(N, C_{in}, D_{in}, H_{in}, W_{in})$
Output: $(N, C_{out}, D_{out}, H_{out}, W_{out})$ where

\[D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor \]

\[H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor \]

\[W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] - \text{dilation}[2] \times (\text{kernel\_size}[2] - 1) - 1}{\text{stride}[2]} + 1\right\rfloor \]

Variables

~Conv3d.weight (Tensor) –
the learnable weights of the module of shape :math:`(text{out_channels}, frac{text{in_channels}}{text{groups}},

text{kernel_size[0]}, text{kernel_size[1]}, text{kernel_size[2]})`.

The values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \prod_{i=0}^{2}\text{kernel\_size}[i]}$
~Conv3d.bias (Tensor) – the learnable bias of the module of shape (out_channels). If bias is True, then the values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \prod_{i=0}^{2}\text{kernel\_size}[i]}$

Examples:

>>> # With square kernels and equal stride
>>> m = nn.Conv3d(16, 33, 3, stride=2)
>>> # non-square kernels and unequal stride and with padding
>>> m = nn.Conv3d(16, 33, (3, 5, 2), stride=(2, 1, 1), padding=(4, 2, 0))
>>> input = torch.randn(20, 16, 10, 50, 100)
>>> output = m(input)

ConvTranspose1d¶

class torch.nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros')¶

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

This module can be seen as the gradient of Conv1d with respect to its input. It is also known as a fractionally-strided convolution or a deconvolution (although it is not an actual deconvolution operation).

stride controls the stride for the cross-correlation.
padding controls the amount of implicit zero-paddings on both sides for dilation * (kernel_size - 1) - padding number of points. See note below for details.
output_padding controls the additional size added to one side of the output shape. See note below for details.
dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.
groups controls the connections between inputs and outputs. in_channels and out_channels must both be divisible by groups. For example,
- At groups=1, all inputs are convolved to all outputs.
- At groups=2, the operation becomes equivalent to having two conv layers side by side, each seeing half the input channels, and producing half the output channels, and both subsequently concatenated.
- At groups= in_channels, each input channel is convolved with its own set of filters (of size $\left\lfloor\frac{out\_channels}{in\_channels}\right\rfloor$).

Note

The padding argument effectively adds dilation * (kernel_size - 1) - padding amount of zero padding to both sizes of the input. This is set so that when a Conv1d and a ConvTranspose1d are initialized with same parameters, they are inverses of each other in regard to the input and output shapes. However, when stride > 1, Conv1d maps multiple input shapes to the same output shape. output_padding is provided to resolve this ambiguity by effectively increasing the calculated output shape on one side. Note that output_padding is only used to find output shape, but does not actually add zero-padding to output.

Note

Parameters

in_channels (int) – Number of channels in the input image
out_channels (int) – Number of channels produced by the convolution
kernel_size (int or tuple) – Size of the convolving kernel
stride (int or tuple, optional) – Stride of the convolution. Default: 1
padding (int or tuple, optional) – dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of the input. Default: 0
output_padding (int or tuple, optional) – Additional size added to one side of the output shape. Default: 0
groups (int, optional) – Number of blocked connections from input channels to output channels. Default: 1
bias (bool, optional) – If True, adds a learnable bias to the output. Default: True
dilation (int or tuple, optional) – Spacing between kernel elements. Default: 1

Shape:

Input: $(N, C_{in}, L_{in})$
Output: $(N, C_{out}, L_{out})$ where

\[L_{out} = (L_{in} - 1) \times \text{stride} - 2 \times \text{padding} + \text{dilation} \times (\text{kernel\_size} - 1) + \text{output\_padding} + 1 \]

Variables

~ConvTranspose1d.weight (Tensor) –
the learnable weights of the module of shape :math:`(text{in_channels}, frac{text{out_channels}}{text{groups}},

text{kernel_size})`. The values of these weights are sampled from

$\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \text{kernel\_size}}$
~ConvTranspose1d.bias (Tensor) – the learnable bias of the module of shape (out_channels). If bias is True, then the values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \text{kernel\_size}}$

ConvTranspose2d¶

class torch.nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros')¶

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

This module can be seen as the gradient of Conv2d with respect to its input. It is also known as a fractionally-strided convolution or a deconvolution (although it is not an actual deconvolution operation).

stride controls the stride for the cross-correlation.
padding controls the amount of implicit zero-paddings on both sides for dilation * (kernel_size - 1) - padding number of points. See note below for details.
output_padding controls the additional size added to one side of the output shape. See note below for details.
dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.
groups controls the connections between inputs and outputs. in_channels and out_channels must both be divisible by groups. For example,
- At groups=1, all inputs are convolved to all outputs.
- At groups=2, the operation becomes equivalent to having two conv layers side by side, each seeing half the input channels, and producing half the output channels, and both subsequently concatenated.
- At groups= in_channels, each input channel is convolved with its own set of filters (of size $\left\lfloor\frac{out\_channels}{in\_channels}\right\rfloor$).

The parameters kernel_size, stride, padding, output_padding can either be:

a single int – in which case the same value is used for the height and width dimensions

a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension

Note

The padding argument effectively adds dilation * (kernel_size - 1) - padding amount of zero padding to both sizes of the input. This is set so that when a Conv2d and a ConvTranspose2d are initialized with same parameters, they are inverses of each other in regard to the input and output shapes. However, when stride > 1, Conv2d maps multiple input shapes to the same output shape. output_padding is provided to resolve this ambiguity by effectively increasing the calculated output shape on one side. Note that output_padding is only used to find output shape, but does not actually add zero-padding to output.

Note

Parameters

in_channels (int) – Number of channels in the input image
out_channels (int) – Number of channels produced by the convolution
kernel_size (int or tuple) – Size of the convolving kernel
stride (int or tuple, optional) – Stride of the convolution. Default: 1
padding (int or tuple, optional) – dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each dimension in the input. Default: 0
output_padding (int or tuple, optional) – Additional size added to one side of each dimension in the output shape. Default: 0
groups (int, optional) – Number of blocked connections from input channels to output channels. Default: 1
bias (bool, optional) – If True, adds a learnable bias to the output. Default: True
dilation (int or tuple, optional) – Spacing between kernel elements. Default: 1

Shape:

Input: $(N, C_{in}, H_{in}, W_{in})$
Output: $(N, C_{out}, H_{out}, W_{out})$ where

\[H_{out} = (H_{in} - 1) \times \text{stride}[0] - 2 \times \text{padding}[0] + \text{dilation}[0] \times (\text{kernel\_size}[0] - 1) + \text{output\_padding}[0] + 1 \]

\[W_{out} = (W_{in} - 1) \times \text{stride}[1] - 2 \times \text{padding}[1] + \text{dilation}[1] \times (\text{kernel\_size}[1] - 1) + \text{output\_padding}[1] + 1 \]

Variables

~ConvTranspose2d.weight (Tensor) –
the learnable weights of the module of shape :math:`(text{in_channels}, frac{text{out_channels}}{text{groups}},

text{kernel_size[0]}, text{kernel_size[1]})`.

The values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}$
~ConvTranspose2d.bias (Tensor) – the learnable bias of the module of shape (out_channels) If bias is True, then the values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}$

Examples:

>>> # With square kernels and equal stride
>>> m = nn.ConvTranspose2d(16, 33, 3, stride=2)
>>> # non-square kernels and unequal stride and with padding
>>> m = nn.ConvTranspose2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
>>> input = torch.randn(20, 16, 50, 100)
>>> output = m(input)
>>> # exact output size can be also specified as an argument
>>> input = torch.randn(1, 16, 12, 12)
>>> downsample = nn.Conv2d(16, 16, 3, stride=2, padding=1)
>>> upsample = nn.ConvTranspose2d(16, 16, 3, stride=2, padding=1)
>>> h = downsample(input)
>>> h.size()
torch.Size([1, 16, 6, 6])
>>> output = upsample(h, output_size=input.size())
>>> output.size()
torch.Size([1, 16, 12, 12])

ConvTranspose3d¶

class torch.nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros')¶

Applies a 3D transposed convolution operator over an input image composed of several input planes. The transposed convolution operator multiplies each input value element-wise by a learnable kernel, and sums over the outputs from all input feature planes.

This module can be seen as the gradient of Conv3d with respect to its input. It is also known as a fractionally-strided convolution or a deconvolution (although it is not an actual deconvolution operation).

stride controls the stride for the cross-correlation.
padding controls the amount of implicit zero-paddings on both sides for dilation * (kernel_size - 1) - padding number of points. See note below for details.
output_padding controls the additional size added to one side of the output shape. See note below for details.
dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.
groups controls the connections between inputs and outputs. in_channels and out_channels must both be divisible by groups. For example,
- At groups=1, all inputs are convolved to all outputs.
- At groups=2, the operation becomes equivalent to having two conv layers side by side, each seeing half the input channels, and producing half the output channels, and both subsequently concatenated.
- At groups= in_channels, each input channel is convolved with its own set of filters (of size $\left\lfloor\frac{out\_channels}{in\_channels}\right\rfloor$).

The parameters kernel_size, stride, padding, output_padding can either be:

a single int – in which case the same value is used for the depth, height and width dimensions

a tuple of three ints – in which case, the first int is used for the depth dimension, the second int for the height dimension and the third int for the width dimension

Note

The padding argument effectively adds dilation * (kernel_size - 1) - padding amount of zero padding to both sizes of the input. This is set so that when a Conv3d and a ConvTranspose3d are initialized with same parameters, they are inverses of each other in regard to the input and output shapes. However, when stride > 1, Conv3d maps multiple input shapes to the same output shape. output_padding is provided to resolve this ambiguity by effectively increasing the calculated output shape on one side. Note that output_padding is only used to find output shape, but does not actually add zero-padding to output.

Note

Parameters

in_channels (int) – Number of channels in the input image
out_channels (int) – Number of channels produced by the convolution
kernel_size (int or tuple) – Size of the convolving kernel
stride (int or tuple, optional) – Stride of the convolution. Default: 1
padding (int or tuple, optional) – dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each dimension in the input. Default: 0
output_padding (int or tuple, optional) – Additional size added to one side of each dimension in the output shape. Default: 0
groups (int, optional) – Number of blocked connections from input channels to output channels. Default: 1
bias (bool, optional) – If True, adds a learnable bias to the output. Default: True
dilation (int or tuple, optional) – Spacing between kernel elements. Default: 1

Shape:

Input: $(N, C_{in}, D_{in}, H_{in}, W_{in})$
Output: $(N, C_{out}, D_{out}, H_{out}, W_{out})$ where

\[D_{out} = (D_{in} - 1) \times \text{stride}[0] - 2 \times \text{padding}[0] + \text{dilation}[0] \times (\text{kernel\_size}[0] - 1) + \text{output\_padding}[0] + 1 \]

\[H_{out} = (H_{in} - 1) \times \text{stride}[1] - 2 \times \text{padding}[1] + \text{dilation}[1] \times (\text{kernel\_size}[1] - 1) + \text{output\_padding}[1] + 1 \]

\[W_{out} = (W_{in} - 1) \times \text{stride}[2] - 2 \times \text{padding}[2] + \text{dilation}[2] \times (\text{kernel\_size}[2] - 1) + \text{output\_padding}[2] + 1 \]

Variables

~ConvTranspose3d.weight (Tensor) –
the learnable weights of the module of shape :math:`(text{in_channels}, frac{text{out_channels}}{text{groups}},

text{kernel_size[0]}, text{kernel_size[1]}, text{kernel_size[2]})`.

The values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \prod_{i=0}^{2}\text{kernel\_size}[i]}$
~ConvTranspose3d.bias (Tensor) – the learnable bias of the module of shape (out_channels) If bias is True, then the values of these weights are sampled from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{C_\text{in} * \prod_{i=0}^{2}\text{kernel\_size}[i]}$

Examples:

>>> # With square kernels and equal stride
>>> m = nn.ConvTranspose3d(16, 33, 3, stride=2)
>>> # non-square kernels and unequal stride and with padding
>>> m = nn.ConvTranspose3d(16, 33, (3, 5, 2), stride=(2, 1, 1), padding=(0, 4, 2))
>>> input = torch.randn(20, 16, 10, 50, 100)
>>> output = m(input)

Unfold¶

class torch.nn.Unfold(kernel_size, dilation=1, padding=0, stride=1)¶

Extracts sliding local blocks from a batched input tensor.

Consider an batched input tensor of shape $(N, C, *)$, where $N$ is the batch dimension, $C$ is the channel dimension, and $*$ represent arbitrary spatial dimensions. This operation flattens each sliding kernel_size-sized block within the spatial dimensions of input into a column (i.e., last dimension) of a 3-D output tensor of shape $(N, C \times \prod(\text{kernel\_size}), L)$, where $C \times \prod(\text{kernel\_size})$ is the total number of values with in each block (a block has $\prod(\text{kernel\_size})$ spatial locations each containing a $C$-channeled vector), and $L$ is the total number of such blocks:

\[L = \prod_d \left\lfloor\frac{\text{spatial\_size}[d] + 2 \times \text{padding}[d] % - \text{dilation}[d] \times (\text{kernel\_size}[d] - 1) - 1}{\text{stride}[d]} + 1\right\rfloor, \]

where $\text{spatial\_size}$ is formed by the spatial dimensions of input ($*$ above), and $d$ is over all spatial dimensions.

Therefore, indexing output at the last dimension (column dimension) gives all values within a certain block.

The padding, stride and dilation arguments specify how the sliding blocks are retrieved.

stride controls the stride for the sliding blocks.
padding controls the amount of implicit zero-paddings on both sides for padding number of points for each dimension before reshaping.
dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.

Parameters

kernel_size (int or tuple) – the size of the sliding blocks
stride (int or tuple, optional) – the stride of the sliding blocks in the input spatial dimensions. Default: 1
padding (int or tuple, optional) – implicit zero padding to be added on both sides of input. Default: 0
dilation (int or tuple, optional) – a parameter that controls the stride of elements within the neighborhood. Default: 1

If kernel_size, dilation, padding or stride is an int or a tuple of length 1, their values will be replicated across all spatial dimensions.
For the case of two input spatial dimensions this operation is sometimes called im2col.

Note

Fold calculates each combined value in the resulting large tensor by summing all values from all containing blocks. Unfold extracts the values in the local blocks by copying from the large tensor. So, if the blocks overlap, they are not inverses of each other.

Warning

Currently, only 4-D input tensors (batched image-like tensors) are supported.

Shape:

Input: $(N, C, *)$
Output: $(N, C \times \prod(\text{kernel\_size}), L)$ as described above

Examples:

>>> unfold = nn.Unfold(kernel_size=(2, 3))
>>> input = torch.randn(2, 5, 3, 4)
>>> output = unfold(input)
>>> # each patch contains 30 values (2x3=6 vectors, each of 5 channels)
>>> # 4 blocks (2x3 kernels) in total in the 3x4 input
>>> output.size()
torch.Size([2, 30, 4])

>>> # Convolution is equivalent with Unfold + Matrix Multiplication + Fold (or view to output shape)
>>> inp = torch.randn(1, 3, 10, 12)
>>> w = torch.randn(2, 3, 4, 5)
>>> inp_unf = torch.nn.functional.unfold(inp, (4, 5))
>>> out_unf = inp_unf.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
>>> out = torch.nn.functional.fold(out_unf, (7, 8), (1, 1))
>>> # or equivalently (and avoiding a copy),
>>> # out = out_unf.view(1, 2, 7, 8)
>>> (torch.nn.functional.conv2d(inp, w) - out).abs().max()
tensor(1.9073e-06)

Fold¶

class torch.nn.Fold(output_size, kernel_size, dilation=1, padding=0, stride=1)¶

Combines an array of sliding local blocks into a large containing tensor.

Consider a batched input tensor containing sliding local blocks, e.g., patches of images, of shape $(N, C \times \prod(\text{kernel\_size}), L)$, where $N$ is batch dimension, $C \times \prod(\text{kernel\_size})$ is the number of values with in a block (a block has $\prod(\text{kernel\_size})$ spatial locations each containing a $C$-channeled vector), and $L$ is the total number of blocks. (This is exacly the same specification as the output shape of Unfold.) This operation combines these local blocks into the large output tensor of shape $(N, C, \text{output\_size}[0], \text{output\_size}[1], \dots)$ by summing the overlapping values. Similar to Unfold, the arguments must satisfy

\[L = \prod_d \left\lfloor\frac{\text{output\_size}[d] + 2 \times \text{padding}[d] % - \text{dilation}[d] \times (\text{kernel\_size}[d] - 1) - 1}{\text{stride}[d]} + 1\right\rfloor, \]

where $d$ is over all spatial dimensions.

output_size describes the spatial shape of the large containing tensor of the sliding local blocks. It is useful to resolve the ambiguity when multiple input shapes map to same number of sliding blocks, e.g., with stride > 0.

The padding, stride and dilation arguments specify how the sliding blocks are retrieved.

stride controls the stride for the sliding blocks.
padding controls the amount of implicit zero-paddings on both sides for padding number of points for each dimension before reshaping.
dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.

Parameters

output_size (int or tuple) – the shape of the spatial dimensions of the output (i.e., output.sizes()[2:])
kernel_size (int or tuple) – the size of the sliding blocks
stride (int or tuple) – the stride of the sliding blocks in the input spatial dimensions. Default: 1
padding (int or tuple, optional) – implicit zero padding to be added on both sides of input. Default: 0
dilation (int or tuple, optional) – a parameter that controls the stride of elements within the neighborhood. Default: 1

If output_size, kernel_size, dilation, padding or stride is an int or a tuple of length 1 then their values will be replicated across all spatial dimensions.
For the case of two output spatial dimensions this operation is sometimes called col2im.

Note

Warning

Currently, only 4-D output tensors (batched image-like tensors) are supported.

Shape:

Input: $(N, C \times \prod(\text{kernel\_size}), L)$
Output: $(N, C, \text{output\_size}[0], \text{output\_size}[1], \dots)$ as described above

Examples:

>>> fold = nn.Fold(output_size=(4, 5), kernel_size=(2, 2))
>>> input = torch.randn(1, 3 * 2 * 2, 12)
>>> output = fold(input)
>>> output.size()
torch.Size([1, 3, 4, 5])

Pooling layers¶

MaxPool1d¶

class torch.nn.MaxPool1d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)¶

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

In the simplest case, the output value of the layer with input size $(N, C, L)$ and output $(N, C, L_{out})$ can be precisely described as:

\[out(N_i, C_j, k) = \max_{m=0, \ldots, \text{kernel\_size} - 1} input(N_i, C_j, stride \times k + m) \]

If padding is non-zero, then the input is implicitly zero-padded on both sides for padding number of points. dilation controls the spacing between the kernel points. It is harder to describe, but this link has a nice visualization of what dilation does.

Parameters

kernel_size – the size of the window to take a max over
stride – the stride of the window. Default value is kernel_size
padding – implicit zero padding to be added on both sides
dilation – a parameter that controls the stride of elements in the window
return_indices – if True, will return the max indices along with the outputs. Useful for torch.nn.MaxUnpool1d later
ceil_mode – when True, will use ceil instead of floor to compute the output shape

Shape:

Input: $(N, C, L_{in})$
Output: $(N, C, L_{out})$, where

\[L_{out} = \left\lfloor \frac{L_{in} + 2 \times \text{padding} - \text{dilation} \times (\text{kernel\_size} - 1) - 1}{\text{stride}} + 1\right\rfloor \]

Examples:

>>> # pool of size=3, stride=2
>>> m = nn.MaxPool1d(3, stride=2)
>>> input = torch.randn(20, 16, 50)
>>> output = m(input)

MaxPool2d¶

class torch.nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)¶

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

In the simplest case, the output value of the layer with input size $(N, C, H, W)$, output $(N, C, H_{out}, W_{out})$ and kernel_size $(kH, kW)$ can be precisely described as:

\[\begin{aligned} out(N_i, C_j, h, w) ={} & \max_{m=0, \ldots, kH-1} \max_{n=0, \ldots, kW-1} \\ & \text{input}(N_i, C_j, \text{stride[0]} \times h + m, \text{stride[1]} \times w + n) \end{aligned} \]

The parameters kernel_size, stride, padding, dilation can either be:

a single int – in which case the same value is used for the height and width dimension

a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension

Parameters

kernel_size – the size of the window to take a max over
stride – the stride of the window. Default value is kernel_size
padding – implicit zero padding to be added on both sides
dilation – a parameter that controls the stride of elements in the window
return_indices – if True, will return the max indices along with the outputs. Useful for torch.nn.MaxUnpool2d later
ceil_mode – when True, will use ceil instead of floor to compute the output shape

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$, where

\[H_{out} = \left\lfloor\frac{H_{in} + 2 * \text{padding[0]} - \text{dilation[0]} \times (\text{kernel\_size[0]} - 1) - 1}{\text{stride[0]}} + 1\right\rfloor \]

\[W_{out} = \left\lfloor\frac{W_{in} + 2 * \text{padding[1]} - \text{dilation[1]} \times (\text{kernel\_size[1]} - 1) - 1}{\text{stride[1]}} + 1\right\rfloor \]

Examples:

>>> # pool of square window of size=3, stride=2
>>> m = nn.MaxPool2d(3, stride=2)
>>> # pool of non-square window
>>> m = nn.MaxPool2d((3, 2), stride=(2, 1))
>>> input = torch.randn(20, 16, 50, 32)
>>> output = m(input)

MaxPool3d¶

class torch.nn.MaxPool3d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)¶

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

In the simplest case, the output value of the layer with input size $(N, C, D, H, W)$, output $(N, C, D_{out}, H_{out}, W_{out})$ and kernel_size $(kD, kH, kW)$ can be precisely described as:

\[\begin{aligned} \text{out}(N_i, C_j, d, h, w) ={} & \max_{k=0, \ldots, kD-1} \max_{m=0, \ldots, kH-1} \max_{n=0, \ldots, kW-1} \\ & \text{input}(N_i, C_j, \text{stride[0]} \times d + k, \text{stride[1]} \times h + m, \text{stride[2]} \times w + n) \end{aligned} \]

The parameters kernel_size, stride, padding, dilation can either be:

a single int – in which case the same value is used for the depth, height and width dimension

a tuple of three ints – in which case, the first int is used for the depth dimension, the second int for the height dimension and the third int for the width dimension

Parameters

kernel_size – the size of the window to take a max over
stride – the stride of the window. Default value is kernel_size
padding – implicit zero padding to be added on all three sides
dilation – a parameter that controls the stride of elements in the window
return_indices – if True, will return the max indices along with the outputs. Useful for torch.nn.MaxUnpool3d later
ceil_mode – when True, will use ceil instead of floor to compute the output shape

Shape:

Input: $(N, C, D_{in}, H_{in}, W_{in})$
Output: $(N, C, D_{out}, H_{out}, W_{out})$, where

\[D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor \]

\[H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor \]

\[W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] - \text{dilation}[2] \times (\text{kernel\_size}[2] - 1) - 1}{\text{stride}[2]} + 1\right\rfloor \]

Examples:

>>> # pool of square window of size=3, stride=2
>>> m = nn.MaxPool3d(3, stride=2)
>>> # pool of non-square window
>>> m = nn.MaxPool3d((3, 2, 2), stride=(2, 1, 2))
>>> input = torch.randn(20, 16, 50,44, 31)
>>> output = m(input)

MaxUnpool1d¶

class torch.nn.MaxUnpool1d(kernel_size, stride=None, padding=0)¶

Computes a partial inverse of MaxPool1d.

MaxPool1d is not fully invertible, since the non-maximal values are lost.

MaxUnpool1d takes in as input the output of MaxPool1d including the indices of the maximal values and computes a partial inverse in which all non-maximal values are set to zero.

Note

MaxPool1d can map several input sizes to the same output sizes. Hence, the inversion process can get ambiguous. To accommodate this, you can provide the needed output size as an additional argument output_size in the forward call. See the Inputs and Example below.

Parameters

kernel_size (int or tuple) – Size of the max pooling window.
stride (int or tuple) – Stride of the max pooling window. It is set to kernel_size by default.
padding (int or tuple) – Padding that was added to the input

Inputs:

input: the input Tensor to invert
indices: the indices given out by MaxPool1d
output_size (optional): the targeted output size

Shape:

Input: $(N, C, H_{in})$
Output: $(N, C, H_{out})$, where

\[H_{out} = (H_{in} - 1) \times \text{stride}[0] - 2 \times \text{padding}[0] + \text{kernel\_size}[0] \]

or as given by output_size in the call operator

Example:

>>> pool = nn.MaxPool1d(2, stride=2, return_indices=True)
>>> unpool = nn.MaxUnpool1d(2, stride=2)
>>> input = torch.tensor([[[1., 2, 3, 4, 5, 6, 7, 8]]])
>>> output, indices = pool(input)
>>> unpool(output, indices)
tensor([[[ 0.,  2.,  0.,  4.,  0.,  6.,  0., 8.]]])

>>> # Example showcasing the use of output_size
>>> input = torch.tensor([[[1., 2, 3, 4, 5, 6, 7, 8, 9]]])
>>> output, indices = pool(input)
>>> unpool(output, indices, output_size=input.size())
tensor([[[ 0.,  2.,  0.,  4.,  0.,  6.,  0., 8.,  0.]]])

>>> unpool(output, indices)
tensor([[[ 0.,  2.,  0.,  4.,  0.,  6.,  0., 8.]]])

MaxUnpool2d¶

class torch.nn.MaxUnpool2d(kernel_size, stride=None, padding=0)¶

Computes a partial inverse of MaxPool2d.

MaxPool2d is not fully invertible, since the non-maximal values are lost.

MaxUnpool2d takes in as input the output of MaxPool2d including the indices of the maximal values and computes a partial inverse in which all non-maximal values are set to zero.

Note

MaxPool2d can map several input sizes to the same output sizes. Hence, the inversion process can get ambiguous. To accommodate this, you can provide the needed output size as an additional argument output_size in the forward call. See the Inputs and Example below.

Parameters

kernel_size (int or tuple) – Size of the max pooling window.
stride (int or tuple) – Stride of the max pooling window. It is set to kernel_size by default.
padding (int or tuple) – Padding that was added to the input

Inputs:

input: the input Tensor to invert
indices: the indices given out by MaxPool2d
output_size (optional): the targeted output size

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$, where

\[H_{out} = (H_{in} - 1) \times \text{stride[0]} - 2 \times \text{padding[0]} + \text{kernel\_size[0]} \]

\[W_{out} = (W_{in} - 1) \times \text{stride[1]} - 2 \times \text{padding[1]} + \text{kernel\_size[1]} \]

or as given by output_size in the call operator

Example:

>>> pool = nn.MaxPool2d(2, stride=2, return_indices=True)
>>> unpool = nn.MaxUnpool2d(2, stride=2)
>>> input = torch.tensor([[[[ 1.,  2,  3,  4],
                            [ 5,  6,  7,  8],
                            [ 9, 10, 11, 12],
                            [13, 14, 15, 16]]]])
>>> output, indices = pool(input)
>>> unpool(output, indices)
tensor([[[[  0.,   0.,   0.,   0.],
          [  0.,   6.,   0.,   8.],
          [  0.,   0.,   0.,   0.],
          [  0.,  14.,   0.,  16.]]]])

>>> # specify a different output size than input size
>>> unpool(output, indices, output_size=torch.Size([1, 1, 5, 5]))
tensor([[[[  0.,   0.,   0.,   0.,   0.],
          [  6.,   0.,   8.,   0.,   0.],
          [  0.,   0.,   0.,  14.,   0.],
          [ 16.,   0.,   0.,   0.,   0.],
          [  0.,   0.,   0.,   0.,   0.]]]])

MaxUnpool3d¶

class torch.nn.MaxUnpool3d(kernel_size, stride=None, padding=0)¶

Computes a partial inverse of MaxPool3d.

MaxPool3d is not fully invertible, since the non-maximal values are lost. MaxUnpool3d takes in as input the output of MaxPool3d including the indices of the maximal values and computes a partial inverse in which all non-maximal values are set to zero.

Note

MaxPool3d can map several input sizes to the same output sizes. Hence, the inversion process can get ambiguous. To accommodate this, you can provide the needed output size as an additional argument output_size in the forward call. See the Inputs section below.

Parameters

kernel_size (int or tuple) – Size of the max pooling window.
stride (int or tuple) – Stride of the max pooling window. It is set to kernel_size by default.
padding (int or tuple) – Padding that was added to the input

Inputs:

input: the input Tensor to invert
indices: the indices given out by MaxPool3d
output_size (optional): the targeted output size

Shape:

Input: $(N, C, D_{in}, H_{in}, W_{in})$
Output: $(N, C, D_{out}, H_{out}, W_{out})$, where

\[D_{out} = (D_{in} - 1) \times \text{stride[0]} - 2 \times \text{padding[0]} + \text{kernel\_size[0]} \]

\[H_{out} = (H_{in} - 1) \times \text{stride[1]} - 2 \times \text{padding[1]} + \text{kernel\_size[1]} \]

\[W_{out} = (W_{in} - 1) \times \text{stride[2]} - 2 \times \text{padding[2]} + \text{kernel\_size[2]} \]

or as given by output_size in the call operator

Example:

>>> # pool of square window of size=3, stride=2
>>> pool = nn.MaxPool3d(3, stride=2, return_indices=True)
>>> unpool = nn.MaxUnpool3d(3, stride=2)
>>> output, indices = pool(torch.randn(20, 16, 51, 33, 15))
>>> unpooled_output = unpool(output, indices)
>>> unpooled_output.size()
torch.Size([20, 16, 51, 33, 15])

AvgPool1d¶

class torch.nn.AvgPool1d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True)¶

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

In the simplest case, the output value of the layer with input size $(N, C, L)$, output $(N, C, L_{out})$ and kernel_size $k$ can be precisely described as:

\[\text{out}(N_i, C_j, l) = \frac{1}{k} \sum_{m=0}^{k-1} \text{input}(N_i, C_j, \text{stride} \times l + m)\]

If padding is non-zero, then the input is implicitly zero-padded on both sides for padding number of points.

The parameters kernel_size, stride, padding can each be an int or a one-element tuple.

Parameters

kernel_size – the size of the window
stride – the stride of the window. Default value is kernel_size
padding – implicit zero padding to be added on both sides
ceil_mode – when True, will use ceil instead of floor to compute the output shape
count_include_pad – when True, will include the zero-padding in the averaging calculation

Shape:

Input: $(N, C, L_{in})$
Output: $(N, C, L_{out})$, where

\[L_{out} = \left\lfloor \frac{L_{in} + 2 \times \text{padding} - \text{kernel\_size}}{\text{stride}} + 1\right\rfloor \]

Examples:

>>> # pool with window of size=3, stride=2
>>> m = nn.AvgPool1d(3, stride=2)
>>> m(torch.tensor([[[1.,2,3,4,5,6,7]]]))
tensor([[[ 2.,  4.,  6.]]])

AvgPool2d¶

class torch.nn.AvgPool2d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True)¶

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

In the simplest case, the output value of the layer with input size $(N, C, H, W)$, output $(N, C, H_{out}, W_{out})$ and kernel_size $(kH, kW)$ can be precisely described as:

\[out(N_i, C_j, h, w) = \frac{1}{kH * kW} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1} input(N_i, C_j, stride[0] \times h + m, stride[1] \times w + n)\]

If padding is non-zero, then the input is implicitly zero-padded on both sides for padding number of points.

The parameters kernel_size, stride, padding can either be:

a single int – in which case the same value is used for the height and width dimension

a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension

Parameters

kernel_size – the size of the window
stride – the stride of the window. Default value is kernel_size
padding – implicit zero padding to be added on both sides
ceil_mode – when True, will use ceil instead of floor to compute the output shape
count_include_pad – when True, will include the zero-padding in the averaging calculation

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$, where

\[H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[0] - \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor \]

\[W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[1] - \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor \]

Examples:

>>> # pool of square window of size=3, stride=2
>>> m = nn.AvgPool2d(3, stride=2)
>>> # pool of non-square window
>>> m = nn.AvgPool2d((3, 2), stride=(2, 1))
>>> input = torch.randn(20, 16, 50, 32)
>>> output = m(input)

AvgPool3d¶

class torch.nn.AvgPool3d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True)¶

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

In the simplest case, the output value of the layer with input size $(N, C, D, H, W)$, output $(N, C, D_{out}, H_{out}, W_{out})$ and kernel_size $(kD, kH, kW)$ can be precisely described as:

\[\begin{aligned} \text{out}(N_i, C_j, d, h, w) ={} & \sum_{k=0}^{kD-1} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1} \\ & \frac{\text{input}(N_i, C_j, \text{stride}[0] \times d + k, \text{stride}[1] \times h + m, \text{stride}[2] \times w + n)} {kD \times kH \times kW} \end{aligned} \]

If padding is non-zero, then the input is implicitly zero-padded on all three sides for padding number of points.

The parameters kernel_size, stride can either be:

a single int – in which case the same value is used for the depth, height and width dimension

a tuple of three ints – in which case, the first int is used for the depth dimension, the second int for the height dimension and the third int for the width dimension

Parameters

kernel_size – the size of the window
stride – the stride of the window. Default value is kernel_size
padding – implicit zero padding to be added on all three sides
ceil_mode – when True, will use ceil instead of floor to compute the output shape
count_include_pad – when True, will include the zero-padding in the averaging calculation

Shape:

Input: $(N, C, D_{in}, H_{in}, W_{in})$
Output: $(N, C, D_{out}, H_{out}, W_{out})$, where

\[D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] - \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor \]

\[H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] - \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor \]

\[W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] - \text{kernel\_size}[2]}{\text{stride}[2]} + 1\right\rfloor \]

Examples:

>>> # pool of square window of size=3, stride=2
>>> m = nn.AvgPool3d(3, stride=2)
>>> # pool of non-square window
>>> m = nn.AvgPool3d((3, 2, 2), stride=(2, 1, 2))
>>> input = torch.randn(20, 16, 50,44, 31)
>>> output = m(input)

FractionalMaxPool2d¶

class torch.nn.FractionalMaxPool2d(kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None)¶

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

Fractional MaxPooling is described in detail in the paper Fractional MaxPooling by Ben Graham

The max-pooling operation is applied in $kH \times kW$ regions by a stochastic step size determined by the target output size. The number of output features is equal to the number of input planes.

Parameters

kernel_size – the size of the window to take a max over. Can be a single number k (for a square kernel of k x k) or a tuple (kh, kw)
output_size – the target output size of the image of the form oH x oW. Can be a tuple (oH, oW) or a single number oH for a square image oH x oH
output_ratio – If one wants to have an output size as a ratio of the input size, this option can be given. This has to be a number or tuple in the range (0, 1)
return_indices – if True, will return the indices along with the outputs. Useful to pass to nn.MaxUnpool2d(). Default: False

Examples

>>> # pool of square window of size=3, and target output size 13x12
>>> m = nn.FractionalMaxPool2d(3, output_size=(13, 12))
>>> # pool of square window and target output size being half of input image size
>>> m = nn.FractionalMaxPool2d(3, output_ratio=(0.5, 0.5))
>>> input = torch.randn(20, 16, 50, 32)
>>> output = m(input)

LPPool1d¶

class torch.nn.LPPool1d(norm_type, kernel_size, stride=None, ceil_mode=False)¶

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

On each window, the function computed is:

\[f(X) = \sqrt[p]{\sum_{x \in X} x^{p}} \]

At p = $\infty$, one gets Max Pooling
At p = 1, one gets Sum Pooling (which is proportional to Average Pooling)

Note

If the sum to the power of p is zero, the gradient of this function is not defined. This implementation will set the gradient to zero in this case.

Parameters

kernel_size – a single int, the size of the window
stride – a single int, the stride of the window. Default value is kernel_size
ceil_mode – when True, will use ceil instead of floor to compute the output shape

Shape:

Input: $(N, C, L_{in})$
Output: $(N, C, L_{out})$, where

\[L_{out} = \left\lfloor\frac{L_{in} + 2 \times \text{padding} - \text{kernel\_size}}{\text{stride}} + 1\right\rfloor \]

Examples::

>>> # power-2 pool of window of length 3, with stride 2.
>>> m = nn.LPPool1d(2, 3, stride=2)
>>> input = torch.randn(20, 16, 50)
>>> output = m(input)

LPPool2d¶

class torch.nn.LPPool2d(norm_type, kernel_size, stride=None, ceil_mode=False)¶

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

On each window, the function computed is:

\[f(X) = \sqrt[p]{\sum_{x \in X} x^{p}} \]

At p = $\infty$, one gets Max Pooling
At p = 1, one gets Sum Pooling (which is proportional to average pooling)

The parameters kernel_size, stride can either be:

a single int – in which case the same value is used for the height and width dimension

a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension

Note

If the sum to the power of p is zero, the gradient of this function is not defined. This implementation will set the gradient to zero in this case.

Parameters

kernel_size – the size of the window
stride – the stride of the window. Default value is kernel_size
ceil_mode – when True, will use ceil instead of floor to compute the output shape

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$, where

\[H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor \]

\[W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor \]

Examples:

>>> # power-2 pool of square window of size=3, stride=2
>>> m = nn.LPPool2d(2, 3, stride=2)
>>> # pool of non-square window of power 1.2
>>> m = nn.LPPool2d(1.2, (3, 2), stride=(2, 1))
>>> input = torch.randn(20, 16, 50, 32)
>>> output = m(input)

AdaptiveMaxPool1d¶

class torch.nn.AdaptiveMaxPool1d(output_size, return_indices=False)¶

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

The output size is H, for any input size. The number of output features is equal to the number of input planes.

Parameters

output_size – the target output size H
return_indices – if True, will return the indices along with the outputs. Useful to pass to nn.MaxUnpool1d. Default: False

Examples

>>> # target output size of 5
>>> m = nn.AdaptiveMaxPool1d(5)
>>> input = torch.randn(1, 64, 8)
>>> output = m(input)

AdaptiveMaxPool2d¶

class torch.nn.AdaptiveMaxPool2d(output_size, return_indices=False)¶

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

The output is of size H x W, for any input size. The number of output features is equal to the number of input planes.

Parameters

output_size – the target output size of the image of the form H x W. Can be a tuple (H, W) or a single H for a square image H x H. H and W can be either a int, or None which means the size will be the same as that of the input.
return_indices – if True, will return the indices along with the outputs. Useful to pass to nn.MaxUnpool2d. Default: False

Examples

>>> # target output size of 5x7
>>> m = nn.AdaptiveMaxPool2d((5,7))
>>> input = torch.randn(1, 64, 8, 9)
>>> output = m(input)
>>> # target output size of 7x7 (square)
>>> m = nn.AdaptiveMaxPool2d(7)
>>> input = torch.randn(1, 64, 10, 9)
>>> output = m(input)
>>> # target output size of 10x7
>>> m = nn.AdaptiveMaxPool2d((None, 7))
>>> input = torch.randn(1, 64, 10, 9)
>>> output = m(input)

AdaptiveMaxPool3d¶

class torch.nn.AdaptiveMaxPool3d(output_size, return_indices=False)¶

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

The output is of size D x H x W, for any input size. The number of output features is equal to the number of input planes.

Parameters

output_size – the target output size of the image of the form D x H x W. Can be a tuple (D, H, W) or a single D for a cube D x D x D. D, H and W can be either a int, or None which means the size will be the same as that of the input.
return_indices – if True, will return the indices along with the outputs. Useful to pass to nn.MaxUnpool3d. Default: False

Examples

>>> # target output size of 5x7x9
>>> m = nn.AdaptiveMaxPool3d((5,7,9))
>>> input = torch.randn(1, 64, 8, 9, 10)
>>> output = m(input)
>>> # target output size of 7x7x7 (cube)
>>> m = nn.AdaptiveMaxPool3d(7)
>>> input = torch.randn(1, 64, 10, 9, 8)
>>> output = m(input)
>>> # target output size of 7x9x8
>>> m = nn.AdaptiveMaxPool3d((7, None, None))
>>> input = torch.randn(1, 64, 10, 9, 8)
>>> output = m(input)

AdaptiveAvgPool1d¶

class torch.nn.AdaptiveAvgPool1d(output_size)¶

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

The output size is H, for any input size. The number of output features is equal to the number of input planes.

Parameters: output_size – the target output size H

Examples

>>> # target output size of 5
>>> m = nn.AdaptiveAvgPool1d(5)
>>> input = torch.randn(1, 64, 8)
>>> output = m(input)

AdaptiveAvgPool2d¶

class torch.nn.AdaptiveAvgPool2d(output_size)¶

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

The output is of size H x W, for any input size. The number of output features is equal to the number of input planes.

Parameters: output_size – the target output size of the image of the form H x W. Can be a tuple (H, W) or a single H for a square image H x H. H and W can be either a int, or None which means the size will be the same as that of the input.

Examples

>>> # target output size of 5x7
>>> m = nn.AdaptiveAvgPool2d((5,7))
>>> input = torch.randn(1, 64, 8, 9)
>>> output = m(input)
>>> # target output size of 7x7 (square)
>>> m = nn.AdaptiveAvgPool2d(7)
>>> input = torch.randn(1, 64, 10, 9)
>>> output = m(input)
>>> # target output size of 10x7
>>> m = nn.AdaptiveMaxPool2d((None, 7))
>>> input = torch.randn(1, 64, 10, 9)
>>> output = m(input)

AdaptiveAvgPool3d¶

class torch.nn.AdaptiveAvgPool3d(output_size)¶

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

The output is of size D x H x W, for any input size. The number of output features is equal to the number of input planes.

Parameters: output_size – the target output size of the form D x H x W. Can be a tuple (D, H, W) or a single number D for a cube D x D x D. D, H and W can be either a int, or None which means the size will be the same as that of the input.

Examples

>>> # target output size of 5x7x9
>>> m = nn.AdaptiveAvgPool3d((5,7,9))
>>> input = torch.randn(1, 64, 8, 9, 10)
>>> output = m(input)
>>> # target output size of 7x7x7 (cube)
>>> m = nn.AdaptiveAvgPool3d(7)
>>> input = torch.randn(1, 64, 10, 9, 8)
>>> output = m(input)
>>> # target output size of 7x9x8
>>> m = nn.AdaptiveMaxPool3d((7, None, None))
>>> input = torch.randn(1, 64, 10, 9, 8)
>>> output = m(input)

Padding layers¶

ReflectionPad1d¶

class torch.nn.ReflectionPad1d(padding)¶

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

For N-dimensional padding, use torch.nn.functional.pad().

Parameters: padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 2-tuple, uses ($\text{padding\_left}$, $\text{padding\_right}$)

Shape:

Input: $(N, C, W_{in})$
Output: $(N, C, W_{out})$ where

$W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}$

Examples:

>>> m = nn.ReflectionPad1d(2)
>>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)
>>> input
tensor([[[0., 1., 2., 3.],
         [4., 5., 6., 7.]]])
>>> m(input)
tensor([[[2., 1., 0., 1., 2., 3., 2., 1.],
         [6., 5., 4., 5., 6., 7., 6., 5.]]])
>>> # using different paddings for different sides
>>> m = nn.ReflectionPad1d((3, 1))
>>> m(input)
tensor([[[3., 2., 1., 0., 1., 2., 3., 2.],
         [7., 6., 5., 4., 5., 6., 7., 6.]]])

ReflectionPad2d¶

class torch.nn.ReflectionPad2d(padding)¶

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

For N-dimensional padding, use torch.nn.functional.pad().

Parameters: padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 4-tuple, uses ($\text{padding\_left}$, $\text{padding\_right}$, $\text{padding\_top}$, $\text{padding\_bottom}$)

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$ where

$H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}$

$W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}$

Examples:

>>> m = nn.ReflectionPad2d(2)
>>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)
>>> input
tensor([[[[0., 1., 2.],
          [3., 4., 5.],
          [6., 7., 8.]]]])
>>> m(input)
tensor([[[[8., 7., 6., 7., 8., 7., 6.],
          [5., 4., 3., 4., 5., 4., 3.],
          [2., 1., 0., 1., 2., 1., 0.],
          [5., 4., 3., 4., 5., 4., 3.],
          [8., 7., 6., 7., 8., 7., 6.],
          [5., 4., 3., 4., 5., 4., 3.],
          [2., 1., 0., 1., 2., 1., 0.]]]])
>>> # using different paddings for different sides
>>> m = nn.ReflectionPad2d((1, 1, 2, 0))
>>> m(input)
tensor([[[[7., 6., 7., 8., 7.],
          [4., 3., 4., 5., 4.],
          [1., 0., 1., 2., 1.],
          [4., 3., 4., 5., 4.],
          [7., 6., 7., 8., 7.]]]])

ReplicationPad1d¶

class torch.nn.ReplicationPad1d(padding)¶

Pads the input tensor using replication of the input boundary.

For N-dimensional padding, use torch.nn.functional.pad().

Parameters: padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 2-tuple, uses ($\text{padding\_left}$, $\text{padding\_right}$)

Shape:

Input: $(N, C, W_{in})$
Output: $(N, C, W_{out})$ where

$W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}$

Examples:

>>> m = nn.ReplicationPad1d(2)
>>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)
>>> input
tensor([[[0., 1., 2., 3.],
         [4., 5., 6., 7.]]])
>>> m(input)
tensor([[[0., 0., 0., 1., 2., 3., 3., 3.],
         [4., 4., 4., 5., 6., 7., 7., 7.]]])
>>> # using different paddings for different sides
>>> m = nn.ReplicationPad1d((3, 1))
>>> m(input)
tensor([[[0., 0., 0., 0., 1., 2., 3., 3.],
         [4., 4., 4., 4., 5., 6., 7., 7.]]])

ReplicationPad2d¶

class torch.nn.ReplicationPad2d(padding)¶

Pads the input tensor using replication of the input boundary.

For N-dimensional padding, use torch.nn.functional.pad().

Parameters: padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 4-tuple, uses ($\text{padding\_left}$, $\text{padding\_right}$, $\text{padding\_top}$, $\text{padding\_bottom}$)

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$ where

$H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}$

$W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}$

Examples:

>>> m = nn.ReplicationPad2d(2)
>>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)
>>> input
tensor([[[[0., 1., 2.],
          [3., 4., 5.],
          [6., 7., 8.]]]])
>>> m(input)
tensor([[[[0., 0., 0., 1., 2., 2., 2.],
          [0., 0., 0., 1., 2., 2., 2.],
          [0., 0., 0., 1., 2., 2., 2.],
          [3., 3., 3., 4., 5., 5., 5.],
          [6., 6., 6., 7., 8., 8., 8.],
          [6., 6., 6., 7., 8., 8., 8.],
          [6., 6., 6., 7., 8., 8., 8.]]]])
>>> # using different paddings for different sides
>>> m = nn.ReplicationPad2d((1, 1, 2, 0))
>>> m(input)
tensor([[[[0., 0., 1., 2., 2.],
          [0., 0., 1., 2., 2.],
          [0., 0., 1., 2., 2.],
          [3., 3., 4., 5., 5.],
          [6., 6., 7., 8., 8.]]]])

ReplicationPad3d¶

class torch.nn.ReplicationPad3d(padding)¶

Pads the input tensor using replication of the input boundary.

For N-dimensional padding, use torch.nn.functional.pad().

Parameters: padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 6-tuple, uses ($\text{padding\_left}$, $\text{padding\_right}$, $\text{padding\_top}$, $\text{padding\_bottom}$, $\text{padding\_front}$, $\text{padding\_back}$)

Shape:

Input: $(N, C, D_{in}, H_{in}, W_{in})$
Output: $(N, C, D_{out}, H_{out}, W_{out})$ where

$D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}$

$H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}$

$W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}$

Examples:

>>> m = nn.ReplicationPad3d(3)
>>> input = torch.randn(16, 3, 8, 320, 480)
>>> output = m(input)
>>> # using different paddings for different sides
>>> m = nn.ReplicationPad3d((3, 3, 6, 6, 1, 1))
>>> output = m(input)

ZeroPad2d¶

class torch.nn.ZeroPad2d(padding)¶

Pads the input tensor boundaries with zero.

For N-dimensional padding, use torch.nn.functional.pad().

Parameters: padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 4-tuple, uses ($\text{padding\_left}$, $\text{padding\_right}$, $\text{padding\_top}$, $\text{padding\_bottom}$)

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$ where

$H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}$

$W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}$

Examples:

>>> m = nn.ZeroPad2d(2)
>>> input = torch.randn(1, 1, 3, 3)
>>> input
tensor([[[[-0.1678, -0.4418,  1.9466],
          [ 0.9604, -0.4219, -0.5241],
          [-0.9162, -0.5436, -0.6446]]]])
>>> m(input)
tensor([[[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000, -0.1678, -0.4418,  1.9466,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.9604, -0.4219, -0.5241,  0.0000,  0.0000],
          [ 0.0000,  0.0000, -0.9162, -0.5436, -0.6446,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]])
>>> # using different paddings for different sides
>>> m = nn.ZeroPad2d((1, 1, 2, 0))
>>> m(input)
tensor([[[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000, -0.1678, -0.4418,  1.9466,  0.0000],
          [ 0.0000,  0.9604, -0.4219, -0.5241,  0.0000],
          [ 0.0000, -0.9162, -0.5436, -0.6446,  0.0000]]]])

ConstantPad1d¶

class torch.nn.ConstantPad1d(padding, value)¶

Pads the input tensor boundaries with a constant value.

For N-dimensional padding, use torch.nn.functional.pad().

Parameters: padding (int, tuple) – the size of the padding. If is int, uses the same padding in both boundaries. If a 2-tuple, uses ($\text{padding\_left}$, $\text{padding\_right}$)

Shape:

Input: $(N, C, W_{in})$
Output: $(N, C, W_{out})$ where

$W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}$

Examples:

>>> m = nn.ConstantPad1d(2, 3.5)
>>> input = torch.randn(1, 2, 4)
>>> input
tensor([[[-1.0491, -0.7152, -0.0749,  0.8530],
         [-1.3287,  1.8966,  0.1466, -0.2771]]])
>>> m(input)
tensor([[[ 3.5000,  3.5000, -1.0491, -0.7152, -0.0749,  0.8530,  3.5000,
           3.5000],
         [ 3.5000,  3.5000, -1.3287,  1.8966,  0.1466, -0.2771,  3.5000,
           3.5000]]])
>>> m = nn.ConstantPad1d(2, 3.5)
>>> input = torch.randn(1, 2, 3)
>>> input
tensor([[[ 1.6616,  1.4523, -1.1255],
         [-3.6372,  0.1182, -1.8652]]])
>>> m(input)
tensor([[[ 3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000,  3.5000],
         [ 3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000,  3.5000]]])
>>> # using different paddings for different sides
>>> m = nn.ConstantPad1d((3, 1), 3.5)
>>> m(input)
tensor([[[ 3.5000,  3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000],
         [ 3.5000,  3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000]]])

ConstantPad2d¶

class torch.nn.ConstantPad2d(padding, value)¶

Pads the input tensor boundaries with a constant value.

For N-dimensional padding, use torch.nn.functional.pad().

Parameters: padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 4-tuple, uses ($\text{padding\_left}$, $\text{padding\_right}$, $\text{padding\_top}$, $\text{padding\_bottom}$)

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$ where

$H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}$

$W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}$

Examples:

>>> m = nn.ConstantPad2d(2, 3.5)
>>> input = torch.randn(1, 2, 2)
>>> input
tensor([[[ 1.6585,  0.4320],
         [-0.8701, -0.4649]]])
>>> m(input)
tensor([[[ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  1.6585,  0.4320,  3.5000,  3.5000],
         [ 3.5000,  3.5000, -0.8701, -0.4649,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000]]])
>>> # using different paddings for different sides
>>> m = nn.ConstantPad2d((3, 0, 2, 1), 3.5)
>>> m(input)
tensor([[[ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  1.6585,  0.4320],
         [ 3.5000,  3.5000,  3.5000, -0.8701, -0.4649],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000]]])

ConstantPad3d¶

class torch.nn.ConstantPad3d(padding, value)¶

Pads the input tensor boundaries with a constant value.

For N-dimensional padding, use torch.nn.functional.pad().

Parameters: padding (int, tuple) – the size of the padding. If is int, uses the same padding in all boundaries. If a 6-tuple, uses ($\text{padding\_left}$, $\text{padding\_right}$, $\text{padding\_top}$, $\text{padding\_bottom}$, $\text{padding\_front}$, $\text{padding\_back}$)

Shape:

Input: $(N, C, D_{in}, H_{in}, W_{in})$
Output: $(N, C, D_{out}, H_{out}, W_{out})$ where

$D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}$

$H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}$

$W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}$

Examples:

>>> m = nn.ConstantPad3d(3, 3.5)
>>> input = torch.randn(16, 3, 10, 20, 30)
>>> output = m(input)
>>> # using different paddings for different sides
>>> m = nn.ConstantPad3d((3, 3, 6, 6, 0, 1), 3.5)
>>> output = m(input)

Non-linear activations (weighted sum, nonlinearity)¶

ELU¶

class torch.nn.ELU(alpha=1.0, inplace=False)¶

Applies the element-wise function:

\[\text{ELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x) - 1)) \]

Parameters

alpha – the $\alpha$ value for the ELU formulation. Default: 1.0
inplace – can optionally do the operation in-place. Default: False

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.ELU()
>>> input = torch.randn(2)
>>> output = m(input)

Hardshrink¶

class torch.nn.Hardshrink(lambd=0.5)¶

Applies the hard shrinkage function element-wise:

\[\text{HardShrink}(x) = \begin{cases} x, & \text{ if } x > \lambda \\ x, & \text{ if } x < -\lambda \\ 0, & \text{ otherwise } \end{cases} \]

Parameters: lambd – the $\lambda$ value for the Hardshrink formulation. Default: 0.5

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

scripts/activation_images/Hardshrink.png

Examples:

>>> m = nn.Hardshrink()
>>> input = torch.randn(2)
>>> output = m(input)

Hardtanh¶

class torch.nn.Hardtanh(min_val=-1.0, max_val=1.0, inplace=False, min_value=None, max_value=None)¶

Applies the HardTanh function element-wise

HardTanh is defined as:

\[\text{HardTanh}(x) = \begin{cases} 1 & \text{ if } x > 1 \\ -1 & \text{ if } x < -1 \\ x & \text{ otherwise } \\ \end{cases} \]

The range of the linear region $[-1, 1]$ can be adjusted using min_val and max_val.

Parameters

min_val – minimum value of the linear region range. Default: -1
max_val – maximum value of the linear region range. Default: 1
inplace – can optionally do the operation in-place. Default: False

Keyword arguments min_value and max_value have been deprecated in favor of min_val and max_val.

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.Hardtanh(-2, 2)
>>> input = torch.randn(2)
>>> output = m(input)

LeakyReLU¶

class torch.nn.LeakyReLU(negative_slope=0.01, inplace=False)¶

Applies the element-wise function:

\[\text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x) \]

\[\text{LeakyRELU}(x) = \begin{cases} x, & \text{ if } x \geq 0 \\ \text{negative\_slope} \times x, & \text{ otherwise } \end{cases} \]

Parameters

negative_slope – Controls the angle of the negative slope. Default: 1e-2
inplace – can optionally do the operation in-place. Default: False

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.LeakyReLU(0.1)
>>> input = torch.randn(2)
>>> output = m(input)

LogSigmoid¶

class torch.nn.LogSigmoid¶

Applies the element-wise function:

\[\text{LogSigmoid}(x) = \log\left(\frac{ 1 }{ 1 + \exp(-x)}\right) \]

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

scripts/activation_images/LogSigmoid.png

Examples:

>>> m = nn.LogSigmoid()
>>> input = torch.randn(2)
>>> output = m(input)

PReLU¶

class torch.nn.PReLU(num_parameters=1, init=0.25)¶

Applies the element-wise function:

\[\text{PReLU}(x) = \max(0,x) + a * \min(0,x) \]

\[\text{PReLU}(x) = \begin{cases} x, & \text{ if } x \geq 0 \\ ax, & \text{ otherwise } \end{cases} \]

Here $a$ is a learnable parameter. When called without arguments, nn.PReLU() uses a single parameter $a$ across all input channels. If called with nn.PReLU(nChannels), a separate $a$ is used for each input channel.

Note

weight decay should not be used when learning $a$ 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.

Parameters

num_parameters (int) – number of $a$ to learn. Although it takes an int as input, there is only two values are legitimate: 1, or the number of channels at input. Default: 1
init (float) – the initial value of $a$. Default: 0.25

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Variables: ~PReLU.weight (Tensor) – the learnable weights of shape (num_parameters).

Examples:

>>> m = nn.PReLU()
>>> input = torch.randn(2)
>>> output = m(input)

ReLU¶

class torch.nn.ReLU(inplace=False)¶

Applies the rectified linear unit function element-wise:

$\text{ReLU}(x)= \max(0, x)$

Parameters: inplace – can optionally do the operation in-place. Default: False

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

  >>> m = nn.ReLU()
  >>> input = torch.randn(2)
  >>> output = m(input)


An implementation of CReLU - https://arxiv.org/abs/1603.05201

  >>> m = nn.ReLU()
  >>> input = torch.randn(2).unsqueeze(0)
  >>> output = torch.cat((m(input),m(-input)))

ReLU6¶

class torch.nn.ReLU6(inplace=False)¶

Applies the element-wise function:

\[\text{ReLU6}(x) = \min(\max(0,x), 6) \]

Parameters: inplace – can optionally do the operation in-place. Default: False

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.ReLU6()
>>> input = torch.randn(2)
>>> output = m(input)

RReLU¶

class torch.nn.RReLU(lower=0.125, upper=0.3333333333333333, inplace=False)¶

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:

\[\text{RReLU}(x) = \begin{cases} x & \text{if } x \geq 0 \\ ax & \text{ otherwise } \end{cases} \]

where $a$ is randomly sampled from uniform distribution $\mathcal{U}(\text{lower}, \text{upper})$.

See: https://arxiv.org/pdf/1505.00853.pdf

Parameters

lower – lower bound of the uniform distribution. Default: $\frac{1}{8}$
upper – upper bound of the uniform distribution. Default: $\frac{1}{3}$
inplace – can optionally do the operation in-place. Default: False

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.RReLU(0.1, 0.3)
>>> input = torch.randn(2)
>>> output = m(input)

SELU¶

class torch.nn.SELU(inplace=False)¶

Applied element-wise, as:

\[\text{SELU}(x) = \text{scale} * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1))) \]

with $\alpha = 1.6732632423543772848170429916717$ and $\text{scale} = 1.0507009873554804934193349852946$.

More details can be found in the paper Self-Normalizing Neural Networks .

Parameters: inplace (bool, optional) – can optionally do the operation in-place. Default: False

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.SELU()
>>> input = torch.randn(2)
>>> output = m(input)

CELU¶

class torch.nn.CELU(alpha=1.0, inplace=False)¶

Applies the element-wise function:

\[\text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1)) \]

More details can be found in the paper Continuously Differentiable Exponential Linear Units .

Parameters

alpha – the $\alpha$ value for the CELU formulation. Default: 1.0
inplace – can optionally do the operation in-place. Default: False

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.CELU()
>>> input = torch.randn(2)
>>> output = m(input)

Sigmoid¶

class torch.nn.Sigmoid¶

Applies the element-wise function:

\[\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)} \]

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.Sigmoid()
>>> input = torch.randn(2)
>>> output = m(input)

Softplus¶

class torch.nn.Softplus(beta=1, threshold=20)¶

Applies the element-wise function:

\[\text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x)) \]

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 for inputs above a certain value.

Parameters

beta – the $\beta$ value for the Softplus formulation. Default: 1
threshold – values above this revert to a linear function. Default: 20

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.Softplus()
>>> input = torch.randn(2)
>>> output = m(input)

Softshrink¶

class torch.nn.Softshrink(lambd=0.5)¶

Applies the soft shrinkage function elementwise:

\[\text{SoftShrinkage}(x) = \begin{cases} x - \lambda, & \text{ if } x > \lambda \\ x + \lambda, & \text{ if } x < -\lambda \\ 0, & \text{ otherwise } \end{cases} \]

Parameters: lambd – the $\lambda$ value for the Softshrink formulation. Default: 0.5

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

scripts/activation_images/Softshrink.png

Examples:

>>> m = nn.Softshrink()
>>> input = torch.randn(2)
>>> output = m(input)

Softsign¶

class torch.nn.Softsign¶

Applies the element-wise function:

\[\text{SoftSign}(x) = \frac{x}{ 1 + |x|} \]

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.Softsign()
>>> input = torch.randn(2)
>>> output = m(input)

Tanh¶

class torch.nn.Tanh¶

Applies the element-wise function:

\[\text{Tanh}(x) = \tanh(x) = \frac{e^x - e^{-x}} {e^x + e^{-x}} \]

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.Tanh()
>>> input = torch.randn(2)
>>> output = m(input)

Tanhshrink¶

class torch.nn.Tanhshrink¶

Applies the element-wise function:

\[\text{Tanhshrink}(x) = x - \text{Tanh}(x) \]

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

scripts/activation_images/Tanhshrink.png

Examples:

>>> m = nn.Tanhshrink()
>>> input = torch.randn(2)
>>> output = m(input)

Threshold¶

class torch.nn.Threshold(threshold, value, inplace=False)¶

Thresholds each element of the input Tensor.

Threshold is defined as:

\[y = \begin{cases} x, &\text{ if } x > \text{threshold} \\ \text{value}, &\text{ otherwise } \end{cases} \]

Parameters

threshold – The value to threshold at
value – The value to replace with
inplace – can optionally do the operation in-place. Default: False

Shape:

Input: $(N, *)$ where * means, any number of additional dimensions
Output: $(N, *)$, same shape as the input

Examples:

>>> m = nn.Threshold(0.1, 20)
>>> input = torch.randn(2)
>>> output = m(input)

Non-linear activations (other)¶

Softmin¶

class torch.nn.Softmin(dim=None)¶

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:

\[\text{Softmin}(x_{i}) = \frac{\exp(-x_i)}{\sum_j \exp(-x_j)} \]

Shape:

Input: $(*)$ where * means, any number of additional dimensions
Output: $(*)$, same shape as the input

Parameters: dim (int) – A dimension along which Softmin will be computed (so every slice along dim will sum to 1).
Returns: a Tensor of the same dimension and shape as the input, with values in the range [0, 1]

Examples:

>>> m = nn.Softmin()
>>> input = torch.randn(2, 3)
>>> output = m(input)

Softmax¶

class torch.nn.Softmax(dim=None)¶

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:

\[\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)} \]

Shape:

Input: $(*)$ where * means, any number of additional dimensions
Output: $(*)$, same shape as the input

Returns: a Tensor of the same dimension and shape as the input with values in the range [0, 1]
Parameters: 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:

>>> m = nn.Softmax()
>>> input = torch.randn(2, 3)
>>> output = m(input)

Softmax2d¶

class torch.nn.Softmax2d¶

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, h_i, w_j)$

Shape:

Input: $(N, C, H, W)$
Output: $(N, C, H, W)$ (same shape as input)

Returns: a Tensor of the same dimension and shape as the input with values in the range [0, 1]

Examples:

>>> m = nn.Softmax2d()
>>> # you softmax over the 2nd dimension
>>> input = torch.randn(2, 3, 12, 13)
>>> output = m(input)

LogSoftmax¶

class torch.nn.LogSoftmax(dim=None)¶

Applies the $\log(\text{Softmax}(x))$ function to an n-dimensional input Tensor. The LogSoftmax formulation can be simplified as:

\[\text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right) \]

Shape:

Input: $(*)$ where * means, any number of additional dimensions
Output: $(*)$, same shape as the input

Parameters: dim (int) – A dimension along which LogSoftmax will be computed.
Returns: a Tensor of the same dimension and shape as the input with values in the range [-inf, 0)

Examples:

>>> m = nn.LogSoftmax()
>>> input = torch.randn(2, 3)
>>> output = m(input)

AdaptiveLogSoftmaxWithLoss¶

class torch.nn.AdaptiveLogSoftmaxWithLoss(in_features, n_classes, cutoffs, div_value=4.0, head_bias=False)¶

Efficient softmax approximation as described in Efficient softmax approximation for GPUs by Edouard Grave, Armand Joulin, Moustapha Cissé, David Grangier, and Hervé Jégou.

Adaptive softmax is an approximate strategy for training models with large output spaces. It is most effective when the label distribution is highly imbalanced, for example in natural language modelling, where the word frequency distribution approximately follows the Zipf’s law.

Adaptive softmax partitions the labels into several clusters, according to their frequency. These clusters may contain different number of targets each. Additionally, clusters containing less frequent labels assign lower dimensional embeddings to those labels, which speeds up the computation. For each minibatch, only clusters for which at least one target is present are evaluated.

The idea is that the clusters which are accessed frequently (like the first one, containing most frequent labels), should also be cheap to compute – that is, contain a small number of assigned labels.

We highly recommend taking a look at the original paper for more details.

cutoffs should be an ordered Sequence of integers sorted in the increasing order. It controls number of clusters and the partitioning of targets into clusters. For example setting cutoffs = [10, 100, 1000] means that first 10 targets will be assigned to the ‘head’ of the adaptive softmax, targets 11, 12, …, 100 will be assigned to the first cluster, and targets 101, 102, …, 1000 will be assigned to the second cluster, while targets 1001, 1002, …, n_classes - 1 will be assigned to the last, third cluster.
div_value is used to compute the size of each additional cluster, which is given as $\left\lfloor\frac{in\_features}{div\_value^{idx}}\right\rfloor$, where $idx$ is the cluster index (with clusters for less frequent words having larger indices, and indices starting from $1$).
head_bias if set to True, adds a bias term to the ‘head’ of the adaptive softmax. See paper for details. Set to False in the official implementation.

Warning

Labels passed as inputs to this module should be sorted accoridng to their frequency. This means that the most frequent label should be represented by the index 0, and the least frequent label should be represented by the index n_classes - 1.

Note

This module returns a NamedTuple with output and loss fields. See further documentation for details.

Note

To compute log-probabilities for all classes, the log_prob method can be used.

Parameters

in_features (int) – Number of features in the input tensor
n_classes (int) – Number of classes in the dataset
cutoffs (Sequence) – Cutoffs used to assign targets to their buckets
div_value (float, optional) – value used as an exponent to compute sizes of the clusters. Default: 4.0
head_bias (bool, optional) – If True, adds a bias term to the ‘head’ of the adaptive softmax. Default: False

Returns

output is a Tensor of size N containing computed target log probabilities for each example
loss is a Scalar representing the computed negative log likelihood loss

Return type

NamedTuple with output and loss fields

Shape:

input: $(N, in\_features)$
target: $(N)$ where each value satisfies $0 <= target[i] <= n\_classes$
output1: $(N)$
output2: Scalar

log_prob(input)¶

Computes log probabilities for all $n\_classes$

Parameters: input (Tensor) – a minibatch of examples
Returns: log-probabilities of for each class $c$ in range $0 <= c <= n\_classes$, where $n\_classes$ is a parameter passed to AdaptiveLogSoftmaxWithLoss constructor.

Shape:

Input: $(N, in\_features)$
Output: $(N, n\_classes)$

predict(input)¶

This is equivalent to self.log_pob(input).argmax(dim=1), but is more efficient in some cases.

Parameters: input (Tensor) – a minibatch of examples
Returns: a class with the highest probability for each example
Return type: output (Tensor)

Shape:

Input: $(N, in\_features)$
Output: $(N)$

Normalization layers¶

BatchNorm1d¶

class torch.nn.BatchNorm1d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)¶

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

\[y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

The mean and standard-deviation are calculated per-dimension over the mini-batches and $\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 $\mathcal{U}(0, 1)$ and the elements of $\beta$ are set to 0.

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 $\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t$, where $\hat{x}$ is the estimated statistic and $x_t$ is the new observed value.

Because the Batch Normalization is done over the C dimension, computing statistics on (N, L) slices, it’s common terminology to call this Temporal Batch Normalization.

Parameters

num_features – $C$ from an expected input of size $(N, C, L)$ or $L$ from input of size $(N, L)$
eps – a value added to the denominator for numerical stability. Default: 1e-5
momentum – the value used for the running_mean and running_var computation. Can be set to None for cumulative moving average (i.e. simple average). Default: 0.1
affine – a boolean value that when set to True, this module has learnable affine parameters. Default: True
track_running_stats – a boolean value that when set to True, this module tracks the running mean and variance, and when set to False, this module does not track such statistics and always uses batch statistics in both training and eval modes. Default: True

Shape:

Input: $(N, C)$ or $(N, C, L)$
Output: $(N, C)$ or $(N, C, L)$ (same shape as input)

Examples:

>>> # With Learnable Parameters
>>> m = nn.BatchNorm1d(100)
>>> # Without Learnable Parameters
>>> m = nn.BatchNorm1d(100, affine=False)
>>> input = torch.randn(20, 100)
>>> output = m(input)

BatchNorm2d¶

class torch.nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)¶

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

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

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

Because the Batch Normalization is done over the C dimension, computing statistics on (N, H, W) slices, it’s common terminology to call this Spatial Batch Normalization.

Parameters

num_features – $C$ from an expected input of size $(N, C, H, W)$
eps – a value added to the denominator for numerical stability. Default: 1e-5
momentum – the value used for the running_mean and running_var computation. Can be set to None for cumulative moving average (i.e. simple average). Default: 0.1
affine – a boolean value that when set to True, this module has learnable affine parameters. Default: True
track_running_stats – a boolean value that when set to True, this module tracks the running mean and variance, and when set to False, this module does not track such statistics and always uses batch statistics in both training and eval modes. Default: True

Shape:

Input: $(N, C, H, W)$
Output: $(N, C, H, W)$ (same shape as input)

Examples:

>>> # With Learnable Parameters
>>> m = nn.BatchNorm2d(100)
>>> # Without Learnable Parameters
>>> m = nn.BatchNorm2d(100, affine=False)
>>> input = torch.randn(20, 100, 35, 45)
>>> output = m(input)

BatchNorm3d¶

class torch.nn.BatchNorm3d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)¶

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

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

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

Because the Batch Normalization is done over the C dimension, computing statistics on (N, D, H, W) slices, it’s common terminology to call this Volumetric Batch Normalization or Spatio-temporal Batch Normalization.

Parameters

num_features – $C$ from an expected input of size $(N, C, D, H, W)$
eps – a value added to the denominator for numerical stability. Default: 1e-5
momentum – the value used for the running_mean and running_var computation. Can be set to None for cumulative moving average (i.e. simple average). Default: 0.1
affine – a boolean value that when set to True, this module has learnable affine parameters. Default: True
track_running_stats – a boolean value that when set to True, this module tracks the running mean and variance, and when set to False, this module does not track such statistics and always uses batch statistics in both training and eval modes. Default: True

Shape:

Input: $(N, C, D, H, W)$
Output: $(N, C, D, H, W)$ (same shape as input)

Examples:

>>> # With Learnable Parameters
>>> m = nn.BatchNorm3d(100)
>>> # Without Learnable Parameters
>>> m = nn.BatchNorm3d(100, affine=False)
>>> input = torch.randn(20, 100, 35, 45, 10)
>>> output = m(input)

GroupNorm¶

class torch.nn.GroupNorm(num_groups, num_channels, eps=1e-05, affine=True)¶

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

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta \]

The input channels are separated into num_groups groups, each containing num_channels / num_groups channels. The mean and standard-deviation are calculated separately over the each group. $\gamma$ and $\beta$ are learnable per-channel affine transform parameter vectorss of size num_channels if affine is True.

This layer uses statistics computed from input data in both training and evaluation modes.

Parameters

num_groups (int) – number of groups to separate the channels into
num_channels (int) – number of channels expected in input
eps – a value added to the denominator for numerical stability. Default: 1e-5
affine – a boolean value that when set to True, this module has learnable per-channel affine parameters initialized to ones (for weights) and zeros (for biases). Default: True.

Shape:

Input: $(N, C, *)$ where $C=\text{num\_channels}$
Output: $(N, C, *)$ (same shape as input)

Examples:

>>> input = torch.randn(20, 6, 10, 10)
>>> # Separate 6 channels into 3 groups
>>> m = nn.GroupNorm(3, 6)
>>> # Separate 6 channels into 6 groups (equivalent with InstanceNorm)
>>> m = nn.GroupNorm(6, 6)
>>> # Put all 6 channels into a single group (equivalent with LayerNorm)
>>> m = nn.GroupNorm(1, 6)
>>> # Activating the module
>>> output = m(input)

InstanceNorm1d¶

class torch.nn.InstanceNorm1d(num_features, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)¶

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

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

The mean and standard-deviation are calculated per-dimension separately for each object in a mini-batch. $\gamma$ and $\beta$ are learnable parameter vectors of size C (where C is the input size) if affine is True.

By default, this layer uses instance statistics computed from input data in both training and evaluation modes.

If track_running_stats is set to True, 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.

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 $\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momemtum} \times x_t$, where $\hat{x}$ is the estimated statistic and $x_t$ is the new observed value.

Note

InstanceNorm1d and LayerNorm are very similar, but have some subtle differences. InstanceNorm1d is applied on each channel of channeled data like multidimensional time series, but LayerNorm is usually applied on entire sample and often in NLP tasks. Additionaly, LayerNorm applies elementwise affine transform, while InstanceNorm1d usually don’t apply affine transform.

Parameters

num_features – $C$ from an expected input of size $(N, C, L)$ or $L$ from input of size $(N, L)$
eps – a value added to the denominator for numerical stability. Default: 1e-5
momentum – the value used for the running_mean and running_var computation. Default: 0.1
affine – a boolean value that when set to True, this module has learnable affine parameters, initialized the same way as done for batch normalization. Default: False.
track_running_stats – a boolean value that when set to True, this module tracks the running mean and variance, and when set to False, this module does not track such statistics and always uses batch statistics in both training and eval modes. Default: False

Shape:

Input: $(N, C, L)$
Output: $(N, C, L)$ (same shape as input)

Examples:

>>> # Without Learnable Parameters
>>> m = nn.InstanceNorm1d(100)
>>> # With Learnable Parameters
>>> m = nn.InstanceNorm1d(100, affine=True)
>>> input = torch.randn(20, 100, 40)
>>> output = m(input)

InstanceNorm2d¶

class torch.nn.InstanceNorm2d(num_features, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)¶

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

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

By default, this layer uses instance statistics computed from input data in both training and evaluation modes.

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 $\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momemtum} \times x_t$, where $\hat{x}$ is the estimated statistic and $x_t$ is the new observed value.

Note

InstanceNorm2d and LayerNorm are very similar, but have some subtle differences. InstanceNorm2d is applied on each channel of channeled data like RGB images, but LayerNorm is usually applied on entire sample and often in NLP tasks. Additionaly, LayerNorm applies elementwise affine transform, while InstanceNorm2d usually don’t apply affine transform.

Parameters

num_features – $C$ from an expected input of size $(N, C, H, W)$
eps – a value added to the denominator for numerical stability. Default: 1e-5
momentum – the value used for the running_mean and running_var computation. Default: 0.1
affine – a boolean value that when set to True, this module has learnable affine parameters, initialized the same way as done for batch normalization. Default: False.
track_running_stats – a boolean value that when set to True, this module tracks the running mean and variance, and when set to False, this module does not track such statistics and always uses batch statistics in both training and eval modes. Default: False

Shape:

Input: $(N, C, H, W)$
Output: $(N, C, H, W)$ (same shape as input)

Examples:

>>> # Without Learnable Parameters
>>> m = nn.InstanceNorm2d(100)
>>> # With Learnable Parameters
>>> m = nn.InstanceNorm2d(100, affine=True)
>>> input = torch.randn(20, 100, 35, 45)
>>> output = m(input)

InstanceNorm3d¶

class torch.nn.InstanceNorm3d(num_features, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)¶

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

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

By default, this layer uses instance statistics computed from input data in both training and evaluation modes.

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 $\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momemtum} \times x_t$, where $\hat{x}$ is the estimated statistic and $x_t$ is the new observed value.

Note

InstanceNorm3d and LayerNorm are very similar, but have some subtle differences. InstanceNorm3d is applied on each channel of channeled data like 3D models with RGB color, but LayerNorm is usually applied on entire sample and often in NLP tasks. Additionaly, LayerNorm applies elementwise affine transform, while InstanceNorm3d usually don’t apply affine transform.

Parameters

num_features – $C$ from an expected input of size $(N, C, D, H, W)$
eps – a value added to the denominator for numerical stability. Default: 1e-5
momentum – the value used for the running_mean and running_var computation. Default: 0.1
affine – a boolean value that when set to True, this module has learnable affine parameters, initialized the same way as done for batch normalization. Default: False.
track_running_stats – a boolean value that when set to True, this module tracks the running mean and variance, and when set to False, this module does not track such statistics and always uses batch statistics in both training and eval modes. Default: False

Shape:

Input: $(N, C, D, H, W)$
Output: $(N, C, D, H, W)$ (same shape as input)

Examples:

>>> # Without Learnable Parameters
>>> m = nn.InstanceNorm3d(100)
>>> # With Learnable Parameters
>>> m = nn.InstanceNorm3d(100, affine=True)
>>> input = torch.randn(20, 100, 35, 45, 10)
>>> output = m(input)

LayerNorm¶

class torch.nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True)¶

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

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta \]

The mean and standard-deviation are calculated separately over the last certain number dimensions which have to be of the shape specified by normalized_shape. $\gamma$ and $\beta$ are learnable affine transform parameters of normalized_shape if elementwise_affine is True.

Note

Unlike Batch Normalization and Instance Normalization, which applies scalar scale and bias for each entire channel/plane with the affine option, Layer Normalization applies per-element scale and bias with elementwise_affine.

This layer uses statistics computed from input data in both training and evaluation modes.

Parameters

normalized_shape (int or list or torch.Size) –
input shape from an expected input of size

\[[* \times \text{normalized\_shape}[0] \times \text{normalized\_shape}[1] \times \ldots \times \text{normalized\_shape}[-1]] \]

If a single integer is used, it is treated as a singleton list, and this module will normalize over the last dimension which is expected to be of that specific size.
eps – a value added to the denominator for numerical stability. Default: 1e-5
elementwise_affine – a boolean value that when set to True, this module has learnable per-element affine parameters initialized to ones (for weights) and zeros (for biases). Default: True.

Shape:

Input: $(N, *)$
Output: $(N, *)$ (same shape as input)

Examples:

>>> input = torch.randn(20, 5, 10, 10)
>>> # With Learnable Parameters
>>> m = nn.LayerNorm(input.size()[1:])
>>> # Without Learnable Parameters
>>> m = nn.LayerNorm(input.size()[1:], elementwise_affine=False)
>>> # Normalize over last two dimensions
>>> m = nn.LayerNorm([10, 10])
>>> # Normalize over last dimension of size 10
>>> m = nn.LayerNorm(10)
>>> # Activating the module
>>> output = m(input)

LocalResponseNorm¶

class torch.nn.LocalResponseNorm(size, alpha=0.0001, beta=0.75, k=1.0)¶

Applies local response normalization over an input signal composed of several input planes, where channels occupy the second dimension. Applies normalization across channels.

\[b_{c} = a_{c}\left(k + \frac{\alpha}{n} \sum_{c'=\max(0, c-n/2)}^{\min(N-1,c+n/2)}a_{c'}^2\right)^{-\beta} \]

Parameters

size – amount of neighbouring channels used for normalization
alpha – multiplicative factor. Default: 0.0001
beta – exponent. Default: 0.75
k – additive factor. Default: 1

Shape:

Input: $(N, C, *)$
Output: $(N, C, *)$ (same shape as input)

Examples:

>>> lrn = nn.LocalResponseNorm(2)
>>> signal_2d = torch.randn(32, 5, 24, 24)
>>> signal_4d = torch.randn(16, 5, 7, 7, 7, 7)
>>> output_2d = lrn(signal_2d)
>>> output_4d = lrn(signal_4d)

Recurrent layers¶

RNN¶

class torch.nn.RNN(*args, **kwargs)¶

Applies a multi-layer Elman RNN with $tanh$ or $ReLU$ non-linearity to an input sequence.

For each element in the input sequence, each layer computes the following function:

\[h_t = \text{tanh}(W_{ih} x_t + b_{ih} + W_{hh} h_{(t-1)} + b_{hh}) \]

where $h_t$ is the hidden state at time t, $x_t$ is the input at time t, and $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 is used instead of tanh.

Parameters

input_size – The number of expected features in the input x
hidden_size – The number of features in the hidden state h
num_layers – Number of recurrent layers. E.g., setting num_layers=2 would mean stacking two RNNs together to form a stacked RNN, with the second RNN taking in outputs of the first RNN and computing the final results. Default: 1
nonlinearity – The non-linearity to use. Can be either 'tanh' or 'relu'. Default: 'tanh'
bias – If False, then the layer does not use bias weights b_ih and b_hh. Default: True
batch_first – If True, then the input and output tensors are provided as (batch, seq, feature). Default: False
dropout – If non-zero, introduces a Dropout layer on the outputs of each RNN layer except the last layer, with dropout probability equal to dropout. Default: 0
bidirectional – If True, becomes a bidirectional RNN. Default: False

Inputs: input, h_0

input of shape (seq_len, batch, input_size): tensor containing the features of the input sequence. The input can also be a packed variable length sequence. See torch.nn.utils.rnn.pack_padded_sequence() or torch.nn.utils.rnn.pack_sequence() for details.
h_0 of shape (num_layers * num_directions, batch, hidden_size): tensor containing the initial hidden state for each element in the batch. Defaults to zero if not provided. If the RNN is bidirectional, num_directions should be 2, else it should be 1.

Outputs: output, h_n

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).

Shape:

Input1: $(L, N, H_{in})$ tensor containing input features where $H_{in}=\text{input\_size}$ and L represents a sequence length.
Input2: $(S, N, H_{out})$ tensor containing the initial hidden state for each element in the batch. $H_{out}=\text{hidden\_size}$ Defaults to zero if not provided. where $S=\text{num\_layers} * \text{num\_directions}$ If the RNN is bidirectional, num_directions should be 2, else it should be 1.
Output1: $(L, N, H_{all})$ where $H_all=\text{num\_directions} * \text{hidden\_size}$
Output2: $(S, N, H_{out})$ tensor containing the next hidden state for each element in the batch

Variables

~RNN.weight_ih_l[k] – the learnable input-hidden weights of the k-th layer, of shape (hidden_size, input_size) for k = 0. Otherwise, the shape is (hidden_size, num_directions * hidden_size)
~RNN.weight_hh_l[k] – the learnable hidden-hidden weights of the k-th layer, of shape (hidden_size, hidden_size)
~RNN.bias_ih_l[k] – the learnable input-hidden bias of the k-th layer, of shape (hidden_size)
~RNN.bias_hh_l[k] – the learnable hidden-hidden bias of the k-th layer, of shape (hidden_size)

Note

All the weights and biases are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{\text{hidden\_size}}$

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:

>>> rnn = nn.RNN(10, 20, 2)
>>> input = torch.randn(5, 3, 10)
>>> h0 = torch.randn(2, 3, 20)
>>> output, hn = rnn(input, h0)

LSTM¶

class torch.nn.LSTM(*args, **kwargs)¶

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

For each element in the input sequence, each layer computes the following function:

\[\begin{array}{ll} \\ i_t = \sigma(W_{ii} x_t + b_{ii} + W_{hi} h_{(t-1)} + b_{hi}) \\ f_t = \sigma(W_{if} x_t + b_{if} + W_{hf} h_{(t-1)} + b_{hf}) \\ g_t = \tanh(W_{ig} x_t + b_{ig} + W_{hg} h_{(t-1)} + b_{hg}) \\ o_t = \sigma(W_{io} x_t + b_{io} + W_{ho} h_{(t-1)} + b_{ho}) \\ c_t = f_t * c_{(t-1)} + i_t * g_t \\ h_t = o_t * \tanh(c_t) \\ \end{array} \]

where $h_t$ is the hidden state at time t, $c_t$ is the cell state at time t, $x_t$ is the input at time t, $h_{(t-1)}$ is the hidden state of the layer at time t-1 or the initial hidden state at time 0, and $i_t$, $f_t$, $g_t$, $o_t$ are the input, forget, cell, and output gates, respectively. $\sigma$ is the sigmoid function, and $*$ is the Hadamard product.

In a multilayer LSTM, the input $x^{(l)}_t$ of the $l$ -th layer ($l >= 2$) is the hidden state $h^{(l-1)}_t$ of the previous layer multiplied by dropout $\delta^{(l-1)}_t$ where each $\delta^{(l-1)}_t$ is a Bernoulli random variable which is $0$ with probability dropout.

Parameters

input_size – The number of expected features in the input x
hidden_size – The number of features in the hidden state h
num_layers – Number of recurrent layers. E.g., setting num_layers=2 would mean stacking two LSTMs together to form a stacked LSTM, with the second LSTM taking in outputs of the first LSTM and computing the final results. Default: 1
bias – If False, then the layer does not use bias weights b_ih and b_hh. Default: True
batch_first – If True, then the input and output tensors are provided as (batch, seq, feature). Default: False
dropout – If non-zero, introduces a Dropout layer on the outputs of each LSTM layer except the last layer, with dropout probability equal to dropout. Default: 0
bidirectional – If True, becomes a bidirectional LSTM. Default: False

Inputs: input, (h_0, c_0)

input of shape (seq_len, batch, input_size): tensor containing the features of the input sequence. The input can also be a packed variable length sequence. See torch.nn.utils.rnn.pack_padded_sequence() or torch.nn.utils.rnn.pack_sequence() for details.
h_0 of shape (num_layers * num_directions, batch, hidden_size): tensor containing the initial hidden state for each element in the batch. If the LSTM is bidirectional, num_directions should be 2, else it should be 1.
c_0 of shape (num_layers * num_directions, batch, hidden_size): tensor containing the initial cell state for each element in the batch.

If (h_0, c_0) is not provided, both h_0 and c_0 default to zero.

Outputs: output, (h_n, c_n)

output of shape (seq_len, batch, num_directions * hidden_size): tensor containing the output features (h_t) from the last layer of the LSTM, 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) and similarly for c_n.
c_n of shape (num_layers * num_directions, batch, hidden_size): tensor containing the cell state for t = seq_len.

Variables

~LSTM.weight_ih_l[k] – the learnable input-hidden weights of the $\text{k}^{th}$ layer (W_ii|W_if|W_ig|W_io), of shape (4*hidden_size, input_size) for k = 0. Otherwise, the shape is (4*hidden_size, num_directions * hidden_size)
~LSTM.weight_hh_l[k] – the learnable hidden-hidden weights of the $\text{k}^{th}$ layer (W_hi|W_hf|W_hg|W_ho), of shape (4*hidden_size, hidden_size)
~LSTM.bias_ih_l[k] – the learnable input-hidden bias of the $\text{k}^{th}$ layer (b_ii|b_if|b_ig|b_io), of shape (4*hidden_size)
~LSTM.bias_hh_l[k] – the learnable hidden-hidden bias of the $\text{k}^{th}$ layer (b_hi|b_hf|b_hg|b_ho), of shape (4*hidden_size)

Note

All the weights and biases are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{\text{hidden\_size}}$

Note

Examples:

>>> rnn = nn.LSTM(10, 20, 2)
>>> input = torch.randn(5, 3, 10)
>>> h0 = torch.randn(2, 3, 20)
>>> c0 = torch.randn(2, 3, 20)
>>> output, (hn, cn) = rnn(input, (h0, c0))

GRU¶

class torch.nn.GRU(*args, **kwargs)¶

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

For each element in the input sequence, each layer computes the following function:

\[\begin{array}{ll} r_t = \sigma(W_{ir} x_t + b_{ir} + W_{hr} h_{(t-1)} + b_{hr}) \\ z_t = \sigma(W_{iz} x_t + b_{iz} + W_{hz} h_{(t-1)} + b_{hz}) \\ n_t = \tanh(W_{in} x_t + b_{in} + r_t * (W_{hn} h_{(t-1)}+ b_{hn})) \\ h_t = (1 - z_t) * n_t + z_t * h_{(t-1)} \end{array} \]

where $h_t$ is the hidden state at time t, $x_t$ is the input at time t, $h_{(t-1)}$ is the hidden state of the layer at time t-1 or the initial hidden state at time 0, and $r_t$, $z_t$, $n_t$ are the reset, update, and new gates, respectively. $\sigma$ is the sigmoid function, and $*$ is the Hadamard product.

In a multilayer GRU, the input $x^{(l)}_t$ of the $l$ -th layer ($l >= 2$) is the hidden state $h^{(l-1)}_t$ of the previous layer multiplied by dropout $\delta^{(l-1)}_t$ where each $\delta^{(l-1)}_t$ is a Bernoulli random variable which is $0$ with probability dropout.

Parameters

input_size – The number of expected features in the input x
hidden_size – The number of features in the hidden state h
num_layers – Number of recurrent layers. E.g., setting num_layers=2 would mean stacking two GRUs together to form a stacked GRU, with the second GRU taking in outputs of the first GRU and computing the final results. Default: 1
bias – If False, then the layer does not use bias weights b_ih and b_hh. Default: True
batch_first – If True, then the input and output tensors are provided as (batch, seq, feature). Default: False
dropout – If non-zero, introduces a Dropout layer on the outputs of each GRU layer except the last layer, with dropout probability equal to dropout. Default: 0
bidirectional – If True, becomes a bidirectional GRU. Default: False

Inputs: input, h_0

input of shape (seq_len, batch, input_size): tensor containing the features of the input sequence. The input can also be a packed variable length sequence. See torch.nn.utils.rnn.pack_padded_sequence() for details.
h_0 of shape (num_layers * num_directions, batch, hidden_size): tensor containing the initial hidden state for each element in the batch. Defaults to zero if not provided. If the RNN is bidirectional, num_directions should be 2, else it should be 1.

Outputs: output, h_n

output of shape (seq_len, batch, num_directions * hidden_size): tensor containing the output features h_t from the last layer of the GRU, 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).

Shape:

Input1: $(L, N, H_{in})$ tensor containing input features where $H_{in}=\text{input\_size}$ and L represents a sequence length.
Input2: $(S, N, H_{out})$ tensor containing the initial hidden state for each element in the batch. $H_{out}=\text{hidden\_size}$ Defaults to zero if not provided. where $S=\text{num\_layers} * \text{num\_directions}$ If the RNN is bidirectional, num_directions should be 2, else it should be 1.
Output1: $(L, N, H_{all})$ where $H_all=\text{num\_directions} * \text{hidden\_size}$
Output2: $(S, N, H_{out})$ tensor containing the next hidden state for each element in the batch

Variables

~GRU.weight_ih_l[k] – the learnable input-hidden weights of the $\text{k}^{th}$ layer (W_ir|W_iz|W_in), of shape (3*hidden_size, input_size) for k = 0. Otherwise, the shape is (3*hidden_size, num_directions * hidden_size)
~GRU.weight_hh_l[k] – the learnable hidden-hidden weights of the $\text{k}^{th}$ layer (W_hr|W_hz|W_hn), of shape (3*hidden_size, hidden_size)
~GRU.bias_ih_l[k] – the learnable input-hidden bias of the $\text{k}^{th}$ layer (b_ir|b_iz|b_in), of shape (3*hidden_size)
~GRU.bias_hh_l[k] – the learnable hidden-hidden bias of the $\text{k}^{th}$ layer (b_hr|b_hz|b_hn), of shape (3*hidden_size)

Note

All the weights and biases are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{\text{hidden\_size}}$

Note

Examples:

>>> rnn = nn.GRU(10, 20, 2)
>>> input = torch.randn(5, 3, 10)
>>> h0 = torch.randn(2, 3, 20)
>>> output, hn = rnn(input, h0)

RNNCell¶

class torch.nn.RNNCell(input_size, hidden_size, bias=True, nonlinearity='tanh')¶

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

\[h' = \tanh(W_{ih} x + b_{ih} + W_{hh} h + b_{hh})\]

If nonlinearity is ‘relu’, then ReLU is used in place of tanh.

Parameters

input_size – The number of expected features in the input x
hidden_size – The number of features in the hidden state h
bias – If False, then the layer does not use bias weights b_ih and b_hh. Default: True
nonlinearity – The non-linearity to use. Can be either 'tanh' or 'relu'. Default: 'tanh'

Inputs: input, hidden

input of shape (batch, input_size): tensor containing input features
hidden of shape (batch, hidden_size): tensor containing the initial hidden state for each element in the batch. Defaults to zero if not provided.

Outputs: h’

h’ of shape (batch, hidden_size): tensor containing the next hidden state for each element in the batch

Shape:

Input1: $(N, H_{in})$ tensor containing input features where $H_{in}$ = input_size
Input2: $(N, H_{out})$ tensor containing the initial hidden state for each element in the batch where $H_{out}$ = hidden_size Defaults to zero if not provided.
Output: $(N, H_{out})$ tensor containing the next hidden state for each element in the batch

Variables

~RNNCell.weight_ih – the learnable input-hidden weights, of shape (hidden_size, input_size)
~RNNCell.weight_hh – the learnable hidden-hidden weights, of shape (hidden_size, hidden_size)
~RNNCell.bias_ih – the learnable input-hidden bias, of shape (hidden_size)
~RNNCell.bias_hh – the learnable hidden-hidden bias, of shape (hidden_size)

Note

All the weights and biases are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{\text{hidden\_size}}$

Examples:

>>> rnn = nn.RNNCell(10, 20)
>>> input = torch.randn(6, 3, 10)
>>> hx = torch.randn(3, 20)
>>> output = []
>>> for i in range(6):
        hx = rnn(input[i], hx)
        output.append(hx)

LSTMCell¶

class torch.nn.LSTMCell(input_size, hidden_size, bias=True)¶

A long short-term memory (LSTM) cell.

\[\begin{array}{ll} i = \sigma(W_{ii} x + b_{ii} + W_{hi} h + b_{hi}) \\ f = \sigma(W_{if} x + b_{if} + W_{hf} h + b_{hf}) \\ g = \tanh(W_{ig} x + b_{ig} + W_{hg} h + b_{hg}) \\ o = \sigma(W_{io} x + b_{io} + W_{ho} h + b_{ho}) \\ c' = f * c + i * g \\ h' = o * \tanh(c') \\ \end{array}\]

where $\sigma$ is the sigmoid function, and $*$ is the Hadamard product.

Parameters

input_size – The number of expected features in the input x
hidden_size – The number of features in the hidden state h
bias – If False, then the layer does not use bias weights b_ih and b_hh. Default: True

Inputs: input, (h_0, c_0)

input of shape (batch, input_size): tensor containing input features
h_0 of shape (batch, hidden_size): tensor containing the initial hidden state for each element in the batch.
c_0 of shape (batch, hidden_size): tensor containing the initial cell state for each element in the batch.

If (h_0, c_0) is not provided, both h_0 and c_0 default to zero.

Outputs: (h_1, c_1)

h_1 of shape (batch, hidden_size): tensor containing the next hidden state for each element in the batch
c_1 of shape (batch, hidden_size): tensor containing the next cell state for each element in the batch

Variables

~LSTMCell.weight_ih – the learnable input-hidden weights, of shape (4*hidden_size, input_size)
~LSTMCell.weight_hh – the learnable hidden-hidden weights, of shape (4*hidden_size, hidden_size)
~LSTMCell.bias_ih – the learnable input-hidden bias, of shape (4*hidden_size)
~LSTMCell.bias_hh – the learnable hidden-hidden bias, of shape (4*hidden_size)

Note

All the weights and biases are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{\text{hidden\_size}}$

Examples:

>>> rnn = nn.LSTMCell(10, 20)
>>> input = torch.randn(6, 3, 10)
>>> hx = torch.randn(3, 20)
>>> cx = torch.randn(3, 20)
>>> output = []
>>> for i in range(6):
        hx, cx = rnn(input[i], (hx, cx))
        output.append(hx)

GRUCell¶

class torch.nn.GRUCell(input_size, hidden_size, bias=True)¶

A gated recurrent unit (GRU) cell

\[\begin{array}{ll} r = \sigma(W_{ir} x + b_{ir} + W_{hr} h + b_{hr}) \\ z = \sigma(W_{iz} x + b_{iz} + W_{hz} h + b_{hz}) \\ n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\ h' = (1 - z) * n + z * h \end{array}\]

where $\sigma$ is the sigmoid function, and $*$ is the Hadamard product.

Parameters

input_size – The number of expected features in the input x
hidden_size – The number of features in the hidden state h
bias – If False, then the layer does not use bias weights b_ih and b_hh. Default: True

Inputs: input, hidden

input of shape (batch, input_size): tensor containing input features
hidden of shape (batch, hidden_size): tensor containing the initial hidden state for each element in the batch. Defaults to zero if not provided.

Outputs: h’

h’ of shape (batch, hidden_size): tensor containing the next hidden state for each element in the batch

Shape:

Input1: $(N, H_{in})$ tensor containing input features where $H_{in}$ = input_size
Input2: $(N, H_{out})$ tensor containing the initial hidden state for each element in the batch where $H_{out}$ = hidden_size Defaults to zero if not provided.
Output: $(N, H_{out})$ tensor containing the next hidden state for each element in the batch

Variables

~GRUCell.weight_ih – the learnable input-hidden weights, of shape (3*hidden_size, input_size)
~GRUCell.weight_hh – the learnable hidden-hidden weights, of shape (3*hidden_size, hidden_size)
~GRUCell.bias_ih – the learnable input-hidden bias, of shape (3*hidden_size)
~GRUCell.bias_hh – the learnable hidden-hidden bias, of shape (3*hidden_size)

Note

All the weights and biases are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{\text{hidden\_size}}$

Examples:

>>> rnn = nn.GRUCell(10, 20)
>>> input = torch.randn(6, 3, 10)
>>> hx = torch.randn(3, 20)
>>> output = []
>>> for i in range(6):
        hx = rnn(input[i], hx)
        output.append(hx)

Linear layers¶

Linear¶

class torch.nn.Linear(in_features, out_features, bias=True)¶

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

Parameters

in_features – size of each input sample
out_features – size of each output sample
bias – If set to False, the layer will not learn an additive bias. Default: True

Shape:

Input: $(N, *, H_{in})$ where $*$ means any number of additional dimensions and $H_{in} = \text{in\_features}$
Output: $(N, *, H_{out})$ where all but the last dimension are the same shape as the input and $H_{out} = \text{out\_features}$.

Variables

~Linear.weight – the learnable weights of the module of shape $(\text{out\_features}, \text{in\_features})$. The values are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$, where $k = \frac{1}{\text{in\_features}}$
~Linear.bias – the learnable bias of the module of shape $(\text{out\_features})$. If bias is True, the values are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$ where $k = \frac{1}{\text{in\_features}}$

Examples:

>>> m = nn.Linear(20, 30)
>>> input = torch.randn(128, 20)
>>> output = m(input)
>>> print(output.size())
torch.Size([128, 30])

Bilinear¶

class torch.nn.Bilinear(in1_features, in2_features, out_features, bias=True)¶

Applies a bilinear transformation to the incoming data: $y = x_1 A x_2 + b$

Parameters

in1_features – size of each first input sample
in2_features – size of each second input sample
out_features – size of each output sample
bias – If set to False, the layer will not learn an additive bias. Default: True

Shape:

Input1: $(N, *, H_{in1})$ where $H_{in1}=\text{in1\_features}$ and $*$ means any number of additional dimensions. All but the last dimension of the inputs should be the same.
Input2: $(N, *, H_{in2})$ where $H_{in2}=\text{in2\_features}$.
Output: $(N, *, H_{out})$ where $H_{out}=\text{out\_features}$ and all but the last dimension are the same shape as the input.

Variables

~Bilinear.weight – the learnable weights of the module of shape $(\text{out\_features}, \text{in1\_features}, \text{in2\_features})$. The values are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$, where $k = \frac{1}{\text{in1\_features}}$
~Bilinear.bias – the learnable bias of the module of shape $(\text{out\_features})$. If bias is True, the values are initialized from $\mathcal{U}(-\sqrt{k}, \sqrt{k})$, where $k = \frac{1}{\text{in1\_features}}$

Examples:

>>> m = nn.Bilinear(20, 30, 40)
>>> input1 = torch.randn(128, 20)
>>> input2 = torch.randn(128, 30)
>>> output = m(input1, input2)
>>> print(output.size())
torch.Size([128, 40])

Dropout layers¶

Dropout¶

class torch.nn.Dropout(p=0.5, inplace=False)¶

During training, randomly zeroes some of the elements of the input tensor with probability p using samples from a Bernoulli distribution. Each channel will be zeroed out independently on every forward call.

This has proven to be an effective technique for regularization and preventing the co-adaptation of neurons as described in the paper Improving neural networks by preventing co-adaptation of feature detectors .

Furthermore, the outputs are scaled by a factor of $\frac{1}{1-p}$ during training. This means that during evaluation the module simply computes an identity function.

Parameters

p – probability of an element to be zeroed. Default: 0.5
inplace – If set to True, will do this operation in-place. Default: False

Shape:

Input: $(*)$. Input can be of any shape
Output: $(*)$. Output is of the same shape as input

Examples:

>>> m = nn.Dropout(p=0.2)
>>> input = torch.randn(20, 16)
>>> output = m(input)

Dropout2d¶

class torch.nn.Dropout2d(p=0.5, inplace=False)¶

Randomly zero out entire channels (a channel is a 2D feature map, e.g., the $j$-th channel of the $i$-th sample in the batched input is a 2D tensor $\text{input}[i, j]$). Each channel will be zeroed out independently on every forward call with probability p using samples from a Bernoulli distribution.

Usually the input comes from nn.Conv2d modules.

As described in the paper Efficient Object Localization Using Convolutional Networks , if adjacent pixels within feature maps are strongly correlated (as is normally the case in early convolution layers) then i.i.d. dropout will not regularize the activations and will otherwise just result in an effective learning rate decrease.

In this case, nn.Dropout2d() will help promote independence between feature maps and should be used instead.

Parameters

p (float, optional) – probability of an element to be zero-ed.
inplace (bool, optional) – If set to True, will do this operation in-place

Shape:

Input: $(N, C, H, W)$
Output: $(N, C, H, W)$ (same shape as input)

Examples:

>>> m = nn.Dropout2d(p=0.2)
>>> input = torch.randn(20, 16, 32, 32)
>>> output = m(input)

Dropout3d¶

class torch.nn.Dropout3d(p=0.5, inplace=False)¶

Randomly zero out entire channels (a channel is a 3D feature map, e.g., the $j$-th channel of the $i$-th sample in the batched input is a 3D tensor $\text{input}[i, j]$). Each channel will be zeroed out independently on every forward call with probability p using samples from a Bernoulli distribution.

Usually the input comes from nn.Conv3d modules.

In this case, nn.Dropout3d() will help promote independence between feature maps and should be used instead.

Parameters

p (float, optional) – probability of an element to be zeroed.
inplace (bool, optional) – If set to True, will do this operation in-place

Shape:

Input: $(N, C, D, H, W)$
Output: $(N, C, D, H, W)$ (same shape as input)

Examples:

>>> m = nn.Dropout3d(p=0.2)
>>> input = torch.randn(20, 16, 4, 32, 32)
>>> output = m(input)

AlphaDropout¶

class torch.nn.AlphaDropout(p=0.5, inplace=False)¶

Applies Alpha Dropout over the input.

Alpha Dropout is a type of Dropout that maintains the self-normalizing property. For an input with zero mean and unit standard deviation, the output of Alpha Dropout maintains the original mean and standard deviation of the input. Alpha Dropout goes hand-in-hand with SELU activation function, which ensures that the outputs have zero mean and unit standard deviation.

During training, it randomly masks some of the elements of the input tensor with probability p using samples from a bernoulli distribution. The elements to masked are randomized on every forward call, and scaled and shifted to maintain zero mean and unit standard deviation.

During evaluation the module simply computes an identity function.

More details can be found in the paper Self-Normalizing Neural Networks .

Parameters

p (float) – probability of an element to be dropped. Default: 0.5
inplace (bool, optional) – If set to True, will do this operation in-place

Shape:

Input: $(*)$. Input can be of any shape
Output: $(*)$. Output is of the same shape as input

Examples:

>>> m = nn.AlphaDropout(p=0.2)
>>> input = torch.randn(20, 16)
>>> output = m(input)

Sparse layers¶

Embedding¶

class torch.nn.Embedding(num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, _weight=None)¶

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

This module is often used to store word embeddings and retrieve them using indices. The input to the module is a list of indices, and the output is the corresponding word embeddings.

Parameters

num_embeddings (int) – size of the dictionary of embeddings
embedding_dim (int) – the size of each embedding vector
padding_idx (int, optional) – If given, pads the output with the embedding vector at padding_idx (initialized to zeros) whenever it encounters the index.
max_norm (float, optional) – If given, each embedding vector with norm larger than max_norm is renormalized to have norm max_norm.
norm_type (float, optional) – The p of the p-norm to compute for the max_norm option. Default 2.
scale_grad_by_freq (boolean, optional) – If given, this will scale gradients by the inverse of frequency of the words in the mini-batch. Default False.
sparse (bool, optional) – If True, gradient w.r.t. weight matrix will be a sparse tensor. See Notes for more details regarding sparse gradients.

Variables

~Embedding.weight (Tensor) – the learnable weights of the module of shape (num_embeddings, embedding_dim) initialized from $\mathcal{N}(0, 1)$

Shape:

Input: $(*)$, LongTensor of arbitrary shape containing the indices to extract
Output: $(*, H)$, where * is the input shape and $H=\text{embedding\_dim}$

Note

Keep in mind that only a limited number of optimizers support sparse gradients: currently it’s optim.SGD (CUDA and CPU), optim.SparseAdam (CUDA and CPU) and optim.Adagrad (CPU)

Note

With padding_idx set, the embedding vector at padding_idx is initialized to all zeros. However, note that this vector can be modified afterwards, e.g., using a customized initialization method, and thus changing the vector used to pad the output. The gradient for this vector from Embedding is always zero.

Examples:

>>> # an Embedding module containing 10 tensors of size 3
>>> embedding = nn.Embedding(10, 3)
>>> # a batch of 2 samples of 4 indices each
>>> input = torch.LongTensor([[1,2,4,5],[4,3,2,9]])
>>> embedding(input)
tensor([[[-0.0251, -1.6902,  0.7172],
         [-0.6431,  0.0748,  0.6969],
         [ 1.4970,  1.3448, -0.9685],
         [-0.3677, -2.7265, -0.1685]],

        [[ 1.4970,  1.3448, -0.9685],
         [ 0.4362, -0.4004,  0.9400],
         [-0.6431,  0.0748,  0.6969],
         [ 0.9124, -2.3616,  1.1151]]])


>>> # example with padding_idx
>>> embedding = nn.Embedding(10, 3, padding_idx=0)
>>> input = torch.LongTensor([[0,2,0,5]])
>>> embedding(input)
tensor([[[ 0.0000,  0.0000,  0.0000],
         [ 0.1535, -2.0309,  0.9315],
         [ 0.0000,  0.0000,  0.0000],
         [-0.1655,  0.9897,  0.0635]]])

classmethod from_pretrained(embeddings, freeze=True, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False)¶

Creates Embedding instance from given 2-dimensional FloatTensor.

Parameters

embeddings (Tensor) – FloatTensor containing weights for the Embedding. First dimension is being passed to Embedding as num_embeddings, second as embedding_dim.
freeze (boolean, optional) – If True, the tensor does not get updated in the learning process. Equivalent to embedding.weight.requires_grad = False. Default: True
padding_idx (int, optional) – See module initialization documentation.
max_norm (float, optional) – See module initialization documentation.
norm_type (float, optional) – See module initialization documentation. Default 2.
scale_grad_by_freq (boolean, optional) – See module initialization documentation. Default False.
sparse (bool, optional) – See module initialization documentation.

Examples:

>>> # FloatTensor containing pretrained weights
>>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])
>>> embedding = nn.Embedding.from_pretrained(weight)
>>> # Get embeddings for index 1
>>> input = torch.LongTensor([1])
>>> embedding(input)
tensor([[ 4.0000,  5.1000,  6.3000]])

EmbeddingBag¶

class torch.nn.EmbeddingBag(num_embeddings, embedding_dim, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, mode='mean', sparse=False, _weight=None)¶

Computes sums or means of ‘bags’ of embeddings, without instantiating the intermediate embeddings.

For bags of constant length, this class

with mode="sum" is equivalent to Embedding followed by torch.sum(dim=0),

with mode="mean" is equivalent to Embedding followed by torch.mean(dim=0),

with mode="max" is equivalent to Embedding followed by torch.max(dim=0).

However, EmbeddingBag is much more time and memory efficient than using a chain of these operations.

Parameters

num_embeddings (int) – size of the dictionary of embeddings
embedding_dim (int) – the size of each embedding vector
max_norm (float, optional) – If given, each embedding vector with norm larger than max_norm is renormalized to have norm max_norm.
norm_type (float, optional) – The p of the p-norm to compute for the max_norm option. Default 2.
scale_grad_by_freq (boolean, optional) – if given, this will scale gradients by the inverse of frequency of the words in the mini-batch. Default False. Note: this option is not supported when mode="max".
mode (string, optional) – "sum", "mean" or "max". Specifies the way to reduce the bag. Default: "mean"
sparse (bool, optional) – if True, gradient w.r.t. weight matrix will be a sparse tensor. See Notes for more details regarding sparse gradients. Note: this option is not supported when mode="max".

Variables

~EmbeddingBag.weight (Tensor) – the learnable weights of the module of shape (num_embeddings, embedding_dim) initialized from $\mathcal{N}(0, 1)$.

Inputs: input (LongTensor) and offsets (LongTensor, optional)

If input is 2D of shape (B, N),

it will be treated as B bags (sequences) each of fixed length N, and this will return B values aggregated in a way depending on the mode. offsets is ignored and required to be None in this case.

If input is 1D of shape (N),

it will be treated as a concatenation of multiple bags (sequences). offsets is required to be a 1D tensor containing the starting index positions of each bag in input. Therefore, for offsets of shape (B), input will be viewed as having B bags. Empty bags (i.e., having 0-length) will have returned vectors filled by zeros.

Output shape: (B, embedding_dim)

Examples:

>>> # an Embedding module containing 10 tensors of size 3
>>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum')
>>> # a batch of 2 samples of 4 indices each
>>> input = torch.LongTensor([1,2,4,5,4,3,2,9])
>>> offsets = torch.LongTensor([0,4])
>>> embedding_sum(input, offsets)
tensor([[-0.8861, -5.4350, -0.0523],
        [ 1.1306, -2.5798, -1.0044]])

classmethod from_pretrained(embeddings, freeze=True, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, mode='mean', sparse=False)¶

Creates EmbeddingBag instance from given 2-dimensional FloatTensor.

Parameters

embeddings (Tensor) – FloatTensor containing weights for the EmbeddingBag. First dimension is being passed to EmbeddingBag as ‘num_embeddings’, second as ‘embedding_dim’.
freeze (boolean, optional) – If True, the tensor does not get updated in the learning process. Equivalent to embeddingbag.weight.requires_grad = False. Default: True
max_norm (float, optional) – See module initialization documentation. Default: None
norm_type (float, optional) – See module initialization documentation. Default 2.
scale_grad_by_freq (boolean, optional) – See module initialization documentation. Default False.
mode (string, optional) – See module initialization documentation. Default: "mean"
sparse (bool, optional) – See module initialization documentation. Default: False.

Examples:

>>> # FloatTensor containing pretrained weights
>>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])
>>> embeddingbag = nn.EmbeddingBag.from_pretrained(weight)
>>> # Get embeddings for index 1
>>> input = torch.LongTensor([[1, 0]])
>>> embeddingbag(input)
tensor([[ 2.5000,  3.7000,  4.6500]])

Distance functions¶

CosineSimilarity¶

class torch.nn.CosineSimilarity(dim=1, eps=1e-08)¶

Returns cosine similarity between $x_1$ and $x_2$, computed along dim.

\[\text{similarity} = \dfrac{x_1 \cdot x_2}{\max(\Vert x_1 \Vert _2 \cdot \Vert x_2 \Vert _2, \epsilon)} \]

Parameters

dim (int, optional) – Dimension where cosine similarity is computed. Default: 1
eps (float, optional) – Small value to avoid division by zero. Default: 1e-8

Shape:

Input1: $(\ast_1, D, \ast_2)$ where D is at position dim
Input2: $(\ast_1, D, \ast_2)$, same shape as the Input1
Output: $(\ast_1, \ast_2)$

Examples:

>>> input1 = torch.randn(100, 128)
>>> input2 = torch.randn(100, 128)
>>> cos = nn.CosineSimilarity(dim=1, eps=1e-6)
>>> output = cos(input1, input2)

PairwiseDistance¶

class torch.nn.PairwiseDistance(p=2.0, eps=1e-06, keepdim=False)¶

Computes the batchwise pairwise distance between vectors $v_1$, $v_2$ using the p-norm:

\[\Vert x \Vert _p = \left( \sum_{i=1}^n \vert x_i \vert ^ p \right) ^ {1/p} \]

Parameters

p (real) – the norm degree. Default: 2
eps (float, optional) – Small value to avoid division by zero. Default: 1e-6
keepdim (bool, optional) – Determines whether or not to keep the batch dimension. Default: False

Shape:

Input1: $(N, D)$ where D = vector dimension
Input2: $(N, D)$, same shape as the Input1
Output: $(N)$. If keepdim is False, then $(N, 1)$.

Examples:

>>> pdist = nn.PairwiseDistance(p=2)
>>> input1 = torch.randn(100, 128)
>>> input2 = torch.randn(100, 128)
>>> output = pdist(input1, input2)

Loss functions¶

L1Loss¶

class torch.nn.L1Loss(size_average=None, reduce=None, reduction='mean')¶

Creates a criterion that measures the mean absolute error (MAE) between each element in the input $x$ and target $y$.

The unreduced (i.e. with reduction set to 'none') loss can be described as:

\[\ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad l_n = \left| x_n - y_n \right|, \]

where $N$ is the batch size. If reduction is not 'none' (default 'mean'), then:

\[\ell(x, y) = \begin{cases} \operatorname{mean}(L), & \text{if reduction} = \text{'mean';}\\ \operatorname{sum}(L), & \text{if reduction} = \text{'sum'.} \end{cases} \]

$x$ and $y$ are tensors of arbitrary shapes with a total of $n$ elements each.

The sum operation still operates over all the elements, and divides by $n$.

The division by $n$ can be avoided if one sets reduction = 'sum'.

Parameters

size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(N, *)$ where $*$ means, any number of additional dimensions
Target: $(N, *)$, same shape as the input
Output: scalar. If reduction is 'none', then $(N, *)$, same shape as the input

Examples:

>>> loss = nn.L1Loss()
>>> input = torch.randn(3, 5, requires_grad=True)
>>> target = torch.randn(3, 5)
>>> output = loss(input, target)
>>> output.backward()

MSELoss¶

class torch.nn.MSELoss(size_average=None, reduce=None, reduction='mean')¶

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

The unreduced (i.e. with reduction set to 'none') loss can be described as:

\[\ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad l_n = \left( x_n - y_n \right)^2, \]

where $N$ is the batch size. If reduction is not 'none' (default 'mean'), then:

\[\ell(x, y) = \begin{cases} \operatorname{mean}(L), & \text{if reduction} = \text{'mean';}\\ \operatorname{sum}(L), & \text{if reduction} = \text{'sum'.} \end{cases} \]

$x$ and $y$ are tensors of arbitrary shapes with a total of $n$ elements each.

The sum operation still operates over all the elements, and divides by $n$.

The division by $n$ can be avoided if one sets reduction = 'sum'.

Parameters

size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(N, *)$ where $*$ means, any number of additional dimensions
Target: $(N, *)$, same shape as the input

Examples:

>>> loss = nn.MSELoss()
>>> input = torch.randn(3, 5, requires_grad=True)
>>> target = torch.randn(3, 5)
>>> output = loss(input, target)
>>> output.backward()

CrossEntropyLoss¶

class torch.nn.CrossEntropyLoss(weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')¶

This criterion combines nn.LogSoftmax() and nn.NLLLoss() in one single class.

It is useful when training 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 is expected to contain raw, unnormalized scores for each class.

input has to be a Tensor of size either $(minibatch, C)$ or $(minibatch, C, d_1, d_2, ..., d_K)$ with $K \geq 1$ for the K-dimensional case (described later).

This criterion expects a class index in the range $[0, C-1]$ as the target for each value of a 1D tensor of size minibatch.

The loss can be described as:

\[\text{loss}(x, class) = -\log\left(\frac{\exp(x[class])}{\sum_j \exp(x[j])}\right) = -x[class] + \log\left(\sum_j \exp(x[j])\right) \]

or in the case of the weight argument being specified:

\[\text{loss}(x, class) = weight[class] \left(-x[class] + \log\left(\sum_j \exp(x[j])\right)\right) \]

The losses are averaged across observations for each minibatch.

Can also be used for higher dimension inputs, such as 2D images, by providing an input of size $(minibatch, C, d_1, d_2, ..., d_K)$ with $K \geq 1$, where $K$ is the number of dimensions, and a target of appropriate shape (see below).

Parameters

weight (Tensor, optional) – a manual rescaling weight given to each class. If given, has to be a Tensor of size C
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
ignore_index (int, optional) – Specifies a target value that is ignored and does not contribute to the input gradient. When size_average is True, the loss is averaged over non-ignored targets.
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(N, C)$ where C = number of classes, or $(N, C, d_1, d_2, ..., d_K)$ with $K \geq 1$ in the case of K-dimensional loss.
Target: $(N)$ where each value is $0 \leq \text{targets}[i] \leq C-1$, or $(N, d_1, d_2, ..., d_K)$ with $K \geq 1$ in the case of K-dimensional loss.
Output: scalar. If reduction is 'none', then the same size as the target: $(N)$, or $(N, d_1, d_2, ..., d_K)$ with $K \geq 1$ in the case of K-dimensional loss.

Examples:

>>> loss = nn.CrossEntropyLoss()
>>> input = torch.randn(3, 5, requires_grad=True)
>>> target = torch.empty(3, dtype=torch.long).random_(5)
>>> output = loss(input, target)
>>> output.backward()

CTCLoss¶

class torch.nn.CTCLoss(blank=0, reduction='mean', zero_infinity=False)¶

The Connectionist Temporal Classification loss.

Parameters

blank (int, optional) – blank label. Default $0$.
reduction (string, optional) – Specifies the reduction to apply to the output: ‘none’ | ‘mean’ | ‘sum’. ‘none’: no reduction will be applied, ‘mean’: the output losses will be divided by the target lengths and then the mean over the batch is taken. Default: ‘mean’
zero_infinity (bool, optional) – Whether to zero infinite losses and the associated gradients. Default: False Infinite losses mainly occur when the inputs are too short to be aligned to the targets.

Inputs:

log_probs: Tensor of size $(T, N, C)$ where C = number of characters in alphabet including blank,: T = input length, and N = batch size. The logarithmized probabilities of the outputs (e.g. obtained with torch.nn.functional.log_softmax()).
targets: Tensor of size $(N, S)$ or (sum(target_lengths)).: Targets (cannot be blank). In the second form, the targets are assumed to be concatenated.
input_lengths: Tuple or tensor of size $(N)$.: Lengths of the inputs (must each be $\leq T$)
target_lengths: Tuple or tensor of size $(N)$.: Lengths of the targets

Example:

>>> ctc_loss = nn.CTCLoss()
>>> log_probs = torch.randn(50, 16, 20).log_softmax(2).detach().requires_grad_()
>>> targets = torch.randint(1, 20, (16, 30), dtype=torch.long)
>>> input_lengths = torch.full((16,), 50, dtype=torch.long)
>>> target_lengths = torch.randint(10,30,(16,), dtype=torch.long)
>>> loss = ctc_loss(log_probs, targets, input_lengths, target_lengths)
>>> loss.backward()

Reference:: A. Graves et al.: Connectionist Temporal Classification: Labelling Unsegmented Sequence Data with Recurrent Neural Networks: https://www.cs.toronto.edu/~graves/icml_2006.pdf

Note

In order to use CuDNN, the following must be satisfied: targets must be in concatenated format, all input_lengths must be T. $blank=0$, target_lengths $\leq 256$, the integer arguments must be of dtype torch.int32.

The regular implementation uses the (more common in PyTorch) torch.long dtype.

Note

NLLLoss¶

class torch.nn.NLLLoss(weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')¶

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)$ or $(minibatch, C, d_1, d_2, ..., d_K)$ with $K \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 is a class index in the range $[0, C-1]$ where C = number of classes.

The unreduced (i.e. with reduction set to 'none') loss can be described as:

\[\ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad l_n = - w_{y_n} x_{n,y_n}, \quad w_{c} = \text{weight}[c] \cdot \mathbb{1}\{c \not= \text{ignore\_index}\}, \]

where $N$ is the batch size. If reduction is not 'none' (default 'mean'), then

\[\ell(x, y) = \begin{cases} \sum_{n=1}^N \frac{1}{\sum_{n=1}^N w_{y_n}} l_n, & \text{if reduction} = \text{'mean';}\\ \sum_{n=1}^N l_n, & \text{if reduction} = \text{'sum'.} \end{cases} \]

Parameters

weight (Tensor, optional) – a manual rescaling weight given to each class. If given, it has to be a Tensor of size C. Otherwise, it is treated as if having all ones.
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
ignore_index (int, optional) – Specifies a target value that is ignored and does not contribute to the input gradient. When size_average is True, the loss is averaged over non-ignored targets.
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(N, C)$ where C = number of classes, or $(N, C, d_1, d_2, ..., d_K)$ with $K \geq 1$ in the case of K-dimensional loss.
Target: $(N)$ where each value is $0 \leq \text{targets}[i] \leq C-1$, or $(N, d_1, d_2, ..., d_K)$ with $K \geq 1$ in the case of K-dimensional loss.
Output: scalar. If reduction is 'none', then the same size as the target: $(N)$, or $(N, d_1, d_2, ..., d_K)$ with $K \geq 1$ in the case of K-dimensional loss.

Examples:

>>> m = nn.LogSoftmax(dim=1)
>>> loss = nn.NLLLoss()
>>> # input is of size N x C = 3 x 5
>>> input = torch.randn(3, 5, requires_grad=True)
>>> # each element in target has to have 0 <= value < C
>>> target = torch.tensor([1, 0, 4])
>>> output = loss(m(input), target)
>>> output.backward()
>>>
>>>
>>> # 2D loss example (used, for example, with image inputs)
>>> N, C = 5, 4
>>> loss = nn.NLLLoss()
>>> # input is of size N x C x height x width
>>> data = torch.randn(N, 16, 10, 10)
>>> conv = nn.Conv2d(16, C, (3, 3))
>>> m = nn.LogSoftmax(dim=1)
>>> # each element in target has to have 0 <= value < C
>>> target = torch.empty(N, 8, 8, dtype=torch.long).random_(0, C)
>>> output = loss(m(conv(data)), target)
>>> output.backward()

PoissonNLLLoss¶

class torch.nn.PoissonNLLLoss(log_input=True, full=False, size_average=None, eps=1e-08, reduce=None, reduction='mean')¶

Negative log likelihood loss with Poisson distribution of target.

The loss can be described as:

\[\text{target} \sim \mathrm{Poisson}(\text{input}) \text{loss}(\text{input}, \text{target}) = \text{input} - \text{target} * \log(\text{input}) + \log(\text{target!})\]

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.

Parameters

log_input (bool, optional) – if True the loss is computed as $\exp(\text{input}) - \text{target}*\text{input}$, if False the loss is $\text{input} - \text{target}*\log(\text{input}+\text{eps})$.
full (bool, optional) –
whether to compute full loss, i. e. to add the Stirling approximation term

\[\text{target}*\log(\text{target}) - \text{target} + 0.5 * \log(2\pi\text{target}). \]
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
eps (float, optional) – Small value to avoid evaluation of $\log(0)$ when log_input = False. Default: 1e-8
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Examples:

>>> loss = nn.PoissonNLLLoss()
>>> log_input = torch.randn(5, 2, requires_grad=True)
>>> target = torch.randn(5, 2)
>>> output = loss(log_input, target)
>>> output.backward()

Shape:

Input: $(N, *)$ where $*$ means, any number of additional dimensions
Target: $(N, *)$, same shape as the input
Output: scalar by default. If reduction is 'none', then $(N, *)$, the same shape as the input

KLDivLoss¶

class torch.nn.KLDivLoss(size_average=None, reduce=None, reduction='mean')¶

The Kullback-Leibler divergence Loss

KL divergence is a useful distance measure for continuous distributions and is often useful when performing direct regression over the space of (discretely sampled) continuous output distributions.

As with NLLLoss, the input given is expected to contain log-probabilities and is not restricted to a 2D Tensor. The targets are given as probabilities (i.e. without taking the logarithm).

This criterion expects a target Tensor of the same size as the input Tensor.

The unreduced (i.e. with reduction set to 'none') loss can be described as:

\[l(x,y) = L = \{ l_1,\dots,l_N \}, \quad l_n = y_n \cdot \left( \log y_n - x_n \right) \]

where the index $N$ spans all dimensions of input and $L$ has the same shape as input. If reduction is not 'none' (default 'mean'), then:

\[\ell(x, y) = \begin{cases} \operatorname{mean}(L), & \text{if reduction} = \text{'mean';} \\ \operatorname{sum}(L), & \text{if reduction} = \text{'sum'.} \end{cases} \]

In default reduction mode 'mean', the losses are averaged for each minibatch over observations as well as over dimensions. 'batchmean' mode gives the correct KL divergence where losses are averaged over batch dimension only. 'mean' mode’s behavior will be changed to the same as 'batchmean' in the next major release.

Parameters

size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'batchmean' | 'sum' | 'mean'. 'none': no reduction will be applied. 'batchmean': the sum of the output will be divided by batchsize. 'sum': the output will be summed. 'mean': the output will be divided by the number of elements in the output. Default: 'mean'

Note

size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction.

Note

:attr:reduction = 'mean' doesn’t return the true kl divergence value, please use :attr:reduction = 'batchmean' which aligns with KL math definition. In the next major release, 'mean' will be changed to be the same as 'batchmean'.

Shape:

Input: $(N, *)$ where $*$ means, any number of additional dimensions
Target: $(N, *)$, same shape as the input
Output: scalar by default. If :attr:reduction is 'none', then $(N, *)$, the same shape as the input

BCELoss¶

class torch.nn.BCELoss(weight=None, size_average=None, reduce=None, reduction='mean')¶

Creates a criterion that measures the Binary Cross Entropy between the target and the output:

The unreduced (i.e. with reduction set to 'none') loss can be described as:

\[\ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad l_n = - w_n \left[ y_n \cdot \log x_n + (1 - y_n) \cdot \log (1 - x_n) \right], \]

where $N$ is the batch size. If reduction is not 'none' (default 'mean'), then

\[\ell(x, y) = \begin{cases} \operatorname{mean}(L), & \text{if reduction} = \text{'mean';}\\ \operatorname{sum}(L), & \text{if reduction} = \text{'sum'.} \end{cases} \]

This is used for measuring the error of a reconstruction in for example an auto-encoder. Note that the targets $y$ should be numbers between 0 and 1.

Parameters

weight (Tensor, optional) – a manual rescaling weight given to the loss of each batch element. If given, has to be a Tensor of size nbatch.
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(N, *)$ where $*$ means, any number of additional dimensions
Target: $(N, *)$, same shape as the input
Output: scalar. If reduction is 'none', then $(N, *)$, same shape as input.

Examples:

>>> m = nn.Sigmoid()
>>> loss = nn.BCELoss()
>>> input = torch.randn(3, requires_grad=True)
>>> target = torch.empty(3).random_(2)
>>> output = loss(m(input), target)
>>> output.backward()

BCEWithLogitsLoss¶

class torch.nn.BCEWithLogitsLoss(weight=None, size_average=None, reduce=None, reduction='mean', pos_weight=None)¶

This loss combines a Sigmoid layer and the BCELoss in one single class. This version is more numerically stable than using a plain Sigmoid followed by a BCELoss as, by combining the operations into one layer, we take advantage of the log-sum-exp trick for numerical stability.

The unreduced (i.e. with reduction set to 'none') loss can be described as:

\[\ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad l_n = - w_n \left[ y_n \cdot \log \sigma(x_n) + (1 - y_n) \cdot \log (1 - \sigma(x_n)) \right], \]

where $N$ is the batch size. If reduction is not 'none' (default 'mean'), then

\[\ell(x, y) = \begin{cases} \operatorname{mean}(L), & \text{if reduction} = \text{'mean';}\\ \operatorname{sum}(L), & \text{if reduction} = \text{'sum'.} \end{cases} \]

This is used for measuring the error of a reconstruction in for example an auto-encoder. Note that the targets t[i] should be numbers between 0 and 1.

It’s possible to trade off recall and precision by adding weights to positive examples. In this case the loss can be described as:

\[\ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad l_n = - w_n \left[ p_n y_n \cdot \log \sigma(x_n) + (1 - y_n) \cdot \log (1 - \sigma(x_n)) \right], \]

where $p_n$ is the weight of the positive class for sample $n$ in the batch. $p_n > 1$ increases the recall, $p_n < 1$ increases the precision.

For example, if a dataset contains 100 positive and 300 negative examples of a single class, then pos_weight for the class should be equal to $\frac{300}{100}=3$. The loss would act as if the dataset contains $3\times 100=300$ positive examples.

Parameters

weight (Tensor, optional) – a manual rescaling weight given to the loss of each batch element. If given, has to be a Tensor of size nbatch.
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'
pos_weight (Tensor, optional) – a weight of positive examples. Must be a vector with length equal to the number of classes.

Shape:

Input: $(N, *)$ where $*$ means, any number of additional dimensions

Target: $(N, *)$, same shape as the input

Output: scalar. If reduction is 'none', then $(N, *)$, same shape as input.

Examples:

>>> loss = nn.BCEWithLogitsLoss()
>>> input = torch.randn(3, requires_grad=True)
>>> target = torch.empty(3).random_(2)
>>> output = loss(input, target)
>>> output.backward()

MarginRankingLoss¶

class torch.nn.MarginRankingLoss(margin=0.0, size_average=None, reduce=None, reduction='mean')¶

Creates a criterion that measures the loss given inputs $x1$, $x2$, two 1D mini-batch Tensor`s, and a label 1D mini-batch tensor :math:`y (containing 1 or -1).

If $y = 1$ then it assumed the first input should be ranked higher (have a larger value) than the second input, and vice-versa for $y = -1$.

The loss function for each sample in the mini-batch is:

\[\text{loss}(x, y) = \max(0, -y * (x1 - x2) + \text{margin}) \]

Parameters

margin (float, optional) – Has a default value of $0$.
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(N, D)$ where N is the batch size and D is the size of a sample.
Target: $(N)$
Output: scalar. If reduction is 'none', then $(N)$.

HingeEmbeddingLoss¶

class torch.nn.HingeEmbeddingLoss(margin=1.0, size_average=None, reduce=None, reduction='mean')¶

Measures the loss given an input tensor $x$ and a labels tensor $y$ (containing 1 or -1). This is usually used for measuring whether two inputs are similar or dissimilar, e.g. using the L1 pairwise distance as $x$, and is typically used for learning nonlinear embeddings or semi-supervised learning.

The loss function for $n$-th sample in the mini-batch is

\[l_n = \begin{cases} x_n, & \text{if}\; y_n = 1,\\ \max \{0, \Delta - x_n\}, & \text{if}\; y_n = -1, \end{cases} \]

and the total loss functions is

\[\ell(x, y) = \begin{cases} \operatorname{mean}(L), & \text{if reduction} = \text{'mean';}\\ \operatorname{sum}(L), & \text{if reduction} = \text{'sum'.} \end{cases} \]

where $L = \{l_1,\dots,l_N\}^\top$.

Parameters

margin (float, optional) – Has a default value of 1.
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(*)$ where $*$ means, any number of dimensions. The sum operation operates over all the elements.
Target: $(*)$, same shape as the input
Output: scalar. If :attr:reduction is 'none', then same shape as the input

MultiLabelMarginLoss¶

class torch.nn.MultiLabelMarginLoss(size_average=None, reduce=None, reduction='mean')¶

Creates a criterion that optimizes a multi-class multi-classification hinge loss (margin-based loss) between input $x$ (a 2D mini-batch Tensor) and output $y$ (which is a 2D Tensor of target class indices). For each sample in the mini-batch:

\[\text{loss}(x, y) = \sum_{ij}\frac{\max(0, 1 - (x[y[j]] - x[i]))}{\text{x.size}(0)} \]

where $x \in \left\{0, \; \cdots , \; \text{x.size}(0) - 1\right\}$, $y \in \left\{0, \; \cdots , \; \text{y.size}(0) - 1\right\}$, $0 \leq y[j] \leq \text{x.size}(0)-1$, and $i \neq y[j]$ for all $i$ and $j$.

$y$ and $x$ 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.

Parameters

size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(C)$ or $(N, C)$ where N is the batch size and C is the number of classes.
Target: $(C)$ or $(N, C)$, label targets padded by -1 ensuring same shape as the input.
Output: scalar. If :attr:reduction is 'none', then $(N)$.

Examples:

>>> loss = nn.MultiLabelMarginLoss()
>>> x = torch.FloatTensor([[0.1, 0.2, 0.4, 0.8]])
>>> # for target y, only consider labels 3 and 0, not after label -1
>>> y = torch.LongTensor([[3, 0, -1, 1]])
>>> loss(x, y)
>>> # 0.25 * ((1-(0.1-0.2)) + (1-(0.1-0.4)) + (1-(0.8-0.2)) + (1-(0.8-0.4)))
tensor(0.8500)

SmoothL1Loss¶

class torch.nn.SmoothL1Loss(size_average=None, reduce=None, reduction='mean')¶

Creates a criterion that uses a squared term if the absolute element-wise error falls below 1 and an L1 term otherwise. It is less sensitive to outliers than the MSELoss and in some cases prevents exploding gradients (e.g. see “Fast R-CNN” paper by Ross Girshick). Also known as the Huber loss:

\[\text{loss}(x, y) = \frac{1}{n} \sum_{i} z_{i} \]

where $z_{i}$ is given by:

\[z_{i} = \begin{cases} 0.5 (x_i - y_i)^2, & \text{if } |x_i - y_i| < 1 \\ |x_i - y_i| - 0.5, & \text{otherwise } \end{cases} \]

$x$ and $y$ arbitrary shapes with a total of $n$ elements each the sum operation still operates over all the elements, and divides by $n$.

The division by $n$ can be avoided if sets reduction = 'sum'.

Parameters

size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(N, *)$ where $*$ means, any number of additional dimensions
Target: $(N, *)$, same shape as the input
Output: scalar. If reduction is 'none', then $(N, *)$, same shape as the input

SoftMarginLoss¶

class torch.nn.SoftMarginLoss(size_average=None, reduce=None, reduction='mean')¶

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

\[\text{loss}(x, y) = \sum_i \frac{\log(1 + \exp(-y[i]*x[i]))}{\text{x.nelement}()} \]

Parameters

size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(*)$ where $*$ means, any number of additional dimensions
Target: $(*)$, same shape as the input
Output: scalar. If reduction is 'none', then same shape as the input

MultiLabelSoftMarginLoss¶

class torch.nn.MultiLabelSoftMarginLoss(weight=None, size_average=None, reduce=None, reduction='mean')¶

Creates a criterion that optimizes a multi-label one-versus-all loss based on max-entropy, between input $x$ and target $y$ of size $(N, C)$. For each sample in the minibatch:

\[loss(x, y) = - \frac{1}{C} * \sum_i y[i] * \log((1 + \exp(-x[i]))^{-1}) + (1-y[i]) * \log\left(\frac{\exp(-x[i])}{(1 + \exp(-x[i]))}\right) \]

where $i \in \left\{0, \; \cdots , \; \text{x.nElement}() - 1\right\}$, $y[i] \in \left\{0, \; 1\right\}$.

Parameters

weight (Tensor, optional) – a manual rescaling weight given to each class. If given, it has to be a Tensor of size C. Otherwise, it is treated as if having all ones.
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(N, C)$ where N is the batch size and C is the number of classes.
Target: $(N, C)$, label targets padded by -1 ensuring same shape as the input.
Output: scalar. If reduction is 'none', then $(N)$.

CosineEmbeddingLoss¶

class torch.nn.CosineEmbeddingLoss(margin=0.0, size_average=None, reduce=None, reduction='mean')¶

Creates a criterion that measures the loss given input tensors $x_1$, $x_2$ and a Tensor label $y$ with values 1 or -1. This is used for measuring whether two inputs are similar or dissimilar, using the cosine distance, and is typically used for learning nonlinear embeddings or semi-supervised learning.

The loss function for each sample is:

\[\text{loss}(x, y) = \begin{cases} 1 - \cos(x_1, x_2), & \text{if } y = 1 \\ \max(0, \cos(x_1, x_2) - \text{margin}), & \text{if } y = -1 \end{cases} \]

Parameters: margin (float, optional) – Should be a number from $-1$ to $1$,

:param $0$ to $0.5$ is suggested. If margin is missing, the: :param default value is $0$.: :param size_average: Deprecated (see reduction). By default,

the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True

Parameters

reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

MultiMarginLoss¶

class torch.nn.MultiMarginLoss(p=1, margin=1.0, weight=None, size_average=None, reduce=None, reduction='mean')¶

Creates a criterion that optimizes a multi-class classification hinge loss (margin-based loss) between input $x$ (a 2D mini-batch Tensor) and output $y$ (which is a 1D tensor of target class indices, $0 \leq y \leq \text{x.size}(1)-1$):

For each mini-batch sample, the loss in terms of the 1D input $x$ and scalar output $y$ is:

\[\text{loss}(x, y) = \frac{\sum_i \max(0, \text{margin} - x[y] + x[i]))^p}{\text{x.size}(0)} \]

where $x \in \left\{0, \; \cdots , \; \text{x.size}(0) - 1\right\}$ and $i \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:

\[\text{loss}(x, y) = \frac{\sum_i \max(0, w[y] * (\text{margin} - x[y] + x[i]))^p)}{\text{x.size}(0)} \]

Parameters

p (int, optional) – Has a default value of $1$. $1$ and $2$ are the only supported values.
margin (float, optional) – Has a default value of $1$.
weight (Tensor, optional) – a manual rescaling weight given to each class. If given, it has to be a Tensor of size C. Otherwise, it is treated as if having all ones.
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

TripletMarginLoss¶

class torch.nn.TripletMarginLoss(margin=1.0, p=2.0, eps=1e-06, swap=False, size_average=None, reduce=None, reduction='mean')¶

Creates a criterion that measures the triplet loss given an input tensors $x1$, $x2$, $x3$ and a margin with a value greater than $0$. This is used for measuring a relative similarity between samples. A triplet is composed by a, p and n: anchor, positive examples and negative examples respectively. The shapes of all input tensors should be $(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:

\[L(a, p, n) = \max \{d(a_i, p_i) - d(a_i, n_i) + {\rm margin}, 0\} \]

where

\[d(x_i, y_i) = \left\lVert {\bf x}_i - {\bf y}_i \right\rVert_p \]

Parameters

margin (float, optional) – Default: $1$.
p (int, optional) – The norm degree for pairwise distance. Default: $2$.
swap (float, optional) – 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. Default: False.
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

Input: $(N, D)$ where $D$ is the vector dimension.
Output: scalar. If reduction is 'none', then $(N)$.

>>> triplet_loss = nn.TripletMarginLoss(margin=1.0, p=2)
>>> input1 = torch.randn(100, 128, requires_grad=True)
>>> input2 = torch.randn(100, 128, requires_grad=True)
>>> input3 = torch.randn(100, 128, requires_grad=True)
>>> output = triplet_loss(input1, input2, input3)
>>> output.backward()

Vision layers¶

PixelShuffle¶

class torch.nn.PixelShuffle(upscale_factor)¶

Rearranges elements in a tensor of shape $(*, C \times r^2, H, W)$ to a tensor of shape $(*, C, H \times r, W \times r)$.

This is useful for implementing efficient sub-pixel convolution with a stride of $1/r$.

Look at 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.

Parameters: upscale_factor (int) – factor to increase spatial resolution by

Shape:

Input: $(N, L, H_{in}, W_{in})$ where $L=C \times \text{upscale\_factor}^2$
Output: $(N, C, H_{out}, W_{out})$ where $H_{out} = H_{in} \times \text{upscale\_factor}$ and $W_{out} = W_{in} \times \text{upscale\_factor}$

Examples:

>>> pixel_shuffle = nn.PixelShuffle(3)
>>> input = torch.randn(1, 9, 4, 4)
>>> output = pixel_shuffle(input)
>>> print(output.size())
torch.Size([1, 1, 12, 12])

Upsample¶

class torch.nn.Upsample(size=None, scale_factor=None, mode='nearest', align_corners=None)¶

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

The input data is assumed to be of the form minibatch x channels x [optional depth] x [optional height] x width. Hence, for spatial inputs, we expect a 4D Tensor and for volumetric inputs, we expect a 5D Tensor.

The algorithms available for upsampling are nearest neighbor and linear, bilinear, bicubic and trilinear for 3D, 4D and 5D input Tensor, respectively.

One can either give a scale_factor or the target output size to calculate the output size. (You cannot give both, as it is ambiguous)

Parameters

(int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int], (size) – optional): output spatial sizes
(float or Tuple[float] or Tuple[float, float] or (scale_factor) – Tuple[float, float, float], optional): multiplier for spatial size. Has to match input size if it is a tuple.
mode (str, optional) – the upsampling algorithm: one of 'nearest', 'linear', 'bilinear', 'bicubic' and 'trilinear'. Default: 'nearest'
align_corners (bool, optional) – if True, the corner pixels of the input and output tensors are aligned, and thus preserving the values at those pixels. This only has effect when mode is 'linear', 'bilinear', or 'trilinear'. Default: False

Shape:

Input: $(N, C, W_{in})$, $(N, C, H_{in}, W_{in})$ or $(N, C, D_{in}, H_{in}, W_{in})$
Output: $(N, C, W_{out})$, $(N, C, H_{out}, W_{out})$ or $(N, C, D_{out}, H_{out}, W_{out})$, where

\[D_{out} = \left\lfloor D_{in} \times \text{scale\_factor} \right\rfloor \]

\[H_{out} = \left\lfloor H_{in} \times \text{scale\_factor} \right\rfloor \]

\[W_{out} = \left\lfloor W_{in} \times \text{scale\_factor} \right\rfloor \]

Warning

With align_corners = True, the linearly interpolating modes (linear, bilinear, bicubic, and trilinear) 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 below for concrete examples on how this affects the outputs.

Note

If you want downsampling/general resizing, you should use interpolate().

Examples:

>>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
>>> input
tensor([[[[ 1.,  2.],
          [ 3.,  4.]]]])

>>> m = nn.Upsample(scale_factor=2, mode='nearest')
>>> m(input)
tensor([[[[ 1.,  1.,  2.,  2.],
          [ 1.,  1.,  2.,  2.],
          [ 3.,  3.,  4.,  4.],
          [ 3.,  3.,  4.,  4.]]]])

>>> m = nn.Upsample(scale_factor=2, mode='bilinear')  # align_corners=False
>>> m(input)
tensor([[[[ 1.0000,  1.2500,  1.7500,  2.0000],
          [ 1.5000,  1.7500,  2.2500,  2.5000],
          [ 2.5000,  2.7500,  3.2500,  3.5000],
          [ 3.0000,  3.2500,  3.7500,  4.0000]]]])

>>> m = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
>>> m(input)
tensor([[[[ 1.0000,  1.3333,  1.6667,  2.0000],
          [ 1.6667,  2.0000,  2.3333,  2.6667],
          [ 2.3333,  2.6667,  3.0000,  3.3333],
          [ 3.0000,  3.3333,  3.6667,  4.0000]]]])

>>> # Try scaling the same data in a larger tensor
>>>
>>> input_3x3 = torch.zeros(3, 3).view(1, 1, 3, 3)
>>> input_3x3[:, :, :2, :2].copy_(input)
tensor([[[[ 1.,  2.],
          [ 3.,  4.]]]])
>>> input_3x3
tensor([[[[ 1.,  2.,  0.],
          [ 3.,  4.,  0.],
          [ 0.,  0.,  0.]]]])

>>> m = nn.Upsample(scale_factor=2, mode='bilinear')  # align_corners=False
>>> # Notice that values in top left corner are the same with the small input (except at boundary)
>>> m(input_3x3)
tensor([[[[ 1.0000,  1.2500,  1.7500,  1.5000,  0.5000,  0.0000],
          [ 1.5000,  1.7500,  2.2500,  1.8750,  0.6250,  0.0000],
          [ 2.5000,  2.7500,  3.2500,  2.6250,  0.8750,  0.0000],
          [ 2.2500,  2.4375,  2.8125,  2.2500,  0.7500,  0.0000],
          [ 0.7500,  0.8125,  0.9375,  0.7500,  0.2500,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]])

>>> m = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
>>> # Notice that values in top left corner are now changed
>>> m(input_3x3)
tensor([[[[ 1.0000,  1.4000,  1.8000,  1.6000,  0.8000,  0.0000],
          [ 1.8000,  2.2000,  2.6000,  2.2400,  1.1200,  0.0000],
          [ 2.6000,  3.0000,  3.4000,  2.8800,  1.4400,  0.0000],
          [ 2.4000,  2.7200,  3.0400,  2.5600,  1.2800,  0.0000],
          [ 1.2000,  1.3600,  1.5200,  1.2800,  0.6400,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]])

UpsamplingNearest2d¶

class torch.nn.UpsamplingNearest2d(size=None, scale_factor=None)¶

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

To specify the scale, it takes either the size or the scale_factor as it’s constructor argument.

When size is given, it is the output size of the image (h, w).

Parameters

size (int or Tuple[int, int], optional) – output spatial sizes
scale_factor (float or Tuple[float, float], optional) – multiplier for spatial size.

Warning

This class is deprecated in favor of interpolate().

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$ where

\[H_{out} = \left\lfloor H_{in} \times \text{scale\_factor} \right\rfloor \]

\[W_{out} = \left\lfloor W_{in} \times \text{scale\_factor} \right\rfloor \]

Examples:

>>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
>>> input
tensor([[[[ 1.,  2.],
          [ 3.,  4.]]]])

>>> m = nn.UpsamplingNearest2d(scale_factor=2)
>>> m(input)
tensor([[[[ 1.,  1.,  2.,  2.],
          [ 1.,  1.,  2.,  2.],
          [ 3.,  3.,  4.,  4.],
          [ 3.,  3.,  4.,  4.]]]])

UpsamplingBilinear2d¶

class torch.nn.UpsamplingBilinear2d(size=None, scale_factor=None)¶

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

To specify the scale, it takes either the size or the scale_factor as it’s constructor argument.

When size is given, it is the output size of the image (h, w).

Parameters

size (int or Tuple[int, int], optional) – output spatial sizes
scale_factor (float or Tuple[float, float], optional) – multiplier for spatial size.

Warning

This class is deprecated in favor of interpolate(). It is equivalent to nn.functional.interpolate(..., mode='bilinear', align_corners=True).

Shape:

Input: $(N, C, H_{in}, W_{in})$
Output: $(N, C, H_{out}, W_{out})$ where

\[H_{out} = \left\lfloor H_{in} \times \text{scale\_factor} \right\rfloor \]

\[W_{out} = \left\lfloor W_{in} \times \text{scale\_factor} \right\rfloor \]

Examples:

>>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
>>> input
tensor([[[[ 1.,  2.],
          [ 3.,  4.]]]])

>>> m = nn.UpsamplingBilinear2d(scale_factor=2)
>>> m(input)
tensor([[[[ 1.0000,  1.3333,  1.6667,  2.0000],
          [ 1.6667,  2.0000,  2.3333,  2.6667],
          [ 2.3333,  2.6667,  3.0000,  3.3333],
          [ 3.0000,  3.3333,  3.6667,  4.0000]]]])

DataParallel layers (multi-GPU, distributed)¶

DataParallel¶

class torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)¶

Implements data parallelism 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 (other objects will be copied once per device). In the forward pass, the module is replicated on each device, and each replica handles a portion of the input. During the backwards pass, gradients from each replica are summed into the original module.

The batch size should be larger than the number of GPUs used.

DistributedDataParallel¶

class torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25, check_reduction=False)¶

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 cuda-nn-dataparallel-instead. 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 can be used in the following two ways:

Single-Process Multi-GPU

In this case, a single process will be spawned on each host/node and each process will operate on all the GPUs of the node where it’s running. To use DistributedDataParallel in this way, you can simply construct the model as the following:

>>> torch.distributed.init_process_group(backend="nccl")
>>> model = DistributedDataParallel(model) # device_ids will include all GPU devices by default

Multi-Process Single-GPU

This is the highly recommended way to use DistributedDataParallel, with multiple processes, each of which operates on a single GPU. This is currently the fastest approach to do data parallel training using PyTorch and applies to both single-node(multi-GPU) and multi-node data parallel training. It is proven to be significantly faster than torch.nn.DataParallel for single-node multi-GPU data parallel training.

Here is how to use it: on each host with N GPUs, you should spawn up N processes, while ensuring that each process individually works on a single GPU from 0 to N-1. Therefore, it is your job to ensure that your training script operates on a single given GPU by calling:

>>> torch.cuda.set_device(i)

where i is from 0 to N-1. In each process, you should refer the following to construct this module:

>>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')
>>> model = DistributedDataParallel(model, device_ids=[i], output_device=i)

In order to spawn up multiple processes per node, you can use either torch.distributed.launch or torch.multiprocessing.spawn

Note

nccl backend is currently the fastest and highly recommended backend to be used with Multi-Process Single-GPU distributed training and 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. Also note that nccl backend is currently the fastest and highly recommended backend for fp16/fp32 mixed-precision training.

Warning

This module works only with the gloo and nccl backends.

Warning

Constructor, forward method, and differentiation of the output (or a function of the output of this module) is a distributed synchronization point. 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 all-reduction 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 assumes all buffers and gradients are dense.

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. In other words, 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 after the DistributedDataParallel construction, this is not supported and unexpected behaviors can happen, since some parameters’ gradient reduction functions might not get called.

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.

Parameters

module (Module) – module to be parallelized
device_ids (list of python:int or torch.device) – CUDA devices (default: all devices)
output_device (int or torch.device) – device location of output (default: device_ids[0])
broadcast_buffers (bool) – flag that enables syncing (broadcasting) buffers of the module at beginning of the forward function. (default: True)
process_group – the process group to be used for distributed data all-reduction. If None, the default process group, which is created by `torch.distributed.init_process_group`, will be used. (default: None)
bucket_cap_mb – DistributedDataParallel will bucket parameters into multiple buckets so that gradient reduction of each bucket can potentially overlap with backward computation. bucket_cap_mb controls the bucket size in MegaBytes (MB) (default: 25)
check_reduction – when setting to True, it enables DistributedDataParallel to automatically check if the previous iteration’s backward reductions were successfully issued at the beginning of every iteration’s forward function. You normally don’t need this option enabled unless you are observing weird behaviors such as different ranks are getting different gradients, which should not happen if DistributedDataParallel is correctly used. (default: False)

Variables

~DistributedDataParallel.module (Module) – the module to be parallelized

Example:

>>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')
>>> net = torch.nn.DistributedDataParallel(model, pg)

DistributedDataParallelCPU¶

class torch.nn.parallel.DistributedDataParallelCPU(module)¶

Implements distributed data parallelism for CPU at the module level.

This module supports the mpi and gloo backends.

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 such replica handles a portion of the input. During the backwards pass, gradients from each node are averaged.

This module could be used in conjunction with the DistributedSampler, (see DistributedSampler) which will load a subset of the original dataset for each node with the same batch size. So strong scaling should be configured like this:

n = 1, batch size = 12

n = 2, batch size = 64

n = 4, batch size = 32

n = 8, batch size = 16

Creation of this class requires the distributed package to be already initialized in the process group mode (see torch.distributed.init_process_group()).

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.

Warning

This module assumes all gradients are dense.

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

Forward and backward hooks defined on module and its submodules won’t be invoked anymore, unless the hooks are initialized in the forward() method.

Note

Parameters are broadcast between nodes in the __init__() function. The module performs an all-reduce step on gradients and assumes that they will be modified by the optimizer in all nodes in the same way.

Parameters: module – module to be parallelized

Example:

>>> torch.distributed.init_process_group(world_size=4, init_method='...')
>>> net = torch.nn.DistributedDataParallelCPU(model)

Utilities¶

clip_grad_norm_¶

torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=2)¶

Clips gradient norm of an iterable of parameters.

The norm is computed over all gradients together, as if they were concatenated into a single vector. Gradients are modified in-place.

Parameters

parameters (Iterable[Tensor] or Tensor) – an iterable of Tensors or a single Tensor that will have gradients normalized
max_norm (float or int) – max norm of the gradients
norm_type (float or int) – type of the used p-norm. Can be 'inf' for infinity norm.

Returns

Total norm of the parameters (viewed as a single vector).

clip_grad_value_¶

torch.nn.utils.clip_grad_value_(parameters, clip_value)¶

Clips gradient of an iterable of parameters at specified value.

Gradients are modified in-place.

Parameters

parameters (Iterable[Tensor] or Tensor) – an iterable of Tensors or a single Tensor that will have gradients normalized
clip_value (float or int) – maximum allowed value of the gradients. The gradients are clipped in the range $\left[\text{-clip\_value}, \text{clip\_value}\right]$

parameters_to_vector¶

torch.nn.utils.parameters_to_vector(parameters)¶

Convert parameters to one vector

Parameters: parameters (Iterable[Tensor]) – an iterator of Tensors that are the parameters of a model.
Returns: The parameters represented by a single vector

vector_to_parameters¶

torch.nn.utils.vector_to_parameters(vec, parameters)¶

Convert one vector to the parameters

Parameters

vec (Tensor) – a single vector represents the parameters of a model.
parameters (Iterable[Tensor]) – an iterator of Tensors that are the parameters of a model.

weight_norm¶

torch.nn.utils.weight_norm(module, name='weight', dim=0)¶

Applies weight normalization to a parameter in the given module.

\[\mathbf{w} = g \dfrac{\mathbf{v}}{\|\mathbf{v}\|} \]

Weight normalization is a reparameterization that decouples the magnitude of a weight tensor from its direction. This replaces the parameter specified by name (e.g. 'weight') with two parameters: one specifying the magnitude (e.g. 'weight_g') and one specifying the direction (e.g. 'weight_v'). Weight normalization is implemented via a hook that recomputes the weight tensor from the magnitude and direction before every forward() call.

By default, with dim=0, the norm is computed independently per output channel/plane. To compute a norm over the entire weight tensor, use dim=None.

See https://arxiv.org/abs/1602.07868

Parameters

module (Module) – containing module
name (str, optional) – name of weight parameter
dim (int, optional) – dimension over which to compute the norm

Returns

The original module with the weight norm hook

Example:

>>> m = weight_norm(nn.Linear(20, 40), name='weight')
>>> m
Linear(in_features=20, out_features=40, bias=True)
>>> m.weight_g.size()
torch.Size([40, 1])
>>> m.weight_v.size()
torch.Size([40, 20])

remove_weight_norm¶

torch.nn.utils.remove_weight_norm(module, name='weight')¶

Removes the weight normalization reparameterization from a module.

Parameters

module (Module) – containing module
name (str, optional) – name of weight parameter

Example

>>> m = weight_norm(nn.Linear(20, 40))
>>> remove_weight_norm(m)

spectral_norm¶

torch.nn.utils.spectral_norm(module, name='weight', n_power_iterations=1, eps=1e-12, dim=None)¶

Applies spectral normalization to a parameter in the given module.

\[\mathbf{W}_{SN} = \dfrac{\mathbf{W}}{\sigma(\mathbf{W})}, \sigma(\mathbf{W}) = \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2} \]

Spectral normalization stabilizes the training of discriminators (critics) in Generative Adversarial Networks (GANs) by rescaling the weight tensor with spectral norm $\sigma$ of the weight matrix calculated using power iteration method. If the dimension of the weight tensor is greater than 2, it is reshaped to 2D in power iteration method to get spectral norm. This is implemented via a hook that calculates spectral norm and rescales weight before every forward() call.

See Spectral Normalization for Generative Adversarial Networks .

Parameters

module (nn.Module) – containing module
name (str, optional) – name of weight parameter
n_power_iterations (int, optional) – number of power iterations to calculate spectral norm
eps (float, optional) – epsilon for numerical stability in calculating norms
dim (int, optional) – dimension corresponding to number of outputs, the default is 0, except for modules that are instances of ConvTranspose{1,2,3}d, when it is 1

Returns

The original module with the spectral norm hook

Example:

>>> m = spectral_norm(nn.Linear(20, 40))
>>> m
Linear(in_features=20, out_features=40, bias=True)
>>> m.weight_u.size()
torch.Size([40])

remove_spectral_norm¶

torch.nn.utils.remove_spectral_norm(module, name='weight')¶

Removes the spectral normalization reparameterization from a module.

Parameters

module (Module) – containing module
name (str, optional) – name of weight parameter

Example

>>> m = spectral_norm(nn.Linear(40, 10))
>>> remove_spectral_norm(m)

PackedSequence¶

torch.nn.utils.rnn.PackedSequence(data, batch_sizes=None, sorted_indices=None, unsorted_indices=None)¶

Holds the data and list of batch_sizes of a packed sequence.

All RNN modules accept packed sequences as inputs.

Note

Instances of this class should never be created manually. They are meant to be instantiated by functions like pack_padded_sequence().

Batch sizes represent the number elements at each sequence step in the batch, not the varying sequence lengths passed to pack_padded_sequence(). For instance, given data abc and x the PackedSequence would contain data axbc with batch_sizes=[2,1,1].

Variables

~PackedSequence.data (Tensor) – Tensor containing packed sequence
~PackedSequence.batch_sizes (Tensor) – Tensor of integers holding information about the batch size at each sequence step

pack_padded_sequence¶

torch.nn.utils.rnn.pack_padded_sequence(input, lengths, batch_first=False, enforce_sorted=True)¶

Packs a Tensor containing padded sequences of variable length.

Input can be of size T x B x * where T is the length of the longest sequence (equal to lengths[0]), B is the batch size, and * is any number of dimensions (including 0). If batch_first is True B x T x * inputs are expected.

For unsorted sequences, use enforce_sorted = False. If enforce_sorted is True, the sequences should be sorted by length in a decreasing order, i.e. input[:,0] should be the longest sequence, and input[:,B-1] the shortest one. enforce_sorted = True is only necessary for ONNX export.

Note

This function accepts any input that has at least two dimensions. You can apply it to pack the labels, and use the output of the RNN with them to compute the loss directly. A Tensor can be retrieved from a PackedSequence object by accessing its .data attribute.

Parameters

input (Tensor) – padded batch of variable length sequences.
lengths (Tensor) – list of sequences lengths of each batch element.
batch_first (bool, optional) – if True, the input is expected in B x T x * format.
enforce_sorted (bool, optional) – if True, the input is expected to contain sequences sorted by length in a decreasing order. If False, this condition is not checked. Default: True.

Returns

a PackedSequence object

pad_packed_sequence¶

torch.nn.utils.rnn.pad_packed_sequence(sequence, batch_first=False, padding_value=0.0, total_length=None)¶

Pads a packed batch of variable length sequences.

It is an inverse operation to pack_padded_sequence().

The returned Tensor’s data will be of size T x B x *, where T is the length of the longest sequence and B is the batch size. If batch_first is True, the data will be transposed into B x T x * format.

Batch elements will be ordered decreasingly by their length.

Note

total_length is useful to implement the pack sequence -> recurrent network -> unpack sequence pattern in a Module wrapped in DataParallel. See this FAQ section for details.

Parameters

sequence (PackedSequence) – batch to pad
batch_first (bool, optional) – if True, the output will be in B x T x * format.
padding_value (float, optional) – values for padded elements.
total_length (int, optional) – if not None, the output will be padded to have length total_length. This method will throw ValueError if total_length is less than the max sequence length in sequence.

Returns

Tuple of Tensor containing the padded sequence, and a Tensor containing the list of lengths of each sequence in the batch.

pad_sequence¶

torch.nn.utils.rnn.pad_sequence(sequences, batch_first=False, padding_value=0)¶

Pad a list of variable length Tensors with padding_value

pad_sequence stacks a list of Tensors along a new dimension, and pads them to equal length. For example, if the input is list of sequences with size L x * and if batch_first is False, and T x B x * otherwise.

B is batch size. It is equal to the number of elements in sequences. T is length of the longest sequence. L is length of the sequence. * is any number of trailing dimensions, including none.

Example

>>> from torch.nn.utils.rnn import pad_sequence
>>> a = torch.ones(25, 300)
>>> b = torch.ones(22, 300)
>>> c = torch.ones(15, 300)
>>> pad_sequence([a, b, c]).size()
torch.Size([25, 3, 300])

Note

This function returns a Tensor of size T x B x * or B x T x * where T is the length of the longest sequence. This function assumes trailing dimensions and type of all the Tensors in sequences are same.

Parameters

sequences (list[Tensor]) – list of variable length sequences.
batch_first (bool, optional) – output will be in B x T x * if True, or in T x B x * otherwise
padding_value (float, optional) – value for padded elements. Default: 0.

Returns

Tensor of size T x B x * if batch_first is False. Tensor of size B x T x * otherwise

pack_sequence¶

torch.nn.utils.rnn.pack_sequence(sequences, enforce_sorted=True)¶

Packs a list of variable length Tensors

sequences should be a list of Tensors of size L x *, where L is the length of a sequence and * is any number of trailing dimensions, including zero.

For unsorted sequences, use enforce_sorted = False. If enforce_sorted is True, the sequences should be sorted in the order of decreasing length. enforce_sorted = True is only necessary for ONNX export.

Example

>>> from torch.nn.utils.rnn import pack_sequence
>>> a = torch.tensor([1,2,3])
>>> b = torch.tensor([4,5])
>>> c = torch.tensor([6])
>>> pack_sequence([a, b, c])
PackedSequence(data=tensor([ 1,  4,  6,  2,  5,  3]), batch_sizes=tensor([ 3,  2,  1]))

Parameters

sequences (list[Tensor]) – A list of sequences of decreasing length.
enforce_sorted (bool, optional) – if True, checks that the input contains sequences sorted by length in a decreasing order. If False, this condition is not checked. Default: True.

Returns

a PackedSequence object

torch.nn.functional¶

Convolution functions¶

conv1d¶

torch.nn.functional.conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros') → Tensor¶

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

See Conv1d for details and output shape.

Note

Parameters

input – input tensor of shape $(\text{minibatch} , \text{in\_channels} , iW)$
weight – filters of shape $(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kW)$
bias – optional bias of shape $(\text{out\_channels})$. Default: None
stride – the stride of the convolving kernel. Can be a single number or a one-element tuple (sW,). Default: 1
padding – implicit paddings on both sides of the input. Can be a single number or a one-element tuple (padW,). Default: 0
dilation – the spacing between kernel elements. Can be a single number or a one-element tuple (dW,). Default: 1
groups – split input into groups, $\text{in\_channels}$ should be divisible by the number of groups. Default: 1
padding_mode – the type of paddings applied to both sided can be: zeros or circular. Default: zeros

Examples:

>>> filters = torch.randn(33, 16, 3)
>>> inputs = torch.randn(20, 16, 50)
>>> F.conv1d(inputs, filters)

conv2d¶

torch.nn.functional.conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros') → Tensor¶

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

See Conv2d for details and output shape.

Note

Parameters

input – input tensor of shape $(\text{minibatch} , \text{in\_channels} , iH , iW)$
weight – filters of shape $(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kH , kW)$
bias – optional bias tensor of shape $(\text{out\_channels})$. Default: None
stride – the stride of the convolving kernel. Can be a single number or a tuple (sH, sW). Default: 1
padding – implicit paddings on both sides of the input. Can be a single number or a tuple (padH, padW). Default: 0
dilation – the spacing between kernel elements. Can be a single number or a tuple (dH, dW). Default: 1
groups – split input into groups, $\text{in\_channels}$ should be divisible by the number of groups. Default: 1
padding_mode – the type of paddings applied to both sided can be: zeros or circular. Default: zeros

Examples:

>>> # With square kernels and equal stride
>>> filters = torch.randn(8,4,3,3)
>>> inputs = torch.randn(1,4,5,5)
>>> F.conv2d(inputs, filters, padding=1)

conv3d¶

torch.nn.functional.conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros') → Tensor¶

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

See Conv3d for details and output shape.

Note

Parameters

input – input tensor of shape $(\text{minibatch} , \text{in\_channels} , iT , iH , iW)$
weight – filters of shape $(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kT , kH , kW)$
bias – optional bias tensor of shape $(\text{out\_channels})$. Default: None
stride – the stride of the convolving kernel. Can be a single number or a tuple (sT, sH, sW). Default: 1
padding – implicit paddings on both sides of the input. Can be a single number or a tuple (padT, padH, padW). Default: 0
dilation – the spacing between kernel elements. Can be a single number or a tuple (dT, dH, dW). Default: 1
groups – split input into groups, $\text{in\_channels}$ should be divisible by the number of groups. Default: 1
padding_mode – the type of paddings applied to both sided can be: zeros or circular. Default: zeros

Examples:

>>> filters = torch.randn(33, 16, 3, 3, 3)
>>> inputs = torch.randn(20, 16, 50, 10, 20)
>>> F.conv3d(inputs, filters)

conv_transpose1d¶

torch.nn.functional.conv_transpose1d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) → Tensor¶

Applies a 1D transposed convolution operator over an input signal composed of several input planes, sometimes also called “deconvolution”.

See ConvTranspose1d for details and output shape.

Note

Parameters

input – input tensor of shape $(\text{minibatch} , \text{in\_channels} , iW)$
weight – filters of shape $(\text{in\_channels} , \frac{\text{out\_channels}}{\text{groups}} , kW)$
bias – optional bias of shape $(\text{out\_channels})$. Default: None
stride – the stride of the convolving kernel. Can be a single number or a tuple (sW,). Default: 1
padding – dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each dimension in the input. Can be a single number or a tuple (padW,). Default: 0
output_padding – additional size added to one side of each dimension in the output shape. Can be a single number or a tuple (out_padW). Default: 0
groups – split input into groups, $\text{in\_channels}$ should be divisible by the number of groups. Default: 1
dilation – the spacing between kernel elements. Can be a single number or a tuple (dW,). Default: 1

Examples:

>>> inputs = torch.randn(20, 16, 50)
>>> weights = torch.randn(16, 33, 5)
>>> F.conv_transpose1d(inputs, weights)

conv_transpose2d¶

torch.nn.functional.conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) → Tensor¶

Applies a 2D transposed convolution operator over an input image composed of several input planes, sometimes also called “deconvolution”.

See ConvTranspose2d for details and output shape.

Note

Parameters

input – input tensor of shape $(\text{minibatch} , \text{in\_channels} , iH , iW)$
weight – filters of shape $(\text{in\_channels} , \frac{\text{out\_channels}}{\text{groups}} , kH , kW)$
bias – optional bias of shape $(\text{out\_channels})$. Default: None
stride – the stride of the convolving kernel. Can be a single number or a tuple (sH, sW). Default: 1
padding – dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each dimension in the input. Can be a single number or a tuple (padH, padW). Default: 0
output_padding – additional size added to one side of each dimension in the output shape. Can be a single number or a tuple (out_padH, out_padW). Default: 0
groups – split input into groups, $\text{in\_channels}$ should be divisible by the number of groups. Default: 1
dilation – the spacing between kernel elements. Can be a single number or a tuple (dH, dW). Default: 1

Examples:

>>> # With square kernels and equal stride
>>> inputs = torch.randn(1, 4, 5, 5)
>>> weights = torch.randn(4, 8, 3, 3)
>>> F.conv_transpose2d(inputs, weights, padding=1)

conv_transpose3d¶

torch.nn.functional.conv_transpose3d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) → Tensor¶

Applies a 3D transposed convolution operator over an input image composed of several input planes, sometimes also called “deconvolution”

See ConvTranspose3d for details and output shape.

Note

Parameters

input – input tensor of shape $(\text{minibatch} , \text{in\_channels} , iT , iH , iW)$
weight – filters of shape $(\text{in\_channels} , \frac{\text{out\_channels}}{\text{groups}} , kT , kH , kW)$
bias – optional bias of shape $(\text{out\_channels})$. Default: None
stride – the stride of the convolving kernel. Can be a single number or a tuple (sT, sH, sW). Default: 1
padding – dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each dimension in the input. Can be a single number or a tuple (padT, padH, padW). Default: 0
output_padding – additional size added to one side of each dimension in the output shape. Can be a single number or a tuple (out_padT, out_padH, out_padW). Default: 0
groups – split input into groups, $\text{in\_channels}$ should be divisible by the number of groups. Default: 1
dilation – the spacing between kernel elements. Can be a single number or a tuple (dT, dH, dW). Default: 1

Examples:

>>> inputs = torch.randn(20, 16, 50, 10, 20)
>>> weights = torch.randn(16, 33, 3, 3, 3)
>>> F.conv_transpose3d(inputs, weights)

unfold¶

torch.nn.functional.unfold(input, kernel_size, dilation=1, padding=0, stride=1)¶

Extracts sliding local blocks from an batched input tensor.

Warning

Currently, only 4-D input tensors (batched image-like tensors) are supported.

Warning

More than one element of the unfolded tensor may refer to a single memory location. As a result, in-place operations (especially ones that are vectorized) may result in incorrect behavior. If you need to write to the tensor, please clone it first.

See torch.nn.Unfold for details

fold¶

torch.nn.functional.fold(input, output_size, kernel_size, dilation=1, padding=0, stride=1)¶

Combines an array of sliding local blocks into a large containing tensor.

Warning

Currently, only 4-D output tensors (batched image-like tensors) are supported.

See torch.nn.Fold for details

Pooling functions¶

avg_pool1d¶

torch.nn.functional.avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True) → Tensor¶

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

See AvgPool1d for details and output shape.

Parameters

input – input tensor of shape $(\text{minibatch} , \text{in\_channels} , iW)$
kernel_size – the size of the window. Can be a single number or a tuple (kW,)
stride – the stride of the window. Can be a single number or a tuple (sW,). Default: kernel_size
padding – implicit zero paddings on both sides of the input. Can be a single number or a tuple (padW,). Default: 0
ceil_mode – when True, will use ceil instead of floor to compute the output shape. Default: False
count_include_pad – when True, will include the zero-padding in the averaging calculation. Default: True

Examples:

>>> # pool of square window of size=3, stride=2
>>> input = torch.tensor([[[1, 2, 3, 4, 5, 6, 7]]], dtype=torch.float32)
>>> F.avg_pool1d(input, kernel_size=3, stride=2)
tensor([[[ 2.,  4.,  6.]]])

avg_pool2d¶

torch.nn.functional.avg_pool2d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True) → Tensor¶

Applies 2D average-pooling operation in $kH \times kW$ regions by step size $sH \times sW$ steps. The number of output features is equal to the number of input planes.

See AvgPool2d for details and output shape.

Parameters

input – input tensor $(\text{minibatch} , \text{in\_channels} , iH , iW)$
kernel_size – size of the pooling region. Can be a single number or a tuple (kH, kW)
stride – stride of the pooling operation. Can be a single number or a tuple (sH, sW). Default: kernel_size
padding – implicit zero paddings on both sides of the input. Can be a single number or a tuple (padH, padW). Default: 0
ceil_mode – when True, will use ceil instead of floor in the formula to compute the output shape. Default: False
count_include_pad – when True, will include the zero-padding in the averaging calculation. Default: True

avg_pool3d¶

torch.nn.functional.avg_pool3d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True) → Tensor¶

Applies 3D average-pooling operation in $kT \times kH \times kW$ regions by step size $sT \times sH \times sW$ steps. The number of output features is equal to $\lfloor\frac{\text{input planes}}{sT}\rfloor$.

See AvgPool3d for details and output shape.

Parameters

input – input tensor $(\text{minibatch} , \text{in\_channels} , iT \times iH , iW)$
kernel_size – size of the pooling region. Can be a single number or a tuple (kT, kH, kW)
stride – stride of the pooling operation. Can be a single number or a tuple (sT, sH, sW). Default: kernel_size
padding – implicit zero paddings on both sides of the input. Can be a single number or a tuple (padT, padH, padW), Default: 0
ceil_mode – when True, will use ceil instead of floor in the formula to compute the output shape
count_include_pad – when True, will include the zero-padding in the averaging calculation

max_pool1d¶

torch.nn.functional.max_pool1d(*args, **kwargs)¶

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

See MaxPool1d for details.

max_pool2d¶

torch.nn.functional.max_pool2d(*args, **kwargs)¶

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

See MaxPool2d for details.

max_pool3d¶

torch.nn.functional.max_pool3d(*args, **kwargs)¶

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

See MaxPool3d for details.

max_unpool1d¶

torch.nn.functional.max_unpool1d(input, indices, kernel_size, stride=None, padding=0, output_size=None)¶

Computes a partial inverse of MaxPool1d.

See MaxUnpool1d for details.

max_unpool2d¶

torch.nn.functional.max_unpool2d(input, indices, kernel_size, stride=None, padding=0, output_size=None)¶

Computes a partial inverse of MaxPool2d.

See MaxUnpool2d for details.

max_unpool3d¶

torch.nn.functional.max_unpool3d(input, indices, kernel_size, stride=None, padding=0, output_size=None)¶

Computes a partial inverse of MaxPool3d.

See MaxUnpool3d for details.

lp_pool1d¶

torch.nn.functional.lp_pool1d(input, norm_type, kernel_size, stride=None, ceil_mode=False)¶

Applies a 1D power-average pooling over an input signal composed of several input planes. If the sum of all inputs to the power of p is zero, the gradient is set to zero as well.

See LPPool1d for details.

lp_pool2d¶

torch.nn.functional.lp_pool2d(input, norm_type, kernel_size, stride=None, ceil_mode=False)¶

Applies a 2D power-average pooling over an input signal composed of several input planes. If the sum of all inputs to the power of p is zero, the gradient is set to zero as well.

See LPPool2d for details.

adaptive_max_pool1d¶

torch.nn.functional.adaptive_max_pool1d(*args, **kwargs)¶

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

See AdaptiveMaxPool1d for details and output shape.

Parameters

output_size – the target output size (single integer)
return_indices – whether to return pooling indices. Default: False

adaptive_max_pool2d¶

torch.nn.functional.adaptive_max_pool2d(*args, **kwargs)¶

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

See AdaptiveMaxPool2d for details and output shape.

Parameters

output_size – the target output size (single integer or double-integer tuple)
return_indices – whether to return pooling indices. Default: False

adaptive_max_pool3d¶

torch.nn.functional.adaptive_max_pool3d(*args, **kwargs)¶

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

See AdaptiveMaxPool3d for details and output shape.

Parameters

output_size – the target output size (single integer or triple-integer tuple)
return_indices – whether to return pooling indices. Default: False

adaptive_avg_pool1d¶

torch.nn.functional.adaptive_avg_pool1d(input, output_size) → Tensor¶

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

See AdaptiveAvgPool1d for details and output shape.

Parameters: output_size – the target output size (single integer)

adaptive_avg_pool2d¶

torch.nn.functional.adaptive_avg_pool2d(input, output_size)¶

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

See AdaptiveAvgPool2d for details and output shape.

Parameters: output_size – the target output size (single integer or double-integer tuple)

adaptive_avg_pool3d¶

torch.nn.functional.adaptive_avg_pool3d(input, output_size)¶

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

See AdaptiveAvgPool3d for details and output shape.

Parameters: output_size – the target output size (single integer or triple-integer tuple)

Non-linear activation functions¶

threshold¶

torch.nn.functional.threshold(input, threshold, value, inplace=False)¶

Thresholds each element of the input Tensor.

See Threshold for more details.

torch.nn.functional.threshold_(input, threshold, value) → Tensor¶: In-place version of threshold().

relu¶

torch.nn.functional.relu(input, inplace=False) → Tensor¶: Applies the rectified linear unit function element-wise. See ReLU for more details.

torch.nn.functional.relu_(input) → Tensor¶: In-place version of relu().

hardtanh¶

torch.nn.functional.hardtanh(input, min_val=-1., max_val=1., inplace=False) → Tensor¶: Applies the HardTanh function element-wise. See Hardtanh for more details.

torch.nn.functional.hardtanh_(input, min_val=-1., max_val=1.) → Tensor¶: In-place version of hardtanh().

relu6¶

torch.nn.functional.relu6(input, inplace=False) → Tensor¶

Applies the element-wise function $\text{ReLU6}(x) = \min(\max(0,x), 6)$.

See ReLU6 for more details.

elu¶

torch.nn.functional.elu(input, alpha=1.0, inplace=False)¶

Applies element-wise, $\text{ELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x) - 1))$.

See ELU for more details.

torch.nn.functional.elu_(input, alpha=1.) → Tensor¶: In-place version of elu().

selu¶

torch.nn.functional.selu(input, inplace=False) → Tensor¶

Applies element-wise, $\text{SELU}(x) = scale * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1)))$, with $\alpha=1.6732632423543772848170429916717$ and $scale=1.0507009873554804934193349852946$.

See SELU for more details.

celu¶

torch.nn.functional.celu(input, alpha=1., inplace=False) → Tensor¶

Applies element-wise, $\text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))$.

See CELU for more details.

leaky_relu¶

torch.nn.functional.leaky_relu(input, negative_slope=0.01, inplace=False) → Tensor¶

Applies element-wise, $\text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)$

See LeakyReLU for more details.

torch.nn.functional.leaky_relu_(input, negative_slope=0.01) → Tensor¶: In-place version of leaky_relu().

prelu¶

torch.nn.functional.prelu(input, weight) → Tensor¶

Applies element-wise the function $\text{PReLU}(x) = \max(0,x) + \text{weight} * \min(0,x)$ where weight is a learnable parameter.

See PReLU for more details.

rrelu¶

torch.nn.functional.rrelu(input, lower=1./8, upper=1./3, training=False, inplace=False) → Tensor¶

Randomized leaky ReLU.

See RReLU for more details.

torch.nn.functional.rrelu_(input, lower=1./8, upper=1./3, training=False) → Tensor¶: In-place version of rrelu().

glu¶

torch.nn.functional.glu(input, dim=-1) → Tensor¶

The gated linear unit. Computes:

\[\text{GLU}(a, b) = a \otimes \sigma(b) \]

where input is split in half along dim to form a and b, $\sigma$ is the sigmoid function and $\otimes$ is the element-wise product between matrices.

See Language Modeling with Gated Convolutional Networks.

Parameters

input (Tensor) – input tensor
dim (int) – dimension on which to split the input. Default: -1

logsigmoid¶

torch.nn.functional.logsigmoid(input) → Tensor¶

Applies element-wise $\text{LogSigmoid}(x_i) = \log \left(\frac{1}{1 + \exp(-x_i)}\right)$

See LogSigmoid for more details.

hardshrink¶

torch.nn.functional.hardshrink(input, lambd=0.5) → Tensor¶

Applies the hard shrinkage function element-wise

See Hardshrink for more details.

tanhshrink¶

torch.nn.functional.tanhshrink(input) → Tensor¶

Applies element-wise, $\text{Tanhshrink}(x) = x - \text{Tanh}(x)$

See Tanhshrink for more details.

softsign¶

torch.nn.functional.softsign(input) → Tensor¶

Applies element-wise, the function $\text{SoftSign}(x) = \frac{x}{1 + |x|}$

See Softsign for more details.

softplus¶

torch.nn.functional.softplus(input, beta=1, threshold=20) → Tensor¶

softmin¶

torch.nn.functional.softmin(input, dim=None, _stacklevel=3, dtype=None)¶

Applies a softmin function.

Note that $\text{Softmin}(x) = \text{Softmax}(-x)$. See softmax definition for mathematical formula.

See Softmin for more details.

Parameters

input (Tensor) – input
dim (int) – A dimension along which softmin will be computed (so every slice along dim will sum to 1).
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

softmax¶

torch.nn.functional.softmax(input, dim=None, _stacklevel=3, dtype=None)¶

Applies a softmax function.

Softmax is defined as:

$\text{Softmax}(x_{i}) = \frac{exp(x_i)}{\sum_j exp(x_j)}$

It is applied to all slices along dim, and will re-scale them so that the elements lie in the range [0, 1] and sum to 1.

See Softmax for more details.

Parameters

input (Tensor) – input
dim (int) – A dimension along which softmax will be computed.
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

Note

This function doesn’t work directly with NLLLoss, which expects the Log to be computed between the Softmax and itself. Use log_softmax instead (it’s faster and has better numerical properties).

softshrink¶

torch.nn.functional.softshrink(input, lambd=0.5) → Tensor¶

Applies the soft shrinkage function elementwise

See Softshrink for more details.

gumbel_softmax¶

torch.nn.functional.gumbel_softmax(logits, tau=1, hard=False, eps=1e-10, dim=-1)¶

Samples from the Gumbel-Softmax distribution and optionally discretizes.

Parameters

logits – […, num_features] unnormalized log probabilities
tau – non-negative scalar temperature
hard – if True, the returned samples will be discretized as one-hot vectors, but will be differentiated as if it is the soft sample in autograd
dim (int) – A dimension along which softmax will be computed. Default: -1.

Returns

Sampled tensor of same shape as logits from the Gumbel-Softmax distribution. If hard=True, the returned samples will be one-hot, otherwise they will be probability distributions that sum to 1 across dim.

Note

This function is here for legacy reasons, may be removed from nn.Functional in the future.

Note

The main trick for hard is to do y_hard - y_soft.detach() + y_soft

It achieves two things: - makes the output value exactly one-hot (since we add then subtract y_soft value) - makes the gradient equal to y_soft gradient (since we strip all other gradients)

Examples::

>>> logits = torch.randn(20, 32)
>>> # Sample soft categorical using reparametrization trick:
>>> F.gumbel_softmax(logits, tau=1, hard=False)
>>> # Sample hard categorical using "Straight-through" trick:
>>> F.gumbel_softmax(logits, tau=1, hard=True)

log_softmax¶

torch.nn.functional.log_softmax(input, dim=None, _stacklevel=3, dtype=None)¶

Applies a softmax followed by a logarithm.

While mathematically equivalent to log(softmax(x)), doing these two operations separately is slower, and numerically unstable. This function uses an alternative formulation to compute the output and gradient correctly.

See LogSoftmax for more details.

Parameters

input (Tensor) – input
dim (int) – A dimension along which log_softmax will be computed.
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

tanh¶

torch.nn.functional.tanh(input) → Tensor¶

Applies element-wise, $\text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)}{\exp(x) + \exp(-x)}$

See Tanh for more details.

sigmoid¶

torch.nn.functional.sigmoid(input) → Tensor¶

Applies the element-wise function $\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)}$

See Sigmoid for more details.

Normalization functions¶

batch_norm¶

torch.nn.functional.batch_norm(input, running_mean, running_var, weight=None, bias=None, training=False, momentum=0.1, eps=1e-05)¶

Applies Batch Normalization for each channel across a batch of data.

See BatchNorm1d, BatchNorm2d, BatchNorm3d for details.

instance_norm¶

torch.nn.functional.instance_norm(input, running_mean=None, running_var=None, weight=None, bias=None, use_input_stats=True, momentum=0.1, eps=1e-05)¶

Applies Instance Normalization for each channel in each data sample in a batch.

See InstanceNorm1d, InstanceNorm2d, InstanceNorm3d for details.

layer_norm¶

torch.nn.functional.layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05)¶

Applies Layer Normalization for last certain number of dimensions.

See LayerNorm for details.

local_response_norm¶

torch.nn.functional.local_response_norm(input, size, alpha=0.0001, beta=0.75, k=1.0)¶

Applies local response normalization over an input signal composed of several input planes, where channels occupy the second dimension. Applies normalization across channels.

See LocalResponseNorm for details.

normalize¶

torch.nn.functional.normalize(input, p=2, dim=1, eps=1e-12, out=None)¶

Performs $L_p$ normalization of inputs over specified dimension.

For a tensor input of sizes $(n_0, ..., n_{dim}, ..., n_k)$, each $n_{dim}$ -element vector $v$ along dimension dim is transformed as

\[v = \frac{v}{\max(\lVert v \rVert_p, \epsilon)}. \]

With the default arguments it uses the Euclidean norm over vectors along dimension $1$ for normalization.

Parameters

input – input tensor of any shape
p (float) – the exponent value in the norm formulation. Default: 2
dim (int) – the dimension to reduce. Default: 1
eps (float) – small value to avoid division by zero. Default: 1e-12
out (Tensor, optional) – the output tensor. If out is used, this operation won’t be differentiable.

Linear functions¶

linear¶

torch.nn.functional.linear(input, weight, bias=None)¶

Applies a linear transformation to the incoming data: $y = xA^T + b$.

Shape:

Input: $(N, *, in\_features)$ where * means any number of additional dimensions

Weight: $(out\_features, in\_features)$

Bias: $(out\_features)$

Output: $(N, *, out\_features)$

bilinear¶

torch.nn.functional.bilinear(input1, input2, weight, bias=None)¶

Dropout functions¶

dropout¶

torch.nn.functional.dropout(input, p=0.5, training=True, inplace=False)¶

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

See Dropout for details.

Parameters

p – probability of an element to be zeroed. Default: 0.5
training – apply dropout if is True. Default: True
inplace – If set to True, will do this operation in-place. Default: False

alpha_dropout¶

torch.nn.functional.alpha_dropout(input, p=0.5, training=False, inplace=False)¶

Applies alpha dropout to the input.

See AlphaDropout for details.

dropout2d¶

torch.nn.functional.dropout2d(input, p=0.5, training=True, inplace=False)¶

Randomly zero out entire channels (a channel is a 2D feature map, e.g., the $j$-th channel of the $i$-th sample in the batched input is a 2D tensor $\text{input}[i, j]$) of the input tensor). Each channel will be zeroed out independently on every forward call with probability p using samples from a Bernoulli distribution.

See Dropout2d for details.

Parameters

p – probability of a channel to be zeroed. Default: 0.5
training – apply dropout if is True. Default: True
inplace – If set to True, will do this operation in-place. Default: False

dropout3d¶

torch.nn.functional.dropout3d(input, p=0.5, training=True, inplace=False)¶

Randomly zero out entire channels (a channel is a 3D feature map, e.g., the $j$-th channel of the $i$-th sample in the batched input is a 3D tensor $\text{input}[i, j]$) of the input tensor). Each channel will be zeroed out independently on every forward call with probability p using samples from a Bernoulli distribution.

See Dropout3d for details.

Parameters

p – probability of a channel to be zeroed. Default: 0.5
training – apply dropout if is True. Default: True
inplace – If set to True, will do this operation in-place. Default: False

Sparse functions¶

embedding¶

torch.nn.functional.embedding(input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False)¶

A simple lookup table that looks up embeddings in a fixed dictionary and size.

This module is often used to retrieve word embeddings using indices. The input to the module is a list of indices, and the embedding matrix, and the output is the corresponding word embeddings.

See torch.nn.Embedding for more details.

Parameters

input (LongTensor) – Tensor containing indices into the embedding matrix
weight (Tensor) – The embedding matrix with number of rows equal to the maximum possible index + 1, and number of columns equal to the embedding size
padding_idx (int, optional) – If given, pads the output with the embedding vector at padding_idx (initialized to zeros) whenever it encounters the index.
max_norm (float, optional) – If given, each embedding vector with norm larger than max_norm is renormalized to have norm max_norm. Note: this will modify weight in-place.
norm_type (float, optional) – The p of the p-norm to compute for the max_norm option. Default 2.
scale_grad_by_freq (boolean, optional) – If given, this will scale gradients by the inverse of frequency of the words in the mini-batch. Default False.
sparse (bool, optional) – If True, gradient w.r.t. weight will be a sparse tensor. See Notes under torch.nn.Embedding for more details regarding sparse gradients.

Shape:

Input: LongTensor of arbitrary shape containing the indices to extract
Weight: Embedding matrix of floating point type with shape (V, embedding_dim),
where V = maximum index + 1 and embedding_dim = the embedding size
Output: (*, embedding_dim), where * is the input shape

Examples:

>>> # a batch of 2 samples of 4 indices each
>>> input = torch.tensor([[1,2,4,5],[4,3,2,9]])
>>> # an embedding matrix containing 10 tensors of size 3
>>> embedding_matrix = torch.rand(10, 3)
>>> F.embedding(input, embedding_matrix)
tensor([[[ 0.8490,  0.9625,  0.6753],
         [ 0.9666,  0.7761,  0.6108],
         [ 0.6246,  0.9751,  0.3618],
         [ 0.4161,  0.2419,  0.7383]],

        [[ 0.6246,  0.9751,  0.3618],
         [ 0.0237,  0.7794,  0.0528],
         [ 0.9666,  0.7761,  0.6108],
         [ 0.3385,  0.8612,  0.1867]]])

>>> # example with padding_idx
>>> weights = torch.rand(10, 3)
>>> weights[0, :].zero_()
>>> embedding_matrix = weights
>>> input = torch.tensor([[0,2,0,5]])
>>> F.embedding(input, embedding_matrix, padding_idx=0)
tensor([[[ 0.0000,  0.0000,  0.0000],
         [ 0.5609,  0.5384,  0.8720],
         [ 0.0000,  0.0000,  0.0000],
         [ 0.6262,  0.2438,  0.7471]]])

embedding_bag¶

torch.nn.functional.embedding_bag(input, weight, offsets=None, max_norm=None, norm_type=2, scale_grad_by_freq=False, mode='mean', sparse=False)¶

Computes sums, means or maxes of bags of embeddings, without instantiating the intermediate embeddings.

See torch.nn.EmbeddingBag for more details.

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour in be backward that is not easily switched off. Please see the notes on /notes/randomness for background.

Parameters

input (LongTensor) – Tensor containing bags of indices into the embedding matrix
weight (Tensor) – The embedding matrix with number of rows equal to the maximum possible index + 1, and number of columns equal to the embedding size
offsets (LongTensor, optional) – Only used when input is 1D. offsets determines the starting index position of each bag (sequence) in input.
max_norm (float, optional) – If given, each embedding vector with norm larger than max_norm is renormalized to have norm max_norm. Note: this will modify weight in-place.
norm_type (float, optional) – The p in the p-norm to compute for the max_norm option. Default 2.
scale_grad_by_freq (boolean, optional) – if given, this will scale gradients by the inverse of frequency of the words in the mini-batch. Default False. Note: this option is not supported when mode="max".
mode (string, optional) – "sum", "mean" or "max". Specifies the way to reduce the bag. Default: "mean"
sparse (bool, optional) – if True, gradient w.r.t. weight will be a sparse tensor. See Notes under torch.nn.Embedding for more details regarding sparse gradients. Note: this option is not supported when mode="max".

Shape:

input (LongTensor) and offsets (LongTensor, optional)

If input is 2D of shape (B, N),

it will be treated as B bags (sequences) each of fixed length N, and this will return B values aggregated in a way depending on the mode. offsets is ignored and required to be None in this case.

If input is 1D of shape (N),

it will be treated as a concatenation of multiple bags (sequences). offsets is required to be a 1D tensor containing the starting index positions of each bag in input. Therefore, for offsets of shape (B), input will be viewed as having B bags. Empty bags (i.e., having 0-length) will have returned vectors filled by zeros.

weight (Tensor): the learnable weights of the module of shape (num_embeddings, embedding_dim)

output: aggregated embedding values of shape (B, embedding_dim)

Examples:

>>> # an Embedding module containing 10 tensors of size 3
>>> embedding_matrix = torch.rand(10, 3)
>>> # a batch of 2 samples of 4 indices each
>>> input = torch.tensor([1,2,4,5,4,3,2,9])
>>> offsets = torch.tensor([0,4])
>>> F.embedding_bag(embedding_matrix, input, offsets)
tensor([[ 0.3397,  0.3552,  0.5545],
        [ 0.5893,  0.4386,  0.5882]])

one_hot¶

torch.nn.functional.one_hot(tensor, num_classes=0) → LongTensor¶

Takes LongTensor with index values of shape (*) and returns a tensor of shape (*, num_classes) that have zeros everywhere except where the index of last dimension matches the corresponding value of the input tensor, in which case it will be 1.

Distance functions¶

pairwise_distance¶

torch.nn.functional.pairwise_distance(x1, x2, p=2.0, eps=1e-06, keepdim=False)¶: See torch.nn.PairwiseDistance for details

cosine_similarity¶

torch.nn.functional.cosine_similarity(x1, x2, dim=1, eps=1e-8) → Tensor¶

Returns cosine similarity between x1 and x2, computed along dim.

\[\text{similarity} = \dfrac{x_1 \cdot x_2}{\max(\Vert x_1 \Vert _2 \cdot \Vert x_2 \Vert _2, \epsilon)} \]

Parameters

x1 (Tensor) – First input.
x2 (Tensor) – Second input (of size matching x1).
dim (int, optional) – Dimension of vectors. Default: 1
eps (float, optional) – Small value to avoid division by zero. Default: 1e-8

Shape:

Input: $(\ast_1, D, \ast_2)$ where D is at position dim.
Output: $(\ast_1, \ast_2)$ where 1 is at position dim.

Example:

>>> input1 = torch.randn(100, 128)
>>> input2 = torch.randn(100, 128)
>>> output = F.cosine_similarity(input1, input2)
>>> print(output)

pdist¶

torch.nn.functional.pdist(input, p=2) → Tensor¶

Computes the p-norm distance between every pair of row vectors in the input. This is identical to the upper triangular portion, excluding the diagonal, of torch.norm(input[:, None] - input, dim=2, p=p). This function will be faster if the rows are contiguous.

If input has shape $N \times M$ then the output will have shape $\frac{1}{2} N (N - 1)$.

This function is equivalent to scipy.spatial.distance.pdist(input, ‘minkowski’, p=p) if $p \in (0, \infty)$. When $p = 0$ it is equivalent to scipy.spatial.distance.pdist(input, ‘hamming’) * M. When $p = \infty$, the closest scipy function is scipy.spatial.distance.pdist(xn, lambda x, y: np.abs(x - y).max()).

Parameters

input – input tensor of shape $N \times M$.
p – p value for the p-norm distance to calculate between each vector pair $\in [0, \infty]$.

Loss functions¶

binary_cross_entropy¶

torch.nn.functional.binary_cross_entropy(input, target, weight=None, size_average=None, reduce=None, reduction='mean')¶

Function that measures the Binary Cross Entropy between the target and the output.

See BCELoss for details.

Parameters

input – Tensor of arbitrary shape
target – Tensor of the same shape as input
weight (Tensor, optional) – a manual rescaling weight if provided it’s repeated to match input tensor shape
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Examples:

>>> input = torch.randn((3, 2), requires_grad=True)
>>> target = torch.rand((3, 2), requires_grad=False)
>>> loss = F.binary_cross_entropy(F.sigmoid(input), target)
>>> loss.backward()

binary_cross_entropy_with_logits¶

torch.nn.functional.binary_cross_entropy_with_logits(input, target, weight=None, size_average=None, reduce=None, reduction='mean', pos_weight=None)¶

Function that measures Binary Cross Entropy between target and output logits.

See BCEWithLogitsLoss for details.

Parameters

input – Tensor of arbitrary shape
target – Tensor of the same shape as input
weight (Tensor, optional) – a manual rescaling weight if provided it’s repeated to match input tensor shape
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'
pos_weight (Tensor, optional) – a weight of positive examples. Must be a vector with length equal to the number of classes.

Examples:

>>> input = torch.randn(3, requires_grad=True)
>>> target = torch.empty(3).random_(2)
>>> loss = F.binary_cross_entropy_with_logits(input, target)
>>> loss.backward()

poisson_nll_loss¶

torch.nn.functional.poisson_nll_loss(input, target, log_input=True, full=False, size_average=None, eps=1e-08, reduce=None, reduction='mean')¶

Poisson negative log likelihood loss.

See PoissonNLLLoss for details.

Parameters

input – expectation of underlying Poisson distribution.
target – random sample $target \sim \text{Poisson}(input)$.
log_input – if True the loss is computed as $\exp(\text{input}) - \text{target} * \text{input}$, if False then loss is $\text{input} - \text{target} * \log(\text{input}+\text{eps})$. Default: True
full – whether to compute full loss, i. e. to add the Stirling approximation term. Default: False $\text{target} * \log(\text{target}) - \text{target} + 0.5 * \log(2 * \pi * \text{target})$.
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
eps (float, optional) – Small value to avoid evaluation of $\log(0)$ when log_input`=``False`. Default: 1e-8
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

cosine_embedding_loss¶

torch.nn.functional.cosine_embedding_loss(input1, input2, target, margin=0, size_average=None, reduce=None, reduction='mean') → Tensor¶: See CosineEmbeddingLoss for details.

cross_entropy¶

torch.nn.functional.cross_entropy(input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')¶

This criterion combines log_softmax and nll_loss in a single function.

See CrossEntropyLoss for details.

Parameters

input (Tensor) – $(N, C)$ where C = number of classes or $(N, C, H, W)$ in case of 2D Loss, or $(N, C, d_1, d_2, ..., d_K)$ where $K \geq 1$ in the case of K-dimensional loss.
target (Tensor) – $(N)$ where each value is $0 \leq \text{targets}[i] \leq C-1$, or $(N, d_1, d_2, ..., d_K)$ where $K \geq 1$ for K-dimensional loss.
weight (Tensor, optional) – a manual rescaling weight given to each class. If given, has to be a Tensor of size C
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
ignore_index (int, optional) – Specifies a target value that is ignored and does not contribute to the input gradient. When size_average is True, the loss is averaged over non-ignored targets. Default: -100
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Examples:

>>> input = torch.randn(3, 5, requires_grad=True)
>>> target = torch.randint(5, (3,), dtype=torch.int64)
>>> loss = F.cross_entropy(input, target)
>>> loss.backward()

ctc_loss¶

torch.nn.functional.ctc_loss(log_probs, targets, input_lengths, target_lengths, blank=0, reduction='mean', zero_infinity=False)¶

The Connectionist Temporal Classification loss.

See CTCLoss for details.

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour in be backward that is not easily switched off. Please see the notes on /notes/randomness for background.

Parameters

log_probs – $(T, N, C)$ where C = number of characters in alphabet including blank, T = input length, and N = batch size. The logarithmized probabilities of the outputs (e.g. obtained with torch.nn.functional.log_softmax()).
targets – $(N, S)$ or (sum(target_lengths)). Targets cannot be blank. In the second form, the targets are assumed to be concatenated.
input_lengths – $(N)$. Lengths of the inputs (must each be $\leq T$)
target_lengths – $(N)$. Lengths of the targets
blank (int, optional) – Blank label. Default $0$.
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the output losses will be divided by the target lengths and then the mean over the batch is taken, 'sum': the output will be summed. Default: 'mean'
zero_infinity (bool, optional) – Whether to zero infinite losses and the associated gradients. Default: False Infinite losses mainly occur when the inputs are too short to be aligned to the targets.

Example:

>>> log_probs = torch.randn(50, 16, 20).log_softmax(2).detach().requires_grad_()
>>> targets = torch.randint(1, 20, (16, 30), dtype=torch.long)
>>> input_lengths = torch.full((16,), 50, dtype=torch.long)
>>> target_lengths = torch.randint(10,30,(16,), dtype=torch.long)
>>> loss = F.ctc_loss(log_probs, targets, input_lengths, target_lengths)
>>> loss.backward()

hinge_embedding_loss¶

torch.nn.functional.hinge_embedding_loss(input, target, margin=1.0, size_average=None, reduce=None, reduction='mean') → Tensor¶: See HingeEmbeddingLoss for details.

kl_div¶

torch.nn.functional.kl_div(input, target, size_average=None, reduce=None, reduction='mean')¶

The Kullback-Leibler divergence Loss.

See KLDivLoss for details.

Parameters

input – Tensor of arbitrary shape
target – Tensor of the same shape as input
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'batchmean' | 'sum' | 'mean'. 'none': no reduction will be applied 'batchmean': the sum of the output will be divided by the batchsize 'sum': the output will be summed 'mean': the output will be divided by the number of elements in the output Default: 'mean'

Note

size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction.

Note

l1_loss¶

torch.nn.functional.l1_loss(input, target, size_average=None, reduce=None, reduction='mean') → Tensor¶

Function that takes the mean element-wise absolute value difference.

See L1Loss for details.

mse_loss¶

torch.nn.functional.mse_loss(input, target, size_average=None, reduce=None, reduction='mean') → Tensor¶

Measures the element-wise mean squared error.

See MSELoss for details.

margin_ranking_loss¶

torch.nn.functional.margin_ranking_loss(input1, input2, target, margin=0, size_average=None, reduce=None, reduction='mean') → Tensor¶: See MarginRankingLoss for details.

multilabel_margin_loss¶

torch.nn.functional.multilabel_margin_loss(input, target, size_average=None, reduce=None, reduction='mean') → Tensor¶: See MultiLabelMarginLoss for details.

multilabel_soft_margin_loss¶

torch.nn.functional.multilabel_soft_margin_loss(input, target, weight=None, size_average=None) → Tensor¶: See MultiLabelSoftMarginLoss for details.

multi_margin_loss¶

torch.nn.functional.multi_margin_loss(input, target, p=1, margin=1.0, weight=None, size_average=None, reduce=None, reduction='mean')¶

multi_margin_loss(input, target, p=1, margin=1, weight=None, size_average=None,: reduce=None, reduction=’mean’) -> Tensor

See MultiMarginLoss for details.

nll_loss¶

torch.nn.functional.nll_loss(input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')¶

The negative log likelihood loss.

See NLLLoss for details.

Parameters

input – $(N, C)$ where C = number of classes or $(N, C, H, W)$ in case of 2D Loss, or $(N, C, d_1, d_2, ..., d_K)$ where $K \geq 1$ in the case of K-dimensional loss.
target – $(N)$ where each value is $0 \leq \text{targets}[i] \leq C-1$, or $(N, d_1, d_2, ..., d_K)$ where $K \geq 1$ for K-dimensional loss.
weight (Tensor, optional) – a manual rescaling weight given to each class. If given, has to be a Tensor of size C
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
ignore_index (int, optional) – Specifies a target value that is ignored and does not contribute to the input gradient. When size_average is True, the loss is averaged over non-ignored targets. Default: -100
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Example:

>>> # input is of size N x C = 3 x 5
>>> input = torch.randn(3, 5, requires_grad=True)
>>> # each element in target has to have 0 <= value < C
>>> target = torch.tensor([1, 0, 4])
>>> output = F.nll_loss(F.log_softmax(input), target)
>>> output.backward()

smooth_l1_loss¶

torch.nn.functional.smooth_l1_loss(input, target, size_average=None, reduce=None, reduction='mean')¶

Function that uses a squared term if the absolute element-wise error falls below 1 and an L1 term otherwise.

See SmoothL1Loss for details.

soft_margin_loss¶

torch.nn.functional.soft_margin_loss(input, target, size_average=None, reduce=None, reduction='mean') → Tensor¶: See SoftMarginLoss for details.

triplet_margin_loss¶

torch.nn.functional.triplet_margin_loss(anchor, positive, negative, margin=1.0, p=2, eps=1e-06, swap=False, size_average=None, reduce=None, reduction='mean')¶: See TripletMarginLoss for details

Vision functions¶

pixel_shuffle¶

torch.nn.functional.pixel_shuffle()¶

Rearranges elements in a tensor of shape $(*, C \times r^2, H, W)$ to a tensor of shape $(*, C, H \times r, W \times r)$.

See PixelShuffle for details.

Parameters

input (Tensor) – the input tensor
upscale_factor (int) – factor to increase spatial resolution by

Examples:

>>> input = torch.randn(1, 9, 4, 4)
>>> output = torch.nn.functional.pixel_shuffle(input, 3)
>>> print(output.size())
torch.Size([1, 1, 12, 12])

pad¶

torch.nn.functional.pad(input, pad, mode='constant', value=0)¶

Pads tensor.

Padding size:: The padding size by which to pad some dimensions of input are described starting from the last dimension and moving forward. $\left\lfloor\frac{\text{len(pad)}}{2}\right\rfloor$ dimensions of input will be padded. For example, to pad only the last dimension of the input tensor, then pad has the form $(\text{padding\_left}, \text{padding\_right})$; to pad the last 2 dimensions of the input tensor, then use $(\text{padding\_left}, \text{padding\_right},$ $\text{padding\_top}, \text{padding\_bottom})$; to pad the last 3 dimensions, use $(\text{padding\_left}, \text{padding\_right},$ $\text{padding\_top}, \text{padding\_bottom}$ $\text{padding\_front}, \text{padding\_back})$.
Padding mode:: See torch.nn.ConstantPad2d, torch.nn.ReflectionPad2d, and torch.nn.ReplicationPad2d for concrete examples on how each of the padding modes works. Constant padding is implemented for arbitrary dimensions. Replicate padding is implemented for padding the last 3 dimensions of 5D input tensor, or the last 2 dimensions of 4D input tensor, or the last dimension of 3D input tensor. Reflect padding is only implemented for padding the last 2 dimensions of 4D input tensor, or the last dimension of 3D input tensor.

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour in be backward that is not easily switched off. Please see the notes on /notes/randomness for background.

Parameters

input (Tensor) – N-dimensional tensor
pad (tuple) – m-elements tuple, where $\frac{m}{2} \leq$ input dimensions and $m$ is even.
mode – 'constant', 'reflect', 'replicate' or 'circular'. Default: 'constant'
value – fill value for 'constant' padding. Default: 0

Examples:

>>> t4d = torch.empty(3, 3, 4, 2)
>>> p1d = (1, 1) # pad last dim by 1 on each side
>>> out = F.pad(t4d, p1d, "constant", 0)  # effectively zero padding
>>> print(out.data.size())
torch.Size([3, 3, 4, 4])
>>> p2d = (1, 1, 2, 2) # pad last dim by (1, 1) and 2nd to last by (2, 2)
>>> out = F.pad(t4d, p2d, "constant", 0)
>>> print(out.data.size())
torch.Size([3, 3, 8, 4])
>>> t4d = torch.empty(3, 3, 4, 2)
>>> p3d = (0, 1, 2, 1, 3, 3) # pad by (0, 1), (2, 1), and (3, 3)
>>> out = F.pad(t4d, p3d, "constant", 0)
>>> print(out.data.size())
torch.Size([3, 9, 7, 3])

interpolate¶

torch.nn.functional.interpolate(input, size=None, scale_factor=None, mode='nearest', align_corners=None)¶

Down/up samples the input to either the given size or the given scale_factor

The algorithm used for interpolation is determined by mode.

Currently temporal, spatial and volumetric sampling are supported, i.e. expected inputs are 3-D, 4-D or 5-D in shape.

The input dimensions are interpreted in the form: mini-batch x channels x [optional depth] x [optional height] x width.

The modes available for resizing are: nearest, linear (3D-only), bilinear, bicubic (4D-only), trilinear (5D-only), area

Parameters

input (Tensor) – the input tensor
size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]) – output spatial size.
scale_factor (float or Tuple[float]) – multiplier for spatial size. Has to match input size if it is a tuple.
mode (str) – algorithm used for upsampling: 'nearest' | 'linear' | 'bilinear' | 'bicubic' | 'trilinear' | 'area'. Default: 'nearest'
align_corners (bool, optional) – Geometrically, we consider the pixels of the input and output as squares rather than points. If set to True, the input and output tensors are aligned by the center points of their corner pixels. If set to False, the input and output tensors are aligned by the corner points of their corner pixels, and the interpolation uses edge value padding for out-of-boundary values. This only has effect when mode is 'linear', 'bilinear', 'bicubic', or 'trilinear'. Default: False

Warning

With align_corners = True, the linearly interpolating modes (linear, bilinear, and trilinear) 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.

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour in be backward that is not easily switched off. Please see the notes on /notes/randomness for background.

upsample¶

torch.nn.functional.upsample(input, size=None, scale_factor=None, mode='nearest', align_corners=None)¶

Upsamples the input to either the given size or the given scale_factor

Warning

This function is deprecated in favor of torch.nn.functional.interpolate(). This is equivalent with nn.functional.interpolate(...).

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour in be backward that is not easily switched off. Please see the notes on /notes/randomness for background.

The algorithm used for upsampling is determined by mode.

Currently temporal, spatial and volumetric upsampling are supported, i.e. expected inputs are 3-D, 4-D or 5-D in shape.

The input dimensions are interpreted in the form: mini-batch x channels x [optional depth] x [optional height] x width.

The modes available for upsampling are: nearest, linear (3D-only), bilinear, bicubic (4D-only), trilinear (5D-only)

Parameters

input (Tensor) – the input tensor
size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]) – output spatial size.
scale_factor (float or Tuple[float]) – multiplier for spatial size. Has to be an integer.
mode (string) – algorithm used for upsampling: 'nearest' | 'linear' | 'bilinear' | 'bicubic' | 'trilinear'. Default: 'nearest'
align_corners (bool, optional) – Geometrically, we consider the pixels of the input and output as squares rather than points. If set to True, the input and output tensors are aligned by the center points of their corner pixels. If set to False, the input and output tensors are aligned by the corner points of their corner pixels, and the interpolation uses edge value padding for out-of-boundary values. This only has effect when mode is 'linear', 'bilinear', 'bicubic' or 'trilinear'. Default: False

Warning

upsample_nearest¶

torch.nn.functional.upsample_nearest(input, size=None, scale_factor=None)¶

Upsamples the input, using nearest neighbours’ pixel values.

Warning

This function is deprecated in favor of torch.nn.functional.interpolate(). This is equivalent with nn.functional.interpolate(..., mode='nearest').

Currently spatial and volumetric upsampling are supported (i.e. expected inputs are 4 or 5 dimensional).

Parameters

input (Tensor) – input
size (int or Tuple[int, int] or Tuple[int, int, int]) – output spatia size.
scale_factor (int) – multiplier for spatial size. Has to be an integer.

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour in be backward that is not easily switched off. Please see the notes on /notes/randomness for background.

upsample_bilinear¶

torch.nn.functional.upsample_bilinear(input, size=None, scale_factor=None)¶

Upsamples the input, using bilinear upsampling.

Warning

This function is deprecated in favor of torch.nn.functional.interpolate(). This is equivalent with nn.functional.interpolate(..., mode='bilinear', align_corners=True).

Expected inputs are spatial (4 dimensional). Use upsample_trilinear fo volumetric (5 dimensional) inputs.

Parameters

input (Tensor) – input
size (int or Tuple[int, int]) – output spatial size.
scale_factor (int or Tuple[int, int]) – multiplier for spatial size

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour in be backward that is not easily switched off. Please see the notes on /notes/randomness for background.

grid_sample¶

torch.nn.functional.grid_sample(input, grid, mode='bilinear', padding_mode='zeros')¶

Given an input and a flow-field grid, computes the output using input values and pixel locations from grid.

Currently, only spatial (4-D) and volumetric (5-D) input are supported.

In the spatial (4-D) case, for input with shape $(N, C, H_\text{in}, W_\text{in})$ and grid with shape $(N, H_\text{out}, W_\text{out}, 2)$, the output will have shape $(N, C, H_\text{out}, W_\text{out})$.

For each output location output[n, :, h, w], the size-2 vector grid[n, h, w] specifies input pixel locations x and y, which are used to interpolate the output value output[n, :, h, w]. In the case of 5D inputs, grid[n, d, h, w] specifies the x, y, z pixel locations for interpolating output[n, :, d, h, w]. mode argument specifies nearest or bilinear interpolation method to sample the input pixels.

grid should have most values in the range of [-1, 1]. This is because the pixel locations are normalized by the input spatial dimensions. For example, values x = -1, y = -1 is the left-top pixel of input, and values x = 1, y = 1 is the right-bottom pixel of input.

If grid has values outside the range of [-1, 1], those locations are handled as defined by padding_mode. Options are

padding_mode="zeros": use 0 for out-of-bound values,

padding_mode="border": use border values for out-of-bound values,

padding_mode="reflection": use values at locations reflected by the border for out-of-bound values. For location far away from the border, it will keep being reflected until becoming in bound, e.g., (normalized) pixel location x = -3.5 reflects by -1 and becomes x' = 1.5, then reflects by border 1 and becomes x'' = -0.5.

Note

This function is often used in building Spatial Transformer Networks.

Note

When using the CUDA backend, this operation may induce nondeterministic behaviour in be backward that is not easily switched off. Please see the notes on /notes/randomness for background.

Parameters

input (Tensor) – input of shape $(N, C, H_\text{in}, W_\text{in})$ (4-D case) or $(N, C, D_\text{in}, H_\text{in}, W_\text{in})$ (5-D case)
grid (Tensor) – flow-field of shape $(N, H_\text{out}, W_\text{out}, 2)$ (4-D case) or $(N, D_\text{out}, H_\text{out}, W_\text{out}, 3)$ (5-D case)
mode (str) – interpolation mode to calculate output values 'bilinear' | 'nearest'. Default: 'bilinear'
padding_mode (str) – padding mode for outside grid values 'zeros' | 'border' | 'reflection'. Default: 'zeros'

Returns

output Tensor

Return type

output (Tensor)

affine_grid¶

torch.nn.functional.affine_grid(theta, size)¶

Generates a 2d flow field, given a batch of affine matrices theta. Generally used in conjunction with grid_sample() to implement Spatial Transformer Networks.

Parameters

theta (Tensor) – input batch of affine matrices ($N \times 2 \times 3$)
size (torch.Size) – the target output image size ($N \times C \times H \times W$). Example: torch.Size((32, 3, 24, 24))

Returns

output Tensor of size ($N \times H \times W \times 2$)

Return type

output (Tensor)

DataParallel functions (multi-GPU, distributed)¶

data_parallel¶

torch.nn.parallel.data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None)¶

Evaluates module(input) in parallel across the GPUs given in device_ids.

This is the functional version of the DataParallel module.

Parameters

module (Module) – the module to evaluate in parallel
inputs (Tensor) – inputs to the module
device_ids (list of python:int or torch.device) – GPU ids on which to replicate module
output_device (list of python:int or torch.device) – GPU location of the output Use -1 to indicate the CPU. (default: device_ids[0])

Returns

a Tensor containing the result of module(input) located on output_device

torch.nn.init¶

torch.nn.init.calculate_gain(nonlinearity, param=None)¶

Return the recommended gain value for the given nonlinearity function. The values are as follows:

nonlinearity	gain
Linear / Identity	$1$
Conv{1,2,3}D	$1$
Sigmoid	$1$
Tanh	$\frac{5}{3}$
ReLU	$\sqrt{2}$
Leaky Relu	$\sqrt{\frac{2}{1 + \text{negative\_slope}^2}}$

Parameters

nonlinearity – the non-linear function (nn.functional name)
param – optional parameter for the non-linear function

Examples

>>> gain = nn.init.calculate_gain('leaky_relu')

torch.nn.init.uniform_(tensor, a=0, b=1)¶

Fills the input Tensor with values drawn from the uniform distribution $\mathcal{U}(a, b)$.

Parameters

tensor – an n-dimensional torch.Tensor
a – the lower bound of the uniform distribution
b – the upper bound of the uniform distribution

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.uniform_(w)

torch.nn.init.normal_(tensor, mean=0, std=1)¶

Fills the input Tensor with values drawn from the normal distribution $\mathcal{N}(\text{mean}, \text{std})$.

Parameters

tensor – an n-dimensional torch.Tensor
mean – the mean of the normal distribution
std – the standard deviation of the normal distribution

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.normal_(w)

torch.nn.init.constant_(tensor, val)¶

Fills the input Tensor with the value $\text{val}$.

Parameters

tensor – an n-dimensional torch.Tensor
val – the value to fill the tensor with

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.constant_(w, 0.3)

torch.nn.init.eye_(tensor)¶

Fills the 2-dimensional input Tensor with the identity matrix. Preserves the identity of the inputs in Linear layers, where as many inputs are preserved as possible.

Parameters: tensor – a 2-dimensional torch.Tensor

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.eye_(w)

torch.nn.init.dirac_(tensor)¶

Fills the {3, 4, 5}-dimensional input Tensor with the Dirac delta function. Preserves the identity of the inputs in Convolutional layers, where as many input channels are preserved as possible.

Parameters: tensor – a {3, 4, 5}-dimensional torch.Tensor

Examples

>>> w = torch.empty(3, 16, 5, 5)
>>> nn.init.dirac_(w)

torch.nn.init.xavier_uniform_(tensor, gain=1)¶

Fills the input Tensor with values according to the method described in Understanding the difficulty of training deep feedforward neural networks - Glorot, X. & Bengio, Y. (2010), using a uniform distribution. The resulting tensor will have values sampled from $\mathcal{U}(-a, a)$ where

\[a = \text{gain} \times \sqrt{\frac{6}{\text{fan\_in} + \text{fan\_out}}} \]

Also known as Glorot initialization.

Parameters

tensor – an n-dimensional torch.Tensor
gain – an optional scaling factor

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.xavier_uniform_(w, gain=nn.init.calculate_gain('relu'))

torch.nn.init.xavier_normal_(tensor, gain=1)¶

Fills the input Tensor with values according to the method described in Understanding the difficulty of training deep feedforward neural networks - Glorot, X. & Bengio, Y. (2010), using a normal distribution. The resulting tensor will have values sampled from $\mathcal{N}(0, \text{std})$ where

\[\text{std} = \text{gain} \times \sqrt{\frac{2}{\text{fan\_in} + \text{fan\_out}}} \]

Also known as Glorot initialization.

Parameters

tensor – an n-dimensional torch.Tensor
gain – an optional scaling factor

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.xavier_normal_(w)

torch.nn.init.kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')¶

Fills the input Tensor with values according to the method described in Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification - He, K. et al. (2015), using a uniform distribution. The resulting tensor will have values sampled from $\mathcal{U}(-\text{bound}, \text{bound})$ where

\[\text{bound} = \sqrt{\frac{6}{(1 + a^2) \times \text{fan\_in}}} \]

Also known as He initialization.

Parameters

tensor – an n-dimensional torch.Tensor
a – the negative slope of the rectifier used after this layer (0 for ReLU by default)
mode – either 'fan_in' (default) or 'fan_out'. Choosing 'fan_in' preserves the magnitude of the variance of the weights in the forward pass. Choosing 'fan_out' preserves the magnitudes in the backwards pass.
nonlinearity – the non-linear function (nn.functional name), recommended to use only with 'relu' or 'leaky_relu' (default).

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.kaiming_uniform_(w, mode='fan_in', nonlinearity='relu')

torch.nn.init.kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')¶

Fills the input Tensor with values according to the method described in Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification - He, K. et al. (2015), using a normal distribution. The resulting tensor will have values sampled from $\mathcal{N}(0, \text{std})$ where

\[\text{std} = \sqrt{\frac{2}{(1 + a^2) \times \text{fan\_in}}} \]

Also known as He initialization.

Parameters

tensor – an n-dimensional torch.Tensor
a – the negative slope of the rectifier used after this layer (0 for ReLU by default)
mode – either 'fan_in' (default) or 'fan_out'. Choosing 'fan_in' preserves the magnitude of the variance of the weights in the forward pass. Choosing 'fan_out' preserves the magnitudes in the backwards pass.
nonlinearity – the non-linear function (nn.functional name), recommended to use only with 'relu' or 'leaky_relu' (default).

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.kaiming_normal_(w, mode='fan_out', nonlinearity='relu')

torch.nn.init.orthogonal_(tensor, gain=1)¶

Fills the input Tensor with a (semi) orthogonal matrix, as described in Exact solutions to the nonlinear dynamics of learning in deep linear neural networks - Saxe, A. et al. (2013). The input tensor must have at least 2 dimensions, and for tensors with more than 2 dimensions the trailing dimensions are flattened.

Parameters

tensor – an n-dimensional torch.Tensor, where $n \geq 2$
gain – optional scaling factor

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.orthogonal_(w)

torch.nn.init.sparse_(tensor, sparsity, std=0.01)¶

Fills the 2D input Tensor as a sparse matrix, where the non-zero elements will be drawn from the normal distribution $\mathcal{N}(0, 0.01)$, as described in Deep learning via Hessian-free optimization - Martens, J. (2010).

Parameters

tensor – an n-dimensional torch.Tensor
sparsity – The fraction of elements in each column to be set to zero
std – the standard deviation of the normal distribution used to generate the non-zero values

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.sparse_(w, sparsity=0.1)

torch.optim¶

torch.optim is a package implementing various optimization algorithms. Most commonly used methods are already supported, and the interface is general enough, so that more sophisticated ones can be also easily integrated in the future.

How to use an optimizer¶

To use torch.optim you have to construct an optimizer object, that will hold the current state and will update the parameters based on the computed gradients.

Constructing it¶

To construct an Optimizer you have to give it an iterable containing the parameters (all should be Variable s) to optimize. Then, you can specify optimizer-specific options such as the learning rate, weight decay, etc.

Note

If you need to move a model to GPU via .cuda(), please do so before constructing optimizers for it. Parameters of a model after .cuda() will be different objects with those before the call.

In general, you should make sure that optimized parameters live in consistent locations when optimizers are constructed and used.

Example:

optimizer = optim.SGD(model.parameters(), lr = 0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr = 0.0001)

Per-parameter options¶

Optimizer s also support specifying per-parameter options. To do this, instead of passing an iterable of Variable s, pass in an iterable of dict s. Each of them will define a separate parameter group, and should contain a params key, containing a list of parameters belonging to it. Other keys should match the keyword arguments accepted by the optimizers, and will be used as optimization options for this group.

Note

You can still pass options as keyword arguments. They will be used as defaults, in the groups that didn’t override them. This is useful when you only want to vary a single option, while keeping all others consistent between parameter groups.

For example, this is very useful when one wants to specify per-layer learning rates:

optim.SGD([
                {'params': model.base.parameters()},
                {'params': model.classifier.parameters(), 'lr': 1e-3}
            ], lr=1e-2, momentum=0.9)

This means that model.base’s parameters will use the default learning rate of 1e-2, model.classifier’s parameters will use a learning rate of 1e-3, and a momentum of 0.9 will be used for all parameters.

Taking an optimization step¶

All optimizers implement a step() method, that updates the parameters. It can be used in two ways:

`optimizer.step()`¶

This is a simplified version supported by most optimizers. The function can be called once the gradients are computed using e.g. backward().

Example:

for input, target in dataset:
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

`optimizer.step(closure)`¶

Some optimization algorithms such as Conjugate Gradient and LBFGS need to reevaluate the function multiple times, so you have to pass in a closure that allows them to recompute your model. The closure should clear the gradients, compute the loss, and return it.

Example:

for input, target in dataset:
    def closure():
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        return loss
    optimizer.step(closure)

Algorithms¶

class torch.optim.Optimizer(params, defaults)¶

Base class for all optimizers.

Warning

Parameters need to be specified as collections that have a deterministic ordering that is consistent between runs. Examples of objects that don’t satisfy those properties are sets and iterators over values of dictionaries.

Parameters

params (iterable) – an iterable of torch.Tensor s or dict s. Specifies what Tensors should be optimized.
defaults – (dict): a dict containing default values of optimization options (used when a parameter group doesn’t specify them).

add_param_group(param_group)¶

Add a param group to the Optimizer s param_groups.

This can be useful when fine tuning a pre-trained network as frozen layers can be made trainable and added to the Optimizer as training progresses.

Parameters

param_group (dict) – Specifies what Tensors should be optimized along with group
optimization options. (specific) –

load_state_dict(state_dict)¶

Loads the optimizer state.

Parameters: state_dict (dict) – optimizer state. Should be an object returned from a call to state_dict().

state_dict()¶

Returns the state of the optimizer as a dict.

It contains two entries:

state - a dict holding current optimization state. Its content
differs between optimizer classes.
param_groups - a dict containing all parameter groups

step(closure)¶

Performs a single optimization step (parameter update).

Parameters: closure (callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

zero_grad()¶: Clears the gradients of all optimized torch.Tensor s.

class torch.optim.Adadelta(params, lr=1.0, rho=0.9, eps=1e-06, weight_decay=0)¶

Implements Adadelta algorithm.

It has been proposed in ADADELTA: An Adaptive Learning Rate Method.

Parameters

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
rho (float, optional) – coefficient used for computing a running average of squared gradients (default: 0.9)
eps (float, optional) – term added to the denominator to improve numerical stability (default: 1e-6)
lr (float, optional) – coefficient that scale delta before it is applied to the parameters (default: 1.0)
weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)

step(closure=None)¶

Performs a single optimization step.

Parameters: closure (callable, optional) – A closure that reevaluates the model and returns the loss.

class torch.optim.Adagrad(params, lr=0.01, lr_decay=0, weight_decay=0, initial_accumulator_value=0)¶

Implements Adagrad algorithm.

It has been proposed in Adaptive Subgradient Methods for Online Learning and Stochastic Optimization.

Parameters

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float, optional) – learning rate (default: 1e-2)
lr_decay (float, optional) – learning rate decay (default: 0)
weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)

step(closure=None)¶

Performs a single optimization step.

Parameters: closure (callable, optional) – A closure that reevaluates the model and returns the loss.

class torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)¶

Implements Adam algorithm.

It has been proposed in Adam: A Method for Stochastic Optimization.

Parameters

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float, optional) – learning rate (default: 1e-3)
betas (Tuple[float, float], optional) – coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional) – term added to the denominator to improve numerical stability (default: 1e-8)
weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)
amsgrad (boolean, optional) – whether to use the AMSGrad variant of this algorithm from the paper On the Convergence of Adam and Beyond (default: False)

step(closure=None)¶

Performs a single optimization step.

Parameters: closure (callable, optional) – A closure that reevaluates the model and returns the loss.

class torch.optim.SparseAdam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08)¶

Implements lazy version of Adam algorithm suitable for sparse tensors.

In this variant, only moments that show up in the gradient get updated, and only those portions of the gradient get applied to the parameters.

Parameters

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float, optional) – learning rate (default: 1e-3)
betas (Tuple[float, float], optional) – coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional) – term added to the denominator to improve numerical stability (default: 1e-8)

step(closure=None)¶

Performs a single optimization step.

Parameters: closure (callable, optional) – A closure that reevaluates the model and returns the loss.

class torch.optim.Adamax(params, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)¶

Implements Adamax algorithm (a variant of Adam based on infinity norm).

It has been proposed in Adam: A Method for Stochastic Optimization.

Parameters

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float, optional) – learning rate (default: 2e-3)
betas (Tuple[float, float], optional) – coefficients used for computing running averages of gradient and its square
eps (float, optional) – term added to the denominator to improve numerical stability (default: 1e-8)
weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)

step(closure=None)¶

Performs a single optimization step.

Parameters: closure (callable, optional) – A closure that reevaluates the model and returns the loss.

class torch.optim.ASGD(params, lr=0.01, lambd=0.0001, alpha=0.75, t0=1000000.0, weight_decay=0)¶

Implements Averaged Stochastic Gradient Descent.

It has been proposed in Acceleration of stochastic approximation by averaging.

Parameters

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float, optional) – learning rate (default: 1e-2)
lambd (float, optional) – decay term (default: 1e-4)
alpha (float, optional) – power for eta update (default: 0.75)
t0 (float, optional) – point at which to start averaging (default: 1e6)
weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)

step(closure=None)¶

Performs a single optimization step.

Parameters: closure (callable, optional) – A closure that reevaluates the model and returns the loss.

class torch.optim.LBFGS(params, lr=1, max_iter=20, max_eval=None, tolerance_grad=1e-05, tolerance_change=1e-09, history_size=100, line_search_fn=None)¶

Implements L-BFGS algorithm.

Warning

This optimizer doesn’t support per-parameter options and parameter groups (there can be only one).

Warning

Right now all parameters have to be on a single device. This will be improved in the future.

Note

This is a very memory intensive optimizer (it requires additional param_bytes * (history_size + 1) bytes). If it doesn’t fit in memory try reducing the history size, or use a different algorithm.

Parameters

lr (float) – learning rate (default: 1)
max_iter (int) – maximal number of iterations per optimization step (default: 20)
max_eval (int) – maximal number of function evaluations per optimization step (default: max_iter * 1.25).
tolerance_grad (float) – termination tolerance on first order optimality (default: 1e-5).
tolerance_change (float) – termination tolerance on function value/parameter changes (default: 1e-9).
history_size (int) – update history size (default: 100).

step(closure)¶

Performs a single optimization step.

Parameters: closure (callable) – A closure that reevaluates the model and returns the loss.

class torch.optim.RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)¶

Implements RMSprop algorithm.

Proposed by G. Hinton in his course.

The centered version first appears in Generating Sequences With Recurrent Neural Networks.

Parameters

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float, optional) – learning rate (default: 1e-2)
momentum (float, optional) – momentum factor (default: 0)
alpha (float, optional) – smoothing constant (default: 0.99)
eps (float, optional) – term added to the denominator to improve numerical stability (default: 1e-8)
centered (bool, optional) – if True, compute the centered RMSProp, the gradient is normalized by an estimation of its variance
weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)

step(closure=None)¶

Performs a single optimization step.

Parameters: closure (callable, optional) – A closure that reevaluates the model and returns the loss.

class torch.optim.Rprop(params, lr=0.01, etas=(0.5, 1.2), step_sizes=(1e-06, 50))¶

Implements the resilient backpropagation algorithm.

Parameters

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float, optional) – learning rate (default: 1e-2)
etas (Tuple[float, float], optional) – pair of (etaminus, etaplis), that are multiplicative increase and decrease factors (default: (0.5, 1.2))
step_sizes (Tuple[float, float], optional) – a pair of minimal and maximal allowed step sizes (default: (1e-6, 50))

step(closure=None)¶

Performs a single optimization step.

Parameters: closure (callable, optional) – A closure that reevaluates the model and returns the loss.

class torch.optim.SGD(params, lr=<required parameter>, momentum=0, dampening=0, weight_decay=0, nesterov=False)¶

Implements stochastic gradient descent (optionally with momentum).

Nesterov momentum is based on the formula from On the importance of initialization and momentum in deep learning.

Parameters

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float) – learning rate
momentum (float, optional) – momentum factor (default: 0)
weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)
dampening (float, optional) – dampening for momentum (default: 0)
nesterov (bool, optional) – enables Nesterov momentum (default: False)

Example

>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> optimizer.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optimizer.step()

Note

The implementation of SGD with Momentum/Nesterov subtly differs from Sutskever et. al. and implementations in some other frameworks.

Considering the specific case of Momentum, the update can be written as

\[v = \rho * v + g \\ p = p - lr * v \]

where p, g, v and $\rho$ denote the parameters, gradient, velocity, and momentum respectively.

This is in contrast to Sutskever et. al. and other frameworks which employ an update of the form

\[v = \rho * v + lr * g \\ p = p - v \]

The Nesterov version is analogously modified.

step(closure=None)¶

Performs a single optimization step.

Parameters: closure (callable, optional) – A closure that reevaluates the model and returns the loss.

How to adjust Learning Rate¶

torch.optim.lr_scheduler provides several methods to adjust the learning rate based on the number of epochs. torch.optim.lr_scheduler.ReduceLROnPlateau allows dynamic learning rate reducing based on some validation measurements.

class torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1)¶

Sets the learning rate of each parameter group to the initial lr times a given function. When last_epoch=-1, sets initial lr as lr.

Parameters

optimizer (Optimizer) – Wrapped optimizer.
lr_lambda (function or list) – A function which computes a multiplicative factor given an integer parameter epoch, or a list of such functions, one for each group in optimizer.param_groups.
last_epoch (int) – The index of last epoch. Default: -1.

Example

>>> # Assuming optimizer has two groups.
>>> lambda1 = lambda epoch: epoch // 30
>>> lambda2 = lambda epoch: 0.95 ** epoch
>>> scheduler = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2])
>>> for epoch in range(100):
>>>     scheduler.step()
>>>     train(...)
>>>     validate(...)

load_state_dict(state_dict)¶

Loads the schedulers state.

Parameters: state_dict (dict) – scheduler state. Should be an object returned from a call to state_dict().

state_dict()¶

Returns the state of the scheduler as a dict.

It contains an entry for every variable in self.__dict__ which is not the optimizer. The learning rate lambda functions will only be saved if they are callable objects and not if they are functions or lambdas.

class torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1)¶

Decays the learning rate of each parameter group by gamma every step_size epochs. Notice that such decay can happen simultaneously with other changes to the learning rate from outside this scheduler. When last_epoch=-1, sets initial lr as lr.

Parameters

optimizer (Optimizer) – Wrapped optimizer.
step_size (int) – Period of learning rate decay.
gamma (float) – Multiplicative factor of learning rate decay. Default: 0.1.
last_epoch (int) – The index of last epoch. Default: -1.

Example

>>> # Assuming optimizer uses lr = 0.05 for all groups
>>> # lr = 0.05     if epoch < 30
>>> # lr = 0.005    if 30 <= epoch < 60
>>> # lr = 0.0005   if 60 <= epoch < 90
>>> # ...
>>> scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
>>> for epoch in range(100):
>>>     scheduler.step()
>>>     train(...)
>>>     validate(...)

class torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma=0.1, last_epoch=-1)¶

Decays the learning rate of each parameter group by gamma once the number of epoch reaches one of the milestones. Notice that such decay can happen simultaneously with other changes to the learning rate from outside this scheduler. When last_epoch=-1, sets initial lr as lr.

Parameters

optimizer (Optimizer) – Wrapped optimizer.
milestones (list) – List of epoch indices. Must be increasing.
gamma (float) – Multiplicative factor of learning rate decay. Default: 0.1.
last_epoch (int) – The index of last epoch. Default: -1.

Example

>>> # Assuming optimizer uses lr = 0.05 for all groups
>>> # lr = 0.05     if epoch < 30
>>> # lr = 0.005    if 30 <= epoch < 80
>>> # lr = 0.0005   if epoch >= 80
>>> scheduler = MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)
>>> for epoch in range(100):
>>>     scheduler.step()
>>>     train(...)
>>>     validate(...)

class torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma, last_epoch=-1)¶

Decays the learning rate of each parameter group by gamma every epoch. When last_epoch=-1, sets initial lr as lr.

Parameters

optimizer (Optimizer) – Wrapped optimizer.
gamma (float) – Multiplicative factor of learning rate decay.
last_epoch (int) – The index of last epoch. Default: -1.

class torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max, eta_min=0, last_epoch=-1)¶

Set the learning rate of each parameter group using a cosine annealing schedule, where $\eta_{max}$ is set to the initial lr and $T_{cur}$ is the number of epochs since the last restart in SGDR:

\[\eta_{t+1} = \eta_{min} + (\eta_t - \eta_{min})\frac{1 + \cos(\frac{T_{cur+1}}{T_{max}}\pi)}{1 + \cos(\frac{T_{cur}}{T_{max}}\pi)} \]

When last_epoch=-1, sets initial lr as lr. Notice that because the schedule is defined recursively, the learning rate can be simultaneously modified outside this scheduler by other operators. If the learning rate is set solely by this scheduler, the learning rate at each step becomes:

\[\eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})(1 + \cos(\frac{T_{cur}}{T_{max}}\pi)) \]

It has been proposed in SGDR: Stochastic Gradient Descent with Warm Restarts. Note that this only implements the cosine annealing part of SGDR, and not the restarts.

Parameters

optimizer (Optimizer) – Wrapped optimizer.
T_max (int) – Maximum number of iterations.
eta_min (float) – Minimum learning rate. Default: 0.
last_epoch (int) – The index of last epoch. Default: -1.

class torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, verbose=False, threshold=0.0001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-08)¶

Reduce learning rate when a metric has stopped improving. Models often benefit from reducing the learning rate by a factor of 2-10 once learning stagnates. This scheduler reads a metrics quantity and if no improvement is seen for a ‘patience’ number of epochs, the learning rate is reduced.

Parameters

optimizer (Optimizer) – Wrapped optimizer.
mode (str) – One of min, max. In min mode, lr will be reduced when the quantity monitored has stopped decreasing; in max mode it will be reduced when the quantity monitored has stopped increasing. Default: ‘min’.
factor (float) – Factor by which the learning rate will be reduced. new_lr = lr * factor. Default: 0.1.
patience (int) – Number of epochs with no improvement after which learning rate will be reduced. For example, if patience = 2, then we will ignore the first 2 epochs with no improvement, and will only decrease the LR after the 3rd epoch if the loss still hasn’t improved then. Default: 10.
verbose (bool) – If True, prints a message to stdout for each update. Default: False.
threshold (float) – Threshold for measuring the new optimum, to only focus on significant changes. Default: 1e-4.
threshold_mode (str) – One of rel, abs. In rel mode, dynamic_threshold = best * ( 1 + threshold ) in ‘max’ mode or best * ( 1 - threshold ) in min mode. In abs mode, dynamic_threshold = best + threshold in max mode or best - threshold in min mode. Default: ‘rel’.
cooldown (int) – Number of epochs to wait before resuming normal operation after lr has been reduced. Default: 0.
min_lr (float or list) – A scalar or a list of scalars. A lower bound on the learning rate of all param groups or each group respectively. Default: 0.
eps (float) – Minimal decay applied to lr. If the difference between new and old lr is smaller than eps, the update is ignored. Default: 1e-8.

Example

>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> scheduler = ReduceLROnPlateau(optimizer, 'min')
>>> for epoch in range(10):
>>>     train(...)
>>>     val_loss = validate(...)
>>>     # Note that step should be called after validate()
>>>     scheduler.step(val_loss)

Automatic differentiation package - torch.autograd¶

torch.autograd provides classes and functions implementing automatic differentiation of arbitrary scalar valued functions. It requires minimal changes to the existing code - you only need to declare Tensor s for which gradients should be computed with the requires_grad=True keyword.

torch.autograd.backward(tensors, grad_tensors=None, retain_graph=None, create_graph=False, grad_variables=None)¶

Computes the sum of gradients of given tensors w.r.t. graph leaves.

The graph is differentiated using the chain rule. If any of tensors are non-scalar (i.e. their data has more than one element) and require gradient, then the Jacobian-vector product would be computed, in this case the function additionally requires specifying grad_tensors. It should be a sequence of matching length, that contains the “vector” in the Jacobian-vector product, usually the gradient of the differentiated function w.r.t. corresponding tensors (None is an acceptable value for all tensors that don’t need gradient tensors).

This function accumulates gradients in the leaves - you might need to zero them before calling it.

Parameters

tensors (sequence of Tensor) – Tensors of which the derivative will be computed.
grad_tensors (sequence of (Tensor or None)) – The “vector” in the Jacobian-vector product, usually gradients w.r.t. each element of corresponding tensors. None values can be specified for scalar Tensors or ones that don’t require grad. If a None value would be acceptable for all grad_tensors, then this argument is optional.
retain_graph (bool, optional) – If False, the graph used to compute the grad will be freed. Note that in nearly all cases setting this option to True is not needed and often can be worked around in a much more efficient way. Defaults to the value of create_graph.
create_graph (bool, optional) – If True, graph of the derivative will be constructed, allowing to compute higher order derivative products. Defaults to False.

torch.autograd.grad(outputs, inputs, grad_outputs=None, retain_graph=None, create_graph=False, only_inputs=True, allow_unused=False)¶

Computes and returns the sum of gradients of outputs w.r.t. the inputs.

grad_outputs should be a sequence of length matching output containing the “vector” in Jacobian-vector product, usually the pre-computed gradients w.r.t. each of the outputs. If an output doesn’t require_grad, then the gradient can be None).

If only_inputs is True, the function will only return a list of gradients w.r.t the specified inputs. If it’s False, then gradient w.r.t. all remaining leaves will still be computed, and will be accumulated into their .grad attribute.

Parameters

outputs (sequence of Tensor) – outputs of the differentiated function.
inputs (sequence of Tensor) – Inputs w.r.t. which the gradient will be returned (and not accumulated into .grad).
grad_outputs (sequence of Tensor) – The “vector” in the Jacobian-vector product. Usually gradients w.r.t. each output. None values can be specified for scalar Tensors or ones that don’t require grad. If a None value would be acceptable for all grad_tensors, then this argument is optional. Default: None.
retain_graph (bool, optional) – If False, the graph used to compute the grad will be freed. Note that in nearly all cases setting this option to True is not needed and often can be worked around in a much more efficient way. Defaults to the value of create_graph.
create_graph (bool, optional) – If True, graph of the derivative will be constructed, allowing to compute higher order derivative products. Default: False.
allow_unused (bool, optional) – If False, specifying inputs that were not used when computing outputs (and therefore their grad is always zero) is an error. Defaults to False.

Locally disabling gradient computation¶

class torch.autograd.no_grad¶

Context-manager that disabled gradient calculation.

Disabling gradient calculation is useful for inference, when you are sure that you will not call Tensor.backward(). It will reduce memory consumption for computations that would otherwise have requires_grad=True. In this mode, the result of every computation will have requires_grad=False, even when the inputs have requires_grad=True.

Also functions as a decorator.

Example:

>>> x = torch.tensor([1], requires_grad=True)
>>> with torch.no_grad():
...   y = x * 2
>>> y.requires_grad
False
>>> @torch.no_grad()
... def doubler(x):
...     return x * 2
>>> z = doubler(x)
>>> z.requires_grad
False

class torch.autograd.enable_grad¶

Context-manager that enables gradient calculation.

Enables gradient calculation inside a no_grad context. This has no effect outside of no_grad.

Also functions as a decorator.

Example:

>>> x = torch.tensor([1], requires_grad=True)
>>> with torch.no_grad():
...   with torch.enable_grad():
...     y = x * 2
>>> y.requires_grad
True
>>> y.backward()
>>> x.grad
>>> @torch.enable_grad()
... def doubler(x):
...     return x * 2
>>> with torch.no_grad():
...     z = doubler(x)
>>> z.requires_grad
True

class torch.autograd.set_grad_enabled(mode)¶

Context-manager that sets gradient calculation to on or off.

set_grad_enabled will enable or disable grads based on its argument mode. It can be used as a context-manager or as a function.

Parameters: mode (bool) – Flag whether to enable grad (True), or disable (False). This can be used to conditionally enable gradients.

Example:

>>> x = torch.tensor([1], requires_grad=True)
>>> is_train = False
>>> with torch.set_grad_enabled(is_train):
...   y = x * 2
>>> y.requires_grad
False
>>> torch.set_grad_enabled(True)
>>> y = x * 2
>>> y.requires_grad
True
>>> torch.set_grad_enabled(False)
>>> y = x * 2
>>> y.requires_grad
False

In-place operations on Tensors¶

Supporting in-place operations in autograd is a hard matter, and we discourage their use in most cases. Autograd’s aggressive buffer freeing and reuse makes it very efficient and there are very few occasions when in-place operations actually lower memory usage by any significant amount. Unless you’re operating under heavy memory pressure, you might never need to use them.

In-place correctness checks¶

All Tensor s keep track of in-place operations applied to them, and if the implementation detects that a tensor was saved for backward in one of the functions, but it was modified in-place afterwards, an error will be raised once backward pass is started. This ensures that if you’re using in-place functions and not seeing any errors, you can be sure that the computed gradients are correct.

Variable (deprecated)¶

Warning

The Variable API has been deprecated: Variables are no longer necessary to use autograd with tensors. Autograd automatically supports Tensors with requires_grad set to True. Below please find a quick guide on what has changed:

Variable(tensor) and Variable(tensor, requires_grad) still work as expected, but they return Tensors instead of Variables.
var.data is the same thing as tensor.data.
Methods such as var.backward(), var.detach(), var.register_hook() now work on tensors with the same method names.

In addition, one can now create tensors with requires_grad=True using factory methods such as torch.randn(), torch.zeros(), torch.ones(), and others like the following:

autograd_tensor = torch.randn((2, 3, 4), requires_grad=True)

Tensor autograd functions¶

class torch.Tensor¶

backward(gradient=None, retain_graph=None, create_graph=False)¶

Computes the gradient of current tensor w.r.t. graph leaves.

The graph is differentiated using the chain rule. If the tensor is non-scalar (i.e. its data has more than one element) and requires gradient, the function additionally requires specifying gradient. It should be a tensor of matching type and location, that contains the gradient of the differentiated function w.r.t. self.

This function accumulates gradients in the leaves - you might need to zero them before calling it.

Parameters

gradient (Tensor or None) – Gradient w.r.t. the tensor. If it is a tensor, it will be automatically converted to a Tensor that does not require grad unless create_graph is True. None values can be specified for scalar Tensors or ones that don’t require grad. If a None value would be acceptable then this argument is optional.
retain_graph (bool, optional) – If False, the graph used to compute the grads will be freed. Note that in nearly all cases setting this option to True is not needed and often can be worked around in a much more efficient way. Defaults to the value of create_graph.
create_graph (bool, optional) – If True, graph of the derivative will be constructed, allowing to compute higher order derivative products. Defaults to False.

detach()¶

Returns a new Tensor, detached from the current graph.

The result will never require gradient.

Note

Returned Tensor shares the same storage with the original one. In-place modifications on either of them will be seen, and may trigger errors in correctness checks. IMPORTANT NOTE: Previously, in-place size / stride / storage changes (such as resize_ / resize_as_ / set_ / transpose_) to the returned tensor also update the original tensor. Now, these in-place changes will not update the original tensor anymore, and will instead trigger an error. For sparse tensors: In-place indices / values changes (such as zero_ / copy_ / add_) to the returned tensor will not update the original tensor anymore, and will instead trigger an error.

detach_()¶: Detaches the Tensor from the graph that created it, making it a leaf. Views cannot be detached in-place.

grad¶: This attribute is None by default and becomes a Tensor the first time a call to backward() computes gradients for self. The attribute will then contain the gradients computed and future calls to backward() will accumulate (add) gradients into it.

is_leaf¶

All Tensors that have requires_grad which is False will be leaf Tensors by convention.

For Tensors that have requires_grad which is True, they will be leaf Tensors if they were created by the user. This means that they are not the result of an operation and so grad_fn is None.

Only leaf Tensors will have their grad populated during a call to backward(). To get grad populated for non-leaf Tensors, you can use retain_grad().

Example:

>>> a = torch.rand(10, requires_grad=True)
>>> a.is_leaf
True
>>> b = torch.rand(10, requires_grad=True).cuda()
>>> b.is_leaf
False
# b was created by the operation that cast a cpu Tensor into a cuda Tensor
>>> c = torch.rand(10, requires_grad=True) + 2
>>> c.is_leaf
False
# c was created by the addition operation
>>> d = torch.rand(10).cuda()
>>> d.is_leaf
True
# d does not require gradients and so has no operation creating it (that is tracked by the autograd engine)
>>> e = torch.rand(10).cuda().requires_grad_()
>>> e.is_leaf
True
# e requires gradients and has no operations creating it
>>> f = torch.rand(10, requires_grad=True, device="cuda")
>>> f.is_leaf
True
# f requires grad, has not operation creating it

register_hook(hook)¶

Registers a backward hook.

The hook will be called every time a gradient with respect to the Tensor is computed. The hook should have the following signature:

hook(grad) -> Tensor or None

The hook should not modify its argument, but it can optionally return a new gradient which will be used in place of grad.

This function returns a handle with a method handle.remove() that removes the hook from the module.

Example:

>>> v = torch.tensor([0., 0., 0.], requires_grad=True)
>>> h = v.register_hook(lambda grad: grad * 2)  # double the gradient
>>> v.backward(torch.tensor([1., 2., 3.]))
>>> v.grad

 2
 4
 6
[torch.FloatTensor of size (3,)]

>>> h.remove()  # removes the hook

requires_grad¶: Is True if gradients need to be computed for this Tensor, False otherwise.

Note

The fact that gradients need to be computed for a Tensor do not mean that the grad attribute will be populated, see is_leaf for more details.

retain_grad()¶: Enables .grad attribute for non-leaf Tensors.

Function¶

class torch.autograd.Function¶

Records operation history and defines formulas for differentiating ops.

Every operation performed on Tensor s creates a new function object, that performs the computation, and records that it happened. The history is retained in the form of a DAG of functions, with edges denoting data dependencies (input <- output). Then, when backward is called, the graph is processed in the topological ordering, by calling backward() methods of each Function object, and passing returned gradients on to next Function s.

Normally, the only way users interact with functions is by creating subclasses and defining new operations. This is a recommended way of extending torch.autograd.

Each function object is meant to be used only once (in the forward pass).

Examples:

>>> class Exp(Function):
>>>
>>>     @staticmethod
>>>     def forward(ctx, i):
>>>         result = i.exp()
>>>         ctx.save_for_backward(result)
>>>         return result
>>>
>>>     @staticmethod
>>>     def backward(ctx, grad_output):
>>>         result, = ctx.saved_tensors
>>>         return grad_output * result

static backward(ctx, *grad_outputs)¶

Defines a formula for differentiating the operation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs did forward() return, and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computated w.r.t. the output.

static forward(ctx, *args, **kwargs)¶

Performs the operation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

The context can be used to store tensors that can be then retrieved during the backward pass.

Numerical gradient checking¶

torch.autograd.gradcheck(func, inputs, eps=1e-06, atol=1e-05, rtol=0.001, raise_exception=True, check_sparse_nnz=False)¶

Check gradients computed via small finite differences against analytical gradients w.r.t. tensors in inputs that are of floating point type and with requires_grad=True.

The check between numerical and analytical gradients uses allclose().

Note

The default values are designed for input of double precision. This check will likely fail if input is of less precision, e.g., FloatTensor.

Warning

If any checked tensor in input has overlapping memory, i.e., different indices pointing to the same memory address (e.g., from torch.expand()), this check will likely fail because the numerical gradients computed by point perturbation at such indices will change values at all other indices that share the same memory address.

Parameters

func (function) – a Python function that takes Tensor inputs and returns a Tensor or a tuple of Tensors
inputs (tuple of Tensor or Tensor) – inputs to the function
eps (float, optional) – perturbation for finite differences
atol (float, optional) – absolute tolerance
rtol (float, optional) – relative tolerance
raise_exception (bool, optional) – indicating whether to raise an exception if the check fails. The exception gives more information about the exact nature of the failure. This is helpful when debugging gradchecks.
check_sparse_nnz (bool, optional) – if True, gradcheck allows for SparseTensor input, and for any SparseTensor at input, gradcheck will perform check at nnz positions only.

Returns

True if all differences satisfy allclose condition

torch.autograd.gradgradcheck(func, inputs, grad_outputs=None, eps=1e-06, atol=1e-05, rtol=0.001, gen_non_contig_grad_outputs=False, raise_exception=True)¶

Check gradients of gradients computed via small finite differences against analytical gradients w.r.t. tensors in inputs and grad_outputs that are of floating point type and with requires_grad=True.

This function checks that backpropagating through the gradients computed to the given grad_outputs are correct.

The check between numerical and analytical gradients uses allclose().

Note

The default values are designed for input and grad_outputs of double precision. This check will likely fail if they are of less precision, e.g., FloatTensor.

Warning

If any checked tensor in input and grad_outputs has overlapping memory, i.e., different indices pointing to the same memory address (e.g., from torch.expand()), this check will likely fail because the numerical gradients computed by point perturbation at such indices will change values at all other indices that share the same memory address.

Parameters

func (function) – a Python function that takes Tensor inputs and returns a Tensor or a tuple of Tensors
inputs (tuple of Tensor or Tensor) – inputs to the function
grad_outputs (tuple of Tensor or Tensor, optional) – The gradients with respect to the function’s outputs.
eps (float, optional) – perturbation for finite differences
atol (float, optional) – absolute tolerance
rtol (float, optional) – relative tolerance
gen_non_contig_grad_outputs (bool, optional) – if grad_outputs is None and gen_non_contig_grad_outputs is True, the randomly generated gradient outputs are made to be noncontiguous
raise_exception (bool, optional) – indicating whether to raise an exception if the check fails. The exception gives more information about the exact nature of the failure. This is helpful when debugging gradchecks.

Returns

True if all differences satisfy allclose condition

Profiler¶

Autograd includes a profiler that lets you inspect the cost of different operators inside your model - both on the CPU and GPU. There are two modes implemented at the moment - CPU-only using profile. and nvprof based (registers both CPU and GPU activity) using emit_nvtx.

class torch.autograd.profiler.profile(enabled=True, use_cuda=False)¶

Context manager that manages autograd profiler state and holds a summary of results.

Parameters

enabled (bool, optional) – Setting this to False makes this context manager a no-op. Default: True.
use_cuda (bool, optional) – Enables timing of CUDA events as well using the cudaEvent API. Adds approximately 4us of overhead to each tensor operation. Default: False

Example

>>> x = torch.randn((1, 1), requires_grad=True)
>>> with torch.autograd.profiler.profile() as prof:
...     y = x ** 2
...     y.backward()
>>> # NOTE: some columns were removed for brevity
... print(prof)
-------------------------------------  ---------------  ---------------
Name                                          CPU time        CUDA time
-------------------------------------  ---------------  ---------------
PowConstant                                  142.036us          0.000us
N5torch8autograd9GraphRootE                   63.524us          0.000us
PowConstantBackward                          184.228us          0.000us
MulConstant                                   50.288us          0.000us
PowConstant                                   28.439us          0.000us
Mul                                           20.154us          0.000us
N5torch8autograd14AccumulateGradE             13.790us          0.000us
N5torch8autograd5CloneE                        4.088us          0.000us

export_chrome_trace(path)¶

Exports an EventList as a Chrome tracing tools file.

The checkpoint can be later loaded and inspected under chrome://tracing URL.

Parameters: path (str) – Path where the trace will be written.

key_averages()¶

Averages all function events over their keys.

Returns: An EventList containing FunctionEventAvg objects.

table(sort_by=None)¶

Prints an EventList as a nicely formatted table.

Parameters: sort_by (str, optional) – Attribute used to sort entries. By default they are printed in the same order as they were registered. Valid keys include: cpu_time, cuda_time, cpu_time_total, cuda_time_total, count.
Returns: A string containing the table.

total_average()¶

Averages all events.

Returns: A FunctionEventAvg object.

class torch.autograd.profiler.emit_nvtx(enabled=True)¶

Context manager that makes every autograd operation emit an NVTX range.

It is useful when running the program under nvprof:

nvprof --profile-from-start off -o trace_name.prof -- <regular command here>

Unfortunately, there’s no way to force nvprof to flush the data it collected to disk, so for CUDA profiling one has to use this context manager to annotate nvprof traces and wait for the process to exit before inspecting them. Then, either NVIDIA Visual Profiler (nvvp) can be used to visualize the timeline, or torch.autograd.profiler.load_nvprof() can load the results for inspection e.g. in Python REPL.

Parameters: enabled (bool, optional) – Setting this to False makes this context manager a no-op. Default: True.

Example

>>> with torch.cuda.profiler.profile():
...     model(x) # Warmup CUDA memory allocator and profiler
...     with torch.autograd.profiler.emit_nvtx():
...         model(x)

Forward-backward correlation

When viewing a profile created using emit_nvtx in the Nvidia Visual Profiler, correlating each backward-pass op with the corresponding forward-pass op can be difficult. To ease this task, emit_nvtx appends sequence number information to the ranges it generates.

During the forward pass, each function range is decorated with seq=<N>. seq is a running counter, incremented each time a new backward Function object is created and stashed for backward. Thus, the seq=<N> annotation associated with each forward function range tells you that if a backward Function object is created by this forward function, the backward object will receive sequence number N. During the backward pass, the top-level range wrapping each C++ backward Function’s apply() call is decorated with stashed seq=<M>. M is the sequence number that the backward object was created with. By comparing stashed seq numbers in backward with seq numbers in forward, you can track down which forward op created each backward Function.

Any functions executed during the backward pass are also decorated with seq=<N>. During default backward (with create_graph=False) this information is irrelevant, and in fact, N may simply be 0 for all such functions. Only the top-level ranges associated with backward Function objects’ apply() methods are useful, as a way to correlate these Function objects with the earlier forward pass.

Double-backward

If, on the other hand, a backward pass with create_graph=True is underway (in other words, if you are setting up for a double-backward), each function’s execution during backward is given a nonzero, useful seq=<N>. Those functions may themselves create Function objects to be executed later during double-backward, just as the original functions in the forward pass did. The relationship between backward and double-backward is conceptually the same as the relationship between forward and backward: The functions still emit current-sequence-number-tagged ranges, the Function objects they create still stash those sequence numbers, and during the eventual double-backward, the Function objects’ apply() ranges are still tagged with stashed seq numbers, which can be compared to seq numbers from the backward pass.

torch.autograd.profiler.load_nvprof(path)¶

Opens an nvprof trace file and parses autograd annotations.

Parameters: path (str) – path to nvprof trace

Anomaly detection¶

class torch.autograd.detect_anomaly¶

Context-manager that enable anomaly detection for the autograd engine.

This does two things: - Running the forward pass with detection enabled will allow the backward pass to print the traceback of the forward operation that created the failing backward function. - Any backward computation that generate “nan” value will raise an error.

Example

>>> import torch
>>> from torch import autograd
>>> class MyFunc(autograd.Function):
...     @staticmethod
...     def forward(ctx, inp):
...         return inp.clone()
...     @staticmethod
...     def backward(ctx, gO):
...         # Error during the backward pass
...         raise RuntimeError("Some error in backward")
...         return gO.clone()
>>> def run_fn(a):
...     out = MyFunc.apply(a)
...     return out.sum()
>>> inp = torch.rand(10, 10, requires_grad=True)
>>> out = run_fn(inp)
>>> out.backward()
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
      File "/your/pytorch/install/torch/tensor.py", line 93, in backward
        torch.autograd.backward(self, gradient, retain_graph, create_graph)
      File "/your/pytorch/install/torch/autograd/__init__.py", line 90, in backward
        allow_unreachable=True)  # allow_unreachable flag
      File "/your/pytorch/install/torch/autograd/function.py", line 76, in apply
        return self._forward_cls.backward(self, *args)
      File "<stdin>", line 8, in backward
    RuntimeError: Some error in backward
>>> with autograd.detect_anomaly():
...     inp = torch.rand(10, 10, requires_grad=True)
...     out = run_fn(inp)
...     out.backward()
    Traceback of forward call that caused the error:
      File "tmp.py", line 53, in <module>
        out = run_fn(inp)
      File "tmp.py", line 44, in run_fn
        out = MyFunc.apply(a)
    Traceback (most recent call last):
      File "<stdin>", line 4, in <module>
      File "/your/pytorch/install/torch/tensor.py", line 93, in backward
        torch.autograd.backward(self, gradient, retain_graph, create_graph)
      File "/your/pytorch/install/torch/autograd/__init__.py", line 90, in backward
        allow_unreachable=True)  # allow_unreachable flag
      File "/your/pytorch/install/torch/autograd/function.py", line 76, in apply
        return self._forward_cls.backward(self, *args)
      File "<stdin>", line 8, in backward
    RuntimeError: Some error in backward

class torch.autograd.set_detect_anomaly(mode)¶

Context-manager that sets the anomaly detection for the autograd engine on or off.

set_detect_anomaly will enable or disable the autograd anomaly detection based on its argument mode. It can be used as a context-manager or as a function.

See detect_anomaly above for details of the anomaly detection behaviour.

Parameters: mode (bool) – Flag whether to enable anomaly detection (True), or disable (False).

Distributed communication package - torch.distributed¶

Backends¶

torch.distributed supports three backends, each with different capabilities. The table below shows which functions are available for use with CPU / CUDA tensors. MPI supports CUDA only if the implementation used to build PyTorch supports it.

Backend	`gloo`		`mpi`		`nccl`
Device	CPU	GPU	CPU	GPU	CPU	GPU
send	✓	✘	✓	?	✘	✘
recv	✓	✘	✓	?	✘	✘
broadcast	✓	✓	✓	?	✘	✓
all_reduce	✓	✓	✓	?	✘	✓
reduce	✓	✘	✓	?	✘	✓
all_gather	✓	✘	✓	?	✘	✓
gather	✓	✘	✓	?	✘	✘
scatter	✓	✘	✓	?	✘	✘
barrier	✓	✘	✓	?	✘	✓

Backends that come with PyTorch¶

PyTorch distributed currently only supports Linux. By default, the Gloo and NCCL backends are built and included in PyTorch distributed (NCCL only when building with CUDA). MPI is an optional backend that can only be included if you build PyTorch from source. (e.g. building PyTorch on a host that has MPI installed.)

Which backend to use?¶

In the past, we were often asked: “which backend should I use?”.

Rule of thumb
- Use the NCCL backend for distributed GPU training
- Use the Gloo backend for distributed CPU training.
GPU hosts with InfiniBand interconnect
- Use NCCL, since it’s the only backend that currently supports InfiniBand and GPUDirect.
GPU hosts with Ethernet interconnect
- Use NCCL, since it currently provides the best distributed GPU training performance, especially for multiprocess single-node or multi-node distributed training. If you encounter any problem with NCCL, use Gloo as the fallback option. (Note that Gloo currently runs slower than NCCL for GPUs.)
CPU hosts with InfiniBand interconnect
- If your InfiniBand has enabled IP over IB, use Gloo, otherwise, use MPI instead. We are planning on adding InfiniBand support for Gloo in the upcoming releases.
CPU hosts with Ethernet interconnect
- Use Gloo, unless you have specific reasons to use MPI.

Common environment variables¶

Choosing the network interface to use¶

By default, both NCCL and Gloo backends will try to find the network interface to use for communication. However, this is not always guaranteed to be successful from our experiences. Therefore, if you encounter any problem on either backend not being able to find the correct network interface. You can try to set the following environment variables (each one applicable to its respective backend):

NCCL_SOCKET_IFNAME, for example export NCCL_SOCKET_IFNAME=eth0
GLOO_SOCKET_IFNAME, for example export GLOO_SOCKET_IFNAME=eth0

Other NCCL environment variables¶

NCCL has also provided a number of environment variables for fine-tuning purposes.

Commonly used ones include the following for debugging purposes:

export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL

For the full list of NCCL environment variables, please refer to NVIDIA NCCL’s official documentation

Basics¶

The torch.distributed package provides PyTorch support and communication primitives for multiprocess parallelism across several computation nodes running on one or more machines. The class torch.nn.parallel.DistributedDataParallel() builds on this functionality to provide synchronous distributed training as a wrapper around any PyTorch model. This differs from the kinds of parallelism provided by Multiprocessing package - torch.multiprocessing and torch.nn.DataParallel() in that it supports multiple network-connected machines and in that the user must explicitly launch a separate copy of the main training script for each process.

In the single-machine synchronous case, torch.distributed or the torch.nn.parallel.DistributedDataParallel() wrapper may still have advantages over other approaches to data-parallelism, including torch.nn.DataParallel():

Each process maintains its own optimizer and performs a complete optimization step with each iteration. While this may appear redundant, since the gradients have already been gathered together and averaged across processes and are thus the same for every process, this means that no parameter broadcast step is needed, reducing time spent transferring tensors between nodes.
Each process contains an independent Python interpreter, eliminating the extra interpreter overhead and “GIL-thrashing” that comes from driving several execution threads, model replicas, or GPUs from a single Python process. This is especially important for models that make heavy use of the Python runtime, including models with recurrent layers or many small components.

Initialization¶

The package needs to be initialized using the torch.distributed.init_process_group() function before calling any other methods. This blocks until all processes have joined.

Currently three initialization methods are supported:

TCP initialization¶

There are two ways to initialize using TCP, both requiring a network address reachable from all processes and a desired world_size. The first way requires specifying an address that belongs to the rank 0 process. This initialization method requires that all processes have manually specified ranks.

Note that multicast address is not supported anymore in the latest distributed package. group_name is deprecated as well.

import torch.distributed as dist

# Use address of one of the machines
dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456',
                        rank=args.rank, world_size=4)

Shared file-system initialization¶

Another initialization method makes use of a file system that is shared and visible from all machines in a group, along with a desired world_size. The URL should start with file:// and contain a path to a non-existent file (in an existing directory) on a shared file system. File-system initialization will automatically create that file if it doesn’t exist, but will not delete the file. Therefore, it is your responsibility to make sure that the file is cleaned up before the next init_process_group() call on the same file path/name.

Note that automatic rank assignment is not supported anymore in the latest distributed package and group_name is deprecated as well.

Warning

This method assumes that the file system supports locking using fcntl - most local systems and NFS support it.

Warning

This method will always create the file and try its best to clean up and remove the file at the end of the program. In other words, each initialization with the file init method will need a brand new empty file in order for the initialization to succeed. If the same file used by the previous initialization (which happens not to get cleaned up) is used again, this is unexpected behavior and can often cause deadlocks and failures. Therefore, even though this method will try its best to clean up the file, if the auto-delete happens to be unsuccessful, it is your responsibility to ensure that the file is removed at the end of the training to prevent the same file to be reused again during the next time. This is especially important if you plan to call init_process_group() multiple times on the same file name. In other words, if the file is not removed/cleaned up and you call init_process_group() again on that file, failures are expected. The rule of thumb here is that, make sure that the file is non-existent or empty everytime init_process_group() is called.

import torch.distributed as dist

# rank should always be specified
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',
                        world_size=4, rank=args.rank)

Environment variable initialization¶

This method will read the configuration from environment variables, allowing one to fully customize how the information is obtained. The variables to be set are:

MASTER_PORT - required; has to be a free port on machine with rank 0
MASTER_ADDR - required (except for rank 0); address of rank 0 node
WORLD_SIZE - required; can be set either here, or in a call to init function
RANK - required; can be set either here, or in a call to init function

The machine with rank 0 will be used to set up all connections.

This is the default method, meaning that init_method does not have to be specified (or can be env://).

Groups¶

By default collectives operate on the default group (also called the world) and require all processes to enter the distributed function call. However, some workloads can benefit from more fine-grained communication. This is where distributed groups come into play. new_group() function can be used to create new groups, with arbitrary subsets of all processes. It returns an opaque group handle that can be given as a group argument to all collectives (collectives are distributed functions to exchange information in certain well-known programming patterns).

Currently torch.distributed does not support creating groups with different backends. In other words, each group being created will use the same backend as you specified in init_process_group().

Point-to-point communication¶

isend() and irecv() return distributed request objects when used. In general, the type of this object is unspecified as they should never be created manually, but they are guaranteed to support two methods:

is_completed() - returns True if the operation has finished
wait() - will block the process until the operation is finished. is_completed() is guaranteed to return True once it returns.

Synchronous and asynchronous collective operations¶

Every collective operation function supports the following two kinds of operations:

synchronous operation - the default mode, when async_op is set to False. when the function returns, it is guaranteed that the collective operation is performed (not necessarily completed if it’s a CUDA op since all CUDA ops are asynchronous), and any further function calls depending on the data of the collective operation can be called. In the synchronous mode, the collective function does not return anything

asynchronous operation - when async_op is set to True. The collective operation function returns a distributed request object. In general, you don’t need to create it manually and it is guaranteed to support two methods:

is_completed() - returns True if the operation has finished
wait() - will block the process until the operation is finished.

Collective functions¶

class torch.distributed.reduce_op¶

Deprecated enum-like class for reduction operations: SUM, PRODUCT, MIN, and MAX.

ReduceOp is recommended to use instead.

Multi-GPU collective functions¶

If you have more than one GPU on each node, when using the NCCL and Gloo backend, broadcast_multigpu() all_reduce_multigpu() reduce_multigpu() and all_gather_multigpu() support distributed collective operations among multiple GPUs within each node. These functions can potentially improve the overall distributed training performance and be easily used by passing a list of tensors. Each Tensor in the passed tensor list needs to be on a separate GPU device of the host where the function is called. Note that the length of the tensor list needs to be identical among all the distributed processes. Also note that currently the multi-GPU collective functions are only supported by the NCCL backend.

For example, if the system we use for distributed training has 2 nodes, each of which has 8 GPUs. On each of the 16 GPUs, there is a tensor that we would like to all-reduce. The following code can serve as a reference:

Code running on Node 0

import torch
import torch.distributed as dist

dist.init_process_group(backend="nccl",
                        init_method="file:///distributed_test",
                        world_size=2,
                        rank=0)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
    tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))

dist.all_reduce_multigpu(tensor_list)

Code running on Node 1

import torch
import torch.distributed as dist

dist.init_process_group(backend="nccl",
                        init_method="file:///distributed_test",
                        world_size=2,
                        rank=1)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
    tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))

dist.all_reduce_multigpu(tensor_list)

After the call, all 16 tensors on the two nodes will have the all-reduced value of 16

Launch utility¶

The torch.distributed package also provides a launch utility in torch.distributed.launch. This helper utility can be used to launch multiple processes per node for distributed training. This utility also supports both python2 and python3.

torch.distributed.launch is a module that spawns up multiple distributed training processes on each of the training nodes.

The utility can be used for single-node distributed training, in which one or more processes per node will be spawned. The utility can be used for either CPU training or GPU training. If the utility is used for GPU training, each distributed process will be operating on a single GPU. This can achieve well-improved single-node training performance. It can also be used in multi-node distributed training, by spawning up multiple processes on each node for well-improved multi-node distributed training performance as well. This will especially be benefitial for systems with multiple Infiniband interfaces that have direct-GPU support, since all of them can be utilized for aggregated communication bandwidth.

In both cases of single-node distributed training or multi-node distributed training, this utility will launch the given number of processes per node (--nproc_per_node). If used for GPU training, this number needs to be less or euqal to the number of GPUs on the current system (nproc_per_node), and each process will be operating on a single GPU from GPU 0 to GPU (nproc_per_node - 1).

How to use this module:

Single-Node multi-process distributed training

>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
           YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other
           arguments of your training script)

Multi-Node multi-process distributed training: (e.g. two nodes)

Node 1: (IP: 192.168.1.1, and has a free port: 1234)

>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
           --nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
           --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
           and all other arguments of your training script)

Node 2:

>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
           --nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
           --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
           and all other arguments of your training script)

To look up what optional arguments this module offers:

>>> python -m torch.distributed.launch --help

Important Notices:

1. This utilty and multi-process distributed (single-node or multi-node) GPU training currently only achieves the best performance using the NCCL distributed backend. Thus NCCL backend is the recommended backend to use for GPU training.

2. In your training program, you must parse the command-line argument: --local_rank=LOCAL_PROCESS_RANK, which will be provided by this module. If your training program uses GPUs, you should ensure that your code only runs on the GPU device of LOCAL_PROCESS_RANK. This can be done by:

Parsing the local_rank argument

>>> import argparse
>>> parser = argparse.ArgumentParser()
>>> parser.add_argument("--local_rank", type=int)
>>> args = parser.parse_args()

Set your device to local rank using either

>>> torch.cuda.set_device(arg.local_rank)  # before your code runs

>>> with torch.cuda.device(arg.local_rank):
>>>    # your code to run

3. In your training program, you are supposed to call the following function at the beginning to start the distributed backend. You need to make sure that the init_method uses env://, which is the only supported init_method by this module.

torch.distributed.init_process_group(backend='YOUR BACKEND',
                                     init_method='env://')

4. In your training program, you can either use regular distributed functions or use torch.nn.parallel.DistributedDataParallel() module. If your training program uses GPUs for training and you would like to use torch.nn.parallel.DistributedDataParallel() module, here is how to configure it.

model = torch.nn.parallel.DistributedDataParallel(model,
                                                  device_ids=[arg.local_rank],
                                                  output_device=arg.local_rank)

Please ensure that device_ids argument is set to be the only GPU device id that your code will be operating on. This is generally the local rank of the process. In other words, the device_ids needs to be [args.local_rank], and output_device needs to be args.local_rank in order to use this utility

5. Another way to pass local_rank to the subprocesses via environment variable LOCAL_RANK. This behavior is enabled when you launch the script with --use_env=True. You must adjust the subprocess example above to replace args.local_rank with os.environ['LOCAL_RANK']; the launcher will not pass --local_rank when you specify this flag.

Warning

local_rank is NOT globally unique: it is only unique per process on a machine. Thus, don’t use it to decide if you should, e.g., write to a networked filesystem. See https://github.com/pytorch/pytorch/issues/12042 for an example of how things can go wrong if you don’t do this correctly.

Spawn utility¶

The torch.multiprocessing package also provides a spawn function in torch.multiprocessing.spawn(). This helper function can be used to spawn multiple processes. It works by passing in the function that you want to run and spawns N processes to run it. This can be used for multiprocess distributed training as well.

For references on how to use it, please refer to PyToch example - ImageNet implementation

Note that this function requires Python 3.4 or higher.

Probability distributions - torch.distributions¶

The distributions package contains parameterizable probability distributions and sampling functions. This allows the construction of stochastic computation graphs and stochastic gradient estimators for optimization. This package generally follows the design of the TensorFlow Distributions package.

It is not possible to directly backpropagate through random samples. However, there are two main methods for creating surrogate functions that can be backpropagated through. These are the score function estimator/likelihood ratio estimator/REINFORCE and the pathwise derivative estimator. REINFORCE is commonly seen as the basis for policy gradient methods in reinforcement learning, and the pathwise derivative estimator is commonly seen in the reparameterization trick in variational autoencoders. Whilst the score function only requires the value of samples $f(x)$, the pathwise derivative requires the derivative $f'(x)$. The next sections discuss these two in a reinforcement learning example. For more details see Gradient Estimation Using Stochastic Computation Graphs .

Score function¶

When the probability density function is differentiable with respect to its parameters, we only need sample() and log_prob() to implement REINFORCE:

\[\Delta\theta = \alpha r \frac{\partial\log p(a|\pi^\theta(s))}{\partial\theta}\]

where $\theta$ are the parameters, $\alpha$ is the learning rate, $r$ is the reward and $p(a|\pi^\theta(s))$ is the probability of taking action $a$ in state $s$ given policy $\pi^\theta$.

In practice we would sample an action from the output of a network, apply this action in an environment, and then use log_prob to construct an equivalent loss function. Note that we use a negative because optimizers use gradient descent, whilst the rule above assumes gradient ascent. With a categorical policy, the code for implementing REINFORCE would be as follows:

probs = policy_network(state)
# Note that this is equivalent to what used to be called multinomial
m = Categorical(probs)
action = m.sample()
next_state, reward = env.step(action)
loss = -m.log_prob(action) * reward
loss.backward()

Pathwise derivative¶

The other way to implement these stochastic/policy gradients would be to use the reparameterization trick from the rsample() method, where the parameterized random variable can be constructed via a parameterized deterministic function of a parameter-free random variable. The reparameterized sample therefore becomes differentiable. The code for implementing the pathwise derivative would be as follows:

params = policy_network(state)
m = Normal(*params)
# Any distribution with .has_rsample == True could work based on the application
action = m.rsample()
next_state, reward = env.step(action)  # Assuming that reward is differentiable
loss = -reward
loss.backward()

Distribution¶

class torch.distributions.distribution.Distribution(batch_shape=torch.Size([]), event_shape=torch.Size([]), validate_args=None)¶

Bases: object

Distribution is the abstract base class for probability distributions.

arg_constraints¶: Returns a dictionary from argument names to Constraint objects that should be satisfied by each argument of this distribution. Args that are not tensors need not appear in this dict.

batch_shape¶: Returns the shape over which parameters are batched.

cdf(value)¶

Returns the cumulative density/mass function evaluated at value.

Parameters: value (Tensor) –

entropy()¶

Returns entropy of distribution, batched over batch_shape.

Returns: Tensor of shape batch_shape.

enumerate_support(expand=True)¶

Returns tensor containing all values supported by a discrete distribution. The result will enumerate over dimension 0, so the shape of the result will be (cardinality,) + batch_shape + event_shape (where event_shape = () for univariate distributions).

Note that this enumerates over all batched tensors in lock-step [[0, 0], [1, 1], …]. With expand=False, enumeration happens along dim 0, but with the remaining batch dimensions being singleton dimensions, [[0], [1], ...

To iterate over the full Cartesian product use itertools.product(m.enumerate_support()).

Parameters: expand (bool) – whether to expand the support over the batch dims to match the distribution’s batch_shape.
Returns: Tensor iterating over dimension 0.

event_shape¶: Returns the shape of a single sample (without batching).

expand(batch_shape, _instance=None)¶

Returns a new distribution instance (or populates an existing instance provided by a derived class) with batch dimensions expanded to batch_shape. This method calls expand on the distribution’s parameters. As such, this does not allocate new memory for the expanded distribution instance. Additionally, this does not repeat any args checking or parameter broadcasting in __init__.py, when an instance is first created.

Parameters

batch_shape (torch.Size) – the desired expanded size.
_instance – new instance provided by subclasses that need to override .expand.

Returns

New distribution instance with batch dimensions expanded to batch_size.

icdf(value)¶

Returns the inverse cumulative density/mass function evaluated at value.

Parameters: value (Tensor) –

log_prob(value)¶

Returns the log of the probability density/mass function evaluated at value.

Parameters: value (Tensor) –

mean¶: Returns the mean of the distribution.

perplexity()¶

Returns perplexity of distribution, batched over batch_shape.

Returns: Tensor of shape batch_shape.

rsample(sample_shape=torch.Size([]))¶: Generates a sample_shape shaped reparameterized sample or sample_shape shaped batch of reparameterized samples if the distribution parameters are batched.

sample(sample_shape=torch.Size([]))¶: Generates a sample_shape shaped sample or sample_shape shaped batch of samples if the distribution parameters are batched.

sample_n(n)¶: Generates n samples or n batches of samples if the distribution parameters are batched.

stddev¶: Returns the standard deviation of the distribution.

support¶: Returns a Constraint object representing this distribution’s support.

variance¶: Returns the variance of the distribution.

ExponentialFamily¶

class torch.distributions.exp_family.ExponentialFamily(batch_shape=torch.Size([]), event_shape=torch.Size([]), validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

ExponentialFamily is the abstract base class for probability distributions belonging to an exponential family, whose probability mass/density function has the form is defined below

\[p_{F}(x; \theta) = \exp(\langle t(x), \theta\rangle - F(\theta) + k(x))\]

where $\theta$ denotes the natural parameters, $t(x)$ denotes the sufficient statistic, $F(\theta)$ is the log normalizer function for a given family and $k(x)$ is the carrier measure.

Note

This class is an intermediary between the Distribution class and distributions which belong to an exponential family mainly to check the correctness of the .entropy() and analytic KL divergence methods. We use this class to compute the entropy and KL divergence using the AD framework and Bregman divergences (courtesy of: Frank Nielsen and Richard Nock, Entropies and Cross-entropies of Exponential Families).

entropy()¶: Method to compute the entropy using Bregman divergence of the log normalizer.

Bernoulli¶

class torch.distributions.bernoulli.Bernoulli(probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.exp_family.ExponentialFamily

Creates a Bernoulli distribution parameterized by probs or logits (but not both).

Samples are binary (0 or 1). They take the value 1 with probability p and 0 with probability 1 - p.

Example:

>>> m = Bernoulli(torch.tensor([0.3]))
>>> m.sample()  # 30% chance 1; 70% chance 0
tensor([ 0.])

Parameters

probs (Number, Tensor) – the probability of sampling 1
logits (Number, Tensor) – the log-odds of sampling 1

arg_constraints = {'logits': Real(), 'probs': Interval(lower_bound=0.0, upper_bound=1.0)}¶

entropy()¶

enumerate_support(expand=True)¶

expand(batch_shape, _instance=None)¶

has_enumerate_support = True¶

log_prob(value)¶

logits¶

mean¶

param_shape¶

probs¶

sample(sample_shape=torch.Size([]))¶

support = Boolean()¶

variance¶

Beta¶

class torch.distributions.beta.Beta(concentration1, concentration0, validate_args=None)¶

Bases: torch.distributions.exp_family.ExponentialFamily

Beta distribution parameterized by concentration1 and concentration0.

Example:

>>> m = Beta(torch.tensor([0.5]), torch.tensor([0.5]))
>>> m.sample()  # Beta distributed with concentration concentration1 and concentration0
tensor([ 0.1046])

Parameters

concentration1 (float or Tensor) – 1st concentration parameter of the distribution (often referred to as alpha)
concentration0 (float or Tensor) – 2nd concentration parameter of the distribution (often referred to as beta)

arg_constraints = {'concentration0': GreaterThan(lower_bound=0.0), 'concentration1': GreaterThan(lower_bound=0.0)}¶

concentration0¶

concentration1¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

log_prob(value)¶

mean¶

rsample(sample_shape=())¶

support = Interval(lower_bound=0.0, upper_bound=1.0)¶

variance¶

Binomial¶

class torch.distributions.binomial.Binomial(total_count=1, probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a Binomial distribution parameterized by total_count and either probs or logits (but not both). total_count must be broadcastable with probs/logits.

Example:

>>> m = Binomial(100, torch.tensor([0 , .2, .8, 1]))
>>> x = m.sample()
tensor([   0.,   22.,   71.,  100.])

>>> m = Binomial(torch.tensor([[5.], [10.]]), torch.tensor([0.5, 0.8]))
>>> x = m.sample()
tensor([[ 4.,  5.],
        [ 7.,  6.]])

Parameters

total_count (int or Tensor) – number of Bernoulli trials
probs (Tensor) – Event probabilities
logits (Tensor) – Event log-odds

arg_constraints = {'logits': Real(), 'probs': Interval(lower_bound=0.0, upper_bound=1.0), 'total_count': IntegerGreaterThan(lower_bound=0)}¶

enumerate_support(expand=True)¶

expand(batch_shape, _instance=None)¶

has_enumerate_support = True¶

log_prob(value)¶

logits¶

mean¶

param_shape¶

probs¶

sample(sample_shape=torch.Size([]))¶

support¶

variance¶

Categorical¶

class torch.distributions.categorical.Categorical(probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a categorical distribution parameterized by either probs or logits (but not both).

Note

It is equivalent to the distribution that torch.multinomial() samples from.

Samples are integers from $\{0, \ldots, K-1\}$ where K is probs.size(-1).

If probs is 1D with length-K, each element is the relative probability of sampling the class at that index.

If probs is 2D, it is treated as a batch of relative probability vectors.

Note

probs must be non-negative, finite and have a non-zero sum, and it will be normalized to sum to 1.

Cauchy¶

class torch.distributions.cauchy.Cauchy(loc, scale, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Samples from a Cauchy (Lorentz) distribution. The distribution of the ratio of independent normally distributed random variables with means 0 follows a Cauchy distribution.

Example:

>>> m = Cauchy(torch.tensor([0.0]), torch.tensor([1.0]))
>>> m.sample()  # sample from a Cauchy distribution with loc=0 and scale=1
tensor([ 2.3214])

Parameters

loc (float or Tensor) – mode or median of the distribution.
scale (float or Tensor) – half width at half maximum.

arg_constraints = {'loc': Real(), 'scale': GreaterThan(lower_bound=0.0)}¶

cdf(value)¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

icdf(value)¶

log_prob(value)¶

mean¶

rsample(sample_shape=torch.Size([]))¶

support = Real()¶

variance¶

Chi2¶

class torch.distributions.chi2.Chi2(df, validate_args=None)¶

Bases: torch.distributions.gamma.Gamma

Creates a Chi2 distribution parameterized by shape parameter df. This is exactly equivalent to Gamma(alpha=0.5*df, beta=0.5)

Example:

>>> m = Chi2(torch.tensor([1.0]))
>>> m.sample()  # Chi2 distributed with shape df=1
tensor([ 0.1046])

Parameters: df (float or Tensor) – shape parameter of the distribution

arg_constraints = {'df': GreaterThan(lower_bound=0.0)}¶

df¶

expand(batch_shape, _instance=None)¶

Dirichlet¶

class torch.distributions.dirichlet.Dirichlet(concentration, validate_args=None)¶

Bases: torch.distributions.exp_family.ExponentialFamily

Creates a Dirichlet distribution parameterized by concentration concentration.

Example:

>>> m = Dirichlet(torch.tensor([0.5, 0.5]))
>>> m.sample()  # Dirichlet distributed with concentrarion concentration
tensor([ 0.1046,  0.8954])

Parameters: concentration (Tensor) – concentration parameter of the distribution (often referred to as alpha)

arg_constraints = {'concentration': GreaterThan(lower_bound=0.0)}¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

log_prob(value)¶

mean¶

rsample(sample_shape=())¶

support = Simplex()¶

variance¶

Exponential¶

class torch.distributions.exponential.Exponential(rate, validate_args=None)¶

Bases: torch.distributions.exp_family.ExponentialFamily

Creates a Exponential distribution parameterized by rate.

Example:

>>> m = Exponential(torch.tensor([1.0]))
>>> m.sample()  # Exponential distributed with rate=1
tensor([ 0.1046])

Parameters: rate (float or Tensor) – rate = 1 / scale of the distribution

arg_constraints = {'rate': GreaterThan(lower_bound=0.0)}¶

cdf(value)¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

icdf(value)¶

log_prob(value)¶

mean¶

rsample(sample_shape=torch.Size([]))¶

stddev¶

support = GreaterThan(lower_bound=0.0)¶

variance¶

FisherSnedecor¶

class torch.distributions.fishersnedecor.FisherSnedecor(df1, df2, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a Fisher-Snedecor distribution parameterized by df1 and df2.

Example:

>>> m = FisherSnedecor(torch.tensor([1.0]), torch.tensor([2.0]))
>>> m.sample()  # Fisher-Snedecor-distributed with df1=1 and df2=2
tensor([ 0.2453])

Parameters

df1 (float or Tensor) – degrees of freedom parameter 1
df2 (float or Tensor) – degrees of freedom parameter 2

arg_constraints = {'df1': GreaterThan(lower_bound=0.0), 'df2': GreaterThan(lower_bound=0.0)}¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

log_prob(value)¶

mean¶

rsample(sample_shape=torch.Size([]))¶

support = GreaterThan(lower_bound=0.0)¶

variance¶

Gamma¶

class torch.distributions.gamma.Gamma(concentration, rate, validate_args=None)¶

Bases: torch.distributions.exp_family.ExponentialFamily

Creates a Gamma distribution parameterized by shape concentration and rate.

Example:

>>> m = Gamma(torch.tensor([1.0]), torch.tensor([1.0]))
>>> m.sample()  # Gamma distributed with concentration=1 and rate=1
tensor([ 0.1046])

Parameters

concentration (float or Tensor) – shape parameter of the distribution (often referred to as alpha)
rate (float or Tensor) – rate = 1 / scale of the distribution (often referred to as beta)

arg_constraints = {'concentration': GreaterThan(lower_bound=0.0), 'rate': GreaterThan(lower_bound=0.0)}¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

log_prob(value)¶

mean¶

rsample(sample_shape=torch.Size([]))¶

support = GreaterThan(lower_bound=0.0)¶

variance¶

Geometric¶

class torch.distributions.geometric.Geometric(probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a Geometric distribution parameterized by probs, where probs is the probability of success of Bernoulli trials. It represents the probability that in $k + 1$ Bernoulli trials, the first $k$ trials failed, before seeing a success.

Samples are non-negative integers [0, $\inf$).

Example:

>>> m = Geometric(torch.tensor([0.3]))
>>> m.sample()  # underlying Bernoulli has 30% chance 1; 70% chance 0
tensor([ 2.])

Parameters

probs (Number, Tensor) – the probability of sampling 1. Must be in range (0, 1]
logits (Number, Tensor) – the log-odds of sampling 1.

arg_constraints = {'logits': Real(), 'probs': Interval(lower_bound=0.0, upper_bound=1.0)}¶

entropy()¶

expand(batch_shape, _instance=None)¶

log_prob(value)¶

logits¶

mean¶

probs¶

sample(sample_shape=torch.Size([]))¶

support = IntegerGreaterThan(lower_bound=0)¶

variance¶

Gumbel¶

class torch.distributions.gumbel.Gumbel(loc, scale, validate_args=None)¶

Bases: torch.distributions.transformed_distribution.TransformedDistribution

Samples from a Gumbel Distribution.

Examples:

>>> m = Gumbel(torch.tensor([1.0]), torch.tensor([2.0]))
>>> m.sample()  # sample from Gumbel distribution with loc=1, scale=2
tensor([ 1.0124])

Parameters

loc (float or Tensor) – Location parameter of the distribution
scale (float or Tensor) – Scale parameter of the distribution

arg_constraints = {'loc': Real(), 'scale': GreaterThan(lower_bound=0.0)}¶

entropy()¶

expand(batch_shape, _instance=None)¶

log_prob(value)¶

mean¶

stddev¶

support = Real()¶

variance¶

HalfCauchy¶

class torch.distributions.half_cauchy.HalfCauchy(scale, validate_args=None)¶

Bases: torch.distributions.transformed_distribution.TransformedDistribution

Creates a half-normal distribution parameterized by scale where:

X ~ Cauchy(0, scale)
Y = |X| ~ HalfCauchy(scale)

Example:

>>> m = HalfCauchy(torch.tensor([1.0]))
>>> m.sample()  # half-cauchy distributed with scale=1
tensor([ 2.3214])

Parameters: scale (float or Tensor) – scale of the full Cauchy distribution

arg_constraints = {'scale': GreaterThan(lower_bound=0.0)}¶

cdf(value)¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

icdf(prob)¶

log_prob(value)¶

mean¶

scale¶

support = GreaterThan(lower_bound=0.0)¶

variance¶

HalfNormal¶

class torch.distributions.half_normal.HalfNormal(scale, validate_args=None)¶

Bases: torch.distributions.transformed_distribution.TransformedDistribution

Creates a half-normal distribution parameterized by scale where:

X ~ Normal(0, scale)
Y = |X| ~ HalfNormal(scale)

Example:

>>> m = HalfNormal(torch.tensor([1.0]))
>>> m.sample()  # half-normal distributed with scale=1
tensor([ 0.1046])

Parameters: scale (float or Tensor) – scale of the full Normal distribution

arg_constraints = {'scale': GreaterThan(lower_bound=0.0)}¶

cdf(value)¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

icdf(prob)¶

log_prob(value)¶

mean¶

scale¶

support = GreaterThan(lower_bound=0.0)¶

variance¶

Independent¶

class torch.distributions.independent.Independent(base_distribution, reinterpreted_batch_ndims, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Reinterprets some of the batch dims of a distribution as event dims.

This is mainly useful for changing the shape of the result of log_prob(). For example to create a diagonal Normal distribution with the same shape as a Multivariate Normal distribution (so they are interchangeable), you can:

>>> loc = torch.zeros(3)
>>> scale = torch.ones(3)
>>> mvn = MultivariateNormal(loc, scale_tril=torch.diag(scale))
>>> [mvn.batch_shape, mvn.event_shape]
[torch.Size(()), torch.Size((3,))]
>>> normal = Normal(loc, scale)
>>> [normal.batch_shape, normal.event_shape]
[torch.Size((3,)), torch.Size(())]
>>> diagn = Independent(normal, 1)
>>> [diagn.batch_shape, diagn.event_shape]
[torch.Size(()), torch.Size((3,))]

Parameters

base_distribution (torch.distributions.distribution.Distribution) – a base distribution
reinterpreted_batch_ndims (int) – the number of batch dims to reinterpret as event dims

arg_constraints = {}¶

entropy()¶

enumerate_support(expand=True)¶

expand(batch_shape, _instance=None)¶

has_enumerate_support¶

has_rsample¶

log_prob(value)¶

mean¶

rsample(sample_shape=torch.Size([]))¶

sample(sample_shape=torch.Size([]))¶

support¶

variance¶

Laplace¶

class torch.distributions.laplace.Laplace(loc, scale, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a Laplace distribution parameterized by loc and :attr:’scale’.

Example:

>>> m = Laplace(torch.tensor([0.0]), torch.tensor([1.0]))
>>> m.sample()  # Laplace distributed with loc=0, scale=1
tensor([ 0.1046])

Parameters

loc (float or Tensor) – mean of the distribution
scale (float or Tensor) – scale of the distribution

arg_constraints = {'loc': Real(), 'scale': GreaterThan(lower_bound=0.0)}¶

cdf(value)¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

icdf(value)¶

log_prob(value)¶

mean¶

rsample(sample_shape=torch.Size([]))¶

stddev¶

support = Real()¶

variance¶

LogNormal¶

class torch.distributions.log_normal.LogNormal(loc, scale, validate_args=None)¶

Bases: torch.distributions.transformed_distribution.TransformedDistribution

Creates a log-normal distribution parameterized by loc and scale where:

X ~ Normal(loc, scale)
Y = exp(X) ~ LogNormal(loc, scale)

Example:

>>> m = LogNormal(torch.tensor([0.0]), torch.tensor([1.0]))
>>> m.sample()  # log-normal distributed with mean=0 and stddev=1
tensor([ 0.1046])

Parameters

loc (float or Tensor) – mean of log of distribution
scale (float or Tensor) – standard deviation of log of the distribution

arg_constraints = {'loc': Real(), 'scale': GreaterThan(lower_bound=0.0)}¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

loc¶

mean¶

scale¶

support = GreaterThan(lower_bound=0.0)¶

variance¶

LowRankMultivariateNormal¶

class torch.distributions.lowrank_multivariate_normal.LowRankMultivariateNormal(loc, cov_factor, cov_diag, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a multivariate normal distribution with covariance matrix having a low-rank form parameterized by cov_factor and cov_diag:

covariance_matrix = cov_factor @ cov_factor.T + cov_diag

Example

>>> m = LowRankMultivariateNormal(torch.zeros(2), torch.tensor([1, 0]), torch.tensor([1, 1]))
>>> m.sample()  # normally distributed with mean=`[0,0]`, cov_factor=`[1,0]`, cov_diag=`[1,1]`
tensor([-0.2102, -0.5429])

Parameters

loc (Tensor) – mean of the distribution with shape batch_shape + event_shape
cov_factor (Tensor) – factor part of low-rank form of covariance matrix with shape batch_shape + event_shape + (rank,)
cov_diag (Tensor) – diagonal part of low-rank form of covariance matrix with shape batch_shape + event_shape

Note

The computation for determinant and inverse of covariance matrix is avoided when cov_factor.shape[1] << cov_factor.shape[0] thanks to Woodbury matrix identity and matrix determinant lemma. Thanks to these formulas, we just need to compute the determinant and inverse of the small size “capacitance” matrix:

capacitance = I + cov_factor.T @ inv(cov_diag) @ cov_factor

arg_constraints = {'cov_diag': GreaterThan(lower_bound=0.0), 'cov_factor': Real(), 'loc': Real()}¶

covariance_matrix¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

log_prob(value)¶

mean¶

precision_matrix¶

rsample(sample_shape=torch.Size([]))¶

scale_tril¶

support = Real()¶

variance¶

Multinomial¶

class torch.distributions.multinomial.Multinomial(total_count=1, probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a Multinomial distribution parameterized by total_count and either probs or logits (but not both). The innermost dimension of probs indexes over categories. All other dimensions index over batches.

Note that total_count need not be specified if only log_prob() is called (see example below)

Note

probs must be non-negative, finite and have a non-zero sum, and it will be normalized to sum to 1.

sample() requires a single shared total_count for all parameters and samples.
log_prob() allows different total_count for each parameter and sample.

Example:

>>> m = Multinomial(100, torch.tensor([ 1., 1., 1., 1.]))
>>> x = m.sample()  # equal probability of 0, 1, 2, 3
tensor([ 21.,  24.,  30.,  25.])

>>> Multinomial(probs=torch.tensor([1., 1., 1., 1.])).log_prob(x)
tensor([-4.1338])

Parameters

total_count (int) – number of trials
probs (Tensor) – event probabilities
logits (Tensor) – event log probabilities

arg_constraints = {'logits': Real(), 'probs': Simplex()}¶

expand(batch_shape, _instance=None)¶

log_prob(value)¶

logits¶

mean¶

param_shape¶

probs¶

sample(sample_shape=torch.Size([]))¶

support¶

variance¶

MultivariateNormal¶

class torch.distributions.multivariate_normal.MultivariateNormal(loc, covariance_matrix=None, precision_matrix=None, scale_tril=None, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a multivariate normal (also called Gaussian) distribution parameterized by a mean vector and a covariance matrix.

The multivariate normal distribution can be parameterized either in terms of a positive definite covariance matrix $\mathbf{\Sigma}$ or a positive definite precision matrix $\mathbf{\Sigma}^{-1}$ or a lower-triangular matrix $\mathbf{L}$ with positive-valued diagonal entries, such that $\mathbf{\Sigma} = \mathbf{L}\mathbf{L}^\top$. This triangular matrix can be obtained via e.g. Cholesky decomposition of the covariance.

Example

>>> m = MultivariateNormal(torch.zeros(2), torch.eye(2))
>>> m.sample()  # normally distributed with mean=`[0,0]` and covariance_matrix=`I`
tensor([-0.2102, -0.5429])

Parameters

loc (Tensor) – mean of the distribution
covariance_matrix (Tensor) – positive-definite covariance matrix
precision_matrix (Tensor) – positive-definite precision matrix
scale_tril (Tensor) – lower-triangular factor of covariance, with positive-valued diagonal

Note

Only one of covariance_matrix or precision_matrix or scale_tril can be specified.

Using scale_tril will be more efficient: all computations internally are based on scale_tril. If covariance_matrix or precision_matrix is passed instead, it is only used to compute the corresponding lower triangular matrices using a Cholesky decomposition.

arg_constraints = {'covariance_matrix': PositiveDefinite(), 'loc': RealVector(), 'precision_matrix': PositiveDefinite(), 'scale_tril': LowerCholesky()}¶

covariance_matrix¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

log_prob(value)¶

mean¶

precision_matrix¶

rsample(sample_shape=torch.Size([]))¶

scale_tril¶

support = Real()¶

variance¶

NegativeBinomial¶

class torch.distributions.negative_binomial.NegativeBinomial(total_count, probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a Negative Binomial distribution, i.e. distribution of the number of independent identical Bernoulli trials needed before total_count failures are achieved. The probability of success of each Bernoulli trial is probs.

Parameters

total_count (float or Tensor) – non-negative number of negative Bernoulli trials to stop, although the distribution is still valid for real valued count
probs (Tensor) – Event probabilities of success in the half open interval [0, 1)
logits (Tensor) – Event log-odds for probabilities of success

arg_constraints = {'logits': Real(), 'probs': HalfOpenInterval(lower_bound=0.0, upper_bound=1.0), 'total_count': GreaterThanEq(lower_bound=0)}¶

expand(batch_shape, _instance=None)¶

log_prob(value)¶

logits¶

mean¶

param_shape¶

probs¶

sample(sample_shape=torch.Size([]))¶

support = IntegerGreaterThan(lower_bound=0)¶

variance¶

Normal¶

class torch.distributions.normal.Normal(loc, scale, validate_args=None)¶

Bases: torch.distributions.exp_family.ExponentialFamily

Creates a normal (also called Gaussian) distribution parameterized by loc and scale.

Example:

>>> m = Normal(torch.tensor([0.0]), torch.tensor([1.0]))
>>> m.sample()  # normally distributed with loc=0 and scale=1
tensor([ 0.1046])

Parameters

loc (float or Tensor) – mean of the distribution (often referred to as mu)
scale (float or Tensor) – standard deviation of the distribution (often referred to as sigma)

arg_constraints = {'loc': Real(), 'scale': GreaterThan(lower_bound=0.0)}¶

cdf(value)¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

icdf(value)¶

log_prob(value)¶

mean¶

rsample(sample_shape=torch.Size([]))¶

sample(sample_shape=torch.Size([]))¶

stddev¶

support = Real()¶

variance¶

OneHotCategorical¶

class torch.distributions.one_hot_categorical.OneHotCategorical(probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a one-hot categorical distribution parameterized by probs or logits.

Samples are one-hot coded vectors of size probs.size(-1).

Note

probs must be non-negative, finite and have a non-zero sum, and it will be normalized to sum to 1.

See also: torch.distributions.Categorical() for specifications of probs and logits.

Example:

>>> m = OneHotCategorical(torch.tensor([ 0.25, 0.25, 0.25, 0.25 ]))
>>> m.sample()  # equal probability of 0, 1, 2, 3
tensor([ 0.,  0.,  0.,  1.])

Parameters

probs (Tensor) – event probabilities
logits (Tensor) – event log probabilities

arg_constraints = {'logits': Real(), 'probs': Simplex()}¶

entropy()¶

enumerate_support(expand=True)¶

expand(batch_shape, _instance=None)¶

has_enumerate_support = True¶

log_prob(value)¶

logits¶

mean¶

param_shape¶

probs¶

sample(sample_shape=torch.Size([]))¶

support = Simplex()¶

variance¶

Pareto¶

class torch.distributions.pareto.Pareto(scale, alpha, validate_args=None)¶

Bases: torch.distributions.transformed_distribution.TransformedDistribution

Samples from a Pareto Type 1 distribution.

Example:

>>> m = Pareto(torch.tensor([1.0]), torch.tensor([1.0]))
>>> m.sample()  # sample from a Pareto distribution with scale=1 and alpha=1
tensor([ 1.5623])

Parameters

scale (float or Tensor) – Scale parameter of the distribution
alpha (float or Tensor) – Shape parameter of the distribution

arg_constraints = {'alpha': GreaterThan(lower_bound=0.0), 'scale': GreaterThan(lower_bound=0.0)}¶

entropy()¶

expand(batch_shape, _instance=None)¶

mean¶

support¶

variance¶

Poisson¶

class torch.distributions.poisson.Poisson(rate, validate_args=None)¶

Bases: torch.distributions.exp_family.ExponentialFamily

Creates a Poisson distribution parameterized by rate, the rate parameter.

Samples are nonnegative integers, with a pmf given by

\[\mathrm{rate}^k \frac{e^{-\mathrm{rate}}}{k!} \]

Example:

>>> m = Poisson(torch.tensor([4]))
>>> m.sample()
tensor([ 3.])

Parameters: rate (Number, Tensor) – the rate parameter

arg_constraints = {'rate': GreaterThan(lower_bound=0.0)}¶

expand(batch_shape, _instance=None)¶

log_prob(value)¶

mean¶

sample(sample_shape=torch.Size([]))¶

support = IntegerGreaterThan(lower_bound=0)¶

variance¶

RelaxedBernoulli¶

class torch.distributions.relaxed_bernoulli.RelaxedBernoulli(temperature, probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.transformed_distribution.TransformedDistribution

Creates a RelaxedBernoulli distribution, parametrized by temperature, and either probs or logits (but not both). This is a relaxed version of the Bernoulli distribution, so the values are in (0, 1), and has reparametrizable samples.

Example:

>>> m = RelaxedBernoulli(torch.tensor([2.2]),
                         torch.tensor([0.1, 0.2, 0.3, 0.99]))
>>> m.sample()
tensor([ 0.2951,  0.3442,  0.8918,  0.9021])

Parameters

temperature (Tensor) – relaxation temperature
probs (Number, Tensor) – the probability of sampling 1
logits (Number, Tensor) – the log-odds of sampling 1

arg_constraints = {'logits': Real(), 'probs': Interval(lower_bound=0.0, upper_bound=1.0)}¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

logits¶

probs¶

support = Interval(lower_bound=0.0, upper_bound=1.0)¶

temperature¶

LogitRelaxedBernoulli¶

class torch.distributions.relaxed_bernoulli.LogitRelaxedBernoulli(temperature, probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a LogitRelaxedBernoulli distribution parameterized by probs or logits (but not both), which is the logit of a RelaxedBernoulli distribution.

Samples are logits of values in (0, 1). See [1] for more details.

Parameters

temperature (Tensor) – relaxation temperature
probs (Number, Tensor) – the probability of sampling 1
logits (Number, Tensor) – the log-odds of sampling 1

[1] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables (Maddison et al, 2017)

[2] Categorical Reparametrization with Gumbel-Softmax (Jang et al, 2017)

arg_constraints = {'logits': Real(), 'probs': Interval(lower_bound=0.0, upper_bound=1.0)}¶

expand(batch_shape, _instance=None)¶

log_prob(value)¶

logits¶

param_shape¶

probs¶

rsample(sample_shape=torch.Size([]))¶

support = Real()¶

RelaxedOneHotCategorical¶

class torch.distributions.relaxed_categorical.RelaxedOneHotCategorical(temperature, probs=None, logits=None, validate_args=None)¶

Bases: torch.distributions.transformed_distribution.TransformedDistribution

Creates a RelaxedOneHotCategorical distribution parametrized by temperature, and either probs or logits. This is a relaxed version of the OneHotCategorical distribution, so its samples are on simplex, and are reparametrizable.

Example:

>>> m = RelaxedOneHotCategorical(torch.tensor([2.2]),
                                 torch.tensor([0.1, 0.2, 0.3, 0.4]))
>>> m.sample()
tensor([ 0.1294,  0.2324,  0.3859,  0.2523])

Parameters

temperature (Tensor) – relaxation temperature
probs (Tensor) – event probabilities
logits (Tensor) – the log probability of each event.

arg_constraints = {'logits': Real(), 'probs': Simplex()}¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

logits¶

probs¶

support = Simplex()¶

temperature¶

StudentT¶

class torch.distributions.studentT.StudentT(df, loc=0.0, scale=1.0, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Creates a Student’s t-distribution parameterized by degree of freedom df, mean loc and scale scale.

Example:

>>> m = StudentT(torch.tensor([2.0]))
>>> m.sample()  # Student's t-distributed with degrees of freedom=2
tensor([ 0.1046])

Parameters

df (float or Tensor) – degrees of freedom
loc (float or Tensor) – mean of the distribution
scale (float or Tensor) – scale of the distribution

arg_constraints = {'df': GreaterThan(lower_bound=0.0), 'loc': Real(), 'scale': GreaterThan(lower_bound=0.0)}¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

log_prob(value)¶

mean¶

rsample(sample_shape=torch.Size([]))¶

support = Real()¶

variance¶

TransformedDistribution¶

class torch.distributions.transformed_distribution.TransformedDistribution(base_distribution, transforms, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Extension of the Distribution class, which applies a sequence of Transforms to a base distribution. Let f be the composition of transforms applied:

X ~ BaseDistribution
Y = f(X) ~ TransformedDistribution(BaseDistribution, f)
log p(Y) = log p(X) + log |det (dX/dY)|

Note that the .event_shape of a TransformedDistribution is the maximum shape of its base distribution and its transforms, since transforms can introduce correlations among events.

An example for the usage of TransformedDistribution would be:

# Building a Logistic Distribution
# X ~ Uniform(0, 1)
# f = a + b * logit(X)
# Y ~ f(X) ~ Logistic(a, b)
base_distribution = Uniform(0, 1)
transforms = [SigmoidTransform().inv, AffineTransform(loc=a, scale=b)]
logistic = TransformedDistribution(base_distribution, transforms)

For more examples, please look at the implementations of Gumbel, HalfCauchy, HalfNormal, LogNormal, Pareto, Weibull, RelaxedBernoulli and RelaxedOneHotCategorical

arg_constraints = {}¶

cdf(value)¶: Computes the cumulative distribution function by inverting the transform(s) and computing the score of the base distribution.

expand(batch_shape, _instance=None)¶

has_rsample¶

icdf(value)¶: Computes the inverse cumulative distribution function using transform(s) and computing the score of the base distribution.

log_prob(value)¶: Scores the sample by inverting the transform(s) and computing the score using the score of the base distribution and the log abs det jacobian.

rsample(sample_shape=torch.Size([]))¶: Generates a sample_shape shaped reparameterized sample or sample_shape shaped batch of reparameterized samples if the distribution parameters are batched. Samples first from base distribution and applies transform() for every transform in the list.

sample(sample_shape=torch.Size([]))¶: Generates a sample_shape shaped sample or sample_shape shaped batch of samples if the distribution parameters are batched. Samples first from base distribution and applies transform() for every transform in the list.

support¶

Uniform¶

class torch.distributions.uniform.Uniform(low, high, validate_args=None)¶

Bases: torch.distributions.distribution.Distribution

Generates uniformly distributed random samples from the half-open interval [low, high).

Example:

>>> m = Uniform(torch.tensor([0.0]), torch.tensor([5.0]))
>>> m.sample()  # uniformly distributed in the range [0.0, 5.0)
tensor([ 2.3418])

Parameters

low (float or Tensor) – lower range (inclusive).
high (float or Tensor) – upper range (exclusive).

arg_constraints = {'high': Dependent(), 'low': Dependent()}¶

cdf(value)¶

entropy()¶

expand(batch_shape, _instance=None)¶

has_rsample = True¶

icdf(value)¶

log_prob(value)¶

mean¶

rsample(sample_shape=torch.Size([]))¶

stddev¶

support¶

variance¶

Weibull¶

class torch.distributions.weibull.Weibull(scale, concentration, validate_args=None)¶

Bases: torch.distributions.transformed_distribution.TransformedDistribution

Samples from a two-parameter Weibull distribution.

Example

>>> m = Weibull(torch.tensor([1.0]), torch.tensor([1.0]))
>>> m.sample()  # sample from a Weibull distribution with scale=1, concentration=1
tensor([ 0.4784])

Parameters

scale (float or Tensor) – Scale parameter of distribution (lambda).
concentration (float or Tensor) – Concentration parameter of distribution (k/shape).

arg_constraints = {'concentration': GreaterThan(lower_bound=0.0), 'scale': GreaterThan(lower_bound=0.0)}¶

entropy()¶

expand(batch_shape, _instance=None)¶

mean¶

support = GreaterThan(lower_bound=0.0)¶

variance¶

KL Divergence¶

torch.distributions.kl.kl_divergence(p, q)¶

Compute Kullback-Leibler divergence $KL(p \| q)$ between two distributions.

\[KL(p \| q) = \int p(x) \log\frac {p(x)} {q(x)} \,dx\]

Parameters

p (Distribution) – A Distribution object.
q (Distribution) – A Distribution object.

Returns

A batch of KL divergences of shape batch_shape.

Return type

Tensor

Raises

NotImplementedError – If the distribution types have not been registered via register_kl().

torch.distributions.kl.register_kl(type_p, type_q)¶

Decorator to register a pairwise function with kl_divergence(). Usage:

@register_kl(Normal, Normal)
def kl_normal_normal(p, q):
    # insert implementation here

Lookup returns the most specific (type,type) match ordered by subclass. If the match is ambiguous, a RuntimeWarning is raised. For example to resolve the ambiguous situation:

@register_kl(BaseP, DerivedQ)
def kl_version1(p, q): ...
@register_kl(DerivedP, BaseQ)
def kl_version2(p, q): ...

you should register a third most-specific implementation, e.g.:

register_kl(DerivedP, DerivedQ)(kl_version1)  # Break the tie.

Parameters

type_p (type) – A subclass of Distribution.
type_q (type) – A subclass of Distribution.

Transforms¶

class torch.distributions.transforms.Transform(cache_size=0)¶

Abstract class for invertable transformations with computable log det jacobians. They are primarily used in torch.distributions.TransformedDistribution.

Caching is useful for tranforms whose inverses are either expensive or numerically unstable. Note that care must be taken with memoized values since the autograd graph may be reversed. For example while the following works with or without caching:

y = t(x)
t.log_abs_det_jacobian(x, y).backward()  # x will receive gradients.

However the following will error when caching due to dependency reversal:

y = t(x)
z = t.inv(y)
grad(z.sum(), [y])  # error because z is x

Derived classes should implement one or both of _call() or _inverse(). Derived classes that set bijective=True should also implement log_abs_det_jacobian().

Parameters

cache_size (int) – Size of cache. If zero, no caching is done. If one, the latest single value is cached. Only 0 and 1 are supported.

Variables

~Transform.domain (Constraint) – The constraint representing valid inputs to this transform.
~Transform.codomain (Constraint) – The constraint representing valid outputs to this transform which are inputs to the inverse transform.
~Transform.bijective (bool) – Whether this transform is bijective. A transform t is bijective iff t.inv(t(x)) == x and t(t.inv(y)) == y for every x in the domain and y in the codomain. Transforms that are not bijective should at least maintain the weaker pseudoinverse properties t(t.inv(t(x)) == t(x) and t.inv(t(t.inv(y))) == t.inv(y).
~Transform.sign (int or Tensor) – For bijective univariate transforms, this should be +1 or -1 depending on whether transform is monotone increasing or decreasing.
~Transform.event_dim (int) – Number of dimensions that are correlated together in the transform event_shape. This should be 0 for pointwise transforms, 1 for transforms that act jointly on vectors, 2 for transforms that act jointly on matrices, etc.

inv¶: Returns the inverse Transform of this transform. This should satisfy t.inv.inv is t.

sign¶: Returns the sign of the determinant of the Jacobian, if applicable. In general this only makes sense for bijective transforms.

log_abs_det_jacobian(x, y)¶: Computes the log det jacobian log |dy/dx| given input and output.

class torch.distributions.transforms.ComposeTransform(parts)¶

Composes multiple transforms in a chain. The transforms being composed are responsible for caching.

Parameters: parts (list of Transform) – A list of transforms to compose.

class torch.distributions.transforms.ExpTransform(cache_size=0)¶: Transform via the mapping $y = \exp(x)$.

class torch.distributions.transforms.PowerTransform(exponent, cache_size=0)¶: Transform via the mapping $y = x^{\text{exponent}}$.

class torch.distributions.transforms.SigmoidTransform(cache_size=0)¶: Transform via the mapping $y = \frac{1}{1 + \exp(-x)}$ and $x = \text{logit}(y)$.

class torch.distributions.transforms.AbsTransform(cache_size=0)¶: Transform via the mapping $y = |x|$.

class torch.distributions.transforms.AffineTransform(loc, scale, event_dim=0, cache_size=0)¶

Transform via the pointwise affine mapping $y = \text{loc} + \text{scale} \times x$.

Parameters

loc (Tensor or float) – Location parameter.
scale (Tensor or float) – Scale parameter.
event_dim (int) – Optional size of event_shape. This should be zero for univariate random variables, 1 for distributions over vectors, 2 for distributions over matrices, etc.

class torch.distributions.transforms.SoftmaxTransform(cache_size=0)¶

Transform from unconstrained space to the simplex via $y = \exp(x)$ then normalizing.

This is not bijective and cannot be used for HMC. However this acts mostly coordinate-wise (except for the final normalization), and thus is appropriate for coordinate-wise optimization algorithms.

class torch.distributions.transforms.StickBreakingTransform(cache_size=0)¶

Transform from unconstrained space to the simplex of one additional dimension via a stick-breaking process.

This transform arises as an iterated sigmoid transform in a stick-breaking construction of the Dirichlet distribution: the first logit is transformed via sigmoid to the first probability and the probability of everything else, and then the process recurses.

This is bijective and appropriate for use in HMC; however it mixes coordinates together and is less appropriate for optimization.

class torch.distributions.transforms.LowerCholeskyTransform(cache_size=0)¶

Transform from unconstrained matrices to lower-triangular matrices with nonnegative diagonal entries.

This is useful for parameterizing positive definite matrices in terms of their Cholesky factorization.

Constraints¶

The following constraints are implemented:

constraints.boolean
constraints.dependent
constraints.greater_than(lower_bound)
constraints.integer_interval(lower_bound, upper_bound)
constraints.interval(lower_bound, upper_bound)
constraints.lower_cholesky
constraints.lower_triangular
constraints.nonnegative_integer
constraints.positive
constraints.positive_definite
constraints.positive_integer
constraints.real
constraints.real_vector
constraints.simplex
constraints.unit_interval

class torch.distributions.constraints.Constraint¶

Abstract base class for constraints.

A constraint object represents a region over which a variable is valid, e.g. within which a variable can be optimized.

check(value)¶: Returns a byte tensor of sample_shape + batch_shape indicating whether each event in value satisfies this constraint.

torch.distributions.constraints.dependent_property¶: alias of torch.distributions.constraints._DependentProperty

torch.distributions.constraints.integer_interval¶: alias of torch.distributions.constraints._IntegerInterval

torch.distributions.constraints.greater_than¶: alias of torch.distributions.constraints._GreaterThan

torch.distributions.constraints.greater_than_eq¶: alias of torch.distributions.constraints._GreaterThanEq

torch.distributions.constraints.less_than¶: alias of torch.distributions.constraints._LessThan

torch.distributions.constraints.interval¶: alias of torch.distributions.constraints._Interval

torch.distributions.constraints.half_open_interval¶: alias of torch.distributions.constraints._HalfOpenInterval

Constraint Registry¶

PyTorch provides two global ConstraintRegistry objects that link Constraint objects to Transform objects. These objects both input constraints and return transforms, but they have different guarantees on bijectivity.

biject_to(constraint) looks up a bijective Transform from constraints.real to the given constraint. The returned transform is guaranteed to have .bijective = True and should implement .log_abs_det_jacobian().
transform_to(constraint) looks up a not-necessarily bijective Transform from constraints.real to the given constraint. The returned transform is not guaranteed to implement .log_abs_det_jacobian().

The transform_to() registry is useful for performing unconstrained optimization on constrained parameters of probability distributions, which are indicated by each distribution’s .arg_constraints dict. These transforms often overparameterize a space in order to avoid rotation; they are thus more suitable for coordinate-wise optimization algorithms like Adam:

loc = torch.zeros(100, requires_grad=True)
unconstrained = torch.zeros(100, requires_grad=True)
scale = transform_to(Normal.arg_constraints['scale'])(unconstrained)
loss = -Normal(loc, scale).log_prob(data).sum()

The biject_to() registry is useful for Hamiltonian Monte Carlo, where samples from a probability distribution with constrained .support are propagated in an unconstrained space, and algorithms are typically rotation invariant.:

dist = Exponential(rate)
unconstrained = torch.zeros(100, requires_grad=True)
sample = biject_to(dist.support)(unconstrained)
potential_energy = -dist.log_prob(sample).sum()

Note

An example where transform_to and biject_to differ is constraints.simplex: transform_to(constraints.simplex) returns a SoftmaxTransform that simply exponentiates and normalizes its inputs; this is a cheap and mostly coordinate-wise operation appropriate for algorithms like SVI. In contrast, biject_to(constraints.simplex) returns a StickBreakingTransform that bijects its input down to a one-fewer-dimensional space; this a more expensive less numerically stable transform but is needed for algorithms like HMC.

The biject_to and transform_to objects can be extended by user-defined constraints and transforms using their .register() method either as a function on singleton constraints:

transform_to.register(my_constraint, my_transform)

or as a decorator on parameterized constraints:

@transform_to.register(MyConstraintClass)
def my_factory(constraint):
    assert isinstance(constraint, MyConstraintClass)
    return MyTransform(constraint.param1, constraint.param2)

You can create your own registry by creating a new ConstraintRegistry object.

class torch.distributions.constraint_registry.ConstraintRegistry¶

Registry to link constraints to transforms.

register(constraint, factory=None)¶

Registers a Constraint subclass in this registry. Usage:

@my_registry.register(MyConstraintClass)
def construct_transform(constraint):
    assert isinstance(constraint, MyConstraint)
    return MyTransform(constraint.arg_constraints)

Parameters

constraint (subclass of Constraint) – A subclass of Constraint, or a singleton object of the desired class.
factory (callable) – A callable that inputs a constraint object and returns a Transform object.

TorchScript is a way to create serializable and optimizable models from PyTorch code. Any code written in TorchScript can be saved from your Python process and loaded in a process where there is no Python dependency.

We provide tools to incrementally transition a model from being a pure Python program to a TorchScript program that can be run independently from Python, for instance, in a standalone C++ program. This makes it possible to train models in PyTorch using familiar tools and then export the model to a production environment where it is not a good idea to run models as Python programs for performance and multi-threading reasons.

Creating TorchScript Code ¶

class torch.jit.ScriptModule(optimize=True)¶

The core data structure in TorchScript is the ScriptModule. It is an analogue of torch’s nn.Module and represents an entire model as a tree of submodules. Like normal modules, each individual module in a ScriptModule can have submodules, parameters, and methods. In nn.Modules methods are implemented as Python functions, but in ScriptModules methods typically implemented as TorchScript functions, a statically-typed subset of Python that contains all of PyTorch’s built-in Tensor operations. This difference allows your ScriptModules code to run without the need for a Python interpreter.

ScriptModules and the TorchScript functions inside of them can be created in two ways:

Tracing:

Using torch.jit.trace, you can take an existing module or python function, provide example inputs, and we run the function, recording the operations performed on all the tensors. We turn the resulting recording into a TorchScript method that is installed as the forward method of a ScriptModule. This module also contains any parameters that the original module had as well.

Example:
import torch
def foo(x, y):
    return 2*x + y
traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))
Note

Tracing a function will produce a ScriptModule with a single forward method that implements that function, and that contains no parameters.

Example:
import torch
import torchvision
traced_net = torch.jit.trace(torchvision.models.resnet18(),
                             torch.rand(1, 3, 224, 224))
Note

Tracing only records operations done when the given function is run on the given tensors. Therefore, the returned ScriptModule will always run the same traced graph on any input. This has some important implications when your module is expected to run different sets of operations, depending on the input and/or the module state. For example,

Tracing will not record any control-flow like if statements or loops. When this control-flow is constant across your module, this is fine and it often just inlines configuration decisions. But sometimes the control-flow is actually part of the model itself. For instance, a recurrent network is a loop over the (possibly dynamic) length of an input sequence.

In the returned ScriptModule, operations that have different behaviors in training and eval modes will always behave as if it is in the mode it was in during tracing, no matter which mode the ScriptModule is in.

In cases like these, tracing would not be appropriate and scripting is a better choice.

Scripting:

You can write TorchScript code directly using Python syntax. You do this using the torch.jit.script annotation (for functions) or torch.jit.script_method annotation (for methods) on subclasses of ScriptModule. With this annotation the body of the annotated function is directly translated into TorchScript. TorchScript itself is a subset of the Python language, so not all features in python work, but we provide enough functionality to compute on tensors and do control-dependent operations.

Example:
import torch
@torch.jit.script
def foo(x, y):
    if x.max() > y.max():
        r = x
    else:
        r = y
    return r
Note

A script function annotation will construct a ScriptModule with a single forward method that implements that function, and that contains no parameters.

Example:
import torch
class MyModule(torch.jit.ScriptModule):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    @torch.jit.script_method
    def forward(self, input):
        return self.weight.mv(input)
Example:
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.jit import ScriptModule, script_method, trace

class MyScriptModule(ScriptModule):
    def __init__(self):
        super(MyScriptModule, self).__init__()
        # trace produces a ScriptModule's conv1 and conv2
        self.conv1 = trace(nn.Conv2d(1, 20, 5), torch.rand(1, 1, 16, 16))
        self.conv2 = trace(nn.Conv2d(20, 20, 5), torch.rand(1, 20, 16, 16))

    @script_method
    def forward(self, input):
      input = F.relu(self.conv1(input))
      input = F.relu(self.conv2(input))
      return input

save(filename)¶

Save an offline version of this module for use in a separate process. The saved module serializes all of the methods and parameters of this module. It can be loaded into the C++ API using torch::jit::load(filename) or into the Python API with torch.jit.load(filename).

To be able to save a module, it must not make any calls to native python functions. This means that all submodules must be subclasses of ScriptModules as well.

Danger

All modules, no matter their device, are always loaded onto the CPU during loading. This is different from torch.load()’s semantics and may change in the future.

torch.jit.load(f, map_location=None, _extra_files=ExtraFilesMap{})¶

Load a ScriptModule previously saved with save

All previously saved modules, no matter their device, are first loaded onto CPU, and then are moved to the devices they were saved from. If this fails (e.g. because the run time system doesn’t have certain devices), an exception is raised. However, storages can be dynamically remapped to an alternative set of devices using the map_location argument. Comparing to torch.load(), map_location in this function is simplified, which only accepts a string (e.g., ‘cpu’, ‘cuda:0’), or torch.device (e.g., torch.device(‘cpu’))

Parameters

f – a file-like object (has to implement read, readline, tell, and seek), or a string containing a file name
map_location – can a string (e.g., ‘cpu’, ‘cuda:0’), a device (e.g., torch.device(‘cpu’))
_extra_files – map from filename to content. The extra filenames given in the map would be loaded and their content would be stored in the provided map.

Returns

A ScriptModule object.

Example

>>> torch.jit.load('scriptmodule.pt')
# Load ScriptModule from io.BytesIO object
>>> with open('scriptmodule.pt', 'rb') as f:
        buffer = io.BytesIO(f.read())
# Load all tensors to the original device
>>> torch.jit.load(buffer)
# Load all tensors onto CPU, using a device
>>> torch.jit.load(buffer, map_location=torch.device('cpu'))
# Load all tensors onto CPU, using a string
>>> torch.jit.load(buffer, map_location='cpu')
# Load with extra files.
>>> files = {'metadata.json' : ''}
>>> torch.jit.load('scriptmodule.pt', _extra_files = files)
>>> print (files['metadata.json'])

torch.jit.trace(func, example_inputs, optimize=True, check_trace=True, check_inputs=None, check_tolerance=1e-05, _force_outplace=False, _module_class=None)¶

Trace a function and return an executable trace that will be optimized using just-in-time compilation.

Warning

Tracing only correctly records functions and modules which are not data dependent (e.g., have conditionals on data in tensors) and do not have any untracked external dependencies (e.g., perform input/output or access global variables). If you trace such models, you may silently get incorrect results on subsequent invocations of the model. The tracer will try to emit warnings when doing something that may cause an incorrect trace to be produced.

Parameters

func (callable or torch.nn.Module) – a python function or torch.nn.Module that will be run with example_inputs. arguments and returns to func must be Tensors or (possibly nested) tuples that contain tensors.
example_inputs (tuple) – a tuple of example inputs that will be passed to the function while tracing. The resulting trace can be run with inputs of different types and shapes assuming the traced operations support those types and shapes. example_inputs may also be a single Tensor in which case it is automatically wrapped in a tuple

Keyword Arguments

optimize (bool, optional) – whether or not to apply optimizations. Default: True.
check_trace (bool, optional) – check if the same inputs run through traced code produce the same outputs. Default: True. You might want to disable this if, for example, your network contains non- deterministic ops or if you are sure that the network is correct despite a checker failure.
check_inputs (list of tuples, optional) – A list of tuples of input arguments that should be used to check the trace against what is expected. Each tuple is equivalent to a seet of input arguments that would be specified in args. For best results, pass in a set of checking inputs representative of the space of shapes and types of inputs you expect the network to see. If not specified, the original args is used for checking
check_tolerance (float, optional) – Floating-point comparison tolerance to use in the checker procedure. This can be used to relax the checker strictness in the event that results diverge numerically for a known reason, such as operator fusion.

Returns

A ScriptModule object with a single forward() method containing the traced code. When func is a torch.nn.Module, the returned ScriptModule will have the same set of sub-modules and parameters as func.

Example

>>> def f(x):
...     return x * 2
>>> traced_f = torch.jit.trace(f, torch.rand(1))

Mixing Tracing and Scripting ¶

In many cases either tracing or script is an easier approach for converting a model. We allow you to compose tracing and scripting to suit the particular requirements of a part of a model.

Scripted functions can call traced ones. This is particularly useful when you need to use control-flow around a simple feed-forward model. For instance the beam search of a sequence to sequence model will typically be written in script but can call an encoder module generated using tracing.

Example:

import torch

def foo(x, y):
    return 2 * x + y
traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))

@torch.jit.script
def bar(x):
    return traced_foo(x, x)

Traced functions can call script functions. This is useful when a small part of a model requires some control-flow even though most of the model is just a feed-forward network. Control-flow inside of a script function called by a traced function is preserved correctly:

Example:

import torch

@torch.jit.script
def foo(x, y):
    if x.max() > y.max():
        r = x
    else:
        r = y
    return r


def bar(x, y, z):
    return foo(x, y) + z

traced_bar = torch.jit.trace(bar, (torch.rand(3), torch.rand(3), torch.rand(3))

This composition also works for modules as well, where it can be used to generate a submodule using tracing that can be called from the methods of a script module:

Example:

import torch
import torchvision

class MyScriptModule(torch.jit.ScriptModule):
    def __init__(self):
        super(MyScriptModule, self).__init__()
        self.means = torch.nn.Parameter(torch.tensor([103.939, 116.779, 123.68])
                                        .resize_(1, 3, 1, 1))
        self.resnet = torch.jit.trace(torchvision.models.resnet18(),
                                      torch.rand(1, 3, 224, 224))

    @torch.jit.script_method
    def forward(self, input):
        return self.resnet(input - self.means)

TorchScript Language Reference ¶

TorchScript is a subset of Python that can either be written directly (using the @script annotations) or generated automatically from Python code via tracing. When using tracing, code is automatically converted into this subset of Python by recording only the actual operators on tensors and simply executing and discarding the other surrounding Python code.

When writing TorchScript directly using @script annotations, the programmer must only use the subset of Python supported in TorchScript. This section documents what is supported in TorchScript as if it were a language reference for a stand alone language. Any features of Python not mentioned in this reference are not part of TorchScript.

As a subset of Python any valid TorchScript function is also a valid Python function. This makes it possible to remove the @script annotations and debug the function using standard Python tools like pdb. The reverse is not true: there are many valid python programs that are not valid TorchScript programs. Instead, TorchScript focuses specifically on the features of Python that are needed to represent neural network models in Torch.

PYTORCH_JIT=1¶: Setting the environment variable PYTORCH_JIT=0 will disable all script and tracing annotations. If there is hard-to-debug error in one of your ScriptModules, you can use this flag to force everything to run using native Python. This allows the use of tools like pdb to debug code.

Types ¶

The largest difference between TorchScript and the full Python language is that TorchScript only support a small set of types that are needed to express neural net models. In particular TorchScript supports:

Tensor: A PyTorch tensor of any dtype, dimension, or backend.
Tuple[T0, T1, ...]: A tuple containing subtypes T0, T1, etc. (e.g. Tuple[Tensor, Tensor])
bool: A boolean value
int: A scalar integer
float: A scalar floating point number
List[T]: A list of which all members are type T
Optional[T]: A value which is either None or type T
`Dict[K, V]: A dict with key type K and value type V. Only str, int, and float are allowed as key types.

Unlike Python, each variable in TorchScript function must have a single static type. This makes it easier to optimize TorchScript functions.

Example:

@torch.jit.script
def an_error(x):
    if x:
        r = torch.rand(1)
    else:
        r = 4
    return r # Type mismatch: r is set to type Tensor in the true branch
             # and type int in the false branch

There are 2 scenarios in which you can annotate:

Function Argument Type Annotation

By default, all parameters to a TorchScript function are assumed to be Tensor because this is the most common type used in modules. To specify that an argument to a TorchScript function is another type, it is possible to use MyPy-style type annotations using the types listed above:

Example:

@torch.jit.script
def foo(x, tup):
    # type: (int, Tuple[Tensor, Tensor]) -> Tensor
    t0, t1 = tup
    return t0 + t1 + x

print(foo(3, (torch.rand(3), torch.rand(3))))

Note

It is also possible to annotate types with Python 3 type annotations. In our examples, we use comment-based annotations to ensure Python 2 compatibility as well.

Variable Type Annotation

A list by default is assumed to be List[Tensor] and empty dicts Dict[str, Tensor]. To instantiate an empty list or dict of other types, use torch.jit.annotate.

Example:

import torch
from torch.jit import Tensor
from typing import List, Tuple

class EmptyDataStructures(torch.jit.ScriptModule):
    def __init__(self):
        super(EmptyDataStructures, self).__init__()

    @torch.jit.script_method
    def forward(self, x):
        # type: (Tensor) -> Tuple[List[Tuple[Tensor, Tensor]], Dict[int, Tensor]]

        # This annotates the list to be a `List[Tuple[Tensor, Tensor]]`
        list_of_tuple = torch.jit.annotate(List[Tuple[Tensor, Tensor]], [])
        for i in range(10):
            list_of_tuple.append((x, x))

            # This annotates the list to be a `Dict[int, Tensor]`
        int_tensor_dict = torch.jit.annotate(Dict[int, Tensor], {})
        return list_of_tuple, int_tensor_dict

Optional Type Refinement:

TorchScript will refine the type of a variable of type Optional[T] when a comparison to None is made inside the conditional of an if statement. The compiler can reason about multiple None checks that are combined with AND, OR, or NOT. Refinement will also occur for else blocks of if statements that are not explicitly written.

The expression must be emitted within the conditional; assigning a None check to a variable and using it in the conditional will not refine types.

Example:

@torch.jit.script
def opt_unwrap(x, y, z):
  # type: (Optional[int], Optional[int], Optional[int]) -> int
  if x is None:
    x = 1
  x = x + 1

  if y is not None and z is not None:
    x = y + z
  return x

Expressions ¶

The following Python Expressions are supported

Literals

True, False, None, 'string literals', "string literals", number literals 3 (interpreted as int) 3.4 (interpreter as a float)

Variables

a

Note

See Variable Resolution for how variables are resolved.

Tuple Construction

(3, 4), (3,)

List Construction

[3, 4], [], [torch.rand(3), torch.rand(4)]

Note

an empty list is assumed have type List[Tensor]. The types of other list literals are derived from the type of the members.

Dict Construction

{'hello': 3}, {}, {'a': torch.rand(3), 'b': torch.rand(4)}

Note

an empty dict is assumed have type Dict[str, Tensor]. The types of other dict literals are derived from the type of the members.

Arithmetic Operators

a + b a - b a * b a / b a ^ b a @ b

Comparison Operators

a == b a != b a < b a > b a <= b a >= b

Logical Operators

a and b a or b not b

Subscripts

t[0] t[-1] t[0:2] t[1:] t[:1] t[:] t[0, 1] t[0, 1:2] t[0, :1] t[-1, 1:, 0] t[1:, -1, 0] t[i:j, i]

Note

TorchScript currently does not support mutating tensors in place, so any tensor indexing can only appear on the right-hand size of an expression.

Function calls

Calls to built-in functions: torch.rand(3, dtype=torch.int)

Calls to other script functions:

import torch

@torch.jit.script
def foo(x):
  return x + 1

@torch.jit.script
def bar(x):
  return foo(x)

Method calls

Calls to methods of builtin types like tensor: x.mm(y)

When defining a Script method inside of a ScriptModule, the @script_method annotation is used. Inside of these methods it is possible to call other methods of this class or access methods on the submodules.

Calling a submodule directly (e.g. self.resnet(input)) is equivalent to calling its forward method (e.g. self.resnet.forward(input))

import torch

class MyScriptModule(torch.jit.ScriptModule):
    def __init__(self):
        super(MyScriptModule, self).__init__()
        self.means = torch.nn.Parameter(torch.tensor([103.939, 116.779, 123.68])
                                        .resize_(1, 3, 1, 1))
        self.resnet = torch.jit.trace(torchvision.models.resnet18(),
                                      torch.rand(1, 3, 224, 224))

    @torch.jit.script_method
    def helper(self, input):
      return self.resnet(input - self.means)

    @torch.jit.script_method
    def forward(self, input):
        return self.helper(input)

If expressions

x if x > y else y

Casts

float(ten), int(3.5), bool(ten)

Accessing Module Parameters

self.my_parameter self.my_submodule.my_parameter

Statements ¶

TorchScript supports the following types of statements:

Simple Assignments

a = b
a += b # short-hand for a = a + b, does not operate in-place on a
a -= b

Pattern Matching Assignments

a, b = tuple_or_list
a, b, *c = a_tuple

Print Statements

print("the result of an add:", a + b)

If Statements

if a < 4:
    r = -a
elif a < 3:
    r = a + a
else:
    r = 3 * a

While Loops

a = 0
while a < 4:
    print(a)
    a += 1

For loops with range

x = 0
for i in range(10):
    x *= i
Note

Script currently does not support iterating over generic iterable objects like lists or tensors. Script currently does not support start or increment parameters to range. These will be added in a future version.

For loops over tuples:

tup = (3, torch.rand(4))
for x in tup:
    print(x)
Note

for loops over tuples will unroll the loop, generating a body for each member of the tuple. The body must type-check correctly for each member.

For loops over constant torch.nn.ModuleList

class SubModule(torch.jit.ScriptModule):
    def __init__(self):
        super(Sub, self).__init__()
        self.weight = nn.Parameter(torch.randn(2))

    @torch.jit.script_method
    def forward(self, input):
        return self.weight + input

class MyModule(torch.jit.ScriptModule):
    __constants__ = ['mods']

    def __init__(self):
        super(MyModule, self).__init__()
        self.mods = torch.nn.ModuleList([SubModule() for i in range(10)])

    @torch.jit.script_method
    def forward(self, v):
        for module in self.mods:
            v = m(v)
        return v

Note

To use a module list inside a @script_method it must be marked constant by adding the name of the attribute to the __constants__ list for the type. For loops over a ModuleList will unroll the body of the loop at compile time, with each member of the constant module list.

Return

return a, b

Note

TorchScript allows returns in the following circumstances:

At the end of a function
In an if-statement where <true> and <false> both return
In an if-statement where <true> returns and <false> is empty (an early return)

Variable Resolution ¶

TorchScript supports a subset of Python’s variable resolution (i.e. scoping) rules. Local variables behave the same as in Python, except for the restriction that a variable must have the same type along all paths through a function. If a variable has a different type on different sides of an if statement, it is an error to use it after the end of the if statement.

Similarly, a variable is not allowed to be used if it is only defined along some paths through the function.

Example:

@torch.jit.script
def foo(x):
    if x < 0:
        y = 4
    print(y) # Error: undefined value y

Non-local variables are resolved to Python values at compile time when the function is defined. These values are then converted into TorchScript values using the rules described in Use of Python Values.

Use of Python Values ¶

To make writing TorchScript more convenient, we allow script code to refer to Python values in the surrounding scope. For instance, any time there is a reference to torch, the TorchScript compiler is actually resolving it to the torch Python module when the function is declared. These Python values are not a first class part of TorchScript. Instead they are desugared at compile-time into the primitive types that TorchScript supports. This section describes the rules that are used when accessing Python values in TorchScript. They depend on the dynamic type of the python valued referenced.

Functions

TorchScript can call python functions. This functionality is very useful when incrementally converting a model into script. The model can be moved function-by-function to script, leaving calls to Python functions in place. This way you can incrementally check the correctness of the model as you go.

Example:

def foo(x):
  print("I am called with {}".format(x))
  import pdb; pdb.set_trace()
  return x

@torch.jit.script
def bar(x)
  return foo(x + 1)

Note

Attempting to call save on a ScriptModule that contains calls to Python functions will fail. The intention is that this pathway is used for debugging and the calls removed or turned into script functions before saving.

Attribute Lookup On Python Modules

TorchScript can lookup attributes on modules. Builtin functions like torch.add are accessed this way. This allows TorchScript to call functions defined in other modules.

Python-defined Constants

TorchScript also provides a way to use constants that are defined in Python. These can be used to hard-code hyper-parameters into the function, or to define universal constants. There are two ways of specifying that a Python value should be treated as a constant.

Values looked up as attributes of a module are assumed to be constant. Example: math.pi

Attributes of a ScriptModule can be marked constant by listing them as a member of the __constants__ property of the class:

Example:

class Foo(torch.jit.ScriptModule):
    __constants__ = ['a']

    def __init__(self):
        super(Foo, self).__init__(False)
        self.a = 1 + 4

   @torch.jit.ScriptModule
   def forward(self, input):
       return self.a + input

Supported constant Python Values are

int
bool
torch.device
torch.layout
torch.dtype
tuples containing supported types
torch.nn.ModuleList which can be used in a TorchScript for loop

Debugging ¶

Disable JIT for Debugging

If you want to disable all JIT modes (tracing and scripting) so you can debug your program in raw Python, you can use the PYTORCH_JIT environment variable. PYTORCH_JIT can be used to globally disable the JIT by setting its value to 0. Given an example script:

@torch.jit.script
def scripted_fn(x : torch.Tensor):
    for i in range(12):
        x = x + x
    return x


def fn(x):
    x = torch.neg(x)
    import pdb; pdb.set_trace()
    return scripted_fn(x)

traced_fn = torch.jit.trace(fn, (torch.rand(4, 5),))

traced_fn(torch.rand(3, 4))

Debugging this script with PDB works except for when we invoke the @script function. We can globally disable JIT, so that we can call the @script function as a normal python function and not compile it. If the above script is called disable_jit_example.py, we can invoke it like so:

$ PYTORCH_JIT=0 python disable_jit_example.py

and we will be able to step into the @script function as a normal Python function.

Interpreting Graphs

TorchScript uses a static single assignment (SSA) intermediate representation (IR) to represent computation. The instructions in this format consist of ATen (the C++ backend of PyTorch) operators and other primitive operators, including control flow operators for loops and conditionals. As an example:

@torch.jit.script
def foo(len):
  # type: (int) -> torch.Tensor
  rv = torch.zeros(3, 4)
  for i in range(len):
    if i < 10:
        rv = rv - 1.0
    else:
        rv = rv + 1.0
  return rv

print(foo.graph)

A ScriptModule with a single forward method will have an attribute graph, which you can use to inspect the IR representing the computation. If the ScriptModule has more than one method, you will need to access .graph on the method itself and not the module. We can inspect the graph of a method named bar on a ScriptModule by accessing .bar.graph.

The example script above produces the graph:

graph(%len : int) {
  %15 : int = prim::Constant[value=1]()
  %9 : bool = prim::Constant[value=1]()
  %7 : Device = prim::Constant[value="cpu"]()
  %6 : int = prim::Constant[value=0]()
  %5 : int = prim::Constant[value=6]()
  %1 : int = prim::Constant[value=3]()
  %2 : int = prim::Constant[value=4]()
  %11 : int = prim::Constant[value=10]()
  %14 : float = prim::Constant[value=1]()
  %4 : int[] = prim::ListConstruct(%1, %2)
  %rv.1 : Tensor = aten::zeros(%4, %5, %6, %7)
  %rv : Tensor = prim::Loop(%len, %9, %rv.1)
    block0(%i : int, %13 : Tensor) {
      %12 : bool = aten::lt(%i, %11)
      %rv.4 : Tensor = prim::If(%12)
        block0() {
          %rv.2 : Tensor = aten::sub(%13, %14, %15)
          -> (%rv.2)
        }
        block1() {
          %rv.3 : Tensor = aten::add(%13, %14, %15)
          -> (%rv.3)
        }
      -> (%9, %rv.4)
    }
  return (%rv);
}

Take the instruction %rv.1 : Dynamic = aten::zeros(%3, %4, %5, %6) for example. %rv.1 : Dynamic means we assign the output to a (unique) value named rv.1, and that value is of Dynamic type, i.e. we do not know its concrete shape. aten::zeros is the operator (equivalent to torch.zeros) and the input list (%3, %4, %5, %6) specifies which values in scope should be passed as inputs. The schema for built-in functions like aten::zeros can be found at Builtin Functions.

Notice that operators can also have associated blocks, namely the prim::Loop and prim::If operators. In the graph print-out, these operators are formatted to reflect their equivalent source code forms to facilitate easy debugging.

Graphs can be inspected as shown to confirm that the computation described by a ScriptModule is correct, in both automated and manual fashion, as described below.

Tracing Edge Cases

There are some edge cases that exist where the trace of a given Python function/module will not be representative of the underlying code. These cases can include:

Tracing of control flow that is dependent on inputs (e.g. tensor shapes)
Tracing of in-place operations of tensor views (e.g. indexing on the left-hand side of an assignment)

Note that these cases may in fact be traceable in the future.

Automatic Trace Checking

One way to automatically catch many errors in traces is by using check_inputs on the torch.jit.trace() API. check_inputs takes a list of tuples of inputs that will be used to re-trace the computation and verify the results. For example:

def loop_in_traced_fn(x):
    result = x[0]
    for i in range(x.size(0)):
        result = result * x[i]
    return result

inputs = (torch.rand(3, 4, 5),)
check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

traced = torch.jit.trace(loop_in_traced_fn, inputs, check_inputs=check_inputs)

Gives us the following diagnostic information::

ERROR: Graphs differed across invocations! Graph diff:

  graph(%x : Tensor) {
    %1 : int = prim::Constant[value=0]()
    %2 : int = prim::Constant[value=0]()
    %result.1 : Tensor = aten::select(%x, %1, %2)
    %4 : int = prim::Constant[value=0]()
    %5 : int = prim::Constant[value=0]()
    %6 : Tensor = aten::select(%x, %4, %5)
    %result.2 : Tensor = aten::mul(%result.1, %6)
    %8 : int = prim::Constant[value=0]()
    %9 : int = prim::Constant[value=1]()
    %10 : Tensor = aten::select(%x, %8, %9)
-   %result : Tensor = aten::mul(%result.2, %10)
+   %result.3 : Tensor = aten::mul(%result.2, %10)
?          ++
    %12 : int = prim::Constant[value=0]()
    %13 : int = prim::Constant[value=2]()
    %14 : Tensor = aten::select(%x, %12, %13)
+   %result : Tensor = aten::mul(%result.3, %14)
+   %16 : int = prim::Constant[value=0]()
+   %17 : int = prim::Constant[value=3]()
+   %18 : Tensor = aten::select(%x, %16, %17)
-   %15 : Tensor = aten::mul(%result, %14)
?     ^                                 ^
+   %19 : Tensor = aten::mul(%result, %18)
?     ^                                 ^
-   return (%15);
?             ^
+   return (%19);
?             ^
  }

This message indicates to us that the computation differed between when we first traced it and when we traced it with the check_inputs. Indeed, the loop within the body of loop_in_traced_fn depends on the shape of the input x, and thus when we try another x with a different shape, the trace differs.

In this case, data-dependent control flow like this can be captured using script instead:

def fn(x):
    result = x[0]
    for i in range(x.size(0)):
        result = result * x[i]
    return result

inputs = (torch.rand(3, 4, 5),)
check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

scripted_fn = torch.jit.script(fn)
print(scripted_fn.graph)

for input_tuple in [inputs] + check_inputs:
    torch.testing.assert_allclose(fn(*input_tuple), scripted_fn(*input_tuple))

Which produces:

graph(%x : Tensor) {
  %5 : bool = prim::Constant[value=1]()
  %1 : int = prim::Constant[value=0]()
  %result.1 : Tensor = aten::select(%x, %1, %1)
  %4 : int = aten::size(%x, %1)
  %result : Tensor = prim::Loop(%4, %5, %result.1)
    block0(%i : int, %7 : Tensor) {
      %10 : Tensor = aten::select(%x, %1, %i)
      %result.2 : Tensor = aten::mul(%7, %10)
      -> (%5, %result.2)
    }
  return (%result);
}

Tracer Warnings

The tracer produces warnings for several problematic patterns in traced computation. As an example, take a trace of a function that contains an in-place assignment on a slice (a view) of a Tensor:

def fill_row_zero(x):
    x[0] = torch.rand(*x.shape[1:2])
    return x

traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
print(traced.graph)

Produces several warnings and a graph which simply returns the input:

fill_row_zero.py:4: TracerWarning: There are 2 live references to the data region being modified when tracing in-place operator copy_ (possibly due to an assignment). This might cause the trace to be incorrect, because all other views that also reference this data will not not reflect this change in the trace! On the other hand, if all other views use the same memory chunk, but are disjoint (e.g. are outputs of torch.split), this might still be safe.
  x[0] = torch.rand(*x.shape[1:2])
fill_row_zero.py:6: TracerWarning: Output nr 1. of the traced function does not match the corresponding output of the Python function. Detailed error:
Not within tolerance rtol=1e-05 atol=1e-05 at input[0, 1] (0.09115803241729736 vs. 0.6782537698745728) and 3 other locations (33.00%)
  traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
graph(%0 : Float(3, 4)) {
  return (%0);
}

We can fix this by modifying the code to not use the in-place update, but rather build up the result tensor out-of-place with torch.cat:

def fill_row_zero(x):
    x = torch.cat((torch.rand(1, *x.shape[1:2]), x[1:2]), dim=0)
    return x

traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
print(traced.graph)

Builtin Functions ¶

TorchScript supports a subset of the builtin tensor and neural network functions that PyTorch provides. Most methods on Tensor as well as functions in the torch namespace, all functions in torch.nn.functional and all modules from torch.nn are supported in TorchScript, excluding those in the table below. For unsupported modules, we suggest using torch.jit.trace().

Unsupported torch.nn Modules

torch.nn.modules.adaptive.AdaptiveLogSoftmaxWithLoss
torch.nn.modules.normalization.CrossMapLRN2d
torch.nn.modules.fold.Fold
torch.nn.modules.fold.Unfold
torch.nn.modules.rnn.GRU
torch.nn.modules.rnn.LSTM
torch.nn.modules.rnn.RNN
torch.nn.modules.rnn.GRUCell
torch.nn.modules.rnn.LSTMCell
torch.nn.modules.rnn.RNNCell

Supported Functions ¶

torch.Size(sizes : List[int]) -> List[int]

torch.abs(self : Tensor) -> Tensor

torch.abs(self : Tensor,
          out : Tensor) -> Tensor

torch.abs_(self : Tensor) -> Tensor

torch.acos(self : Tensor,
           out : Tensor) -> Tensor

torch.acos(self : Tensor) -> Tensor

torch.acos_(self : Tensor) -> Tensor

torch.adaptive_avg_pool1d(self : Tensor,
                          output_size : List[int]) -> Tensor

torch.adaptive_max_pool1d(self : Tensor,
                          output_size : List[int]) -> Tuple[Tensor, Tensor]

torch.add(self : Tensor,
          other : Tensor,
          alpha : number=1) -> Tensor

torch.add(self : Tensor,
          other : number,
          alpha : number=1) -> Tensor

torch.add(self : Tensor,
          other : Tensor,
          alpha : number=1,
          out : Tensor) -> Tensor

torch.add(a : str,
          b : str) -> str

torch.add(a : List[int],
          b : List[int]) -> List[int]

torch.add(a : List[float],
          b : List[float]) -> List[float]

torch.add(a : List[bool],
          b : List[bool]) -> List[bool]

torch.add(a : List[Tensor],
          b : List[Tensor]) -> List[Tensor]

torch.add(a : List[t],
          b : List[t]) -> List[t]

torch.add(a : int,
          b : int) -> int

torch.add(a : float,
          b : float) -> float

torch.add(a : int,
          b : float) -> float

torch.add(a : float,
          b : int) -> float

torch.addbmm(self : Tensor,
             batch1 : Tensor,
             batch2 : Tensor,
             beta : number=1,
             alpha : number=1) -> Tensor

torch.addbmm(self : Tensor,
             batch1 : Tensor,
             batch2 : Tensor,
             beta : number=1,
             alpha : number=1,
             out : Tensor) -> Tensor

torch.addcdiv(self : Tensor,
              tensor1 : Tensor,
              tensor2 : Tensor,
              value : number=1,
              out : Tensor) -> Tensor

torch.addcdiv(self : Tensor,
              tensor1 : Tensor,
              tensor2 : Tensor,
              value : number=1) -> Tensor

torch.addcmul(self : Tensor,
              tensor1 : Tensor,
              tensor2 : Tensor,
              value : number=1) -> Tensor

torch.addcmul(self : Tensor,
              tensor1 : Tensor,
              tensor2 : Tensor,
              value : number=1,
              out : Tensor) -> Tensor

torch.addmm(self : Tensor,
            mat1 : Tensor,
            mat2 : Tensor,
            beta : number=1,
            alpha : number=1,
            out : Tensor) -> Tensor

torch.addmm(self : Tensor,
            mat1 : Tensor,
            mat2 : Tensor,
            beta : number=1,
            alpha : number=1) -> Tensor

torch.addmv(self : Tensor,
            mat : Tensor,
            vec : Tensor,
            beta : number=1,
            alpha : number=1,
            out : Tensor) -> Tensor

torch.addmv(self : Tensor,
            mat : Tensor,
            vec : Tensor,
            beta : number=1,
            alpha : number=1) -> Tensor

torch.addmv_(self : Tensor,
             mat : Tensor,
             vec : Tensor,
             beta : number=1,
             alpha : number=1) -> Tensor

torch.addr(self : Tensor,
           vec1 : Tensor,
           vec2 : Tensor,
           beta : number=1,
           alpha : number=1) -> Tensor

torch.addr(self : Tensor,
           vec1 : Tensor,
           vec2 : Tensor,
           beta : number=1,
           alpha : number=1,
           out : Tensor) -> Tensor

torch.affine_grid_generator(theta : Tensor,
                            size : List[int]) -> Tensor

torch.all(self : Tensor) -> Tensor

torch.all(self : Tensor,
          dim : int,
          keepdim : bool=False) -> Tensor

torch.all(self : Tensor,
          dim : int,
          keepdim : bool=False,
          out : Tensor) -> Tensor

torch.allclose(self : Tensor,
               other : Tensor,
               rtol : float=1e-05,
               atol : float=1e-08,
               equal_nan : bool=False) -> bool

torch.alpha_dropout(input : Tensor,
                    p : float,
                    train : bool) -> Tensor

torch.alpha_dropout_(self : Tensor,
                     p : float,
                     train : bool) -> Tensor

torch.any(self : Tensor) -> Tensor

torch.any(self : Tensor,
          dim : int,
          keepdim : bool=False) -> Tensor

torch.any(self : Tensor,
          dim : int,
          keepdim : bool=False,
          out : Tensor) -> Tensor

torch.arange(end : number,
             dtype : Optional[int],
             layout : Optional[int],
             device : Optional[Device]) -> Tensor

torch.arange(start : number,
             end : number,
             dtype : Optional[int],
             layout : Optional[int],
             device : Optional[Device]) -> Tensor

torch.arange(start : number,
             end : number,
             step : number,
             dtype : Optional[int],
             layout : Optional[int],
             device : Optional[Device]) -> Tensor

torch.arange(end : number,
             out : Tensor) -> Tensor

torch.arange(start : number,
             end : number,
             step : number=1,
             out : Tensor) -> Tensor

torch.argmax(self : Tensor) -> Tensor

torch.argmax(self : Tensor,
             dim : int,
             keepdim : bool=False) -> Tensor

torch.argmin(self : Tensor) -> Tensor

torch.argmin(self : Tensor,
             dim : int,
             keepdim : bool=False) -> Tensor

torch.argsort(self : Tensor,
              dim : int=-1,
              descending : bool=False) -> Tensor

torch.as_strided(self : Tensor,
                 size : List[int],
                 stride : List[int],
                 storage_offset : Optional[int]) -> Tensor

torch.as_strided_(self : Tensor,
                  size : List[int],
                  stride : List[int],
                  storage_offset : Optional[int]) -> Tensor

torch.asin(self : Tensor) -> Tensor

torch.asin(self : Tensor,
           out : Tensor) -> Tensor

torch.asin_(self : Tensor) -> Tensor

torch.atan(self : Tensor) -> Tensor

torch.atan(self : Tensor,
           out : Tensor) -> Tensor

torch.atan2(self : Tensor,
            other : Tensor,
            out : Tensor) -> Tensor

torch.atan2(self : Tensor,
            other : Tensor) -> Tensor

torch.atan_(self : Tensor) -> Tensor

torch.avg_pool1d(self : Tensor,
                 kernel_size : List[int],
                 stride : List[int]=[],
                 padding : List[int]=[0],
                 ceil_mode : bool=False,
                 count_include_pad : bool=True) -> Tensor

torch.baddbmm(self : Tensor,
              batch1 : Tensor,
              batch2 : Tensor,
              beta : number=1,
              alpha : number=1) -> Tensor

torch.baddbmm(self : Tensor,
              batch1 : Tensor,
              batch2 : Tensor,
              beta : number=1,
              alpha : number=1,
              out : Tensor) -> Tensor

torch.bartlett_window(window_length : int,
                      dtype : Optional[int],
                      layout : Optional[int],
                      device : Optional[Device]) -> Tensor

torch.bartlett_window(window_length : int,
                      periodic : bool,
                      dtype : Optional[int],
                      layout : Optional[int],
                      device : Optional[Device]) -> Tensor

torch.batch_norm(input : Tensor,
                 weight : Optional[Tensor],
                 bias : Optional[Tensor],
                 running_mean : Optional[Tensor],
                 running_var : Optional[Tensor],
                 training : bool,
                 momentum : float,
                 eps : float,
                 cudnn_enabled : bool) -> Tensor

torch.batch_norm_backward_elemt(grad_out : Tensor,
                                input : Tensor,
                                mean : Tensor,
                                invstd : Tensor,
                                weight : Optional[Tensor],
                                mean_dy : Tensor,
                                mean_dy_xmu : Tensor) -> Tensor

torch.batch_norm_backward_reduce(grad_out : Tensor,
                                 input : Tensor,
                                 mean : Tensor,
                                 invstd : Tensor,
                                 input_g : bool,
                                 weight_g : bool,
                                 bias_g : bool) -> Tuple[Tensor, Tensor, Tensor, Tensor]

torch.batch_norm_elemt(input : Tensor,
                       weight : Optional[Tensor],
                       bias : Optional[Tensor],
                       mean : Tensor,
                       invstd : Tensor,
                       eps : float) -> Tensor

torch.batch_norm_gather_stats(input : Tensor,
                              mean : Tensor,
                              invstd : Tensor,
                              running_mean : Optional[Tensor],
                              running_var : Optional[Tensor],
                              momentum : float,
                              eps : float,
                              count : int) -> Tuple[Tensor, Tensor]

torch.batch_norm_stats(input : Tensor,
                       eps : float) -> Tuple[Tensor, Tensor]

torch.batch_norm_update_stats(input : Tensor,
                              running_mean : Optional[Tensor],
                              running_var : Optional[Tensor],
                              momentum : float) -> Tuple[Tensor, Tensor]

torch.bernoulli(self : Tensor,
                generator : Optional[Generator]) -> Tensor

torch.bernoulli(self : Tensor,
                p : float,
                generator : Optional[Generator]) -> Tensor

torch.bernoulli(self : Tensor,
                generator : Optional[Generator],
                out : Tensor) -> Tensor

torch.bilinear(input1 : Tensor,
               input2 : Tensor,
               weight : Tensor,
               bias : Optional[Tensor]) -> Tensor

torch.binary_cross_entropy_with_logits(self : Tensor,
                                       target : Tensor,
                                       weight : Optional[Tensor],
                                       pos_weight : Optional[Tensor],
                                       reduction : int) -> Tensor

torch.bincount(self : Tensor,
               weights : Optional[Tensor],
               minlength : int=0) -> Tensor

torch.blackman_window(window_length : int,
                      dtype : Optional[int],
                      layout : Optional[int],
                      device : Optional[Device]) -> Tensor

torch.blackman_window(window_length : int,
                      periodic : bool,
                      dtype : Optional[int],
                      layout : Optional[int],
                      device : Optional[Device]) -> Tensor

torch.bmm(self : Tensor,
          mat2 : Tensor) -> Tensor

torch.bmm(self : Tensor,
          mat2 : Tensor,
          out : Tensor) -> Tensor

torch.broadcast_tensors(tensors : List[Tensor]) -> List[Tensor]

torch.btrifact(self : Tensor,
               pivot : bool=True) -> Tuple[Tensor, Tensor]

torch.btrifact_with_info(self : Tensor,
                         pivot : bool=True) -> Tuple[Tensor, Tensor, Tensor]

torch.btrisolve(self : Tensor,
                LU_data : Tensor,
                LU_pivots : Tensor,
                out : Tensor) -> Tensor

torch.btrisolve(self : Tensor,
                LU_data : Tensor,
                LU_pivots : Tensor) -> Tensor

torch.cartesian_prod(tensors : List[Tensor]) -> Tensor

torch.cat(tensors : List[Tensor],
          dim : int=0) -> Tensor

torch.cat(tensors : List[Tensor],
          dim : int=0,
          out : Tensor) -> Tensor

torch.cdist(x1 : Tensor,
            x2 : Tensor,
            p : float=2.0) -> Tensor

torch.ceil(self : Tensor,
           out : Tensor) -> Tensor

torch.ceil(self : Tensor) -> Tensor

torch.ceil_(self : Tensor) -> Tensor

torch.celu(self : Tensor,
           alpha : number=1.0) -> Tensor

torch.celu_(self : Tensor,
            alpha : number=1.0) -> Tensor

torch.chain_matmul(matrices : List[Tensor]) -> Tensor

torch.cholesky(self : Tensor,
               upper : bool=False,
               out : Tensor) -> Tensor

torch.cholesky(self : Tensor,
               upper : bool=False) -> Tensor

torch.cholesky_solve(self : Tensor,
                     input2 : Tensor,
                     upper : bool=False,
                     out : Tensor) -> Tensor

torch.cholesky_solve(self : Tensor,
                     input2 : Tensor,
                     upper : bool=False) -> Tensor

torch.chunk(self : Tensor,
            chunks : int,
            dim : int=0) -> List[Tensor]

torch.clamp(self : Tensor,
            min : Optional[number],
            max : Optional[number]) -> Tensor

torch.clamp(self : Tensor,
            min : Optional[number],
            max : Optional[number],
            out : Tensor) -> Tensor

torch.clamp_(self : Tensor,
             min : Optional[number],
             max : Optional[number]) -> Tensor

torch.clamp_max(self : Tensor,
                max : number) -> Tensor

torch.clamp_max(self : Tensor,
                max : number,
                out : Tensor) -> Tensor

torch.clamp_max_(self : Tensor,
                 max : number) -> Tensor

torch.clamp_min(self : Tensor,
                min : number,
                out : Tensor) -> Tensor

torch.clamp_min(self : Tensor,
                min : number) -> Tensor

torch.clamp_min_(self : Tensor,
                 min : number) -> Tensor

torch.clone(self : Tensor) -> Tensor

torch.combinations(self : Tensor,
                   r : int=2,
                   with_replacement : bool=False) -> Tensor

torch.constant_pad_nd(self : Tensor,
                      pad : List[int],
                      value : number=0) -> Tensor

torch.conv1d(input : Tensor,
             weight : Tensor,
             bias : Optional[Tensor],
             stride : List[int]=[1],
             padding : List[int]=[0],
             dilation : List[int]=[1],
             groups : int=1) -> Tensor

torch.conv2d(input : Tensor,
             weight : Tensor,
             bias : Optional[Tensor],
             stride : List[int]=[1, 1],
             padding : List[int]=[0, 0],
             dilation : List[int]=[1, 1],
             groups : int=1) -> Tensor

torch.conv3d(input : Tensor,
             weight : Tensor,
             bias : Optional[Tensor],
             stride : List[int]=[1, 1, 1],
             padding : List[int]=[0, 0, 0],
             dilation : List[int]=[1, 1, 1],
             groups : int=1) -> Tensor

torch.conv_tbc(self : Tensor,
               weight : Tensor,
               bias : Tensor,
               pad : int=0) -> Tensor

torch.conv_transpose1d(input : Tensor,
                       weight : Tensor,
                       bias : Optional[Tensor],
                       stride : List[int]=[1],
                       padding : List[int]=[0],
                       output_padding : List[int]=[0],
                       groups : int=1,
                       dilation : List[int]=[1]) -> Tensor

torch.conv_transpose2d(input : Tensor,
                       weight : Tensor,
                       bias : Optional[Tensor],
                       stride : List[int]=[1, 1],
                       padding : List[int]=[0, 0],
                       output_padding : List[int]=[0, 0],
                       groups : int=1,
                       dilation : List[int]=[1, 1]) -> Tensor

torch.conv_transpose3d(input : Tensor,
                       weight : Tensor,
                       bias : Optional[Tensor],
                       stride : List[int]=[1, 1, 1],
                       padding : List[int]=[0, 0, 0],
                       output_padding : List[int]=[0, 0, 0],
                       groups : int=1,
                       dilation : List[int]=[1, 1, 1]) -> Tensor

torch.convolution(input : Tensor,
                  weight : Tensor,
                  bias : Optional[Tensor],
                  stride : List[int],
                  padding : List[int],
                  dilation : List[int],
                  transposed : bool,
                  output_padding : List[int],
                  groups : int) -> Tensor

torch.cos(self : Tensor) -> Tensor

torch.cos(self : Tensor,
          out : Tensor) -> Tensor

torch.cos_(self : Tensor) -> Tensor

torch.cosh(self : Tensor) -> Tensor

torch.cosh(self : Tensor,
           out : Tensor) -> Tensor

torch.cosh_(self : Tensor) -> Tensor

torch.cosine_embedding_loss(input1 : Tensor,
                            input2 : Tensor,
                            target : Tensor,
                            margin : float=0.0,
                            reduction : int=1) -> Tensor

torch.cosine_similarity(x1 : Tensor,
                        x2 : Tensor,
                        dim : int=1,
                        eps : float=1e-08) -> Tensor

torch.cross(self : Tensor,
            other : Tensor,
            dim : int=-1,
            out : Tensor) -> Tensor

torch.cross(self : Tensor,
            other : Tensor,
            dim : int=-1) -> Tensor

torch.ctc_loss(log_probs : Tensor,
               targets : Tensor,
               input_lengths : Tensor,
               target_lengths : Tensor,
               blank : int=0,
               reduction : int=1,
               zero_infinity : bool=False) -> Tensor

torch.ctc_loss(log_probs : Tensor,
               targets : Tensor,
               input_lengths : List[int],
               target_lengths : List[int],
               blank : int=0,
               reduction : int=1,
               zero_infinity : bool=False) -> Tensor

torch.cudnn_affine_grid_generator(theta : Tensor,
                                  N : int,
                                  C : int,
                                  H : int,
                                  W : int) -> Tensor

torch.cudnn_batch_norm(input : Tensor,
                       weight : Tensor,
                       bias : Optional[Tensor],
                       running_mean : Optional[Tensor],
                       running_var : Optional[Tensor],
                       training : bool,
                       exponential_average_factor : float,
                       epsilon : float) -> Tuple[Tensor, Tensor, Tensor]

torch.cudnn_convolution(self : Tensor,
                        weight : Tensor,
                        bias : Optional[Tensor],
                        padding : List[int],
                        stride : List[int],
                        dilation : List[int],
                        groups : int,
                        benchmark : bool,
                        deterministic : bool) -> Tensor

torch.cudnn_convolution_transpose(self : Tensor,
                                  weight : Tensor,
                                  bias : Optional[Tensor],
                                  padding : List[int],
                                  output_padding : List[int],
                                  stride : List[int],
                                  dilation : List[int],
                                  groups : int,
                                  benchmark : bool,
                                  deterministic : bool) -> Tensor

torch.cudnn_grid_sampler(self : Tensor,
                         grid : Tensor) -> Tensor

torch.cudnn_is_acceptable(self : Tensor) -> bool

torch.cumprod(self : Tensor,
              dim : int) -> Tensor

torch.cumprod(self : Tensor,
              dim : int,
              dtype : int) -> Tensor

torch.cumprod(self : Tensor,
              dim : int,
              out : Tensor) -> Tensor

torch.cumprod(self : Tensor,
              dim : int,
              dtype : int,
              out : Tensor) -> Tensor

torch.cumsum(self : Tensor,
             dim : int) -> Tensor

torch.cumsum(self : Tensor,
             dim : int,
             dtype : int) -> Tensor

torch.cumsum(self : Tensor,
             dim : int,
             out : Tensor) -> Tensor

torch.cumsum(self : Tensor,
             dim : int,
             dtype : int,
             out : Tensor) -> Tensor

torch.det(self : Tensor) -> Tensor

torch.detach(self : Tensor) -> Tensor

torch.detach_(self : Tensor) -> Tensor

torch.device(a : str) -> Device

torch.diag(self : Tensor,
           diagonal : int=0) -> Tensor

torch.diag(self : Tensor,
           diagonal : int=0,
           out : Tensor) -> Tensor

torch.diag_embed(self : Tensor,
                 offset : int=0,
                 dim1 : int=-2,
                 dim2 : int=-1) -> Tensor

torch.diagflat(self : Tensor,
               offset : int=0) -> Tensor

torch.diagonal(self : Tensor,
               offset : int=0,
               dim1 : int=0,
               dim2 : int=1) -> Tensor

torch.digamma(self : Tensor) -> Tensor

torch.digamma(self : Tensor,
              out : Tensor) -> Tensor

torch.dist(self : Tensor,
           other : Tensor,
           p : number=2) -> Tensor

torch.div(self : Tensor,
          other : Tensor,
          out : Tensor) -> Tensor

torch.div(self : Tensor,
          other : Tensor) -> Tensor

torch.div(self : Tensor,
          other : number) -> Tensor

torch.div(a : int,
          b : int) -> float

torch.div(a : float,
          b : float) -> float

torch.dot(self : Tensor,
          tensor : Tensor) -> Tensor

torch.dot(self : Tensor,
          tensor : Tensor,
          out : Tensor) -> Tensor

torch.dropout(input : Tensor,
              p : float,
              train : bool) -> Tensor

torch.dropout_(self : Tensor,
               p : float,
               train : bool) -> Tensor

torch.eig(self : Tensor,
          eigenvectors : bool=False) -> Tuple[Tensor, Tensor]

torch.einsum(equation : str,
             tensors : List[Tensor]) -> Tensor

torch.embedding(weight : Tensor,
                indices : Tensor,
                padding_idx : int=-1,
                scale_grad_by_freq : bool=False,
                sparse : bool=False) -> Tensor

torch.embedding_bag(weight : Tensor,
                    indices : Tensor,
                    offsets : Tensor,
                    scale_grad_by_freq : bool=False,
                    mode : int=0,
                    sparse : bool=False) -> Tuple[Tensor, Tensor, Tensor, Tensor]

torch.embedding_renorm_(self : Tensor,
                        indices : Tensor,
                        max_norm : float,
                        norm_type : float) -> Tensor

torch.empty(size : List[int],
            dtype : Optional[int],
            layout : Optional[int],
            device : Optional[Device]) -> Tensor

torch.empty(size : List[int],
            out : Tensor) -> Tensor

torch.empty_like(self : Tensor) -> Tensor

torch.empty_like(self : Tensor,
                 dtype : int,
                 layout : int,
                 device : Device) -> Tensor

torch.empty_strided(size : List[int],
                    stride : List[int],
                    dtype : Optional[int],
                    layout : Optional[int],
                    device : Optional[Device]) -> Tensor

torch.eq(self : Tensor,
         other : Tensor) -> Tensor

torch.eq(self : Tensor,
         other : number) -> Tensor

torch.eq(self : Tensor,
         other : Tensor,
         out : Tensor) -> Tensor

torch.eq(self : Tensor,
         other : number,
         out : Tensor) -> Tensor

torch.eq(a : Device,
         b : Device) -> bool

torch.eq(a : str,
         b : str) -> bool

torch.eq(a : List[int],
         b : List[int]) -> bool

torch.eq(a : List[float],
         b : List[float]) -> bool

torch.eq(a : List[Tensor],
         b : List[Tensor]) -> bool

torch.eq(a : List[bool],
         b : List[bool]) -> bool

torch.eq(a : int,
         b : int) -> bool

torch.eq(a : float,
         b : float) -> bool

torch.eq(a : int,
         b : float) -> bool

torch.eq(a : float,
         b : int) -> bool

torch.equal(self : Tensor,
            other : Tensor) -> bool

torch.erf(self : Tensor,
          out : Tensor) -> Tensor

torch.erf(self : Tensor) -> Tensor

torch.erf_(self : Tensor) -> Tensor

torch.erfc(self : Tensor,
           out : Tensor) -> Tensor

torch.erfc(self : Tensor) -> Tensor

torch.erfc_(self : Tensor) -> Tensor

torch.erfinv(self : Tensor,
             out : Tensor) -> Tensor

torch.erfinv(self : Tensor) -> Tensor

torch.exp(self : Tensor) -> Tensor

torch.exp(self : Tensor,
          out : Tensor) -> Tensor

torch.exp_(self : Tensor) -> Tensor

torch.expm1(self : Tensor,
            out : Tensor) -> Tensor

torch.expm1(self : Tensor) -> Tensor

torch.expm1_(self : Tensor) -> Tensor

torch.eye(n : int,
          out : Tensor) -> Tensor

torch.eye(n : int,
          m : int,
          out : Tensor) -> Tensor

torch.eye(n : int,
          dtype : Optional[int],
          layout : Optional[int],
          device : Optional[Device]) -> Tensor

torch.eye(n : int,
          m : int,
          dtype : Optional[int],
          layout : Optional[int],
          device : Optional[Device]) -> Tensor

torch.fbgemm_is_cpu_supported() -> bool

torch.fbgemm_linear_int8_weight(input : Tensor,
                                weight : Tensor,
                                packed : Tensor,
                                col_offsets : Tensor,
                                weight_scale : number,
                                weight_zero_point : number,
                                bias : Tensor) -> Tensor

torch.fbgemm_linear_quantize_weight(input : Tensor) -> Tuple[Tensor, Tensor, float, int]

torch.fbgemm_pack_quantized_matrix(input : Tensor,
                                   K : int,
                                   N : int) -> Tensor

torch.feature_alpha_dropout(input : Tensor,
                            p : float,
                            train : bool) -> Tensor

torch.feature_alpha_dropout_(self : Tensor,
                             p : float,
                             train : bool) -> Tensor

torch.feature_dropout(input : Tensor,
                      p : float,
                      train : bool) -> Tensor

torch.feature_dropout_(self : Tensor,
                       p : float,
                       train : bool) -> Tensor

torch.fft(self : Tensor,
          signal_ndim : int,
          normalized : bool=False) -> Tensor

torch.fill_(self : Tensor,
            value : Tensor) -> Tensor

torch.fill_(self : Tensor,
            value : number) -> Tensor

torch.flatten(self : Tensor,
              start_dim : int=0,
              end_dim : int=-1) -> Tensor

torch.flip(self : Tensor,
           dims : List[int]) -> Tensor

torch.floor(self : Tensor) -> Tensor

torch.floor(self : Tensor,
            out : Tensor) -> Tensor

torch.floor(a : float) -> int

torch.floor_(self : Tensor) -> Tensor

torch.fmod(self : Tensor,
           other : Tensor,
           out : Tensor) -> Tensor

torch.fmod(self : Tensor,
           other : number,
           out : Tensor) -> Tensor

torch.fmod(self : Tensor,
           other : Tensor) -> Tensor

torch.fmod(self : Tensor,
           other : number) -> Tensor

torch.frac(self : Tensor) -> Tensor

torch.frac(self : Tensor,
           out : Tensor) -> Tensor

torch.frobenius_norm(self : Tensor) -> Tensor

torch.frobenius_norm(self : Tensor,
                     dim : List[int],
                     keepdim : bool=False) -> Tensor

torch.frobenius_norm(self : Tensor,
                     dim : List[int],
                     keepdim : bool=False,
                     out : Tensor) -> Tensor

torch.full(size : List[int],
           fill_value : number,
           dtype : Optional[int],
           layout : Optional[int],
           device : Optional[Device]) -> Tensor

torch.full(size : List[int],
           fill_value : number,
           out : Tensor) -> Tensor

torch.full_like(self : Tensor,
                fill_value : number) -> Tensor

torch.full_like(self : Tensor,
                fill_value : number,
                dtype : int,
                layout : int,
                device : Device) -> Tensor

torch.gather(self : Tensor,
             dim : int,
             index : Tensor,
             sparse_grad : bool=False,
             out : Tensor) -> Tensor

torch.gather(self : Tensor,
             dim : int,
             index : Tensor,
             sparse_grad : bool=False) -> Tensor

torch.ge(self : Tensor,
         other : Tensor) -> Tensor

torch.ge(self : Tensor,
         other : number) -> Tensor

torch.ge(self : Tensor,
         other : Tensor,
         out : Tensor) -> Tensor

torch.ge(self : Tensor,
         other : number,
         out : Tensor) -> Tensor

torch.ge(a : int,
         b : int) -> bool

torch.ge(a : float,
         b : float) -> bool

torch.ge(a : int,
         b : float) -> bool

torch.ge(a : float,
         b : int) -> bool

torch.gels(self : Tensor,
           A : Tensor) -> Tuple[Tensor, Tensor]

torch.geqrf(self : Tensor) -> Tuple[Tensor, Tensor]

torch.ger(self : Tensor,
          vec2 : Tensor) -> Tensor

torch.ger(self : Tensor,
          vec2 : Tensor,
          out : Tensor) -> Tensor

torch.get_device(self : Tensor) -> int

torch.get_device(self : Tensor) -> int

torch.get_device(self : Tensor) -> int

torch.grid_sampler(input : Tensor,
                   grid : Tensor,
                   interpolation_mode : int,
                   padding_mode : int) -> Tensor

torch.grid_sampler_2d(input : Tensor,
                      grid : Tensor,
                      interpolation_mode : int,
                      padding_mode : int) -> Tensor

torch.grid_sampler_3d(input : Tensor,
                      grid : Tensor,
                      interpolation_mode : int,
                      padding_mode : int) -> Tensor

torch.group_norm(input : Tensor,
                 num_groups : int,
                 weight : Optional[Tensor],
                 bias : Optional[Tensor],
                 eps : float=1e-05,
                 cudnn_enabled : bool=True) -> Tensor

torch.gru(data : Tensor,
          batch_sizes : Tensor,
          hx : Tensor,
          params : List[Tensor],
          has_biases : bool,
          num_layers : int,
          dropout : float,
          train : bool,
          bidirectional : bool) -> Tuple[Tensor, Tensor]

torch.gru(input : Tensor,
          hx : Tensor,
          params : List[Tensor],
          has_biases : bool,
          num_layers : int,
          dropout : float,
          train : bool,
          bidirectional : bool,
          batch_first : bool) -> Tuple[Tensor, Tensor]

torch.gru_cell(input : Tensor,
               hx : Tensor,
               w_ih : Tensor,
               w_hh : Tensor,
               b_ih : Optional[Tensor],
               b_hh : Optional[Tensor]) -> Tensor

torch.gt(self : Tensor,
         other : Tensor) -> Tensor

torch.gt(self : Tensor,
         other : number) -> Tensor

torch.gt(self : Tensor,
         other : Tensor,
         out : Tensor) -> Tensor

torch.gt(self : Tensor,
         other : number,
         out : Tensor) -> Tensor

torch.gt(a : int,
         b : int) -> bool

torch.gt(a : float,
         b : float) -> bool

torch.gt(a : int,
         b : float) -> bool

torch.gt(a : float,
         b : int) -> bool

torch.hamming_window(window_length : int,
                     dtype : Optional[int],
                     layout : Optional[int],
                     device : Optional[Device]) -> Tensor

torch.hamming_window(window_length : int,
                     periodic : bool,
                     dtype : Optional[int],
                     layout : Optional[int],
                     device : Optional[Device]) -> Tensor

torch.hamming_window(window_length : int,
                     periodic : bool,
                     alpha : float,
                     dtype : Optional[int],
                     layout : Optional[int],
                     device : Optional[Device]) -> Tensor

torch.hamming_window(window_length : int,
                     periodic : bool,
                     alpha : float,
                     beta : float,
                     dtype : Optional[int],
                     layout : Optional[int],
                     device : Optional[Device]) -> Tensor

torch.hann_window(window_length : int,
                  dtype : Optional[int],
                  layout : Optional[int],
                  device : Optional[Device]) -> Tensor

torch.hann_window(window_length : int,
                  periodic : bool,
                  dtype : Optional[int],
                  layout : Optional[int],
                  device : Optional[Device]) -> Tensor

torch.hardshrink(self : Tensor,
                 lambd : number=0.5) -> Tensor

torch.hinge_embedding_loss(self : Tensor,
                           target : Tensor,
                           margin : float=1.0,
                           reduction : int=1) -> Tensor

torch.histc(self : Tensor,
            bins : int=100,
            min : number=0,
            max : number=0,
            out : Tensor) -> Tensor

torch.histc(self : Tensor,
            bins : int=100,
            min : number=0,
            max : number=0) -> Tensor

torch.hspmm(mat1 : Tensor,
            mat2 : Tensor) -> Tensor

torch.hspmm(mat1 : Tensor,
            mat2 : Tensor,
            out : Tensor) -> Tensor

torch.ifft(self : Tensor,
           signal_ndim : int,
           normalized : bool=False) -> Tensor

torch.index_add(self : Tensor,
                dim : int,
                index : Tensor,
                source : Tensor) -> Tensor

torch.index_copy(self : Tensor,
                 dim : int,
                 index : Tensor,
                 source : Tensor) -> Tensor

torch.index_fill(self : Tensor,
                 dim : int,
                 index : Tensor,
                 value : Tensor) -> Tensor

torch.index_fill(self : Tensor,
                 dim : int,
                 index : Tensor,
                 value : number) -> Tensor

torch.index_put(self : Tensor,
                indices : List[Optional[Tensor]],
                values : Tensor,
                accumulate : bool=False) -> Tensor

torch.index_put(self : Tensor,
                indices : List[Tensor],
                values : Tensor,
                accumulate : bool=False) -> Tensor

torch.index_put_(self : Tensor,
                 indices : List[Optional[Tensor]],
                 values : Tensor,
                 accumulate : bool=False) -> Tensor

torch.index_put_(self : Tensor,
                 indices : List[Tensor],
                 values : Tensor,
                 accumulate : bool=False) -> Tensor

torch.index_select(self : Tensor,
                   dim : int,
                   index : Tensor,
                   out : Tensor) -> Tensor

torch.index_select(self : Tensor,
                   dim : int,
                   index : Tensor) -> Tensor

torch.instance_norm(input : Tensor,
                    weight : Optional[Tensor],
                    bias : Optional[Tensor],
                    running_mean : Optional[Tensor],
                    running_var : Optional[Tensor],
                    use_input_stats : bool,
                    momentum : float,
                    eps : float,
                    cudnn_enabled : bool) -> Tensor

torch.inverse(self : Tensor,
              out : Tensor) -> Tensor

torch.inverse(self : Tensor) -> Tensor

torch.irfft(self : Tensor,
            signal_ndim : int,
            normalized : bool=False,
            onesided : bool=True,
            signal_sizes : List[int]=[]) -> Tensor

torch.is_complex(self : Tensor) -> bool

torch.is_distributed(self : Tensor) -> bool

torch.is_floating_point(self : Tensor) -> bool

torch.is_nonzero(self : Tensor) -> bool

torch.is_same_size(self : Tensor,
                   other : Tensor) -> bool

torch.is_signed(self : Tensor) -> bool

torch.isclose(self : Tensor,
              other : Tensor,
              rtol : float=1e-05,
              atol : float=1e-08,
              equal_nan : bool=False) -> Tensor

torch.isnan(self : Tensor) -> Tensor

torch.kl_div(self : Tensor,
             target : Tensor,
             reduction : int=1) -> Tensor

torch.kthvalue(self : Tensor,
               k : int,
               dim : int=-1,
               keepdim : bool=False) -> Tuple[Tensor, Tensor]

torch.layer_norm(input : Tensor,
                 normalized_shape : List[int],
                 weight : Optional[Tensor],
                 bias : Optional[Tensor],
                 eps : float=1e-05,
                 cudnn_enable : bool=True) -> Tensor

torch.le(self : Tensor,
         other : Tensor,
         out : Tensor) -> Tensor

torch.le(self : Tensor,
         other : number,
         out : Tensor) -> Tensor

torch.le(self : Tensor,
         other : Tensor) -> Tensor

torch.le(self : Tensor,
         other : number) -> Tensor

torch.le(a : int,
         b : int) -> bool

torch.le(a : float,
         b : float) -> bool

torch.le(a : int,
         b : float) -> bool

torch.le(a : float,
         b : int) -> bool

torch.lerp(self : Tensor,
           end : Tensor,
           weight : Tensor) -> Tensor

torch.lerp(self : Tensor,
           end : Tensor,
           weight : number) -> Tensor

torch.lerp(self : Tensor,
           end : Tensor,
           weight : Tensor,
           out : Tensor) -> Tensor

torch.lerp(self : Tensor,
           end : Tensor,
           weight : number,
           out : Tensor) -> Tensor

torch.lgamma(self : Tensor,
             out : Tensor) -> Tensor

torch.lgamma(self : Tensor) -> Tensor

torch.linspace(start : number,
               end : number,
               steps : int=100,
               dtype : Optional[int],
               layout : Optional[int],
               device : Optional[Device]) -> Tensor

torch.linspace(start : number,
               end : number,
               steps : int=100,
               out : Tensor) -> Tensor

torch.log(self : Tensor) -> Tensor

torch.log(self : Tensor,
          out : Tensor) -> Tensor

torch.log10(self : Tensor,
            out : Tensor) -> Tensor

torch.log10(self : Tensor) -> Tensor

torch.log10_(self : Tensor) -> Tensor

torch.log1p(self : Tensor) -> Tensor

torch.log1p(self : Tensor,
            out : Tensor) -> Tensor

torch.log1p_(self : Tensor) -> Tensor

torch.log2(self : Tensor) -> Tensor

torch.log2(self : Tensor,
           out : Tensor) -> Tensor

torch.log2_(self : Tensor) -> Tensor

torch.log_(self : Tensor) -> Tensor

torch.log_softmax(self : Tensor,
                  dim : int) -> Tensor

torch.log_softmax(self : Tensor,
                  dim : int,
                  dtype : int) -> Tensor

torch.logdet(self : Tensor) -> Tensor

torch.logspace(start : number,
               end : number,
               steps : int=100,
               dtype : Optional[int],
               layout : Optional[int],
               device : Optional[Device]) -> Tensor

torch.logspace(start : number,
               end : number,
               steps : int=100,
               out : Tensor) -> Tensor

torch.logsumexp(self : Tensor,
                dim : List[int],
                keepdim : bool=False) -> Tensor

torch.logsumexp(self : Tensor,
                dim : List[int],
                keepdim : bool=False,
                out : Tensor) -> Tensor

torch.lstm(data : Tensor,
           batch_sizes : Tensor,
           hx : List[Tensor],
           params : List[Tensor],
           has_biases : bool,
           num_layers : int,
           dropout : float,
           train : bool,
           bidirectional : bool) -> Tuple[Tensor, Tensor, Tensor]

torch.lstm(input : Tensor,
           hx : List[Tensor],
           params : List[Tensor],
           has_biases : bool,
           num_layers : int,
           dropout : float,
           train : bool,
           bidirectional : bool,
           batch_first : bool) -> Tuple[Tensor, Tensor, Tensor]

torch.lstm_cell(input : Tensor,
                hx : List[Tensor],
                w_ih : Tensor,
                w_hh : Tensor,
                b_ih : Optional[Tensor],
                b_hh : Optional[Tensor]) -> Tuple[Tensor, Tensor]

torch.lt(self : Tensor,
         other : Tensor,
         out : Tensor) -> Tensor

torch.lt(self : Tensor,
         other : number,
         out : Tensor) -> Tensor

torch.lt(self : Tensor,
         other : Tensor) -> Tensor

torch.lt(self : Tensor,
         other : number) -> Tensor

torch.lt(a : int,
         b : int) -> bool

torch.lt(a : float,
         b : float) -> bool

torch.lt(a : int,
         b : float) -> bool

torch.lt(a : float,
         b : int) -> bool

torch.margin_ranking_loss(input1 : Tensor,
                          input2 : Tensor,
                          target : Tensor,
                          margin : float=0.0,
                          reduction : int=1) -> Tensor

torch.masked_fill(self : Tensor,
                  mask : Tensor,
                  value : Tensor) -> Tensor

torch.masked_fill(self : Tensor,
                  mask : Tensor,
                  value : number) -> Tensor

torch.masked_scatter(self : Tensor,
                     mask : Tensor,
                     source : Tensor) -> Tensor

torch.masked_select(self : Tensor,
                    mask : Tensor,
                    out : Tensor) -> Tensor

torch.masked_select(self : Tensor,
                    mask : Tensor) -> Tensor

torch.matmul(self : Tensor,
             other : Tensor,
             out : Tensor) -> Tensor

torch.matmul(self : Tensor,
             other : Tensor) -> Tensor

torch.matrix_power(self : Tensor,
                   n : int) -> Tensor

torch.matrix_rank(self : Tensor,
                  symmetric : bool=False) -> Tensor

torch.matrix_rank(self : Tensor,
                  tol : float,
                  symmetric : bool=False) -> Tensor

torch.max(self : Tensor,
          other : Tensor,
          out : Tensor) -> Tensor

torch.max(self : Tensor) -> Tensor

torch.max(self : Tensor,
          other : Tensor) -> Tensor

torch.max(self : Tensor,
          dim : int,
          keepdim : bool=False) -> Tuple[Tensor, Tensor]

torch.max_pool1d_with_indices(self : Tensor,
                              kernel_size : List[int],
                              stride : List[int]=[],
                              padding : List[int]=[0],
                              dilation : List[int]=[1],
                              ceil_mode : bool=False) -> Tuple[Tensor, Tensor]

torch.mean(self : Tensor) -> Tensor

torch.mean(self : Tensor,
           dtype : int) -> Tensor

torch.mean(self : Tensor,
           dim : List[int],
           keepdim : bool=False) -> Tensor

torch.mean(self : Tensor,
           dim : List[int],
           dtype : int) -> Tensor

torch.mean(self : Tensor,
           dim : List[int],
           keepdim : bool,
           dtype : int) -> Tensor

torch.mean(self : Tensor,
           dim : List[int],
           keepdim : bool=False,
           out : Tensor) -> Tensor

torch.mean(self : Tensor,
           dim : List[int],
           dtype : int,
           out : Tensor) -> Tensor

torch.mean(self : Tensor,
           dim : List[int],
           keepdim : bool,
           dtype : int,
           out : Tensor) -> Tensor

torch.median(self : Tensor) -> Tensor

torch.median(self : Tensor,
             dim : int,
             keepdim : bool=False) -> Tuple[Tensor, Tensor]

torch.meshgrid(tensors : List[Tensor]) -> List[Tensor]

torch.min(self : Tensor) -> Tensor

torch.min(self : Tensor,
          other : Tensor) -> Tensor

torch.min(self : Tensor,
          dim : int,
          keepdim : bool=False) -> Tuple[Tensor, Tensor]

torch.min(self : Tensor,
          other : Tensor,
          out : Tensor) -> Tensor

torch.miopen_batch_norm(input : Tensor,
                        weight : Tensor,
                        bias : Optional[Tensor],
                        running_mean : Optional[Tensor],
                        running_var : Optional[Tensor],
                        training : bool,
                        exponential_average_factor : float,
                        epsilon : float) -> Tuple[Tensor, Tensor, Tensor]

torch.miopen_convolution(self : Tensor,
                         weight : Tensor,
                         bias : Optional[Tensor],
                         padding : List[int],
                         stride : List[int],
                         dilation : List[int],
                         groups : int,
                         benchmark : bool,
                         deterministic : bool) -> Tensor

torch.miopen_convolution_transpose(self : Tensor,
                                   weight : Tensor,
                                   bias : Optional[Tensor],
                                   padding : List[int],
                                   output_padding : List[int],
                                   stride : List[int],
                                   dilation : List[int],
                                   groups : int,
                                   benchmark : bool,
                                   deterministic : bool) -> Tensor

torch.miopen_depthwise_convolution(self : Tensor,
                                   weight : Tensor,
                                   bias : Optional[Tensor],
                                   padding : List[int],
                                   stride : List[int],
                                   dilation : List[int],
                                   groups : int,
                                   benchmark : bool,
                                   deterministic : bool) -> Tensor

torch.mkldnn_convolution(self : Tensor,
                         weight : Tensor,
                         bias : Optional[Tensor],
                         padding : List[int],
                         stride : List[int],
                         dilation : List[int],
                         groups : int) -> Tensor

torch.mkldnn_convolution_backward_weights(weight_size : List[int],
                                          grad_output : Tensor,
                                          self : Tensor,
                                          padding : List[int],
                                          stride : List[int],
                                          dilation : List[int],
                                          groups : int,
                                          bias_defined : bool) -> Tuple[Tensor, Tensor]

torch.mm(self : Tensor,
         mat2 : Tensor,
         out : Tensor) -> Tensor

torch.mm(self : Tensor,
         mat2 : Tensor) -> Tensor

torch.mode(self : Tensor,
           dim : int=-1,
           keepdim : bool=False) -> Tuple[Tensor, Tensor]

torch.mul(self : Tensor,
          other : Tensor) -> Tensor

torch.mul(self : Tensor,
          other : number) -> Tensor

torch.mul(self : Tensor,
          other : Tensor,
          out : Tensor) -> Tensor

torch.mul(l : List[int],
          n : int) -> List[int]

torch.mul(n : int,
          l : List[int]) -> List[int]

torch.mul(l : List[float],
          n : int) -> List[float]

torch.mul(n : int,
          l : List[float]) -> List[float]

torch.mul(l : List[bool],
          n : int) -> List[bool]

torch.mul(n : int,
          l : List[bool]) -> List[bool]

torch.mul(l : List[Tensor],
          n : int) -> List[Tensor]

torch.mul(n : int,
          l : List[Tensor]) -> List[Tensor]

torch.mul(l : List[t],
          n : int) -> List[t]

torch.mul(n : int,
          l : List[t]) -> List[t]

torch.mul(a : int,
          b : int) -> int

torch.mul(a : float,
          b : float) -> float

torch.mul(a : int,
          b : float) -> float

torch.mul(a : float,
          b : int) -> float

torch.multinomial(self : Tensor,
                  num_samples : int,
                  replacement : bool=False,
                  generator : Optional[Generator]) -> Tensor

torch.multinomial(self : Tensor,
                  num_samples : int,
                  replacement : bool=False,
                  generator : Optional[Generator],
                  out : Tensor) -> Tensor

torch.mv(self : Tensor,
         vec : Tensor,
         out : Tensor) -> Tensor

torch.mv(self : Tensor,
         vec : Tensor) -> Tensor

torch.mvlgamma(self : Tensor,
               p : int) -> Tensor

torch.narrow(self : Tensor,
             dim : int,
             start : int,
             length : int) -> Tensor

torch.native_batch_norm(input : Tensor,
                        weight : Optional[Tensor],
                        bias : Optional[Tensor],
                        running_mean : Optional[Tensor],
                        running_var : Optional[Tensor],
                        training : bool,
                        momentum : float,
                        eps : float) -> Tuple[Tensor, Tensor, Tensor]

torch.native_clone(self : Tensor) -> Tensor

torch.native_norm(self : Tensor,
                  p : number=2) -> Tensor

torch.native_pow(self : Tensor,
                 exponent : number) -> Tensor

torch.native_pow(self : Tensor,
                 exponent : number,
                 out : Tensor) -> Tensor

torch.native_resize_as_(self : Tensor,
                        the_template : Tensor) -> Tensor

torch.native_zero_(self : Tensor) -> Tensor

torch.ne(self : Tensor,
         other : Tensor) -> Tensor

torch.ne(self : Tensor,
         other : number) -> Tensor

torch.ne(self : Tensor,
         other : Tensor,
         out : Tensor) -> Tensor

torch.ne(self : Tensor,
         other : number,
         out : Tensor) -> Tensor

torch.ne(a : str,
         b : str) -> bool

torch.ne(a : List[int],
         b : List[int]) -> bool

torch.ne(a : List[float],
         b : List[float]) -> bool

torch.ne(a : List[Tensor],
         b : List[Tensor]) -> bool

torch.ne(a : List[bool],
         b : List[bool]) -> bool

torch.ne(a : int,
         b : int) -> bool

torch.ne(a : float,
         b : float) -> bool

torch.ne(a : int,
         b : float) -> bool

torch.ne(a : float,
         b : int) -> bool

torch.neg(self : Tensor,
          out : Tensor) -> Tensor

torch.neg(self : Tensor) -> Tensor

torch.neg(self : int) -> int

torch.neg(self : float) -> float

torch.nonzero(self : Tensor,
              out : Tensor) -> Tensor

torch.nonzero(self : Tensor) -> Tensor

torch.norm(self : Tensor,
           p : number=2) -> Tensor

torch.norm(self : Tensor,
           p : Optional[number],
           dtype : int) -> Tensor

torch.norm(self : Tensor,
           p : Optional[number],
           dim : List[int],
           keepdim : bool=False) -> Tensor

torch.norm(self : Tensor,
           p : Optional[number],
           dim : List[int],
           keepdim : bool,
           dtype : int) -> Tensor

torch.norm(self : Tensor,
           p : Optional[number],
           dim : List[int],
           keepdim : bool=False,
           out : Tensor) -> Tensor

torch.norm(self : Tensor,
           p : Optional[number],
           dim : List[int],
           keepdim : bool,
           dtype : int,
           out : Tensor) -> Tensor

torch.norm_except_dim(v : Tensor,
                      pow : int=2,
                      dim : int=0) -> Tensor

torch.normal(mean : Tensor,
             std : Tensor,
             generator : Optional[Generator]) -> Tensor

torch.normal(mean : float,
             std : Tensor,
             generator : Optional[Generator]) -> Tensor

torch.normal(mean : Tensor,
             std : float=1.0,
             generator : Optional[Generator]) -> Tensor

torch.normal(mean : Tensor,
             std : Tensor,
             generator : Optional[Generator],
             out : Tensor) -> Tensor

torch.normal(mean : float,
             std : Tensor,
             generator : Optional[Generator],
             out : Tensor) -> Tensor

torch.normal(mean : Tensor,
             std : float=1.0,
             generator : Optional[Generator],
             out : Tensor) -> Tensor

torch.nuclear_norm(self : Tensor,
                   keepdim : bool=False) -> Tensor

torch.nuclear_norm(self : Tensor,
                   keepdim : bool=False,
                   out : Tensor) -> Tensor

torch.numel(self : Tensor) -> int

torch.ones(size : List[int],
           out : Tensor) -> Tensor

torch.ones(size : List[int],
           dtype : Optional[int],
           layout : Optional[int],
           device : Optional[Device]) -> Tensor

torch.ones_like(self : Tensor) -> Tensor

torch.ones_like(self : Tensor,
                dtype : int,
                layout : int,
                device : Device) -> Tensor

torch.orgqr(self : Tensor,
            input2 : Tensor) -> Tensor

torch.orgqr(self : Tensor,
            input2 : Tensor,
            out : Tensor) -> Tensor

torch.ormqr(self : Tensor,
            input2 : Tensor,
            input3 : Tensor,
            left : bool=True,
            transpose : bool=False) -> Tensor

torch.ormqr(self : Tensor,
            input2 : Tensor,
            input3 : Tensor,
            left : bool=True,
            transpose : bool=False,
            out : Tensor) -> Tensor

torch.pairwise_distance(x1 : Tensor,
                        x2 : Tensor,
                        p : float=2.0,
                        eps : float=1e-06,
                        keepdim : bool=False) -> Tensor

torch.pdist(self : Tensor,
            p : float=2.0) -> Tensor

torch.pin_memory(self : Tensor) -> Tensor

torch.pinverse(self : Tensor,
               rcond : float=1e-15) -> Tensor

torch.pixel_shuffle(self : Tensor,
                    upscale_factor : int) -> Tensor

torch.poisson(self : Tensor,
              generator : Optional[Generator]) -> Tensor

torch.polygamma(n : int,
                self : Tensor) -> Tensor

torch.polygamma(n : int,
                self : Tensor,
                out : Tensor) -> Tensor

torch.potri(self : Tensor,
            upper : bool=True) -> Tensor

torch.potri(self : Tensor,
            upper : bool=True,
            out : Tensor) -> Tensor

torch.pow(self : Tensor,
          exponent : Tensor) -> Tensor

torch.pow(self : number,
          exponent : Tensor) -> Tensor

torch.pow(self : Tensor,
          exponent : number) -> Tensor

torch.pow(self : Tensor,
          exponent : Tensor,
          out : Tensor) -> Tensor

torch.pow(self : number,
          exponent : Tensor,
          out : Tensor) -> Tensor

torch.pow(self : Tensor,
          exponent : number,
          out : Tensor) -> Tensor

torch.pow(a : int,
          b : int) -> int

torch.pow(a : float,
          b : float) -> float

torch.pow(a : int,
          b : float) -> float

torch.pow(a : float,
          b : int) -> float

torch.prelu(self : Tensor,
            weight : Tensor) -> Tensor

torch.prod(self : Tensor,
           dim : int,
           keepdim : bool=False,
           out : Tensor) -> Tensor

torch.prod(self : Tensor,
           dim : int,
           dtype : int,
           out : Tensor) -> Tensor

torch.prod(self : Tensor,
           dim : int,
           keepdim : bool,
           dtype : int,
           out : Tensor) -> Tensor

torch.prod(self : Tensor) -> Tensor

torch.prod(self : Tensor,
           dtype : int) -> Tensor

torch.prod(self : Tensor,
           dim : int,
           keepdim : bool=False) -> Tensor

torch.prod(self : Tensor,
           dim : int,
           dtype : int) -> Tensor

torch.prod(self : Tensor,
           dim : int,
           keepdim : bool,
           dtype : int) -> Tensor

torch.pstrf(self : Tensor,
            upper : bool=True,
            tol : number=-1) -> Tuple[Tensor, Tensor]

torch.qr(self : Tensor) -> Tuple[Tensor, Tensor]

torch.quantized_gru_cell(input : Tensor,
                         hx : Tensor,
                         w_ih : Tensor,
                         w_hh : Tensor,
                         b_ih : Tensor,
                         b_hh : Tensor,
                         packed_ih : Tensor,
                         packed_hh : Tensor,
                         col_offsets_ih : Tensor,
                         col_offsets_hh : Tensor,
                         scale_ih : number,
                         scale_hh : number,
                         zero_point_ih : number,
                         zero_point_hh : number) -> Tensor

torch.quantized_lstm(input : Tensor,
                     hx : List[Tensor],
                     params : List[Tensor],
                     has_biases : bool,
                     num_layers : int,
                     dropout : float,
                     train : bool,
                     bidirectional : bool,
                     batch_first : bool) -> Tuple[Tensor, Tensor, Tensor]

torch.quantized_lstm_cell(input : Tensor,
                          hx : List[Tensor],
                          w_ih : Tensor,
                          w_hh : Tensor,
                          b_ih : Tensor,
                          b_hh : Tensor,
                          packed_ih : Tensor,
                          packed_hh : Tensor,
                          col_offsets_ih : Tensor,
                          col_offsets_hh : Tensor,
                          scale_ih : number,
                          scale_hh : number,
                          zero_point_ih : number,
                          zero_point_hh : number) -> Tuple[Tensor, Tensor]

torch.quantized_rnn_relu_cell(input : Tensor,
                              hx : Tensor,
                              w_ih : Tensor,
                              w_hh : Tensor,
                              b_ih : Tensor,
                              b_hh : Tensor,
                              packed_ih : Tensor,
                              packed_hh : Tensor,
                              col_offsets_ih : Tensor,
                              col_offsets_hh : Tensor,
                              scale_ih : number,
                              scale_hh : number,
                              zero_point_ih : number,
                              zero_point_hh : number) -> Tensor

torch.quantized_rnn_tanh_cell(input : Tensor,
                              hx : Tensor,
                              w_ih : Tensor,
                              w_hh : Tensor,
                              b_ih : Tensor,
                              b_hh : Tensor,
                              packed_ih : Tensor,
                              packed_hh : Tensor,
                              col_offsets_ih : Tensor,
                              col_offsets_hh : Tensor,
                              scale_ih : number,
                              scale_hh : number,
                              zero_point_ih : number,
                              zero_point_hh : number) -> Tensor

torch.rand(size : List[int],
           dtype : Optional[int],
           layout : Optional[int],
           device : Optional[Device]) -> Tensor

torch.rand(size : List[int],
           out : Tensor) -> Tensor

torch.rand_like(self : Tensor) -> Tensor

torch.rand_like(self : Tensor,
                dtype : int,
                layout : int,
                device : Device) -> Tensor

torch.randint(high : int,
              size : List[int],
              out : Tensor) -> Tensor

torch.randint(low : int,
              high : int,
              size : List[int],
              out : Tensor) -> Tensor

torch.randint(high : int,
              size : List[int],
              dtype : Optional[int],
              layout : Optional[int],
              device : Optional[Device]) -> Tensor

torch.randint(low : int,
              high : int,
              size : List[int],
              dtype : Optional[int],
              layout : Optional[int],
              device : Optional[Device]) -> Tensor

torch.randint_like(self : Tensor,
                   high : int) -> Tensor

torch.randint_like(self : Tensor,
                   low : int,
                   high : int) -> Tensor

torch.randint_like(self : Tensor,
                   high : int,
                   dtype : int,
                   layout : int,
                   device : Device) -> Tensor

torch.randint_like(self : Tensor,
                   low : int,
                   high : int,
                   dtype : int,
                   layout : int,
                   device : Device) -> Tensor

torch.randn(size : List[int],
            dtype : Optional[int],
            layout : Optional[int],
            device : Optional[Device]) -> Tensor

torch.randn(size : List[int],
            out : Tensor) -> Tensor

torch.randn_like(self : Tensor) -> Tensor

torch.randn_like(self : Tensor,
                 dtype : int,
                 layout : int,
                 device : Device) -> Tensor

torch.randperm(n : int,
               out : Tensor) -> Tensor

torch.randperm(n : int,
               dtype : Optional[int],
               layout : Optional[int],
               device : Optional[Device]) -> Tensor

torch.range(start : number,
            end : number,
            dtype : Optional[int],
            layout : Optional[int],
            device : Optional[Device]) -> Tensor

torch.range(start : number,
            end : number,
            step : number=1,
            dtype : Optional[int],
            layout : Optional[int],
            device : Optional[Device]) -> Tensor

torch.range(start : number,
            end : number,
            step : number=1,
            out : Tensor) -> Tensor

torch.reciprocal(self : Tensor) -> Tensor

torch.reciprocal(self : Tensor,
                 out : Tensor) -> Tensor

torch.relu(self : Tensor) -> Tensor

torch.relu_(self : Tensor) -> Tensor

torch.remainder(self : Tensor,
                other : Tensor) -> Tensor

torch.remainder(self : Tensor,
                other : number) -> Tensor

torch.remainder(self : Tensor,
                other : Tensor,
                out : Tensor) -> Tensor

torch.remainder(self : Tensor,
                other : number,
                out : Tensor) -> Tensor

torch.remainder(a : int,
                b : int) -> int

torch.remainder(a : float,
                b : float) -> float

torch.remainder(a : int,
                b : float) -> float

torch.remainder(a : float,
                b : int) -> float

torch.renorm(self : Tensor,
             p : number,
             dim : int,
             maxnorm : number,
             out : Tensor) -> Tensor

torch.renorm(self : Tensor,
             p : number,
             dim : int,
             maxnorm : number) -> Tensor

torch.reshape(self : Tensor,
              shape : List[int]) -> Tensor

torch.resize_as_(self : Tensor,
                 the_template : Tensor) -> Tensor

torch.rfft(self : Tensor,
           signal_ndim : int,
           normalized : bool=False,
           onesided : bool=True) -> Tensor

torch.rnn_relu(data : Tensor,
               batch_sizes : Tensor,
               hx : Tensor,
               params : List[Tensor],
               has_biases : bool,
               num_layers : int,
               dropout : float,
               train : bool,
               bidirectional : bool) -> Tuple[Tensor, Tensor]

torch.rnn_relu(input : Tensor,
               hx : Tensor,
               params : List[Tensor],
               has_biases : bool,
               num_layers : int,
               dropout : float,
               train : bool,
               bidirectional : bool,
               batch_first : bool) -> Tuple[Tensor, Tensor]

torch.rnn_relu_cell(input : Tensor,
                    hx : Tensor,
                    w_ih : Tensor,
                    w_hh : Tensor,
                    b_ih : Optional[Tensor],
                    b_hh : Optional[Tensor]) -> Tensor

torch.rnn_tanh(data : Tensor,
               batch_sizes : Tensor,
               hx : Tensor,
               params : List[Tensor],
               has_biases : bool,
               num_layers : int,
               dropout : float,
               train : bool,
               bidirectional : bool) -> Tuple[Tensor, Tensor]

torch.rnn_tanh(input : Tensor,
               hx : Tensor,
               params : List[Tensor],
               has_biases : bool,
               num_layers : int,
               dropout : float,
               train : bool,
               bidirectional : bool,
               batch_first : bool) -> Tuple[Tensor, Tensor]

torch.rnn_tanh_cell(input : Tensor,
                    hx : Tensor,
                    w_ih : Tensor,
                    w_hh : Tensor,
                    b_ih : Optional[Tensor],
                    b_hh : Optional[Tensor]) -> Tensor

torch.roll(self : Tensor,
           shifts : List[int],
           dims : List[int]=[]) -> Tensor

torch.rot90(self : Tensor,
            k : int=1,
            dims : List[int]=[0, 1]) -> Tensor

torch.round(self : Tensor) -> Tensor

torch.round(self : Tensor,
            out : Tensor) -> Tensor

torch.round_(self : Tensor) -> Tensor

torch.rrelu(self : Tensor,
            lower : number=0.125,
            upper : number=0.3333333333333333,
            training : bool=False,
            generator : Optional[Generator]) -> Tensor

torch.rrelu_(self : Tensor,
             lower : number=0.125,
             upper : number=0.3333333333333333,
             training : bool=False,
             generator : Optional[Generator]) -> Tensor

torch.rsqrt(self : Tensor,
            out : Tensor) -> Tensor

torch.rsqrt(self : Tensor) -> Tensor

torch.rsqrt_(self : Tensor) -> Tensor

torch.rsub(self : Tensor,
           other : Tensor,
           alpha : number=1) -> Tensor

torch.rsub(self : Tensor,
           other : number,
           alpha : number=1) -> Tensor

torch.s_copy_(self : Tensor,
              src : Tensor,
              non_blocking : bool=False) -> Tensor

torch.s_native_addmm(self : Tensor,
                     mat1 : Tensor,
                     mat2 : Tensor,
                     beta : number=1,
                     alpha : number=1) -> Tensor

torch.s_native_addmm(self : Tensor,
                     mat1 : Tensor,
                     mat2 : Tensor,
                     beta : number=1,
                     alpha : number=1,
                     out : Tensor) -> Tensor

torch.s_native_addmm_(self : Tensor,
                      mat1 : Tensor,
                      mat2 : Tensor,
                      beta : number=1,
                      alpha : number=1) -> Tensor

torch.scalar_tensor(s : number,
                    dtype : Optional[int],
                    layout : Optional[int],
                    device : Optional[Device]) -> Tensor

torch.scatter(self : Tensor,
              dim : int,
              index : Tensor,
              src : Tensor) -> Tensor

torch.scatter(self : Tensor,
              dim : int,
              index : Tensor,
              value : number) -> Tensor

torch.scatter_add(self : Tensor,
                  dim : int,
                  index : Tensor,
                  src : Tensor) -> Tensor

torch.select(self : Tensor,
             dim : int,
             index : int) -> Tensor

torch.select(list : List[Tensor],
             idx : int) -> Tensor

torch.select(a : List[int],
             b : int) -> int

torch.select(a : List[float],
             b : int) -> float

torch.select(a : List[bool],
             b : int) -> bool

torch.select(list : List[t],
             idx : int) -> t

torch.selu(self : Tensor) -> Tensor

torch.selu_(self : Tensor) -> Tensor

torch.sigmoid(self : Tensor) -> Tensor

torch.sigmoid(self : Tensor,
              out : Tensor) -> Tensor

torch.sigmoid_(self : Tensor) -> Tensor

torch.sign(self : Tensor) -> Tensor

torch.sign(self : Tensor,
           out : Tensor) -> Tensor

torch.sin(self : Tensor) -> Tensor

torch.sin(self : Tensor,
          out : Tensor) -> Tensor

torch.sin_(self : Tensor) -> Tensor

torch.sinh(self : Tensor,
           out : Tensor) -> Tensor

torch.sinh(self : Tensor) -> Tensor

torch.sinh_(self : Tensor) -> Tensor

torch.slogdet(self : Tensor) -> Tuple[Tensor, Tensor]

torch.smm(self : Tensor,
          mat2 : Tensor) -> Tensor

torch.softmax(self : Tensor,
              dim : int) -> Tensor

torch.softmax(self : Tensor,
              dim : int,
              dtype : int) -> Tensor

torch.solve(self : Tensor,
            A : Tensor) -> Tuple[Tensor, Tensor]

torch.sort(self : Tensor,
           dim : int=-1,
           descending : bool=False) -> Tuple[Tensor, Tensor]

torch.sparse_coo_tensor(size : List[int],
                        dtype : int,
                        layout : int,
                        device : Device) -> Tensor

torch.sparse_coo_tensor(indices : Tensor,
                        values : Tensor,
                        dtype : Optional[int],
                        layout : Optional[int],
                        device : Optional[Device]) -> Tensor

torch.sparse_coo_tensor(indices : Tensor,
                        values : Tensor,
                        size : List[int],
                        dtype : Optional[int],
                        layout : Optional[int],
                        device : Optional[Device]) -> Tensor

torch.split(self : Tensor,
            split_size : int,
            dim : int=0) -> List[Tensor]

torch.split(self : Tensor,
            split_sizes : List[int],
            dim : int=0) -> List[Tensor]

torch.split_with_sizes(self : Tensor,
                       split_sizes : List[int],
                       dim : int=0) -> List[Tensor]

torch.sqrt(self : Tensor,
           out : Tensor) -> Tensor

torch.sqrt(self : Tensor) -> Tensor

torch.sqrt_(self : Tensor) -> Tensor

torch.squeeze(self : Tensor) -> Tensor

torch.squeeze(self : Tensor,
              dim : int) -> Tensor

torch.sspaddmm(self : Tensor,
               mat1 : Tensor,
               mat2 : Tensor,
               beta : number=1,
               alpha : number=1,
               out : Tensor) -> Tensor

torch.sspaddmm(self : Tensor,
               mat1 : Tensor,
               mat2 : Tensor,
               beta : number=1,
               alpha : number=1) -> Tensor

torch.stack(tensors : List[Tensor],
            dim : int=0) -> Tensor

torch.stack(tensors : List[Tensor],
            dim : int=0,
            out : Tensor) -> Tensor

torch.std(self : Tensor,
          unbiased : bool=True) -> Tensor

torch.std(self : Tensor,
          dim : List[int],
          unbiased : bool=True,
          keepdim : bool=False) -> Tensor

torch.std(self : Tensor,
          dim : List[int],
          unbiased : bool=True,
          keepdim : bool=False,
          out : Tensor) -> Tensor

torch.stft(self : Tensor,
           n_fft : int,
           hop_length : Optional[int],
           win_length : Optional[int],
           window : Optional[Tensor],
           normalized : bool=False,
           onesided : bool=True) -> Tensor

torch.sub(self : Tensor,
          other : Tensor,
          alpha : number=1) -> Tensor

torch.sub(self : Tensor,
          other : number,
          alpha : number=1) -> Tensor

torch.sub(self : Tensor,
          other : Tensor,
          alpha : number=1,
          out : Tensor) -> Tensor

torch.sub(a : int,
          b : int) -> int

torch.sub(a : float,
          b : float) -> float

torch.sub(a : int,
          b : float) -> float

torch.sub(a : float,
          b : int) -> float

torch.sum(self : Tensor,
          dim : List[int],
          keepdim : bool=False,
          out : Tensor) -> Tensor

torch.sum(self : Tensor,
          dim : List[int],
          dtype : int,
          out : Tensor) -> Tensor

torch.sum(self : Tensor,
          dim : List[int],
          keepdim : bool,
          dtype : int,
          out : Tensor) -> Tensor

torch.sum(self : Tensor) -> Tensor

torch.sum(self : Tensor,
          dtype : int) -> Tensor

torch.sum(self : Tensor,
          dim : List[int],
          keepdim : bool=False) -> Tensor

torch.sum(self : Tensor,
          dim : List[int],
          dtype : int) -> Tensor

torch.sum(self : Tensor,
          dim : List[int],
          keepdim : bool,
          dtype : int) -> Tensor

torch.svd(self : Tensor,
          some : bool=True,
          compute_uv : bool=True) -> Tuple[Tensor, Tensor, Tensor]

torch.symeig(self : Tensor,
             eigenvectors : bool=False,
             upper : bool=True) -> Tuple[Tensor, Tensor]

torch.t(self : Tensor) -> Tensor

torch.take(self : Tensor,
           index : Tensor) -> Tensor

torch.take(self : Tensor,
           index : Tensor,
           out : Tensor) -> Tensor

torch.tan(self : Tensor,
          out : Tensor) -> Tensor

torch.tan(self : Tensor) -> Tensor

torch.tan_(self : Tensor) -> Tensor

torch.tanh(self : Tensor) -> Tensor

torch.tanh(self : Tensor,
           out : Tensor) -> Tensor

torch.tanh_(self : Tensor) -> Tensor

torch.tensor(t : float,
             dtype : Optional[int],
             device : Optional[Device]) -> Tensor

torch.tensor(t : int,
             dtype : Optional[int],
             device : Optional[Device]) -> Tensor

torch.tensor(t : bool,
             dtype : Optional[int],
             device : Optional[Device]) -> Tensor

torch.tensor(data : List[t],
             dtype : Optional[int],
             device : Optional[Device]) -> Tensor

torch.tensordot(self : Tensor,
                other : Tensor,
                dims_self : List[int],
                dims_other : List[int]) -> Tensor

torch.threshold(self : Tensor,
                threshold : number,
                value : number) -> Tensor

torch.threshold(self : Tensor,
                threshold : number,
                value : number,
                out : Tensor) -> Tensor

torch.threshold_(self : Tensor,
                 threshold : number,
                 value : number) -> Tensor

torch.topk(self : Tensor,
           k : int,
           dim : int=-1,
           largest : bool=True,
           sorted : bool=True) -> Tuple[Tensor, Tensor]

torch.trace(self : Tensor) -> Tensor

torch.transpose(self : Tensor,
                dim0 : int,
                dim1 : int) -> Tensor

torch.tril(self : Tensor,
           diagonal : int=0,
           out : Tensor) -> Tensor

torch.tril(self : Tensor,
           diagonal : int=0) -> Tensor

torch.tril_indices(row : int,
                   col : int,
                   offset : int=0,
                   dtype : Optional[int]=4,
                   layout : Optional[int],
                   device : Optional[Device]) -> Tensor

torch.triplet_margin_loss(anchor : Tensor,
                          positive : Tensor,
                          negative : Tensor,
                          margin : float=1.0,
                          p : float=2.0,
                          eps : float=1e-06,
                          swap : bool=False,
                          reduction : int=1) -> Tensor

torch.triu(self : Tensor,
           diagonal : int=0,
           out : Tensor) -> Tensor

torch.triu(self : Tensor,
           diagonal : int=0) -> Tensor

torch.triu_indices(row : int,
                   col : int,
                   offset : int=0,
                   dtype : Optional[int]=4,
                   layout : Optional[int],
                   device : Optional[Device]) -> Tensor

torch.trtrs(self : Tensor,
            A : Tensor,
            upper : bool=True,
            transpose : bool=False,
            unitriangular : bool=False) -> Tuple[Tensor, Tensor]

torch.trunc(self : Tensor) -> Tensor

torch.trunc(self : Tensor,
            out : Tensor) -> Tensor

torch.trunc_(self : Tensor) -> Tensor

torch.unbind(self : Tensor,
             dim : int=0) -> List[Tensor]

torch.unsqueeze(self : Tensor,
                dim : int) -> Tensor

torch.var(self : Tensor,
          dim : List[int],
          unbiased : bool=True,
          keepdim : bool=False,
          out : Tensor) -> Tensor

torch.var(self : Tensor,
          unbiased : bool=True) -> Tensor

torch.var(self : Tensor,
          dim : List[int],
          unbiased : bool=True,
          keepdim : bool=False) -> Tensor

torch.wait(self : Future[t]) -> t

torch.where(condition : Tensor,
            self : Tensor,
            other : Tensor) -> Tensor

torch.zero_(self : Tensor) -> Tensor

torch.zeros(size : List[int],
            out : Tensor) -> Tensor

torch.zeros(size : List[int],
            dtype : Optional[int],
            layout : Optional[int],
            device : Optional[Device]) -> Tensor

torch.zeros_like(self : Tensor) -> Tensor

torch.zeros_like(self : Tensor,
                 dtype : int,
                 layout : int,
                 device : Device) -> Tensor

torch._C._nn.adaptive_avg_pool2d(self : Tensor,
                                 output_size : List[int],
                                 out : Tensor) -> Tensor

torch._C._nn.adaptive_avg_pool2d(self : Tensor,
                                 output_size : List[int]) -> Tensor

torch._C._nn.adaptive_avg_pool3d(self : Tensor,
                                 output_size : List[int]) -> Tensor

torch._C._nn.adaptive_avg_pool3d(self : Tensor,
                                 output_size : List[int],
                                 out : Tensor) -> Tensor

torch._C._nn.adaptive_max_pool2d(self : Tensor,
                                 output_size : List[int]) -> Tuple[Tensor, Tensor]

torch._C._nn.adaptive_max_pool3d(self : Tensor,
                                 output_size : List[int]) -> Tuple[Tensor, Tensor]

torch._C._nn.avg_pool2d(self : Tensor,
                        kernel_size : List[int],
                        stride : List[int]=[],
                        padding : List[int]=[0, 0],
                        ceil_mode : bool=False,
                        count_include_pad : bool=True) -> Tensor

torch._C._nn.avg_pool2d(self : Tensor,
                        kernel_size : List[int],
                        stride : List[int]=[],
                        padding : List[int]=[0, 0],
                        ceil_mode : bool=False,
                        count_include_pad : bool=True,
                        out : Tensor) -> Tensor

torch._C._nn.avg_pool3d(self : Tensor,
                        kernel_size : List[int],
                        stride : List[int]=[],
                        padding : List[int]=[0, 0, 0],
                        ceil_mode : bool=False,
                        count_include_pad : bool=True) -> Tensor

torch._C._nn.avg_pool3d(self : Tensor,
                        kernel_size : List[int],
                        stride : List[int]=[],
                        padding : List[int]=[0, 0, 0],
                        ceil_mode : bool=False,
                        count_include_pad : bool=True,
                        out : Tensor) -> Tensor

torch._C._nn.binary_cross_entropy(self : Tensor,
                                  target : Tensor,
                                  weight : Optional[Tensor],
                                  reduction : int=1,
                                  out : Tensor) -> Tensor

torch._C._nn.binary_cross_entropy(self : Tensor,
                                  target : Tensor,
                                  weight : Optional[Tensor],
                                  reduction : int=1) -> Tensor

torch._C._nn.elu(self : Tensor,
                 alpha : number=1,
                 scale : number=1,
                 input_scale : number=1,
                 out : Tensor) -> Tensor

torch._C._nn.elu(self : Tensor,
                 alpha : number=1,
                 scale : number=1,
                 input_scale : number=1) -> Tensor

torch._C._nn.elu_(self : Tensor,
                  alpha : number=1,
                  scale : number=1,
                  input_scale : number=1) -> Tensor

torch._C._nn.fractional_max_pool2d(self : Tensor,
                                   kernel_size : List[int],
                                   output_size : List[int],
                                   random_samples : Tensor) -> Tuple[Tensor, Tensor]

torch._C._nn.fractional_max_pool3d(self : Tensor,
                                   kernel_size : List[int],
                                   output_size : List[int],
                                   random_samples : Tensor) -> Tuple[Tensor, Tensor]

torch._C._nn.glu(self : Tensor,
                 dim : int=-1) -> Tensor

torch._C._nn.glu(self : Tensor,
                 dim : int=-1,
                 out : Tensor) -> Tensor

torch._C._nn.hardtanh(self : Tensor,
                      min_val : number=-1,
                      max_val : number=1,
                      out : Tensor) -> Tensor

torch._C._nn.hardtanh(self : Tensor,
                      min_val : number=-1,
                      max_val : number=1) -> Tensor

torch._C._nn.hardtanh_(self : Tensor,
                       min_val : number=-1,
                       max_val : number=1) -> Tensor

torch._C._nn.l1_loss(self : Tensor,
                     target : Tensor,
                     reduction : int=1,
                     out : Tensor) -> Tensor

torch._C._nn.l1_loss(self : Tensor,
                     target : Tensor,
                     reduction : int=1) -> Tensor

torch._C._nn.leaky_relu(self : Tensor,
                        negative_slope : number=0.01) -> Tensor

torch._C._nn.leaky_relu(self : Tensor,
                        negative_slope : number=0.01,
                        out : Tensor) -> Tensor

torch._C._nn.leaky_relu_(self : Tensor,
                         negative_slope : number=0.01) -> Tensor

torch._C._nn.log_sigmoid(self : Tensor) -> Tensor

torch._C._nn.log_sigmoid(self : Tensor,
                         out : Tensor) -> Tensor

torch._C._nn.max_pool2d_with_indices(self : Tensor,
                                     kernel_size : List[int],
                                     stride : List[int]=[],
                                     padding : List[int]=[0, 0],
                                     dilation : List[int]=[1, 1],
                                     ceil_mode : bool=False) -> Tuple[Tensor, Tensor]

torch._C._nn.max_pool3d_with_indices(self : Tensor,
                                     kernel_size : List[int],
                                     stride : List[int]=[],
                                     padding : List[int]=[0, 0, 0],
                                     dilation : List[int]=[1, 1, 1],
                                     ceil_mode : bool=False) -> Tuple[Tensor, Tensor]

torch._C._nn.max_unpool2d(self : Tensor,
                          indices : Tensor,
                          output_size : List[int]) -> Tensor

torch._C._nn.max_unpool2d(self : Tensor,
                          indices : Tensor,
                          output_size : List[int],
                          out : Tensor) -> Tensor

torch._C._nn.max_unpool3d(self : Tensor,
                          indices : Tensor,
                          output_size : List[int],
                          stride : List[int],
                          padding : List[int],
                          out : Tensor) -> Tensor

torch._C._nn.max_unpool3d(self : Tensor,
                          indices : Tensor,
                          output_size : List[int],
                          stride : List[int],
                          padding : List[int]) -> Tensor

torch._C._nn.mse_loss(self : Tensor,
                      target : Tensor,
                      reduction : int=1,
                      out : Tensor) -> Tensor

torch._C._nn.mse_loss(self : Tensor,
                      target : Tensor,
                      reduction : int=1) -> Tensor

torch._C._nn.multi_margin_loss(self : Tensor,
                               target : Tensor,
                               p : number=1,
                               margin : number=1,
                               weight : Optional[Tensor],
                               reduction : int=1,
                               out : Tensor) -> Tensor

torch._C._nn.multi_margin_loss(self : Tensor,
                               target : Tensor,
                               p : number=1,
                               margin : number=1,
                               weight : Optional[Tensor],
                               reduction : int=1) -> Tensor

torch._C._nn.multilabel_margin_loss(self : Tensor,
                                    target : Tensor,
                                    reduction : int=1) -> Tensor

torch._C._nn.multilabel_margin_loss(self : Tensor,
                                    target : Tensor,
                                    reduction : int=1,
                                    out : Tensor) -> Tensor

torch._C._nn.nll_loss(self : Tensor,
                      target : Tensor,
                      weight : Optional[Tensor],
                      reduction : int=1,
                      ignore_index : int=-100) -> Tensor

torch._C._nn.nll_loss(self : Tensor,
                      target : Tensor,
                      weight : Optional[Tensor],
                      reduction : int=1,
                      ignore_index : int=-100,
                      out : Tensor) -> Tensor

torch._C._nn.nll_loss2d(self : Tensor,
                        target : Tensor,
                        weight : Optional[Tensor],
                        reduction : int=1,
                        ignore_index : int=-100) -> Tensor

torch._C._nn.nll_loss2d(self : Tensor,
                        target : Tensor,
                        weight : Optional[Tensor],
                        reduction : int=1,
                        ignore_index : int=-100,
                        out : Tensor) -> Tensor

torch._C._nn.one_hot(self : Tensor,
                     num_classes : int=-1) -> Tensor

torch._C._nn.reflection_pad1d(self : Tensor,
                              padding : List[int]) -> Tensor

torch._C._nn.reflection_pad1d(self : Tensor,
                              padding : List[int],
                              out : Tensor) -> Tensor

torch._C._nn.reflection_pad2d(self : Tensor,
                              padding : List[int]) -> Tensor

torch._C._nn.reflection_pad2d(self : Tensor,
                              padding : List[int],
                              out : Tensor) -> Tensor

torch._C._nn.replication_pad1d(self : Tensor,
                               padding : List[int]) -> Tensor

torch._C._nn.replication_pad1d(self : Tensor,
                               padding : List[int],
                               out : Tensor) -> Tensor

torch._C._nn.replication_pad2d(self : Tensor,
                               padding : List[int]) -> Tensor

torch._C._nn.replication_pad2d(self : Tensor,
                               padding : List[int],
                               out : Tensor) -> Tensor

torch._C._nn.replication_pad3d(self : Tensor,
                               padding : List[int],
                               out : Tensor) -> Tensor

torch._C._nn.replication_pad3d(self : Tensor,
                               padding : List[int]) -> Tensor

torch._C._nn.rrelu_with_noise(self : Tensor,
                              noise : Tensor,
                              lower : number=0.125,
                              upper : number=0.3333333333333333,
                              training : bool=False,
                              generator : Optional[Generator],
                              out : Tensor) -> Tensor

torch._C._nn.rrelu_with_noise(self : Tensor,
                              noise : Tensor,
                              lower : number=0.125,
                              upper : number=0.3333333333333333,
                              training : bool=False,
                              generator : Optional[Generator]) -> Tensor

torch._C._nn.rrelu_with_noise_(self : Tensor,
                               noise : Tensor,
                               lower : number=0.125,
                               upper : number=0.3333333333333333,
                               training : bool=False,
                               generator : Optional[Generator]) -> Tensor

torch._C._nn.smooth_l1_loss(self : Tensor,
                            target : Tensor,
                            reduction : int=1) -> Tensor

torch._C._nn.smooth_l1_loss(self : Tensor,
                            target : Tensor,
                            reduction : int=1,
                            out : Tensor) -> Tensor

torch._C._nn.soft_margin_loss(self : Tensor,
                              target : Tensor,
                              reduction : int=1,
                              out : Tensor) -> Tensor

torch._C._nn.soft_margin_loss(self : Tensor,
                              target : Tensor,
                              reduction : int=1) -> Tensor

torch._C._nn.softplus(self : Tensor,
                      beta : number=1,
                      threshold : number=20,
                      out : Tensor) -> Tensor

torch._C._nn.softplus(self : Tensor,
                      beta : number=1,
                      threshold : number=20) -> Tensor

torch._C._nn.softshrink(self : Tensor,
                        lambd : number=0.5) -> Tensor

torch._C._nn.softshrink(self : Tensor,
                        lambd : number=0.5,
                        out : Tensor) -> Tensor

torch._C._nn.thnn_col2im(self : Tensor,
                         output_size : List[int],
                         kernel_size : List[int],
                         dilation : List[int],
                         padding : List[int],
                         stride : List[int]) -> Tensor

torch._C._nn.thnn_conv2d(self : Tensor,
                         weight : Tensor,
                         kernel_size : List[int],
                         bias : Optional[Tensor],
                         stride : List[int]=[1, 1],
                         padding : List[int]=[0, 0]) -> Tensor

torch._C._nn.thnn_conv2d(self : Tensor,
                         weight : Tensor,
                         kernel_size : List[int],
                         bias : Optional[Tensor],
                         stride : List[int]=[1, 1],
                         padding : List[int]=[0, 0],
                         out : Tensor) -> Tensor

torch._C._nn.thnn_conv3d(self : Tensor,
                         weight : Tensor,
                         kernel_size : List[int],
                         bias : Optional[Tensor],
                         stride : List[int]=[1, 1, 1],
                         padding : List[int]=[0, 0, 0],
                         out : Tensor) -> Tensor

torch._C._nn.thnn_conv3d(self : Tensor,
                         weight : Tensor,
                         kernel_size : List[int],
                         bias : Optional[Tensor],
                         stride : List[int]=[1, 1, 1],
                         padding : List[int]=[0, 0, 0]) -> Tensor

torch._C._nn.thnn_conv_depthwise2d(self : Tensor,
                                   weight : Tensor,
                                   kernel_size : List[int],
                                   bias : Optional[Tensor],
                                   stride : List[int]=[1, 1],
                                   padding : List[int]=[0, 0],
                                   dilation : List[int]=[1, 1]) -> Tensor

torch._C._nn.thnn_conv_depthwise2d(self : Tensor,
                                   weight : Tensor,
                                   kernel_size : List[int],
                                   bias : Optional[Tensor],
                                   stride : List[int]=[1, 1],
                                   padding : List[int]=[0, 0],
                                   dilation : List[int]=[1, 1],
                                   out : Tensor) -> Tensor

torch._C._nn.thnn_conv_dilated2d(self : Tensor,
                                 weight : Tensor,
                                 kernel_size : List[int],
                                 bias : Optional[Tensor],
                                 stride : List[int]=[1, 1],
                                 padding : List[int]=[0, 0],
                                 dilation : List[int]=[1, 1]) -> Tensor

torch._C._nn.thnn_conv_dilated2d(self : Tensor,
                                 weight : Tensor,
                                 kernel_size : List[int],
                                 bias : Optional[Tensor],
                                 stride : List[int]=[1, 1],
                                 padding : List[int]=[0, 0],
                                 dilation : List[int]=[1, 1],
                                 out : Tensor) -> Tensor

torch._C._nn.thnn_conv_dilated3d(self : Tensor,
                                 weight : Tensor,
                                 kernel_size : List[int],
                                 bias : Optional[Tensor],
                                 stride : List[int]=[1, 1, 1],
                                 padding : List[int]=[0, 0, 0],
                                 dilation : List[int]=[1, 1, 1],
                                 out : Tensor) -> Tensor

torch._C._nn.thnn_conv_dilated3d(self : Tensor,
                                 weight : Tensor,
                                 kernel_size : List[int],
                                 bias : Optional[Tensor],
                                 stride : List[int]=[1, 1, 1],
                                 padding : List[int]=[0, 0, 0],
                                 dilation : List[int]=[1, 1, 1]) -> Tensor

torch._C._nn.thnn_conv_transpose2d(self : Tensor,
                                   weight : Tensor,
                                   kernel_size : List[int],
                                   bias : Optional[Tensor],
                                   stride : List[int]=[1, 1],
                                   padding : List[int]=[0, 0],
                                   output_padding : List[int]=[0, 0],
                                   dilation : List[int]=[1, 1]) -> Tensor

torch._C._nn.thnn_conv_transpose2d(self : Tensor,
                                   weight : Tensor,
                                   kernel_size : List[int],
                                   bias : Optional[Tensor],
                                   stride : List[int]=[1, 1],
                                   padding : List[int]=[0, 0],
                                   output_padding : List[int]=[0, 0],
                                   dilation : List[int]=[1, 1],
                                   out : Tensor) -> Tensor

torch._C._nn.thnn_conv_transpose3d(self : Tensor,
                                   weight : Tensor,
                                   kernel_size : List[int],
                                   bias : Optional[Tensor],
                                   stride : List[int]=[1, 1, 1],
                                   padding : List[int]=[0, 0, 0],
                                   output_padding : List[int]=[0, 0, 0],
                                   dilation : List[int]=[1, 1, 1],
                                   out : Tensor) -> Tensor

torch._C._nn.thnn_conv_transpose3d(self : Tensor,
                                   weight : Tensor,
                                   kernel_size : List[int],
                                   bias : Optional[Tensor],
                                   stride : List[int]=[1, 1, 1],
                                   padding : List[int]=[0, 0, 0],
                                   output_padding : List[int]=[0, 0, 0],
                                   dilation : List[int]=[1, 1, 1]) -> Tensor

torch._C._nn.thnn_im2col(self : Tensor,
                         kernel_size : List[int],
                         dilation : List[int],
                         padding : List[int],
                         stride : List[int]) -> Tensor

torch._C._nn.upsample_bicubic2d(self : Tensor,
                                output_size : List[int],
                                align_corners : bool) -> Tensor

torch._C._nn.upsample_bicubic2d(self : Tensor,
                                output_size : List[int],
                                align_corners : bool,
                                out : Tensor) -> Tensor

torch._C._nn.upsample_bilinear2d(self : Tensor,
                                 output_size : List[int],
                                 align_corners : bool) -> Tensor

torch._C._nn.upsample_bilinear2d(self : Tensor,
                                 output_size : List[int],
                                 align_corners : bool,
                                 out : Tensor) -> Tensor

torch._C._nn.upsample_linear1d(self : Tensor,
                               output_size : List[int],
                               align_corners : bool) -> Tensor

torch._C._nn.upsample_linear1d(self : Tensor,
                               output_size : List[int],
                               align_corners : bool,
                               out : Tensor) -> Tensor

torch._C._nn.upsample_nearest1d(self : Tensor,
                                output_size : List[int]) -> Tensor

torch._C._nn.upsample_nearest1d(self : Tensor,
                                output_size : List[int],
                                out : Tensor) -> Tensor

torch._C._nn.upsample_nearest2d(self : Tensor,
                                output_size : List[int]) -> Tensor

torch._C._nn.upsample_nearest2d(self : Tensor,
                                output_size : List[int],
                                out : Tensor) -> Tensor

torch._C._nn.upsample_nearest3d(self : Tensor,
                                output_size : List[int],
                                out : Tensor) -> Tensor

torch._C._nn.upsample_nearest3d(self : Tensor,
                                output_size : List[int]) -> Tensor

torch._C._nn.upsample_trilinear3d(self : Tensor,
                                  output_size : List[int],
                                  align_corners : bool,
                                  out : Tensor) -> Tensor

torch._C._nn.upsample_trilinear3d(self : Tensor,
                                  output_size : List[int],
                                  align_corners : bool) -> Tensor

torch.nn.functional.adaptive_avg_pool2d(input : Tensor,
                                        output_size : List[int]) -> Tensor

torch.nn.functional.adaptive_avg_pool3d(input : Tensor,
                                        output_size : List[int]) -> Tensor

torch.nn.functional.adaptive_max_pool1d_with_indices(input : Tensor,
                                                     output_size : List[int],
                                                     return_indices : bool=False) -> Tuple[Tensor, Tensor]

torch.nn.functional.adaptive_max_pool2d_with_indices(input : Tensor,
                                                     output_size : List[int],
                                                     return_indices : bool=False) -> Tuple[Tensor, Tensor]

torch.nn.functional.adaptive_max_pool3d_with_indices(input : Tensor,
                                                     output_size : List[int],
                                                     return_indices : bool=False) -> Tuple[Tensor, Tensor]

torch.nn.functional.affine_grid(theta : Tensor,
                                size : List[int]) -> Tensor

torch.nn.functional.alpha_dropout(input : Tensor,
                                  p : float=0.5,
                                  training : bool=False,
                                  inplace : bool=False) -> Tensor

torch.nn.functional.batch_norm(input : Tensor,
                               running_mean : Optional[Tensor],
                               running_var : Optional[Tensor],
                               weight : Optional[Tensor],
                               bias : Optional[Tensor],
                               training : bool=False,
                               momentum : float=0.1,
                               eps : float=1e-05) -> Tensor

torch.nn.functional.bilinear(input1 : Tensor,
                             input2 : Tensor,
                             weight : Tensor,
                             bias : Optional[Tensor]) -> Tensor

torch.nn.functional.binary_cross_entropy(input : Tensor,
                                         target : Tensor,
                                         weight : Optional[Tensor],
                                         size_average : Optional[bool],
                                         reduce : Optional[bool],
                                         reduction : str=mean) -> Tensor

torch.nn.functional.binary_cross_entropy_with_logits(input : Tensor,
                                                     target : Tensor,
                                                     weight : Optional[Tensor],
                                                     size_average : Optional[bool],
                                                     reduce : Optional[bool],
                                                     reduction : str=mean,
                                                     pos_weight : Optional[Tensor]) -> Tensor

torch.nn.functional.celu(input : Tensor,
                         alpha : float=1.0,
                         inplace : bool=False) -> Tensor

torch.nn.functional.cosine_embedding_loss(input1 : Tensor,
                                          input2 : Tensor,
                                          target : Tensor,
                                          margin : float=0.0,
                                          size_average : Optional[bool],
                                          reduce : Optional[bool],
                                          reduction : str=mean) -> Tensor

torch.nn.functional.cross_entropy(input : Tensor,
                                  target : Tensor,
                                  weight : Optional[Tensor],
                                  size_average : Optional[bool],
                                  ignore_index : int=-100,
                                  reduce : Optional[bool],
                                  reduction : str=mean) -> Tensor

torch.nn.functional.ctc_loss(log_probs : Tensor,
                             targets : Tensor,
                             input_lengths : Tensor,
                             target_lengths : Tensor,
                             blank : int=0,
                             reduction : str=mean,
                             zero_infinity : bool=False) -> Tensor

torch.nn.functional.dropout(input : Tensor,
                            p : float=0.5,
                            training : bool=True,
                            inplace : bool=False) -> Tensor

torch.nn.functional.dropout2d(input : Tensor,
                              p : float=0.5,
                              training : bool=True,
                              inplace : bool=False) -> Tensor

torch.nn.functional.dropout3d(input : Tensor,
                              p : float=0.5,
                              training : bool=True,
                              inplace : bool=False) -> Tensor

torch.nn.functional.elu(input : Tensor,
                        alpha : float=1.0,
                        inplace : bool=False) -> Tensor

torch.nn.functional.embedding(input : Tensor,
                              weight : Tensor,
                              padding_idx : Optional[int],
                              max_norm : Optional[float],
                              norm_type : float=2.0,
                              scale_grad_by_freq : bool=False,
                              sparse : bool=False) -> Tensor

torch.nn.functional.embedding_bag(input : Tensor,
                                  weight : Tensor,
                                  offsets : Optional[Tensor],
                                  max_norm : Optional[float],
                                  norm_type : float=2.0,
                                  scale_grad_by_freq : bool=False,
                                  mode : str=mean,
                                  sparse : bool=False) -> Tensor

torch.nn.functional.feature_alpha_dropout(input : Tensor,
                                          p : float=0.5,
                                          training : bool=False,
                                          inplace : bool=False) -> Tensor

torch.nn.functional.fold(input : Tensor,
                         output_size : List[int],
                         kernel_size : List[int],
                         dilation : List[int]=1,
                         padding : List[int]=0,
                         stride : List[int]=1) -> Tensor

torch.nn.functional.fractional_max_pool2d_with_indices(input : Tensor,
                                                       kernel_size : List[int],
                                                       output_size : Optional[List[int]],
                                                       output_ratio : Optional[List[float]],
                                                       return_indices : bool=False,
                                                       _random_samples : Optional[Tensor]) -> Tuple[Tensor, Tensor]

torch.nn.functional.fractional_max_pool3d_with_indices(input : Tensor,
                                                       kernel_size : List[int],
                                                       output_size : Optional[List[int]],
                                                       output_ratio : Optional[List[float]],
                                                       return_indices : bool=False,
                                                       _random_samples : Optional[Tensor]) -> Tuple[Tensor, Tensor]

torch.nn.functional.glu(input : Tensor,
                        dim : int=-1) -> Tensor

torch.nn.functional.grid_sample(input : Tensor,
                                grid : Tensor,
                                mode : str=bilinear,
                                padding_mode : str=zeros) -> Tensor

torch.nn.functional.group_norm(input : Tensor,
                               num_groups : int,
                               weight : Optional[Tensor],
                               bias : Optional[Tensor],
                               eps : float=1e-05) -> Tensor

torch.nn.functional.gumbel_softmax(logits : Tensor,
                                   tau : float=1.0,
                                   hard : bool=False,
                                   eps : float=1e-10,
                                   dim : int=-1) -> Tensor

torch.nn.functional.hardshrink(input : Tensor,
                               lambd : float=0.5) -> Tensor

torch.nn.functional.hardtanh(input : Tensor,
                             min_val : float=-1.0,
                             max_val : float=1.0,
                             inplace : bool=False) -> Tensor

torch.nn.functional.hinge_embedding_loss(input : Tensor,
                                         target : Tensor,
                                         margin : float=1.0,
                                         size_average : Optional[bool],
                                         reduce : Optional[bool],
                                         reduction : str=mean) -> Tensor

torch.nn.functional.instance_norm(input : Tensor,
                                  running_mean : Optional[Tensor],
                                  running_var : Optional[Tensor],
                                  weight : Optional[Tensor],
                                  bias : Optional[Tensor],
                                  use_input_stats : bool=True,
                                  momentum : float=0.1,
                                  eps : float=1e-05) -> Tensor

torch.nn.functional.kl_div(input : Tensor,
                           target : Tensor,
                           size_average : Optional[bool],
                           reduce : Optional[bool],
                           reduction : str=mean) -> Tensor

torch.nn.functional.l1_loss(input : Tensor,
                            target : Tensor,
                            size_average : Optional[bool],
                            reduce : Optional[bool],
                            reduction : str=mean) -> Tensor

torch.nn.functional.layer_norm(input : Tensor,
                               normalized_shape : List[int],
                               weight : Optional[Tensor],
                               bias : Optional[Tensor],
                               eps : float=1e-05) -> Tensor

torch.nn.functional.leaky_relu(input : Tensor,
                               negative_slope : float=0.01,
                               inplace : bool=False) -> Tensor

torch.nn.functional.linear(input : Tensor,
                           weight : Tensor,
                           bias : Optional[Tensor]) -> Tensor

torch.nn.functional.local_response_norm(input : Tensor,
                                        size : int,
                                        alpha : float=0.0001,
                                        beta : float=0.75,
                                        k : float=1.0) -> Tensor

torch.nn.functional.log_softmax(input : Tensor,
                                dim : Optional[int],
                                _stacklevel : int=3,
                                dtype : Optional[int]) -> Tensor

torch.nn.functional.lp_pool1d(input : Tensor,
                              norm_type : float,
                              kernel_size : int,
                              stride : Optional[List[int]],
                              ceil_mode : bool=False) -> Tensor

torch.nn.functional.lp_pool2d(input : Tensor,
                              norm_type : float,
                              kernel_size : int,
                              stride : Optional[List[int]],
                              ceil_mode : bool=False) -> Tensor

torch.nn.functional.margin_ranking_loss(input1 : Tensor,
                                        input2 : Tensor,
                                        target : Tensor,
                                        margin : float=0.0,
                                        size_average : Optional[bool],
                                        reduce : Optional[bool],
                                        reduction : str=mean) -> Tensor

torch.nn.functional.max_pool1d_with_indices(input : Tensor,
                                            kernel_size : List[int],
                                            stride : Optional[List[int]],
                                            padding : List[int]=0,
                                            dilation : List[int]=1,
                                            ceil_mode : bool=False,
                                            return_indices : bool=False) -> Tuple[Tensor, Tensor]

torch.nn.functional.max_pool2d_with_indices(input : Tensor,
                                            kernel_size : List[int],
                                            stride : Optional[List[int]],
                                            padding : List[int]=0,
                                            dilation : List[int]=1,
                                            ceil_mode : bool=False,
                                            return_indices : bool=False) -> Tuple[Tensor, Tensor]

torch.nn.functional.max_pool3d_with_indices(input : Tensor,
                                            kernel_size : List[int],
                                            stride : Optional[List[int]],
                                            padding : List[int]=0,
                                            dilation : List[int]=1,
                                            ceil_mode : bool=False,
                                            return_indices : bool=False) -> Tuple[Tensor, Tensor]

torch.nn.functional.max_unpool1d(input : Tensor,
                                 indices : Tensor,
                                 kernel_size : List[int],
                                 stride : Optional[List[int]],
                                 padding : List[int]=0,
                                 output_size : Optional[List[int]]) -> Tensor

torch.nn.functional.max_unpool2d(input : Tensor,
                                 indices : Tensor,
                                 kernel_size : List[int],
                                 stride : Optional[List[int]],
                                 padding : List[int]=0,
                                 output_size : Optional[List[int]]) -> Tensor

torch.nn.functional.max_unpool3d(input : Tensor,
                                 indices : Tensor,
                                 kernel_size : List[int],
                                 stride : Optional[List[int]],
                                 padding : List[int]=0,
                                 output_size : Optional[List[int]]) -> Tensor

torch.nn.functional.mse_loss(input : Tensor,
                             target : Tensor,
                             size_average : Optional[bool],
                             reduce : Optional[bool],
                             reduction : str=mean) -> Tensor

torch.nn.functional.multi_margin_loss(input : Tensor,
                                      target : Tensor,
                                      p : int=1,
                                      margin : float=1.0,
                                      weight : Optional[Tensor],
                                      size_average : Optional[bool],
                                      reduce : Optional[bool],
                                      reduction : str=mean) -> Tensor

torch.nn.functional.multilabel_margin_loss(input : Tensor,
                                           target : Tensor,
                                           size_average : Optional[bool],
                                           reduce : Optional[bool],
                                           reduction : str=mean) -> Tensor

torch.nn.functional.multilabel_soft_margin_loss(input : Tensor,
                                                target : Tensor,
                                                weight : Optional[Tensor],
                                                size_average : Optional[bool],
                                                reduce : Optional[bool],
                                                reduction : str=mean) -> Tensor

torch.nn.functional.nll_loss(input : Tensor,
                             target : Tensor,
                             weight : Optional[Tensor],
                             size_average : Optional[bool],
                             ignore_index : int=-100,
                             reduce : Optional[bool],
                             reduction : str=mean) -> Tensor

torch.nn.functional.normalize(input : Tensor,
                              p : float=2.0,
                              dim : int=1,
                              eps : float=1e-12,
                              out : Optional[Tensor]) -> Tensor

torch.nn.functional.pad(input : Tensor,
                        pad : List[int],
                        mode : str=constant,
                        value : float=0.0) -> Tensor

torch.nn.functional.pad_circular(input : Tensor,
                                 padding : List[int]) -> Tensor

torch.nn.functional.pairwise_distance(x1 : Tensor,
                                      x2 : Tensor,
                                      p : float=2.0,
                                      eps : float=1e-06,
                                      keepdim : bool=False) -> Tensor

torch.nn.functional.poisson_nll_loss(input : Tensor,
                                     target : Tensor,
                                     log_input : bool=True,
                                     full : bool=False,
                                     size_average : Optional[bool],
                                     eps : float=1e-08,
                                     reduce : Optional[bool],
                                     reduction : str=mean) -> Tensor

torch.nn.functional.prelu(input : Tensor,
                          weight : Tensor) -> Tensor

torch.nn.functional.relu(input : Tensor,
                         inplace : bool=False) -> Tensor

torch.nn.functional.relu6(input : Tensor,
                          inplace : bool=False) -> Tensor

torch.nn.functional.rrelu(input : Tensor,
                          lower : float=0.125,
                          upper : float=0.3333333333333333,
                          training : bool=False,
                          inplace : bool=False) -> Tensor

torch.nn.functional.selu(input : Tensor,
                         inplace : bool=False) -> Tensor

torch.nn.functional.sigmoid(input : Tensor) -> Tensor

torch.nn.functional.smooth_l1_loss(input : Tensor,
                                   target : Tensor,
                                   size_average : Optional[bool],
                                   reduce : Optional[bool],
                                   reduction : str=mean) -> Tensor

torch.nn.functional.soft_margin_loss(input : Tensor,
                                     target : Tensor,
                                     size_average : Optional[bool],
                                     reduce : Optional[bool],
                                     reduction : str=mean) -> Tensor

torch.nn.functional.softmax(input : Tensor,
                            dim : Optional[int],
                            _stacklevel : int=3,
                            dtype : Optional[int]) -> Tensor

torch.nn.functional.softmin(input : Tensor,
                            dim : Optional[int],
                            _stacklevel : int=3,
                            dtype : Optional[int]) -> Tensor

torch.nn.functional.softsign(input : Tensor) -> Tensor

torch.nn.functional.tanh(input : Tensor) -> Tensor

torch.nn.functional.tanhshrink(input : Tensor) -> Tensor

torch.nn.functional.threshold(input : Tensor,
                              threshold : float,
                              value : float,
                              inplace : bool=False) -> Tensor

torch.nn.functional.triplet_margin_loss(anchor : Tensor,
                                        positive : Tensor,
                                        negative : Tensor,
                                        margin : float=1.0,
                                        p : float=2.0,
                                        eps : float=1e-06,
                                        swap : bool=False,
                                        size_average : Optional[bool],
                                        reduce : Optional[bool],
                                        reduction : str=mean) -> Tensor

torch.nn.functional.unfold(input : Tensor,
                           kernel_size : List[int],
                           dilation : List[int]=1,
                           padding : List[int]=0,
                           stride : List[int]=1) -> Tensor

Supported Methods ¶

Tensor.__and__(other : Tensor) -> Tensor

Tensor.__and__(other : number) -> Tensor

Tensor.__iand__(other : Tensor) -> Tensor

Tensor.__iand__(other : number) -> Tensor

Tensor.__ilshift__(other : Tensor) -> Tensor

Tensor.__ilshift__(other : number) -> Tensor

Tensor.__ior__(other : Tensor) -> Tensor

Tensor.__ior__(other : number) -> Tensor

Tensor.__irshift__(other : Tensor) -> Tensor

Tensor.__irshift__(other : number) -> Tensor

Tensor.__ixor__(other : Tensor) -> Tensor

Tensor.__ixor__(other : number) -> Tensor

Tensor.__lshift__(other : Tensor) -> Tensor

Tensor.__lshift__(other : number) -> Tensor

Tensor.__or__(other : Tensor) -> Tensor

Tensor.__or__(other : number) -> Tensor

Tensor.__rshift__(other : Tensor) -> Tensor

Tensor.__rshift__(other : number) -> Tensor

Tensor.__xor__(other : Tensor) -> Tensor

Tensor.__xor__(other : number) -> Tensor

Tensor.abs() -> Tensor

Tensor.abs(out : Tensor) -> Tensor

Tensor.abs_() -> Tensor

Tensor.acos(out : Tensor) -> Tensor

Tensor.acos() -> Tensor

Tensor.acos_() -> Tensor

Tensor.add(other : Tensor,
           alpha : number=1) -> Tensor

Tensor.add(other : number,
           alpha : number=1) -> Tensor

Tensor.add(other : Tensor,
           alpha : number=1,
           out : Tensor) -> Tensor

Tensor.add_(other : Tensor,
            alpha : number=1) -> Tensor

Tensor.add_(other : number,
            alpha : number=1) -> Tensor

Tensor.addbmm(batch1 : Tensor,
              batch2 : Tensor,
              beta : number=1,
              alpha : number=1) -> Tensor

Tensor.addbmm(batch1 : Tensor,
              batch2 : Tensor,
              beta : number=1,
              alpha : number=1,
              out : Tensor) -> Tensor

Tensor.addbmm_(batch1 : Tensor,
               batch2 : Tensor,
               beta : number=1,
               alpha : number=1) -> Tensor

Tensor.addcdiv(tensor1 : Tensor,
               tensor2 : Tensor,
               value : number=1,
               out : Tensor) -> Tensor

Tensor.addcdiv(tensor1 : Tensor,
               tensor2 : Tensor,
               value : number=1) -> Tensor

Tensor.addcdiv_(tensor1 : Tensor,
                tensor2 : Tensor,
                value : number=1) -> Tensor

Tensor.addcmul(tensor1 : Tensor,
               tensor2 : Tensor,
               value : number=1) -> Tensor

Tensor.addcmul(tensor1 : Tensor,
               tensor2 : Tensor,
               value : number=1,
               out : Tensor) -> Tensor

Tensor.addcmul_(tensor1 : Tensor,
                tensor2 : Tensor,
                value : number=1) -> Tensor

Tensor.addmm(mat1 : Tensor,
             mat2 : Tensor,
             beta : number=1,
             alpha : number=1,
             out : Tensor) -> Tensor

Tensor.addmm(mat1 : Tensor,
             mat2 : Tensor,
             beta : number=1,
             alpha : number=1) -> Tensor

Tensor.addmm_(mat1 : Tensor,
              mat2 : Tensor,
              beta : number=1,
              alpha : number=1) -> Tensor

Tensor.addmv(mat : Tensor,
             vec : Tensor,
             beta : number=1,
             alpha : number=1,
             out : Tensor) -> Tensor

Tensor.addmv(mat : Tensor,
             vec : Tensor,
             beta : number=1,
             alpha : number=1) -> Tensor

Tensor.addmv_(mat : Tensor,
              vec : Tensor,
              beta : number=1,
              alpha : number=1) -> Tensor

Tensor.addr(vec1 : Tensor,
            vec2 : Tensor,
            beta : number=1,
            alpha : number=1) -> Tensor

Tensor.addr(vec1 : Tensor,
            vec2 : Tensor,
            beta : number=1,
            alpha : number=1,
            out : Tensor) -> Tensor

Tensor.addr_(vec1 : Tensor,
             vec2 : Tensor,
             beta : number=1,
             alpha : number=1) -> Tensor

Tensor.all() -> Tensor

Tensor.all(dim : int,
           keepdim : bool=False) -> Tensor

Tensor.all(dim : int,
           keepdim : bool=False,
           out : Tensor) -> Tensor

Tensor.allclose(other : Tensor,
                rtol : float=1e-05,
                atol : float=1e-08,
                equal_nan : bool=False) -> bool

Tensor.any() -> Tensor

Tensor.any(dim : int,
           keepdim : bool=False) -> Tensor

Tensor.any(dim : int,
           keepdim : bool=False,
           out : Tensor) -> Tensor

Tensor.argmax() -> Tensor

Tensor.argmax(dim : int,
              keepdim : bool=False) -> Tensor

Tensor.argmin() -> Tensor

Tensor.argmin(dim : int,
              keepdim : bool=False) -> Tensor

Tensor.argsort(dim : int=-1,
               descending : bool=False) -> Tensor

Tensor.as_strided(size : List[int],
                  stride : List[int],
                  storage_offset : Optional[int]) -> Tensor

Tensor.as_strided_(size : List[int],
                   stride : List[int],
                   storage_offset : Optional[int]) -> Tensor

Tensor.asin() -> Tensor

Tensor.asin(out : Tensor) -> Tensor

Tensor.asin_() -> Tensor

Tensor.atan() -> Tensor

Tensor.atan(out : Tensor) -> Tensor

Tensor.atan2(other : Tensor,
             out : Tensor) -> Tensor

Tensor.atan2(other : Tensor) -> Tensor

Tensor.atan2_(other : Tensor) -> Tensor

Tensor.atan_() -> Tensor

Tensor.baddbmm(batch1 : Tensor,
               batch2 : Tensor,
               beta : number=1,
               alpha : number=1) -> Tensor

Tensor.baddbmm(batch1 : Tensor,
               batch2 : Tensor,
               beta : number=1,
               alpha : number=1,
               out : Tensor) -> Tensor

Tensor.baddbmm_(batch1 : Tensor,
                batch2 : Tensor,
                beta : number=1,
                alpha : number=1) -> Tensor

Tensor.bernoulli(generator : Optional[Generator]) -> Tensor

Tensor.bernoulli(p : float,
                 generator : Optional[Generator]) -> Tensor

Tensor.bernoulli(generator : Optional[Generator],
                 out : Tensor) -> Tensor

Tensor.bernoulli_(p : Tensor,
                  generator : Optional[Generator]) -> Tensor

Tensor.bernoulli_(p : float=0.5,
                  generator : Optional[Generator]) -> Tensor

Tensor.bincount(weights : Optional[Tensor],
                minlength : int=0) -> Tensor

Tensor.bmm(mat2 : Tensor) -> Tensor

Tensor.bmm(mat2 : Tensor,
           out : Tensor) -> Tensor

Tensor.btrifact(pivot : bool=True) -> Tuple[Tensor, Tensor]

Tensor.btrifact_with_info(pivot : bool=True) -> Tuple[Tensor, Tensor, Tensor]

Tensor.btrisolve(LU_data : Tensor,
                 LU_pivots : Tensor,
                 out : Tensor) -> Tensor

Tensor.btrisolve(LU_data : Tensor,
                 LU_pivots : Tensor) -> Tensor

Tensor.cauchy_(median : float=0.0,
               sigma : float=1.0,
               generator : Optional[Generator]) -> Tensor

Tensor.ceil(out : Tensor) -> Tensor

Tensor.ceil() -> Tensor

Tensor.ceil_() -> Tensor

Tensor.cholesky(upper : bool=False,
                out : Tensor) -> Tensor

Tensor.cholesky(upper : bool=False) -> Tensor

Tensor.cholesky_solve(input2 : Tensor,
                      upper : bool=False,
                      out : Tensor) -> Tensor

Tensor.cholesky_solve(input2 : Tensor,
                      upper : bool=False) -> Tensor

Tensor.chunk(chunks : int,
             dim : int=0) -> List[Tensor]

Tensor.clamp(min : Optional[number],
             max : Optional[number]) -> Tensor

Tensor.clamp(min : Optional[number],
             max : Optional[number],
             out : Tensor) -> Tensor

Tensor.clamp_(min : Optional[number],
              max : Optional[number]) -> Tensor

Tensor.clamp_max(max : number) -> Tensor

Tensor.clamp_max(max : number,
                 out : Tensor) -> Tensor

Tensor.clamp_max_(max : number) -> Tensor

Tensor.clamp_min(min : number,
                 out : Tensor) -> Tensor

Tensor.clamp_min(min : number) -> Tensor

Tensor.clamp_min_(min : number) -> Tensor

Tensor.clone() -> Tensor

Tensor.coalesce() -> Tensor

Tensor.contiguous() -> Tensor

Tensor.copy_(other : Tensor) -> Tensor

Tensor.copy_(other : int) -> Tensor

Tensor.copy_(other : float) -> Tensor

Tensor.cos() -> Tensor

Tensor.cos(out : Tensor) -> Tensor

Tensor.cos_() -> Tensor

Tensor.cosh() -> Tensor

Tensor.cosh(out : Tensor) -> Tensor

Tensor.cosh_() -> Tensor

Tensor.cpu() -> Tensor

Tensor.cross(other : Tensor,
             dim : int=-1,
             out : Tensor) -> Tensor

Tensor.cross(other : Tensor,
             dim : int=-1) -> Tensor

Tensor.cuda() -> Tensor

Tensor.cumprod(dim : int) -> Tensor

Tensor.cumprod(dim : int,
               dtype : int) -> Tensor

Tensor.cumprod(dim : int,
               out : Tensor) -> Tensor

Tensor.cumprod(dim : int,
               dtype : int,
               out : Tensor) -> Tensor

Tensor.cumsum(dim : int) -> Tensor

Tensor.cumsum(dim : int,
              dtype : int) -> Tensor

Tensor.cumsum(dim : int,
              out : Tensor) -> Tensor

Tensor.cumsum(dim : int,
              dtype : int,
              out : Tensor) -> Tensor

Tensor.dense_dim() -> int

Tensor.det() -> Tensor

Tensor.detach() -> Tensor

Tensor.detach_() -> Tensor

Tensor.diag(diagonal : int=0) -> Tensor

Tensor.diag(diagonal : int=0,
            out : Tensor) -> Tensor

Tensor.diag_embed(offset : int=0,
                  dim1 : int=-2,
                  dim2 : int=-1) -> Tensor

Tensor.diagflat(offset : int=0) -> Tensor

Tensor.diagonal(offset : int=0,
                dim1 : int=0,
                dim2 : int=1) -> Tensor

Tensor.digamma() -> Tensor

Tensor.digamma(out : Tensor) -> Tensor

Tensor.digamma_() -> Tensor

Tensor.dim() -> int

Tensor.dist(other : Tensor,
            p : number=2) -> Tensor

Tensor.div(other : Tensor,
           out : Tensor) -> Tensor

Tensor.div(other : Tensor) -> Tensor

Tensor.div(other : number) -> Tensor

Tensor.div_(other : Tensor) -> Tensor

Tensor.div_(other : number) -> Tensor

Tensor.dot(tensor : Tensor) -> Tensor

Tensor.dot(tensor : Tensor,
           out : Tensor) -> Tensor

Tensor.eig(eigenvectors : bool=False) -> Tuple[Tensor, Tensor]

Tensor.eq(other : Tensor) -> Tensor

Tensor.eq(other : number) -> Tensor

Tensor.eq(other : Tensor,
          out : Tensor) -> Tensor

Tensor.eq(other : number,
          out : Tensor) -> Tensor

Tensor.eq_(other : Tensor) -> Tensor

Tensor.eq_(other : number) -> Tensor

Tensor.equal(other : Tensor) -> bool

Tensor.erf(out : Tensor) -> Tensor

Tensor.erf() -> Tensor

Tensor.erf_() -> Tensor

Tensor.erfc(out : Tensor) -> Tensor

Tensor.erfc() -> Tensor

Tensor.erfc_() -> Tensor

Tensor.erfinv(out : Tensor) -> Tensor

Tensor.erfinv() -> Tensor

Tensor.erfinv_() -> Tensor

Tensor.exp() -> Tensor

Tensor.exp(out : Tensor) -> Tensor

Tensor.exp_() -> Tensor

Tensor.expand(size : List[int],
              implicit : bool=False) -> Tensor

Tensor.expand_as(other : Tensor) -> Tensor

Tensor.expm1(out : Tensor) -> Tensor

Tensor.expm1() -> Tensor

Tensor.expm1_() -> Tensor

Tensor.exponential_(lambd : float=1.0,
                    generator : Optional[Generator]) -> Tensor

Tensor.fft(signal_ndim : int,
           normalized : bool=False) -> Tensor

Tensor.fill_(value : Tensor) -> Tensor

Tensor.fill_(value : number) -> Tensor

Tensor.flatten(start_dim : int=0,
               end_dim : int=-1) -> Tensor

Tensor.flip(dims : List[int]) -> Tensor

Tensor.floor() -> Tensor

Tensor.floor(out : Tensor) -> Tensor

Tensor.floor_() -> Tensor

Tensor.fmod(other : Tensor,
            out : Tensor) -> Tensor

Tensor.fmod(other : number,
            out : Tensor) -> Tensor

Tensor.fmod(other : Tensor) -> Tensor

Tensor.fmod(other : number) -> Tensor

Tensor.fmod_(other : Tensor) -> Tensor

Tensor.fmod_(other : number) -> Tensor

Tensor.frac() -> Tensor

Tensor.frac(out : Tensor) -> Tensor

Tensor.frac_() -> Tensor

Tensor.gather(dim : int,
              index : Tensor,
              sparse_grad : bool=False,
              out : Tensor) -> Tensor

Tensor.gather(dim : int,
              index : Tensor,
              sparse_grad : bool=False) -> Tensor

Tensor.ge(other : Tensor) -> Tensor

Tensor.ge(other : number) -> Tensor

Tensor.ge(other : Tensor,
          out : Tensor) -> Tensor

Tensor.ge(other : number,
          out : Tensor) -> Tensor

Tensor.ge_(other : Tensor) -> Tensor

Tensor.ge_(other : number) -> Tensor

Tensor.gels(A : Tensor) -> Tuple[Tensor, Tensor]

Tensor.geometric_(p : float,
                  generator : Optional[Generator]) -> Tensor

Tensor.geqrf() -> Tuple[Tensor, Tensor]

Tensor.ger(vec2 : Tensor) -> Tensor

Tensor.ger(vec2 : Tensor,
           out : Tensor) -> Tensor

Tensor.get_device() -> int

Tensor.get_device() -> int

Tensor.get_device() -> int

Tensor.gt(other : Tensor) -> Tensor

Tensor.gt(other : number) -> Tensor

Tensor.gt(other : Tensor,
          out : Tensor) -> Tensor

Tensor.gt(other : number,
          out : Tensor) -> Tensor

Tensor.gt_(other : Tensor) -> Tensor

Tensor.gt_(other : number) -> Tensor

Tensor.hardshrink(lambd : number=0.5) -> Tensor

Tensor.histc(bins : int=100,
             min : number=0,
             max : number=0,
             out : Tensor) -> Tensor

Tensor.histc(bins : int=100,
             min : number=0,
             max : number=0) -> Tensor

Tensor.ifft(signal_ndim : int,
            normalized : bool=False) -> Tensor

Tensor.index_add(dim : int,
                 index : Tensor,
                 source : Tensor) -> Tensor

Tensor.index_add_(dim : int,
                  index : Tensor,
                  source : Tensor) -> Tensor

Tensor.index_copy(dim : int,
                  index : Tensor,
                  source : Tensor) -> Tensor

Tensor.index_copy_(dim : int,
                   index : Tensor,
                   source : Tensor) -> Tensor

Tensor.index_fill(dim : int,
                  index : Tensor,
                  value : Tensor) -> Tensor

Tensor.index_fill(dim : int,
                  index : Tensor,
                  value : number) -> Tensor

Tensor.index_fill_(dim : int,
                   index : Tensor,
                   value : Tensor) -> Tensor

Tensor.index_fill_(dim : int,
                   index : Tensor,
                   value : number) -> Tensor

Tensor.index_put(indices : List[Optional[Tensor]],
                 values : Tensor,
                 accumulate : bool=False) -> Tensor

Tensor.index_put(indices : List[Tensor],
                 values : Tensor,
                 accumulate : bool=False) -> Tensor

Tensor.index_put_(indices : List[Optional[Tensor]],
                  values : Tensor,
                  accumulate : bool=False) -> Tensor

Tensor.index_put_(indices : List[Tensor],
                  values : Tensor,
                  accumulate : bool=False) -> Tensor

Tensor.index_select(dim : int,
                    index : Tensor,
                    out : Tensor) -> Tensor

Tensor.index_select(dim : int,
                    index : Tensor) -> Tensor

Tensor.indices() -> Tensor

Tensor.inverse(out : Tensor) -> Tensor

Tensor.inverse() -> Tensor

Tensor.irfft(signal_ndim : int,
             normalized : bool=False,
             onesided : bool=True,
             signal_sizes : List[int]=[]) -> Tensor

Tensor.is_coalesced() -> bool

Tensor.is_complex() -> bool

Tensor.is_distributed() -> bool

Tensor.is_floating_point() -> bool

Tensor.is_nonzero() -> bool

Tensor.is_same_size(other : Tensor) -> bool

Tensor.is_set_to(tensor : Tensor) -> bool

Tensor.is_signed() -> bool

Tensor.isclose(other : Tensor,
               rtol : float=1e-05,
               atol : float=1e-08,
               equal_nan : bool=False) -> Tensor

Tensor.item() -> number

Tensor.kthvalue(k : int,
                dim : int=-1,
                keepdim : bool=False) -> Tuple[Tensor, Tensor]

Tensor.le(other : Tensor,
          out : Tensor) -> Tensor

Tensor.le(other : number,
          out : Tensor) -> Tensor

Tensor.le(other : Tensor) -> Tensor

Tensor.le(other : number) -> Tensor

Tensor.le_(other : Tensor) -> Tensor

Tensor.le_(other : number) -> Tensor

Tensor.lerp(end : Tensor,
            weight : Tensor) -> Tensor

Tensor.lerp(end : Tensor,
            weight : number) -> Tensor

Tensor.lerp(end : Tensor,
            weight : Tensor,
            out : Tensor) -> Tensor

Tensor.lerp(end : Tensor,
            weight : number,
            out : Tensor) -> Tensor

Tensor.lerp_(end : Tensor,
             weight : Tensor) -> Tensor

Tensor.lerp_(end : Tensor,
             weight : number) -> Tensor

Tensor.lgamma(out : Tensor) -> Tensor

Tensor.lgamma() -> Tensor

Tensor.lgamma_() -> Tensor

Tensor.log() -> Tensor

Tensor.log(out : Tensor) -> Tensor

Tensor.log10(out : Tensor) -> Tensor

Tensor.log10() -> Tensor

Tensor.log10_() -> Tensor

Tensor.log1p() -> Tensor

Tensor.log1p(out : Tensor) -> Tensor

Tensor.log1p_() -> Tensor

Tensor.log2() -> Tensor

Tensor.log2(out : Tensor) -> Tensor

Tensor.log2_() -> Tensor

Tensor.log_() -> Tensor

Tensor.log_normal_(mean : float=1.0,
                   std : float=2.0,
                   generator : Optional[Generator]) -> Tensor

Tensor.log_softmax(dim : int) -> Tensor

Tensor.log_softmax(dim : int,
                   dtype : int) -> Tensor

Tensor.logdet() -> Tensor

Tensor.logsumexp(dim : List[int],
                 keepdim : bool=False) -> Tensor

Tensor.logsumexp(dim : List[int],
                 keepdim : bool=False,
                 out : Tensor) -> Tensor

Tensor.lt(other : Tensor,
          out : Tensor) -> Tensor

Tensor.lt(other : number,
          out : Tensor) -> Tensor

Tensor.lt(other : Tensor) -> Tensor

Tensor.lt(other : number) -> Tensor

Tensor.lt_(other : Tensor) -> Tensor

Tensor.lt_(other : number) -> Tensor

Tensor.masked_fill(mask : Tensor,
                   value : Tensor) -> Tensor

Tensor.masked_fill(mask : Tensor,
                   value : number) -> Tensor

Tensor.masked_fill_(mask : Tensor,
                    value : Tensor) -> Tensor

Tensor.masked_fill_(mask : Tensor,
                    value : number) -> Tensor

Tensor.masked_scatter(mask : Tensor,
                      source : Tensor) -> Tensor

Tensor.masked_scatter_(mask : Tensor,
                       source : Tensor) -> Tensor

Tensor.masked_select(mask : Tensor,
                     out : Tensor) -> Tensor

Tensor.masked_select(mask : Tensor) -> Tensor

Tensor.matmul(other : Tensor,
              out : Tensor) -> Tensor

Tensor.matmul(other : Tensor) -> Tensor

Tensor.matrix_power(n : int) -> Tensor

Tensor.max(other : Tensor,
           out : Tensor) -> Tensor

Tensor.max() -> Tensor

Tensor.max(other : Tensor) -> Tensor

Tensor.max(dim : int,
           keepdim : bool=False) -> Tuple[Tensor, Tensor]

Tensor.mean() -> Tensor

Tensor.mean(dtype : int) -> Tensor

Tensor.mean(dim : List[int],
            keepdim : bool=False) -> Tensor

Tensor.mean(dim : List[int],
            dtype : int) -> Tensor

Tensor.mean(dim : List[int],
            keepdim : bool,
            dtype : int) -> Tensor

Tensor.mean(dim : List[int],
            keepdim : bool=False,
            out : Tensor) -> Tensor

Tensor.mean(dim : List[int],
            dtype : int,
            out : Tensor) -> Tensor

Tensor.mean(dim : List[int],
            keepdim : bool,
            dtype : int,
            out : Tensor) -> Tensor

Tensor.median() -> Tensor

Tensor.median(dim : int,
              keepdim : bool=False) -> Tuple[Tensor, Tensor]

Tensor.min() -> Tensor

Tensor.min(other : Tensor) -> Tensor

Tensor.min(dim : int,
           keepdim : bool=False) -> Tuple[Tensor, Tensor]

Tensor.min(other : Tensor,
           out : Tensor) -> Tensor

Tensor.mm(mat2 : Tensor,
          out : Tensor) -> Tensor

Tensor.mm(mat2 : Tensor) -> Tensor

Tensor.mode(dim : int=-1,
            keepdim : bool=False) -> Tuple[Tensor, Tensor]

Tensor.mul(other : Tensor) -> Tensor

Tensor.mul(other : number) -> Tensor

Tensor.mul(other : Tensor,
           out : Tensor) -> Tensor

Tensor.mul_(other : Tensor) -> Tensor

Tensor.mul_(other : number) -> Tensor

Tensor.multinomial(num_samples : int,
                   replacement : bool=False,
                   generator : Optional[Generator]) -> Tensor

Tensor.multinomial(num_samples : int,
                   replacement : bool=False,
                   generator : Optional[Generator],
                   out : Tensor) -> Tensor

Tensor.mv(vec : Tensor,
          out : Tensor) -> Tensor

Tensor.mv(vec : Tensor) -> Tensor

Tensor.mvlgamma(p : int) -> Tensor

Tensor.mvlgamma_(p : int) -> Tensor

Tensor.narrow(dim : int,
              start : int,
              length : int) -> Tensor

Tensor.narrow_copy(dim : int,
                   start : int,
                   length : int) -> Tensor

Tensor.ne(other : Tensor) -> Tensor

Tensor.ne(other : number) -> Tensor

Tensor.ne(other : Tensor,
          out : Tensor) -> Tensor

Tensor.ne(other : number,
          out : Tensor) -> Tensor

Tensor.ne_(other : Tensor) -> Tensor

Tensor.ne_(other : number) -> Tensor

Tensor.neg(out : Tensor) -> Tensor

Tensor.neg() -> Tensor

Tensor.neg_() -> Tensor

Tensor.nonzero(out : Tensor) -> Tensor

Tensor.nonzero() -> Tensor

Tensor.norm(p : number=2) -> Tensor

Tensor.norm(p : Optional[number],
            dtype : int) -> Tensor

Tensor.norm(p : Optional[number],
            dim : List[int],
            keepdim : bool=False) -> Tensor

Tensor.norm(p : Optional[number],
            dim : List[int],
            keepdim : bool,
            dtype : int) -> Tensor

Tensor.norm(p : Optional[number],
            dim : List[int],
            keepdim : bool=False,
            out : Tensor) -> Tensor

Tensor.norm(p : Optional[number],
            dim : List[int],
            keepdim : bool,
            dtype : int,
            out : Tensor) -> Tensor

Tensor.normal_(mean : float=0.0,
               std : float=1.0,
               generator : Optional[Generator]) -> Tensor

Tensor.numel() -> int

Tensor.orgqr(input2 : Tensor) -> Tensor

Tensor.orgqr(input2 : Tensor,
             out : Tensor) -> Tensor

Tensor.ormqr(input2 : Tensor,
             input3 : Tensor,
             left : bool=True,
             transpose : bool=False) -> Tensor

Tensor.ormqr(input2 : Tensor,
             input3 : Tensor,
             left : bool=True,
             transpose : bool=False,
             out : Tensor) -> Tensor

Tensor.permute(dims : List[int]) -> Tensor

Tensor.pin_memory() -> Tensor

Tensor.pinverse(rcond : float=1e-15) -> Tensor

Tensor.polygamma_(n : int) -> Tensor

Tensor.potri(upper : bool=True) -> Tensor

Tensor.potri(upper : bool=True,
             out : Tensor) -> Tensor

Tensor.pow(exponent : Tensor) -> Tensor

Tensor.pow(exponent : number) -> Tensor

Tensor.pow(exponent : Tensor,
           out : Tensor) -> Tensor

Tensor.pow(exponent : number,
           out : Tensor) -> Tensor

Tensor.pow_(exponent : Tensor) -> Tensor

Tensor.pow_(exponent : number) -> Tensor

Tensor.prelu(weight : Tensor) -> Tensor

Tensor.prod(dim : int,
            keepdim : bool=False,
            out : Tensor) -> Tensor

Tensor.prod(dim : int,
            dtype : int,
            out : Tensor) -> Tensor

Tensor.prod(dim : int,
            keepdim : bool,
            dtype : int,
            out : Tensor) -> Tensor

Tensor.prod() -> Tensor

Tensor.prod(dtype : int) -> Tensor

Tensor.prod(dim : int,
            keepdim : bool=False) -> Tensor

Tensor.prod(dim : int,
            dtype : int) -> Tensor

Tensor.prod(dim : int,
            keepdim : bool,
            dtype : int) -> Tensor

Tensor.pstrf(upper : bool=True,
             tol : number=-1) -> Tuple[Tensor, Tensor]

Tensor.put_(index : Tensor,
            source : Tensor,
            accumulate : bool=False) -> Tensor

Tensor.qr() -> Tuple[Tensor, Tensor]

Tensor.random_(generator : Optional[Generator]) -> Tensor

Tensor.random_(to : int,
               generator : Optional[Generator]) -> Tensor

Tensor.random_(from : int,
               to : int,
               generator : Optional[Generator]) -> Tensor

Tensor.reciprocal() -> Tensor

Tensor.reciprocal(out : Tensor) -> Tensor

Tensor.reciprocal_() -> Tensor

Tensor.relu() -> Tensor

Tensor.relu_() -> Tensor

Tensor.remainder(other : Tensor) -> Tensor

Tensor.remainder(other : number) -> Tensor

Tensor.remainder(other : Tensor,
                 out : Tensor) -> Tensor

Tensor.remainder(other : number,
                 out : Tensor) -> Tensor

Tensor.remainder_(other : Tensor) -> Tensor

Tensor.remainder_(other : number) -> Tensor

Tensor.renorm(p : number,
              dim : int,
              maxnorm : number,
              out : Tensor) -> Tensor

Tensor.renorm(p : number,
              dim : int,
              maxnorm : number) -> Tensor

Tensor.renorm_(p : number,
               dim : int,
               maxnorm : number) -> Tensor

Tensor.repeat(repeats : List[int]) -> Tensor

Tensor.reshape(shape : List[int]) -> Tensor

Tensor.reshape_as(other : Tensor) -> Tensor

Tensor.resize_(size : List[int]) -> Tensor

Tensor.resize_as_(the_template : Tensor) -> Tensor

Tensor.rfft(signal_ndim : int,
            normalized : bool=False,
            onesided : bool=True) -> Tensor

Tensor.roll(shifts : List[int],
            dims : List[int]=[]) -> Tensor

Tensor.rot90(k : int=1,
             dims : List[int]=[0, 1]) -> Tensor

Tensor.round() -> Tensor

Tensor.round(out : Tensor) -> Tensor

Tensor.round_() -> Tensor

Tensor.rsqrt(out : Tensor) -> Tensor

Tensor.rsqrt() -> Tensor

Tensor.rsqrt_() -> Tensor

Tensor.scatter(dim : int,
               index : Tensor,
               src : Tensor) -> Tensor

Tensor.scatter(dim : int,
               index : Tensor,
               value : number) -> Tensor

Tensor.scatter_(dim : int,
                index : Tensor,
                src : Tensor) -> Tensor

Tensor.scatter_(dim : int,
                index : Tensor,
                value : number) -> Tensor

Tensor.scatter_add(dim : int,
                   index : Tensor,
                   src : Tensor) -> Tensor

Tensor.scatter_add_(dim : int,
                    index : Tensor,
                    src : Tensor) -> Tensor

Tensor.select(dim : int,
              index : int) -> Tensor

Tensor.set_() -> Tensor

Tensor.set_(source : Tensor) -> Tensor

Tensor.sigmoid() -> Tensor

Tensor.sigmoid(out : Tensor) -> Tensor

Tensor.sigmoid_() -> Tensor

Tensor.sign() -> Tensor

Tensor.sign(out : Tensor) -> Tensor

Tensor.sign_() -> Tensor

Tensor.sin() -> Tensor

Tensor.sin(out : Tensor) -> Tensor

Tensor.sin_() -> Tensor

Tensor.sinh(out : Tensor) -> Tensor

Tensor.sinh() -> Tensor

Tensor.sinh_() -> Tensor

Tensor.size(dim : int) -> int

Tensor.size() -> List[int]

Tensor.slogdet() -> Tuple[Tensor, Tensor]

Tensor.smm(mat2 : Tensor) -> Tensor

Tensor.softmax(dim : int) -> Tensor

Tensor.softmax(dim : int,
               dtype : int) -> Tensor

Tensor.solve(A : Tensor) -> Tuple[Tensor, Tensor]

Tensor.sort(dim : int=-1,
            descending : bool=False) -> Tuple[Tensor, Tensor]

Tensor.sparse_dim() -> int

Tensor.sparse_resize_(size : List[int],
                      sparse_dim : int,
                      dense_dim : int) -> Tensor

Tensor.sparse_resize_and_clear_(size : List[int],
                                sparse_dim : int,
                                dense_dim : int) -> Tensor

Tensor.split(split_size : int,
             dim : int=0) -> List[Tensor]

Tensor.split(split_sizes : List[int],
             dim : int=0) -> List[Tensor]

Tensor.split_with_sizes(split_sizes : List[int],
                        dim : int=0) -> List[Tensor]

Tensor.sqrt(out : Tensor) -> Tensor

Tensor.sqrt() -> Tensor

Tensor.sqrt_() -> Tensor

Tensor.squeeze() -> Tensor

Tensor.squeeze(dim : int) -> Tensor

Tensor.squeeze_() -> Tensor

Tensor.squeeze_(dim : int) -> Tensor

Tensor.sspaddmm(mat1 : Tensor,
                mat2 : Tensor,
                beta : number=1,
                alpha : number=1,
                out : Tensor) -> Tensor

Tensor.sspaddmm(mat1 : Tensor,
                mat2 : Tensor,
                beta : number=1,
                alpha : number=1) -> Tensor

Tensor.std(unbiased : bool=True) -> Tensor

Tensor.std(dim : List[int],
           unbiased : bool=True,
           keepdim : bool=False) -> Tensor

Tensor.std(dim : List[int],
           unbiased : bool=True,
           keepdim : bool=False,
           out : Tensor) -> Tensor

Tensor.stft(n_fft : int,
            hop_length : Optional[int],
            win_length : Optional[int],
            window : Optional[Tensor],
            normalized : bool=False,
            onesided : bool=True) -> Tensor

Tensor.storage_offset() -> int

Tensor.storage_offset() -> int

Tensor.storage_offset() -> int

Tensor.stride(dim : int) -> int

Tensor.sub(other : Tensor,
           alpha : number=1) -> Tensor

Tensor.sub(other : number,
           alpha : number=1) -> Tensor

Tensor.sub(other : Tensor,
           alpha : number=1,
           out : Tensor) -> Tensor

Tensor.sub_(other : Tensor,
            alpha : number=1) -> Tensor

Tensor.sub_(other : number,
            alpha : number=1) -> Tensor

Tensor.sum(dim : List[int],
           keepdim : bool=False,
           out : Tensor) -> Tensor

Tensor.sum(dim : List[int],
           dtype : int,
           out : Tensor) -> Tensor

Tensor.sum(dim : List[int],
           keepdim : bool,
           dtype : int,
           out : Tensor) -> Tensor

Tensor.sum() -> Tensor

Tensor.sum(dtype : int) -> Tensor

Tensor.sum(dim : List[int],
           keepdim : bool=False) -> Tensor

Tensor.sum(dim : List[int],
           dtype : int) -> Tensor

Tensor.sum(dim : List[int],
           keepdim : bool,
           dtype : int) -> Tensor

Tensor.sum_to_size(size : List[int]) -> Tensor

Tensor.svd(some : bool=True,
           compute_uv : bool=True) -> Tuple[Tensor, Tensor, Tensor]

Tensor.symeig(eigenvectors : bool=False,
              upper : bool=True) -> Tuple[Tensor, Tensor]

Tensor.t() -> Tensor

Tensor.t_() -> Tensor

Tensor.take(index : Tensor) -> Tensor

Tensor.take(index : Tensor,
            out : Tensor) -> Tensor

Tensor.tan(out : Tensor) -> Tensor

Tensor.tan() -> Tensor

Tensor.tan_() -> Tensor

Tensor.tanh() -> Tensor

Tensor.tanh(out : Tensor) -> Tensor

Tensor.tanh_() -> Tensor

Tensor.to(other : Tensor,
          non_blocking : bool=False,
          copy : bool=False) -> Tensor

Tensor.to(dtype : int,
          non_blocking : bool=False,
          copy : bool=False) -> Tensor

Tensor.to(device : Device,
          dtype : int,
          non_blocking : bool=False,
          copy : bool=False) -> Tensor

Tensor.to(dtype : int,
          layout : int,
          device : Device,
          non_blocking : bool=False,
          copy : bool=False) -> Tensor

Tensor.to(device : Optional[Device],
          dtype : Optional[int],
          non_blocking : bool=False,
          copy : bool=False) -> Tensor

Tensor.to(dtype : Optional[int],
          non_blocking : bool=False,
          copy : bool=False) -> Tensor

Tensor.to(non_blocking : bool=False,
          copy : bool=False) -> Tensor

Tensor.to_dense() -> Tensor

Tensor.to_sparse() -> Tensor

Tensor.to_sparse(sparse_dim : int) -> Tensor

Tensor.topk(k : int,
            dim : int=-1,
            largest : bool=True,
            sorted : bool=True) -> Tuple[Tensor, Tensor]

Tensor.trace() -> Tensor

Tensor.transpose(dim0 : int,
                 dim1 : int) -> Tensor

Tensor.transpose_(dim0 : int,
                  dim1 : int) -> Tensor

Tensor.tril(diagonal : int=0,
            out : Tensor) -> Tensor

Tensor.tril(diagonal : int=0) -> Tensor

Tensor.tril_(diagonal : int=0) -> Tensor

Tensor.triu(diagonal : int=0,
            out : Tensor) -> Tensor

Tensor.triu(diagonal : int=0) -> Tensor

Tensor.triu_(diagonal : int=0) -> Tensor

Tensor.trtrs(A : Tensor,
             upper : bool=True,
             transpose : bool=False,
             unitriangular : bool=False) -> Tuple[Tensor, Tensor]

Tensor.trunc() -> Tensor

Tensor.trunc(out : Tensor) -> Tensor

Tensor.trunc_() -> Tensor

Tensor.type_as(other : Tensor) -> Tensor

Tensor.unbind(dim : int=0) -> List[Tensor]

Tensor.unfold(dimension : int,
              size : int,
              step : int) -> Tensor

Tensor.uniform_(from : float=0.0,
                to : float=1.0,
                generator : Optional[Generator]) -> Tensor

Tensor.unsqueeze(dim : int) -> Tensor

Tensor.unsqueeze_(dim : int) -> Tensor

Tensor.values() -> Tensor

Tensor.var(dim : List[int],
           unbiased : bool=True,
           keepdim : bool=False,
           out : Tensor) -> Tensor

Tensor.var(unbiased : bool=True) -> Tensor

Tensor.var(dim : List[int],
           unbiased : bool=True,
           keepdim : bool=False) -> Tensor

Tensor.view(size : List[int]) -> Tensor

Tensor.view_as(other : Tensor) -> Tensor

Tensor.zero_() -> Tensor

Frequently Asked Questions ¶

Q: I would like to train a model on GPU and do inference on CPU. What are the best practices?

First convert your model from GPU to CPU and then save it, like so:
cpu_model = gpu_model.cpu()
sample_input_cpu = sample_input_gpu.cpu()
traced_cpu = torch.jit.trace(traced_cpu, sample_input_cpu)
torch.jit.save(traced_cpu, "cpu.pth")

traced_gpu = torch.jit.trace(traced_gpu, sample_input_gpu)
torch.jit.save(traced_gpu, "gpu.pth")

# ... later, when using the model:

if use_gpu:
   model = torch.jit.load("gpu.pth")
else:
   model = torch.jit.load("cpu.pth")

model(input)
This is recommended because the tracer may witness tensor creation on a specific device, so casting an already-loaded model may have unexpected effects. Casting the model before saving it ensures that the tracer has the correct device information.

Multiprocessing package - torch.multiprocessing¶

torch.multiprocessing is a wrapper around the native multiprocessing module. It registers custom reducers, that use shared memory to provide shared views on the same data in different processes. Once the tensor/storage is moved to shared_memory (see share_memory_()), it will be possible to send it to other processes without making any copies.

The API is 100% compatible with the original module - it’s enough to change import multiprocessing to import torch.multiprocessing to have all the tensors sent through the queues or shared via other mechanisms, moved to shared memory.

Because of the similarity of APIs we do not document most of this package contents, and we recommend referring to very good docs of the original module.

Warning

If the main process exits abruptly (e.g. because of an incoming signal), Python’s multiprocessing sometimes fails to clean up its children. It’s a known caveat, so if you’re seeing any resource leaks after interrupting the interpreter, it probably means that this has just happened to you.

Strategy management¶

torch.multiprocessing.get_all_sharing_strategies()¶: Returns a set of sharing strategies supported on a current system.

torch.multiprocessing.get_sharing_strategy()¶: Returns the current strategy for sharing CPU tensors.

torch.multiprocessing.set_sharing_strategy(new_strategy)¶

Sets the strategy for sharing CPU tensors.

Parameters: new_strategy (str) – Name of the selected strategy. Should be one of the values returned by get_all_sharing_strategies().

Spawning subprocesses¶

Note

Available for Python >= 3.4.

This depends on the spawn start method in Python’s multiprocessing package.

Spawning a number of subprocesses to perform some function can be done by creating Process instances and calling join to wait for their completion. This approach works fine when dealing with a single subprocess but presents potential issues when dealing with multiple processes.

Namely, joining processes sequentially implies they will terminate sequentially. If they don’t, and the first process does not terminate, the process termination will go unnoticed. Also, there are no native facilities for error propagation.

The spawn function below addresses these concerns and takes care of error propagation, out of order termination, and will actively terminate processes upon detecting an error in one of them.

torch.multiprocessing.spawn(fn, args=(), nprocs=1, join=True, daemon=False)¶

Spawns nprocs processes that run fn with args.

If one of the processes exits with a non-zero exit status, the remaining processes are killed and an exception is raised with the cause of termination. In the case an exception was caught in the child process, it is forwarded and its traceback is included in the exception raised in the parent process.

Parameters

fn (function) –
Function is called as the entrypoint of the spawned process. This function must be defined at the top level of a module so it can be pickled and spawned. This is a requirement imposed by multiprocessing.

The function is called as fn(i, *args), where i is the process index and args is the passed through tuple of arguments.
args (tuple) – Arguments passed to fn.
nprocs (int) – Number of processes to spawn.
join (bool) – Perform a blocking join on all processes.
daemon (bool) – The spawned processes’ daemon flag. If set to True, daemonic processes will be created.

Returns

None if join is True, SpawnContext if join is False

class torch.multiprocessing.SpawnContext¶

Returned by spawn() when called with join=False.

join(timeout=None)¶

Tries to join one or more processes in this spawn context. If one of them exited with a non-zero exit status, this function kills the remaining processes and raises an exception with the cause of the first process exiting.

Returns True if all processes have been joined successfully, False if there are more processes that need to be joined.

Parameters: timeout (float) – Wait this long before giving up on waiting.

torch.utils.bottleneck¶

torch.utils.bottleneck is a tool that can be used as an initial step for debugging bottlenecks in your program. It summarizes runs of your script with the Python profiler and PyTorch’s autograd profiler.

Run it on the command line with

python -m torch.utils.bottleneck /path/to/source/script.py [args]

where [args] are any number of arguments to script.py, or run python -m torch.utils.bottleneck -h for more usage instructions.

Warning

Because your script will be profiled, please ensure that it exits in a finite amount of time.

Warning

Due to the asynchronous nature of CUDA kernels, when running against CUDA code, the cProfile output and CPU-mode autograd profilers may not show correct timings: the reported CPU time reports the amount of time used to launch the kernels but does not include the time the kernel spent executing on a GPU unless the operation does a synchronize. Ops that do synchronize appear to be extremely expensive under regular CPU-mode profilers. In these case where timings are incorrect, the CUDA-mode autograd profiler may be helpful.

Note

To decide which (CPU-only-mode or CUDA-mode) autograd profiler output to look at, you should first check if your script is CPU-bound (“CPU total time is much greater than CUDA total time”). If it is CPU-bound, looking at the results of the CPU-mode autograd profiler will help. If on the other hand your script spends most of its time executing on the GPU, then it makes sense to start looking for responsible CUDA operators in the output of the CUDA-mode autograd profiler.

Of course the reality is much more complicated and your script might not be in one of those two extremes depending on the part of the model you’re evaluating. If the profiler outputs don’t help, you could try looking at the result of torch.autograd.profiler.emit_nvtx() with nvprof. However, please take into account that the NVTX overhead is very high and often gives a heavily skewed timeline.

Warning

If you are profiling CUDA code, the first profiler that bottleneck runs (cProfile) will include the CUDA startup time (CUDA buffer allocation cost) in its time reporting. This should not matter if your bottlenecks result in code much slower than the CUDA startup time.

For more complicated uses of the profilers (like in a multi-GPU case), please see https://docs.python.org/3/library/profile.html or torch.autograd.profiler.profile() for more information.

torch.utils.checkpoint¶

Note

Checkpointing is implemented by rerunning a forward-pass segment for each checkpointed segment during backward. This can cause persistent states like the RNG state to be advanced than they would without checkpointing. By default, checkpointing includes logic to juggle the RNG state such that checkpointed passes making use of RNG (through dropout for example) have deterministic output as compared to non-checkpointed passes. The logic to stash and restore RNG states can incur a moderate performance hit depending on the runtime of checkpointed operations. If deterministic output compared to non-checkpointed passes is not required, supply preserve_rng_state=False to checkpoint or checkpoint_sequential to omit stashing and restoring the RNG state during each checkpoint.

The stashing logic saves and restores the RNG state for the current device and the device of all cuda Tensor arguments to the run_fn. However, the logic has no way to anticipate if the user will move Tensors to a new device within the run_fn itself. Therefore, if you move Tensors to a new device (“new” meaning not belonging to the set of [current device + devices of Tensor arguments]) within run_fn, deterministic output compared to non-checkpointed passes is never guaranteed.

torch.utils.checkpoint.checkpoint(function, *args, **kwargs)¶

Checkpoint a model or part of the model

Checkpointing works by trading compute for memory. Rather than storing all intermediate activations of the entire computation graph for computing backward, the checkpointed part does not save intermediate activations, and instead recomputes them in backward pass. It can be applied on any part of a model.

Specifically, in the forward pass, function will run in torch.no_grad() manner, i.e., not storing the intermediate activations. Instead, the forward pass saves the inputs tuple and the function parameter. In the backwards pass, the saved inputs and function is retreived, and the forward pass is computed on function again, now tracking the intermediate activations, and then the gradients are calculated using these activation values.

Warning

Checkpointing doesn’t work with torch.autograd.grad(), but only with torch.autograd.backward().

Warning

If function invocation during backward does anything different than the one during forward, e.g., due to some global variable, the checkpointed version won’t be equivalent, and unfortunately it can’t be detected.

Parameters

function – describes what to run in the forward pass of the model or part of the model. It should also know how to handle the inputs passed as the tuple. For example, in LSTM, if user passes (activation, hidden), function should correctly use the first input as activation and the second input as hidden
preserve_rng_state (bool, optional, default=True) – Omit stashing and restoring the RNG state during each checkpoint.
args – tuple containing inputs to the function

Returns

Output of running function on *args

torch.utils.checkpoint.checkpoint_sequential(functions, segments, *inputs, **kwargs)¶

A helper function for checkpointing sequential models.

Sequential models execute a list of modules/functions in order (sequentially). Therefore, we can divide such a model in various segments and checkpoint each segment. All segments except the last will run in torch.no_grad() manner, i.e., not storing the intermediate activations. The inputs of each checkpointed segment will be saved for re-running the segment in the backward pass.

See checkpoint() on how checkpointing works.

Warning

Checkpointing doesn’t work with torch.autograd.grad(), but only with torch.autograd.backward().

Parameters

functions – A torch.nn.Sequential or the list of modules or functions (comprising the model) to run sequentially.
segments – Number of chunks to create in the model
inputs – tuple of Tensors that are inputs to functions
preserve_rng_state (bool, optional, default=True) – Omit stashing and restoring the RNG state during each checkpoint.

Returns

Output of running functions sequentially on *inputs

Example

>>> model = nn.Sequential(...)
>>> input_var = checkpoint_sequential(model, chunks, input_var)

torch.utils.cpp_extension¶

torch.utils.cpp_extension.CppExtension(name, sources, *args, **kwargs)¶

Creates a setuptools.Extension for C++.

Convenience method that creates a setuptools.Extension with the bare minimum (but often sufficient) arguments to build a C++ extension.

All arguments are forwarded to the setuptools.Extension constructor.

Example

>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CppExtension
>>> setup(
        name='extension',
        ext_modules=[
            CppExtension(
                name='extension',
                sources=['extension.cpp'],
                extra_compile_args=['-g']),
        ],
        cmdclass={
            'build_ext': BuildExtension
        })

torch.utils.cpp_extension.CUDAExtension(name, sources, *args, **kwargs)¶

Creates a setuptools.Extension for CUDA/C++.

Convenience method that creates a setuptools.Extension with the bare minimum (but often sufficient) arguments to build a CUDA/C++ extension. This includes the CUDA include path, library path and runtime library.

All arguments are forwarded to the setuptools.Extension constructor.

Example

>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CUDAExtension
>>> setup(
        name='cuda_extension',
        ext_modules=[
            CUDAExtension(
                    name='cuda_extension',
                    sources=['extension.cpp', 'extension_kernel.cu'],
                    extra_compile_args={'cxx': ['-g'],
                                        'nvcc': ['-O2']})
        ],
        cmdclass={
            'build_ext': BuildExtension
        })

torch.utils.cpp_extension.BuildExtension(*args, **kwargs)¶

A custom setuptools build extension .

This setuptools.build_ext subclass takes care of passing the minimum required compiler flags (e.g. -std=c++11) as well as mixed C++/CUDA compilation (and support for CUDA files in general).

When using BuildExtension, it is allowed to supply a dictionary for extra_compile_args (rather than the usual list) that maps from languages (cxx or cuda) to a list of additional compiler flags to supply to the compiler. This makes it possible to supply different flags to the C++ and CUDA compiler during mixed compilation.

torch.utils.cpp_extension.load(name, sources, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, is_python_module=True)¶

Loads a PyTorch C++ extension just-in-time (JIT).

To load an extension, a Ninja build file is emitted, which is used to compile the given sources into a dynamic library. This library is subsequently loaded into the current Python process as a module and returned from this function, ready for use.

By default, the directory to which the build file is emitted and the resulting library compiled to is <tmp>/torch_extensions/<name>, where <tmp> is the temporary folder on the current platform and <name> the name of the extension. This location can be overridden in two ways. First, if the TORCH_EXTENSIONS_DIR environment variable is set, it replaces <tmp>/torch_extensions and all extensions will be compiled into subfolders of this directory. Second, if the build_directory argument to this function is supplied, it overrides the entire path, i.e. the library will be compiled into that folder directly.

To compile the sources, the default system compiler (c++) is used, which can be overridden by setting the CXX environment variable. To pass additional arguments to the compilation process, extra_cflags or extra_ldflags can be provided. For example, to compile your extension with optimizations, pass extra_cflags=['-O3']. You can also use extra_cflags to pass further include directories.

CUDA support with mixed compilation is provided. Simply pass CUDA source files (.cu or .cuh) along with other sources. Such files will be detected and compiled with nvcc rather than the C++ compiler. This includes passing the CUDA lib64 directory as a library directory, and linking cudart. You can pass additional flags to nvcc via extra_cuda_cflags, just like with extra_cflags for C++. Various heuristics for finding the CUDA install directory are used, which usually work fine. If not, setting the CUDA_HOME environment variable is the safest option.

Parameters

name – The name of the extension to build. This MUST be the same as the name of the pybind11 module!
sources – A list of relative or absolute paths to C++ source files.
extra_cflags – optional list of compiler flags to forward to the build.
extra_cuda_cflags – optional list of compiler flags to forward to nvcc when building CUDA sources.
extra_ldflags – optional list of linker flags to forward to the build.
extra_include_paths – optional list of include directories to forward to the build.
build_directory – optional path to use as build workspace.
verbose – If True, turns on verbose logging of load steps.
with_cuda – Determines whether CUDA headers and libraries are added to the build. If set to None (default), this value is automatically determined based on the existence of .cu or .cuh in sources. Set it to True` to force CUDA headers and libraries to be included.
is_python_module – If True (default), imports the produced shared library as a Python module. If False, loads it into the process as a plain dynamic library.

Returns

If is_python_module is True, returns the loaded PyTorch extension as a Python module. If is_python_module is False returns nothing (the shared library is loaded into the process as a side effect).

Example

>>> from torch.utils.cpp_extension import load
>>> module = load(
        name='extension',
        sources=['extension.cpp', 'extension_kernel.cu'],
        extra_cflags=['-O2'],
        verbose=True)

torch.utils.cpp_extension.load_inline(name, cpp_sources, cuda_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, is_python_module=True)¶

Loads a PyTorch C++ extension just-in-time (JIT) from string sources.

This function behaves exactly like load(), but takes its sources as strings rather than filenames. These strings are stored to files in the build directory, after which the behavior of load_inline() is identical to load().

See the tests for good examples of using this function.

Sources may omit two required parts of a typical non-inline C++ extension: the necessary header includes, as well as the (pybind11) binding code. More precisely, strings passed to cpp_sources are first concatenated into a single .cpp file. This file is then prepended with #include <torch/extension.h>.

Furthermore, if the functions argument is supplied, bindings will be automatically generated for each function specified. functions can either be a list of function names, or a dictionary mapping from function names to docstrings. If a list is given, the name of each function is used as its docstring.

The sources in cuda_sources are concatenated into a separate .cu file and prepended with torch/types.h, cuda.h and cuda_runtime.h includes. The .cpp and .cu files are compiled separately, but ultimately linked into a single library. Note that no bindings are generated for functions in cuda_sources per se. To bind to a CUDA kernel, you must create a C++ function that calls it, and either declare or define this C++ function in one of the cpp_sources (and include its name in functions).

See load() for a description of arguments omitted below.

Parameters

cpp_sources – A string, or list of strings, containing C++ source code.
cuda_sources – A string, or list of strings, containing CUDA source code.
functions – A list of function names for which to generate function bindings. If a dictionary is given, it should map function names to docstrings (which are otherwise just the function names).
with_cuda – Determines whether CUDA headers and libraries are added to the build. If set to None (default), this value is automatically determined based on whether cuda_sources is provided. Set it to True` to force CUDA headers and libraries to be included.

Example

>>> from torch.utils.cpp_extension import load_inline
>>> source = '''
at::Tensor sin_add(at::Tensor x, at::Tensor y) {
  return x.sin() + y.sin();
}
'''
>>> module = load_inline(name='inline_extension',
                         cpp_sources=[source],
                         functions=['sin_add'])

torch.utils.cpp_extension.include_paths(cuda=False)¶

Get the include paths required to build a C++ or CUDA extension.

Parameters: cuda – If True, includes CUDA-specific include paths.
Returns: A list of include path strings.

torch.utils.cpp_extension.check_compiler_abi_compatibility(compiler)¶

Verifies that the given compiler is ABI-compatible with PyTorch.

Parameters: compiler (str) – The compiler executable name to check (e.g. g++). Must be executable in a shell process.
Returns: False if the compiler is (likely) ABI-incompatible with PyTorch, else True.

torch.utils.cpp_extension.verify_ninja_availability()¶: Returns True if the ninja build system is available on the system.

torch.utils.data¶

class torch.utils.data.Dataset¶

An abstract class representing a Dataset.

All other datasets should subclass it. All subclasses should override __len__, that provides the size of the dataset, and __getitem__, supporting integer indexing in range from 0 to len(self) exclusive.

class torch.utils.data.TensorDataset(*tensors)¶

Dataset wrapping tensors.

Each sample will be retrieved by indexing tensors along the first dimension.

Parameters: *tensors (Tensor) – tensors that have the same size of the first dimension.

class torch.utils.data.ConcatDataset(datasets)¶

Dataset to concatenate multiple datasets. Purpose: useful to assemble different existing datasets, possibly large-scale datasets as the concatenation operation is done in an on-the-fly manner.

Parameters: datasets (sequence) – List of datasets to be concatenated

class torch.utils.data.Subset(dataset, indices)¶

Subset of a dataset at specified indices.

Parameters

dataset (Dataset) – The whole Dataset
indices (sequence) – Indices in the whole set selected for subset

class torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, collate_fn=<function default_collate>, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None)¶

Data loader. Combines a dataset and a sampler, and provides single- or multi-process iterators over the dataset.

Parameters

dataset (Dataset) – dataset from which to load the data.
batch_size (int, optional) – how many samples per batch to load (default: 1).
shuffle (bool, optional) – set to True to have the data reshuffled at every epoch (default: False).
sampler (Sampler, optional) – defines the strategy to draw samples from the dataset. If specified, shuffle must be False.
batch_sampler (Sampler, optional) – like sampler, but returns a batch of indices at a time. Mutually exclusive with batch_size, shuffle, sampler, and drop_last.
num_workers (int, optional) – how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: 0)
collate_fn (callable, optional) – merges a list of samples to form a mini-batch.
pin_memory (bool, optional) – If True, the data loader will copy tensors into CUDA pinned memory before returning them. If your data elements are a custom type, or your collate_fn returns a batch that is a custom type see the example below.
drop_last (bool, optional) – set to True to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If False and the size of dataset is not divisible by the batch size, then the last batch will be smaller. (default: False)
timeout (numeric, optional) – if positive, the timeout value for collecting a batch from workers. Should always be non-negative. (default: 0)
worker_init_fn (callable, optional) – If not None, this will be called on each worker subprocess with the worker id (an int in [0, num_workers - 1]) as input, after seeding and before data loading. (default: None)

Note

By default, each worker will have its PyTorch seed set to base_seed + worker_id, where base_seed is a long generated by main process using its RNG. However, seeds for other libraies may be duplicated upon initializing workers (w.g., NumPy), causing each worker to return identical random numbers. (See dataloader-workers-random-seed section in FAQ.) You may use torch.initial_seed() to access the PyTorch seed for each worker in worker_init_fn, and use it to set other seeds before data loading.

Warning

If spawn start method is used, worker_init_fn cannot be an unpicklable object, e.g., a lambda function.

The default memory pinning logic only recognizes Tensors and maps and iterables containg Tensors. By default, if the pinning logic sees a batch that is a custom type (which will occur if you have a collate_fn that returns a custom batch type), or if each element of your batch is a custom type, the pinning logic will not recognize them, and it will return that batch (or those elements) without pinning the memory. To enable memory pinning for custom batch or data types, define a pin_memory method on your custom type(s).

Example:

class SimpleCustomBatch:
    def __init__(self, data):
        transposed_data = list(zip(*data))
        self.inp = torch.stack(transposed_data[0], 0)
        self.tgt = torch.stack(transposed_data[1], 0)

    def pin_memory(self):
        self.inp = self.inp.pin_memory()
        self.tgt = self.tgt.pin_memory()
        return self

def collate_wrapper(batch):
    return SimpleCustomBatch(batch)

inps = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
tgts = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
dataset = TensorDataset(inps, tgts)

loader = DataLoader(dataset, batch_size=2, collate_fn=collate_wrapper,
                    pin_memory=True)

for batch_ndx, sample in enumerate(loader):
    print(sample.inp.is_pinned())
    print(sample.tgt.is_pinned())

torch.utils.data.random_split(dataset, lengths)¶

Randomly split a dataset into non-overlapping new datasets of given lengths.

Parameters

dataset (Dataset) – Dataset to be split
lengths (sequence) – lengths of splits to be produced

class torch.utils.data.Sampler(data_source)¶

Base class for all Samplers.

Every Sampler subclass has to provide an __iter__ method, providing a way to iterate over indices of dataset elements, and a __len__ method that returns the length of the returned iterators.

class torch.utils.data.SequentialSampler(data_source)¶

Samples elements sequentially, always in the same order.

Parameters: data_source (Dataset) – dataset to sample from

class torch.utils.data.RandomSampler(data_source, replacement=False, num_samples=None)¶

Samples elements randomly. If without replacement, then sample from a shuffled dataset. If with replacement, then user can specify num_samples to draw.

Parameters

data_source (Dataset) – dataset to sample from
num_samples (int) – number of samples to draw, default=len(dataset)
replacement (bool) – samples are drawn with replacement if True, default=False

class torch.utils.data.SubsetRandomSampler(indices)¶

Samples elements randomly from a given list of indices, without replacement.

Parameters: indices (sequence) – a sequence of indices

class torch.utils.data.WeightedRandomSampler(weights, num_samples, replacement=True)¶

Samples elements from [0,..,len(weights)-1] with given probabilities (weights).

Parameters

weights (sequence) – a sequence of weights, not necessary summing up to one
num_samples (int) – number of samples to draw
replacement (bool) – if True, samples are drawn with replacement. If not, they are drawn without replacement, which means that when a sample index is drawn for a row, it cannot be drawn again for that row.

Example

>>> list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True))
[0, 0, 0, 1, 0]
>>> list(WeightedRandomSampler([0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False))
[0, 1, 4, 3, 2]

class torch.utils.data.BatchSampler(sampler, batch_size, drop_last)¶

Wraps another sampler to yield a mini-batch of indices.

Parameters

sampler (Sampler) – Base sampler.
batch_size (int) – Size of mini-batch.
drop_last (bool) – If True, the sampler will drop the last batch if its size would be less than batch_size

Example

>>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False))
[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
>>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True))
[[0, 1, 2], [3, 4, 5], [6, 7, 8]]

class torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None)¶

Sampler that restricts data loading to a subset of the dataset.

It is especially useful in conjunction with torch.nn.parallel.DistributedDataParallel. In such case, each process can pass a DistributedSampler instance as a DataLoader sampler, and load a subset of the original dataset that is exclusive to it.

Note

Dataset is assumed to be of constant size.

Parameters

dataset – Dataset used for sampling.
num_replicas (optional) – Number of processes participating in distributed training.
rank (optional) – Rank of the current process within num_replicas.

torch.utils.dlpack¶

torch.utils.dlpack.from_dlpack(dlpack) → Tensor¶

Decodes a DLPack to a tensor.

Parameters: dlpack – a PyCapsule object with the dltensor

The tensor will share the memory with the object represented in the dlpack. Note that each dlpack can only be consumed once.

torch.utils.dlpack.to_dlpack(tensor) → PyCapsule¶

Returns a DLPack representing the tensor.

Parameters: tensor – a tensor to be exported

The dlpack shares the tensors memory. Note that each dlpack can only be consumed once.

torch.hub¶

Pytorch Hub is a pre-trained model repository designed to facilitate research reproducibility.

Publishing models¶

Pytorch Hub supports publishing pre-trained models(model definitions and pre-trained weights) to a github repository by adding a simple hubconf.py file;

hubconf.py can have multiple entrypoints. Each entrypoint is defined as a python function with the following signature.

def entrypoint_name(pretrained=False, *args, **kwargs):
    ...

How to implement an entrypoint?¶

Here is a code snippet from pytorch/vision repository, which specifies an entrypoint for resnet18 model. You can see a full script in pytorch/vision repo

dependencies = ['torch', 'math']

def resnet18(pretrained=False, *args, **kwargs):
    """
    Resnet18 model
    pretrained (bool): a recommended kwargs for all entrypoints
    args & kwargs are arguments for the function
    """
    ######## Call the model in the repo ###############
    from torchvision.models.resnet import resnet18 as _resnet18
    model = _resnet18(*args, **kwargs)
    ######## End of call ##############################
    # The following logic is REQUIRED
    if pretrained:
        # For weights saved in local repo
        # model.load_state_dict(<path_to_saved_file>)

        # For weights saved elsewhere
        checkpoint = 'https://download.pytorch.org/models/resnet18-5c106cde.pth'
        model.load_state_dict(model_zoo.load_url(checkpoint, progress=False))
    return model

dependencies variable is a list of package names required to to run the model.
Pretrained weights can either be stored local in the github repo, or loadable by model_zoo.load().
pretrained controls whether to load the pre-trained weights provided by repo owners.
args and kwargs are passed along to the real callable function.
Docstring of the function works as a help message, explaining what does the model do and what are the allowed arguments.
Entrypoint function should ALWAYS return a model(nn.module).

Important Notice¶

The published models should be at least in a branch/tag. It can’t be a random commit.

Loading models from Hub¶

Users can load the pre-trained models using torch.hub.load() API.

torch.hub.load(github, model, force_reload=False, *args, **kwargs)¶

Load a model from a github repo, with pretrained weights.

Parameters

github – Required, a string with format “repo_owner/repo_name[:tag_name]” with an optional tag/branch. The default branch is master if not specified. Example: ‘pytorch/vision[:hub]’
model – Required, a string of entrypoint name defined in repo’s hubconf.py
force_reload – Optional, whether to discard the existing cache and force a fresh download. Default is False.
*args – Optional, the corresponding args for callable model.
**kwargs – Optional, the corresponding kwargs for callable model.

Returns

a single model with corresponding pretrained weights.

Here’s an example loading resnet18 entrypoint from pytorch/vision repo.

hub_model = hub.load(
    'pytorch/vision:master', # repo_owner/repo_name:branch
    'resnet18', # entrypoint
    1234, # args for callable [not applicable to resnet]
    pretrained=True) # kwargs for callable

Where are my downloaded model & weights saved?¶

The locations are used in the order of

hub_dir: user specified path. It can be set in the following ways: - Setting the environment variable TORCH_HUB_DIR - Calling hub.set_dir(<PATH_TO_HUB_DIR>)
~/.torch/hub

torch.hub.set_dir(d)¶

Optionally set hub_dir to a local dir to save downloaded models & weights.

If this argument is not set, env variable TORCH_HUB_DIR will be searched first, ~/.torch/hub will be created and used as fallback.

Parameters: d – path to a local folder to save downloaded models & weights.

Caching logic¶

By default, we don’t clean up files after loading it. Hub uses the cache by default if it already exists in hub_dir.

Users can force a reload by calling hub.load(..., force_reload=True). This will delete the existing github folder and downloaded weights, reinitialize a fresh download. This is useful when updates are published to the same branch, users can keep up with the latest release.

torch.utils.model_zoo¶

torch.utils.model_zoo.load_url(url, model_dir=None, map_location=None, progress=True)¶

Loads the Torch serialized object at the given URL.

If the object is already present in model_dir, it’s deserialized and returned. The filename part of the URL should follow the naming convention filename-<sha256>.ext where <sha256> is the first eight or more digits of the SHA256 hash of the contents of the file. The hash is used to ensure unique names and to verify the contents of the file.

The default value of model_dir is $TORCH_HOME/models where $TORCH_HOME defaults to ~/.torch. The default directory can be overridden with the $TORCH_MODEL_ZOO environment variable.

Parameters

url (string) – URL of the object to download
model_dir (string, optional) – directory in which to save the object
map_location (optional) – a function or a dict specifying how to remap storage locations (see torch.load)
progress (bool, optional) – whether or not to display a progress bar to stderr

Example

>>> state_dict = torch.utils.model_zoo.load_url('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth')

torch.onnx¶

Example: End-to-end AlexNet from PyTorch to Caffe2¶

Here is a simple script which exports a pretrained AlexNet as defined in torchvision into ONNX. It runs a single round of inference and then saves the resulting traced model to alexnet.onnx:

import torch
import torchvision

dummy_input = torch.randn(10, 3, 224, 224, device='cuda')
model = torchvision.models.alexnet(pretrained=True).cuda()

# Providing input and output names sets the display names for values
# within the model's graph. Setting these does not change the semantics
# of the graph; it is only for readability.
#
# The inputs to the network consist of the flat list of inputs (i.e.
# the values you would pass to the forward() method) followed by the
# flat list of parameters. You can partially specify names, i.e. provide
# a list here shorter than the number of inputs to the model, and we will
# only set that subset of names, starting from the beginning.
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]

torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)

The resulting alexnet.onnx is a binary protobuf file which contains both the network structure and parameters of the model you exported (in this case, AlexNet). The keyword argument verbose=True causes the exporter to print out a human-readable representation of the network:

# These are the inputs and parameters to the network, which have taken on
# the names we specified earlier.
graph(%actual_input_1 : Float(10, 3, 224, 224)
      %learned_0 : Float(64, 3, 11, 11)
      %learned_1 : Float(64)
      %learned_2 : Float(192, 64, 5, 5)
      %learned_3 : Float(192)
      # ---- omitted for brevity ----
      %learned_14 : Float(1000, 4096)
      %learned_15 : Float(1000)) {
  # Every statement consists of some output tensors (and their types),
  # the operator to be run (with its attributes, e.g., kernels, strides,
  # etc.), its input tensors (%actual_input_1, %learned_0, %learned_1)
  %17 : Float(10, 64, 55, 55) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[11, 11], pads=[2, 2, 2, 2], strides=[4, 4]](%actual_input_1, %learned_0, %learned_1), scope: AlexNet/Sequential[features]/Conv2d[0]
  %18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: AlexNet/Sequential[features]/ReLU[1]
  %19 : Float(10, 64, 27, 27) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%18), scope: AlexNet/Sequential[features]/MaxPool2d[2]
  # ---- omitted for brevity ----
  %29 : Float(10, 256, 6, 6) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%28), scope: AlexNet/Sequential[features]/MaxPool2d[12]
  # Dynamic means that the shape is not known. This may be because of a
  # limitation of our implementation (which we would like to fix in a
  # future release) or shapes which are truly dynamic.
  %30 : Dynamic = onnx::Shape(%29), scope: AlexNet
  %31 : Dynamic = onnx::Slice[axes=[0], ends=[1], starts=[0]](%30), scope: AlexNet
  %32 : Long() = onnx::Squeeze[axes=[0]](%31), scope: AlexNet
  %33 : Long() = onnx::Constant[value={9216}](), scope: AlexNet
  # ---- omitted for brevity ----
  %output1 : Float(10, 1000) = onnx::Gemm[alpha=1, beta=1, broadcast=1, transB=1](%45, %learned_14, %learned_15), scope: AlexNet/Sequential[classifier]/Linear[6]
  return (%output1);
}

You can also verify the protobuf using the onnx library. You can install onnx with conda:

conda install -c conda-forge onnx

Then, you can run:

import onnx

# Load the ONNX model
model = onnx.load("alexnet.onnx")

# Check that the IR is well formed
onnx.checker.check_model(model)

# Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph)

To run the exported script with caffe2, you will need to install caffe2: If you don’t have one already, Please follow the install instructions.

Once these are installed, you can use the backend for Caffe2:

# ...continuing from above
import caffe2.python.onnx.backend as backend
import numpy as np

rep = backend.prepare(model, device="CUDA:0") # or "CPU"
# For the Caffe2 backend:
#     rep.predict_net is the Caffe2 protobuf for the network
#     rep.workspace is the Caffe2 workspace for the network
#       (see the class caffe2.python.onnx.backend.Workspace)
outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))
# To run networks with more than one input, pass a tuple
# rather than a single numpy ndarray.
print(outputs[0])

In the future, there will be backends for other frameworks as well.

Limitations¶

The ONNX exporter is a trace-based exporter, which means that it operates by executing your model once, and exporting the operators which were actually run during this run. This means that if your model is dynamic, e.g., changes behavior depending on input data, the export won’t be accurate. Similarly, a trace is likely to be valid only for a specific input size (which is one reason why we require explicit inputs on tracing.) We recommend examining the model trace and making sure the traced operators look reasonable.
PyTorch and Caffe2 often have implementations of operators with some numeric differences. Depending on model structure, these differences may be negligible, but they can also cause major divergences in behavior (especially on untrained models.) In a future release, we plan to allow Caffe2 to call directly to Torch implementations of operators, to help you smooth over these differences when precision is important, and to also document these differences.

Supported operators¶

The following operators are supported:

add (nonzero alpha not supported)
sub (nonzero alpha not supported)
mul
div
cat
mm
addmm
neg
sqrt
tanh
sigmoid
mean
sum
prod
t
expand (only when used before a broadcasting ONNX operator; e.g., add)
transpose
view
split
squeeze
prelu (single weight shared among input channels not supported)
threshold (non-zero threshold/non-zero value not supported)
leaky_relu
glu
softmax (only dim=-1 supported)
avg_pool2d (ceil_mode not supported)
log_softmax
unfold (experimental support with ATen-Caffe2 integration)
elu
concat
abs
index_select
pow
clamp
max
min
eq
gt
lt
ge
le
exp
sin
cos
tan
asin
acos
atan
permute
Conv
BatchNorm
MaxPool1d (ceil_mode not supported)
MaxPool2d (ceil_mode not supported)
MaxPool3d (ceil_mode not supported)
Embedding (no optional arguments supported)
RNN
ConstantPadNd
Dropout
FeatureDropout (training mode not supported)
Index (constant integer and tuple indices supported)

The operator set above is sufficient to export the following models:

AlexNet
DCGAN
DenseNet
Inception (warning: this model is highly sensitive to changes in operator implementation)
ResNet
SuperResolution
VGG
word_language_model

Adding export support for operators is an advance usage. To achieve this, developers need to touch the source code of PyTorch. Please follow the instructions for installing PyTorch from source. If the wanted operator is standardized in ONNX, it should be easy to add support for exporting such operator (adding a symbolic function for the operator). To confirm whether the operator is standardized or not, please check the ONNX operator list.

If the operator is an ATen operator, which means you can find the declaration of the function in torch/csrc/autograd/generated/VariableType.h (available in generated code in PyTorch install dir), you should add the symbolic function in torch/onnx/symbolic.py and follow the instructions listed as below:

Define the symbolic function in torch/onnx/symbolic.py. Make sure the function has the same name as the ATen operator/function defined in VariableType.h.
The first parameter is always the exported ONNX graph. Parameter names must EXACTLY match the names in VariableType.h, because dispatch is done with keyword arguments.
Parameter ordering does NOT necessarily match what is in VariableType.h, tensors (inputs) are always first, then non-tensor arguments.
In the symbolic function, if the operator is already standardized in ONNX, we only need to create a node to represent the ONNX operator in the graph.
If the input argument is a tensor, but ONNX asks for a scalar, we have to explicitly do the conversion. The helper function _scalar can convert a scalar tensor into a python scalar, and _if_scalar_type_as can turn a Python scalar into a PyTorch tensor.

If the operator is a non-ATen operator, the symbolic function has to be added in the corresponding PyTorch Function class. Please read the following instructions:

Create a symbolic function named symbolic in the corresponding Function class.
The first parameter is always the exported ONNX graph.
Parameter names except the first must EXACTLY match the names in forward.
The output tuple size must match the outputs of forward.
In the symbolic function, if the operator is already standardized in ONNX, we just need to create a node to represent the ONNX operator in the graph.

Symbolic functions should be implemented in Python. All of these functions interact with Python methods which are implemented via C++-Python bindings, but intuitively the interface they provide looks like this:

def operator/symbolic(g, *inputs):
  """
  Modifies Graph (e.g., using "op"), adding the ONNX operations representing
  this PyTorch function, and returning a Value or tuple of Values specifying the
  ONNX outputs whose values correspond to the original PyTorch return values
  of the autograd Function (or None if an output is not supported by ONNX).

  Arguments:
    g (Graph): graph to write the ONNX representation into
    inputs (Value...): list of values representing the variables which contain
        the inputs for this function
  """

class Value(object):
  """Represents an intermediate tensor value computed in ONNX."""
  def type(self):
    """Returns the Type of the value."""

class Type(object):
  def sizes(self):
    """Returns a tuple of ints representing the shape of a tensor this describes."""

class Graph(object):
  def op(self, opname, *inputs, **attrs):
    """
    Create an ONNX operator 'opname', taking 'args' as inputs
    and attributes 'kwargs' and add it as a node to the current graph,
    returning the value representing the single output of this
    operator (see the `outputs` keyword argument for multi-return
    nodes).

    The set of operators and the inputs/attributes they take
    is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md

    Arguments:
        opname (string): The ONNX operator name, e.g., `Abs` or `Add`.
        args (Value...): The inputs to the operator; usually provided
            as arguments to the `symbolic` definition.
        kwargs: The attributes of the ONNX operator, with keys named
            according to the following convention: `alpha_f` indicates
            the `alpha` attribute with type `f`.  The valid type specifiers are
            `f` (float), `i` (int), `s` (string) or `t` (Tensor).  An attribute
            specified with type float accepts either a single float, or a
            list of floats (e.g., you would say `dims_i` for a `dims` attribute
            that takes a list of integers).
        outputs (int, optional):  The number of outputs this operator returns;
            by default an operator is assumed to return a single output.
            If `outputs` is greater than one, this functions returns a tuple
            of output `Value`, representing each output of the ONNX operator
            in positional.
    """

The ONNX graph C++ definition is in torch/csrc/jit/ir.h.

Here is an example of handling missing symbolic function for elu operator. We try to export the model and see the error message as below:

UserWarning: ONNX export failed on elu because torch.onnx.symbolic.elu does not exist
RuntimeError: ONNX export failed: Couldn't export operator elu

The export fails because PyTorch does not support exporting elu operator. We find virtual Tensor elu(const Tensor & input, Scalar alpha, bool inplace) const override; in VariableType.h. This means elu is an ATen operator. We check the ONNX operator list, and confirm that Elu is standardized in ONNX. We add the following lines to symbolic.py:

def elu(g, input, alpha, inplace=False):
    return g.op("Elu", input, alpha_f=_scalar(alpha))

Now PyTorch is able to export elu operator.

There are more examples in symbolic.py, tensor.py, padding.py.

The interface for specifying operator definitions is experimental; adventurous users should note that the APIs will probably change in a future interface.

Functions¶

torch.onnx.export(*args, **kwargs)¶

Distributed communication package (deprecated) - torch.distributed.deprecated¶

Warning

torch.distributed.deprecated is the older version of torch.distributed and currently deprecated. It will be removed soon. Please use and refer the doc for torch.distributed, which is the latest distributed communication package for PyTorch

torch.distributed.deprecated provides an MPI-like interface for exchanging tensor data across multi-machine networks. It supports a few different backends and initialization methods.

Currently torch.distributed.deprecated supports four backends, each with different capabilities. The table below shows which functions are available for use with CPU / CUDA tensors. MPI supports cuda only if the implementation used to build PyTorch supports it.

Backend	`tcp`		`gloo`		`mpi`		`nccl`
Device	CPU	GPU	CPU	GPU	CPU	GPU	CPU	GPU
send	✓	✘	✘	✘	✓	?	✘	✘
recv	✓	✘	✘	✘	✓	?	✘	✘
broadcast	✓	✘	✓	✓	✓	?	✘	✓
all_reduce	✓	✘	✓	✓	✓	?	✘	✓
reduce	✓	✘	✘	✘	✓	?	✘	✓
all_gather	✓	✘	✘	✘	✓	?	✘	✓
gather	✓	✘	✘	✘	✓	?	✘	✘
scatter	✓	✘	✘	✘	✓	?	✘	✘
barrier	✓	✘	✓	✓	✓	?	✘	✘

Basics¶

The torch.distributed.deprecated package provides PyTorch support and communication primitives for multiprocess parallelism across several computation nodes running on one or more machines. The class torch.nn.parallel.deprecated.DistributedDataParallel() builds on this functionality to provide synchronous distributed training as a wrapper around any PyTorch model. This differs from the kinds of parallelism provided by Multiprocessing package - torch.multiprocessing and torch.nn.DataParallel() in that it supports multiple network-connected machines and in that the user must explicitly launch a separate copy of the main training script for each process.

In the single-machine synchronous case, torch.distributed.deprecated or the torch.nn.parallel.deprecated.DistributedDataParallel() wrapper may still have advantages over other approaches to data-parallelism, including torch.nn.DataParallel():

Each process maintains its own optimizer and performs a complete optimization step with each iteration. While this may appear redundant, since the gradients have already been gathered together and averaged across processes and are thus the same for every process, this means that no parameter broadcast step is needed, reducing time spent transferring tensors between nodes.
Each process contains an independent Python interpreter, eliminating the extra interpreter overhead and “GIL-thrashing” that comes from driving several execution threads, model replicas, or GPUs from a single Python process. This is especially important for models that make heavy use of the Python runtime, including models with recurrent layers or many small components.

Initialization¶

The package needs to be initialized using the torch.distributed.deprecated.init_process_group() function before calling any other methods. This blocks until all processes have joined.

torch.distributed.deprecated.init_process_group(backend, init_method='env://', **kwargs)¶

Initializes the distributed package.

Parameters

backend (str) – Name of the backend to use. Depending on build-time configuration valid values include: tcp, mpi, gloo and nccl.
init_method (str, optional) – URL specifying how to initialize the package.
world_size (int, optional) – Number of processes participating in the job.
rank (int, optional) – Rank of the current process.
group_name (str, optional) – Group name. See description of init methods.

To enable backend == mpi, PyTorch needs to built from source on a system that supports MPI. If you want to use Open MPI with CUDA-aware support, please use Open MPI major version 2 and above.

Note

This method initializes CUDA context. Therefore, if multiple processes run on a single machine but use different GPUs, make sure to use torch.cuda.set_device() before this method to avoid unnecessarily creating context on the first visible device.

torch.distributed.deprecated.get_rank()¶

Returns the rank of current process.

Rank is a unique identifier assigned to each process within a distributed group. They are always consecutive integers ranging from 0 to world_size - 1 (inclusive).

torch.distributed.deprecated.get_world_size()¶: Returns the number of processes in the distributed group.

Currently three initialization methods are supported:

TCP initialization¶

Alternatively, the address has to be a valid IP multicast address, in which case ranks can be assigned automatically. Multicast initialization also supports a group_name argument, which allows you to use the same address for multiple jobs, as long as they use different group names.

import torch.distributed.deprecated as dist

# Use address of one of the machines
dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456', rank=args.rank, world_size=4)

# or a multicast address - rank will be assigned automatically if unspecified
dist.init_process_group(backend, init_method='tcp://[ff15:1e18:5d4c:4cf0:d02d:b659:53ba:b0a7]:23456',
                        world_size=4)

Shared file-system initialization¶

Another initialization method makes use of a file system that is shared and visible from all machines in a group, along with a desired world_size. The URL should start with file:// and contain a path to a non-existent file (in an existing directory) on a shared file system. This initialization method also supports a group_name argument, which allows you to use the same shared file path for multiple jobs, as long as they use different group names.

Warning

This method assumes that the file system supports locking using fcntl - most local systems and NFS support it.

import torch.distributed.deprecated as dist

# Rank will be assigned automatically if unspecified
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',
                        world_size=4, group_name=args.group)

Environment variable initialization¶

This method will read the configuration from environment variables, allowing one to fully customize how the information is obtained. The variables to be set are:

MASTER_PORT - required; has to be a free port on machine with rank 0
MASTER_ADDR - required (except for rank 0); address of rank 0 node
WORLD_SIZE - required; can be set either here, or in a call to init function
RANK - required; can be set either here, or in a call to init function

The machine with rank 0 will be used to set up all connections.

This is the default method, meaning that init_method does not have to be specified (or can be env://).

Groups¶

torch.distributed.deprecated.new_group(ranks=None)¶

Creates a new distributed group.

This function requires that all processes in the main group (i.e., all processes that are part of the distributed job) enter this function, even if they are not going to be members of the group. Additionally, groups should be created in the same order in all processes.

Parameters: ranks (list[int]) – List of ranks of group members.
Returns: A handle of distributed group that can be given to collective calls.

Point-to-point communication¶

torch.distributed.deprecated.send(tensor, dst)¶

Sends a tensor synchronously.

Parameters

tensor (Tensor) – Tensor to send.
dst (int) – Destination rank.

torch.distributed.deprecated.recv(tensor, src=None)¶

Receives a tensor synchronously.

Parameters

tensor (Tensor) – Tensor to fill with received data.
src (int, optional) – Source rank. Will receive from any process if unspecified.

Returns

Sender rank.

is_completed() - returns True if the operation has finished
wait() - will block the process until the operation is finished. is_completed() is guaranteed to return True once it returns.

When using the MPI backend, isend() and irecv() support non-overtaking, which has some guarantees on supporting message order. For more detail, see http://mpi-forum.org/docs/mpi-2.2/mpi22-report/node54.htm#Node54

torch.distributed.deprecated.isend(tensor, dst)¶

Sends a tensor asynchronously.

Parameters

tensor (Tensor) – Tensor to send.
dst (int) – Destination rank.

Returns

A distributed request object.

torch.distributed.deprecated.irecv(tensor, src)¶

Receives a tensor asynchronously.

Parameters

tensor (Tensor) – Tensor to fill with received data.
src (int) – Source rank.

Returns

A distributed request object.

Collective functions¶

torch.distributed.deprecated.broadcast(tensor, src, group=<object object>)¶

Broadcasts the tensor to the whole group.

tensor must have the same number of elements in all processes participating in the collective.

Parameters

tensor (Tensor) – Data to be sent if src is the rank of current process, and tensor to be used to save received data otherwise.
src (int) – Source rank.
group (optional) – Group of the collective.

torch.distributed.deprecated.all_reduce(tensor, op=<object object>, group=<object object>)¶

Reduces the tensor data across all machines in such a way that all get the final result.

After the call tensor will be bitwise identical in all processes.

Parameters

tensor (Tensor) – Input and output of the collective. The function operates in-place.
op (optional) – One of the values from torch.distributed.deprecated.reduce_op enum. Specifies an operation used for element-wise reductions.
group (optional) – Group of the collective.

torch.distributed.deprecated.reduce(tensor, dst, op=<object object>, group=<object object>)¶

Reduces the tensor data across all machines.

Only the process with rank dst is going to receive the final result.

Parameters

tensor (Tensor) – Input and output of the collective. The function operates in-place.
dst (int) – Destination rank
op (optional) – One of the values from torch.distributed.deprecated.reduce_op enum. Specifies an operation used for element-wise reductions.
group (optional) – Group of the collective.

torch.distributed.deprecated.all_gather(tensor_list, tensor, group=<object object>)¶

Gathers tensors from the whole group in a list.

Parameters

tensor_list (list[Tensor]) – Output list. It should contain correctly-sized tensors to be used for output of the collective.
tensor (Tensor) – Tensor to be broadcast from current process.
group (optional) – Group of the collective.

torch.distributed.deprecated.gather(tensor, **kwargs)¶

Gathers a list of tensors in a single process.

Parameters

tensor (Tensor) – Input tensor.
dst (int) – Destination rank. Required in all processes except the one that is receiveing the data.
gather_list (list[Tensor]) – List of appropriately-sized tensors to use for received data. Required only in the receiving process.
group (optional) – Group of the collective.

torch.distributed.deprecated.scatter(tensor, **kwargs)¶

Scatters a list of tensors to all processes in a group.

Each process will receive exactly one tensor and store its data in the tensor argument.

Parameters

tensor (Tensor) – Output tensor.
src (int) – Source rank. Required in all processes except the one that is sending the data.
scatter_list (list[Tensor]) – List of tensors to scatter. Required only in the process that is sending the data.
group (optional) – Group of the collective.

torch.distributed.deprecated.barrier(group=<object object>)¶

Synchronizes all processes.

This collective blocks processes until the whole group enters this function.

Parameters: group (optional) – Group of the collective.

Multi-GPU collective functions¶

If you have more than one GPU on each node, when using the NCCL backend, broadcast_multigpu() all_reduce_multigpu() reduce_multigpu() and all_gather_multigpu() support distributed collective operations among multiple GPUs within each node. These functions can potentially improve the overall distributed training performance and be easily used by passing a list of tensors. Each Tensor in the passed tensor list needs to be on a separate GPU device of the host where the function is called. Note that the length of the tensor list needs to be identical among all the distributed processes. Also note that currently the multi-GPU collective functions are only supported by the NCCL backend.

Code running on Node 0

import torch
import torch.distributed.deprecated as dist

dist.init_process_group(backend="nccl",
                        init_method="file:///distributed_test",
                        world_size=2,
                        rank=0)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
    tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))

dist.all_reduce_multigpu(tensor_list)

Code running on Node 1

import torch
import torch.distributed.deprecated as dist

dist.init_process_group(backend="nccl",
                        init_method="file:///distributed_test",
                        world_size=2,
                        rank=1)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
    tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))

dist.all_reduce_multigpu(tensor_list)

After the call, all 16 tensors on the two nodes will have the all-reduced value of 16

torch.distributed.deprecated.broadcast_multigpu(tensor_list, src, group=<object object>)¶

Broadcasts the tensor to the whole group with multiple GPU tensors per node.

tensor must have the same number of elements in all the GPUs from all processes participating in the collective. each tensor in the list must be on a different GPU.

Note

Only NCCL backend is currently supported. tensor_list should only contain GPU tensors.

Parameters

tensor_list (List[Tensor]) – Tensors that participate in the collective operation. if src is the rank, then the first element of tensor_list (tensor_list[0]) will be broadcasted to all other tensors (on different GPUs) in the src process and all tensors in tensor_list of other non-src processes. You also need to make sure that len(tensor_list) is the same for all the distributed processes calling this function.
src (int) – Source rank.
group (optional) – Group of the collective.

torch.distributed.deprecated.all_reduce_multigpu(tensor_list, op=<object object>, group=<object object>)¶

Reduces the tensor data across all machines in such a way that all get the final result. This function reduces a number of tensors on every node, while each tensor resides on a different GPU. Therefore, the input tensor in the tensor list needs to be GPU tensors. Also, each tensor in the tensor list needs to reside on a different GPU.

After the call, all tensors in tensor_list will be bitwise identical in all processes.

Note

Only NCCL backend is currently supported. tensor_list should only contain GPU tensors.

Parameters

tensor_list (List[Tensor]) – List of input and output tensors of the collective. The function operates in-place and requires that each tensor to be a GPU tensor on different GPUs. You also need to make sure that len(tensor_list) is the same for all the distributed processes calling this function.
op (optional) – One of the values from torch.distributed.deprecated.reduce_op enum. Specifies an operation used for element-wise reductions.
group (optional) – Group of the collective.

torch.distributed.deprecated.reduce_multigpu(tensor_list, dst, op=<object object>, group=<object object>)¶

Reduces the tensor data on multiple GPUs across all machines. Each tensor in tensor_list should reside on a separate GPU.

Only the GPU of tensor_list[0] on the process with rank dst is going to receive the final result.

Note

Only NCCL backend is currently supported. tensor_list should only contain GPU tensors.

Parameters

tensor_list (List[Tensor]) – Input and output GPU tensors of the collective. The function operates in-place. You also need to make sure that len(tensor_list) is the same for all the distributed processes calling this function.
dst (int) – Destination rank
op (optional) – One of the values from torch.distributed.deprecated.reduce_op enum. Specifies an operation used for element-wise reductions.
group (optional) – Group of the collective.

torch.distributed.deprecated.all_gather_multigpu(output_tensor_lists, input_tensor_list, group=<object object>)¶

Gathers tensors from the whole group in a list. Each tensor in input_tensor_list should reside on a separate GPU.

Note

Only NCCL backend is currently supported. output_tensor_lists and input_tensor_list should only contain GPU tensors.

Parameters

output_tensor_lists (List[List[Tensor]]) – Output lists. It should contain correctly-sized tensors on each GPU to be used for output of the collective. e.g. output_tensor_lists[i] contains the all_gather result that resides on the GPU of input_tensor_list[i]. Note that each element of output_tensor_lists[i] has the size of world_size * len(input_tensor_list), since the function all gathers the result from every single GPU in the group. To interpret each element of output_tensor_list[i], note that input_tensor_list[j] of rank k will be appear in output_tensor_list[i][rank * world_size + j] Also note that len(output_tensor_lists), and the size of each element in output_tensor_lists (each element is a list, therefore len(output_tensor_lists[i])) need to be the same for all the distributed processes calling this function.
input_tensor_list (List[Tensor]) – List of tensors (on different GPUs) to be broadcast from current process. Note that len(input_tensor_list) needs to be the same for all the distributed processes calling this function.
group (optional) – Group of the collective.

Launch utility¶

The torch.distributed.deprecated package also provides a launch utility in torch.distributed.deprecated.launch.

torch.distributed.launch is a module that spawns up multiple distributed training processes on each of the training nodes.

How to use this module:

Single-Node multi-process distributed training

>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
           YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other
           arguments of your training script)

Multi-Node multi-process distributed training: (e.g. two nodes)

Node 1: (IP: 192.168.1.1, and has a free port: 1234)

>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
           --nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
           --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
           and all other arguments of your training script)

Node 2:

>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
           --nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
           --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
           and all other arguments of your training script)

To look up what optional arguments this module offers:

>>> python -m torch.distributed.launch --help

Important Notices:

Parsing the local_rank argument

>>> import argparse
>>> parser = argparse.ArgumentParser()
>>> parser.add_argument("--local_rank", type=int)
>>> args = parser.parse_args()

Set your device to local rank using either

>>> torch.cuda.set_device(arg.local_rank)  # before your code runs

>>> with torch.cuda.device(arg.local_rank):
>>>    # your code to run

torch.distributed.init_process_group(backend='YOUR BACKEND',
                                     init_method='env://')

model = torch.nn.parallel.DistributedDataParallel(model,
                                                  device_ids=[arg.local_rank],
                                                  output_device=arg.local_rank)

Warning

nonlinearity	gain
Linear / Identity	\(1\)
Conv{1,2,3}D	\(1\)
Sigmoid	\(1\)
Tanh	\(\frac{5}{3}\)
ReLU	\(\sqrt{2}\)
Leaky Relu	\(\sqrt{\frac{2}{1 + \text{negative\_slope}^2}}\)