5.9. Custom passes

In PopRT, a pass is used for various processes such as model conversion and optimisation. The functionality of PopRT can be expanded by adding custom passes.

The implementation/use of custom passes is subject to the general specification for PopRT pass, as detailed in the Passes chapter.

5.9.1. Implementing custom passes

To add custom passes in PopRT, at least one Python file needs to be written. The example code is as follows:

Listing 5.12 replace_add_with_sub.py

# Copyright (c) 2022 Graphcore Ltd. All rights reserved.
import onnx

from poprt.passes import Pass, register


@register('replace_add_with_sub')
class ReplaceAddwithSub(Pass):
    """Replace Add with Sub."""

    def __init__(self):
        super().__init__()

    # define the transform
    def run_transform(
        self,
        graph: onnx.GraphProto,
        is_main_graph: bool,
    ) -> onnx.GraphProto:
        for node in graph.node:
            if node.op_type == 'Add':
                node.op_type = 'Sub'
        return graph

    # define the run method
    def run(self, onnx_model: onnx.ModelProto) -> onnx.ModelProto:
        onnx_model.graph.CopyFrom(
            self.traverse_graph(onnx_model.graph, self.run_transform)
        )
        return onnx_model

Download custom_shape_inference.py

As shown in the example code, the ReplaceAddwithSub pass is implemented. The function of this pass is to replace all Add nodes in the ONNX model with Sub nodes. The pass is registered as replace_add_with_sub.

Note

A Python file can contain multiple custom passes, but the registered name of each custom pass must be unique.

5.9.2. Using custom passes

Custom passes can be used in the PopRT CLI or the Python API.

Using custom passes in the PopRT CLI

The Python file (multiple files shall be separated by ,) containing custom passes can be specified by using the parameter --custom_pass_config. The pass to be applied can be specified with --passes.

Listing the registered custom passes:

Listing 5.13 load_custom_passes_in_cli.sh

poprt \
    --list_all_passes \
    --custom_pass_config replace_add_with_sub.py |
    grep replace_add_with_

# replace_add_with_sub

Download custom_shape_inference.py

Applying custom passes:

poprt \
    --input_model model.onnx \
    --passes replace_add_with_sub \
    --custom_pass_config replace_add_with_sub.py

Note

• Passing of parameters to custom passes is not supported when using custom passes via CLI.

• There may be dependencies among multiple custom passes, and PopRT will execute them in the pass order specified by --passes.

• The passes specified with --passes will be executed after some of the PopRT default conversion processes are completed. This means that the input ONNX model is not the original model of the input.

Using custom passes in the Python API

Custom passes can be registered by importing the custom pass modules into the application code.

Listing 5.14 load_custom_passes.py

import replace_add_with_sub  # noqa

import poprt

assert 'replace_add_with_sub' in poprt.get_registered_passes()

Download custom_shape_inference.py

A custom pass is used in the same way as a built-in pass. Please refer to the Passes chapter for more details.