【学习笔记】python2读取python3训练的模型(pth文件)
背景:前面跑强化学习用的Python3的环境进行训练的,现在要结合ROS部署,方便起见使用Python2的环境,发现没法在Python2下直接读取之前的
.pth
文件。
参考:
一、Python3转Python2
-
在保存模型文件的文件夹内打开终端,并切换至训练模型时使用的环境(Python),打开Python:
python
-
加载模型文件(policy.pth):
import torch state_dict = torch.load('policy.pth') # 模型可以保存为pth文件,也可以为pt文件。
-
读取模型文件,并保存至有序字典中:
from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): new_state_dict[k] = v
-
保存:
torch.save(new_state_dict, 'new_policy.pth', _use_new_zipfile_serialization=False)