pytorch新特性：用名称操作tensor

PyTorch现在允许用“”名称”命名张量指定的维数，每个tensor接受一个新的名称参数，该参数将名称与每个维度关联起来。这适用于大多数函数：

tensor
empty
ones
zeros
randn
rand

1.新建tensor时，创建维度名称

import torch
imgs = torch.randn(1, 2, 2, 3, names=('N', 'C', 'H', 'W'))
print(imgs.names)

--------out:
('N', 'C', 'H', 'W')

2.对命名tensor进行名称修改

# Method #1: set the .names attribute (this changes name in-place)
imgs.names = ['batch', 'channel', 'width', 'height']
print(imgs.names)

# Method #2: specify new names (this changes names out-of-place)
imgs = imgs.rename(channel='C', width='W', height='H')
print(imgs.names)

-------out:
('batch', 'channel', 'width', 'height')
('batch', 'C', 'W', 'H')

3.删除命名维度的tensor的名称

imgs = imgs.rename(None)
print(imgs.names)

-------out:
(None, None, None, None)

4.已命名张量 and 未命名张量

imgs = torch.randn(3, 1, 1, 2, names=('N', None, None, None))
print(imgs.names)

-------out:
('N', None, None, None)

如上述代码所示，已命名的张量可以与未命名的张量共存，所以我们需要一种很好的方式来编写与已命名张量和未命名张量都兼容的已命名张量感知代码。

使用tensor.refine_names(*names)来细化维度，并将未命名的dims提升为命名dims。

imgs = torch.randn(3, 1, 1, 2)
named_imgs = imgs.refine_names('N', 'C', 'H', 'W')
print(named_imgs.names)

# Refine the last two dims to 'H' and 'W'. In Python 2, use the string '...'
# instead of ...
named_imgs = imgs.refine_names(..., 'H', 'W')
print(named_imgs.names)

-------out:
('N', 'C', 'H', 'W')
(None, None, 'H', 'W')

5.使用命名维度操作tensor

output = named_imgs.sum('C')  # 把特征图按通道相加
print(output.names)

img0 = named_imgs.select('N', 0)  # 获取batch中的第一张图片
print(img0.names)

------out:
('N', 'H', 'W')
('C', 'H', 'W')

6.名称不同的维度不能进行直接计算

x = torch.randn(3, names=('X',))
y = torch.randn(3)
z = torch.randn(3, names=('Z',))
catch_error(lambda: x + z)

-------out:
 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.

7.命名维度可传递

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

-------out:
('X',)

8.命名维度防止出现维度操作失误的情况

如下述代码所示，若没有命名维度，最终会将per_batch_scale这个值乘到W维度上，从而发生难以预测的错误。而使用命名维度能够很好的避免这个问题。

imgs = torch.randn(2, 2, 2, 2, names=('N', 'C', 'H', 'W'))
per_batch_scale = torch.rand(2, names=('N',))
catch_error(lambda: imgs * per_batch_scale)

-------out:
Error when attempting to broadcast dims ['N', 'C', 'H', 'W'] and dims ['N']: dim 'W' and dim 'N' are at the same position from the right but do not match.

对于未命名的张量，若对多个维度进行相乘时，需要对维度进行扩充才能做到想要的操作。如下述代码所示：

imgs = torch.randn(2, 2, 2, 2)  # N, C, H, W
per_batch_scale = torch.rand(2)  # N

correct_result = imgs * per_batch_scale.view(2, 1, 1, 1)  # N, C, H, W
incorrect_result = imgs * per_batch_scale.expand_as(imgs)
assert not torch.allclose(correct_result, incorrect_result)

但如下述代码所示，有了命名的维度后，只要用tensor.align_as(other)这个函数很方便的进行维度的操作。

imgs = imgs.refine_names('N', 'C', 'H', 'W')
per_batch_scale = per_batch_scale.refine_names('N')

named_result = imgs * per_batch_scale.align_as(imgs)
# note: named tensors do not yet work with allclose
assert torch.allclose(named_result.rename(None), correct_result)

9.tensor矩阵相乘时命名维度的变化

markov_states = torch.randn(128, 5, names=('batch', 'D'))
transition_matrix = torch.randn(5, 5, names=('in', 'out'))

# Apply one transition
new_state = markov_states @ transition_matrix
print(new_state.names)

-------out:
('batch', 'out')

如上述代码所示，在进行tensor矩阵相乘时，不会进行命名维度检查，只要两个维度大小相同就会进行融合，命名维度也会进行上述变化。

10.flatten 和 unflatten

如下述代码所示，利用命名维度进行flatten，要比直接使用reshape,view,flatten要方便的多。

imgs = imgs.flatten(['C', 'H', 'W'], 'features')
print(imgs.names)

imgs = imgs.unflatten('features', (('C', 2), ('H', 2), ('W', 2)))
print(imgs.names)

-------out:
('N', 'features')
('N', 'C', 'H', 'W')