Pytorch: 解决因pytorch版本不同 导致训练ckpt加载失败

大家都会遇到在工程项目实施阶段，如果训练的模型文件在不同的torch版本环境下部署时，会报错~。
报错举例

# 查看torch环境
import torch
print(torch.__version__)

# 训练时环境：torch 1.8.2+cu111
# 部署时环境：torch 1.4.0

torch.load(checkpoint, map_location = 'cpu')
# 报错如下：
RuntimeError: version_ <= kMaxSupportedFileFormatVersion INTERNAL ASSERT FAILED at /pytorch/caffe2/serialize/inline_container.cc:132, please report a bug to PyTorch. 
Attempted to read a PyTorch file with version 3, but the maximum supported version for reading is 2. 
Your PyTorch installation may be too old. (init at /pytorch/caffe2/serialize/inline_container.cc:132)

报错的原因分析：
在torch1.6版本及其以后，torch.save函数使用了一种新的文件格式。torch.load任然保持着对旧版本的兼容，如果高版本环境下想要保存低版本兼容的模型文件格式，可以使用_use_new_zipfile_serialization=False参数设定。

The 1.6 release of PyTorch switched torch.save to use a new zipfile-based file format. 
torch.load still retains the ability to load files in the old format. 
If for any reason you want torch.save to use the old format, 
pass the kwarg _use_new_zipfile_serialization=False.

解决方案

import torch
state_dict = torch.load(checkpoint_old_file, map_location="cpu")
torch.save(state_dict, checkpoint_new_file, _use_new_zipfile_serialization=False)

转载自SSDesign的知乎文章
https://zhuanlan.zhihu.com/p/465807214