Program Listing for File output-archive.h

↰ Return to documentation for file (torch/csrc/api/include/torch/serialize/output-archive.h)

#pragma once

#include <torch/csrc/Export.h>
#include <torch/csrc/jit/api/module.h>

#include <iosfwd>
#include <memory>
#include <string>
#include <utility>

namespace at {
class Tensor;
} // namespace at

namespace torch {
using at::Tensor;
namespace jit {
struct Module;
} // namespace jit
} // namespace torch

namespace torch {
namespace serialize {
class TORCH_API OutputArchive final {
 public:
  explicit OutputArchive(std::shared_ptr<jit::CompilationUnit> cu);
  explicit OutputArchive()
      : cu_(std::make_shared<jit::CompilationUnit>()),
        module_("__torch__.Module", cu_) {}

  // Move is allowed.
  OutputArchive(OutputArchive&&) = default;
  OutputArchive& operator=(OutputArchive&&) = default;

  // Copy is disallowed.
  OutputArchive(OutputArchive&) = delete;
  OutputArchive& operator=(OutputArchive&) = delete;

  std::shared_ptr<jit::CompilationUnit> compilation_unit() const {
    return cu_;
  }

  void write(const std::string& key, const c10::IValue& ivalue);

  void write(
      const std::string& key,
      const Tensor& tensor,
      bool is_buffer = false);

  void write(const std::string& key, OutputArchive& nested_archive);

  void save_to(const std::string& filename);

  void save_to(std::ostream& stream);

  void save_to(const std::function<size_t(const void*, size_t)>& func);

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

 private:
  std::shared_ptr<jit::CompilationUnit> cu_;
  jit::Module module_;
};
} // namespace serialize
} // namespace torch