1、模型定义

1.1 一般步骤(构建线性回归为例)

1.1.1 构建数据集

import torch
from sklearn.datasets import make_regression
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import random

plt.rcParams['font.sans-serif']=['SimHei'] #用来正常显示中文标签
plt.rcParams['axes.unicode_minus']=False #用来正常显示负号

# 数据集函数
def create_dataset():

    x, y, coef = make_regression(n_samples=100,
                                 n_features=1,
                                 noise=10,
                                 coef=True,
                                 bias=14.5,
                                 random_state=0)

    # 转换为张量
    x = torch.tensor(x)
    y = torch.tensor(y)

    return x, y, coef


# 构建数据加载器
def data_loader(x, y, batch_size):

    data_len = len(y)
    data_index = list(range(data_len))
    random.shuffle(data_index)
    batch_number = data_len // batch_size

    for idx in range(batch_number):

        start = idx * batch_size
        end = start + batch_size

        batch_train_x = x[start: end]
        batch_train_y = y[start: end]

        yield batch_train_x, batch_train_y

1.1.2 定义假设函数

# 模型参数
w = torch.tensor(0.1, requires_grad=True, dtype=torch.float64)
b = torch.tensor(0.0, requires_grad=True, dtype=torch.float64)


def linear_regression(x):
    return w * x + b

1.1.3 定义损失函数

def square_loss(y_pred, y_true):
    return (y_pred - y_true) ** 2

1.1.4 定义优化方法

def sgd(lr=0.01):
    # 使用批量样本的平均梯度
    w.data = w.data - lr * w.grad.data / 16
    b.data = b.data - lr * b.grad.data / 16

1.1.5 训练函数

def train():

    # 加载数据集
    x, y, coef = create_dataset()
    # 定义训练参数
    epochs = 100
    learning_rate = 0.01
    # 存储损失
    epoch_loss = []
    total_loss = 0.0
    train_sample = 0

    for _ in range(epochs):

        for train_x, train_y in data_loader(x, y, 16):

            # 训练数据送入模型
            y_pred = linear_regression(train_x)

            # 计算损失值
            loss = square_loss(y_pred, train_y.reshape(-1, 1)).sum()
            total_loss += loss.item()
            train_sample += len(train_y)

            # 梯度清零
            if w.grad is not None:
                w.grad.data.zero_()

            if b.grad is not None:
                b.grad.data.zero_()

            # 反向传播
            loss.backward()

            # 更新参数
            sgd(learning_rate)

            print('loss: %.10f' % (total_loss / train_sample))

        epoch_loss.append(total_loss / train_sample)


    # 绘制拟合直线
    print(coef, w.data.item())
    plt.scatter(x, y)

    x = torch.linspace(x.min(), x.max(), 1000)
    y1 = torch.tensor([v * w + 14.5 for v in x])
    y2 = torch.tensor([v * coef + 14.5 for v in x])

    plt.plot(x, y1, label='训练')
    plt.plot(x, y2, label='真实')
    plt.grid()
    plt.legend()
    plt.show()

    # 打印损失变化曲线
    plt.plot(range(epochs), epoch_loss)
    plt.title('损失变化曲线')
    plt.grid()
    plt.show()

调用 train 函数,最后输出的结果为:

PyTorch创建和保存模型_模型保存

PyTorch创建和保存模型_模型创建_02


1.2 使用PyTorch构建线性回归

前面我们使用手动的方式来构建了一个简单的线性回归模型,如果碰到一些较大的网络设计,手动构建过于繁琐。所以,我们需要学会使用 PyTorch 的各个组件来搭建网络。接下来,我们使用 PyTorch 提供的接口来定义线性回归。

  • 使用 PyTorch 的 nn.MSELoss() 代替自定义的平方损失函数
  • 使用 PyTorch 的 data.DataLoader 代替自定义的数据加载器
  • 使用 PyTorch 的 optim.SGD 代替自定义的优化器
  • 使用 PyTorch 的 nn.Linear 代替自定义的假设函数
import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import make_regression
import matplotlib.pyplot as plt


# 构建数据集
def create_dataset():

    x, y, coef = make_regression(n_samples=100,
                                 n_features=1,
                                 noise=10,
                                 coef=True,
                                 bias=14.5,
                                 random_state=0)

    # 将构建数据转换为张量类型
    x = torch.tensor(x)
    y = torch.tensor(y)

    return x, y, coef


