flambe.nn.module

Module Contents

class flambe.nn.module.Module[source]

Bases: flambe.compile.Component, torch.nn.Module

Base Flambé Module inteface.

Provides the exact same interface as Pytorch’s nn.Module, but extends it with a useful set of methods to access and clip parameters, as well as gradients.

This abstraction allows users to convert their modules with a single line change, by importing from Flambé instead. Just like every Pytorch module, a forward method should be implemented.

named_trainable_params :Iterator[Tuple[str, nn.Parameter]][source]

Get all the named parameters with requires_grad=True.

Returns:Iterator over the parameters and their name.
Return type:Iterator[Tuple[str, nn.Parameter]]
trainable_params :Iterator[nn.Parameter][source]

Get all the parameters with requires_grad=True.

Returns:Iterator over the parameters
Return type:Iterator[nn.Parameter]
gradient_norm :float[source]

Compute the average gradient norm.

Returns:The current average gradient norm
Return type:float
parameter_norm :float[source]

Compute the average parameter norm.

Returns:The current average parameter norm
Return type:float
num_parameters(self, trainable=False)[source]

Gets the number of parameters in the model.

Returns:number of model params
Return type:int
clip_params(self, threshold: float)[source]

Clip the parameters to the given range.

Parameters:float – Values are clipped between -threshold, threshold
clip_gradient_norm(self, threshold: float)[source]

Clip the norm of the gradient by the given value.

Parameters:float – Threshold to clip at