You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
47 lines
1.3 KiB
47 lines
1.3 KiB
import argparse |
|
import pickle |
|
from collections import OrderedDict |
|
|
|
from ultralytics import YOLO |
|
|
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('-w', |
|
'--weights', |
|
type=str, |
|
required=True, |
|
help='YOLOv8 pytorch weights') |
|
parser.add_argument('-o', |
|
'--output', |
|
type=str, |
|
required=True, |
|
help='Output file') |
|
args = parser.parse_args() |
|
return args |
|
|
|
|
|
args = parse_args() |
|
|
|
model = YOLO(args.weights) |
|
model.model.fuse() |
|
YOLOv8 = model.model.model |
|
|
|
strides = YOLOv8[-1].stride.detach().cpu().numpy() |
|
reg_max = YOLOv8[-1].dfl.conv.weight.shape[1] |
|
|
|
state_dict = OrderedDict(GD=model.model.yaml['depth_multiple'], |
|
GW=model.model.yaml['width_multiple'], |
|
strides=strides, |
|
reg_max=reg_max) |
|
|
|
for name, value in YOLOv8.state_dict().items(): |
|
value = value.detach().cpu().numpy() |
|
i = int(name.split('.')[0]) |
|
layer = YOLOv8[i] |
|
module_name = layer.type.split('.')[-1] |
|
stem = module_name + '.' + name |
|
state_dict[stem] = value |
|
|
|
with open(args.output, 'wb') as f: |
|
pickle.dump(state_dict, f)
|
|
|