Export the network into a "ONNX-like" file for visualization. Parameters: path: the path to the output file
(self, path: str = None)
| 534 | dot.save(path) |
| 535 | |
| 536 | def to_onnx(self, path: str = None) -> None: |
| 537 | ''' |
| 538 | Export the network into a "ONNX-like" file for visualization. |
| 539 | |
| 540 | Parameters: |
| 541 | path: the path to the output file |
| 542 | ''' |
| 543 | if path is None: |
| 544 | return |
| 545 | |
| 546 | network = self.trt_network |
| 547 | b_onnx_type = True # Use ONNX style operator names |
| 548 | |
| 549 | def layer_type_to_class(layer: trt.ILayer = None) -> trt.ILayer: |
| 550 | ''' |
| 551 | Convert trt.ILayer into a exact TensorRT layer |
| 552 | ''' |
| 553 | layer_type_name = str(layer.type)[10:] |
| 554 | # Some special cases |
| 555 | if layer_type_name == "ELEMENTWISE": |
| 556 | return trt.IElementWiseLayer |
| 557 | if layer_type_name == "LRN": |
| 558 | return trt.ILRNLayer |
| 559 | if layer_type_name == "NMS": |
| 560 | return trt.INMSLayer |
| 561 | if layer_type_name == "PARAMETRIC_RELU": |
| 562 | return trt.IParametricReLULayer |
| 563 | if layer_type_name == "PLUGIN": |
| 564 | return None # IPluginLayer is not supported any more |
| 565 | if layer_type_name == "RAGGED_SOFTMAX": |
| 566 | return trt.IRaggedSoftMaxLayer |
| 567 | if layer_type_name == "SOFTMAX": |
| 568 | return trt.ISoftMaxLayer |
| 569 | if layer_type_name == "TOPK": |
| 570 | return trt.ITopKLayer |
| 571 | |
| 572 | # general cases, e.g. MATRIX_MULTIPLY -> MatrixMultiply |
| 573 | name = "".join(name[0] + name[1:].lower() |
| 574 | for name in layer_type_name.split("_")) |
| 575 | return trt.__builtins__["getattr"](trt, f"I{name}Layer") |
| 576 | |
| 577 | def convert_type_to_onnx(node_type: str = "", |
| 578 | attribution: OrderedDict = OrderedDict()): |
| 579 | ''' |
| 580 | Convert layer names of TensorRT style into ONNX style |
| 581 | ''' |
| 582 | if node_type == "ACTIVATION": |
| 583 | convert_list = { |
| 584 | "RELU": "Relu", |
| 585 | "SIGMOID": "Sigmoid", |
| 586 | "TANH": "Tanh", |
| 587 | "LEAKY_RELU": "LeakyRelu", |
| 588 | "ELU": "Elu", |
| 589 | "SELU": "Selu", |
| 590 | "SOFTSIGN": "Softsign", |
| 591 | "SOFTPLUS": "Softplus", |
| 592 | "CLIP": "Clip", |
| 593 | "HARD_SIGMOID": "HardSigmoid", |