基于pytorch的图像分类网络搭建思路

基于pytorch的图像分类网络搭建思路
本文的基本内容为:
实现了一个基于pytorch框架，Resnet18网络结构的图像分类网络。经训练后可以良好的实现宝可梦（皮卡丘，小火龙，妙蛙种子，杰尼龟，超梦）图片的分类.
key words:pytorch，Resnet18，迁移学习(transfer learning),图像分类(image classification)
本文基于龙良曲的pytorch深度学习代码进行总结
Tips：
1.代码参考龙良曲老师的github 63节迁移学习
https://github.com/dragen1860/Deep-Learning-with-PyTorch-Tutorials
有改动，其中检测部分是自己写的
2.本文中引入visdom进行实时训练情况监测，所以在运行前请在终端启动visdom

python -m visdom.server

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):     # 定义Pokemon数据集，继承torch的Dataset类

    def __init__(self, root, resize, mode):
        super(Pokemon, self).__init__()

        self.root = root
        self.resize = resize

        self.name2label = {
   }

        for name in sorted(os.listdir(os.path.join(root))):   #  通过目录的安排规律获取预测值位置与Pokemon名字的对应关系字典
            if not os.path.isdir(os.path.join(root, name)):
                continue

            self.name2label[name] = len(self.name2label.keys())

        # print(self.name2label)
        # print(self.name2label.keys())

        self.images, self.labels = self.load_csv('images.csv')   # 读取csv文件，获取images路径和对应的label

        if mode == 'train':       # 训练 验证 测试集划分 6:2:2
            self.images = self.images[:int(0.6*len(self.images))]
            self.labels = self.labels[:int(0.6*len(self.labels))]
        elif mode == 'val':
            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))]
        else:
            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):  # 函数输入：需要读取的csv文件名字  函数输出：images，label的列表

        if not os.path.exists(os.path.join(self.root,filename)):  # 如果不存在csv文件，那么创建它

            images = []
            for name in self.name2label.keys():  # 用glob.glob()获取root目录下三种格式的图片
                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')