frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

stream pkl (#79931)

Browse Source 
Summary: Let the stream version of pickle loader to short circuit early

Test Plan:
this might affects loading peak mem
before:

Differential Revision: D37304855

Pull Request resolved: https://github.com/pytorch/pytorch/pull/79931
Approved by: https://github.com/junesg, https://github.com/davidberard98
This commit is contained in:
Han Qi (qihqi)
2022-06-22 19:30:12 +00:00
committed by PyTorch MergeBot
parent e3d0a3ca88
commit 9f866b8dbb
1 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
torch/csrc/jit/mobile/import.cpp
View File

@ -560,11 +560,17 @@ mobile::Module _load_for_mobile(
std::istream& in,
c10::optional<at::Device> device,
ExtraFilesMap& extra_files) {
std::shared_ptr<char> data;
size_t size = 0;
std::tie(data, size) = get_stream_content(in);
return _load_mobile_from_bytes(
data, size, device, extra_files, kDefaultMobileLoadOptions);
if (getFileFormat(in) == FileFormat::FlatbufferFileFormat) {
std::shared_ptr<char> data;
size_t size = 0;
std::tie(data, size) = get_stream_content(in);
return _load_mobile_from_bytes(
data, size, device, extra_files, kDefaultMobileLoadOptions);
}
std::unique_ptr<IStreamAdapter> rai = std::make_unique<IStreamAdapter>(&in);
auto module = _load_for_mobile_impl(
std::move(rai), device, extra_files, kDefaultMobileLoadOptions);
return module;
}
mobile::Module _load_for_mobile(