Caffe2 - C++ API
A deep learning, cross platform ML framework
output-archive.cpp
1 #include <torch/serialize/output-archive.h>
2 
3 #include <torch/types.h>
4 #include <torch/utils.h>
5 
6 #include <torch/csrc/jit/export.h>
7 #include <torch/csrc/jit/script/module.h>
8 
9 #include <c10/util/Exception.h>
10 
11 #include <memory>
12 #include <ostream>
13 #include <string>
14 
15 namespace torch {
16 namespace serialize {
18  : module_(std::make_shared<jit::script::Module>()) {}
19 
21  const std::string& key,
22  const Tensor& tensor,
23  bool is_buffer) {
24  module_->register_parameter(key, tensor, is_buffer);
25 }
26 
28  const std::string& key,
29  OutputArchive& nested_archive) {
30  module_->register_module(key, nested_archive.module_);
31 }
32 
33 void OutputArchive::save_to(const std::string& filename) {
34  AT_ASSERT(module_ != nullptr);
35  jit::ExportModule(*module_, filename);
36 }
37 
38 void OutputArchive::save_to(std::ostream& stream) {
39  AT_ASSERT(module_ != nullptr);
40  jit::ExportModule(*module_, stream);
41 }
42 } // namespace serialize
43 } // namespace torch
Definition: jit_type.h:17
void write(const std::string &key, const Tensor &tensor, bool is_buffer=false)
Writes a (key, tensor) pair to the OutputArchive, and marks it as being or not being a buffer (non-di...
OutputArchive()
Default-constructs the OutputArchive.
void save_to(const std::string &filename)
Saves the OutputArchive into a serialized representation in a file at filename.