Expose some functions out of ENABLE_FLATBUFFER.

Summary: Because they don't depend on flatbuffer

Test Plan: existing unittests

Differential Revision: D35291748

Pull Request resolved: pytorch#75700
Approved by: https://github.com/iseeyuan

Author: qihqi

torch/csrc/jit/mobile/import_data.cpp

@@ -7,6 +7,7 @@
 #include <caffe2/serialize/inline_container.h>
 #include <torch/csrc/jit/api/compilation_unit.h>
 #include <torch/csrc/jit/mobile/file_format.h>
+#include <torch/csrc/jit/mobile/import_export_common.h>
 #include <torch/csrc/jit/mobile/module.h>
 #include <torch/csrc/jit/mobile/observer.h>
 #include <torch/csrc/jit/mobile/type_parser.h>
@@ -16,7 +17,6 @@
 
 #if defined(ENABLE_FLATBUFFER)
 #include <torch/csrc/jit/mobile/flatbuffer_loader.h>
-#include <torch/csrc/jit/mobile/import_export_common.h>
 #endif // defined(ENABLE_FLATBUFFER)
 
 #include <exception>
@@ -182,7 +182,7 @@ std::map<std::string, at::Tensor> load_parameters_from_zip(
   return map;
 }
 
-#if defined(ENABLE_FLATBUFFER)
+} // namespace
 
 /**
  * Extracts the parameter map stored in @p module. Expects a layout
@@ -238,10 +238,6 @@ std::map<std::string, at::Tensor> mobile_module_to_parameter_map(
       "' in deserialized mobile::Module");
 }
 
-#endif // defined(ENABLE_FLATBUFFER)
-
-} // namespace
-
 std::map<std::string, at::Tensor> _load_parameters(
     std::istream& in,
     c10::optional<at::Device> device) {

torch/csrc/jit/mobile/import_data.h

@@ -3,6 +3,7 @@
 #include <ATen/core/TensorBase.h>
 #include <c10/core/Device.h>
 #include <c10/util/Optional.h>
+#include <torch/csrc/jit/mobile/module.h>
 
 #include <istream>
 #include <map>
@@ -29,5 +30,9 @@ TORCH_API std::map<std::string, at::Tensor> _load_parameters(
     const std::string& filename,
     c10::optional<at::Device> device = c10::nullopt);
 
+// NOTE: Please prefer using _load_parameters over using the function below.
+TORCH_API std::map<std::string, at::Tensor> mobile_module_to_parameter_map(
+    const mobile::Module& module);
+
 } // namespace jit
 } // namespace torch

torch/csrc/jit/mobile/import_export_common.h

@@ -9,8 +9,6 @@ namespace torch {
 namespace jit {
 namespace mobile {
 
-#if defined(ENABLE_FLATBUFFER)
-
 namespace internal {
 /**
  * The name of the mobile::Module attribute which contains saved parameters, as
@@ -20,8 +18,6 @@ namespace internal {
 constexpr char kSavedParametersAttributeName[] = "data";
 } // namespace internal
 
-#endif // defined(ENABLE_FLATBUFFER)
-
 } // namespace mobile
 } // namespace jit
 } // namespace torch

torch/csrc/jit/mobile/train/export_data.cpp

@@ -1,5 +1,6 @@
 #include <torch/csrc/jit/mobile/train/export_data.h>
 
+#include <torch/csrc/jit/mobile/import_export_common.h>
 #include <torch/csrc/jit/mobile/module.h>
 #include <torch/csrc/jit/runtime/instruction.h>
 #include <torch/csrc/jit/serialization/pickler.h>
@@ -12,7 +13,6 @@
 
 #if defined(ENABLE_FLATBUFFER)
 #include <flatbuffers/flatbuffers.h>
-#include <torch/csrc/jit/mobile/import_export_common.h>
 #include <torch/csrc/jit/serialization/flatbuffer_serializer.h>
 #endif // defined(ENABLE_FLATBUFFER)
 
@@ -76,6 +76,8 @@ class IValuePickler final {
   TypeNameUniquer type_name_uniquer_;
 };
 
+} // namespace
+
 /**
  * Converts a map of named tensors to a c10::Dict.
  */
@@ -88,8 +90,6 @@ c10::Dict<std::string, at::Tensor> tensor_map_to_dict(
   return dict;
 }
 
-#if defined(ENABLE_FLATBUFFER)
-
 /**
  * Returns a Module with a single attribute, with the attribute name specified
  * by #internal::kSavedParametersAttributeName, whose value is the provided
@@ -117,9 +117,6 @@ mobile::Module tensor_dict_to_mobile(
   return mobile::Module(object, mcu);
 }
 
-#endif // defined(ENABLE_FLATBUFFER)
-
-} // namespace
 } // namespace mobile
 
 void _save_parameters(

torch/csrc/jit/mobile/train/export_data.h

@@ -33,5 +33,17 @@ TORCH_API void _save_parameters(
     const std::string& filename,
     bool use_flatbuffer = false);
 
+namespace mobile {
+
+// NOTE: Please prefer using _save_parameters directly over using the 2
+// functions below.
+TORCH_API mobile::Module tensor_dict_to_mobile(
+    const c10::Dict<std::string, at::Tensor>& dict);
+
+c10::Dict<std::string, at::Tensor> tensor_map_to_dict(
+    const std::map<std::string, at::Tensor>& map);
+
+} // namespace mobile
+
 } // namespace jit
 } // namespace torch