【学习笔记】python2读取python3训练的模型（pth文件）

背景：前面跑强化学习用的Python3的环境进行训练的，现在要结合ROS部署，方便起见使用Python2的环境，发现没法在Python2下直接读取之前的.pth文件。

参考：

1. PyTorch加载模型model.load_state_dict()问题，Unexpected key(s) in state_dict: “module.features…，Expected .
2. Attempted to read a PyTorch file with version 3, but the maximum supported version for reading is 2.

一、Python3转Python2

1. 在保存模型文件的文件夹内打开终端，并切换至训练模型时使用的环境（Python），打开Python：

   python
   

2. 加载模型文件（policy.pth）：

   import torch
   state_dict = torch.load('policy.pth') # 模型可以保存为pth文件，也可以为pt文件。
   

3. 读取模型文件，并保存至有序字典中：

   from collections import OrderedDict
   new_state_dict = OrderedDict()
   for k, v in state_dict.items():
       new_state_dict[k] = v
   

4. 保存：

   torch.save(new_state_dict, 'new_policy.pth', _use_new_zipfile_serialization=False)