Pytorch中的保存与加载

path_model = 'model.pkl'
path_state_dict = 'model_state_dict.pkl'
1、保存(序列化)
    # obj:对象
    # f:输出路径
    torch.save(obj,f)
    
2、保存模型参数 # 数据以字典形式存储
    # 获取模型的可学习参数
    state_dict = net.state_dict()
    # 将参数保存
    torch.save(state_dict,f)

3、加载(反序列化)
    # f:文件路径
    @当使用gpu时用map_location
    # map_location:指定存放位置(cpu或gpu)
    net = torch.load(f,map_location)

4、加载模型参数
    # 加载数据
    state_dict_load = torch.load(f,map_location)
    # 将数据加载至模型
    net.load_state_dict(state_dict_load)
    
5、断点续训练/定时保存模型及优化器等训练好的数据
    @网络模型和优化器随学习而不断变化
    # 需要保存的数据
    checkpoint = {
        # 模型数据
        'model_state_dict':net.state_dict(),
        # 优化器数据
        'optimizer_state_dict':optimizer.state_dict(),
        # epoch
        'epoch':epoch
        }
    path = 'checkpoint_{}_epoch.pkl'.format(epoch)
    torch.save(checkpoint,path)
    
    # 加载
    checkpoint = torch.load(path)
    net.load_state_dict(checkpoint[model_state_dict'])
    opimzer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_epoch = checkpoint['epoch']

 

你可能感兴趣的:(深度学习基础---综合,PyTorch)