命名张量 ¶
命名张量允许用户为张量维度指定明确的名称。在大多数情况下,接受维度参数的操作将接受维度名称,从而避免需要通过位置跟踪维度。此外,命名张量使用名称在运行时自动检查 API 是否被正确使用,提供额外的安全性。名称还可以用于重新排列维度,例如,支持“按名称广播”而不是“按位置广播”。
警告
命名张量 API 是一个原型功能,可能会发生变化。
创建命名张量
工厂函数现在接受一个新参数 names
,该参数将每个维度与一个名称关联。
>>> torch.zeros(2, 3, names=('N', 'C'))
tensor([[0., 0., 0.],
[0., 0., 0.]], names=('N', 'C'))
命名维度,就像常规张量维度一样,是有序的。 tensor.names[i]
是 i
的名称,属于 tensor
。
以下工厂函数支持命名张量:
命名维度
请参阅 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 调用是否正确。这发生在名为名称推理的过程中。更正式地说,名称推理包括以下两个步骤:
检查名称:操作符可以在运行时执行自动检查,以确保某些维度名称必须匹配。
传播名称:名称推理将名称传播到输出张量。
所有支持命名张量的操作都会传播名称。
>>> x = torch.randn(3, 3, names=('N', 'C'))
>>> x.abs().names
('N', 'C')
匹配语义 ¶
两个名称匹配,如果它们相等(字符串相等)或者至少有一个是 None
。None 实质上是一个特殊的“通配符”名称。
unify(A, B)
决定将哪个名称 A
和 B
传播到输出。如果它们匹配,则返回两个名称中更具体的一个。如果名称不匹配,则报错。
注意
在实践中,当与命名张量一起工作时,应避免使用未命名的维度,因为它们的处理可能很复杂。建议使用 refine_names()
将所有未命名的维度提升为命名维度。
基本命名推断规则 ¶
让我们看看在添加两个没有广播的一维张量时,如何使用 match
和 unify
进行命名推断。
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.
传播名称:统一名称以选择要传播的名称。在 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')
自动微分支持
自动微分目前以有限的方式支持命名张量:自动微分忽略了所有张量的名称。梯度计算仍然正确,但我们失去了名称带来的安全性。
>>> 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 和张量运算的完整列表。我们尚未支持以下未在链接中涵盖的内容:
索引、高级索引。
对于 torch.nn.functional
操作符,我们支持以下内容:
子系统 §
支持自动微分,请参阅自动微分支持。由于梯度目前未命名,优化器可能可用但未经测试。
目前不支持神经网络模块。当使用具有命名张量输入的模块时,可能会导致以下情况:
神经网络模块的参数未命名,因此输出可能部分命名。
神经网络模块的前向传播代码不支持命名张量,将适当地报错。
我们也不支持以下子系统,尽管某些子系统可能直接运行:
分布式
序列化(
torch.load()
,torch.save()
)多进程
JIT
分布式
ONNX
如果这些中的任何一项能帮助您的用例,请搜索是否已提交相关问题,如果没有,请提交一个问题。
命名张量 API 参考
在本节中,您可以找到命名张量特定 API 的文档。有关命名如何在其他 PyTorch 运算符中传播的全面参考,请参阅命名张量运算符覆盖范围。
- class torch.Tensor
- names
存储每个张量维度的名称。
names[idx]
对应于张量维度idx
的名称。名称可以是字符串,如果维度有名称,或者None
,如果维度没有名称。维度名称可以包含字符或下划线。此外,维度名称必须是一个有效的 Python 变量名(即不以下划线开头)。
张量不能有两个具有相同名称的命名维度。
警告
命名张量 API 是实验性的,可能会发生变化。
- rename(*names, **rename_map)[source][source]¶
重命名
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 是实验性的,可能会更改。
- 重命名_(*names, **rename_map)[source][source] ¶
rename()
的就地版本 .
- refine_names(*names)[source][source]¶
根据
names
精炼self
的维度名称 .精炼是重命名的一种特殊情况,可以“提升”未命名的维度。一个
None
维度可以被精炼为任何名称;一个已命名的维度只能被精炼为相同的名称。由于命名张量可以与未命名的张量共存,精炼名称为编写既适用于命名张量又适用于未命名张量的命名张量感知代码提供了一种很好的方式。
names
可能包含最多一个省略号(...
)。省略号贪婪地展开;它就地展开以填充names
,使其长度与self.dim()
相同,并使用self.names
相应索引的名称。Python 2 不支持省略号,但可以使用字符串字面量代替(
'...'
)。- 参数:
names(可迭代字符串)- 输出张量的期望名称。最多可包含一个省略号。
示例:
>>> 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) → 张量
将
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)[source][source]¶
将
self
张量的维度重新排列以匹配names
中指定的顺序,为任何新的名称添加大小为 1 的维度。self
的所有维度都必须命名才能使用此方法。结果张量是原始张量的视图。self
中所有维度名称都必须存在于names
中。names
可以包含不在self.names
中的额外名称;输出张量对于这些新名称中的每一个都有一个大小为 1 的维度。names
可以包含最多一个省略号(...
)。省略号扩展为等于self
中未提及的所有维度名称,顺序与它们在names
中出现的顺序相同。Python 2 不支持省略号,但可以使用字符串字面量代替(
'...'
)。- 参数:
names(字符串的迭代器)- 输出张量所需的维度顺序。可以包含最多一个省略号,该省略号扩展为
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
的单维度。所有 dims 在
self
张量中必须按顺序连续,但在内存中不必连续。示例:
>>> 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 是实验性的,可能会发生变化。