[C++ API] Protobuf serialization by goldsborough · Pull Request #11619 · pytorch/pytorch

goldsborough · 2018-09-13T02:49:46Z

This PR serves two purposes:

Design an abstraction over a serialization scheme for C++ modules, optimizers and tensors in general,
Add serialization to the ONNX/PyTorch proto format.

This is currently a rough prototype I coded up today, to get quick feedback.

For this I propose the following serialization interface within the C++ API:

namespace torch { namespace serialize {
class Reader {
 public:
  virtual ~Reader() = default;
  virtual void read(const std::string& key, Tensor& tensor, bool is_buffer = false) = 0;
  virtual void finish() { }
};

class Writer {
 public:
  virtual ~Reader() = default;
  virtual void writer(const std::string& key, const Tensor& tensor, bool is_buffer = false) = 0;
  virtual void finish() { }
};
}} // namespace torch::serialize

There are then subclasses of these two for (1) Cereal and (2) Protobuf (called the "DefaultWriter" and "DefaultReader" to hide the implementation details). See torch/serialize/cereal.h and torch/serialize/default.h. This abstraction and subclassing for these two allows us to:

Provide a cereal-less serialization forward that we can ship and iterate on going forward,
Provide no-friction backwards compatibility with existing C++ API uses, mainly StarCraft.

The user-facing API is (conceptually):

void torch::save(const Module& module, Writer& writer);
void torch::save(const Optimizer& optimizer, Writer& writer);
void torch::read(Module& module, Reader& reader);
void torch::read(Optimizer& optimizer, Reader& reader);

with implementations for both optimizers and modules that write into the Writer and read from the Reader

@ebetica @ezyang @zdevito @dzhulgakov

ebetica · 2018-09-13T02:56:43Z

I wouldn't worry about (2). How would you change this if you didn't have to support Cereal at all?

ezyang · 2018-09-13T04:23:01Z

Is the problem this is solving primarily "how to serialize to both protobuf and cereal", or are there other problems?

dzhulgakov

Looks neat!

For future we should have some mechanism of encoding at least typename of the c++ module in the serialized file for reference. Since for now modules are different from jit modules I guess it's ok

P.S. reusing jit module serialization is... cute :)

torch/csrc/api/include/torch/serialize/reader.h

torch/csrc/api/include/torch/optim/adam.h

torch/csrc/api/src/serialize/default.cpp

goldsborough · 2018-09-18T16:58:58Z

After discussing with @ebetica we found that backwards compatibility with cereal is not as important and that StarCraft would be happy to switch to protobuf. For this reason I've opted to specialize my design a little more towards our Protobuf serialization format. The new primitives are InputArchive and OutputArchive, which both wrap a ScriptModule, which we use as a vehicle for serialization. The API for e.g. InputArchive is:

class InputArchive {
 public:
  void read(const std::string& key, Tensor& tensor, bool is_buffer = false);
  void read(const std::string& key, InputArchive& archive);

  template <typename... Ts>
  void operator()(Ts&&... ts) {
    read(std::forward<Ts>(ts)...);
  }

 private:
  std::shared_ptr<jit::script::Module> module_;
};

which gets populated e.g. in nn::Module::load:

void Module::load(serialize::InputArchive& archive) {
  for (auto& parameter : parameters_) {
    archive.read(parameter.key, parameter.value);
  }
  for (auto& buffer : buffers_) {
    archive.read(buffer.key, buffer.value, /*is_buffer=*/true);
  }
  for (const auto& child : children_) {
      serialize::InputArchive child_archive;
      archive.read(child.key, child_archive);
      child.value->load(child_archive);
  }
}

This API is not intended for inheritence. The reason why is that InputArchives are recursively built (1 InputArchive = 1 ScriptModule) and have to be re-instantiated inside of save() or load(). This re-instantiation wouldn't work (nicely) if InputArchive were an abstract class.

This means we only support protobuf serialization out of the box. However, I've made sure that all parameters and buffers in modules and optimizers that one would want to store are properly exposed. This means if one wanted to, one could write a new serialization library on top of the Module base class. I believe this is more the strategy we have in Python, where we have one blessed serialization API (Pickle), but generally modules and optimizers expose enough state that anyone could dump their models into a Thrift or JSON interface.

CC @ebetica @dzhulgakov @zdevito @apaszke @ezyang

ebetica · 2018-09-18T19:05:54Z

Do you think overloading torch::save for a stream is nice? It just feels more C++-y to me, and is more flexible for users.

ebetica

Happy with the API changes, I didn't do a deep reading of this PR.

torch/csrc/api/include/torch/serialize.h

.gitmodules

test/cpp/api/util.h

ezyang

Please don't leak file descriptors

facebook-github-bot

goldsborough has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

Summary: This PR serves two purposes: 1. Design an abstraction over a serialization scheme for C++ modules, optimizers and tensors in general, 2. Add serialization to the ONNX/PyTorch proto format. This is currently a rough prototype I coded up today, to get quick feedback. For this I propose the following serialization interface within the C++ API: ```cpp namespace torch { namespace serialize { class Reader { public: virtual ~Reader() = default; virtual void read(const std::string& key, Tensor& tensor, bool is_buffer = false) = 0; virtual void finish() { } }; class Writer { public: virtual ~Reader() = default; virtual void writer(const std::string& key, const Tensor& tensor, bool is_buffer = false) = 0; virtual void finish() { } }; }} // namespace torch::serialize ``` There are then subclasses of these two for (1) Cereal and (2) Protobuf (called the "DefaultWriter" and "DefaultReader" to hide the implementation details). See `torch/serialize/cereal.h` and `torch/serialize/default.h`. This abstraction and subclassing for these two allows us to: 1. Provide a cereal-less serialization forward that we can ship and iterate on going forward, 2. Provide no-friction backwards compatibility with existing C++ API uses, mainly StarCraft. The user-facing API is (conceptually): ```cpp void torch::save(const Module& module, Writer& writer); void torch::save(const Optimizer& optimizer, Writer& writer); void torch::read(Module& module, Reader& reader); void torch::read(Optimizer& optimizer, Reader& reader); ``` with implementations for both optimizers and modules that write into the `Writer` and read from the `Reader` ebetica ezyang zdevito dzhulgakov Pull Request resolved: pytorch/pytorch#11619 Differential Revision: D9984664 Pulled By: goldsborough fbshipit-source-id: e03afaa646221546e7f93bb8dfe3558e384a5847

