Python学习之路：Dataset、DataLoader自定义数据集

#encoding:UTF-8
import torch
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
import cv2
import os
import numpy as np

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# step1: 定义MyDataset类， 继承Dataset, 重写抽象方法：__len()__, __getitem()__
class MyDataset(Dataset):
    def __init__(self, root_dir):
        self.root_dir = root_dir

        self.filelist = sorted(os.listdir(self.root_dir))   # 得到所有图像。
        self.size = len(self.filelist)

        self.transform = transforms.Compose([
            np.float32,
            #transforms.ToTensor(),
            #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ])  # 将数据转换为tensor类型，同时归一化。

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        img = cv2.imread(self.filelist[idx])   # use opencv
        return self.transform(img)

train_dataset = MyDataset(root_dir='./imgs/')
train_loader = DataLoader(dataset=train_dataset, batch_size=4, shuffle=True)

for epoch in range(2):
    for i, data in enumerate(train_loader):
        # 将数据从 train_loader 中读出来,一次读取的样本数是32个
        inputs = data

        # 将这些数据转换成Variable类型
        #inputs, labels = Variable(inputs), Variable(labels)

        # 接下来就是跑模型的环节了，我们这里使用print来代替
        print("epoch：", epoch, "的第" , i, "个inputs", inputs.data.size())