Pytorch教程（十四）：stack、flatten

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操作时应该从哪个轴开始（第二个轴，也就是彩色通道轴，我们跳过了批轴，保持原样）