Yolov3模型从pth转换为pt格式

import torch
from nets.yolo import YoloBody
from utils.utils import get_classes

classes_path    = 'model_data/voc_classes.txt'
class_names, num_classes = get_classes(classes_path)

model = YoloBody(num_classes)

model_path = 'logs/last_epoch_weights.pth'
device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(torch.load(model_path, map_location=device))

example = torch.rand(1, 3, 640, 640)
model.eval()
output = torch.jit.trace(model, example)
torch.jit.save(output, "yolo.pt")