pytorch保存训练模型参数并实现继续训练

最近的想法是在推荐模型中考虑根据用户对推荐结果的后续选择，利用已训练的offline预训练模型参数来更新新的结果。简单记录一下中途保存参数和后续使用不同数据训练的方法。

简单模型和训练数据

先准备一个简单模型，简单两层linear出个分类结果。

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(64, 32)
        self.linear1 = nn.Linear(32, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        x = self.linear1(x)
        return x

准备训练用数据，这里直接随机两份，同时给出配套的十个类的分类label结果。要注意的是crossEntropy交叉熵只认long以上的tensor，label记得转一下类型。

    rand1 = torch.rand((100, 64)).to(torch.float)
    label1 = np.random.randint(0, 10, size=100)
    label1 = torch.from_numpy(label1).to(torch.long)
    rand2 = torch.rand((100, 64)).to(torch.float)
    label2 = np.random.randint(0, 10, size=100)
    label2 = torch.from_numpy(label2).to(torch.long)

训练简单使用交叉熵，优化器Adam。

    model = MyModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    loss = nn.CrossEntropyLoss()
    iteration = 100
    for i in range(iteration):
        output = model(rand1)
        my_loss = loss(output, label1)
        optimizer.zero_grad()
        my_loss.backward()
        optimizer.step()
        print("iteration:{} loss:{}".format(i, my_loss))

反正能跑起来：

保存与读取训练参数结果的方法

关键的保存方法，可以分为两种，一种是直接把模型整体保存：

torch.save(model, save_path)

两个参数，模型和保存目录。不过这种不常用，如果模型变化或者只需要其中一部分参数就不太灵活。
常用方法的是将需要的模型或优化器参数取出以字典形式存储，这样可以在使用时初始化相关模型，读入对应参数即可。

def save_model(save_path, iteration, optimizer, model):
    torch.save({'iteration': iteration,
                'optimizer_dict': optimizer.state_dict(),
                'model_dict': model.state_dict()},
                save_path)
    print("model save success")

分别存储训练循环次数，优化器设置和模型参数结果。
初始化模型，读取参数并设置：

def load_model(save_name, optimizer, model):
    model_data = torch.load(save_name)
    model.load_state_dict(model_data['model_dict'])
    optimizer.load_state_dict(model_data['optimizer_dict'])
    print("model load success")

初始化新模型：

    path = "net.dict"
    save_model(path, iteration, optimizer, model)
    print(model.state_dict()['linear.weight'])

    new_model = MyModel()
    new_optimizer = torch.optim.Adam(new_model.parameters(), lr=0.01)
    load_model(path, new_optimizer, new_model)
    print(new_model.state_dict()['linear.weight'])

输出第一个linear层的参数看看，确实相同，参数成功读取上了。注意optimizer的初始化对应模型别写错了。

之后用新模型继续训练试试：

    for i in range(iteration):
        output = new_model(rand2)
        my_loss = loss(output, label2)
        new_optimizer.zero_grad()
        my_loss.backward()
        new_optimizer.step()
        print("iteration:{} loss:{}".format(i, my_loss))

能成功训练。

变化学习率的保存

上面的demo只用了固定的学习率来做实验。如果使用了scheduler来变化步长，只要保存scheduler的state_dict，之后对新初始化的scheduler设置对应的当前循环次数即可。

# 存储时
'scheduler': scheduler.state_dict()
# 读取时
scheduler.load_state_dict(checkpoint['lr_schedule'])

scheduler的使用可以看看我之前整理的文章：
利用scheduler实现learning-rate学习率动态变化

总结

这次主要是整理了一下pytorch模型参数的整体保存方法，来实现新数据的后续训练或直接作为offline预训练模型来使用。
不过后续数据分布不同的话感觉效果会很差啊…也不知道能不能用什么算法修改下权重来贴合新的数据，找点多次训练优化论文看看好了。