pytorch 中的 unsqueeze() 函数详解
本文将详细解释 PyTorch 中 unsqueeze() 函数的功能和用法。unsqueeze() 函数可以为张量添加一个维度,其大小为 1。这在处理神经网络中的数据时非常有用,例如将批大小为 1 的样本添加到批处理维度中。

unsqueeze() 函数参数:
- input (Tensor): 输入张量。 支持 int、float、complex 或 bool 类型。
- dim (int): 要插入新维度的维度索引。 dim 可以是负数,表示从后往前计数。例如,dim=-1 表示在最后一个维度之后添加新维度。
unsqueeze() 函数用法:
unsqueeze() 函数可以以两种方式使用:
-
作为 torch.unsqueeze() 函数调用:
import torch tensor = torch.tensor([1, 2, 3]) new_tensor = torch.unsqueeze(input=tensor, dim=0) # 在维度 0 添加一个维度 print(new_tensor) # 输出: tensor([[1, 2, 3]])
-
作为张量的方法调用:
import torch tensor = torch.tensor([1, 2, 3]) new_tensor = tensor.unsqueeze(0) # 在维度 0 添加一个维度 print(new_tensor) # 输出: tensor([[1, 2, 3]])
示例:
以下示例演示了 unsqueeze() 函数在不同维度上添加新维度的效果:
import torch
my_tensor = torch.tensor([[0, 1, 2],
[3, 4, 5],
[6, 7, 8],
[10, 11, 12]])
print("原始张量:
", my_tensor)
print("
在维度 0 添加维度:
", torch.unsqueeze(input=my_tensor, dim=0))
print("
在维度 1 添加维度:
", torch.unsqueeze(input=my_tensor, dim=1))
print("
在维度 2 添加维度:
", torch.unsqueeze(input=my_tensor, dim=2))
print("
在维度 -1 添加维度:
", torch.unsqueeze(input=my_tensor, dim=-1)) # 等同于dim=2
# 测试不同数据类型
my_tensor_float = torch.tensor([[0., 1., 2.], [3., 4., 5.]])
my_tensor_complex = torch.tensor([[0.+0.j, 1.+0.j], [2.+0.j, 3.+0.j]])
my_tensor_bool = torch.tensor([[True, False], [False, True]])
print("
浮点数张量在维度0添加维度:
", torch.unsqueeze(input=my_tensor_float, dim=0))
print("
复数张量在维度0添加维度:
", torch.unsqueeze(input=my_tensor_complex, dim=0))
print("
布尔型张量在维度0添加维度:
", torch.unsqueeze(input=my_tensor_bool, dim=0))
总结:
unsqueeze() 函数是 PyTorch 中一个非常实用的工具,可以方便地为张量添加维度,从而适应不同的模型输入要求。 记住 dim 参数的正负索引方式,可以灵活地控制新维度的插入位置。
以上就是在 PyTorch 中解压的详细内容,更多请关注知识资源分享宝库其它相关文章!







发表评论:
◎欢迎参与讨论,请在这里发表您的看法、交流您的观点。