Module: add dump function that recursively prints contents of the mod…

…ule. (pytorch#24356)

Summary:
Pull Request resolved: pytorch#24356

Pull Request resolved: pytorch#24356

Test Plan: Imported from OSS

Differential Revision: D16864133

Pulled By: ZolotukhinM

fbshipit-source-id: 1af757334bc8e156427783bc37500de3c934378b

torch/csrc/jit/script/init.cpp

@@ -401,6 +401,12 @@ void initJitScriptBindings(PyObject* module) {
           },
           py::arg("_extra_files") = ExtraFilesMap())
       .def("_set_optimized", &Module::set_optimized)
+      .def(
+          "_dump",
+          &Module::dump,
+          py::arg("omit_method_bodies") = true,
+          py::arg("omit_attr_values") = true,
+          py::arg("omit_param_values") = true)
       .def(
           "_define",
           [](Module& m,

torch/csrc/jit/script/module.cpp

@@ -1,11 +1,12 @@
+#include <torch/csrc/jit/script/module.h>
 #include <c10/util/Exception.h>
 #include <torch/csrc/autograd/generated/variable_factories.h>
 #include <torch/csrc/jit/export.h>
+#include <torch/csrc/jit/jit_log.h>
 #include <torch/csrc/jit/operator.h>
 #include <torch/csrc/jit/passes/dead_code_elimination.h>
 #include <torch/csrc/jit/script/compiler.h>
 #include <torch/csrc/jit/script/error_report.h>
-#include <torch/csrc/jit/script/module.h>
 #include <torch/csrc/jit/script/schema_matching.h>
 
 namespace torch {
@@ -400,6 +401,78 @@ void Module::apply(const std::function<void(Module&)>& fn) {
   fn(*this);
 }
 
+std::string Module::_dump_to_string(
+    bool omit_method_bodies,
+    bool omit_attr_values,
+    bool omit_param_values,
+    int level) const {
+  std::stringstream ss;
+  std::stringstream parameters_ss;
+  std::stringstream attributes_ss;
+  std::stringstream methods_ss;
+  std::stringstream submodules_ss;
+
+  for (Slot param : get_parameters()) {
+    parameters_ss << param.name() << " = ";
+    if (!omit_param_values) {
+      parameters_ss << param.value().toTensor() << std::endl;
+    } else {
+      parameters_ss << "..." << std::endl;
+    }
+  }
+
+  for (Slot attr : get_attributes()) {
+    attributes_ss << attr.name() << " = ";
+    if (!attr.value().isTensor() || !omit_attr_values) {
+      attributes_ss << attr.value() << std::endl;
+    } else {
+      attributes_ss << "..." << std::endl;
+    }
+  }
+
+  for (const Method& method : get_methods()) {
+    methods_ss << "  method " << method.name() << " {" << std::endl;
+    if (!omit_method_bodies) {
+      methods_ss << torch::jit::jit_log_prefix(
+                        "    ", method.graph()->toString())
+                 << std::endl;
+    }
+    methods_ss << "  }" << std::endl;
+  }
+
+  ss << "module " << name().qualifiedName() << " {" << std::endl;
+  ss << "  parameters {" << std::endl;
+  ss << torch::jit::jit_log_prefix("    ", parameters_ss.str());
+  ss << "  }" << std::endl;
+  ss << "  attributes {" << std::endl;
+  ss << torch::jit::jit_log_prefix("    ", attributes_ss.str());
+  ss << "  }" << std::endl;
+  ss << "  methods {" << std::endl;
+  ss << torch::jit::jit_log_prefix("  ", methods_ss.str());
+  ss << "  }" << std::endl;
+  ss << "  submodules {" << std::endl;
+  for (const Module& submodule : get_modules()) {
+    // We do level + 2, because one level of indentation comes from 'submodules'
+    // scope and the other one goes from a specific submodule we're printing.
+    ss << submodule._dump_to_string(
+        omit_method_bodies, omit_attr_values, omit_param_values, level + 2);
+  }
+  ss << "  }" << std::endl;
+  ss << "}" << std::endl;
+
+  std::string indent(2 * level, ' ');
+  return torch::jit::jit_log_prefix(indent, ss.str());
+}
+
+void Module::dump(
+    bool omit_method_bodies = true,
+    bool omit_attr_values = true,
+    bool omit_param_values = true) const {
+  std::cout << _dump_to_string(
+                   omit_method_bodies, omit_attr_values, omit_param_values, 0)
+            << std::endl;
+}
+
 } // namespace script
 } // namespace jit
 } // namespace torch

torch/csrc/jit/script/module.h

@@ -205,6 +205,11 @@ struct TORCH_API Module {
   slot_list get_attributes() const;
   slot_list get_module_slots() const;
 
+  void dump(
+      bool omit_method_bodies,
+      bool omit_attr_values,
+      bool omit_param_values) const;
+
   const std::vector<Method> get_methods() const {
     return fmap(
         type()->methods(),
@@ -356,6 +361,12 @@ struct TORCH_API Module {
  private:
   Module clone_impl(std::unordered_map<TypePtr, TypePtr>& type_remap) const;
 
+  std::string _dump_to_string(
+      bool omit_method_bodies,
+      bool omit_attr_values,
+      bool omit_param_values,
+      int level) const;
+
   void clone_method(
       const Module& orig,
       const Function& method,