zishuo/crnn_plate_recognition
1
0
Fork 0
mirror of https://github.com/we0091234/crnn_plate_recognition.git synced 2025-09-27 07:52:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

onnx-trt support

Browse Source
This commit is contained in:
we0091234
2022-12-13 13:52:53 +08:00
parent 79c609f2b5
commit 3688ccc532
3 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
export.py
View File

@@ -14,6 +14,7 @@ if __name__=="__main__":
parser.add_argument('--batch_size', type=int, default=1, help='batch size')
parser.add_argument('--dynamic', action='store_true', default=False, help='enable dynamic axis in onnx model')
parser.add_argument('--simplify', action='store_true', default=False, help='simplified onnx')
parser.add_argument('--trt', action='store_true', default=False, help='support trt')
@@ -21,7 +22,7 @@ if __name__=="__main__":
print(opt)
checkpoint = torch.load(opt.weights)
cfg = checkpoint['cfg']
model = myNet_ocr(num_classes=len(plate_chr),cfg=cfg,export=True)
model = myNet_ocr(num_classes=len(plate_chr),cfg=cfg,export=True,trt=opt.trt)
model.load_state_dict(checkpoint['state_dict'])
model.eval()