pytorch数据预处理和卷积神经网络训练

数据集建立

在实际情况中，我们拥有的图片数据往往不是那么规整的，图片的大小以及图片的格式往往都会不同。我们要将这些在不同文件中间的图片，加载到一个固定的数据集的类别中，以便我们之后调用。
这里利用Pokemon数据，可以从百度上搜到，我们定义一个class类来存放这些数据。具体代码如下所示。
注：main（）函数仅仅用来查看当前class是否有误，检查其运行状况。
注：实例中用到了torch，cuda， visdom等需要手动下载的包； visdom也需要在cmd中调用，否则main（）函数报错。

# 这个是用来加载数据集的，数据预处理
# 1 创建自定义数据集的类

import torch
import os, glob
import random, csv
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image


class Pokemon(Dataset):

    def __init__(self, root, resize, model):
        # root接收图片存储的位置
        # resize 图片输出的size
        # model 图片的作用,train or test or validation
        super(Pokemon, self).__init__()
        self.root = root
        self.resize = resize
        # 下面给每个类做一个映射
        self.name2label = {
   }
        for name in sorted(os.listdir((os.path.join(root)))):  # 找到当前数据集的目录,sorted是防止乱序
            if not os.path.isdir(os.path.join(root, name)):
                continue  # 如果目录不对的话就过滤掉
            # 如果目录是的话就保存到下面这个字典上面去,# 上面name作为key，当前元素存储的数量作为其值
            self.name2label[name] = len(self.name2label.keys())
        print(self.name2label)  # 看表格样子

        # 往表格里面输入image, label

        self.images, self.labels = self.load_csv('images.csv')

        # 按6：2：2的比例从图片中读取数据.切片操作

        if model == 'train':  # 取60%，从零开始取
            self.images = self.images[:int(0.6*len(self.images))]
            self.labels = self.labels[:int(0.6*len(self.labels))]

        if model == 'val':  # 取20%，从60%处取到80%处
            self.images = self.images[int(0.6*len(self.images)): int(0.8*len(self.images))]
            self.labels = self.labels[int(0.6*len(self.labels)): int(0.8*len(self.labels))]

        if model == 'test':  # 取20%，从80%处取到末尾
            self.images = self.images[int(0.8 * len(self.images)):]
            self.labels = self.labels[int(0.8 * len(self.labels)):]

    def load_csv(self, filename):  # 读取图片的函数

        if not os.path.exists(os.path.join(self.root, filename)):
            images = []
            for name in self.name2label.keys():
                # 'pokemo\\newtwo\\0001.png'
                images += glob.glob(os.path.join(self.root, name, '*png'))
                images += glob.glob(os.path.join(self.root, name, '*jpg'))
                images += glob.glob(os.path.join(self.root, name, '*jpeg'))

            print(len(images), images)  # 1167 ['pokemon\\bulbasaur\\00000000.png',

            random.shuffle(images)
            with open(os.path.join(self.root, filename), mode="w", newline='') as f:
                writer = csv.writer(f)
                for img in images:  # 'pokemon\\bulbasaur\\00000000.png'
                    name = img.split(os.sep)[-2]  # os.sep就是‘\’.
                    label = self.name2label[name]
                    # 以“路径+label”形式存储，如下
                    writer.writerow([img, label])
                print("writen into csv file", filename)

        # 下一次再打开的时候，就直接读取上面写好的文件了

        images = []
        labels = []
        with open(os.path.join(self.root,