pytorch resnet18网络动物分类完整的模型训练、验证、测试过程

先重写resnet18的模型，其实不用写也行，可以直接在models中调用

import torch
import torch.nn as nn
from torch.nn import functional as F


class RestNetBasicBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride):
        super(RestNetBasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        output = self.conv1(x)
        output = F.relu(self.bn1(output))
        output = self.conv2(output)
        output = self.bn2(output)
        return F.relu(x + output)


class RestNetDownBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride):
        super(RestNetDownBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride[0], padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride[1], padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.extra = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride[0], padding=0),
            nn.BatchNorm2d(out_channels)
        )

    def forward(self, x):
        extra_x = self.extra(x)
        output = self.conv1(x)
        out = F.relu(self.bn1(output))

        out = self.conv2(out)
        out = self.bn2(out)
        return F.relu(extra_x + out)


class RestNet18(nn.Module):
    def __init__(self):
        super(RestNet18, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
        self.bn1 = nn.BatchNorm2d(64)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        self.layer1 = nn.Sequential(RestNetBasicBlock(64, 64, 1),
                                    RestNetBasicBlock(64, 64, 1))

        self.layer2 = nn.Sequential(RestNetDownBlock(64, 128, [2, 1]),
                                    RestNetBasicBlock(128, 128, 1))

        self.layer3 = nn.Sequential(RestNetDownBlock(128, 256, [2, 1]),
                                    RestNetBasicBlock(256, 256, 1))

        self.layer4 = nn.Sequential(RestNetDownBlock(256, 512, [2, 1]),
                                    RestNetBasicBlock(512, 512, 1))

        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))

        self.fc = nn.Linear(512, 2)

    def forward(self, x):
        out = self.conv1(x)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.avgpool(out)
        out = out.view(x.shape[0], -1)
        out = self.fc(out)
        return out

然后再开始加载数据、训练、验证

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

class Mydata(Dataset):
    def __init__(self, root,  mode):
        super(Mydata, self).__init__()

        self.root = root
        # ----------------编号-----------------------------------
        self.name2label = {}  # 给不同类别的图片编号
        for name in sorted(os.listdir(os.path.join(root))):
            os.path.isdir(os.path.join(root, name))

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

        #print(self.name2label)
        # -----------------------------------------------------

        # ------------------------裁剪-----------------------
        # image, label
        self.images, self.labels = self.load_csv('images.csv')

        if mode == 'train':  # 选60%的数据用于train
            self.images = self.images[:int(0.6 * len(self.images))]
            self.labels = self.labels[:int(0.6 * len(self.labels))]
        else:  # 20%的数据用于test
            self.images = self.images[int(0.6 * len(self.images)):]
            self.labels = self.labels[int(0.6 * len(self.labels)):]
            # -----------------------------------------------------------

    def load_csv(self, filename):
        # print(self.root)
        if not os.path.exists(os.path.join(self.root, filename)):
            images = []
            for name in self.name2label.keys():
                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)

            # 打乱一下
            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('write into csv file:', filename)

        # read from csv file
        images, labels = [], []
        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

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

    def __getitem__(self, idx):
        # idx的范围 [0~len(images)
        # label : 0/1/2/3/4
        img, label = self.images[idx], self.labels[idx]

        tf = transforms.Compose([
            lambda x: Image.open(x).convert('RGB'),
            transforms.Resize((224, 224)),# string path => image data
            transforms.ToTensor()
        ])

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

def main():
    # 建一个对象并#数据加载
    TrainData = Mydata('data', 'train')#训练集
    TestData=Mydata('data', 'test')#测试集

    TrainLoader = DataLoader(TrainData, batch_size=32, shuffle=True)#数据打包
    TestLoader=DataLoader(TestData, batch_size=16, shuffle=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    #初始化模型
    model=resnet.RestNet18()
    #修改线性层的输出
    model.to(device)#加载到模型上
    #设置损失函数
    loss_fun=torch.nn.CrossEntropyLoss()#交叉熵损失函数
    #设置优化器
    optimizer= torch.optim.SGD(model.parameters(),lr=0.01, momentum=0.9)
    #设置学习下降策略
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
    #训练网络
    Epoch=200
    total_test_step=0
    for epoch in range(Epoch):
        model.train()
        for data in TrainLoader:
            img,label=data
            img,label=img.to(device),label.to(device)
            output=model(img)#放入模型训练
            #梯度清零
            optimizer.zero_grad()
            #计算损失函数
            loss=loss_fun(output,label)
            loss.backward()
            #梯度更新
            optimizer.step()
        #打印信息
        test_data_size=len(TestData)
        print('数据长度',test_data_size)
        total_test_loss=0
        total_accuracy=0
        model.eval()
        with torch.no_grad():
            for data in TestLoader:
                imgs,label=data
                imgs, label = imgs.to(device), label.to(device)
                output = model(imgs)  # 放入模型训练
                loss=loss_fun(output,label)
                total_test_loss = total_test_loss + loss.item()#整体测试集上的损失和
                accuracy = (output.argmax(1) == label).sum()#计算对应位置相等的个数
                print('AAA',accuracy)
                total_accuracy = total_accuracy + accuracy
        print('整体测试集上的loss:{}'.format(total_test_loss))
        print('整体测试集上的正确率为:{}'.format(total_accuracy.item() / test_data_size))
        total_test_step+=1
        #更新学习率
        #scheduler.step()  # 更新学习率
        print('train finish!')
    # Step 7: 存储权重
        torch.save(model.state_dict(), './resnet18_Cat_Dog.pth')#保存权重
if __name__ == '__main__':
    main()

在我们保存训练好的模型权重之后，我们要拿这个模型和训练好的权重去随便照一张图片测试

import torch
import torchvision
from PIL import Image
import resnet

device=torch.device('cuda')#在GPU上测试
image=Image.open('D:\\project\\learn\\train\\cat.4.jpg')#打开测试图片
tf=torchvision.transforms.Compose([torchvision.transforms.Resize((224,224)),
                                           torchvision.transforms.ToTensor()])#改变图片尺寸大小与训练时的一样，转变为张量类型
image=tf(image)
image=torch.reshape(image,(1,3,224,224))#改变通道数，尺寸大小与训练时一致
model=resnet.RestNet18()#加载网络模型
model.load_state_dict(torch.load("resnet18_Cat_Dog.pth"))#网络模型加载训练好的权重
model.to(device)
model.eval()
with torch.no_grad():
    image=image.to(device)
    out=model(image)
    _,predicted=torch.max(out,1)
    print(predicted)
#predict=out.argmax(1)
#print(predict)

上边就是测试过程，要主要，我们测试时图片的尺寸要和训练时图片的尺寸保持一致。有啥么错误欢迎指正