diff --git a/gen_pkl.py b/gen_pkl.py new file mode 100644 index 0000000..27a20ea --- /dev/null +++ b/gen_pkl.py @@ -0,0 +1,47 @@ +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)