#pragma once
|
|
#include <torch/detail/static.h>
|
#include <torch/nn/module.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 <memory>
|
#include <type_traits>
|
#include <typeinfo>
|
#include <utility>
|
#include <vector>
|
|
namespace torch {
|
namespace nn {
|
|
/// Stores a type erased `Module`.
|
///
|
/// The PyTorch C++ API does not impose an interface on the signature of
|
/// `forward()` in `Module` subclasses. This gives you complete freedom to
|
/// design your `forward()` methods to your liking. However, this also means
|
/// there is no unified base type you could store in order to call `forward()`
|
/// polymorphically for any module. This is where the `AnyModule` comes in.
|
/// Instead of inheritance, it relies on type erasure for polymorphism.
|
///
|
/// An `AnyModule` can store any `nn::Module` subclass that provides a
|
/// `forward()` method. This `forward()` may accept any types and return any
|
/// type. Once stored in an `AnyModule`, you can invoke the underlying module's
|
/// `forward()` by calling `AnyModule::forward()` with the arguments you would
|
/// supply to the stored module (though see one important limitation below).
|
/// Example:
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// struct GenericTrainer {
|
/// torch::nn::AnyModule module;
|
///
|
/// void train(torch::Tensor input) {
|
/// module.forward(input);
|
/// }
|
/// };
|
///
|
/// GenericTrainer trainer1{torch::nn::Linear(3, 4)};
|
/// GenericTrainer trainer2{torch::nn::Conv2d(3, 4, 2)};
|
/// \endrst
|
///
|
/// As `AnyModule` erases the static type of the stored module (and its
|
/// `forward()` method) to achieve polymorphism, type checking of arguments is
|
/// moved to runtime. That is, passing an argument with an incorrect type to an
|
/// `AnyModule` will compile, but throw an exception at runtime:
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// torch::nn::AnyModule module(torch::nn::Linear(3, 4));
|
/// // Linear takes a tensor as input, but we are passing an integer.
|
/// // This will compile, but throw a `torch::Error` exception at runtime.
|
/// module.forward(123);
|
/// \endrst
|
///
|
/// \rst
|
/// .. attention::
|
/// One noteworthy limitation of `AnyModule` is that its `forward()` method
|
/// does not support implicit conversion of argument types. For example, if
|
/// the stored module's `forward()` method accepts a `float` and you call
|
/// `any_module.forward(3.4)` (where `3.4` is a `double`), this will throw
|
/// an exception.
|
/// \endrst
|
///
|
/// The return type of the `AnyModule`'s `forward()` method is controlled via
|
/// the first template argument to `AnyModule::forward()`. It defaults to
|
/// `torch::Tensor`. To change it, you can write `any_module.forward<int>()`,
|
/// for example.
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// torch::nn::AnyModule module(torch::nn::Linear(3, 4));
|
/// auto output = module.forward(torch::ones({2, 3}));
|
///
|
/// struct IntModule {
|
/// int forward(int x) { return x; }
|
/// };
|
/// torch::nn::AnyModule module(IntModule{});
|
/// int output = module.forward<int>(5);
|
/// \endrst
|
///
|
/// The only other method an `AnyModule` provides access to on the stored
|
/// module is `clone()`. However, you may acquire a handle on the module via
|
/// `.ptr()`, which returns a `shared_ptr<nn::Module>`. Further, if you know
|
/// the concrete type of the stored module, you can get a concrete handle to it
|
/// using `.get<T>()` where `T` is the concrete module type.
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// torch::nn::AnyModule module(torch::nn::Linear(3, 4));
|
/// std::shared_ptr<nn::Module> ptr = module.ptr();
|
/// torch::nn::Linear linear(module.get<torch::nn::Linear>());
|
/// \endrst
|
class AnyModule {
|
public:
|
/// A type-erased value.
|
class Value;
|
|
/// A default-constructed `AnyModule` is in an empty state.
|
AnyModule() = default;
|
|
/// Constructs an `AnyModule` from a `shared_ptr` to concrete module object.
|
template <typename ModuleType>
|
explicit AnyModule(std::shared_ptr<ModuleType> module);
|
|
/// Constructs an `AnyModule` from a concrete module object.
|
template <
|
typename ModuleType,
|
typename = torch::detail::enable_if_module_t<ModuleType>>
|
explicit AnyModule(ModuleType&& module);
|
|
/// Constructs an `AnyModule` from a module holder.
|
template <typename ModuleType>
|
explicit AnyModule(const ModuleHolder<ModuleType>& module_holder);
|
|
/// Move construction and assignment is allowed, and follows the default
|
/// behavior of move for `std::unique_ptr`.
|
AnyModule(AnyModule&&) = default;
|
AnyModule& operator=(AnyModule&&) = default;
|
|
/// Creates a shallow copy of an `AnyModule`.
|
AnyModule(const AnyModule& other);
|
AnyModule& operator=(const AnyModule& other);
|
|
/// Creates a deep copy of an `AnyModule` if it contains a module, else an
|
/// empty `AnyModule` if it is empty.
|
AnyModule clone(optional<Device> device = nullopt) const;
|
|
/// Assigns a module to the `AnyModule` (to circumvent the explicit
|
/// constructor).
|
template <typename ModuleType>
|
AnyModule& operator=(std::shared_ptr<ModuleType> module);
|
|
/// Invokes `forward()` on the contained module with the given arguments, and
|
/// returns the return value as an `Value`. Use this method when chaining
|
/// `AnyModule`s in a loop.
|
template <typename... ArgumentTypes>
|
Value any_forward(ArgumentTypes&&... arguments);
|
|
/// Invokes `forward()` on the contained module with the given arguments, and
|
/// casts the returned `Value` to the supplied `ReturnType` (which defaults to
|
/// `torch::Tensor`).
|
template <typename ReturnType = torch::Tensor, typename... ArgumentTypes>
|
ReturnType forward(ArgumentTypes&&... arguments);
|
|
/// Attempts to cast the underlying module to the given module type. Throws an
|
/// exception if the types do not match.
|
template <typename T, typename = torch::detail::enable_if_module_t<T>>
|
T& get();
|
|
/// Attempts to cast the underlying module to the given module type. Throws an
|
/// exception if the types do not match.
|
template <typename T, typename = torch::detail::enable_if_module_t<T>>
|
const T& get() const;
|
|
/// Returns the contained module in a `nn::ModuleHolder` subclass if possible
|
/// (i.e. if `T` has a constructor for the underlying module type).
|
template <typename T, typename ContainedType = typename T::ContainedType>
|
T get() const;
|
|
/// Returns a `std::shared_ptr` whose dynamic type is that of the underlying
|
/// module.
|
std::shared_ptr<Module> ptr() const;
|
|
/// Like `ptr()`, but casts the pointer to the given type.
|
template <typename T, typename = torch::detail::enable_if_module_t<T>>
|
std::shared_ptr<T> ptr() const;
|
|
/// Returns the `type_info` object of the contained value.
|
const std::type_info& type_info() const;
|
|
/// Returns true if the `AnyModule` does not contain a module.
|
bool is_empty() const noexcept;
|
|
private:
|
/// \internal
|
/// The static type of the object we store in the `AnyModule`, which erases
|
/// the actual type, but allows us to call `forward()` on the underlying
|
/// module.
|
struct Placeholder;
|
|
/// \internal
|
/// The dynamic type of the object stored in the `AnyModule`. It contains the
|
/// concrete instance to which all calls are forwarded. It is parameterized
|
/// over the concrete type of the module, and the types of the arguments the
|
/// module takes in its `forward()` method.
|
template <typename ModuleType, typename... ArgumentTypes>
|
struct Holder;
|
|
/// Creates a `unique_ptr<Placeholder>` pointing to a `Holder` of the correct
|
/// type. This method is used to deduce the arguments of the module's
|
/// `forward()` method.
|
template <
|
typename ModuleType,
|
typename Class,
|
typename ReturnType,
|
typename... ArgumentTypes>
|
std::unique_ptr<Placeholder> make_holder(
|
std::shared_ptr<ModuleType>&& module,
|
ReturnType (Class::*)(ArgumentTypes...));
|
|
/// Helper method invoked by const and non-const `get()`.
|
template <typename ModuleType, typename ReturnType, typename... ArgumentTypes>
|
ModuleType& get_(ReturnType (ModuleType::*)(ArgumentTypes...)) const;
|
|
/// Helper method invoked by const and non-const `get()`.
|
template <typename ModuleType>
|
ModuleType& get_() const;
|
|
/// The type erased module.
|
std::unique_ptr<Placeholder> content_;
|
};
|
|
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ AnyModule::Value ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
/// A simplified implementation of `std::any` which stores
|
/// a type erased object, whose concrete value can be retrieved at runtime by
|
/// checking if the `typeid()` of a requested type matches the `typeid()` of
|
/// the object stored. It is simplified in that it does not handle copying, as
|
/// we do not require it for our use cases. Moves are sufficient.
|
class AnyModule::Value {
|
public:
|
/// Move construction and assignment is allowed, and follows the default
|
/// behavior of move for `std::unique_ptr`.
|
Value(Value&&) = default;
|
Value& operator=(Value&&) = default;
|
|
/// Copy is disallowed, because we don't need it.
|
Value(const Value& other) = delete;
|
Value& operator=(const Value& other) = delete;
|
|
/// Returns a pointer to the value contained in the `Value` if the type passed
|
/// as template parameter matches the type of the value stored, and returns a
|
/// null pointer otherwise.
|
template <typename T>
|
T* try_get() {
|
static_assert(
|
!std::is_reference<T>::value,
|
"Value stores decayed types, you cannot cast it to a reference type");
|
static_assert(
|
!std::is_array<T>::value,
|
"Value stores decayed types, you must cast it to T* instead of T[]");
|
if (typeid(T).hash_code() == type_info().hash_code()) {
|
return &static_cast<Holder<T>&>(*content_).value;
|
}
|
return nullptr;
|
}
|
|
/// Returns the value contained in the `Value` if the type passed as template
|
/// parameter matches the type of the value stored, and throws an exception
|
/// otherwise.
|
template <typename T>
|
T get() {
|
if (auto* maybe_value = try_get<T>()) {
|
return *maybe_value;
|
}
|
AT_ERROR(
|
"Attempted to cast Value to ",
|
c10::demangle(typeid(T).name()),
|
", but its actual type is ",
|
c10::demangle(type_info().name()));
|
}
|
|
/// Returns the `type_info` object of the contained value.
|
const std::type_info& type_info() const noexcept {
|
return content_->type_info;
|
}
|
|
private:
|
friend class AnyModule;
|
friend struct TestValue;
|
|
/// Constructs the `Value` from value type.
|
template <
|
typename T,
|
typename =
|
torch::disable_if_t<std::is_same<autograd::Variable, T>::value>>
|
explicit Value(T&& value)
|
: content_(
|
torch::make_unique<Holder<decay_t<T>>>(std::forward<T>(value))) {}
|
|
/// Constructs the `Value` from an `autograd::Variable`, first converting it
|
/// to a `torch::Tensor`.
|
explicit Value(autograd::Variable variable)
|
: Value(Tensor(std::move(variable))) {}
|
|
/// \internal
|
/// The static type of the object we store in the `Value`, which erases the
|
/// actual object's type, allowing us only to check the `type_info` of the
|
/// type stored in the dynamic type.
|
struct Placeholder {
|
explicit Placeholder(const std::type_info& type_info_) noexcept
|
: type_info(type_info_) {}
|
virtual ~Placeholder() = default;
|
const std::type_info& type_info;
|
};
|
|
/// \internal
|
/// The dynamic type of the object we store in the `Value`, which hides the
|
/// actual object we have erased in this `Value`.
|
template <typename T>
|
struct Holder : public Placeholder {
|
/// A template because T&& would not be universal reference here.
|
template <typename U>
|
explicit Holder(U&& value_) noexcept
|
: Placeholder(typeid(T)), value(std::forward<U>(value_)) {}
|
T value;
|
};
|
|
/// The type erased object.
|
std::unique_ptr<Placeholder> content_;
|
};
|
|
// ~~~~~~~~~~~~~~~~~~~~~~~~~~ AnyModule::Placeholder ~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
struct AnyModule::Placeholder : public AnyModule::Value::Placeholder {
|
using AnyModule::Value::Placeholder::Placeholder;
|
|
/// The "erased" `forward()` method.
|
virtual Value forward(std::vector<Value>&& arguments) = 0;
|
|
/// Returns std::shared_ptr<Module> pointing to the erased module.
|
virtual std::shared_ptr<Module> ptr() = 0;
|
|
/// Returns a `Placeholder` with a shallow copy of this `AnyModule`.
|
virtual std::unique_ptr<Placeholder> copy() const = 0;
|
|
/// Returns a `Placeholder` with a deep copy of this `AnyModule`.
|
virtual std::unique_ptr<Placeholder> clone(optional<Device> device) const = 0;
|
};
|
|
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ AnyModule::Holder ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
template <typename ModuleType, typename... ArgumentTypes>
|
struct AnyModule::Holder : public AnyModule::Placeholder {
|
/// \internal
|
struct CheckedGetter {
|
template <typename T>
|
decay_t<T>&& operator()(size_t index) {
|
AT_ASSERT(index < arguments_.size());
|
auto& value = arguments_[index];
|
if (auto* maybe_value = value.template try_get<decay_t<T>>()) {
|
return std::move(*maybe_value);
|
}
|
AT_ERROR(
|
"Expected argument #",
|
index,
|
" to be of type ",
|
c10::demangle(typeid(T).name()),
|
", but received value of type ",
|
c10::demangle(value.type_info().name()));
|
}
|
std::vector<Value>& arguments_;
|
};
|
|
/// \internal
|
struct InvokeForward {
|
template <typename... Ts>
|
Value operator()(Ts&&... ts) {
|
return Value(module_->forward(std::forward<Ts>(ts)...));
|
}
|
std::shared_ptr<ModuleType>& module_;
|
};
|
|
/// Constructs the `Holder` from a concrete module.
|
explicit Holder(std::shared_ptr<ModuleType>&& module_)
|
: Placeholder(typeid(ModuleType)), module(std::move(module_)) {}
|
|
/// Calls `forward()` on the underlying module, casting each `Value` in the
|
/// argument vector to a concrete value.
|
Value forward(std::vector<Value>&& arguments) override {
|
TORCH_CHECK(
|
arguments.size() == sizeof...(ArgumentTypes),
|
c10::demangle(type_info.name()),
|
"'s forward() method expects ",
|
sizeof...(ArgumentTypes),
|
" arguments, but received ",
|
arguments.size());
|
// FYI: During invocation of a module's `forward()` method, the values live
|
// in the `arguments` vector inside this function.
|
return torch::unpack<Value, ArgumentTypes...>(
|
InvokeForward{module}, CheckedGetter{arguments});
|
}
|
|
std::shared_ptr<Module> ptr() override {
|
return module;
|
}
|
|
std::unique_ptr<Placeholder> copy() const override {
|
return torch::make_unique<Holder>(*this);
|
}
|
|
std::unique_ptr<Placeholder> clone(optional<Device> device) const override {
|
return torch::make_unique<Holder>(
|
std::dynamic_pointer_cast<ModuleType>(module->clone(device)));
|
}
|
|
/// The actual concrete module instance.
|
std::shared_ptr<ModuleType> module;
|
};
|
|
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ AnyModule ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
template <typename ModuleType>
|
AnyModule::AnyModule(std::shared_ptr<ModuleType> module)
|
: content_(make_holder(
|
std::move(module),
|
&std::remove_reference<ModuleType>::type::forward)) {
|
// `AnyModule` can only store an `nn::Module` subclass object that provides
|
// a `forward()` method that has a non-templatized return type.
|
// (e.g. `AnyModule` cannot store `nn::Sequential`, because `nn::Sequential`'s
|
// `forward()` method has a templatized return type.)
|
static_assert(
|
torch::detail::is_module<ModuleType>::value,
|
"Can only store object derived from nn::Module into AnyModule");
|
static_assert(
|
torch::detail::has_forward<ModuleType>::value,
|
"Can only store module with a forward() method that has a non-templatized"
|
" argument type and return type into AnyModule (e.g. we cannot store nn::Sequential"
|
"into AnyModule, because its forward() method's argument type and return type are templatized."
|
" If you need to use nn::Sequentials inside each other you can subclass "
|
"nn::Sequential and write a non-templatized forward function for it. You can checkout "
|
"https://github.com/pytorch/vision/blob/2f46070f3cb1ea894d82578f3dc5677f82f34958/torchvision/csrc/models/mnasnet.cpp#L59 "
|
"for an example on how to do this.).");
|
}
|
|
template <typename ModuleType, typename>
|
AnyModule::AnyModule(ModuleType&& module)
|
: AnyModule(
|
std::make_shared<ModuleType>(std::forward<ModuleType>(module))) {}
|
|
template <typename ModuleType>
|
AnyModule::AnyModule(const ModuleHolder<ModuleType>& module_holder)
|
: AnyModule(module_holder.ptr()) {}
|
|
inline AnyModule::AnyModule(const AnyModule& other)
|
: content_(other.content_ ? other.content_->copy() : nullptr) {}
|
|
inline AnyModule& AnyModule::operator=(const AnyModule& other) {
|
if (this != &other) {
|
content_ = other.content_ ? other.content_->copy() : nullptr;
|
}
|
return *this;
|
}
|
|
inline AnyModule AnyModule::clone(optional<Device> device) const {
|
AnyModule clone;
|
clone.content_ = content_ ? content_->clone(device) : nullptr;
|
return clone;
|
}
|
|
template <typename ModuleType>
|
AnyModule& AnyModule::operator=(std::shared_ptr<ModuleType> module) {
|
return (*this = AnyModule(std::move(module)));
|
}
|
|
template <typename... ArgumentTypes>
|
AnyModule::Value AnyModule::any_forward(ArgumentTypes&&... arguments) {
|
TORCH_CHECK(!is_empty(), "Cannot call forward() on an empty AnyModule");
|
std::vector<Value> values;
|
values.reserve(sizeof...(ArgumentTypes));
|
torch::apply(
|
[&values](Value&& value) { values.push_back(std::move(value)); },
|
Value(std::forward<ArgumentTypes>(arguments))...);
|
return content_->forward(std::move(values));
|
}
|
|
template <typename ReturnType, typename... ArgumentTypes>
|
ReturnType AnyModule::forward(ArgumentTypes&&... arguments) {
|
return any_forward(std::forward<ArgumentTypes>(arguments)...)
|
.template get<ReturnType>();
|
}
|
|
template <typename T, typename>
|
T& AnyModule::get() {
|
TORCH_CHECK(!is_empty(), "Cannot call get() on an empty AnyModule");
|
return get_<T>();
|
}
|
|
template <typename T, typename>
|
const T& AnyModule::get() const {
|
TORCH_CHECK(!is_empty(), "Cannot call get() on an empty AnyModule");
|
return get_<T>();
|
}
|
|
template <typename T, typename ContainedType>
|
T AnyModule::get() const {
|
return T(ptr<ContainedType>());
|
}
|
|
inline std::shared_ptr<Module> AnyModule::ptr() const {
|
TORCH_CHECK(!is_empty(), "Cannot call ptr() on an empty AnyModule");
|
return content_->ptr();
|
}
|
|
template <typename T, typename>
|
std::shared_ptr<T> AnyModule::ptr() const {
|
TORCH_CHECK(!is_empty(), "Cannot call ptr() on an empty AnyModule");
|
// Call get() but discard the value, just to do the type checking.
|
get_<T>();
|
return std::dynamic_pointer_cast<T>(ptr());
|
}
|
|
inline const std::type_info& AnyModule::type_info() const {
|
TORCH_CHECK(!is_empty(), "Cannot call type_info() on an empty AnyModule");
|
return content_->type_info;
|
}
|
|
inline bool AnyModule::is_empty() const noexcept {
|
return content_ == nullptr;
|
}
|
|
// Private Methods
|
|
template <
|
typename ModuleType,
|
typename Class,
|
typename ReturnType,
|
typename... ArgumentTypes>
|
std::unique_ptr<AnyModule::Placeholder> AnyModule::make_holder(
|
std::shared_ptr<ModuleType>&& module,
|
ReturnType (Class::*)(ArgumentTypes...)) {
|
static_assert(
|
torch::detail::check_not_lvalue_references<ArgumentTypes...>(),
|
"Modules stored inside AnyModule must not take references. "
|
"Use pointers instead.");
|
static_assert(
|
!std::is_void<ReturnType>::value,
|
"AnyModule cannot store modules that return void "
|
"(you can return a dummy value).");
|
return torch::make_unique<Holder<decay_t<ModuleType>, ArgumentTypes...>>(
|
std::move(module));
|
}
|
|
template <typename ModuleType>
|
ModuleType& AnyModule::get_() const {
|
using M = typename std::remove_reference<ModuleType>::type;
|
static_assert(
|
torch::detail::has_forward<M>::value,
|
"Can only call AnyModule::get<T> with a type T that has a forward method");
|
return get_(&M::forward);
|
}
|
|
template <typename ModuleType, typename ReturnType, typename... ArgumentTypes>
|
ModuleType& AnyModule::get_(
|
ReturnType (ModuleType::*)(ArgumentTypes...)) const {
|
if (typeid(ModuleType).hash_code() == type_info().hash_code()) {
|
return *static_cast<Holder<ModuleType, ArgumentTypes...>&>(*content_)
|
.module;
|
}
|
AT_ERROR(
|
"Attempted to cast module of type ",
|
c10::demangle(type_info().name()),
|
" to type ",
|
c10::demangle(typeid(ModuleType).name()));
|
}
|
|
} // namespace nn
|
} // namespace torch
|
