Program Listing for File input-archive.h

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

#pragma once

#include <c10/core/Device.h>
#include <c10/util/Optional.h>
#include <torch/csrc/Export.h>
#include <torch/csrc/jit/api/module.h>
#include <torch/types.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 InputArchive final {
 public:
  InputArchive();

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

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

  ~InputArchive() = default;

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

  bool try_read(const std::string& key, c10::IValue& ivalue);

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

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

  bool try_read(const std::string& key, InputArchive& archive);

  void read(const std::string& key, InputArchive& archive);

  void load_from(
      const std::string& filename,
      c10::optional<torch::Device> device = c10::nullopt);

  void load_from(
      std::istream& stream,
      c10::optional<torch::Device> device = c10::nullopt);

  // Loads given the specified flat array.
  void load_from(
      const char* data,
      size_t size,
      c10::optional<torch::Device> device = c10::nullopt);

  // Loads given the specified read and size functions.
  void load_from(
      const std::function<size_t(uint64_t pos, void* buf, size_t nbytes)>&
          read_func,
      const std::function<size_t(void)>& size_func,
      c10::optional<torch::Device> device = c10::nullopt);

  // Returns the vector of keys in the input archive.
  std::vector<std::string> keys();

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

 private:
  jit::Module module_;
  std::string hierarchy_prefix_;
};
} // namespace serialize
} // namespace torch