torch编写dataloader记录

项目场景：

需要编写一个dataloader批量读取图片数据

解决方案：

from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import transforms
import cv2
import torch


def load_cifar100(path):
    img = cv2.imread(path)
    img = cv2.resize(img, (224, 224))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    img = transform(img).to(device)
    print('img1 : {}'.format(img.shape))
    # img = img.unsqueeze(0)
    # print('img2 : {}',format(img.shape))
    return img

class trainset(Dataset):
    def __init__(self, file_train, loader=load_cifar100):
        #定义好 image 的路径
        self.images = file_train
        self.loader = loader

    def __getitem__(self, index):
        fn = self.images[index]
        img = self.loader(fn)
        return img

    def __len__(self):
        return len(self.images)

file_list = ['../image2_100/33375.jpg', '../image2_100/101661.jpg','../image2_100/103665.jpg','../image2_100/105421.jpg','../image2_100/126197.jpg','../image2_100/121875.jpg','../image2_100/113066.jpg','../image2_100/110776.jpg','../image2_100/110266.jpg']
train_data  = trainset(file_list)
trainloader = DataLoader(train_data, batch_size=3,shuffle=True)
for batch_id,data in enumerate(trainloader):
    print('batch_id:{}'.format(batch_id))
    print('data.shape:{}'.format(data.shape))