#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 {
|
|
/// The base class for all modules in PyTorch.
|
///
|
/// \rst
|
/// .. note::
|
/// The design and implementation of this class is largely based on the Python
|
/// API. You may want to consult the python documentation for
|
/// :py:class:`pytorch:torch.nn.Module` for further clarification on certain
|
/// methods or behavior.
|
/// \endrst
|
///
|
/// A `Module` is an abstraction over the implementation of some function or
|
/// algorithm, possibly associated with some persistent data. A `Module` may
|
/// contain further `Module`s ("submodules"), each with their own
|
/// implementation, persistent data and further submodules. `Module`s can thus
|
/// be said to form a recursive tree structure. A `Module` is registered as a
|
/// submodule to another `Module` by calling `register_module()`, typically from
|
/// within a parent module's constructor.
|
///
|
/// A distinction is made between three kinds of persistent data that may be
|
/// associated with a `Module`:
|
///
|
/// 1. *Parameters*: tensors that record gradients, typically weights updated
|
/// during the backward step (e.g. the `weight` of a `Linear` module),
|
/// 2. *Buffers*: tensors that do not record gradients, typically updated during
|
/// the forward step, such as running statistics (e.g. `mean` and `variance`
|
/// in the `BatchNorm` module),
|
/// 3. Any additional state, not necessarily tensors, required for the
|
/// implementation or configuration of a `Module`.
|
///
|
/// The first two kinds of state are special in that they may be registered
|
/// with the `Module` system to allow convenient access and batch configuration.
|
/// For example, registered parameters in any `Module` may be iterated over via
|
/// the `parameters()` accessor. Further, changing the data type of a `Module`'s
|
/// registered parameters can be done conveniently via `Module::to()`, e.g.
|
/// `module->to(torch::kCUDA)` to move all parameters to GPU memory. Lastly,
|
/// registered parameters and buffers are handled specially during a `clone()`
|
/// operation, which performs a deepcopy of a cloneable `Module` hierarchy.
|
///
|
/// Parameters are registered with a `Module` via `register_parameter`. Buffers
|
/// are registered separately via `register_buffer`. These methods are part of
|
/// the public API of `Module` and are typically invoked from within a
|
/// concrete `Module`s constructor.
|
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>&)>;
|
|
/// Tells the base `Module` about the name of the submodule.
|
explicit Module(std::string name);
|
|
/// Constructs the module without immediate knowledge of the submodule's name.
|
/// The name of the submodule is inferred via RTTI (if possible) the first
|
/// time `.name()` is invoked.
|
Module();
|
|
virtual ~Module() = default;
|
|
/// Returns the name of the `Module`.
|
///
|
/// A `Module` has an associated `name`, which is a string representation of
|
/// the kind of concrete `Module` it represents, such as `"Linear"` for the
|
/// `Linear` module. Under most circumstances, this name is automatically
|
/// inferred via runtime type information (RTTI). In the unusual circumstance
|
/// that you have this feature disabled, you may want to manually name your
|
/// `Module`s by passing the string name to the `Module` base class'
|
/// constructor.
|
const std::string& name() const noexcept;
|
|
/// Performs a recursive deep copy of the module and all its registered
|
/// parameters, buffers and submodules.
|
///
|
/// Optionally, this method sets the current device
|
/// to the one supplied before cloning. If no device is given, each
|
/// parameter and buffer will be moved to the device of its source.
|
///
|
/// \rst
|
/// .. attention::
|
/// Attempting to call the `clone()` method inherited from the base `Module`
|
/// class (the one documented here) will fail. To inherit an actual
|
/// implementation of `clone()`, you must subclass `Cloneable`. `Cloneable`
|
/// is templatized on the concrete module type, and can thus properly copy a
|
/// `Module`. This method is provided on the base class' API solely for an
|
/// easier-to-use polymorphic interface.
|
/// \endrst
|
virtual std::shared_ptr<Module> clone(
|
const optional<Device>& device = nullopt) const;
|
|
/// Applies the `function` to the `Module` and recursively to every submodule.
|
/// The function must accept a `Module&`.
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
/// MyModule module;
|
/// module->apply([](nn::Module& module) {
|
/// std::cout << module.name() << std::endl;
|
/// });
|
/// \endrst
|
void apply(const ModuleApplyFunction& function);
|
|
/// Applies the `function` to the `Module` and recursively to every submodule.
|
/// The function must accept a `const Module&`.
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
/// MyModule module;
|
/// module->apply([](const nn::Module& module) {
|
/// std::cout << module.name() << std::endl;
|
/// });
|
/// \endrst
|
void apply(const ConstModuleApplyFunction& function) const;
|
|
/// Applies the `function` to the `Module` and recursively to every submodule.
|
/// The function must accept a `const std::string&` for the key of the module,
|
/// and a `Module&`. The key of the module itself is the empty string. If
|
/// `name_prefix` is given, it is prepended to every key as
|
/// `<name_prefix>.<key>` (and just `name_prefix` for the module itself).
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
/// MyModule module;
|
/// module->apply([](const std::string& key, nn::Module& module) {
|
/// std::cout << key << ": " << module.name() << std::endl;
|
/// });
|
/// \endrst
|
void apply(
|
const NamedModuleApplyFunction& function,
|
const std::string& name_prefix = std::string());
|
|
/// Applies the `function` to the `Module` and recursively to every submodule.
|
/// The function must accept a `const std::string&` for the key of the module,
|
/// and a `const Module&`. The key of the module itself is the empty string.
|
/// If `name_prefix` is given, it is prepended to every key as
|
/// `<name_prefix>.<key>` (and just `name_prefix` for the module itself).
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
/// MyModule module;
|
/// module->apply([](const std::string& key, const nn::Module& module) {
|
/// std::cout << key << ": " << module.name() << std::endl;
|
/// });
|
/// \endrst
|
void apply(
|
const ConstNamedModuleApplyFunction& function,
|
const std::string& name_prefix = std::string()) const;
|
|
/// Applies the `function` to the `Module` and recursively to every submodule.
|
/// The function must accept a `const std::shared_ptr<Module>&`.
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
/// MyModule module;
|
/// module->apply([](const std::shared_ptr<nn::Module>& module) {
|
/// std::cout << module->name() << std::endl;
|
/// });
|
/// \endrst
|
void apply(const ModulePointerApplyFunction& function) const;
|
|
/// Applies the `function` to the `Module` and recursively to every submodule.
|
/// The function must accept a `const std::string&` for the key of the module,
|
/// and a `const std::shared_ptr<Module>&`. The key of the module itself is
|
/// the empty string. If `name_prefix` is given, it is prepended to every key
|
/// as
|
/// `<name_prefix>.<key>` (and just `name_prefix` for the module itself).
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
/// MyModule module;
|
/// module->apply([](const std::string& key,
|
/// const std::shared_ptr<nn::Module>& module) {
|
/// std::cout << key << ": " << module->name() << std::endl;
|
/// });
|
/// \endrst
|
void apply(
|
const NamedModulePointerApplyFunction& function,
|
const std::string& name_prefix = std::string()) const;
|
|
/// Returns the parameters of this `Module` and if `recurse` is true, also
|
/// recursively of every submodule.
|
std::vector<Tensor> parameters(bool recurse = true) const;
|
|
/// Returns an `OrderedDict` with the parameters of this `Module` along with
|
/// their keys, and if `recurse` is true also recursively of every submodule.
|
OrderedDict<std::string, Tensor> named_parameters(bool recurse = true) const;
|
|
/// Returns the buffers of this `Module` and if `recurse` is true, also
|
/// recursively of every submodule.
|
std::vector<Tensor> buffers(bool recurse = true) const;
|
|
/// Returns an `OrderedDict` with the buffers of this `Module` along with
|
/// their keys, and if `recurse` is true also recursively of every submodule.
|
OrderedDict<std::string, Tensor> named_buffers(bool recurse = true) const;
|
|
/// Returns the submodules of this `Module` (the entire submodule hierarchy)
|
/// and if `include_self` is true, also inserts a `shared_ptr` to this module
|
/// in the first position.
|
///
|
/// \rst
|
/// .. warning::
|
/// Only pass `include_self` as `true` if this `Module` is stored in a
|
/// `shared_ptr`! Otherwise an exception will be thrown. You may still call
|
/// this method with `include_self` set to false if your `Module` is not
|
/// stored in a `shared_ptr`.
|
/// \endrst
|
std::vector<std::shared_ptr<Module>> modules(bool include_self = true) const;
|
|
/// Returns an `OrderedDict` of the submodules of this `Module` (the entire
|
/// submodule hierarchy) and their keys, and if `include_self` is true, also
|
/// inserts a `shared_ptr` to this module in the first position. If
|
/// `name_prefix` is given, it is prepended to every key as
|
/// `<name_prefix>.<key>` (and just `name_prefix` for the module itself).
|
///
|
/// \rst
|
/// .. warning::
|
/// Only pass `include_self` as `true` if this `Module` is stored in a
|
/// `shared_ptr`! Otherwise an exception will be thrown. You may still call
|
/// this method with `include_self` set to false if your `Module` is not
|
/// stored in a `shared_ptr`.
|
/// \endrst
|
OrderedDict<std::string, std::shared_ptr<Module>> named_modules(
|
const std::string& name_prefix = std::string(),
|
bool include_self = true) const;
|
|
/// Returns the direct submodules of this `Module`.
|
std::vector<std::shared_ptr<Module>> children() const;
|
|
/// Returns an `OrderedDict` of the direct submodules of this `Module` and
|
/// their keys.
|
OrderedDict<std::string, std::shared_ptr<Module>> named_children() const;
|
|
/// Enables "training" mode.
|
virtual void train(bool on = true);
|
|
/// Calls train(false) to enable "eval" mode.
|
/// Do not override this method, override `train()` instead.
|
void eval();
|
|
/// True if the module is in training mode.
|
///
|
/// Every `Module` has a boolean associated with it that determines whether
|
/// the `Module` is currently in *training* mode (set via `.train()`) or in
|
/// *evaluation* (inference) mode (set via `.eval()`). This property is
|
/// exposed via `is_training()`, and may be used by the implementation of a
|
/// concrete module to modify its runtime behavior. See the `BatchNorm` or
|
/// `Dropout` modules for examples of `Module`s that use different code paths
|
/// depending on this property.
|
virtual bool is_training() const noexcept;
|
|
/// Recursively casts all parameters to the given `dtype` and `device`.
|
///
|
/// If `non_blocking` is true and the source is in pinned memory and
|
/// destination is on the GPU or vice versa, the copy is performed
|
/// asynchronously with respect to the host. Otherwise, the argument has no
|
/// effect.
|
virtual void to(
|
torch::Device device,
|
torch::Dtype dtype,
|
bool non_blocking = false);
|
|
/// Recursively casts all parameters to the given dtype.
|
///
|
/// If `non_blocking` is true and the source is in pinned memory and
|
/// destination is on the GPU or vice versa, the copy is performed
|
/// asynchronously with respect to the host. Otherwise, the argument has no
|
/// effect.
|
virtual void to(torch::Dtype dtype, bool non_blocking = false);
|
|
/// Recursively moves all parameters to the given device.
|
///
|
/// If `non_blocking` is true and the source is in pinned memory and
|
/// destination is on the GPU or vice versa, the copy is performed
|
/// asynchronously with respect to the host. Otherwise, the argument has no
|
/// effect.
|
virtual void to(torch::Device device, bool non_blocking = false);
|
|
/// Recursively zeros out the `grad` value of each registered parameter.
|
virtual void zero_grad();
|
|
/// Attempts to cast this `Module` to the given `ModuleType`.
|
///
|
/// This method is useful when calling `apply()`.
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// void initialize_weights(nn::Module& module) {
|
/// torch::NoGradGuard no_grad;
|
/// if (auto* linear = module.as<nn::Linear>()) {
|
/// linear->weight.normal_(0.0, 0.02);
|
/// }
|
/// }
|
///
|
/// MyModule module;
|
/// module->apply(initialize_weights);
|
/// \endrst
|
template <typename ModuleType>
|
typename ModuleType::ContainedType* as() noexcept;
|
|
/// Attempts to cast this `Module` to the given `ModuleType`.
|
///
|
/// This method is useful when calling `apply()`.
|
/// \rst
|
/// .. code-block:: cpp
|
/// void initialize_weights(nn::Module& module) {
|
/// torch::NoGradGuard no_grad;
|
/// if (auto* linear = module.as<nn::Linear>()) {
|
/// linear->weight.normal_(0.0, 0.02);
|
/// }
|
/// }
|
///
|
/// MyModule module;
|
/// module->apply(initialize_weights);
|
/// \endrst
|
template <typename ModuleType>
|
const typename ModuleType::ContainedType* as() const noexcept;
|
|
/// Attempts to cast this `Module` to the given `ModuleType`.
|
///
|
/// This method is useful when calling `apply()`.
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// void initialize_weights(nn::Module& module) {
|
/// torch::NoGradGuard no_grad;
|
/// if (auto* linear = module.as<nn::Linear>()) {
|
/// linear->weight.normal_(0.0, 0.02);
|
/// }
|
/// }
|
///
|
/// MyModule module;
|
/// module.apply(initialize_weights);
|
/// \endrst
|
template <
|
typename ModuleType,
|
typename = torch::detail::disable_if_module_holder_t<ModuleType>>
|
ModuleType* as() noexcept;
|
|
/// Attempts to cast this `Module` to the given `ModuleType`.
|
///
|
/// This method is useful when calling `apply()`.
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// void initialize_weights(nn::Module& module) {
|
/// torch::NoGradGuard no_grad;
|
/// if (auto* linear = module.as<nn::Linear>()) {
|
/// linear->weight.normal_(0.0, 0.02);
|
/// }
|
/// }
|
///
|
/// MyModule module;
|
/// module.apply(initialize_weights);
|
/// \endrst
|
template <
|
typename ModuleType,
|
typename = torch::detail::disable_if_module_holder_t<ModuleType>>
|
const ModuleType* as() const noexcept;
|
|
/// Serializes the `Module` into the given `OutputArchive`.
|
///
|
/// If the `Module` contains unserializable submodules (e.g. `nn::Functional`),
|
/// those submodules are skipped when serializing.
|
virtual void save(serialize::OutputArchive& archive) const;
|
|
/// Deserializes the `Module` from the given `InputArchive`.
|
///
|
/// If the `Module` contains unserializable submodules (e.g. `nn::Functional`),
|
/// we don't check the existence of those submodules in the `InputArchive` when
|
/// deserializing.
|
virtual void load(serialize::InputArchive& archive);
|
|
/// Streams a pretty representation of the `Module` into the given `stream`.
|
/// By default, this representation will be the name of the module (taken from
|
/// `name()`), followed by a recursive pretty print of all of the `Module`'s
|
/// submodules.
|
///
|
/// Override this method to change the pretty print. The input
|
/// `stream` should be returned from the method, to allow easy chaining.
|
virtual void pretty_print(std::ostream& stream) const;
|
|
/// Returns whether the `Module` is serializable.
|
virtual bool is_serializable() const;
|
|
/// Registers a parameter with this `Module`.
|
///
|
/// A parameter should be any gradient-recording tensor used in the
|
/// implementation of your `Module`. Registering it makes it available to
|
/// methods such as `parameters()`, `clone()` or `to().`
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// MyModule::MyModule() {
|
/// weight_ = register_parameter("weight", torch::randn({A, B}));
|
/// }
|
/// \endrst
|
Tensor& register_parameter(
|
std::string name,
|
Tensor tensor,
|
bool requires_grad = true);
|
|
/// Registers a buffer with this `Module`.
|
///
|
/// A buffer is intended to be state in your module that does not record
|
/// gradients, such as running statistics. Registering it makes it available
|
/// to methods such as `buffers()`, `clone()` or `to().
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// MyModule::MyModule() {
|
/// mean_ = register_buffer("mean", torch::empty({num_features_}));
|
/// }
|
/// \endrst
|
Tensor& register_buffer(std::string name, Tensor tensor);
|
|
/// Registers a submodule with this `Module`.
|
///
|
/// Registering a module makes it available to methods such as `modules()`,
|
/// `clone()` or `to()`.
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// MyModule::MyModule() {
|
/// submodule_ = register_module("linear", torch::nn::Linear(3, 4));
|
/// }
|
/// \endrst
|
template <typename ModuleType>
|
std::shared_ptr<ModuleType> register_module(
|
std::string name,
|
std::shared_ptr<ModuleType> module);
|
|
/// Registers a submodule with this `Module`.
|
///
|
/// This method deals with `ModuleHolder`s.
|
///
|
/// Registering a module makes it available to methods such as `modules()`,
|
/// `clone()` or `to()`.
|
///
|
/// \rst
|
/// .. code-block:: cpp
|
///
|
/// MyModule::MyModule() {
|
/// submodule_ = register_module("linear", torch::nn::Linear(3, 4));
|
/// }
|
/// \endrst
|
template <typename ModuleType>
|
std::shared_ptr<ModuleType> register_module(
|
std::string name,
|
ModuleHolder<ModuleType> module_holder);
|
|
/// Replaces a registered submodule with this `Module`.
|
///
|
/// This takes care of the registration, if you used submodule members, you should
|
// assign the submodule as well, i.e. use as
|
/// module->submodule_ = module->replace_module("linear", torch::nn::Linear(3, 4));
|
/// It only works when a module of the name is already registered.
|
///
|
/// This is useful for replacing a module after initialization, e.g.
|
/// for finetuning.
|
template <typename ModuleType>
|
std::shared_ptr<ModuleType> replace_module(
|
const std::string& name,
|
std::shared_ptr<ModuleType> module);
|
|
/// Replaces a registered submodule with this `Module`.
|
/// This method deals with `ModuleHolder`s.
|
///
|
/// This takes care of the registration, if you used submodule members, you should
|
// assign the submodule as well, i.e. use as
|
/// module->submodule_ = module->replace_module("linear", linear_holder);
|
/// It only works when a module of the name is already registered.
|
///
|
/// This is useful for replacing a module after initialization, e.g.
|
/// for finetuning.
|
template <typename ModuleType>
|
std::shared_ptr<ModuleType> replace_module(
|
const std::string& name,
|
ModuleHolder<ModuleType> module_holder);
|
|
/// Unregisters a submodule from this `Module`. If there is no such module
|
/// with `name` an exception is thrown.
|
void unregister_module(const std::string& name);
|
|
private:
|
// Friend classes.
|
|
template <typename Derived>
|
friend class Cloneable;
|
|
/// Pretty prints the given `Module` into the `ostream`.
|
TORCH_API friend std::ostream& operator<<(
|
std::ostream& stream,
|
const nn::Module& module);
|
|
// data parallel using this method to configure gradient edges during the
|
// replicate step.
|
template <typename ModuleType>
|
friend void replicate_grad_edges(
|
const std::shared_ptr<Module>& module,
|
const std::vector<std::shared_ptr<ModuleType>>& replicas,
|
const std::vector<Device>& devices);
|
|
// Private methods.
|
|
/// Used in the implementation of `Cloneable`.
|
virtual void clone_(Module& other, const optional<Device>& device);
|
|
/// The implementation of the various `to()` methods.
|
template <typename... Ts>
|
void to_impl(Ts&&... ts);
|
|
/// Implements pretty printing the module hierarchy.
|
void pretty_print_recursive(
|
std::ostream& stream,
|
const std::string& indentation) const;
|
|
/// Applies the `function` to every submodule recursively, starting at this
|
/// `Module`'s children (thus not including the module itself).
|
void apply_to_submodules(
|
const NamedModulePointerApplyFunction& function,
|
const std::string& name_prefix = std::string()) const;
|
|
/// Returns a shared_ptr to `this` in a safe (checked) way.
|
std::shared_ptr<Module> shared_from_this_checked() const;
|
|
/// The registered parameters of this `Module`.
|
OrderedDict<std::string, Tensor> parameters_;
|
|
/// The registered buffers of this `Module`.
|
OrderedDict<std::string, Tensor> buffers_;
|
|
/// The registered (direct) submodules of this `Module`.
|
OrderedDict<std::string, std::shared_ptr<Module>> children_;
|
|
/// The module's name (e.g. "LSTM").
|
mutable optional<std::string> name_;
|
|
/// Whether the module is in training mode.
|
bool is_training_{true};
|
};
|
|
/// Serialize a `Module` pointer into an `OutputArchive`.
|
TORCH_API serialize::OutputArchive& operator<<(
|
serialize::OutputArchive& archive,
|
const std::shared_ptr<nn::Module>& module);
|
|
/// Deserializes a `Module` from an `InputArchive`.
|
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) {
|
TORCH_CHECK(!name.empty(), "Submodule name must not be empty");
|
TORCH_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 ModuleType>
|
std::shared_ptr<ModuleType> Module::replace_module(
|
const std::string& name,
|
std::shared_ptr<ModuleType> module) {
|
auto& base_module = (children_[name] = std::move(module));
|
return std::dynamic_pointer_cast<ModuleType>(base_module);
|
}
|
|
template <typename ModuleType>
|
std::shared_ptr<ModuleType> Module::replace_module(
|
const std::string& name,
|
ModuleHolder<ModuleType> module_holder) {
|
return replace_module(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).to(ts...));
|
}
|
// Then move every buffer to the new dtype/device.
|
for (auto& buffer : buffers_) {
|
buffer->set_data(autograd::Variable(*buffer).to(ts...));
|
}
|
}
|
|
} // namespace nn
|
} // namespace torch
|
