如何将PyTorch Lightning的.ckpt文件转化为.pth文件

最新推荐文章于 2024-09-05 12:02:43 发布

原创最新推荐文章于 2024-09-05 12:02:43 发布 · 1.7k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

其他专栏收录该内容

9 篇文章

订阅专栏

文章讲述了如何在使用原生PyTorch的代码中加载PyTorchLightning训练的模型，重点在于处理state_dict的加载以及不同框架间参数命名的调整。

记录碰到的一个问题：

使用pytorch lightning框架下训练好的一个分割模型，需要在另外一套使用原生pytorch写的代码中，不知道怎么加载训练好的.ckpt模型

解决方案：

ckpt里保存了更多信息，而我们只需要state_dict，同时观察到两者的参数命名存在略微差异

# 创建一个新的模型状态字典
new_state_dict = {}

# 遍历原始字典中的键值对
for old_key, old_value in xx['state_dict'].items():
    # 删除键中的 'net.' 部分
    new_key = old_key.replace('net.', '')

    # 使用新的键来保存值
    new_state_dict[new_key] = old_value

# 将新的状态字典保存到一个新的文件
torch.save(new_state_dict, 'xxx.pth')