Pytorch教程（十四）：stack、flatten

最新推荐文章于 2025-09-24 22:40:58 发布

原创最新推荐文章于 2025-09-24 22:40:58 发布 · 518 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #python #深度学习

pytorch教程专栏收录该内容

21 篇文章

订阅专栏

本文介绍了如何使用PyTorch的stack函数将多个相同形状的张量组合成一个新张量，并讨论了如何调整张量的形状以适应CNN的需求。通过将张量reshape，增加了额外的彩色通道轴，然后使用flatten方法将张量部分拍平，保留批处理轴，使得张量更适合CNN的输入格式。

部署运行你感兴趣的模型镜像

stack

t1 = torch.tensor([
    [1,1,1,1],
    [1,1,1,1],
    [1,1,1,1],
    [1,1,1,1]
])
t2 = torch.tensor([
    [2,2,2,2],
    [2,2,2,2],
    [2,2,2,2],
    [2,2,2,2]
])
t3 = torch.tensor([
    [3,3,3,3],
    [3,3,3,3],
    [3,3,3,3],
    [3,3,3,3]
])
t = torch.stack((t1,t2,t3))
t.shape  # torch.Size([3, 4, 4])

我们使用一个堆栈函数将三个张量的序列连接到一个新的轴上，如果把t1,t2,t3看成图片的话，我们可以看到他们的高度和宽度是4，长度为3的轴表示批大小，t表示一个批次的张量输出。

现在我们有一个秩为3的张量，其中包含了3个4乘4的图像。我们现在需要做的就是把这个张量变成一个CNN期望的形式，为彩色通道增加一个轴，对于每个图像张量，我们基本上都有一个隐式的单色通道，而CNN将会看到一个显式的彩色通道轴。
在这里插入图片描述
然后我们将其拍平：

我们将整个批次都拍平了，这把所有的图像都映射到一个单独的轴上，而我们往往需要单独预测，这里的解决办法是在保持批轴的同时，将每个图像边平，这意味着我们只想把张量的一部分边平，这就可以通过使用Pytorch的flatten来完成：
在这里插入图片描述
注意，我们在调用中指定start_dim参数，这告诉了flatten方法当它开始flatten操作时应该从哪个轴开始（第二个轴，也就是彩色通道轴，我们跳过了批轴，保持原样）