thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update gen_wts

Browse Source
This commit is contained in:
Marcos Luciano
2022-06-28 01:14:44 -03:00
parent fdd9092284
commit 416a8e0108
3 changed files with 169 additions and 170 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
readme.md
View File

@@ -478,7 +478,7 @@ wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5n.pt
#### 4. Generate cfg and wts files (example for YOLOv5n) #### 4. Generate cfg and wts files (example for YOLOv5n)
``` ```
python3 gen_wts_yoloV5.py -w yolov5n.pt python3 gen_wts_yoloV5.py -w yolov5n.pt -c models/yolov5n.yaml
``` ```
#### 5. Copy generated cfg and wts files to DeepStream-Yolo folder #### 5. Copy generated cfg and wts files to DeepStream-Yolo folder
@@ -567,28 +567,22 @@ deepstream-app -c deepstream_app_config.txt
-c or --yaml -c or --yaml
``` ```
* Model width **(default = 640 / 1280 [P6])** * Inference size [size] or [height , weight] **(default [640] / [1280] if --p6)**
``` ```
-mw or --width -s or --size
``` ```
* Model height **(default = 640 / 1280 [P6])** Example for 1280
``` ```
-mh or --height -s 1280
``` ```
* Model channels **(default = 3)** or
``` ```
-mc or --channels -s 1280 1280
```
* P6 model
```
--p6
``` ```
## ##
@@ -687,7 +681,9 @@ To change the `iou-threshold`, `score-threshold` and `topk` values, modify the `
**NOTE**: Lower `topk` values will result in more performance. **NOTE**: Lower `topk` values will result in more performance.
**NOTE**: Make sure to set cluster-mode=4 and pre-cluster-threshold=0 in config_infer file. **NOTE**: Make sure to set cluster-mode=4 in config_infer file.
**NOTE**: You are still able to change the `pre-cluster-threshold` values in the `config_infer.txt` file.
``` ```
[property] [property]

270
utils/gen_wts_yoloV5.py
View File

@@ -17,70 +17,62 @@ class YoloLayers():
break break
return route return route
def route(self, layers=""): def route(self, layers=''):
return "\n[route]\n" + \ return '\n[route]\n' + \
"layers=%s\n" % layers 'layers=%s\n' % layers
def reorg(self): def reorg(self):
return "\n[reorg]\n" return '\n[reorg]\n'
def shortcut(self, route=-1, activation="linear"): def shortcut(self, route=-1, activation='linear'):
return "\n[shortcut]\n" + \ return '\n[shortcut]\n' + \
"from=%d\n" % route + \ 'from=%d\n' % route + \
"activation=%s\n" % activation 'activation=%s\n' % activation
def maxpool(self, stride=1, size=1): def maxpool(self, stride=1, size=1):
return "\n[maxpool]\n" + \ return '\n[maxpool]\n' + \
"stride=%d\n" % stride + \ 'stride=%d\n' % stride + \
"size=%d\n" % size 'size=%d\n' % size
def upsample(self, stride=1): def upsample(self, stride=1):
return "\n[upsample]\n" + \ return '\n[upsample]\n' + \
"stride=%d\n" % stride 'stride=%d\n' % stride
def convolutional(self, bn=False, size=1, stride=1, pad=1, filters=1, groups=1, activation="linear"): def convolutional(self, bn=False, size=1, stride=1, pad=1, filters=1, groups=1, activation='linear'):
b = "batch_normalize=1\n" if bn is True else "" b = 'batch_normalize=1\n' if bn is True else ''
g = "groups=%d\n" % groups if groups > 1 else "" g = 'groups=%d\n' % groups if groups > 1 else ''
return "\n[convolutional]\n" + \ return '\n[convolutional]\n' + \
b + \ b + \
"filters=%d\n" % filters + \ 'filters=%d\n' % filters + \
"size=%d\n" % size + \ 'size=%d\n' % size + \
"stride=%d\n" % stride + \ 'stride=%d\n' % stride + \
"pad=%d\n" % pad + \ 'pad=%d\n' % pad + \
g + \ g + \
"activation=%s\n" % activation 'activation=%s\n' % activation
def yolo(self, mask="", anchors="", classes=80, num=3): def yolo(self, mask='', anchors='', classes=80, num=3):
return "\n[yolo]\n" + \ return '\n[yolo]\n' + \
"mask=%s\n" % mask + \ 'mask=%s\n' % mask + \
"anchors=%s\n" % anchors + \ 'anchors=%s\n' % anchors + \
"classes=%d\n" % classes + \ 'classes=%d\n' % classes + \
"num=%d\n" % num + \ 'num=%d\n' % num + \
"scale_x_y=2.0\n" + \ 'scale_x_y=2.0\n' + \
"beta_nms=0.6\n" + \ 'beta_nms=0.6\n' + \
"new_coords=1\n" 'new_coords=1\n'
def parse_args(): def parse_args():
parser = argparse.ArgumentParser(description="PyTorch YOLOv5 conversion") parser = argparse.ArgumentParser(description='PyTorch YOLOv5 conversion')
parser.add_argument("-w", "--weights", required=True, help="Input weights (.pt) file path (required)") parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
parser.add_argument("-c", "--yaml", help="Input cfg (.yaml) file path") parser.add_argument('-c', '--yaml', help='Input cfg (.yaml) file path')
parser.add_argument("-mw", "--width", type=int, help="Model width (default = 640 / 1280 [P6])") parser.add_argument(
parser.add_argument("-mh", "--height", type=int, help="Model height (default = 640 / 1280 [P6])") '-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
parser.add_argument("-mc", "--channels", type=int, help="Model channels (default = 3)")
parser.add_argument("--p6", action="store_true", help="P6 model")
args = parser.parse_args() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit("Invalid weights file") raise SystemExit('Invalid weights file')
if not args.yaml: if not args.yaml:
args.yaml = "" args.yaml = ''
if not args.width: return args.weights, args.yaml, args.size
args.width = 1280 if args.p6 else 640
if not args.height:
args.height = 1280 if args.p6 else 640
if not args.channels:
args.channels = 3
return args.weights, args.yaml, args.width, args.height, args.channels, args.p6
def get_width(x, gw, divisor=8): def get_width(x, gw, divisor=8):
@@ -96,40 +88,40 @@ def get_depth(x, gd):
return max(r, 1) return max(r, 1)
pt_file, yaml_file, model_width, model_height, model_channels, p6 = parse_args() pt_file, yaml_file, inference_size = parse_args()
model_name = pt_file.split(".pt")[0] model_name = os.path.basename(pt_file).split('.pt')[0]
wts_file = model_name + ".wts" if "yolov5" in model_name else "yolov5_" + model_name + ".wts" wts_file = model_name + '.wts' if 'yolov5' in model_name else 'yolov5_' + model_name + '.wts'
cfg_file = model_name + ".cfg" if "yolov5" in model_name else "yolov5_" + model_name + ".cfg" cfg_file = model_name + '.cfg' if 'yolov5' in model_name else 'yolov5_' + model_name + '.cfg'
if yaml_file == "": if yaml_file == '':
yaml_file = "models/" + model_name + ".yaml" yaml_file = 'models/' + model_name + '.yaml'
if not os.path.isfile(yaml_file): if not os.path.isfile(yaml_file):
yaml_file = "models/hub/" + model_name + ".yaml" yaml_file = 'models/hub/' + model_name + '.yaml'
if not os.path.isfile(yaml_file): if not os.path.isfile(yaml_file):
raise SystemExit("YAML file not found") raise SystemExit('YAML file not found')
elif not os.path.isfile(yaml_file): elif not os.path.isfile(yaml_file):
raise SystemExit("Invalid YAML file") raise SystemExit('Invalid YAML file')
device = select_device("cpu") device = select_device('cpu')
model = torch.load(pt_file, map_location=device)["model"].float() model = torch.load(pt_file, map_location=device)['model'].float()
anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None] anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None]
delattr(model.model[-1], "anchor_grid") delattr(model.model[-1], 'anchor_grid')
model.model[-1].register_buffer("anchor_grid", anchor_grid) model.model[-1].register_buffer('anchor_grid', anchor_grid)
model.to(device).eval() model.to(device).eval()
nc = 0 nc = 0
anchors = "" anchors = ''
masks = [] masks = []
yolo_idx = 0 yolo_idx = 0
spp_idx = 0 spp_idx = 0
for k, v in model.state_dict().items(): for k, v in model.state_dict().items():
if "anchor_grid" in k: if 'anchor_grid' in k:
yolo_idx = int(k.split(".")[1]) yolo_idx = int(k.split('.')[1])
vr = v.cpu().numpy().tolist() vr = v.cpu().numpy().tolist()
a = v.reshape(-1).cpu().numpy().astype(float).tolist() a = v.reshape(-1).cpu().numpy().astype(float).tolist()
anchors = str(a)[1:-1] anchors = str(a)[1:-1]
@@ -140,116 +132,116 @@ for k, v in model.state_dict().items():
mask.append(num) mask.append(num)
num += 1 num += 1
masks.append(mask) masks.append(mask)
elif ".%d.m.0.weight" % yolo_idx in k: elif '.%d.m.0.weight' % yolo_idx in k:
vr = v.cpu().numpy().tolist() vr = v.cpu().numpy().tolist()
nc = int((len(vr) / len(masks[0])) - 5) nc = int((len(vr) / len(masks[0])) - 5)
with open(cfg_file, "w") as c: with open(cfg_file, 'w') as c:
with open(yaml_file, "r", encoding="utf-8") as f: with open(yaml_file, 'r', encoding='utf-8') as f:
c.write("[net]\n") c.write('[net]\n')
c.write("width=%d\n" % model_width) c.write('width=%d\n' % (inference_size[0] if len(inference_size) == 1 else inference_size[1]))
c.write("height=%d\n" % model_height) c.write('height=%d\n' % inference_size[0])
c.write("channels=%d\n" % model_channels) c.write('channels=3\n')
c.write("letter_box=1\n") c.write('letter_box=1\n')
depth_multiple = 0 depth_multiple = 0
width_multiple = 0 width_multiple = 0
layers = [] layers = []
yoloLayers = YoloLayers() yoloLayers = YoloLayers()
f = yaml.load(f, Loader=yaml.FullLoader) f = yaml.load(f, Loader=yaml.FullLoader)
for topic in f: for topic in f:
if topic == "depth_multiple": if topic == 'depth_multiple':
depth_multiple = f[topic] depth_multiple = f[topic]
elif topic == "width_multiple": elif topic == 'width_multiple':
width_multiple = f[topic] width_multiple = f[topic]
elif topic == "backbone" or topic == "head": elif topic == 'backbone' or topic == 'head':
for v in f[topic]: for v in f[topic]:
if v[2] == "Focus": if v[2] == 'Focus':
layer = "\n# Focus\n" layer = '\n# Focus\n'
blocks = 0 blocks = 0
layer += yoloLayers.reorg() layer += yoloLayers.reorg()
blocks += 1 blocks += 1
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1], layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1],
activation="silu") activation='silu')
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
if v[2] == "Conv": if v[2] == 'Conv':
layer = "\n# Conv\n" layer = '\n# Conv\n'
blocks = 0 blocks = 0
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1], layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1],
stride=v[3][2], activation="silu") stride=v[3][2], activation='silu')
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "C3": elif v[2] == 'C3':
layer = "\n# C3\n" layer = '\n# C3\n'
blocks = 0 blocks = 0
# SPLIT # SPLIT
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
activation="silu") activation='silu')
blocks += 1 blocks += 1
layer += yoloLayers.route(layers="-2") layer += yoloLayers.route(layers='-2')
blocks += 1 blocks += 1
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
activation="silu") activation='silu')
blocks += 1 blocks += 1
# Residual Block # Residual Block
if len(v[3]) == 1 or v[3][1] is True: if len(v[3]) == 1 or v[3][1] is True:
for _ in range(get_depth(v[1], depth_multiple)): for _ in range(get_depth(v[1], depth_multiple)):
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
activation="silu") activation='silu')
blocks += 1 blocks += 1
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
size=3, activation="silu") size=3, activation='silu')
blocks += 1 blocks += 1
layer += yoloLayers.shortcut(route=-3) layer += yoloLayers.shortcut(route=-3)
blocks += 1 blocks += 1
# Merge # Merge
layer += yoloLayers.route(layers="-1, -%d" % (3 * get_depth(v[1], depth_multiple) + 3)) layer += yoloLayers.route(layers='-1, -%d' % (3 * get_depth(v[1], depth_multiple) + 3))
blocks += 1 blocks += 1
else: else:
for _ in range(get_depth(v[1], depth_multiple)): for _ in range(get_depth(v[1], depth_multiple)):
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
activation="silu") activation='silu')
blocks += 1 blocks += 1
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
size=3, activation="silu") size=3, activation='silu')
blocks += 1 blocks += 1
# Merge # Merge
layer += yoloLayers.route(layers="-1, -%d" % (2 * get_depth(v[1], depth_multiple) + 3)) layer += yoloLayers.route(layers='-1, -%d' % (2 * get_depth(v[1], depth_multiple) + 3))
blocks += 1 blocks += 1
# Transition # Transition
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
activation="silu") activation='silu')
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "SPP": elif v[2] == 'SPP':
spp_idx = len(layers) spp_idx = len(layers)
layer = "\n# SPP\n" layer = '\n# SPP\n'
blocks = 0 blocks = 0
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
activation="silu") activation='silu')
blocks += 1 blocks += 1
layer += yoloLayers.maxpool(size=v[3][1][0]) layer += yoloLayers.maxpool(size=v[3][1][0])
blocks += 1 blocks += 1
layer += yoloLayers.route(layers="-2") layer += yoloLayers.route(layers='-2')
blocks += 1 blocks += 1
layer += yoloLayers.maxpool(size=v[3][1][1]) layer += yoloLayers.maxpool(size=v[3][1][1])
blocks += 1 blocks += 1
layer += yoloLayers.route(layers="-4") layer += yoloLayers.route(layers='-4')
blocks += 1 blocks += 1
layer += yoloLayers.maxpool(size=v[3][1][2]) layer += yoloLayers.maxpool(size=v[3][1][2])
blocks += 1 blocks += 1
layer += yoloLayers.route(layers="-6, -5, -3, -1") layer += yoloLayers.route(layers='-6, -5, -3, -1')
blocks += 1 blocks += 1
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
activation="silu") activation='silu')
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "SPPF": elif v[2] == 'SPPF':
spp_idx = len(layers) spp_idx = len(layers)
layer = "\n# SPPF\n" layer = '\n# SPPF\n'
blocks = 0 blocks = 0
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
activation="silu") activation='silu')
blocks += 1 blocks += 1
layer += yoloLayers.maxpool(size=v[3][1]) layer += yoloLayers.maxpool(size=v[3][1])
blocks += 1 blocks += 1
@@ -257,35 +249,35 @@ with open(cfg_file, "w") as c:
blocks += 1 blocks += 1
layer += yoloLayers.maxpool(size=v[3][1]) layer += yoloLayers.maxpool(size=v[3][1])
blocks += 1 blocks += 1
layer += yoloLayers.route(layers="-4, -3, -2, -1") layer += yoloLayers.route(layers='-4, -3, -2, -1')
blocks += 1 blocks += 1
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
activation="silu") activation='silu')
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "nn.Upsample": elif v[2] == 'nn.Upsample':
layer = "\n# nn.Upsample\n" layer = '\n# nn.Upsample\n'
blocks = 0 blocks = 0
layer += yoloLayers.upsample(stride=v[3][1]) layer += yoloLayers.upsample(stride=v[3][1])
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "Concat": elif v[2] == 'Concat':
route = v[0][1] route = v[0][1]
route = yoloLayers.get_route(route, layers) if route > 0 else \ route = yoloLayers.get_route(route, layers) if route > 0 else \
yoloLayers.get_route(len(layers) + route, layers) yoloLayers.get_route(len(layers) + route, layers)
layer = "\n# Concat\n" layer = '\n# Concat\n'
blocks = 0 blocks = 0
layer += yoloLayers.route(layers="-1, %d" % (route - 1)) layer += yoloLayers.route(layers='-1, %d' % (route - 1))
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "Detect": elif v[2] == 'Detect':
for i, n in enumerate(v[0]): for i, n in enumerate(v[0]):
route = yoloLayers.get_route(n, layers) route = yoloLayers.get_route(n, layers)
layer = "\n# Detect\n" layer = '\n# Detect\n'
blocks = 0 blocks = 0
layer += yoloLayers.route(layers="%d" % (route - 1)) layer += yoloLayers.route(layers='%d' % (route - 1))
blocks += 1 blocks += 1
layer += yoloLayers.convolutional(filters=((nc + 5) * len(masks[i])), activation="logistic") layer += yoloLayers.convolutional(filters=((nc + 5) * len(masks[i])), activation='logistic')
blocks += 1 blocks += 1
layer += yoloLayers.yolo(mask=str(masks[i])[1:-1], anchors=anchors, classes=nc, num=num) layer += yoloLayers.yolo(mask=str(masks[i])[1:-1], anchors=anchors, classes=nc, num=num)
blocks += 1 blocks += 1
@@ -293,44 +285,44 @@ with open(cfg_file, "w") as c:
for layer in layers: for layer in layers:
c.write(layer[0]) c.write(layer[0])
with open(wts_file, "w") as f: with open(wts_file, 'w') as f:
wts_write = "" wts_write = ''
conv_count = 0 conv_count = 0
cv1 = "" cv1 = ''
cv3 = "" cv3 = ''
cv3_idx = 0 cv3_idx = 0
for k, v in model.state_dict().items(): for k, v in model.state_dict().items():
if "num_batches_tracked" not in k and "anchors" not in k and "anchor_grid" not in k: if 'num_batches_tracked' not in k and 'anchors' not in k and 'anchor_grid' not in k:
vr = v.reshape(-1).cpu().numpy() vr = v.reshape(-1).cpu().numpy()
idx = int(k.split(".")[1]) idx = int(k.split('.')[1])
if ".cv1." in k and ".m." not in k and idx != spp_idx: if '.cv1.' in k and '.m.' not in k and idx != spp_idx:
cv1 += "{} {} ".format(k, len(vr)) cv1 += '{} {} '.format(k, len(vr))
for vv in vr: for vv in vr:
cv1 += " " cv1 += ' '
cv1 += struct.pack(">f", float(vv)).hex() cv1 += struct.pack('>f', float(vv)).hex()
cv1 += "\n" cv1 += '\n'
conv_count += 1 conv_count += 1
elif cv1 != "" and ".m." in k: elif cv1 != '' and '.m.' in k:
wts_write += cv1 wts_write += cv1
cv1 = "" cv1 = ''
if ".cv3." in k: if '.cv3.' in k:
cv3 += "{} {} ".format(k, len(vr)) cv3 += '{} {} '.format(k, len(vr))
for vv in vr: for vv in vr:
cv3 += " " cv3 += ' '
cv3 += struct.pack(">f", float(vv)).hex() cv3 += struct.pack('>f', float(vv)).hex()
cv3 += "\n" cv3 += '\n'
cv3_idx = idx cv3_idx = idx
conv_count += 1 conv_count += 1
elif cv3 != "" and cv3_idx != idx: elif cv3 != '' and cv3_idx != idx:
wts_write += cv3 wts_write += cv3
cv3 = "" cv3 = ''
cv3_idx = 0 cv3_idx = 0
if ".cv3." not in k and not (".cv1." in k and ".m." not in k and idx != spp_idx): if '.cv3.' not in k and not ('.cv1.' in k and '.m.' not in k and idx != spp_idx):
wts_write += "{} {} ".format(k, len(vr)) wts_write += '{} {} '.format(k, len(vr))
for vv in vr: for vv in vr:
wts_write += " " wts_write += ' '
wts_write += struct.pack(">f", float(vv)).hex() wts_write += struct.pack('>f', float(vv)).hex()
wts_write += "\n" wts_write += '\n'
conv_count += 1 conv_count += 1
f.write("{}\n".format(conv_count)) f.write('{}\n'.format(conv_count))
f.write(wts_write) f.write(wts_write)

45
utils/gen_wts_yolor.py
View File

@@ -7,39 +7,50 @@ from models.models import Darknet
def parse_args(): def parse_args():
parser = argparse.ArgumentParser(description="PyTorch YOLOR conversion (main branch)") parser = argparse.ArgumentParser(description='PyTorch YOLOR conversion (main branch)')
parser.add_argument("-w", "--weights", required=True, help="Input weights (.pt) file path (required)") parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
parser.add_argument("-c", "--cfg", required=True, help="Input cfg (.cfg) file path (required)") parser.add_argument('-c', '--cfg', help='Input cfg (.cfg) file path')
args = parser.parse_args() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit("Invalid weights file") raise SystemExit('Invalid weights file')
if not os.path.isfile(args.cfg): if not os.path.isfile(args.cfg):
raise SystemExit("Invalid cfg file") raise SystemExit('Invalid cfg file')
return args.weights, args.cfg return args.weights, args.cfg
pt_file, cfg_file = parse_args() pt_file, cfg_file = parse_args()
wts_file = "%s.wts" % cfg_file.rsplit("/")[1].split(".cfg")[0]
device = select_device("cpu") model_name = os.path.basename(pt_file).split('.pt')[0]
wts_file = model_name + '.wts' if 'yolor' in model_name else 'yolor_' + model_name + '.wts'
new_cfg_file = model_name + '.cfg' if 'yolor' in model_name else 'yolor_' + model_name + '.cfg'
if cfg_file == '':
cfg_file = 'cfg/' + model_name + '.cfg'
if not os.path.isfile(cfg_file):
raise SystemExit('CFG file not found')
elif not os.path.isfile(cfg_file):
raise SystemExit('Invalid CFG file')
device = select_device('cpu')
model = Darknet(cfg_file).to(device) model = Darknet(cfg_file).to(device)
model.load_state_dict(torch.load(pt_file, map_location=device)["model"]) model.load_state_dict(torch.load(pt_file, map_location=device)['model'])
model.to(device).eval() model.to(device).eval()
with open(wts_file, "w") as f: with open(wts_file, 'w') as f:
wts_write = "" wts_write = ''
conv_count = 0 conv_count = 0
for k, v in model.state_dict().items(): for k, v in model.state_dict().items():
if "num_batches_tracked" not in k: if 'num_batches_tracked' not in k:
vr = v.reshape(-1).cpu().numpy() vr = v.reshape(-1).cpu().numpy()
wts_write += "{} {} ".format(k, len(vr)) wts_write += '{} {} '.format(k, len(vr))
for vv in vr: for vv in vr:
wts_write += " " wts_write += ' '
wts_write += struct.pack(">f", float(vv)).hex() wts_write += struct.pack('>f', float(vv)).hex()
wts_write += "\n" wts_write += '\n'
conv_count += 1 conv_count += 1
f.write("{}\n".format(conv_count)) f.write('{}\n'.format(conv_count))
f.write(wts_write) f.write(wts_write)
os.system("cp %s ./" % cfg_file) if not os.path.isfile(new_cfg_file):
os.system('cp %s %s' % (cfg_file, new_cfg_file))