Pytorch 学习笔记（一）：加载数据集

一、Pytorch加载数据

1. Dataset：提供一种方式去获取数据label

• 如何获取每一个数据及其label
• 告诉我们总共有多少数据

# 读取图片
from PIL import Image
img_path = "C:/Users/chengzhou/Desktop/Python_Machine_Learning/train/ants/0013035.jpg"
img = Image.open(img_path)
img.show()

• 获取图片名称列

#-获得图片名称列表
import os
dir_path = "C:/Users/chengzhou/Desktop/Python_Machine_Learning/train/ants"
img_path_list = os.listdir(dir_path)

• 获取图片

from PIL import Image
import os
from torch.utils.data import Dataset #采样器dataset工具导入
class MyData(Dataset):
    def __init__(self, root_dir, label_dir): #根目录和标签目录
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir,self.label_dir)
        self.img_path = os.listdir(self.path)

    def __getitem__(self, idx): # 获取每一个图片
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)
        img = Image.open(img_item_path)
        label = self.label_dir
        return img,label

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

ants_dataset = MyData("train","ants")
bees_dataset = MyData("train","bees")
img, label = ants_dataset[0]
img.show()
train_dataset = ants_dataset + bees_dataset

2. Dataloader：为后面的网络提供不同的数据形式

# DataLoader
from torch.utils.data import DataLoader
test_loader = DataLoader(dataset = test_set, batch_size = 64, shuffle = True, num_workers = 0,drop_last = False)
# batch size 将所选取的数据集打包
# drop last 不满足64张时 舍去最后那些图片
# 返回batch_size方式
writer = SummaryWriter("dataloader")
step = 0
for data in test_loader:
    imgs, targets = data
    # print(imgs.shape)
    # print(targets)
    writer.add_images("test_data1",imgs,step)
    step += 1
writer.close()
# shuffle 打乱顺序取样