Pytorch:当批量大小大于1时，如何绘制分割任务的预测输出？

import torch
import matplotlib.pyplot as plt

# 假设有一个批量的预测输出，形状为(batch_size, height, width)
batch_size = 4
height = 256
width = 256

# 生成随机的预测输出
predictions = torch.randn(batch_size, height, width)

# 将预测输出转换为图像格式
# 这里假设预测输出是二值化的分割结果，可以使用torch.where函数将其转换为RGB图像
# 例如，将预测输出中大于0的像素设为红色，小于等于0的像素设为绿色
predictions_image = torch.zeros(batch_size, height, width, 3)
predictions_image[predictions > 0] = torch.tensor([1, 0, 0])
predictions_image[predictions <= 0] = torch.tensor([0, 1, 0])

# 使用Matplotlib绘制预测输出
fig, axes = plt.subplots(nrows=batch_size, ncols=1, figsize=(6, 6))
for i, ax in enumerate(axes):
    ax.imshow(predictions_image[i])
    ax.axis('off')

plt.show()