|
Mila
Deep Neural Network Library
|
Unified elementwise activation component. More...
Public Types | |
| using | ComponentBase = Component<TDeviceType, TPrecision> |
| using | MR = typename DeviceTypeTraits<TDeviceType>::memory_resource |
| using | TensorType = Tensor<TPrecision, MR> |
Public Member Functions | |
| Activation (const std::string &name, const ActivationConfig &config=ActivationConfig(TFn), std::optional< DeviceId > device_id=std::nullopt) | |
| TensorType & | backward (const TensorType &input, const TensorType &output_grad) |
| TensorType & | forward (const TensorType &input) |
| DeviceId | getDeviceId () const override |
| Get the compute device id associated with this component. | |
| std::vector< ITensor * > | getGradients () const override |
| Return non-owning pointers to parameter gradient tensors. | |
| MemoryStats | getMemoryStats () const override |
| Return the current memory allocation breakdown for this component. | |
| std::vector< ITensor * > | getParameters () const override |
| Return non-owning pointers to parameter tensors. | |
| const ComponentType | getType () const override |
| Get the component type identifier. | |
| dim_t | parameterCount () const override |
| Return number of trainable parameters. | |
| void | save_ (ModelArchive &archive, SerializationMode mode) const override |
| void | synchronize () override |
| std::string | toString () const override |
| Produce a short, human-readable description of the component. | |
| Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision > | |
| Component (const std::string &name) | |
| Construct component with required name identifier. | |
| virtual void | build (const BuildContext &context) final |
| Build the component with the provided BuildContext (canonical overload). | |
| const std::string | getName () const |
| Get the component's name identifier. | |
| virtual std::vector< std::string > | getParameterNames () const |
| List all available parameter names for this component. | |
| virtual MemoryStats | getRequiredMemory (const BuildContext &context) const |
| Report what build( context ) would allocate, without allocating it. | |
| TrainingMode | getTrainingMode () const noexcept |
| The current runtime behavioral mode of this Component. | |
| virtual bool | isBuilt () const final |
| Returns true if build() has completed successfully. | |
| virtual void | load_ (ModelArchive &archive, SerializationMode mode) |
| Restore this component's parameters from its archive scope. | |
| virtual void | loadParameter (const std::string &, const Serialization::ITensorBlob &) |
| Load a parameter from serialized tensor data. | |
| virtual void | requireSerializableParameters () const |
| Verify this component can serialize whatever parameters it owns. | |
| virtual void | saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const |
| Drive this component's tensors through one pass of a flat safetensors save. | |
| void | setTrainingMode (TrainingMode mode) |
| Set the runtime behavioral mode for this Component. | |
| virtual void | zeroGradients () |
| Clear all model-owned gradients for this component. | |
Static Public Member Functions | |
| static std::unique_ptr< Activation > | fromArchive_ (ModelArchive &archive, const std::string &component_name, IExecutionContext *exec_context) |
| static constexpr ActivationType | getActivationType () noexcept |
| The compile-time activation function realized by this component. | |
| Static Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision > | |
| static constexpr DeviceType | getDeviceType () |
| Compile-time device type for this component instance. | |
| static constexpr TensorDataType | getPrecision () noexcept |
| Compile-time tensor precision for this component instance. | |
Protected Member Functions | |
| void | onBuilding (const BuildContext &build_context) override |
| Hook invoked by build() to allocate component buffers. | |
| void | onExecutionContextSet () override |
| Lifecycle hook: Called immediately after ExecutionContext is set. | |
| void | onTrainingModeChanging (TrainingMode training_mode) override |
| Hook called before TrainingMode transitions. | |
| Protected Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision > | |
| IExecutionContext * | getExecutionContext () const |
| Get the shared execution context. | |
| bool | hasExecutionContext () const noexcept |
| Check if execution context has been set. | |
| template<TensorDataType TParameterPrecision, typename TMemoryResource> | |
| void | loadParameterFromBlob (const std::string ¶m_name, const Serialization::ITensorBlob &blob, Tensor< TParameterPrecision, TMemoryResource > &target, const shape_t &expected_shape) |
| Load a tensor blob into a parameter tensor with validation. | |
| template<TensorDataType TParameterPrecision, typename TMemoryResource> | |
| void | saveParameterToArchive (ModelArchive &archive, const std::string ¶meter_name, const Tensor< TParameterPrecision, TMemoryResource > ¶meter) const |
| Write one parameter tensor into the archive under "tensors/<name>". | |
| template<TensorDataType TParameterPrecision, typename TMemoryResource> | |
| void | saveParameterToWriter (Serialization::SafeTensorsWriter &writer, const std::string &flat_name, const Tensor< TParameterPrecision, TMemoryResource > ¶meter, Serialization::TensorSavePass pass) const |
| Drive one parameter through one pass of a flat safetensors save. | |
| void | setExecutionContext (IExecutionContext *context) |
| Set the execution context for this component. | |
Additional Inherited Members | |
| Protected Types inherited from Mila::Dnn::Component< TDeviceType, TPrecision > | |
| using | HostStagingMemoryResource |
| Host memory a device-resident parameter stages through. | |
| Protected Attributes inherited from Mila::Dnn::Component< TDeviceType, TPrecision > | |
| BuildContext | build_context_ { shape_t{ 1 }, RuntimeMode::Training } |
| The BuildContext stored at build time. | |
Unified elementwise activation component.
| TDeviceType | Compile-time device (Cpu or Cuda). |
| TPrecision | Activation/compute precision. |
| TFn | Elementwise activation function, fixed at compile time. |
Construction follows the standalone/shared pattern (mirrors Gelu): pass a DeviceId to own an ExecutionContext, or omit it and let a parent inject one.
|
inlineoverridevirtual |
Get the compute device id associated with this component.
Must return the device on which parameters and operations execute.
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverridevirtual |
Return non-owning pointers to parameter gradient tensors.
Gradient buffers are allocated only when the component is built in training mode, so a component built for inference returns an empty vector. Stateless components return empty in either mode. This is the accessor counterpart to getParameters() and does not throw on mode.
| std::runtime_error | if called before the component has been built. |
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverridevirtual |
Return the current memory allocation breakdown for this component.
Reflects allocations at the moment of the call. The returned stats naturally track the component lifecycle:
After construction – nothing; construction allocates none After build( Inference ) – parameters + T=1 state buffers After build( Training ) – parameters + T=full state buffers After setTrainingMode( Train ) – parameters + state + gradients
For CompositeComponent and Network, the returned stats are the recursive aggregate of all child components.
May be called at any time – no lifecycle preconditions.
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverridevirtual |
Return non-owning pointers to parameter tensors.
The returned tensor pointers remain valid for the lifetime of the component. Order should be canonical (weights before biases).
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverridevirtual |
Get the component type identifier.
Used for serialization and runtime type identification.
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverrideprotectedvirtual |
Hook invoked by build() to allocate component buffers.
Receives the stored BuildContext. Implementations must use config.allocationSeqLen() when sizing output buffers – this is the single call that makes Inference and Training allocate the correct buffer sizes automatically without per-component logic.
The default implementation forwards to the legacy onBuilding( const shape_t& ) overload for backwards compatibility. New components should override this overload directly.
Takes the build-time configuration; use its allocationSeqLen() to obtain the correct output buffer sequence dimension.
Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverrideprotectedvirtual |
Lifecycle hook: Called immediately after ExecutionContext is set.
Override this to perform initialization that requires a valid ExecutionContext. At the time this is called, getExecutionContext() is guaranteed to return a valid context.
Common uses:
Default implementation does nothing.
| Any | exception thrown will cause setExecutionContext() to fail and restore the component to a "context not set" state. |
Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverrideprotectedvirtual |
Hook called before TrainingMode transitions.
Called by setTrainingMode() after validation and lock acquisition, before the internal state is updated. Derived classes override to respond to the transition – e.g. zeroing gradient buffers on transition to Eval, or re-enabling dropout on transition to Training.
The default implementation is a no-op.
Takes the incoming TrainingMode.
Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverridevirtual |
Return number of trainable parameters.
For leaf components this is the element count of owned parameter tensors. CompositeComponent and Network implementations should return the recursive aggregate across all children.
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverridevirtual |
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverridevirtual |
@brief Convenience accessor -- true if currently in Eval mode.
Equivalent to getTrainingMode() == TrainingMode::Eval.
Valid for both RuntimeMode::Inference and RuntimeMode::Training
built components.
@return true if in Eval mode.
‍/
bool isEvalMode() const noexcept { return getTrainingMode() == TrainingMode::Eval; }
RuntimeMode getRuntimeMode() const noexcept
{
return build_context_.getRuntimeMode();
}
bool isInferenceMode() const noexcept
{
return build_context_.isInferenceMode();
}
bool isTrainingMode() const noexcept
{
return build_context_.isTrainingMode();
}
====================================================================
/**
@brief Wait for outstanding device work submitted by this component.
On CPU this may be a no-op. Use to ensure results are visible to
the host or to measure synchronous timings.
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverridevirtual |
Produce a short, human-readable description of the component.
Implementations should keep output concise and avoid throwing.
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.