利用自己数据集来训练神经网络pytorch，重写Dataset类

很多小伙伴在刚刚结束深度学习算法的时候&＃xff0c;肯定想用自己的数据来进行训练网络&＃xff0c;但是不知到怎么写代码&＃xff0c;下面这个代码就会为你解惑&＃xff0c;自己可以根据实际情况来更改代码&＃xff0c;训练自己的图片数据集。

下面我用蚂蚁和蜜蜂数据集为例&＃xff0c;我的数据格式是这样的&＃xff0c;如下图&＃xff1a;

 每个类别都会有相应的图片


from torch.utils.data import Dataset,DataLoader
import os
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
transform &＃61; transforms.Compose([transforms.Resize([500, 500]), # 图像预处理transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize(mean&＃61;[.5, .5, .5], std&＃61;[.5, .5, .5])])
class_list&＃61;{0:"ants",1:"bees"}#用于后续预测的时候可以使用&＃xff0c;用预测到的标签来直接获取相应的类别
class Mydataset(Dataset):def __init__(self,file_path&＃61;"D:/PycharmProjects/pythonProject/classification-pytorch-main1/datasets",formate&＃61;"train",transform&＃61;False):self.transforms&＃61;transformself.file_path&＃61;file_pathself.formate&＃61;formateself.file_train&＃61;os.path.join(self.file_path,self.formate)print(self.file_train)files_class &＃61; os.listdir(self.file_train)self.imgs&＃61;[]for i, j in enumerate(files_class):data &＃61; os.path.join(self.file_train, j)print(data)data_1 &＃61; os.listdir(data)data_all &＃61; [[os.path.join(data, k), i] for k in data_1]self.imgs &＃43;&＃61; data_allprint(self.imgs)def __len__(self):return len(self.imgs)def __getitem__(self, index):img_path, label &＃61; self.imgs[index] # 选择文件路径pil_img &＃61; Image.open(img_path).convert(&＃39;RGB&＃39;) # 利用PIL打开文件路径if self.transforms:img&＃61;transform(pil_img)else:pil_img &＃61; np.asarray(pil_img)img &＃61; torch.from_numpy(pil_img)return img, labelif __name__ &＃61;&＃61; &＃39;__main__&＃39;:train_data&＃61; Mydataset(transform&＃61;True)print(train_data.__getitem__(0)[0])print(train_data.__getitem__(0)[1])#验证能否传进模型中train_dataloder&＃61;DataLoader(train_data,batch_size&＃61;8,shuffle&＃61;True)for data in train_dataloder:print(data[0].shape)print(data[1])#结果不唯一&＃xff0c;其中的结果如下&＃xff1a;# torch.Size([8, 3, 500, 500])# tensor([1, 0, 0, 1, 1, 0, 1, 1])break

如果觉得有帮助&＃xff0c;就点个赞吧&＃xff0c;祝大家学业有成&＃xff01;