| .. currentmodule:: torch |
| |
| .. _named_tensors-doc: |
| |
| Named Tensors |
| ============= |
| |
| Named Tensors allow users to give explicit names to tensor dimensions. |
| In most cases, operations that take dimension parameters will accept |
| dimension names, avoiding the need to track dimensions by position. |
| In addition, named tensors use names to automatically check that APIs |
| are being used correctly at runtime, providing extra safety. Names can |
| also be used to rearrange dimensions, for example, to support |
| "broadcasting by name" rather than "broadcasting by position". |
| |
| |
| .. warning:: |
| The named tensor API is a prototype feature and subject to change. |
| |
| Creating named tensors |
| ---------------------- |
| |
| Factory functions now take a new :attr:`names` argument that associates a name |
| with each dimension. |
| |
| :: |
| |
| >>> torch.zeros(2, 3, names=('N', 'C')) |
| tensor([[0., 0., 0.], |
| [0., 0., 0.]], names=('N', 'C')) |
| |
| Named dimensions, like regular Tensor dimensions, are ordered. |
| ``tensor.names[i]`` is the name of dimension ``i`` of ``tensor``. |
| |
| The following factory functions support named tensors: |
| |
| - :func:`torch.empty` |
| - :func:`torch.rand` |
| - :func:`torch.randn` |
| - :func:`torch.ones` |
| - :func:`torch.tensor` |
| - :func:`torch.zeros` |
| |
| Named dimensions |
| ---------------- |
| |
| See :attr:`~Tensor.names` for restrictions on tensor names. |
| |
| Use :attr:`~Tensor.names` to access the dimension names of a tensor and |
| :meth:`~Tensor.rename` to rename named dimensions. |
| |
| :: |
| |
| >>> imgs = torch.randn(1, 2, 2, 3 , names=('N', 'C', 'H', 'W')) |
| >>> imgs.names |
| ('N', 'C', 'H', 'W') |
| |
| >>> renamed_imgs = imgs.rename(H='height', W='width') |
| >>> renamed_imgs.names |
| ('N', 'C', 'height', 'width) |
| |
| |
| Named tensors can coexist with unnamed tensors; named tensors are instances of |
| :class:`torch.Tensor`. Unnamed tensors have ``None``-named dimensions. Named |
| tensors do not require all dimensions to be named. |
| |
| :: |
| |
| >>> imgs = torch.randn(1, 2, 2, 3 , names=(None, 'C', 'H', 'W')) |
| >>> imgs.names |
| (None, 'C', 'H', 'W') |
| |
| Name propagation semantics |
| -------------------------- |
| |
| Named tensors use names to automatically check that APIs are being called |
| correctly at runtime. This occurs in a process called *name inference*. |
| More formally, name inference consists of the following two steps: |
| |
| - **Check names**: an operator may perform automatic checks at runtime that |
| check that certain dimension names must match. |
| - **Propagate names**: name inference propagates names to output tensors. |
| |
| All operations that support named tensors propagate names. |
| |
| :: |
| |
| >>> x = torch.randn(3, 3, names=('N', 'C')) |
| >>> x.abs().names |
| ('N', 'C') |
| |
| |
| .. _match_semantics-doc: |
| |
| match semantics |
| ^^^^^^^^^^^^^^^ |
| |
| Two names *match* if they are equal (string equality) or if at least one is ``None``. |
| Nones are essentially a special "wildcard" name. |
| |
| ``unify(A, B)`` determines which of the names ``A`` and ``B`` to propagate to the outputs. |
| It returns the more *specific* of the two names, if they match. If the names do not match, |
| then it errors. |
| |
| .. note:: |
| In practice, when working with named tensors, one should avoid having unnamed |
| dimensions because their handling can be complicated. It is recommended to lift |
| all unnamed dimensions to be named dimensions by using :meth:`~Tensor.refine_names`. |
| |
| |
| Basic name inference rules |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| |
| Let's see how ``match`` and ``unify`` are used in name inference in the case of |
| adding two one-dim tensors with no broadcasting. |
| |
| :: |
| |
| x = torch.randn(3, names=('X',)) |
| y = torch.randn(3) |
| z = torch.randn(3, names=('Z',)) |
| |
| **Check names**: check that the names of the two tensors *match*. |
| |
| For the following examples: |
| |
| :: |
| |
| >>> # x + y # match('X', None) is True |
| >>> # x + z # match('X', 'Z') is False |
| >>> # x + x # match('X', 'X') is True |
| |
| >>> x + z |
| Error when attempting to broadcast dims ['X'] and dims ['Z']: dim 'X' and dim 'Z' are at the same position from the right but do not match. |
| |
| **Propagate names**: *unify* the names to select which one to propagate. |
| In the case of ``x + y``, ``unify('X', None) = 'X'`` because ``'X'`` is more |
| specific than ``None``. |
| |
| :: |
| |
| >>> (x + y).names |
| ('X',) |
| >>> (x + x).names |
| ('X',) |
| |
| For a comprehensive list of name inference rules, see :ref:`name_inference_reference-doc`. |
| Here are two common operations that may be useful to go over: |
| |
| - Binary arithmetic ops: :ref:`unifies_names_from_inputs-doc` |
| - Matrix multiplication ops: :ref:`contracts_away_dims-doc` |
| |
| Explicit alignment by names |
| --------------------------- |
| |
| Use :meth:`~Tensor.align_as` or :meth:`~Tensor.align_to` to align tensor dimensions |
| by name to a specified ordering. This is useful for performing "broadcasting by names". |
| |
| :: |
| |
| # This function is agnostic to the dimension ordering of `input`, |
| # as long as it has a `C` dimension somewhere. |
| def scale_channels(input, scale): |
| scale = scale.refine_names('C') |
| return input * scale.align_as(input) |
| |
| >>> num_channels = 3 |
| >>> scale = torch.randn(num_channels, names=('C',)) |
| >>> imgs = torch.rand(3, 3, 3, num_channels, names=('N', 'H', 'W', 'C')) |
| >>> more_imgs = torch.rand(3, num_channels, 3, 3, names=('N', 'C', 'H', 'W')) |
| >>> videos = torch.randn(3, num_channels, 3, 3, 3, names=('N', 'C', 'H', 'W', 'D') |
| |
| >>> scale_channels(imgs, scale) |
| >>> scale_channels(more_imgs, scale) |
| >>> scale_channels(videos, scale) |
| |
| Manipulating dimensions |
| ----------------------- |
| |
| Use :meth:`~Tensor.align_to` to permute large amounts of dimensions without |
| mentioning all of them as in required by :meth:`~Tensor.permute`. |
| |
| :: |
| |
| >>> tensor = torch.randn(2, 2, 2, 2, 2, 2) |
| >>> named_tensor = tensor.refine_names('A', 'B', 'C', 'D', 'E', 'F') |
| |
| # Move the F (dim 5) and E dimension (dim 4) to the front while keeping |
| # the rest in the same order |
| >>> tensor.permute(5, 4, 0, 1, 2, 3) |
| >>> named_tensor.align_to('F', 'E', ...) |
| |
| Use :meth:`~Tensor.flatten` and :meth:`~Tensor.unflatten` to flatten and unflatten |
| dimensions, respectively. These methods are more verbose than :meth:`~Tensor.view` |
| and :meth:`~Tensor.reshape`, but have more semantic meaning to someone reading the code. |
| |
| :: |
| |
| >>> imgs = torch.randn(32, 3, 128, 128) |
| >>> named_imgs = imgs.refine_names('N', 'C', 'H', 'W') |
| |
| >>> flat_imgs = imgs.view(32, -1) |
| >>> named_flat_imgs = named_imgs.flatten(['C', 'H', 'W'], 'features') |
| >>> named_flat_imgs.names |
| ('N', 'features') |
| |
| >>> unflattened_imgs = imgs.view(32, 3, 128, 128) |
| >>> unflattened_named_imgs = named_flat_imgs.unflatten( |
| 'features', [('C', 3), ('H', 128), ('W', 128)]) |
| |
| .. _named_tensors_autograd-doc: |
| |
| Autograd support |
| ---------------- |
| |
| Autograd currently supports named tensors in a limited manner: autograd ignores |
| names on all tensors. Gradient computation is still correct but we lose the |
| safety that names give us. |
| |
| :: |
| |
| >>> x = torch.randn(3, names=('D',)) |
| >>> weight = torch.randn(3, names=('D',), requires_grad=True) |
| >>> loss = (x - weight).abs() |
| >>> grad_loss = torch.randn(3) |
| >>> loss.backward(grad_loss) |
| >>> weight.grad # Unnamed for now. Will be named in the future |
| tensor([-1.8107, -0.6357, 0.0783]) |
| |
| >>> weight.grad.zero_() |
| >>> grad_loss = grad_loss.refine_names('C') |
| >>> loss = (x - weight).abs() |
| # Ideally we'd check that the names of loss and grad_loss match but we don't yet. |
| >>> loss.backward(grad_loss) |
| >>> weight.grad |
| tensor([-1.8107, -0.6357, 0.0783]) |
| |
| Currently supported operations and subsystems |
| --------------------------------------------- |
| |
| Operators |
| ^^^^^^^^^ |
| |
| See :ref:`name_inference_reference-doc` for a full list of the supported torch and |
| tensor operations. We do not yet support the following that is not covered by the link: |
| |
| - indexing, advanced indexing. |
| |
| For ``torch.nn.functional`` operators, we support the following: |
| |
| - :func:`torch.nn.functional.relu` |
| - :func:`torch.nn.functional.softmax` |
| - :func:`torch.nn.functional.log_softmax` |
| - :func:`torch.nn.functional.tanh` |
| - :func:`torch.nn.functional.sigmoid` |
| - :func:`torch.nn.functional.dropout` |
| |
| Subsystems |
| ^^^^^^^^^^ |
| |
| Autograd is supported, see :ref:`named_tensors_autograd-doc`. |
| Because gradients are currently unnamed, optimizers may work but are untested. |
| |
| NN modules are currently unsupported. This can lead to the following when calling |
| modules with named tensor inputs: |
| |
| - NN module parameters are unnamed, so outputs may be partially named. |
| - NN module forward passes have code that don't support named tensors and will |
| error out appropriately. |
| |
| We also do not support the following subsystems, though some may work out |
| of the box: |
| |
| - distributions |
| - serialization (:func:`torch.load`, :func:`torch.save`) |
| - multiprocessing |
| - JIT |
| - distributed |
| - ONNX |
| |
| If any of these would help your use case, please |
| `search if an issue has already been filed <https://github.com/pytorch/pytorch/issues?q=is%3Aopen+is%3Aissue+label%3A%22module%3A+named+tensor%22>`_ |
| and if not, `file one <https://github.com/pytorch/pytorch/issues/new/choose>`_. |
| |
| Named tensor API reference |
| -------------------------- |
| |
| In this section please find the documentation for named tensor specific APIs. |
| For a comprehensive reference for how names are propagated through other PyTorch |
| operators, see :ref:`name_inference_reference-doc`. |
| |
| .. class:: Tensor() |
| :noindex: |
| |
| .. autoattribute:: names |
| .. automethod:: rename |
| .. automethod:: rename_ |
| .. automethod:: refine_names |
| |
| .. automethod:: align_as |
| .. automethod:: align_to |
| |
| .. py:method:: flatten(dims, out_dim) -> Tensor |
| :noindex: |
| |
| Flattens :attr:`dims` into a single dimension with name :attr:`out_dim`. |
| |
| All of `dims` must be consecutive in order in the :attr:`self` tensor, |
| but not necessary contiguous in memory. |
| |
| Examples:: |
| |
| >>> imgs = torch.randn(32, 3, 128, 128, names=('N', 'C', 'H', 'W')) |
| >>> flat_imgs = imgs.flatten(['C', 'H', 'W'], 'features') |
| >>> flat_imgs.names, flat_imgs.shape |
| (('N', 'features'), torch.Size([32, 49152])) |
| |
| .. warning:: |
| The named tensor API is experimental and subject to change. |
