重塑 Pytorch 张量

Question

我在 Pytorch 中有一个大小为 (24, 2, 224, 224) 的张量。

• 24 = 批量大小
• 2 = 代表前景的矩阵和 背景
• 224 = 图像高度尺寸
• 224 = 图像宽度 维度

这是执行二进制分割的 CNN 的输出。在 2 个矩阵的每个单元格中存储了该像素是前景还是背景的概率：[n][0][h][w] + [n][1][h][w] = 1 对于每个坐标

我想将它重塑为大小为(24, 1, 224, 224) 的张量。根据概率较高的矩阵，新层中的值应为0或1。

我该怎么做？我应该使用哪个功能？

Answer 1

使用torch.argmax()（对于 PyTorch +0.4）：

prediction = torch.argmax(tensor, dim=1) # with 'dim' the considered dimension 
prediction = prediction.unsqueeze(1) # to reshape from (24, 224, 224) to (24, 1, 224, 224)

---

如果 PyTorch 版本低于 0.4.0，则可以使用 tensor.max() 来返回最大值及其索引（但不能通过索引值区分）：

_, prediction = tensor.max(dim=1)
prediction = prediction.unsqueeze(1) # to reshape from (24, 224, 224) to (24, 1, 224, 224)