评价此页

命名张量#

创建日期:2019年10月8日 | 最近更新:2025年6月14日

命名张量允许用户为张量维度赋予显式名称。在大多数情况下,接受维度参数的操作将接受维度名称,从而避免了按位置跟踪维度的需要。此外,命名张量使用名称在运行时自动检查 API 是否被正确使用,从而提供了额外的安全性。名称还可用于重新排列维度,例如,支持“按名称广播”而不是“按位置广播”。

警告

The named tensor API is a prototype feature and subject to change.

创建命名张量#

工厂函数现在接受一个新的 names 参数,该参数将名称与每个维度关联起来。

    >>> torch.zeros(2, 3, names=('N', 'C'))
    tensor([[0., 0., 0.],
            [0., 0., 0.]], names=('N', 'C'))

命名维度与常规张量维度一样,是有序的。tensor.names[i]tensori 个维度的名称。

以下工厂函数支持命名张量

命名维度#

有关张量名称的限制,请参阅 names

使用 names 访问张量的维度名称,使用 rename() 重命名命名维度。

    >>> 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)

命名张量可以与未命名张量共存;命名张量是 torch.Tensor 的实例。未命名张量的维度名称为 None。命名张量不需要所有维度都被命名。

    >>> imgs = torch.randn(1, 2, 2, 3 , names=(None, 'C', 'H', 'W'))
    >>> imgs.names
    (None, 'C', 'H', 'W')

名称传播语义#

命名张量使用名称在运行时自动检查 API 调用是否正确。这发生在一个称为“名称推导”(name inference)的过程中。更正式地说,名称推导由以下两个步骤组成

  • 检查名称:算子可以在运行时执行自动检查,以验证某些维度名称必须匹配。

  • 传播名称:名称推导将名称传播到输出张量。

所有支持命名张量的操作都会传播名称。

    >>> x = torch.randn(3, 3, names=('N', 'C'))
    >>> x.abs().names
    ('N', 'C')

匹配语义#

如果两个名称相等(字符串相等)或者其中至少有一个是 None,则这两个名称“匹配”。None 本质上是一个特殊的“通配符”名称。

unify(A, B) 决定将 AB 中的哪一个名称传播到输出。如果它们匹配,它会返回两者中更“具体”的一个名称。如果名称不匹配,则会报错。

注意

在实践中,使用命名张量时应避免使用未命名维度,因为它们的处理可能会很复杂。建议使用 refine_names() 将所有未命名维度提升为命名维度。

基本名称推导规则#

让我们看看在两个一维张量相加且不涉及广播的情况下,matchunify 是如何用于名称推导的。

    x = torch.randn(3, names=('X',))
    y = torch.randn(3)
    z = torch.randn(3, names=('Z',))

检查名称:检查两个张量的名称是否“匹配”。

对于以下示例

    >>> # 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.

传播名称:“统一”(unify)名称以选择传播哪一个。在 x + y 的情况下,unify('X', None) = 'X',因为 'X'None 更具体。

    >>> (x + y).names
    ('X',)
    >>> (x + x).names
    ('X',)

有关名称推导规则的完整列表,请参阅 命名张量算子覆盖范围。这里有两个常用的操作,可能值得复习:

按名称显式对齐#

使用 align_as()align_to() 按名称将张量维度对齐到指定的顺序。这对于执行“按名称广播”非常有用。

    # 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)

操作维度#

使用 align_to() 可以置换大量维度,而不需要像 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', ...)

使用 flatten()unflatten() 分别用于展平和恢复维度。这些方法比 view()reshape() 更冗长,但对代码阅读者来说更具语义化。

    >>> 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_named_imgs = named_flat_imgs.unflatten('features', [('C', 3), ('H', 128), ('W', 128)])
    >>> unflattened_named_imgs.names
    ('N', 'C', 'H', 'W')

Autograd 支持#

Autograd 目前以有限的方式支持命名张量:Autograd 会忽略所有张量上的名称。梯度计算仍然是正确的,但我们失去了名称带来的安全性。

    >>> 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])

目前支持的操作和子系统#

算子#

有关支持的 torch 和 tensor 操作的完整列表,请参阅 命名张量算子覆盖范围。目前我们还不支持链接中未涵盖的以下内容:

  • 索引、高级索引。

对于 torch.nn.functional 算子,我们支持以下内容:

子系统#

支持 Autograd,请参阅 Autograd 支持。由于梯度目前是未命名的,优化器可能可以工作,但尚未经过测试。

目前不支持 NN 模块。使用命名张量输入调用模块时可能会导致以下情况:

  • NN 模块参数是未命名的,因此输出可能只是部分命名的。

  • NN 模块的前向传递代码不支持命名张量,并会产生相应的错误。