def train():

    # 构建数据集
    x, y, coef = create_dataset()
    # 构建数据集对象
    dataset = TensorDataset(x, y)
    # 构建数据加载器
    dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
    # 构建模型
    model = nn.Linear(in_features=1, out_features=1)
    # 构建损失函数
    criterion = nn.MSELoss()
    # 优化方法
    optimizer = optim.SGD(model.parameters(), lr=1e-2)
    # 初始化训练参数
    epochs = 100

    for _ in range(epochs):

        for train_x, train_y in dataloader:

            # 将一个batch的训练数据送入模型
            y_pred = model(train_x.type(torch.float32))
            # 计算损失值
            loss = criterion(y_pred, train_y.reshape(-1, 1).type(torch.float32))
            # 梯度清零
            optimizer.zero_grad()
            # 自动微分(反向传播)
            loss.backward()
            # 更新参数
            optimizer.step()


    # 绘制拟合直线
    plt.scatter(x, y)
    x = torch.linspace(x.min(), x.max(), 1000)
    y1 = torch.tensor([v * model.weight + model.bias for v in x])
    y2 = torch.tensor([v * coef + 14.5 for v in x])

    plt.plot(x, y1, label='训练')
    plt.plot(x, y1, label='真实')
    plt.grid()
    plt.legend()
    plt.show()


if __name__ == '__main__':
    train()

2、模型的保存加载

2.1 PyTorch保存模型的方法

神经网络的训练有时需要几天、几周、甚至几个月,为了在每次使用模型时避免高代价的重复训练,我们就需要将模型序列化到磁盘中,使用的时候反序列化到内存中。

PyTorch 提供了两种保存模型的方法:

  • 序列化模型对象
  • 存储模型的网络参数

2.2 序列化模型对象

import torch
import torch.nn as nn
import pickle


class Model(nn.Module):

    def __init__(self, input_size, output_size):

        super(Model, self).__init__()
        self.linear1 = nn.Linear(input_size, input_size * 2)
        self.linear2 = nn.Linear(input_size * 2, output_size)

    def forward(self, inputs):

        inputs = self.linear1(inputs)
        output = self.linear2(inputs)
        return output


def test01():

    model = Model(128, 10)

    # 第一个参数: 存储的模型
    # 第二个参数: 存储的路径
    # 第三个参数: 使用的模块
    # 第四个参数: 存储的协议
    torch.save(model, 'model/test_model_save.pth', pickle_module=pickle, pickle_protocol=2)


def test02():

    # 第一个参数: 加载的路径
    # 第二个参数: 模型加载的设备
    # 第三个参数: 加载的模块
    model = torch.load('model/test_model_save.pth', map_location='cpu', pickle_module=pickle)


if __name__ == '__main__':
    test01()
    test02()

Python 的 Pickle 序列化协议有多种,详细可查看官网: https://www.python.org/search/?q=pickle+protocol

注意: 当我们训练的模型在 GPU 中时,torch.save 函数将其存储到磁盘中。当再次加载该模型时,会将该模型从磁盘先加载到 CPU 中,再移动到指定的 GPU 中,例如: cuda:0、cuda:1。但是,当重新加载的机器不存在 GPU 时,模型加载可能会出错,这时,可通过 map_localtion=’CPU’ 将其加载到 CPU 中。

2.3 存储模型的网络参数

import torch
import torch.nn as nn
import torch.optim as optim


class Model(nn.Module):

    def __init__(self, input_size, output_size):

        super(Model, self).__init__()
        self.linear1 = nn.Linear(input_size, input_size * 2)
        self.linear2 = nn.Linear(input_size * 2, output_size)

    def forward(self, inputs):

        inputs = self.linear1(inputs)
        output = self.linear2(inputs)
        return output



def test01():

    model = Model(128, 10)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    # 定义存储参数
    save_params = {
        'init_params': {
            'input_size': 128,
            'output_size': 10
        },
        'acc_score': 0.98,
        'avg_loss': 0.86,
        'iter_numbers': 100,
        'optim_params': optimizer.state_dict(),
        'model_params': model.state_dict()
    }

    # 存储模型参数
    torch.save(save_params, 'model/model_params.pth')


def test02():

    # 加载模型参数
    model_params = torch.load('model/model_params.pth')
    # 初始化模型
    model = Model(model_params['init_params']['input_size'], model_params['init_params']['output_size'])
    # 初始化优化器
    optimizer = optim.Adam(model.parameters())
    optimizer.load_state_dict(model_params['optim_params'])
    # 显示其他参数
    print('迭代次数:', model_params['iter_numbers'])
    print('准确率:', model_params['acc_score'])
    print('平均损失:', model_params['avg_loss'])


if __name__ == '__main__':
    test01()
    test02()

在上面代码中,我们把模型的一些初始化参数、模型的权重参数、训练的迭代次数、以及优化器的参数等都进行了存储。