resnet18实现猫狗分类

先简单说一下整体流程，利用pytorch训练模型并转化为onnx格式，然后配置好dlinfer，利用cv22infer在cv22平台量化序列化模型，展开推理

1训练模型

1.1处理数据集

参考图片下载地址:

https://www.kaggle.com/c/dogs-vs-cats-redux-kernels-edition/data

1.1.1首先继承写一个继承自dataset的类

#继承了Dataset的类

class DogCat(data.Dataset):

    def __init__(self, root, transform=None, train=True, test=False):

        '''root表示用于训练的图片地址，前70%用于训练，后30%用于测试

        '''

        self.test = test

        self.train = train

        self.transform = transform

        imgs = [os.path.join(root, img) for img in os.listdir(root)]  #imgs是一个list，list中放(编号，图片地址)

 

        if self.test:#如果是test集

            imgs = sorted(imgs, key=lambda x: int(x.split('.')[-2].split('/')[-1]))

        else:        #如果是train集

            imgs = sorted(imgs, key=lambda x: int(x.split('.')[-2]))

         

        imgs_num = len(imgs)

 

        if self.test:

            self.imgs = imgs

        else:

            random.shuffle(imgs)

            if self.train:

                self.imgs = imgs[:int(0.7 * imgs_num)]#self.imgs表示前百分之七十的图片

            else:

                self.imgs = imgs[int(0.7 * imgs_num):]#self.imgs表示除了前百分之七十的图片

 

 

    # 作为迭代器必须有的方法，可以用[]符号读取数据

    def __getitem__(self, index):

        img_path = self.imgs[index]

        if self.test:

            label = int(self.imgs[index].split('.')[-2].split('/')[-1])

        else:

            label = 1 if 'dog' in img_path.split('/')[-1] else 0  # 狗的label设为1，猫的设为0

        data = Image.open(img_path)

        data = self.transform(data)

        return data, label

 

 

    def __len__(self):

        return len(self.imgs)

1.1.2把传入的图片集转换为Tensor并且改变格式

# 对数据集训练集的处理

transform_train = transforms.Compose([

    transforms.Resize((256, 256)),  # 先调整图片大小至256x256

    transforms.RandomCrop((224, 224)),  # 再随机裁剪到224x224

    transforms.RandomHorizontalFlip(),  # 随机的图像水平翻转，通俗讲就是图像的左右对调

    transforms.ToTensor(),              #图片转换为Tensor

    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.2225))  # 归一化，数值是用ImageNet给出的数值

])

 

 

# 对数据集验证集的处理

transform_val = transforms.Compose([

    transforms.Resize((224, 224)),                         

    transforms.ToTensor(),                     

    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),

])

 

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 若能使用cuda，则使用cuda

 

trainset = DogCat('/mnt/ssd0/zhangwentao/train', transform=transform_train)#这是我图片的地址

valset = DogCat('/mnt/ssd0/zhangwentao/train', transform=transform_val)

 

 

#继承了dataset                            dataset     一捆有多大      数据顺序      多线程输入，=0表示单线程

trainloader = torch.utils.data.DataLoader(trainset, batch_size=20, shuffle=True, num_workers=0)

valloader = torch.utils.data.DataLoader(valset, batch_size=20, shuffle=False, num_workers=0)

1.2修改网络模型

class Net(nn.Module):

    '''pytorch的resnet18接口的最后一层fc层的输出维度是1000。这明显不符合猫狗大战数据集，

    因为猫狗大战数据集是二分类的，所以最后一层fc输出的维度已经是2才对。

    因此我们需要对resnet18进行最后一层的修改。

    '''

    def __init__(self, model):

        super(Net, self).__init__()

        # 取掉model的后1层

        self.resnet_layer = nn.Sequential(*list(model.children())[:-1])

        self.Linear_layer = nn.Linear(512, 2)  # 加上一层参数修改好的全连接层

 

    def forward(self, x):

        x = self.resnet_layer(x)

        x = x.view(x.size(0), -1)

        x = self.Linear_layer(x)

        return x

 

'''具体参考完整代码，这里仅展示如何使用网络模型

model = Net(resnet18(pretrained=True))

model = model.to(device)#转到gpu上运行

'''

