Pytorch数据加载与使用
前言
在训练的时候通常使用Dataset来处理数据集。
Dataset的作用
提供一个方式获取数据内容和标签(label)。
实战
from torch.utils.data import Dataset
from PIL import Image
import os
class get_data(Dataset):
def __init__(self,root_dir, label_dir):
self.root_dir = root_dir
self.label_dir = label_dir
self.img_dir = os.path.join(root_dir,label_dir)
self.img_list = os.listdir(self.img_dir)
def __getitem__(self, indx):
img_path = os.path.join(self.img_dir,self.img_list[indx])
img_label = self.label_dir
img_data = Image.open(img_path)
return img_data,img_label
def __len__(self):
return len(self.img_list)
root_dir = "C:\\Users\\Traveler\\Pictures"
label_dir = "Screenshots"
test = get_data(root_dir,label_dir)
img , label = test[1]
# img.show()
print(label)
print(len(test))
此代码定义了一个fet_data类,继承了Dataset,主要提供两个方法,获取数据(getitem)和获取大小(len)。
然而这两个方法使用的是内置的类,当达到一定条件时自动触发,比如__getitem__当需要获取数据时自动触发这个方法。
getitem
返回两个数据,一个是data,一个是label,实现原理就是主要看这两个函数,
os.listdir()
是获取一个路径下的文件名(包含后缀)列表。类似于[‘1.txt’,'2.jpg']
Image.open()
是打开图片文件的,打开一个图片后会赋值很多属性:如下图
使用img.show()就可以打开,img.size()就可以获取大小。
另外os的其他函数也挺重要,比如os.path.join()
就是拼接路径,这个的好处是防止Linux和Windows之间的路径不匹配问题。