Mila
Deep Neural Network Library
Loading...
Searching...
No Matches
Mila::Dnn::Activation< TDeviceType, TPrecision, TFn > Class Template Referenceexport

Unified elementwise activation component. More...

Inheritance diagram for Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >:
Mila::Dnn::Component< TDeviceType, TPrecision >

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 >
IExecutionContextgetExecutionContext () 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 &param_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 &parameter_name, const Tensor< TParameterPrecision, TMemoryResource > &parameter) 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 > &parameter, 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.

Detailed Description

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
class Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >

Unified elementwise activation component.

Template Parameters
TDeviceTypeCompile-time device (Cpu or Cuda).
TPrecisionActivation/compute precision.
TFnElementwise 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.

Member Function Documentation

◆ getDeviceId()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
DeviceId Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::getDeviceId ( ) const
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 >.

◆ getGradients()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
std::vector< ITensor * > Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::getGradients ( ) const
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.

Returns
Vector of gradient pointers; empty when built for inference.
Exceptions
std::runtime_errorif called before the component has been built.

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ getMemoryStats()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
MemoryStats Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::getMemoryStats ( ) const
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.

Returns
MemoryStats reflecting current allocations.

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ getParameters()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
std::vector< ITensor * > Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::getParameters ( ) const
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 >.

◆ getType()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
const ComponentType Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::getType ( ) const
inlineoverridevirtual

Get the component type identifier.

Used for serialization and runtime type identification.

Returns
Component type enum value.

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ onBuilding()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
void Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::onBuilding ( const BuildContext & build_context)
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.

// Example -- Linear component:
shape_t out_shape =
{
config.batchSize(),
config.allocationSeqLen(), // 1 for Inference, T for Training
config_.getOutputFeatures()
};
output_ = std::make_unique<TensorType>( device, out_shape,
this->getName() + ".output" );
TensorShape shape_t
Row-major shape descriptor for tensor dimensional sizes.
Definition Tensor.Types.ixx:173
const std::string getName() const
Get the component's name identifier.
Definition Component.ixx:533

The default implementation forwards to the legacy onBuilding( const shape_t& ) overload for backwards compatibility. New components should override this overload directly.

Note
Do not call build() or onBuilding() from within this hook.
Implementations should either succeed fully or leave no partial state, as a failed build() may be retried.

Takes the build-time configuration; use its allocationSeqLen() to obtain the correct output buffer sequence dimension.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ onExecutionContextSet()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
void Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::onExecutionContextSet ( )
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:

  • Composite components: Create and configure child components.
  • Device resource allocation: Query device capabilities.

Default implementation does nothing.

Exceptions
Anyexception thrown will cause setExecutionContext() to fail and restore the component to a "context not set" state.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ onTrainingModeChanging()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
void Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::onTrainingModeChanging ( TrainingMode training_mode)
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 >.

◆ parameterCount()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
dim_t Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::parameterCount ( ) const
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 >.

◆ save_()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
void Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::save_ ( ModelArchive & archive,
SerializationMode mode ) const
inlineoverridevirtual

◆ synchronize()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
void Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::synchronize ( )
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.
        &zwj;/

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();
    }

====================================================================

Synchronization

    /**
       @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 >.

◆ toString()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TFn = ActivationType::Gelu>
std::string Mila::Dnn::Activation< TDeviceType, TPrecision, TFn >::toString ( ) const
inlineoverridevirtual

Produce a short, human-readable description of the component.

Implementations should keep output concise and avoid throwing.

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.


The documentation for this class was generated from the following file: