[ONNX] Track and list model params for scripting

tools/build_variables.bzl

@@ -503,6 +503,7 @@ libtorch_python_core_sources = [
     "torch/csrc/jit/passes/onnx/constant_fold.cpp",
     "torch/csrc/jit/passes/onnx/eliminate_unused_items.cpp",
     "torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.cpp",
+    "torch/csrc/jit/passes/onnx/list_model_parameters.cpp",
     "torch/csrc/jit/passes/onnx/function_substitution.cpp",
     "torch/csrc/jit/passes/onnx/helper.cpp",
     "torch/csrc/jit/passes/onnx/peephole.cpp",

torch/csrc/jit/passes/onnx/list_model_parameters.cpp

@@ -0,0 +1,164 @@
+#include <torch/csrc/jit/passes/onnx/list_model_parameters.h>
+#include <torch/csrc/jit/jit_log.h>
+#include <torch/csrc/jit/passes/dead_code_elimination.h>
+
+namespace torch {
+namespace jit {
+
+// findSubModuleAttr function chases getAttr chains to locate the submodules.
+// For example:
+// module M {
+//   attributes {
+//     A = <SubModule at ...>
+//   }
+//   ...
+//   %A = prim::GetAttr[name="A"](%self)
+//   ...
+//   %B = prim::GetAttr[name="B"](%A)
+//   ...
+//   %weight = prim::GetAttr[name="scale"](%B)
+//   ...
+std::deque<std::string> findSubModuleAttr(
+    Value* input,
+    std::string& name,
+    Module& attrModule,
+    std::shared_ptr<Graph>& graph) {
+  Node* node = input->node();
+  std::deque<std::string> moduleNames;
+
+  while (!(node->outputs()[0]->type() == graph->inputs()[0]->type())) {
+    if (node->kind() == prim::GetAttr) {
+      moduleNames.push_front(node->s(attr::name));
+      node = node->inputs()[0]->node();
+    } else {
+      return moduleNames;
+    }
+  }
+
+  for (auto& moduleName : moduleNames) {
+    attrModule = attrModule.attr(moduleName).toModule();
+  }
+  return moduleNames;
+}
+
+Value* addParamAsArgument(Function* function, std::string& name, IValue& attr) {
+  auto schema = function->getSchema();
+  auto args = schema.arguments();
+  args.emplace_back(Argument(name, nullptr, c10::nullopt, attr));
+  auto new_schema = FunctionSchema(
+      schema.name(),
+      schema.overload_name(),
+      args,
+      schema.returns(),
+      schema.is_vararg(),
+      schema.is_varret());
+  function->setSchema(new_schema);
+  return function->graph()->addInput(name)->setType(attr.type());
+}
+
+std::vector<IValue> getParamAttributes(
+    std::shared_ptr<Graph>& graph,
+    const Module& module_,
+    Function* function_) {
+  std::vector<IValue> attrValues;
+  auto isEval = !module_.hasattr("training") || !module_.is_training();
+  auto block = graph->block();
+  std::vector<Block*> blocks({block});
+
+  Node* m = *block->nodes().begin();
+  WithInsertPoint guard(m);
+
+  while (!blocks.empty()) {
+    Block* block = blocks.back();
+    blocks.pop_back();
+    for (auto it = block->nodes().begin(); it != block->nodes().end();) {
+      Node* n = *it;
+      it++; // node n can be destroyed
+
+      for (Block* sub_block : n->blocks()) {
+        blocks.emplace_back(sub_block);
+      }
+      if (n->kind() == prim::SetAttr &&
+          n->s(attr::name) == "num_batches_tracked") {
+        n->destroy();
+      } else if (n->kind() == prim::GetAttr) {
+        auto name = n->s(attr::name);
+        auto attrModule = module_;
+        auto input = n->inputs()[0];
+
+        auto moduleNames = findSubModuleAttr(input, name, attrModule, graph);
+
+        if (!attrModule.hasattr(name)) {
+          continue;
+        }
+        Value* paramConst = nullptr;
+
+        auto attr = attrModule.attr(name);
+
+        std::string fullName("self_");
+        for (auto& name : moduleNames) {
+          fullName += name + '_';
+        }
+        fullName += name;
+
+        auto type = attrModule.type();
+        auto slot = *type->findAttributeSlot(name);
+
+        if (type->is_parameter(slot) || type->is_buffer(slot) ||
+            name == "training") {
+          if (type->is_parameter(slot) || type->is_buffer(slot) ||
+              name == "training") {
+            if (attr.isTensor()) {
+              TORCH_INTERNAL_ASSERT(attr.isTensor());
+              auto tensor_ = attr.toTensor();
+              if (isEval && tensor_.requires_grad()) {
+                tensor_ = tensor_.detach();
+                tensor_.set_requires_grad(false);
+                attr = IValue(tensor_);
+              }
+              attrValues.emplace_back(attr.toTensor());
+              paramConst = addParamAsArgument(function_, fullName, attr);
+            } else if (attr.isNone() || name == "training") {
+              auto attrVal = tryInsertConstant(*graph, attr);
+              paramConst = *attrVal;
+            }
+            n->output()->replaceAllUsesWith(paramConst);
+            n->removeAllInputs();
+
+            GRAPH_UPDATE(
+                "Folding GetAttr %",
+                n->outputs()[0]->debugName(),
+                " with ",
+                paramConst->debugName());
+          }
+        }
+      }
+    }
+  }
+  return attrValues;
+}
+
+std::pair<Module, std::vector<IValue>> list_module_parameters(
+    const Module& module) {
+  Module moduleClone = module.clone(true);
+  Method method = moduleClone.get_method("forward");
+  std::unordered_set<Function*> preservedMethods_;
+  preservedMethods_.insert(&method.function());
+
+  std::vector<IValue> modelParams;
+  for (auto function : preservedMethods_) {
+    GRAPH_DEBUG("List attributes for function: " + function->name());
+    auto graph = function->graph();
+    auto attributes = getParamAttributes(graph, moduleClone, function);
+    for (auto& attr_ : attributes) {
+      modelParams.push_back(attr_);
+    }
+    GRAPH_DEBUG("Cleaning up module");
+    EliminateDeadCode(graph->block());
+  }
+
+  return std::make_pair(moduleClone, modelParams);
+}
+
+} // namespace jit
+} // namespace torch

torch/csrc/jit/passes/onnx/list_model_parameters.h

@@ -0,0 +1,13 @@
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch {
+namespace jit {
+
+TORCH_API std::pair<Module, std::vector<IValue>> list_module_parameters(
+    const Module& module);
+
+} // namespace jit
+} // namespace torch

torch/csrc/jit/python/init.cpp

@@ -39,6 +39,7 @@
 #include <torch/csrc/jit/passes/onnx/eval_peephole.h>
 #include <torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.h>
 #include <torch/csrc/jit/passes/onnx/function_substitution.h>
+#include <torch/csrc/jit/passes/onnx/list_model_parameters.h>
 #include <torch/csrc/jit/passes/onnx/peephole.h>
 #include <torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h>
 #include <torch/csrc/jit/passes/onnx/preprocess_for_onnx.h>
@@ -279,6 +280,9 @@ void initJITBindings(PyObject* module) {
           "_jit_pass_quant_fusion",
           [](std::shared_ptr<Graph>& g) { return QuantFusion(g); })
       .def("_jit_pass_fold_convbn", &FoldConvBatchNorm)
+      .def(
+          "_jit_onnx_list_model_parameters",
+          [](Module& module) { return list_module_parameters(module); })
       .def(
           "_freeze_module",
           [](Module& module,

torch/onnx/utils.py

@@ -358,10 +358,10 @@ def _create_jit_graph(model, args, _retain_param_name, use_new_jit_passes):
             if not use_new_jit_passes:
                 method_graph, params = torch._C._jit_pass_lower_graph(graph, model._c)
             else:
-                freezed_m = torch._C._freeze_module(model._c)
+                freezed_m = torch._C._freeze_module(model._c, preserveParameters=True)
+                freezed_m, params = torch._C._jit_onnx_list_model_parameters(freezed_m)
                 method_graph = freezed_m._get_method('forward').graph
                 method_graph.eraseInput(0)  # Remove 'self' from model inputs
-                params = []
 
             in_vars, in_desc = torch.jit._flatten(tuple(args) + tuple(params))
             graph = _propagate_and_assign_input_shapes(