load - Pytorch 数据加载器在批处理中随机播放示例

Question

我有以下 Dataloader 用于 Pytorch 中的自动编码器

class DataLoader(data.DataLoader):
    def __init__(self,  *args, **kwargs):
        super(DataLoader, self).__init__(*args, **kwargs)
        self.iterator = self.__iter__()

    def __next__(self):
        try:
            return next(self.iterator)
        except:
            self.iterator = self.__iter__()
            return next(self.iterator)

在运行损失函数计算时，我想将预测显示为数据集中第一个文件的图像，但即使我只有一批 2 张图像用于训练，这些图像也会被交换，所以我的代码：

class CrossEntropyLoss2d(nn.Module):
    def __init__(self, weight=None):
        super(CrossEntropyLoss2d, self).__init__()

    def forward(self, pred, target, weight=None):
        pred = pred.clamp(min=1e-16)

        plt.clf()
        predpic = (np.uint8(np.argmax(pred.data.cpu().numpy(), 1)))
        plt.imshow(predpic[0]) ### here I want only the first picture! 
        plt.title('Example prediction')
        plt.savefig(os.path.join(results_path, "exampleprediction.png"))
        plt.show()

        ##### COMPUTE LOSS
        loss = -(pred.log() * target)
        loss = loss.sum(1)
        if weight is not None:
            loss = loss * weight
        loss = loss.mean()
        return loss 

不起作用。

我怎样才能解决这个问题？谢谢！