pytorch实现自己制作图像训练集测试集代码（二）

我在网上查资料时，发现了另一种用pytorch自己制作图像集的方法，操作更加简便的，现在分享给大家。
第二篇pytorch实现自己制作图像集代码：

from torchvision.datasets import ImageFolder
from torchvision import transforms

ROOT_TRAIN = r'D:/pycharm/pychanrm项目文件/data/已知类别' #训练路径
ROOT_TEST = r'D:/pycharm/pychanrm项目文件/data/待分类'    #测试路径

train_transform = transforms.Compose([
    transforms.Resize((1000)), 
    #T.RandomVerticalFlip(), # 随机垂直翻转
    transforms.CenterCrop(1000),
    transforms.ToTensor(), 
    transforms.Normalize([.5, .5, .5], [.5, .5, .5])])
    
test_transform = transforms.Compose([
    transforms.Resize((1000)),
    #T.RandomVerticalFlip(), # 随机垂直翻转
    transforms.CenterCrop(1000),
    transforms.ToTensor(), 
    transforms.Normalize([.5, .5, .5], [.5, .5, .5])])

#读取数据
train_dataset = ImageFolder(ROOT_TRAIN, transform=train_transform)
test_dataset = ImageFolder(ROOT_TEST, transform=test_transform)

#加载数据
data_loader_train = t.utils.data.DataLoader(dataset=train_dataset,
                                            batch_size=1,   
                                            shuffle=True,
                                            num_workers=0)

data_loader_test = t.utils.data.DataLoader(dataset=test_dataset,
                                            batch_size=1,
                                            shuffle=True,
                                            num_workers=0)

这样就可以利用自己的数据集来训练和测试了。

对ImageFolder是一个通用的数据加载器，想要更深的理解，这篇文章写的很不错。
ImageFolder的理解