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

[deploy] torch::deploy API (#51754)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51754

This API allows you to manage multiple python interpreters in a single
process to deploy PyTorch models packaged with torch.package.

torch/csrc/deploy/deploy.h contains the API definition
torch/csrc/deploy/test_deploy.cpp has some examples.

Notes:
* mutex is added to PyTorchStreamReader to make it safe to use from multiple threads at once.
* USE_DEPLOY is only true for the special libtorch_deployinterpreter.so library, when enabled
  we use a hash table to maintain PyObject <> at::Tensor mappping rather than the internal pointer
  in Tensor since >1 interpreter may have a reference to the tensor.
* serialization.py has some additional functions for creating pickle objects
  but keeping storages in memory for use transfering tensors between interpreters

Test Plan: Imported from OSS

Reviewed By: wconstab

Differential Revision: D26329468

Pulled By: zdevito

fbshipit-source-id: d75f4ebb9a27f1d911179d9996041bcb3ca04a07
This commit is contained in:
Zachary DeVito
2021-02-18 02:28:08 -08:00
committed by Facebook GitHub Bot
parent 9cf6be6b3e
commit 60518d10f6
26 changed files with 1735 additions and 513 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
caffe2/serialize/inline_container.cc
View File

@ -189,6 +189,7 @@ size_t getPadding(
}
bool PyTorchStreamReader::hasRecord(const std::string& name) {
std::lock_guard<std::mutex> guard(reader_lock_);
std::string ss = archive_name_plus_slash_ + name;
mz_zip_reader_locate_file(ar_.get(), ss.c_str(), nullptr, 0);
bool result = ar_->m_last_error != MZ_ZIP_FILE_NOT_FOUND;
@ -200,6 +201,7 @@ bool PyTorchStreamReader::hasRecord(const std::string& name) {
}
std::vector<std::string> PyTorchStreamReader::getAllRecords() {
std::lock_guard<std::mutex> guard(reader_lock_);
mz_uint num_files = mz_zip_reader_get_num_files(ar_.get());
std::vector<std::string> out;
char buf[MZ_ZIP_MAX_ARCHIVE_FILENAME_SIZE];
@ -232,6 +234,7 @@ size_t PyTorchStreamReader::getRecordID(const std::string& name) {
// return dataptr, size
std::tuple<at::DataPtr, size_t> PyTorchStreamReader::getRecord(const std::string& name) {
std::lock_guard<std::mutex> guard(reader_lock_);
size_t key = getRecordID(name);
mz_zip_archive_file_stat stat;
mz_zip_reader_file_stat(ar_.get(), key, &stat);
@ -248,6 +251,7 @@ static int64_t read_le_16(uint8_t* buf) {
}
size_t PyTorchStreamReader::getRecordOffset(const std::string& name) {
std::lock_guard<std::mutex> guard(reader_lock_);
mz_zip_archive_file_stat stat;
mz_zip_reader_file_stat(ar_.get(), getRecordID(name), &stat);
valid("retrieving file meta-data for ", name.c_str());