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))

反正能跑起来:
pytorch保存训练模型参数并实现继续训练_第1张图片

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

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

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的初始化对应模型别写错了。
pytorch保存训练模型参数并实现继续训练_第2张图片
之后用新模型继续训练试试:

    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))

能成功训练。
pytorch保存训练模型参数并实现继续训练_第3张图片

变化学习率的保存

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

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

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

总结

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

你可能感兴趣的:(dl,python,python,pytorch,深度学习)