nip.trainers.trainer_base.attach_progress_bar

nip.trainers.trainer_base.attach_progress_bar(num_iterations_func: Callable[[Trainer], int]) → Callable[[Callable], Callable]

Decorate a Trainer method to attach a progress bar.

Decorate a method of a Trainer subclass with this decorator to have it run with a progress bar and declare the number of iterations that it runs for.

The supplied function should take a Trainer instance as input and return the number of iterations.

The intention is that once a Trainer subclass is instantiated, the number of iterations can be determined using num_iterations_func.

This decorator wraps the decorated method in an IterationContext and assigns the iteration_context keyword argument to this context. This allows the method to interact with the progress bar and access the number of iterations.

Note

The number of iterations must be calculable as soon as the Trainer subclass is instantiated, so it should not depend on any other state.

Parameters:

num_iterations_func (Callable[[Trainer], int]) – A function that takes a Trainer instance as input and returns the number of iterations which the decorated method will run for.

Returns:

decorator (Callable[[Callable], Callable]) – The decorator to apply to the method of a Trainer subclass.

Example

>>> class MyTrainer(Trainer):
...     @attach_progress_bar(lambda self: self.hyper_params.num_iterations)
...     def train(self, iteration_context: IterationContext):
...         for i in range(iteration_context.num_iterations):
...             iteration_context.step()