1.3开始训练

''' for epoch in range(1):

        train(epoch)训练

        val(epoch)  验证，算损失函数

'''

 

#训练epoch次

def train(epoch):

    scheduler.step() #按照Pytorch的定义是用来更新优化器的学习率的

    model.train()    #训练模式

    train_acc = 0.0#准确率

    for batch_idx, (img, label) in enumerate(trainloader):

        image = Variable(img.cuda())#放到gpu上

        label = Variable(label.cuda())#放到ppu上

 

        optimizer.zero_grad()#意思是把梯度置零，也就是把loss关于weight的导数变成0.

        out = model(image)   #投喂图片

         

        loss = criterion(out, label)#利用交叉熵损失函数算出out和label的差别

         

        loss.backward()#反向传播

        optimizer.step()

 

        train_acc = get_acc(out, label)#获得准确率

        print("Epoch:%d [%d|%d] loss:%f acc:%f" % (epoch, batch_idx, len(trainloader), loss.mean(), train_acc))

 

#验证准确率

def val(epoch):

    print("\nValidation Epoch: %d" % epoch)

    model.eval()#进入推理模式

    total = 0

    correct = 0

    with torch.no_grad():

        for batch_idx, (img, label) in enumerate(valloader):

            image = Variable(img.cuda())

            label = Variable(label.cuda())

            out = model(image)

 

            _, predicted = torch.max(out.data, 1)

 

            total += image.size(0)

            correct += predicted.data.eq(label.data).cpu().sum()

 

    print("Acc: %f " % ((1.0 * correct.numpy()) / total))

 

 

1.4完整代码

from torchvision.models.resnet import resnet18

import os

import random

from PIL import Image

import torch.utils.data as data

import numpy as np

import torchvision.transforms as transforms

import torch

import torch.nn as nn

import torch.optim as optim

from torch.autograd import Variable

from torch.optim.lr_scheduler import *

 

#继承了Dataset的类

class DogCat(data.Dataset):

    def __init__(self, root, transform=None, train=True, test=False):

        '''root表示用于训练的图片地址，前70%用于训练，后30%用于测试

        '''

        self.test = test

        self.train = train

        self.transform = transform

        imgs = [os.path.join(root, img) for img in os.listdir(root)]  #imgs是一个list，list中放(编号，图片地址)

 

        if self.test:#如果是test集

            imgs = sorted(imgs, key=lambda x: int(x.split('.')[-2].split('/')[-1]))

        else:        #如果是train集

            imgs = sorted(imgs, key=lambda x: int(x.split('.')[-2]))

         

        imgs_num = len(imgs)

 

        if self.test:

            self.imgs = imgs

        else:

            random.shuffle(imgs)

            if self.train:

                self.imgs = imgs[:int(0.7 * imgs_num)]#self.imgs表示前百分之七十

            else:

                self.imgs = imgs[int(0.7 * imgs_num):]#self.imgs表示除了前百分之七十

 

 

    # 作为迭代器必须有的方法，可以用[]符号读取数据

    def __getitem__(self, index):

        img_path = self.imgs[index]

        if self.test:

            label = int(self.imgs[index].split('.')[-2].split('/')[-1])

        else:

            label = 1 if 'dog' in img_path.split('/')[-1] else 0  # 狗的label设为1，猫的设为0

        data = Image.open(img_path)

        data = self.transform(data)

        return data, label

 

 

    def __len__(self):

        return len(self.imgs)

 

 

 

 

 

 

 

# 对数据集训练集的处理

transform_train = transforms.Compose([

    transforms.Resize((256, 256)),  # 先调整图片大小至256x256

    transforms.RandomCrop((224, 224)),  # 再随机裁剪到224x224

    transforms.RandomHorizontalFlip(),  # 随机的图像水平翻转，通俗讲就是图像的左右对调

    transforms.ToTensor(),              #图片转换为Tensor

    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.2225))  # 归一化，数值是用ImageNet给出的数值

])

 

 

# 对数据集验证集的处理

transform_val = transforms.Compose([

    transforms.Resize((224, 224)),                         

    transforms.ToTensor(),                     

    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),

])

 

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 若能使用cuda，则使用cuda

 

