-
Notifications
You must be signed in to change notification settings - Fork 0
Revise serialize #73
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Revise serialize #73
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,5 @@ | ||
| // torch | ||
| #include <torch/script.h> | ||
| #include <torch/serialize.h> | ||
|
|
||
| // kintera | ||
|
|
@@ -15,17 +16,31 @@ void save_tensors(const std::map<std::string, torch::Tensor>& tensor_map, | |
| archive.save_to(filename); | ||
| } | ||
|
|
||
| void load_tensors(std::map<std::string, torch::Tensor>& tensor_map, | ||
| const std::string& filename) { | ||
| torch::serialize::InputArchive archive; | ||
| std::map<std::string, torch::Tensor> load_tensors(const std::string& filename) { | ||
| std::map<std::string, torch::Tensor> data; | ||
|
|
||
| // get keys | ||
| torch::jit::Module m = torch::jit::load(filename); | ||
|
|
||
| for (const auto& p : m.named_parameters(/*recurse=*/true)) { | ||
| data[p.name] = p.value; | ||
| } | ||
|
|
||
| for (const auto& p : m.named_buffers(/*recurse=*/true)) { | ||
| data[p.name] = p.value; | ||
| } | ||
|
|
||
| /*torch::serialize::InputArchive archive; | ||
| archive.load_from(filename); | ||
| for (auto& pair : tensor_map) { | ||
| for (auto& pair : data) { | ||
| try { | ||
| archive.read(pair.first, pair.second); | ||
| } catch (const c10::Error& e) { | ||
| // skip missing tensors | ||
| } | ||
| } | ||
| }*/ | ||
|
Comment on lines
+33
to
+41
|
||
|
|
||
| return data; | ||
| } | ||
|
|
||
| } // namespace kintera | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -11,7 +11,6 @@ namespace kintera { | |
| void save_tensors(const std::map<std::string, torch::Tensor>& tensor_map, | ||
| const std::string& filename); | ||
|
|
||
| void load_tensors(std::map<std::string, torch::Tensor>& tensor_map, | ||
| const std::string& filename); | ||
| std::map<std::string, torch::Tensor> load_tensors(const std::string& filename); | ||
|
||
|
|
||
| } // namespace kintera | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
load_tensorsnow callstorch::jit::loadand returnsnamed_parameters/named_buffers, butsave_tensorswrites atorch::serialize::OutputArchive. These serialization formats are not compatible, soload_tensorswill fail (or silently return unrelated data) when given files produced bysave_tensors. Consider either (a) switchingload_tensorsback totorch::serialize::InputArchiveto read the same keys that were written, or (b) updatingsave_tensorsto write a TorchScriptModule(register tensors as buffers/parameters) if TorchScript loading is the new intended format.