【pytorch算子】torch.squeeze 和 torch.unsqueeze 的用法

最新推荐文章于 2024-02-01 14:21:45 发布

ctrl A_ctrl C_ctrl V

最新推荐文章于 2024-02-01 14:21:45 发布

阅读量1.1k

点赞数 1

分类专栏： # PyTorch算子/函数用法文章标签： python pytorch 深度学习机器学习人工智能

本文链接：https://blog.csdn.net/qq_43799400/article/details/119679063

版权

PyTorch算子/函数用法专栏收录该内容

12 篇文章 3 订阅

订阅专栏

本文介绍了PyTorch中torch.squeeze()和unsqueeze()两个函数的用法。squeeze()用于移除Tensor中维度为1的元素，unsqueeze()则在指定位置插入一个维度为1的新维度。通过代码示例展示了这两个函数如何改变Tensor的形状，并解释了在不同情况下的操作效果。

摘要由CSDN通过智能技术生成

torch.squeeze() 用法：

对 Tensor 进行降维，去掉维数为 1 的的维度。

a.squeeze() 是将 a 中所有为 1 的维度删掉，不为 1 的维度没有影响。
a.squeeze(n) 是将 a 中第 n 个维度中维数为 1 的维度删掉，即若 dim(n) = 1，则删掉第 n 维；若 dim(n) != 1，则保留第 n 维，不做任何操作。

torch.unsqueeze() 用法：

对 Tensor 进行扩维，扩大的维数为 1 。

a.squeeze(n) 是将 a 中第 n 个维度的维数改为 1 ，其他维度依次增加 1，即扩充第 n 维。

代码示例

a=torch.tensor([
    [[1,2,3,5]],
    [[2,3,6,5]],
    [[5,9,9,6]]])      # a.shape=torch.Size([3, 1, 4])
b=a.squeeze()    # 去掉维度为1的维度， b.shape=torch.Size([3, 4])
c=b.unsqueeze(0)   # 扩充第0维， b.shape=torch.Size([1, 3, 4])
c=b.unsqueeze(1)   # 扩充第1维， b.shape=torch.Size([3, 1, 4])
c=b.unsqueeze(2)   # 扩充第2维， b.shape=torch.Size([3, 4, 1])


d=torch.tensor([[[1,2],[2,3]],
                [[5,6],[5,6]],
                [[7,5],[7,5]]])    # d.shape=torch.Size([3, 2, 2])
e=d.squeeze()    # d中没有维数为1的维度，因此shape不变，e.shape=torch.Size([3, 2, 2])
f=d.squeeze(1)   # d中第1维的维数不为1，因此无法降维，shape不变，e.shape=torch.Size([3, 2, 2])