完整的模型训练套路

保存模型方式1：

vgg_16 = torchvision.models.vgg16(pretrained=False)
torch.save(vgg_16, "vgg16_method1.pth")

利用torch.save方法进行保存，参数2是保存的路径及其名称，该方法保存模型，保存了模型的结构及其权重参数。

加载模型：

model = torch.load("vgg16_method1.pth")
print(model)

保存模型方式2：

torch.save(vgg_16.state_dict(), "vgg16_method2.pth")

该方法保存模型，是只保存了模型的权重参数，这样的保存方式所占存储更小，比较推荐。我们加载时可以加载随意权重的该模型，然后把权重赋值给该模型就可以了。

加载模型：

vgg16 = torchvision.models.vgg16(pretrained=False)
model = torch.load("vgg16_method2.pth")
vgg16.load_state_dict(model)

可以看出我们这里加载模型，首先是随意得到一个vgg16的模型，然后将我们保存的权重参数加载出来然后复制给vgg16模型，就得到了我们想要的vgg16模型。

两个函数 一个是torch.save另一个是torch.load

利用第一个方式保存自己的模型的时候，有一个小trick，就是在另一个python文件下，加载模型时，必须要把模型的那个类也导入过来，不然就会报错。

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter

from nn_module import *

# 准备数据集
from torch.utils.data import DataLoader

train_set = torchvision.datasets.CIFAR10("./datasets", train=True, transform=torchvision.transforms.ToTensor(),
                                         download=True)
test_set = torchvision.datasets.CIFAR10("./datasets", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)

# 数据集的长度
train_set_length = len(train_set)
test_set_length = len(test_set)

# 将训练集批量化
train_dataloader = DataLoader(train_set, batch_size=64)
test_dataloader = DataLoader(test_set, batch_size=64)

# 创建网络模型
my_model = Model()

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

# 优化器
learning_rate = 1e-3
optimizer = torch.optim.Adam(my_model.parameters(), lr=learning_rate)

# 训练次数
train_step = 0
test_step = 0

writer = SummaryWriter("logs")
# 迭代
epoch = 10
for i in range(epoch):
    my_model.train()  # 切换训练和评估的模式， 有些层的作用只在训练中起作用如dropout，
    print("---第{}轮训练开始---".format(i + 1))
    for data in train_dataloader:
        imgs, targets = data
        outputs = my_model(imgs)
        result_loss = loss_cross(outputs, targets)
        optimizer.zero_grad()
        result_loss.backward()
        optimizer.step()
        train_step = train_step + 1
        if train_step % 50 == 0:
            print("第{}次训练后的损失为：{}".format(train_step, result_loss.item()))
            writer.add_scalar("train_loss", result_loss, train_step)
    total_loss = 0
    total_accuracy = 0
    my_model.eval()
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            outputs = my_model(imgs)
            loss = loss_cross(outputs, targets)
            total_loss = total_loss + loss
            accuracy = (outputs.argmax(1) == targets).sum()  # argmax函数能返回最大位置的索引，1横向计算，0是纵向计算，返回类型不变
            total_accuracy = total_accuracy + accuracy
        print("整个数据集上的损失为：{}".format(total_loss.item()))
        print("在整个数据集上的正确率为：{}".format(total_accuracy/test_set_length))
        writer.add_scalar("test_loss", total_loss, test_step+1)
        writer.add_scalar("accuracy", total_accuracy/test_set_length, test_step+1)
        test_step = test_step + 1
        torch.save(my_model, "my_model{}.pth".format(test_step))
        print("模型已保存")
writer.close()

完整的训练了一个CIFAR10的网络，神经网络的训练步骤大致如上，先准备数据集（训练集和测试集），然后将数据集进行批量化处理，创建我们的模型（也就是我们的网络），然后设置损失函数，设置优化器，然后采用mini-batch梯度下降法，套用两个for循环，第一个for循环是epoch，第二个for循环是完整的遍历一次训练集，我们可以在每完整遍历完一次训练集（也就是1个eopch），可以对模型进行推断，也就是用测试集进行推断，可以看看损失的变化，也可以看看准确率的变化。

在训练过程中，我们也可以利用tensorboard去可视化我们的损失的变化，准确率的变化。