반응형

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의 내용을 참조하여 추가하면 된다.

반응형

'Deep learning > 소스' 카테고리의 다른 글

yolov5 coco data set training  (0) 2021.09.17
classification mnist-LeNet-5  (0) 2021.04.03
classification model code  (0) 2021.01.14
python 한국 시간으로 설정  (0) 2020.11.17
20201110-yolov5로 데이터 학습  (2) 2020.11.04

+ Recent posts