#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())
|