trainset = DogCat('/mnt/ssd0/zhangwentao/train', transform=transform_train)

valset = DogCat('/mnt/ssd0/zhangwentao/train', transform=transform_val)

 

 

#继承了dataset                            dataset     一捆有多大      数据顺序      多线程输入，=0表示单线程

trainloader = torch.utils.data.DataLoader(trainset, batch_size=20, shuffle=True, num_workers=0)

valloader = torch.utils.data.DataLoader(valset, batch_size=20, shuffle=False, num_workers=0)

 

 

#acc为准确率，该函数计算准确率

def get_acc(output, label):

    total = output.shape[0]

    _, pred_label = output.max(1)

    num_correct = (pred_label == label).sum().item()

    return num_correct / total

 

 

#训练epoch次

def train(epoch):

    scheduler.step() #按照Pytorch的定义是用来更新优化器的学习率的

    model.train()    #训练模式

    train_acc = 0.0#准确率

    for batch_idx, (img, label) in enumerate(trainloader):

        image = Variable(img.cuda())#放到gpu上

        label = Variable(label.cuda())#放到ppu上

 

        optimizer.zero_grad()#意思是把梯度置零，也就是把loss关于weight的导数变成0.

        out = model(image)   #投喂图片

         

        loss = criterion(out, label)#利用交叉熵损失函数算出out和label的差别

         

        loss.backward()#反向传播

        optimizer.step()

 

        train_acc = get_acc(out, label)#获得准确率

        print("Epoch:%d [%d|%d] loss:%f acc:%f" % (epoch, batch_idx, len(trainloader), loss.mean(), train_acc))

 

#验证准确率

def val(epoch):

    print("\nValidation Epoch: %d" % epoch)

    model.eval()#进入推理模式

    total = 0

    correct = 0

    with torch.no_grad():

        for batch_idx, (img, label) in enumerate(valloader):

            image = Variable(img.cuda())

            label = Variable(label.cuda())

            out = model(image)

 

            _, predicted = torch.max(out.data, 1)

 

            total += image.size(0)

            correct += predicted.data.eq(label.data).cpu().sum()

 

    print("Acc: %f " % ((1.0 * correct.numpy()) / total))

 

 

class Net(nn.Module):

    ''' 继承了nn.Module便于操作

     

        pytorch的resnet18接口的最后一层fc层的输出维度是1000。这明显不符合猫狗大战数据集，

    因为猫狗大战数据集是二分类的，所以最后一层fc输出的维度已经是2才对。

    因此我们需要对resnet18进行最后一层的修改。

    '''

    def __init__(self, model):

        super(Net, self).__init__()

        # 取掉model的后1层

        self.resnet_layer = nn.Sequential(*list(model.children())[:-1])

        self.Linear_layer = nn.Linear(512, 2)  # 加上一层参数修改好的全连接层

 

    def forward(self, x):

        x = self.resnet_layer(x)

        x = x.view(x.size(0), -1)

        x = self.Linear_layer(x)

        return x

 

 

if __name__ == '__main__':

    model = Net(resnet18(pretrained=True))

    model = model.to(device)#转到gpu上运行

 

                                #model.parameters()为当前的网络模型的参数空间，lr为学习率，

    optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)  # 设置训练细节

    '''torch.optim是一个实现各种优化算法的包，优化器

 

    优化器就是需要根据网络反向传播的梯度信息来更新网络的参数，以起到降低loss函数计算值的作用

    '''

 

    scheduler = StepLR(optimizer, step_size=3)

    criterion = nn.CrossEntropyLoss()#交叉熵损失函数

 

    for epoch in range(1):

        train(epoch)

        val(epoch)

 

    #torch.save(model, 'modelcatdog.pth')  # 保存模型

 

    model.to("cpu")#转换回cpu

 

    model.eval()#注意，必须进入推理模式。详细看pytorch官方文档

 

    dummy_input = Variable(torch.randn(1, 3, 224, 224)) #设置入口的shape

 

    #把模型转换为onnx模型，

    torch.onnx.export(model,

                       dummy_input,

                       "ans.onnx", input_names=["input"],

                        opset_version=9)