pytorch保存save模型的坑

pytorch中保存模型相关的函数有3个：

• torch.save：利用python的pickle模块实现序列化并保存序列化后的object
• torch.load：利用pickle将保存的object反序列化
• torch.nn.Module.load_state_dict：通过反序列化得到的state_dict读取保存的训练参数

有两种方法保存模型：

1. torch.save(model, path) # 直接保存整个模型
2. torch.save(model.state_dict(), path) # 保存模型的参数

   
   

   
   
1
2

相应地有两种方法加载保存的模型：

1. model = torch.load(path) # 直接加载模型
2. model = Model()                         # 先初始化一个模型
   model.load_state_dict(torch.load(path)) # 再加载模型参数

   
   

   
   
1
2
3

看起来第一种方法更加简单方便，但官方推荐的确实第二种方法，这是因为第一种方法容易产生问题：

设备错误

当我们在cuda:0上训练好一个模型并保存时，读取出来的模型也是默认在cuda:0上的，如果训练过程的其他数据被放到了如cuda:1上，那么就会发生错误:

RuntimeError: arguments are located on different GPUs at /opt/conda/conda-bld/pytorch_1503966894950/work/torch/lib/THC/generated/../generic/THCTensorMathPointwise.cu:215

   
   

   
   
1

这个错误的解决办法：

1. 把其他所有数据也都放到cuda:0上

2.  device = torch.device("cuda:1")
    model = torch.load(PATH, map_location=device)
   
        
        
   
        
        
   1
   2

此外，当我们在使用pytorch的预训练模型时，默认会把预训练模型加载到gpu0上去，一般来说不会报错，但是却额外占用了gpu0的显存，所以有必要利用CUDA_VISIBLE_DEVICES来指定gpu：

1. 终端中直接指定：CUDA_VISIBLE_DEVICES=1 python main.py
2. pytorch代码中指定：

   import os
   os.environ["CUDA_VISIBLE_DEVICES"] = "5"
   
        
        
   
        
        
   1
   2

要注意的是，这样就不能用device = torch.device('cuda:5')来指定gpu了，因为默认的gpu个数只有1，即程序已经只能看到cuda:5这张显卡了，建议直接使用device = torch.device('cuda')。这样既能够指定显卡，又防止了显存的浪费。

模型错误

这个错误的出现主要是因为框架的更新换代，比如当我们用pytorch1.0训练并保存了一个CNN模型，再用pytorch1.1去读取模型的话，就会出现这个错误：

AttributeError: 'Conv2d' object has no attribute 'padding_mode'

   
   

   
   
1

解决方法：

1.  model = Model()  # 先初始化一个模型
    model_state = torch.load(model_path).state_dict()
    model.load_state_dict(model_state)  # 加载模型参数
   
        
        
   
        
        
   1
   2
   3

2.  model = torch.load(PATH)   # 直接加载模型
    for m in model.modules():  # 手动给这个object加一个属性
    	if 'Conv' in str(type(m)):
        	setattr(m, 'padding_mode', 'zeros')
   
        
        
   
        
        
   1
   2
   3
   4

同样的错误不仅仅会出现在加载模型时，保存时也可能会出现错误。
我使用了别人用pytorch0.4训练好的一个SSD模型，并想要用pytorch1.1对它进行微调，加载模型出现的问题都用上述方法解决了，但在保存训练好的模型时却出现了问题：

AttributeError: 'SGD' object has no attribute 'defaults'

   
   

   
   
1

这是因为我用的优化器是SGD，而pytorch版本的更新导致了这个问题，解决办法则是保存优化器的参数

在恢复优化器时，可能会出现这个问题：

Traceback (most recent call last):
  File "train.py", line 156, in <module>
    print_freq=print_freq)
  File "train.py", line 52, in train
    optimizer.step()
  File "/home/chenfan/anaconda3/envs/chenf/lib/python3.7/site-packages/torch/optim/sgd.py", line 100, in step
    buf.mul_(momentum).add_(1 - dampening, d_p)
RuntimeError: expected backend CPU and dtype Float but got backend CUDA and dtype Float

   
   

   
   
1
2
3
4
5
6
7
8

这是因为优化器保存的状态是cuda类型的tensor，再加载时会自动使用map to cpu，因此需要再次转成cuda类型的tensor:

for state in optimizer.state.values():
    for k, v in state.items():
        if torch.is_tensor(v):
            state[k] = v.to(device)

   
   

   
   
1
2
3
4

总结

需要保存模型/优化器时，建议保存参数，并将tensor转换到CPU上，以免之后使用时带来不必要的麻烦。当需要继续之前的训练，就必须保存优化器的状态；如果只是想在之前的基础上微调模型，就只用保存模型的weights。