PyTorch读取自己的数据集

方法一：用ImageFloder类

图片结构如下所示： 子文件夹的名字为label

transform = T.Compose([
    T.RandomHorizontalFlip(),
    T.RandomVerticalFlip(),
   # T.RandomCrop(c.cropsize),
    T.Resize([c.imageSize, c.imageSize]), 
    T.ToTensor()
])

train_dataset = ImageFolder(
            c.TRAIN_PATH,
            transform)

# Training data loader
trainloader = DataLoader(
    train_dataset,
    batch_size=c.batch_size,
    shuffle=True,
    pin_memory=True,
    num_workers=8,
    drop_last=True
)

trainloader返回的是一个列表[image,label]

 for i_batch, datalist in enumerate(datasets.trainloader):
           
            #print(np.array(datalist).shape) #datalist是列表
           
            data, labels = datalist
            data = data.to(device)
            #print(data.shape) #[batchsize,3,128,128]

方法二：自定义Dataset类

用glob.glob读取图片（不能读子文件夹）

图片结构为：

class Hinet_Dataset(Dataset):
    def __init__(self, transforms_=None, mode="train"):

        self.transform = transforms_
        self.mode = mode
        if mode == 'train':
            # train
            self.files = natsorted(sorted(glob.glob(c.TRAIN_PATH + "/*." + c.format_train, recursive=True)))
        else:
            # test
            self.files = sorted(glob.glob(c.VAL_PATH + "/*." + c.format_val,recursive=True))
        print("Total training examples:", len(self.files))

    def __getitem__(self, index):
        try:
            image = Image.open(self.files[index])
            image = to_rgb(image)
            item = self.transform(image)
            return item  #返回的形式为一张图片

        except:
            return self.__getitem__(index + 1)

    def __len__(self):
        if self.mode == 'shuffle':
            return max(len(self.files_cover), len(self.files_secret))

        else:
            return len(self.files)

 _getitem_(self,index)函数定义放回的形式  即trainloader加载的形式，这里返回是只有图片没有标签

train_datset = Hinet_Dataset(transforms_=transform, mode="train")

trainloader = DataLoader(
    train_dataset,
    batch_size=c.batch_size,
    shuffle=True,
    pin_memory=True,
    num_workers=8,
    drop_last=True
)

使用时：

      for i_batch, data in enumerate(datasets.trainloader):
           
            #print(data.shape) #[batchsize,3,128,128] 
            data = data.to(device)
            cover = data[data.shape[0] // 2:] #data里的前一半
            secret = data[:data.shape[0] // 2]#data里的后一半