#pragma once
|
|
#include <torch/detail/static.h>
|
#include <torch/nn/module.h>
|
#include <torch/nn/modules/container/any.h>
|
#include <torch/nn/pimpl.h>
|
#include <torch/types.h>
|
|
#include <torch/csrc/autograd/variable.h>
|
#include <torch/csrc/utils/memory.h>
|
#include <torch/csrc/utils/variadic.h>
|
|
#include <ATen/Device.h>
|
|
#include <initializer_list>
|
#include <memory>
|
#include <type_traits>
|
#include <typeinfo>
|
#include <utility>
|
#include <vector>
|
|
namespace torch {
|
namespace nn {
|
|
/// Stores a type erased `Module` with name.
|
///
|
/// The `NamedAnyModule` class and the `modules_ordered_dict(...)` function enables
|
/// the following API for constructing `nn::Sequential` with named submodules:
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// struct M : torch::nn::Module {
|
/// explicit M(int value_) : value(value_) {}
|
/// int value;
|
/// int forward() {
|
/// return value;
|
/// }
|
/// };
|
///
|
/// Sequential sequential(modules_ordered_dict({
|
/// {"m1", std::make_shared<M>(1)}, // shared pointer to `Module` is supported
|
/// {std::string("m2"), M(2)}, // `Module` is supported
|
/// {"linear1", Linear(10, 3)} // `ModuleHolder` is supported
|
/// }));
|
/// \endrst
|
///
|
/// Specifically, we design the signature of `modules_ordered_dict(...)` to be
|
/// `modules_ordered_dict(std::initializer_list<NamedAnyModule> named_modules)`, as
|
/// a result of evaluating the following possible approaches:
|
///
|
/// Approach 1:
|
/// `modules_ordered_dict(std::initializer_list<
|
/// torch::OrderedDict<std::string, ModuleType>::Item> named_modules)`
|
///
|
/// Why it doens't work:
|
/// When we pass in a braced-init list such as
|
/// `modules_ordered_dict({{"m1", M(1)}, {"m2", M(2)}})`, at the template argument
|
/// deduction step the compiler is not able to deduce the type of `ModuleType` to
|
/// the type of `M(1)` or `M(2)`, since the compiler doesn't actually look into the
|
/// braced-init list `{"m1", M(1)}` and figure out what the types of its elements are.
|
///
|
/// Approach 2:
|
/// `modules_ordered_dict(std::initializer_list<
|
/// std::pair<std::string, AnyModule> named_modules)`
|
///
|
/// Why it doens't work:
|
/// When we pass in a braced-init list such as
|
/// `modules_ordered_dict({{"m1", M(1)}, {"m2", M(2)}})`, the compiler is not able to
|
/// match `std::initializer_list<std::pair<std::string, AnyModule>>` to the nested
|
/// braced-init list `{{"m1", M(1)}, {"m2", M(2)}}`, and results in a "could not
|
/// convert" error.
|
///
|
/// Approach 3:
|
/// `modules_ordered_dict(std::initializer_list<NamedAnyModule> named_modules)`
|
///
|
/// Why it works:
|
/// When we pass in a braced-init list such as
|
/// `modules_ordered_dict({{"m1", M(1)}, {"m2", M(2)}})`, the compiler is passing the
|
/// braced-init lists {"m1", M(1)} and {"m2", M(2)} to the `NamedAnyModule`
|
/// constructors, and the constructors are able to figure out the types of the
|
/// braced-init lists' elements and match to the correct module type.
|
|
class NamedAnyModule {
|
public:
|
/// Creates a `NamedAnyModule` from a (boxed) `Module`.
|
template <typename ModuleType>
|
NamedAnyModule(std::string name, std::shared_ptr<ModuleType> module_ptr)
|
: NamedAnyModule(std::move(name), AnyModule(std::move(module_ptr))) {}
|
|
/// Creates a `NamedAnyModule` from a `Module`, moving or copying it
|
/// into a `shared_ptr` internally.
|
// NOTE: We need to use `std::remove_reference<M>::type` to get rid of
|
// any reference components for make_unique.
|
template <typename M, typename = torch::detail::enable_if_module_t<M>>
|
NamedAnyModule(std::string name, M&& module)
|
: NamedAnyModule(
|
std::move(name),
|
std::make_shared<typename std::remove_reference<M>::type>(
|
std::forward<M>(module))) {}
|
|
/// Creates a `NamedAnyModule` from a `Module` that is unwrapped from
|
/// a `ModuleHolder`.
|
template <typename M>
|
NamedAnyModule(std::string name, const ModuleHolder<M>& module_holder)
|
: NamedAnyModule(std::move(name), module_holder.ptr()) {}
|
|
/// Returns a reference to the name.
|
const std::string& name() const noexcept {
|
return name_;
|
}
|
|
/// Returns a reference to the module.
|
AnyModule& module() noexcept {
|
return module_;
|
}
|
|
private:
|
/// Creates a `NamedAnyModule` from a type-erased `AnyModule`.
|
NamedAnyModule(std::string name, AnyModule any_module)
|
: name_(std::move(name)), module_(std::move(any_module)) {}
|
|
std::string name_;
|
AnyModule module_;
|
};
|
|
TORCH_API torch::OrderedDict<std::string, AnyModule> modules_ordered_dict(
|
std::initializer_list<NamedAnyModule> named_modules);
|
|
} // namespace nn
|
} // namespace torch
|
