[Enhance] YOLOv6 model convert (#165)

* add yolov5 voc training

* fix bug

* fix bug

* fix yolov6 model convert

* update yolov5 modelconvert

* rollback yolov5 model convert
pull/249/head
Nioolek 2022-10-18 10:08:06 +08:00 committed by Haian Huang(深度眸)
parent 936bcf8d86
commit 5708864d5b
1 changed files with 9 additions and 1 deletions

View File

@ -5,7 +5,15 @@ import torch
def convert(src, dst):
ckpt = torch.load(src, map_location=torch.device('cpu'))
import sys
sys.path.append('yolov6')
try:
ckpt = torch.load(src, map_location=torch.device('cpu'))
except ModuleNotFoundError:
raise RuntimeError(
'This script must be placed under the meituan/YOLOv6 repo,'
' because loading the official pretrained model need'
' some python files to build model.')
# The saved model is the model before reparameterization
model = ckpt['ema' if ckpt.get('ema') else 'model'].float()
new_state_dict = OrderedDict()