pytorch学习

pytorch学习第一天 （跟着B站小土堆学习 附代码及详细注释

1、pytorch’加载数据及dataset类实战

from torch.utils.data import Dataset
from PIL import Image
import os #os模块：语义操作系统，用于处理文件和目录


class MyData(Dataset):
    def __init__(self,root_dir,label_dir): #提供一些全局变量
        self.root_dir=root_dir #self相当于指定了一个类中间的全局变量，这个函数变量可以给其他函数使用
        self.label_dir=label_dir
        self.path=os.path.join(self.root_dir,self.label_dir) #os.path.join用于路径拼接文件路径
        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)

root_dir="dataset"
ants_label_dir="ants_img"
bees_label_dir="bees_img"
ants_dataset=MyData(root_dir,ants_label_dir) #mydata里填初始化里的变量？ 用类创建一个实例
bees_dataset=MyData(root_dir,bees_label_dir)

train_dataset=ants_dataset+bees_dataset

2、提取图片的标签写入记事本

import os

root_dir="dataset"
target_dir="bees_img"
img_path=os.listdir(os.path.join(root_dir,target_dir))
label=target_dir.split('_')[0]
out_dir="bees_label"
for i in img_path:
    file_name=i.split('.jpg')[0]
    with open(os.path.join(root_dir,out_dir,"{}.txt".format(file_name)),'w')as f:
        f.write(label)