下面使用一个二维矩阵看下dim不同时呈现出的效果:
# 创建一个3*4的全1二维tensor
a = torch.ones(3,4)
'''
运行结果
tensor([[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]])
'''
在0维度上插入一个维度,可以看到现在a的形状变为[1, 3, 4],第0维度的大小默认是1
a = a.unsqueeze(0)
print(a.shape)
'''
运行结果
tensor([[[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]]])
torch.Size([1, 3, 4])
'''
在最后一个维度上插入一个维度,形状变为[3, 4, 1]
a = a.unsqueeze(a.dim())
print(a.shape)
'''
运行结果
tensor([[[1.],
[1.],
[1.],
[1.]],
[[1.],
[1.],
[1.],
[1.]],
[[1.],
[1.],
[1.],
[1.]]])
torch.Size([3, 4, 1])
'''