fix model manager with various batchsize

This commit is contained in:
jario
2024-01-19 21:12:13 +08:00
parent 334c2055d9
commit 212c9f18fa
15 changed files with 72 additions and 19 deletions

View File

@@ -51,10 +51,16 @@ def main():
if os.path.exists(model_file):
print("[1] Model:", line, "EXIST!")
if line.startswith('Nv'):
engine_fn = os.path.splitext(model_file)[0] + '.engine'
online_fn = os.path.splitext(model_file)[0] + '-online.engine'
net = line.split('-')[2]
if len(net.split('_')) == 3:
name, seg, ncls = net.split('_')
engine_fn = os.path.splitext(model_file)[0].replace(net, name + "_" + seg + '_b1_' + ncls) + '.engine'
online_fn = os.path.splitext(model_file)[0].replace(net, name + "_" + seg + '_b1_' + ncls) + '-online.engine'
else:
name, ncls = net.split('_')
engine_fn = os.path.splitext(model_file)[0].replace(net, name + '_b1_' + ncls) + '.engine'
online_fn = os.path.splitext(model_file)[0].replace(net, name + '_b1_' + ncls) + '-online.engine'
if not os.path.exists(engine_fn) and not os.path.exists(online_fn):
net = line.split('-')[2]
if net.startswith("yolov5"):
if len(net.split('_')) == 3:
name, seg, ncls = net.split('_')