Update bp export example (#633)

Summary:
Pull Request resolved: #633

Update the export example for bundled program, including api update, comments and file name suffix.

Reviewed By: tarun292

Differential Revision: D49849371

fbshipit-source-id: ba0faaecf4b6af1ae9d02d5d7f86a00b79379f5a

Author: Gasoonjia

examples/export/export_bundled_program.py

@@ -21,6 +21,7 @@
 
 
 def save_bundled_program(
+    method_names,
     inputs,
     exec_prog,
     graph_module,
@@ -30,20 +31,20 @@ def save_bundled_program(
     # set for the model. If we wish to test the model with multiple inputs then they can be
     # appended to this list. len(inputs) == number of test sets we want to run.
     #
-    # If we have multiple execution plans in this program then we add another list of tuples
-    # to test that corresponding execution plan. Index of list of tuples will match the index
-    # of the execution plan against which it will be tested.
+    # If we have multiple methods in this program then we add another list of tuples to test
+    # that corresponding method. Index of list of tuples will match the index of the method's name
+    # in the method_names list forwarded to BundledConfig against which it will be tested.
     bundled_inputs = [inputs for _ in range(len(exec_prog.program.execution_plan))]
 
     # For each input tuple we run the graph module and put the resulting output in a list. This
-    # is repeated over all the tuples present in the input list and then repeated for each execution
-    # plan we want to test against.
+    # is repeated over all the tuples present in the input list and then repeated for each method
+    # name we want to test against.
     expected_outputs = [
         [[graph_module(*x)] for x in inputs]
         for i in range(len(exec_prog.program.execution_plan))
     ]
 
-    bundled_config = BundledConfig(bundled_inputs, expected_outputs)
+    bundled_config = BundledConfig(method_names, bundled_inputs, expected_outputs)
 
     bundled_program = create_bundled_program(exec_prog.program, bundled_config)
     bundled_program_buffer = serialize_from_bundled_program_to_flatbuffer(
@@ -54,16 +55,18 @@ def save_bundled_program(
         file.write(bundled_program_buffer)
 
 
-def export_to_pte(model_name, model, example_inputs):
+def export_to_pte(model_name, model, method_names, example_inputs):
     exec_prog = export_to_exec_prog(model, example_inputs)
     save_pte_program(exec_prog.buffer, model_name)
 
     # Just as an example to show how multiple input sets can be bundled along, here we
     # create a list with the example_inputs tuple used twice. Each instance of example_inputs
     # is a Tuple[Union[torch.tenor, int, bool]] which represents one test set for the model.
     bundled_inputs = [example_inputs, example_inputs]
-    print(f"Saving exported program to {model_name}_bundled.pte")
-    save_bundled_program(bundled_inputs, exec_prog, model, f"{model_name}_bundled.pte")
+    print(f"Saving exported program to {model_name}_bundled.bp")
+    save_bundled_program(
+        method_names, bundled_inputs, exec_prog, model, f"{model_name}_bundled.bp"
+    )
 
 
 if __name__ == "__main__":
@@ -87,4 +90,6 @@ def export_to_pte(model_name, model, example_inputs):
         *MODEL_NAME_TO_MODEL[args.model_name]
     )
 
-    export_to_pte(args.model_name, model, example_inputs)
+    method_names = ["forward"]
+
+    export_to_pte(args.model_name, model, method_names, example_inputs)