ckpt、pt、pth的相互转换

ckpt转pth

import torch

c = torch.load("imagenet100-epoch=100-step=102515.ckpt")
# 创建一个新的模型状态字典
new_state_dict = {}

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

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

# 将新的状态字典保存到一个新的文件
torch.save(new_state_dict, 'imagenet100-epoch=100-step=102515.pth')

pt转pth

import torch
from ultralytics import YOLO
# 加载pt文件
# model = YOLO('weights-v8/yolov8n.pt')
model = torch.load('weights-v8/yolov8n.pt')
# 创建一个新的模型状态字典
parameters = {}
# 提取爹数
parameters =  model['model'].state_dict()  # 会多一个model.

# 保存参数为pth文件
torch.save(parameters, 'weights-v8/yolov8n.pth')