Pytorch数据加载器在一批中随机播放示例

Question

我在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 

不起作用。

我该如何解决？谢谢！