extension.rst

Create your own trainer extension

chainer

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

• basic

• How to create your own trainer extension

  • by defining a simple function<function>
  • by defining a function decorated with @make_extension <decorator>
  • by defining a class inherited from Extension class <class>

What is trainer Extension?

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

A ~chainer.training.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 ~chainer.training.Extension by writing a simple function which takes ~chainer.training.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 ~chainer.training.Trainer object via ~chainer.training.Trainer.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

~chainer.training.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 ~chainer.training.Trainer.extend method, the trigger specified in ~chainer.training.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 ~chainer.training.make_extension decorator.

1. trigger

trigger is an object that takes a ~chainer.training.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 ~chainer.training.triggers.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 ~chainer.training.trainer.extend method that adds an extension to a ~chainer.training.Trainer object. The priority of triggers is as follows:

• When both ~chainer.training.Trainer.extend and a given ~chainer.training.Extension have triggers, the trigger given to ~chainer.training.Trainer.extend is used.
• When None is given to ~chainer.training.Trainer.extend as the trigger argument and a given ~chainer.training.Extension has trigger, the trigger given to the ~chainer.training.Extension is used.
• When both trigger attributes in ~chainer.training.Trainer.extend and ~chainer.training.Extension are None, the ~chainer.training.Extension will be fired every iteration.

See the details in the documentation of ~chainer.training.get_trigger.

2. default_name

An ~chainer.training.Extension is kept in a dictionary which is a property in a ~chainer.training.Trainer. This argument gives the name of the ~chainer.training.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 ~chainer.training.Trainer object. There are three standard values for the priorities:

• ~chainer.training.extension.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.
• ~chainer.training.extension.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.
• ~chainer.training.extension.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 ~chainer.training.Trainer.

4. finalizer

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

5. initializer

You can specify a function which takes ~chainer.training.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$:

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 ~chainer.traininig.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 ~chainer.training.Trainer object. ~chainer.traininig.serialize method stores or recovers the properties, _t (number of iterations) and _last_value (the latest learning rate), which this extension has.