持续创作，加速成长！这是我参与「掘金日新计划 · 10 月更文挑战」的第15天，点击查看活动详情

1 现有模型的使用和更改

VGG16模型+ImageNet数据集

需要scipy包，在虚拟环境下查看已安装的包：pip list，没有安装就安装一下：pip install scipy ImageNet数据集接近150G，VGG16模型最终输入1000个分类。

现在使用CIFAR 10数据集，需要将1000个分类改为10个分类

import torchvision
from torch import nn

train_data = torchvision.datasets.CIFAR10('./cifar10', train=True , transform=torchvision.transforms.ToTensor(),download=True)

vgg16_true = torchvision.models.vgg16(pretrained=True) #pretrained参数表示网络是否训练好
vgg16_false = torchvision.models.vgg16(pretrained=False)

# 改造vgg16
print(vgg16_false)
# 方式1：更改 vgg16， 增加一层
vgg16_true.classifier.add_module('add_linear', nn.Linear(1000, 10))
print(vgg16_true)

# 方式2：更改 vgg16， 改动某层
vgg16_false.classifier[6] = nn.Linear(4096, 10)
print(vgg16_false)

2 网络模型的保存与读取

模型的保存

import torch
import torchvision
vgg16=torchvision.models.vgg16(pretrained=False)
#保存方式1（"保存路径"）模型结构+模型参数
torch.save(vgg16,"vgg16_method1.pth")

#保存方式2：将参数保存为字典格式（官方推荐） 模型参数
torch.save(vgg16.state_dict(),"vgg16_method2.pth")

模型的读取

import torch
import torchvision
vgg16=torchvision.models.vgg16(pretrained=False)
#方式一 =》保存方式一(路径)
model=torch.load("vgg16_method1.pth")
print(model)

#方式2 =》字典格式
#注意方式2要print vgg16这个模型，不能model=xxx，print（model）
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
print(vgg16)

3 完整的模型训练套路

以CIFAR10数据集为例

model.py

import torch
from torch import nn

#搭建神经网络
class Anke(nn.Module): 
    def __init__(self): 
        super(Anke, self).__init__() 
        self.model1 = Sequential( 
            Conv2d(3, 32, 5, 1，2), 
            MaxPool2d(2), 
            Conv2d(32, 32 ,5,1,,2), 
            MaxPool2d(2), 
            Conv2d(32, 64, 5, 1,2), 
            MaxPool2d(2), 
            Flatten(), 
            Linear(64*4*4, 64), 
            Linear(64, 10) 
         )
    def forward(self, x): 
        x = self.model1(x) 
        return x
        
if __name__=='__main__':
    anke=Anke()
    input=torch.ones((64,3,32,32))
    output=anke(input)
    print(output.shape)

train.py

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
#import自己写的models
from model import *

#准备数据集
train_data=torchvision.datasets.CIFAR10("./data",train=True,transform=torchvision.transforms.ToTensor(),download=True)
test_data=torchvision.datasets.CIFAR10("./data",train=False,transform=torchvision.transforms.ToTensor(),download=True)

#查看数据集长度
train_size=len(train_data)
test_size=len(test_data)
print("size of train,test is:{},{}".format(train_size,test_size))

#利用Dataloader加载
train_dataloader=DataLoader(train_data,64)
test_dataloader=DataLoader(test_data,64)


#创建网络模型
anke=Anke()

#损失函数
loss_f = nn.CrossEntropyLoss()

#优化器
#1e-2=0.01
learning_rate=1e-2
optimizer = torch.optim.SGD(anke.parameters(),lr=learning_rate,)
#设置训练网络的参数
#记录训练次数
train_step=0
#测试次数
test_step=0
#训练轮数
epoch=10
#添加tensoeboard
writer=SummaryWriter("train_log")
 
 
for i in range(epoch):
    print("-----------第{}轮训练开始---------".format(i+1))
    #训练步骤开始
    #有时不必要：anke.train()
    for data in train_dataloader:
        imgs,t=data
        output=anke(imgs)
        loss=loss_f(output,t)
        #优化器优化模型
        opt.zero_grad()
        loss.backward()
        opt.step()
        train_step=train_step+1
        #loss.item更加规范（.item不会打印数据类型，例如tensor(5))
        if train_step % 100 == 0:
            print("训练次数{}，loss值为{}".format(train_step,loss.item()))
            writer.add_scalar("train_loss",loss.item(),train_step)
    loss_total=0
    #测试步骤开始
    #有时不必要：anke.eval()
    total_correct=0
    with torch.no_grad():
        for data in test_dataloader:
            imgs,t=data
            output=test1(imgs)
            loss=loss_f(output,t)
            loss_total=loss_total+loss.item()
            test_step=test_step+1
            #argmax参数：1为横向比较，2为纵向比较，output为64,10的矩阵
            #output.argmax(1)==t是为了得到[Ture,False,True....]这种形式
            #.sum:T为1，F为0
            corect=(output.argmax(1)==t).sum()
            total_correct=total_correct+corect
            accuracy=total_correct/test_size
    print("测试集总loss{}".format(loss_total))
    writer.add_scalar("test_loss",loss_total,test_step)
    writer.add_scalar("accuracy",accuracy,test_step)
 
    torch.save(anke,"anke{}.pth".format(i))
    print("模型已保存")