我们也不支持以下子系统,尽管有些可能开箱即用:

  • distributions(分布)

  • serialization(序列化:torch.load(), torch.save()

  • multiprocessing(多进程)

  • JIT

  • 分布式

  • ONNX

如果其中任何一个对您的用例有帮助,请搜索是否已提交过相关 Issue,如果没有,请提交一个

命名张量 API 参考#

在本节中,请查找命名张量特定 API 的文档。有关名称如何通过其他 PyTorch 算子传播的综合参考,请参阅 命名张量算子覆盖范围

class torch.Tensor
names#

存储此张量每个维度的名称。

names[idx] 对应于张量维度 idx 的名称。如果维度已命名,则名称为字符串;如果维度未命名,则为 None

维度名称可以包含字符或下划线。此外,维度名称必须是有效的 Python 变量名(即,不能以下划线开头)。

张量不能有两个具有相同名称的命名维度。

警告

命名张量 API 处于实验阶段,可能会发生变化。

rename(*names, **rename_map)[源码]#

重命名 self 的维度名称。

主要有两种用法:

self.rename(**rename_map) 返回张量的一个视图,其维度按照映射 rename_map 中的指定进行重命名。

self.rename(*names) 返回张量的一个视图,使用 names 按位置重命名所有维度。使用 self.rename(None) 可以丢弃张量上的名称。

不能同时指定位置参数 names 和关键字参数 rename_map

示例

>>> imgs = torch.rand(2, 3, 5, 7, names=('N', 'C', 'H', 'W'))
>>> renamed_imgs = imgs.rename(N='batch', C='channels')
>>> renamed_imgs.names
('batch', 'channels', 'H', 'W')

>>> renamed_imgs = imgs.rename(None)
>>> renamed_imgs.names
(None, None, None, None)

>>> renamed_imgs = imgs.rename('batch', 'channel', 'height', 'width')
>>> renamed_imgs.names
('batch', 'channel', 'height', 'width')

警告

命名张量 API 处于实验阶段,可能会发生变化。

rename_(*names, **rename_map)[源码]#

rename() 的就地(in-place)版本。

refine_names(*names)[源码]#

根据 names 细化 self 的维度名称。

细化(Refining)是重命名的一种特殊情况,它“提升”了未命名维度。一个 None 维度可以被细化为具有任何名称;一个已命名的维度只能被细化为具有相同的名称。

由于命名张量可以与未命名张量共存,因此细化名称提供了一种很好的方式来编写既适用于命名张量又适用于未命名张量的命名张量感知代码。

names 最多可以包含一个省略号(...)。省略号会贪婪地扩展;它被就地展开,使用来自 self.names 对应索引的名称,将 names 填充到与 self.dim() 相同的长度。

Python 2 不支持 Ellipsis 对象,但可以使用字符串字面量('...')代替。

参数:

names (iterable of str) – 输出张量的期望名称。最多可包含一个省略号。

示例

>>> imgs = torch.randn(32, 3, 128, 128)
>>> named_imgs = imgs.refine_names('N', 'C', 'H', 'W')
>>> named_imgs.names
('N', 'C', 'H', 'W')

>>> tensor = torch.randn(2, 3, 5, 7, 11)
>>> tensor = tensor.refine_names('A', ..., 'B', 'C')
>>> tensor.names
('A', None, None, 'B', 'C')

警告

命名张量 API 处于实验阶段,可能会发生变化。

align_as(other) Tensor#

置换 self 张量的维度以匹配 other 张量中的维度顺序,并为任何新名称添加大小为 1 的维度。

此操作对于按名称进行显式广播非常有用(见示例)。

要使用此方法,self 的所有维度必须已命名。结果张量是原始张量的视图。

self 的所有维度名称必须存在于 other.names 中。other 可能包含不在 self.names 中的命名维度;对于每一个这些新名称,输出张量都有一个大小为 1 的维度。

要将张量对齐到特定顺序,请使用 align_to()

示例

# Example 1: Applying a mask
>>> mask = torch.randint(2, [127, 128], dtype=torch.bool).refine_names('W', 'H')
>>> imgs = torch.randn(32, 128, 127, 3, names=('N', 'H', 'W', 'C'))
>>> imgs.masked_fill_(mask.align_as(imgs), 0)


# Example 2: Applying a per-channel-scale
>>> 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(32, 128, 128, num_channels, names=('N', 'H', 'W', 'C'))
>>> more_imgs = torch.rand(32, num_channels, 128, 128, names=('N', 'C', 'H', 'W'))
>>> videos = torch.randn(3, num_channels, 128, 128, 128, names=('N', 'C', 'H', 'W', 'D'))

# scale_channels is agnostic to the dimension order of the input
>>> scale_channels(imgs, scale)
>>> scale_channels(more_imgs, scale)
>>> scale_channels(videos, scale)

警告

命名张量 API 处于实验阶段,可能会发生变化。

align_to(*names)[源码]#

置换 self 张量的维度以匹配 names 中指定的顺序,并为任何新名称添加大小为 1 的维度。

要使用此方法,self 的所有维度必须已命名。结果张量是原始张量的视图。

self 的所有维度名称必须存在于 names 中。names 可能包含不在 self.names 中的额外名称;对于每一个这些新名称,输出张量都有一个大小为 1 的维度。

names 最多可以包含一个省略号(...)。省略号会被展开为 self 中所有未在 names 中提及的维度名称,并保持它们在 self 中出现的顺序。

Python 2 不支持 Ellipsis 对象,但可以使用字符串字面量('...')代替。

参数:

names (iterable of str) – 输出张量期望的维度顺序。最多可包含一个省略号,该省略号会扩展为 self 中所有未提及的维度名称。

示例

>>> tensor = torch.randn(2, 2, 2, 2, 2, 2)
>>> named_tensor = tensor.refine_names('A', 'B', 'C', 'D', 'E', 'F')

# Move the F and E dims to the front while keeping the rest in order
>>> named_tensor.align_to('F', 'E', ...)

警告

命名张量 API 处于实验阶段,可能会发生变化。

flatten(dims, out_dim) Tensor

dims 展平为一个名称为 out_dim 的单一维度。

self 张量中,所有 dims 必须是连续的顺序,但在内存中不必是连续的。

示例

>>> 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]))

警告

命名张量 API 处于实验阶段,可能会发生变化。