【pytorch】维度的扩展和压缩

• [:, None] 扩展维度
• … 省略之间的维度
• torch.unsqueeze(input, dim=None) 扩展维度
• torch.squeeze(input, dim=None) 压缩维度

import torch

tensor = torch.randn(3, 4)
print('tensor size:', tensor.size()) # tensor size: torch.Size([3, 4])

# 末尾维度的扩展与压缩
tensor_1 = tensor[:, :, None]
print('tensor_1 size', tensor_1.size()) # tensor_1 size torch.Size([3, 4, 1])
tensor_11 = tensor[..., None]
print('tensor_11 size', tensor_11.size()) # tensor_11 size torch.Size([3, 4, 1])

tensor1 = tensor.unsqueeze(-1)
print('tensor1 size', tensor1.size()) # tensor1 size torch.Size([3, 4, 1])

tensor_0 = tensor1[:,:,0]
print('tensor_0 size', tensor_0.size())  # tensor_0 size torch.Size([3, 4])

tensor0 = tensor1.squeeze(-1)
print('tensor0 size', tensor0.size())  # tensor0 size torch.Size([3, 4])

# 中间维度的扩展与压缩
tensor_2 = tensor[:, None, :]
print('tensor_2 size', tensor_2.size()) # tensor_2 size torch.Size([3, 1, 4])

tensor_22 = tensor[:, None]
print('tensor_22 size:', tensor_22.size()) # tensor_22 size: torch.Size([3, 1, 4])

tensor_222 = tensor[..., None, :]
print('tensor_222 size', tensor_222.size()) # tensor_222 size torch.Size([3, 1, 4])

tensor2 = tensor.unsqueeze(1)
print('tensor2 size', tensor2.size()) # tensor2 size torch.Size([3, 1, 4])

tensor_0 = tensor2[:,0,:]
print('tensor_0 size', tensor_0.size())  # tensor_0 size torch.Size([3, 4])

tensor0 = tensor2.squeeze(1)
print('tensor0 size', tensor0.size())  # tensor0 size torch.Size([3, 4])

# 扩展压缩两个维度
tensor_3 = tensor[:, None, None, :]
print('tensor_3 size', tensor_3.size())  # tensor_3 size torch.Size([3, 1, 1, 4])

tensor_0 = tensor_3[:,0, 0,:]
print('tensor_0 size', tensor_0.size())  # tensor_0 size torch.Size([3, 4])

【参考】
python numpy 里面的[:, None]是个什么鬼?[…, None]呢?