Add final plugin call on training end

If you want your plugin to be call at the whole training end (very
useful for one last save of progress, e.g. network snapshots or output
generated), just add it to the end queue with interval 1. Implemented by
default for SaverPlugin and OutputGenerator.

plugins.py

@@ -144,7 +144,7 @@ class SaverPlugin(Plugin):
     last_pattern = 'network-snapshot-{}-{}.dat'
 
     def __init__(self, checkpoints_path, keep_old_checkpoints=False, network_snapshot_ticks=40):
-        super().__init__([(network_snapshot_ticks, 'epoch')])
+        super().__init__([(network_snapshot_ticks, 'epoch'), (1, 'end')])
         self.checkpoints_path = checkpoints_path
         self.keep_old_checkpoints = keep_old_checkpoints
         self._best_val_loss = float('+inf')
@@ -165,6 +165,9 @@ def epoch(self, epoch_index):
                 )
             )
 
+    def end(self, *args):
+        self.epoch(*args)
+
     def _clear(self, pattern):
         pattern = os.path.join(self.checkpoints_path, pattern)
         for file_name in glob(pattern):
@@ -174,7 +177,7 @@ def _clear(self, pattern):
 class OutputGenerator(Plugin):
 
     def __init__(self, sample_fn, output_postprocessors, samples_count=6, output_snapshot_ticks=3):
-        super(OutputGenerator, self).__init__([(output_snapshot_ticks, 'epoch')])
+        super(OutputGenerator, self).__init__([(output_snapshot_ticks, 'epoch'), (1, 'end')])
         self.sample_fn = sample_fn
         self.output_postprocessors = output_postprocessors
         self.samples_count = samples_count
@@ -188,6 +191,9 @@ def epoch(self, epoch_index):
         for proc in self.output_postprocessors:
             proc(out, self.trainer.cur_nimg // 1000)
 
+    def end(self, *args):
+        self.epoch(*args)
+
 
 class CometPlugin(Plugin):
 

trainer.py

@@ -40,7 +40,8 @@ def __init__(self,
         self.plugin_queues = {
             'iteration': [],
             'epoch': [],
-            's':    []
+            's':    [],
+            'end':  []
         }
 
     def register_plugin(self, plugin):
@@ -79,6 +80,7 @@ def run(self, total_kimg=1):
                 self.stats['kimg_stat']['val'] = self.cur_nimg / 1000.
                 self.stats['tick_stat']['val'] = self.cur_tick
                 self.call_plugins('epoch', self.cur_tick)
+        self.call_plugins('end', 1)
 
     def train(self):
         fake_latents_in = self.random_latents_generator().cuda()