|
|
@ -314,10 +314,10 @@ class Exporter:
|
|
|
|
if dynamic:
|
|
|
|
if dynamic:
|
|
|
|
dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
|
|
|
|
dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
|
|
|
|
if isinstance(self.model, SegmentationModel):
|
|
|
|
if isinstance(self.model, SegmentationModel):
|
|
|
|
dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
|
|
|
|
dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
|
|
|
|
dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
|
|
|
|
dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
|
|
|
|
elif isinstance(self.model, DetectionModel):
|
|
|
|
elif isinstance(self.model, DetectionModel):
|
|
|
|
dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
|
|
|
|
dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400)
|
|
|
|
|
|
|
|
|
|
|
|
torch.onnx.export(
|
|
|
|
torch.onnx.export(
|
|
|
|
self.model.cpu() if dynamic else self.model, # --dynamic only compatible with cpu
|
|
|
|
self.model.cpu() if dynamic else self.model, # --dynamic only compatible with cpu
|
|
|
|