mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
284b2b7135
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/74209 Extending `_save_for_mobile` and `_load_for_mobile` to support faltbuffer format with additional optional argument which is set to pick pickle by default. Adding new binary target with suffix `_pickle_and_flatbuffer` to help migration. Size test in D34909502 shows the size has regressed by ~40K but after removing pickle and comparing lite_predictors we have ~120K size measure that we will achieve when deprecating pickle and moving to flatbuffer **BEFORE:** ```lang=mermaid graph TD; torch_core-->torch_mobile_deserialize; torch_mobile_core-->torch_mobile_deserialize; jit_module_saving-->torch_core; jit_module_saving-->torch_mobile_core; torch_mobile_deserialize-->caffe2_serialize; torch_mobile_deserialize-->torch_mobile_module; caffe2_serialize-->miniz; flatbuffer_loader-->mobile_bytecode; flatbuffer_serializer-->mobile_bytecode; mobile_bytecode-->flatbuffer_2.0; flatbuffer_loader-->torch_mobile_module; flatbuffer_serializer-->torch_mobile_module; ``` **AFTER:** ```lang=mermaid graph TD; torch_core-->torch_mobile_deserialize; torch_mobile_core-->torch_mobile_deserialize; jit_module_saving-->torch_core; jit_module_saving-->torch_mobile_core; torch_mobile_deserialize-->caffe2_serialize; torch_mobile_deserialize-->torch_mobile_module; caffe2_serialize-->miniz; flatbuffer_loader-->mobile_bytecode; flatbuffer_serializer-->mobile_bytecode; mobile_bytecode-->flatbuffer_2.0; torch_mobile_deserialize_pickle_and_flatbuffer-->|new| flatbuffer_loader; torch_mobile_deserialize_pickle_and_flatbuffer-->|new| torch_mobile_deserialize; torch_mobile_core_pickle_and_flatbuffer-->|new| torch_mobile_deserialize_pickle_and_flatbuffer; torch_core_pickle_and_flatbuffer-->|new| torch_mobile_deserialize_pickle_and_flatbuffer; jit_module_saving_pickle_and_flatbuffer-->|new| torch_core_pickle_and_flatbuffer; jit_module_saving_pickle_and_flatbuffer-->|new| torch_mobile_core_pickle_and_flatbuffer; flatbuffer_serializer-->torch_mobile_module; jit_module_saving_pickle_and_flatbuffer-->|new|jit_module_saving; jit_module_saving_pickle_and_flatbuffer-->|new|flatbuffer_serializer; flatbuffer_loader-->torch_mobile_module; ``` ghstack-source-id: 151744258 Test Plan: Similar Build Deps Dags ``` [pavithran@devvm5216.vll0 /data/users/pavithran/fbsource] buck query 'allpaths(//xplat/caffe2:torch_mobile_all_ops_pickle_and_flatbuffer, //xplat/caffe2:torch_mobile_deserialize_pickle_and_flatbuffer)' --output-format dot-compact | pastry P486770901: https://www.internalfb.com/intern/paste/P486770901/ [pavithran@devvm5216.vll0 /data/users/pavithran/fbsource] buck query 'allpaths(//xplat/caffe2:torch_mobile_all_ops, //xplat/caffe2:torch_mobile_deserialize)' --output-format dot-compact | pastry P486771278: https://www.internalfb.com/intern/paste/P486771278/ ``` pickle_and_flatbuffer: https://www.internalfb.com/intern/dgw/graph/?build_id=P486770901 pickle: https://www.internalfb.com/intern/dgw/graph/?build_id=P486771278 Reviewed By: iseeyuan Differential Revision: D34805092 fbshipit-source-id: 57f3fc81d68fce941a050c35bd8e6f05951183b3 (cherry picked from commit 671ae4ed29e65b86ffe507a503548d3e86ab0ea4)
46 lines
1.1 KiB
C++
46 lines
1.1 KiB
C++
#include <torch/csrc/jit/api/module.h>
|
|
#include <torch/csrc/jit/serialization/export.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
void Module::save(std::ostream& out, const ExtraFilesMap& extra_files) const {
|
|
ExportModule(*this, out, extra_files, false /* bytecode_format */);
|
|
}
|
|
|
|
void Module::save(const std::string& filename, const ExtraFilesMap& extra_files)
|
|
const {
|
|
ExportModule(*this, filename, extra_files, false /* bytecode_format */);
|
|
}
|
|
|
|
void Module::_save_for_mobile(
|
|
std::ostream& out,
|
|
const ExtraFilesMap& extra_files,
|
|
bool save_mobile_debug_info,
|
|
bool use_flatbuffer) const {
|
|
ExportModule(
|
|
*this,
|
|
out,
|
|
extra_files,
|
|
true /* bytecode_format */,
|
|
save_mobile_debug_info,
|
|
use_flatbuffer);
|
|
}
|
|
|
|
void Module::_save_for_mobile(
|
|
const std::string& filename,
|
|
const ExtraFilesMap& extra_files,
|
|
bool save_mobile_debug_info,
|
|
bool use_flatbuffer) const {
|
|
ExportModule(
|
|
*this,
|
|
filename,
|
|
extra_files,
|
|
true /* bytecode_format */,
|
|
save_mobile_debug_info,
|
|
use_flatbuffer);
|
|
}
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|