mirrors
/
mmyolo
mirror of https://github.com/open-mmlab/mmyolo.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tools/model_converters/yolov6_to_mmyolo.py
View File

@ -5,7 +5,15 @@ import torch
def convert(src, dst): def convert(src, dst):
import sys
sys.path.append('yolov6')
try:
ckpt = torch.load(src, map_location=torch.device('cpu')) 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 # The saved model is the model before reparameterization
model = ckpt['ema' if ckpt.get('ema') else 'model'].float() model = ckpt['ema' if ckpt.get('ema') else 'model'].float()
new_state_dict = OrderedDict() new_state_dict = OrderedDict()