module.h

#pragma once

#include <torch/nn/pimpl.h>
#include <torch/ordered_dict.h>
#include <torch/serialize/archive.h>
#include <torch/types.h>

#include <ATen/ATen.h>

#include <functional>
#include <iosfwd>
#include <map>
#include <memory>
#include <string>
#include <type_traits>

namespace torch {
namespace nn {

class TORCH_API Module : public std::enable_shared_from_this<Module> {
 public:
  using ModuleApplyFunction = std::function<void(Module&)>;
  using ConstModuleApplyFunction = std::function<void(const Module&)>;
  using NamedModuleApplyFunction =
      std::function<void(const std::string&, Module&)>;
  using ConstNamedModuleApplyFunction =
      std::function<void(const std::string&, const Module&)>;
  using ModulePointerApplyFunction =
      std::function<void(const std::shared_ptr<Module>&)>;
  using NamedModulePointerApplyFunction =
      std::function<void(const std::string&, const std::shared_ptr<Module>&)>;

  explicit Module(std::string name);

  Module();

  virtual ~Module() = default;

  const std::string& name() const noexcept;

  virtual std::shared_ptr<Module> clone(
      const optional<Device>& device = nullopt) const;

  void apply(const ModuleApplyFunction& function);

  void apply(const ConstModuleApplyFunction& function) const;

  void apply(
      const NamedModuleApplyFunction& function,
      const std::string& name_prefix = std::string());

  void apply(
      const ConstNamedModuleApplyFunction& function,
      const std::string& name_prefix = std::string()) const;

  void apply(const ModulePointerApplyFunction& function) const;

  void apply(
      const NamedModulePointerApplyFunction& function,
      const std::string& name_prefix = std::string()) const;

  std::vector<Tensor> parameters(bool recurse = true) const;

  OrderedDict<std::string, Tensor> named_parameters(bool recurse = true) const;

  std::vector<Tensor> buffers(bool recurse = true) const;

  OrderedDict<std::string, Tensor> named_buffers(bool recurse = true) const;

  std::vector<std::shared_ptr<Module>> modules(bool include_self = true) const;

  OrderedDict<std::string, std::shared_ptr<Module>> named_modules(
      const std::string& name_prefix = std::string(),
      bool include_self = true) const;

  std::vector<std::shared_ptr<Module>> children() const;

  OrderedDict<std::string, std::shared_ptr<Module>> named_children() const;

  virtual void train(bool on = true);

  void eval();

  virtual bool is_training() const noexcept;

  virtual void to(
      torch::Device device,
      torch::Dtype dtype,
      bool non_blocking = false);

  virtual void to(torch::Dtype dtype, bool non_blocking = false);

  virtual void to(torch::Device device, bool non_blocking = false);

  virtual void zero_grad();

  template <typename ModuleType>
  typename ModuleType::ContainedType* as() noexcept;

  template <typename ModuleType>
  const typename ModuleType::ContainedType* as() const noexcept;

  template <
      typename ModuleType,
      typename = torch::detail::disable_if_module_holder_t<ModuleType>>
  ModuleType* as() noexcept;

  template <
      typename ModuleType,
      typename = torch::detail::disable_if_module_holder_t<ModuleType>>
  const ModuleType* as() const noexcept;

  virtual void save(serialize::OutputArchive& archive) const;

  virtual void load(serialize::InputArchive& archive);

  virtual void pretty_print(std::ostream& stream) const;

 protected:
  Tensor& register_parameter(
      std::string name,
      Tensor tensor,
      bool requires_grad = true);

  Tensor& register_buffer(std::string name, Tensor tensor);

  template <typename ModuleType>
  std::shared_ptr<ModuleType> register_module(
      std::string name,
      std::shared_ptr<ModuleType> module);

  template <typename ModuleType>
  std::shared_ptr<ModuleType> register_module(
      std::string name,
      ModuleHolder<ModuleType> module_holder);

 private:
  // Friend classes.

  template <typename Derived>
  friend class Cloneable;

  TORCH_API friend std::ostream& operator<<(
      std::ostream& stream,
      const nn::Module& module);

  // Private methods.

  virtual void clone_(Module& other, const optional<Device>& device);

  template <typename... Ts>
  void to_impl(Ts&&... ts);

  void pretty_print_recursive(
      std::ostream& stream,
      const std::string& indentation) const;

  void apply_to_submodules(
      const NamedModulePointerApplyFunction& function,
      const std::string& name_prefix = std::string()) const;

  std::shared_ptr<Module> shared_from_this_checked() const;

  OrderedDict<std::string, Tensor> parameters_;

  OrderedDict<std::string, Tensor> buffers_;

  OrderedDict<std::string, std::shared_ptr<Module>> children_;

  mutable optional<std::string> name_;

  bool is_training_{true};
};

TORCH_API serialize::OutputArchive& operator<<(
    serialize::OutputArchive& archive,
    const std::shared_ptr<nn::Module>& module);

TORCH_API serialize::InputArchive& operator>>(
    serialize::InputArchive& archive,
    const std::shared_ptr<nn::Module>& module);

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ nn::Module ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

template <typename ModuleType>
typename ModuleType::ContainedType* Module::as() noexcept {
  // Use the contained type of the `ModuleHolder`, e.g. `LinearImpl` for
  // `Linear`, since `LinearImpl` inherits `nn::Module`.
  return as<typename ModuleType::ContainedType>();
}

template <typename ModuleType>
const typename ModuleType::ContainedType* Module::as() const noexcept {
  // Use the contained type of the `ModuleHolder`, e.g. `LinearImpl` for
  // `Linear`, since `LinearImpl` inherits `nn::Module`.
  return as<typename ModuleType::ContainedType>();
}

template <typename ModuleType, typename>
ModuleType* Module::as() noexcept {
  return dynamic_cast<ModuleType*>(this);
}

template <typename ModuleType, typename>
const ModuleType* Module::as() const noexcept {
  return dynamic_cast<const ModuleType*>(this);
}

template <typename ModuleType>
std::shared_ptr<ModuleType> Module::register_module(
    std::string name,
    std::shared_ptr<ModuleType> module) {
  AT_CHECK(!name.empty(), "Submodule name must not be empty");
  AT_CHECK(
      name.find('.') == std::string::npos,
      "Submodule name must not contain a dot (got '",
      name,
      "')");
  auto& base_module = children_.insert(std::move(name), std::move(module));
  return std::dynamic_pointer_cast<ModuleType>(base_module);
}

template <typename ModuleType>
std::shared_ptr<ModuleType> Module::register_module(
    std::string name,
    ModuleHolder<ModuleType> module_holder) {
  return register_module(std::move(name), module_holder.ptr());
}

template <typename... Ts>
void Module::to_impl(Ts&&... ts) {
  // First call `to()` on every child module.
  for (auto& child : children_) {
    child.value()->to(ts...);
  }
  // Then move every parameter to the new dtype/device.
  for (auto& parameter : parameters_) {
    parameter->set_data(autograd::Variable(*parameter).data().to(ts...));
  }
  // Then move every buffer to the new dtype/device.
  for (auto& buffer : buffers_) {
    buffer->set_data(autograd::Variable(*buffer).data().to(ts...));
  }
}

} // namespace nn
} // namespace torch