Design

This module was started to experiment around function torch.export.export() and see what kind of issues occur when leveraging that function to convert a torch.nn.Module into ONNX. The design is organized around the following main pieces:

• GraphBuilder: simplifies the creation of an ONNX model with the inner API available in onnx. It also tracks type and shape information. It supports shape inference and constant computation with numpy arrays or torch tensors.

• DynamoInterpreter: walks through a torch model and converts every piece into ONNX nodes, the first step is to get an torch.fx.Graph, then convert every of its nodes into ONNX operators.

• _aten_functions and _aten_methods: a collection of functions converting every node from torch.fx.Graph instance into ONNX operators. There is one function per node type.

• GraphBuilderPatternOptimization: once the model is converted, this class looks for sequence of ONNX nodes which can be rewritten in a more efficient way. This can be applied on graphs produced by the exporter and any existing graphs.

Function to_onnx calls these pieces to produce an ONNX models. Next sections gives more details about how it is implemented. It may not be fully up to date.

• Custom Exporter
• Pattern Optimizer
• Dynamo Backends