mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
Back out "Revert D34805092: Extend _save_for_mobile and _load_for_mobile to support flatbuffer format; Default format is pickle + Change buck targets to support only pickle and pickle + flatbuffer for migration" (#74594)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74594
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;
```
Original commit changeset: 780dfb6fd6ba
Original Phabricator Diff: D34805092 (284b2b7135)
ghstack-source-id: 152044801
(Note: this ignores all push blocking failures!)
Test Plan:
CI
```
~/fbsource/fbcode] cd ~/fbsource/fbcode/ && buck test -c fbcode.caffe2_enable_flatbuffer=1 //caffe2/test/cpp/jit:jit -- FlatbufferTest.ExtraFiles
Parsing buck files: finished in 0.9 sec
Building: finished in 5.3 sec (100%) 12992/54304 jobs, 0/54304 updated
Total time: 6.2 sec
More details at https://www.internalfb.com/intern/buck/build/2b387fff-f813-4cfa-b53f-eb2378630d4e
BUILD SUCCEEDED
Tpx test run coordinator for Facebook. See https://fburl.com/tpx for details.
Running with tpx session id: f93a84d6-e7ce-41a0-a97f-0ef3fa6d199d
Trace available for this run at /tmp/tpx-20220323-134108.766518-f93a84d6-e7ce-41a0-a97f-0ef3fa6d199d/trace.log
RemoteExecution session id: reSessionID-f93a84d6-e7ce-41a0-a97f-0ef3fa6d199d-tpx
Started reporting to test run: https://www.internalfb.com/intern/testinfra/testrun/4503599723101693
✓ ListingSuccess: caffe2/test/cpp/jit:jit : 486 tests discovered (19.122)
✓ Pass: caffe2/test/cpp/jit:jit - FlatbufferTest.ExtraFiles (0.187)
Summary
Pass: 1
ListingSuccess: 1
If you need help understanding your runs, please follow the wiki: https://fburl.com/posting_in_tpx_users
Finished test run: https://www.internalfb.com/intern/testinfra/testrun/4503599723101693
```
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: D35067157
fbshipit-source-id: 9044259c17a2e0da79bd6aedb28efbdfd57e23e0
(cherry picked from commit f738069ec3a72e79da56172741d027de514e9e5f)
This commit is contained in:
Pavithran Ramachandran
committed by
PyTorch MergeBot
PyTorch MergeBot
parent
d64e7634ff
commit
fc2cf3d26f
@ -10,6 +10,10 @@
|
||||
#include <caffe2/serialize/inline_container.h>
|
||||
#include <caffe2/serialize/versions.h>
|
||||
#include <torch/csrc/jit/api/compilation_unit.h>
|
||||
#include <torch/csrc/jit/mobile/file_format.h>
|
||||
#if defined(ENABLE_FLATBUFFER)
|
||||
#include <torch/csrc/jit/mobile/flatbuffer_loader.h>
|
||||
#endif
|
||||
#include <torch/csrc/jit/mobile/interpreter.h>
|
||||
#include <torch/csrc/jit/mobile/observer.h>
|
||||
#include <torch/csrc/jit/mobile/type_parser.h>
|
||||
@ -536,18 +540,47 @@ mobile::Module _load_for_mobile(
|
||||
std::istream& in,
|
||||
c10::optional<at::Device> device,
|
||||
ExtraFilesMap& extra_files) {
|
||||
std::unique_ptr<IStreamAdapter> rai = std::make_unique<IStreamAdapter>(&in);
|
||||
auto module = _load_for_mobile(std::move(rai), device, extra_files);
|
||||
return module;
|
||||
auto format = getFileFormat(in);
|
||||
switch (format) {
|
||||
case FileFormat::ZipFileFormat: {
|
||||
std::unique_ptr<IStreamAdapter> rai =
|
||||
std::make_unique<IStreamAdapter>(&in);
|
||||
auto module = _load_for_mobile(std::move(rai), device, extra_files);
|
||||
return module;
|
||||
}
|
||||
#if defined(ENABLE_FLATBUFFER)
|
||||
case FileFormat::FlatbufferFileFormat: {
|
||||
std::shared_ptr<char> data;
|
||||
size_t size = 0;
|
||||
std::tie(data, size) = get_stream_content(in);
|
||||
auto* flatbuffer_module =
|
||||
mobile::serialization::GetMutableModule(data.get());
|
||||
mobile::Module m = initialize_mobile_module(flatbuffer_module);
|
||||
parseExtraFiles(flatbuffer_module, extra_files);
|
||||
return m;
|
||||
}
|
||||
#else
|
||||
case FileFormat::FlatbufferFileFormat: {
|
||||
TORCH_CHECK(
|
||||
false,
|
||||
"Flatbuffer input file but the build hasn't enabled flatbuffer");
|
||||
}
|
||||
#endif
|
||||
default: {
|
||||
TORCH_CHECK(false, "Format error");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mobile::Module _load_for_mobile(
|
||||
const std::string& filename,
|
||||
c10::optional<at::Device> device,
|
||||
ExtraFilesMap& extra_files) {
|
||||
std::unique_ptr<FileAdapter> rai = std::make_unique<FileAdapter>(filename);
|
||||
auto module = _load_for_mobile(std::move(rai), device, extra_files);
|
||||
return module;
|
||||
return _load_for_mobile(
|
||||
filename,
|
||||
device,
|
||||
extra_files,
|
||||
/*module_load_options=*/_default_mobile_module_load_options);
|
||||
}
|
||||
|
||||
mobile::Module _load_for_mobile(
|
||||
@ -555,10 +588,37 @@ mobile::Module _load_for_mobile(
|
||||
c10::optional<at::Device> device,
|
||||
ExtraFilesMap& extra_files,
|
||||
uint64_t module_load_options) {
|
||||
std::unique_ptr<FileAdapter> rai = std::make_unique<FileAdapter>(filename);
|
||||
auto module = _load_for_mobile_impl(
|
||||
std::move(rai), device, extra_files, module_load_options);
|
||||
return module;
|
||||
auto format = getFileFormat(filename);
|
||||
switch (format) {
|
||||
case FileFormat::ZipFileFormat: {
|
||||
std::unique_ptr<FileAdapter> rai =
|
||||
std::make_unique<FileAdapter>(filename);
|
||||
auto module = _load_for_mobile_impl(
|
||||
std::move(rai), device, extra_files, module_load_options);
|
||||
return module;
|
||||
}
|
||||
#if defined(ENABLE_FLATBUFFER)
|
||||
case FileFormat::FlatbufferFileFormat: {
|
||||
std::shared_ptr<char> data;
|
||||
size_t size = 0;
|
||||
std::tie(data, size) = get_file_content(filename.c_str());
|
||||
auto* flatbuffer_module =
|
||||
mobile::serialization::GetMutableModule(data.get());
|
||||
mobile::Module m = initialize_mobile_module(flatbuffer_module);
|
||||
parseExtraFiles(flatbuffer_module, extra_files);
|
||||
return m;
|
||||
}
|
||||
#else
|
||||
case FileFormat::FlatbufferFileFormat: {
|
||||
TORCH_CHECK(
|
||||
false,
|
||||
"Flatbuffer input file but the build hasn't enabled flatbuffer");
|
||||
}
|
||||
#endif
|
||||
default: {
|
||||
TORCH_CHECK(false, "Format error");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mobile::Module _load_for_mobile(
|
||||
|
||||
Reference in New Issue
Block a user