Create your own trainer extension

In this section, you will learn about the following things:

What is trainer Extension?

Extension is a callable object that takes a Trainer object as an argument. Adding an Extension to a Trainer using extend() method, the Extension will be called at the given timing you specified by using trigger object (See the details in 1. trigger)

A Trainer object has all information used in a training loop, e.g., models, optimizers, updaters, iterators, and datasets, etc. So you can change the settings of optimizers

Write a simple function

You can make a new Extension by writing a simple function which takes Trainer object as its argument. For example, when you want to reduce the learning rate at specified timing during training, lr_drop extension can be written as follows:

def lr_drop(trainer):
    trainer.updater.get_optimizer('main').lr *= 0.1

Then you can add this function to a Trainer object via extend() method.

trainer.extend(lr_drop, trigger=(10, 'epoch'))

It lowers the learning rate every 10 epochs by multiplying 0.1 with the current learning rate.

Write a function decorated with @make_extension

make_extension() is a decorator that adds some attributes to a given function. For example, the simple extension we created above can be written in this form:

@training.make_extension(trigger=(10, 'epoch'))
def lr_drop(trainer):
    trainer.updater.get_optimizer('main').lr *= 0.1

The difference between the above one and this is whether it has a default trigger or not. In this latter case, lr_drop() has its default trigger so that unless another trigger is specified via extend() method, the trigger specified in make_extension() is used as default. So the code below acts the same as the former example, i.e., it reduces the learning rate every 10 epochs.

trainer.extend(lr_drop)

There are several attributes you can add using the make_extension() decorator.

1. trigger

trigger is an object that takes a Trainer object as an argument and returns a boolean value. If a tuple in a form (period, unit) is given as a trigger, it will be considered as an IntervalTrigger that invokes the extension at every period unit. For example, when the given tuple is (10, 'epoch'), the extension will be fired at every 10 epochs.

trigger can also be given to the extend() method that adds an extension to a Trainer object. The priority of triggers is as follows:

See the details in the documentation of get_trigger().

2. default_name

An Extension is kept in a dictionary which is a property in a Trainer. This argument gives the name of the Extension. Users will see this name in keys of the snapshot which is a dictionary generated by serialization.

3. priority

The priority that is used to determine the order of execution of extensions in a Trainer object. There are three standard values for the priorities:

  • PRIORITY_WRITER: The priority for extensions that write some records to the observation dictionary. It includes cases that the extension directly adds values to the observation dictionary, or the extension uses the chainer.report() function to report values to the observation dictionary. Extensions which write something to reporter should go first because the other Extensions which read those values could be added.
  • PRIORITY_EDITOR: The priority for extensions that edit the observation dictionary based on already reported values. Extensions which edit some values of reported ones should go after the extensions which write values to reporter but before extensions which read the final values.
  • PRIORITY_READER: The priority for extensions that only read records from the observation dictionary. This is also suitable for extensions that do not use the observation dictionary at all. Extensions which read the reported values should be fired after all the extensions which have other priorities, e.g, PRIORITY_WRITER and PRIORITY_EDITOR because it should read the final values.

See the details in the documentation of Trainer.

4. finalizer

You can specify a function which takes Trainer object to finalize the extension. It is called once at the end of the whole training loops, namely, the run() finished.

5. initializer

You can specify a function which takes Trainer object to initialize the extension. It is called once at the beginning of the training loop, namely, before starting the actual loop.

Write a class inherited from Extension class

This is the way to define your own extension with the maximum degree of freedom. You can keep any values inside of the extension and serialize them.

As an example, let’s make an extension that drops the learning rate polynomially. It calculates the learning rate by this equation:

\[\eta = \eta_{\rm init} \left( 1 - \frac{t}{t_{\rm max}} \right)^{\rm power}\]

The learning rate will be dropped like the curve below with \({\rm power} = 0.5\):

../_images/polynomial.png
class PolynomialShift(training.Extension):

    def __init__(self, attr, power, stop_trigger, batchsize=None,
                 len_dataset=None):
        self._attr = attr
        self._power = power
        self._init = None
        self._t = 0
        self._last_value = 0

        if stop_trigger[1] == 'iteration':
            self._maxiter = stop_trigger[0]
        elif stop_trigger[1] == 'epoch':
            if batchsize is None or len_dataset is None:
                raise ValueError(
                    'When the unit of \'stop_trigger\' is \'epoch\', '
                    '\'batchsize\' and \'len_dataset\' should be '
                    'specified to calculate the maximum iteration.')
            n_iter_per_epoch = len_dataset / float(batchsize)
            self._maxiter = float(stop_trigger[0] * n_iter_per_epoch)

    def initialize(self, trainer):
        optimizer = trainer.updater.get_optimizer('main')
        # ensure that _init is set
        if self._init is None:
            self._init = getattr(optimizer, self._attr)

    def __call__(self, trainer):
        self._t += 1

        optimizer = trainer.updater.get_optimizer('main')
        value = self._init * ((1 - (self._t / self._maxiter)) ** self._power)
        setattr(optimizer, self._attr, value)
        self._last_value = value

    def serialize(self, serializer):
        self._t = serializer('_t', self._t)
        self._last_value = serializer('_last_value', self._last_value)
        if isinstance(self._last_value, np.ndarray):
            self._last_value = np.asscalar(self._last_value)
stop_trigger = (10000, 'iteration')
trainer.extend(PolynomialShift('lr', 0.5, stop_trigger)

This extension PolynomialShift takes five arguments.

  • attr: The name of optimizer property you want to update by this extension.
  • power: The power of the above equation to calculate the learning rate.
  • stop_trigger: The trigger given to the Trainer object to specify when to stop the training loop.
  • batchsize: The training mini-batchsize.
  • len_dataset: The length of dataset, i.e., the number of data in the training dataset.

This extension calculates the number of iterations which will be performed in training by using stop_trigger, batchsize, and len_dataset, then store it as a property _maxiter. This property will be used in __call__() method to update the learning rate. initialize() method obtains the initial learning rate from the optimizer set to give Trainer object. serialize() method stores or recovers the properties, _t (number of iterations) and _last_value (the latest learning rate), which this extension has.