Summary: This PR serves two purposes: 1. Design an abstraction over a serialization scheme for C++ modules, optimizers and tensors in general, 2. Add serialization to the ONNX/PyTorch proto format. This is currently a rough prototype I coded up today, to get quick feedback. For this I propose the following serialization interface within the C++ API: ```cpp namespace torch { namespace serialize { class Reader { public: virtual ~Reader() = default; virtual void read(const std::string& key, Tensor& tensor, bool is_buffer = false) = 0; virtual void finish() { } }; class Writer { public: virtual ~Reader() = default; virtual void writer(const std::string& key, const Tensor& tensor, bool is_buffer = false) = 0; virtual void finish() { } }; }} // namespace torch::serialize ``` There are then subclasses of these two for (1) Cereal and (2) Protobuf (called the "DefaultWriter" and "DefaultReader" to hide the implementation details). See `torch/serialize/cereal.h` and `torch/serialize/default.h`. This abstraction and subclassing for these two allows us to: 1. Provide a cereal-less serialization forward that we can ship and iterate on going forward, 2. Provide no-friction backwards compatibility with existing C++ API uses, mainly StarCraft. The user-facing API is (conceptually): ```cpp void torch::save(const Module& module, Writer& writer); void torch::save(const Optimizer& optimizer, Writer& writer); void torch::read(Module& module, Reader& reader); void torch::read(Optimizer& optimizer, Reader& reader); ``` with implementations for both optimizers and modules that write into the `Writer` and read from the `Reader` ebetica ezyang zdevito dzhulgakov Pull Request resolved: pytorch#11619 Differential Revision: D9984664 Pulled By: goldsborough fbshipit-source-id: e03afaa646221546e7f93bb8dfe3558e384a5847

goldsborough requested review from apaszke, colesbury, ebetica, ezyang, gchanan, soumith and zdevito as code owners September 13, 2018 02:49

dzhulgakov requested a review from houseroad September 13, 2018 06:04

dzhulgakov reviewed Sep 13, 2018

View reviewed changes

goldsborough force-pushed the cpp-serialization branch 2 times, most recently from 36d2455 to 3f83954 Compare September 18, 2018 06:33

goldsborough changed the title ~~[C++ API][Prototype] Abstract serialization and add Protobuf serialization~~ [C++ API][Prototype] Add Protobuf serialization Sep 18, 2018

goldsborough changed the title ~~[C++ API][Prototype] Add Protobuf serialization~~ [C++ API] Protobuf serialization Sep 18, 2018

goldsborough force-pushed the cpp-serialization branch 3 times, most recently from 6379c68 to 7881ad9 Compare September 18, 2018 18:01

goldsborough force-pushed the cpp-serialization branch from 7881ad9 to be74fad Compare September 18, 2018 20:10

ebetica approved these changes Sep 20, 2018

View reviewed changes

goldsborough force-pushed the cpp-serialization branch from be74fad to 95ac9b7 Compare September 20, 2018 18:55

ezyang reviewed Sep 20, 2018

View reviewed changes

torch/csrc/api/include/torch/serialize.h Outdated

This comment was marked as off-topic.

Sign in to view

ezyang reviewed Sep 20, 2018

View reviewed changes

.gitmodules Outdated

This comment was marked as off-topic.

Sign in to view

ezyang reviewed Sep 20, 2018

View reviewed changes

test/cpp/api/util.h Outdated

This comment was marked as off-topic.

Sign in to view

ezyang requested changes Sep 20, 2018

View reviewed changes

goldsborough added 3 commits September 20, 2018 15:33

[C++ API] Abstract serialization

28e187c

Support exclusively protobuf

ca24c1b

Dont remove cereal submodule

7ff03b9

goldsborough force-pushed the cpp-serialization branch 2 times, most recently from 7bb46aa to 03d08f6 Compare September 20, 2018 23:18

Make torch::load() return a tensor and dont leak tempfiles

1e92bf3

goldsborough force-pushed the cpp-serialization branch from 03d08f6 to 1e92bf3 Compare September 21, 2018 00:00

facebook-github-bot reviewed Sep 21, 2018

View reviewed changes

ezyang approved these changes Sep 21, 2018

View reviewed changes

facebook-github-bot closed this in d712a71 Sep 21, 2018

ezyang added the merged label Jun 26, 2019

Conversation

goldsborough commented Sep 13, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ebetica commented Sep 13, 2018

Uh oh!

ezyang commented Sep 13, 2018

Uh oh!

dzhulgakov left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

goldsborough commented Sep 18, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ebetica commented Sep 18, 2018

Uh oh!

ebetica left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

ezyang left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

goldsborough commented Sep 13, 2018 •

edited

Loading

goldsborough commented Sep 18, 2018 •

edited

Loading