基于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'))

            random.shuffle(images)  # 打散图片顺序，不再根据名字排列
            with open(os.path.join(self.root,filename),mode='w',newline='') as f:
                writer = csv.writer(f)
                for img in images:
                    name = img.split(os.sep)[-2]
                    label = self.name2label[name]

                    writer.writerow([img,label])
                print('written into csv file',filename)   # 写入csv文件，其中利用路径中的名字和name2label字典获取索引

        images, labels = [],[]      # 如果存在csv文件，那么读取它，此处先将images和labels清零防止受到前面写入csv时的影响
        with open(os.path.join(self.root,filename)) as f:
            reader = csv.reader(f)
            for row in reader:
                img, label = row
                label = int(label)

                images.append(img)
                labels.append(label)

        assert len(images) == len(labels)

        return images , labels    # 按行读取images和labels，值返回

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


    def denormalize(self,x_hat):

        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]

        # x_hat = (x-mean)/std
        # x= x_hat*std+mean
        mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1)
        std = torch.tensor(std).unsqueeze(1).unsqueeze(1)

        x = x_hat * std + mean

        return x



    def __getitem__(self, idx):
        # idx~[0~len(images)]
        # self.images, self.labels
        # img: 'pokemon\\bulbasaur\\00000000.png'
        # label: 0
        img, label = self.images[idx], self.labels[idx]

        tf = transforms.Compose([
            lambda x:Image.open(x).convert('RGB'), # string path= > image data
            transforms.Resize((int(self.resize*1.25), int(self.resize*1.25))),
            transforms.RandomRotation(15),
            transforms.CenterCrop(self.resize),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])

        img = tf(img)
        label = torch.tensor(label)


        return img, label

def main():
    import visdom
    import time

    viz = visdom.Visdom()
    db = Pokemon('pokemon', 64, 'train')
    x,y = next(iter(db))
    print('sample:',x.shape,y.shape,y)

    viz.image(db.denormalize(x),win='sample_x',opts=dict(title='sample_x'))

    loader = DataLoader(db,batch_size=32,shuffle=True,num_workers=4)

    for x,y in loader:
        viz.images(db.denormalize(x),nrow=8,win='batch',opts=dict(title='batch'))
        viz.text(str(y.numpy()),win='label',opts=dict(title='batch-y'))
        time.sleep(10)



if __name__ == '__main__':
    main()

2.网络训练

import  torch
from    torch import optim, nn
import  visdom
import  torchvision
from    torch.utils.data import DataLoader

from    pokemon import Pokemon
# from    resnet import ResNet18
from    torchvision.models import resnet18

from    utils import Flatten

batchsz = 32
lr = 5e-4
epochs = 20

device = torch.device('cuda')
torch.manual_seed(1234)


train_db = Pokemon('pokemon', 224, mode='train')
val_db = Pokemon('pokemon', 224, mode='val')
test_db = Pokemon('pokemon', 224, mode='test')
train_loader = DataLoader(train_db, batch_size=batchsz, shuffle=True,
                          num_workers=4)
val_loader = DataLoader(val_db, batch_size=batchsz, num_workers=2)
test_loader = DataLoader(test_db, batch_size=batchsz, num_workers=2)


viz = visdom.Visdom()

def evalute(model, loader):
    model.eval()
    
    correct = 0
    total = len(loader.dataset)

    for x,y in loader:
        x,y = x.to(device), y.to(device)
        with torch.no_grad():
            logits = model(x)
            pred = logits.argmax(dim=1)
        correct += torch.eq(pred, y).sum().float().item()

    return correct / total

