Pytorch 定义MyDatasets实现多通道分别输入不同数据

最近在做一个项目,用双通道神经网络,每个通道输入不同数据训练,具有相同label。开始没想到如何实现,网上很多例子都是单通道,即便找到双通道的例子,两个通道的输入也相同。
最后,终于想到了一个办法。多输入和单输入其实是一样的,只需要重新改写torch.utils.data.Datasets. 需要改写class Dataset里面的initlengetitem
一个例子:

  class MyDataset(data.Dataset):
    def __init__(self, data1,data2, labels):
        self.data1= data1
        self.data2= data2
        self.labels = labels  # 我的例子中label是一样的,如果你的不同,再增加一个即可

    def __getitem__(self, index):    
        img1,img2, target = self.data1[index], self.data2[index], self.labels[index]
        return img1,img2, target

    def __len__(self):
        return len(self.data1) # 我的例子中len(self.data1) = len(self.data2)

觉得有帮助请手动点个赞,谢谢!

你可能感兴趣的:(pytorch)