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
-