def main():

    # model = ResNet18(5).to(device)
    trained_model = resnet18(pretrained=True)
    model = nn.Sequential(*list(trained_model.children())[:-1], #[b, 512, 1, 1]
                          Flatten(), # [b, 512, 1, 1] => [b, 512]
                          nn.Linear(512, 5)
                          ).to(device)
    # x = torch.randn(2, 3, 224, 224)
    # print(model(x).shape)

    optimizer = optim.Adam(model.parameters(), lr=lr)
    criteon = nn.CrossEntropyLoss()


    best_acc, best_epoch = 0, 0
    global_step = 0
    viz.line([0], [-1], win='loss', opts=dict(title='loss'))
    viz.line([0], [-1], win='val_acc', opts=dict(title='val_acc'))
    for epoch in range(epochs):

        for step, (x,y) in enumerate(train_loader):

            # x: [b, 3, 224, 224], y: [b]
            x, y = x.to(device), y.to(device)

            model.train()
            print(x.size())
            logits = model(x)
            loss = criteon(logits, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            viz.line([loss.item()], [global_step], win='loss', update='append')
            global_step += 1

        if epoch % 1 == 0:

            val_acc = evalute(model, val_loader)
            if val_acc> best_acc:
                best_epoch = epoch
                best_acc = val_acc

                torch.save(model.state_dict(), 'best.mdl')

                viz.line([val_acc], [global_step], win='val_acc', update='append')


    print('best acc:', best_acc, 'best epoch:', best_epoch)

    model.load_state_dict(torch.load('best.mdl'))
    print('loaded from ckpt!')

    test_acc = evalute(model, test_loader)
    print('test acc:', test_acc)





if __name__ == '__main__':
    main()

3.验证代码

import torch
from torch import optim, nn
from PIL import Image, ImageDraw, ImageFont
from torchvision import transforms
from torchvision.models import resnet18
import os

resize = 224  # 设置resize参数，保持与事先设置的网络入口图片大小参数一致。Resnet18的图片大小参数为3*224*224
device = torch.device('cuda')  # 设置device，试运行的电脑没有安装cuda，为cpu运行
weights = 'F:/Deep-Learning-with-PyTorch-Tutorials-master/lesson63-transfer learning/best.mdl'  # 输入权重的路径，此处权重为前面的迁移学习获得
path = 'images'  # 需要检测的文件夹名字


# imgsz=224

def label2name(root, label):  # 函数输入：数据集路径，网络预测类别数字  函数输出：网络预测类别名字
    namedic = {}
    for name in sorted(os.listdir(os.path.join(root))):   # 通过对pokemon文件夹的摆放获取预测数字值和pokemon名字的字典
        if not os.path.isdir(os.path.join(root, name)):
            continue
        namedic[len(namedic.keys())] = name
    # namedic = {0: 'bulbasaur', 1: 'charmander', 2: 'mewtwo', 3: 'pikachu', 4: 'squirtle'}
    # 也可以通过这行代码直接输入数字与pokemon名字的字典，此时不需要root参数
    # print(namedic)
    pokemon_name = namedic[label]
    return pokemon_name


def creatmodel():   # 函数输出：返回网络结构模型
    trained_model = resnet18(pretrained=True) # 利用torch预训练的resnet18模型
    model = nn.Sequential(*list(trained_model.children())[:-1],  # [b, 512, 1, 1]  利用.children方法，将模型的前17层设置与resnet18一致
                          nn.Flatten(),  # [b, 512, 1, 1] => [b, 512]
                          nn.Linear(512, 5)  # 全连接层，输出为最后5种不同的pokemon类型
                          ).to(device)
    return model


def getimagelist(source):   # 函数输入：检测文件夹路径  函数输出：所有检测文件的路径列表
    rlist = []
    for dir, folder, file in os.walk(source):
        for i in file:
            t = "%s\%s" % (dir, i)
            rlist.append(t)
    # print(rlist)
    return rlist


def drawimagetext(image, context):   # 函数输入：需要写入内容的图片，需要写入的内容  函数输出：写入内容后的图片
    # get an image
    # make a blank image for the text, initialized to transparent text color
    img = Image.open(image)
    # get a font
    fnt = ImageFont.truetype("arial.ttf", 40)
    # get a drawing context
    d = ImageDraw.Draw(img)
    context = str(context)
    d.text((10, 10), context, font=fnt, fill=(0, 0, 0))

    return img


def main():
    model = creatmodel() # 创建模型
    model.load_state_dict(torch.load('best.mdl'))  # load the parameters
    data_transform = transforms.Compose([  # 检测数据预处理步骤
        lambda x: Image.open(x).convert('RGB'),  # string path= > image data
        transforms.Resize((resize, resize)),   # Resize 到网络输入大小
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])     # 由于网络训练时进行了Normalize，检测时也应当进行Normalize再送入网络
    ])

    for image in getimagelist(path):  # 进行图片预测的for循环
        img = data_transform(image)   # 取出图片+预处理
        img = torch.unsqueeze(img, dim=0)  # 扩张batch维度，与网络入口对齐
        img = img.to(device)           # 转移到GPU上
        # print(img.size())
        model.eval()                    # 进入验证模式
        with torch.no_grad():
            logits = model(img)  # 送入模型的img为一个tensor，得到模型输出的五种神奇宝贝的分别的概率
            pred = logits.argmax(dim=1)   # 预测值为概率最大值所处的位置
            # print(int(pred[0]))
            image_afterdraw = drawimagetext(image, label2name('pokemon', int(pred[0])))  # 画图，将信息写入
        image_afterdraw.show()  # 展示图片
    # 图片保存模块
    #     image_afterdraw.save('./imagesave/{}.png'.format(i))
    #     i +=1

if __name__ == '__main__':
    main()