frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-11-04 16:04:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

Expose _export_data from C++ to Python (#79207)

Browse Source 
Summary:
https://www.internalfb.com/code/fbsource/[477a5768452957f87e56044169de47f051197567]/fbcode/caffe2/torch/csrc/jit/mobile/train/export_data.cpp
export_data is used to serialize data.

I binded this method to Python with PyBind11

Test Plan:
Wrote a file pybind_check.py which checks if the binding works.

Then, tried to read the produced data file from C++ with "torch::jit::_load_parameters" and checked that content matched.

Differential Revision: D37029253

Pull Request resolved: https://github.com/pytorch/pytorch/pull/79207
Approved by: https://github.com/qihqi
This commit is contained in:
Luka Mushkudiani
2022-06-10 00:41:33 +00:00
committed by PyTorch MergeBot
parent 338bfe6315
commit c0a7c1d02e
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
torch/csrc/jit/python/script_init.cpp
View File

@ -59,6 +59,7 @@
#include <pybind11/pybind11.h> #include <pybind11/pybind11.h>
#include <pybind11/stl.h> #include <pybind11/stl.h>
#include <pybind11/stl_bind.h> #include <pybind11/stl_bind.h>
#include <torch/csrc/jit/mobile/train/export_data.h>
#include <chrono> #include <chrono>
#include <cstddef> #include <cstddef>
#include <memory> #include <memory>
@ -2294,6 +2295,14 @@ void initJitScriptBindings(PyObject* module) {
} }
}); });
m.def(
"_save_parameters",
[](const std::map<std::string, at::Tensor>& map,
const std::string& filename,
bool use_flatbuffer = false) {
_save_parameters(map, filename, use_flatbuffer);
});
initScriptDictBindings(module); initScriptDictBindings(module);
initScriptListBindings(module); initScriptListBindings(module);
} }