Pytorch Dataset 准备图片数据

将一堆图片载入dataset中

class NewDataset(Dataset):
    def __init__(self, fnames, transform):
        self.transform = transform
        self.fnames = fnames
        self.num_samples = len(self.fnames)

    def __getitem__(self, idx):
        fname = self.fnames[idx]
        img = torchvision.io.read_image(fname)
        img = self.transform(img)
        return img

    def __len__(self):
        return self.num_samples

fnames = glob.glob(os.path.join('./data/faces','*'))
        # 1. Resize the image to (64, 64)
        # 2. Linearly map [0, 1] to [-1, 1]
transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((64, 64)),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5)),
        ])
dataset = NewDataset(fnames,transform)