[ONNX][dynamo] Parameter to export flat graphs

[asfiyab-nvidia]

🚀 The feature, motivation and pitch

I'm exporting a model with a single linear layer to ONNX. Each layer of the generated graph is an ONNX function with an underlying function body composed of other functions and operators. The feature I'm requesting is a configurable option to generate a flat graph with no ONNX functions.

The motivation for this request is to enable optimizations like constant folding. With function nodes, important information is not passed down to the function body resulting in fewer optimizations than possible with a flat graph. Additionally, the TorchScript-based ONNX export does provide an argument to export modules as functions. It would be beneficial for users if the TorchDynamo-based ONNX exporter had a similar feature.

Code to reproduce:

import torch

class LinearModel(torch.nn.Module):
  def __init__(self):
      super().__init__()
      self.fc0 = torch.nn.Linear(5, 7)

  def forward(self, tensor_x: torch.Tensor):
      output = self.fc0(tensor_x)
      return output

def linearDataloader():
    yield torch.randn(3, 5).cuda()

# Get model and input
model = LinearModel()
data = next(linearDataloader())

# ONNX Export
export_output = torch.onnx.dynamo_export(
    model.eval().to('cuda'),
    data
)
export_output.save('linear_dynamo.onnx')

Alternatives

No response

Additional context

No response

[justinchuby]

Curious to know what are the important information that is not passed down? Thanks!

[thiagocrepaldi]

Have you tried something like

import torch

class LinearModel(torch.nn.Module):
  def __init__(self):
      super().__init__()
      self.fc0 = torch.nn.Linear(5, 7)

  def forward(self, tensor_x: torch.Tensor):
      output = self.fc0(tensor_x)
      return output

def linearDataloader():
    yield torch.randn(3, 5).cuda()

# Get model and input
model = LinearModel()
data = next(linearDataloader())

# ONNX Export
export_output = torch.onnx.dynamo_export(
    model.eval().to('cuda'),
    data
)

import onnx
flat_graph = onnx.inliner.inline_local_functions(export_output.model_proto)

Check out https://github.com/onnx/onnx/blob/main/onnx/inliner.py or https://github.com/onnx/onnx/blob/main/onnx/test/inliner_test.py

[MaanavD]

Hey @asfiyab-nvidia , any updates?

[thiagocrepaldi]

Closing as there is a mechanism to inline the resulting ONNX graph in an easy way