关于pytorch的Dataloader和Datasets类的理解

关于Dataloader和Datasets类的理解,示例代码:

继承自torch_data.Dataset类,其中的 __len__(self)和__getitem__(self,index)方法必须重写,在使用DataLoader加载数据时,会根据__len__(self)的长度作为遍历的上限,实际返回的个数为__len__(self)除以batch_size,同时返回的结果为调用__getitem__(self,index)这个函数batch_size次,然后拼接的结果,返回的数据可以是各种类型,而具体返回的结果是__getitem__(self,index)的return写的结果.

import torch
import torch.nn as nn
import torch.utils.data as torch_data
import cv2
import os
from torchvision import transforms
from torch.utils.data import DataLoader
import numpy as np
from PIL import Image
dir="/home/jjuv/anaconda3/envs/PointRCNNTest/point__rcnn/data/KITTI/object/training/image_2"
transform2 = transforms.Compose([transforms.RandomCrop([200,200],padding=10),transforms.ToTensor()])
class KittiDataset(torch_data.Dataset):
    def __init__(self, root_dir):
        self.img_datasets=[]
        self.root_dir=root_dir
        for img_index in range(0,16):
            print("img_index:{}".format(img_index))
            file_path = os.path.join(self.root_dir, "%06d.png" % img_index)
            img=Image.open(file_path)
            img = transform2(img)
            self.img_datasets.append(img)
    def __len__(self):
        return len(self.img_datasets)

    def __getitem__(self,index):
        # cv2.imshow("image",img)
        return self.img_datasets[index]
datasets=KittiDataset(dir)
train_loader=DataLoader(datasets,batch_size=4,shuffle=False,num_workers=1)
for i,img in enumerate(train_loader):
    print("i:{}".format(i))