yolov5 모델pt

yolov5의 모델 그냥 저장시 무거워서 dict까지 저장하면 가벼워진다.

from models.yolo import Model
import torch
device = torch.device("cpu")
yaml_path='./dataset/test.yaml'
#yaml_path = './models/yolov5m.yaml'
new_weights='./best.pt'
model = Model(yaml_path).to(device)
trained_model = torch.load(new_weights,map_location=device)['model'].to(device)

torch.save(trained_model, 'best_1.pt')

 

오류 처리시 사용한 yaml의 내용을 참조하여 추가하면 된다.