Load tensors directly from pickle archive #23281
Conversation
This allows tensors to be loaded directly from the pickle archive, using the same encoding as the eager serializer. The one difference is that the tensor data is still stored in a zip file rather than appended to the archive.
[WIP] Load tensors directly from pickle archive This allows tensors to be loaded directly from the pickle archive, using the same encoding as the eager serializer. The one difference is that the tensor data is still stored in a zip file rather than appended to the archive. gh-metadata: pytorch pytorch 23281 gh/zdevito/78/head
| result = autograd::make_variable(result, requires_grad); | ||
| stack_.push_back(std::move(result)); | ||
| }); | ||
| } else if (module_name == "collections" && class_name == "OrderedDict") { |
There was a problem hiding this comment.
this is a bit scary. Is there are stricter condition we can enforce here (e.g. that it's indeed in Tensor deserialization only)
torch/csrc/jit/pickler.cpp
Outdated
| } | ||
| }); | ||
| } else if ( | ||
| module_name == "torch._utils" && class_name == "_rebuild_tensor_v2") { |
There was a problem hiding this comment.
would be nice to extract to a subfunction and explicitly document in both python and C++ that those functions should match semantically
Load tensors directly from pickle archive This allows tensors to be loaded directly from the pickle archive, using the same encoding as the eager serializer. The one difference is that the tensor data is still stored in a zip file rather than appended to the archive. gh-metadata: pytorch pytorch 23281 gh/zdevito/78/head
pytorchbot
added
caffe2
oncall: distributed
zdevito
changed the title
Load tensors directly from pickle archive This allows tensors to be loaded directly from the pickle archive, using the same encoding as the eager serializer. The one difference is that the tensor data is still stored in a zip file rather than appended to the archive. gh-metadata: pytorch pytorch 23281 gh/zdevito/78/head
|
|
||
| ScriptModuleSerializer2(std::ostream* ofs) : ofs_(), writer_(ofs) {} | ||
|
|
||
| void writeArchive(const std::string& archive_name, const IValue& value) { |
There was a problem hiding this comment.
Doesn't look like this needs to be public?
torch/csrc/jit/pickle.h
Outdated
| /// `bounds_checker` is a function that returns `true` if the reader can read | ||
| /// more data, and `false` if it cannot (i.e. if a stream has hit its end of | ||
| /// file) | ||
| /// |
There was a problem hiding this comment.
This comment is now out of date, references the bounds_checker arg
|
|
||
| public: | ||
| // tensors inside the pickle are references to the tensor_table | ||
| Unpickler( |
There was a problem hiding this comment.
Is this version legacy-only? If so, we should note it in the comments
There was a problem hiding this comment.
It is possible that RPC will be using this version, and completely handling tensor serialization on its own. Not sure though.
| module_name == "torch._utils" && | ||
| (class_name == "_rebuild_tensor_v2" || | ||
| class_name == "_rebuild_qtensor")) { | ||
| bool quantized = class_name == "_rebuild_qtensor"; |
There was a problem hiding this comment.
ugh, I don't like that quantized tensors are somehow special case for the serialization format
There was a problem hiding this comment.
I agree, but this is how it works in eager serialization, so we are stuck with it if we want torch.save from torchscript to work with torch.load
| /// file) | ||
| /// | ||
| /// See `torch::pickle` for details. | ||
| TORCH_API IValue unpickle( |
There was a problem hiding this comment.
In general we should think a bit about how we want the Pickle C++ API to look at the end of this. @driazati, any thoughts? At least I think we should get rid of tensor_table if possible
There was a problem hiding this comment.
Any use of the tensor_table would be covered by pickle's persistent_id/load functions (which we could implement if we need to in the future). This is already how eager's serialization works anyways, so it's fine to get rid of the tensor_table
|
|
||
| void writeArchive(const std::string& archive_name, const IValue& value) { | ||
| std::vector<char> data; | ||
| Pickler data_pickle( |
There was a problem hiding this comment.
It looks like jit::pickle(...) would be fine here instead of the manual use of the Pickler
There was a problem hiding this comment.
Right after this it uses data_pickle.tensorData which is not exposed by those functions.
| ss << archive_name << "/" << name; | ||
| return std::get<0>(reader_->getRecord(ss.str())); | ||
| }; | ||
| Unpickler unpickler( |
There was a problem hiding this comment.
Same here, this can be jit::unpickle
There was a problem hiding this comment.
user-exposed pickle functions do not take the read_record or device_ arguments needed here to load the tensors.
Load tensors directly from pickle archive This allows tensors to be loaded directly from the pickle archive, using the same encoding as the eager serializer. The one difference is that the tensor data is still stored in a zip file rather than appended to the archive. gh-metadata: pytorch pytorch 23281 gh/zdevito/78/head
Stack from ghstack:
This allows tensors to be loaded directly from the pickle archive,
using the same encoding as the eager serializer. The one difference
is that the tensor data is still stored in a zip file rather
than appended to the archive, which still allows for aligning the tensor data.
With this change, we no longer rely on the json file for any information, so this PR also stops writing the json file and loads modules directly from the pickle.
New layout is:
Differential Revision: D16452815