Fix graph building to exclude input, output and initializer from value_info

onnxscript/function_libs/torch_lib/graph_building.py

@@ -822,9 +822,15 @@
         # nn.Modules exported by dynamo exporter have unique call sites, their function
         # op_type name can serve to form the unique identifier for value info.
         # Store inside top level GraphProto.
-        existing_value_info.update(self.generate_subgraphs_value_info_proto())
-        # Insert value info for nodes in top level graph.
-        existing_value_info.update(self.generate_maingraph_value_info_proto())
+        new_value_info = self.generate_maingraph_value_info_proto()
+        # Do not store input, output or initializer into value_info
+        for name in onnx_model.graph.input:
+            new_value_info.pop(name.name, None)
+        for name in onnx_model.graph.output:
+            new_value_info.pop(name.name, None)
+        for name in self.initializers:
+            new_value_info.pop(name, None)
+        existing_value_info.update(new_value_info)
         onnx_model.graph.value_info.extend(existing_value_info.values())
 
         return onnx_model

onnxscript/function_libs/torch_lib/graph_building_test.py

@@ -3,7 +3,6 @@
 # mypy: disable-error-code="arg-type,type-arg,valid-type"
 from __future__ import annotations
 
-import os
 import unittest
 
 import torch
@@ -140,7 +139,6 @@
 
 
 class TestModelSaving(unittest.TestCase):
-    @unittest.skipIf(os.getenv("CI") == "true", "CI is not ready to run dyanmo_export.")
     def test_save_initializer_to_files_for_large_model(self):
         class MLP(torch.nn.Module):
             def __init__(self, input_size, hidden_size, output_size):
@@ -167,6 +165,34 @@
         # Assert model is larger than 2GB (~=3GB)
         self.assertGreater(model_proto.ByteSize(), 2**31)
 
+    def test_input_output_and_initializer_are_not_stored_in_value_info(self):
+        class MLP(torch.nn.Module):
+            def __init__(self, input_size, hidden_size, output_size):
+                super().__init__()
+                self.fc1 = torch.nn.Linear(input_size, hidden_size)
+                self.fc2 = torch.nn.Linear(hidden_size, output_size)
+                self.relu = torch.nn.ReLU()
+
+            def forward(self, x):
+                out = self.fc1(x)
+                out = self.relu(out)
+                out = self.fc2(out)
+                return out
+
+        batch_size, input_size, hidden_size, output_size = 1, 4, 5, 10
+        model = MLP(input_size, hidden_size, output_size)
+        x = torch.randn(batch_size, input_size)
+
+        model_proto = torch.onnx.dynamo_export(model, x).model_proto
+        v_names = set(v.name for v in model_proto.graph.value_info)
+        print(v_names)
+        for i in model_proto.graph.input:
+            self.assertNotIn(i.name, v_names)
+        for o in model_proto.graph.output:
+            self.assertNotIn(o.name, v_names)
+        for i in model_proto.graph.initializer:
+            self.assertNotIn(i.name, v_names)
+
 
 if __name__ == "__main__":
     unittest.main()