keyboard interrupt for adversarial training

toolkit.py

@@ -2280,38 +2280,41 @@ class ClassificationProjectDecorr(ClassificationProject):
         self.model.stop_training = False
         callbacks.on_train_begin()
         epoch_logs = {}
-        for epoch in range(epochs):
-            callbacks.on_epoch_begin(epoch)
-            logger.info("Fitting epoch {}".format(epoch))
-            metrics = None
-            avg_metrics = None
-            for batch_id in tqdm(range(self.steps_per_epoch)):
-                x, y, w = next(batch_generator)
-                batch_logs = {}
-                batch_logs['batch'] = batch_id
-                batch_logs['size'] = len(x)
-                callbacks.on_batch_begin(batch_id, batch_logs)
-
-                # fit the classifier
-                batch_metrics = self.model.train_on_batch(
-                    x, y, sample_weight=w
-                )
+        try:
+            for epoch in range(epochs):
+                callbacks.on_epoch_begin(epoch)
+                logger.info("Fitting epoch {}".format(epoch))
+                metrics = None
+                avg_metrics = None
+                for batch_id in tqdm(range(self.steps_per_epoch)):
+                    x, y, w = next(batch_generator)
+                    batch_logs = {}
+                    batch_logs['batch'] = batch_id
+                    batch_logs['size'] = len(x)
+                    callbacks.on_batch_begin(batch_id, batch_logs)
+
+                    # fit the classifier
+                    batch_metrics = self.model.train_on_batch(
+                        x, y, sample_weight=w
+                    )
 
-                # fit the adversary
-                self.model_adv.train_on_batch(
-                    x, y[1:], sample_weight=w[1:]
-                )
-                outs = list(batch_metrics)
-                for l, o in zip(out_labels, outs):
-                    batch_logs[l] = float(o)
-                callbacks.on_batch_end(batch_id, batch_logs)
-            val_metrics = self.model.test_on_batch(*self.validation_data)
-            val_outs = list(val_metrics)
-            for l, o in zip(out_labels, val_outs):
-                epoch_logs['val_' + l] = float(o)
-            callbacks.on_epoch_end(epoch, epoch_logs)
-            if self.model.stop_training:
-                break
+                    # fit the adversary
+                    self.model_adv.train_on_batch(
+                        x, y[1:], sample_weight=w[1:]
+                    )
+                    outs = list(batch_metrics)
+                    for l, o in zip(out_labels, outs):
+                        batch_logs[l] = float(o)
+                    callbacks.on_batch_end(batch_id, batch_logs)
+                val_metrics = self.model.test_on_batch(*self.validation_data)
+                val_outs = list(val_metrics)
+                for l, o in zip(out_labels, val_outs):
+                    epoch_logs['val_' + l] = float(o)
+                callbacks.on_epoch_end(epoch, epoch_logs)
+                if self.model.stop_training:
+                    break
+        except KeyboardInterrupt:
+            pass
         callbacks.on_train_end()
 
         if not skip_checkpoint: