Added automatic train logs plotting after training is finished or user interrupts. Added sts augmentation as script parameter. Other improvements and bugfixes

Author: juan coria

common.py

@@ -5,7 +5,9 @@
 import torch
 import numpy as np
 import random
+
 import losses.config as cf
+from losses.triplet import SemiHardNegative
 from distances import CosineDistance, EuclideanDistance
 
 
@@ -86,9 +88,10 @@ def get_config(loss: str, nfeat: int, nclass: int, task: str, margin: float) ->
         print(f"[Margin: {margin}]")
         return cf.TripletConfig(DEVICE,
                                 margin=margin,
-                                distance=EuclideanDistance(),
+                                distance=CosineDistance(),
                                 size_average=False,
-                                online=task != 'sts')
+                                online=task != 'sts',
+                                sampling=SemiHardNegative(margin, deviation=0.02))
     elif loss == 'arcface':
         print(f"[Margin: {margin}]")
         return cf.ArcFaceConfig(DEVICE, nfeat, nclass, margin=margin)
@@ -134,3 +137,22 @@ def dump_params(filepath: str, args):
     with open(filepath, 'w') as out:
         for k, v in sorted(vars(args).items()):
             out.write(f"{k}={v}\n")
+
+
+def get_basic_plots(lr: float, batch_size: int, eval_metric: str, eval_metric_color: str) -> list:
+    return [
+        {
+            'log_file': 'loss.log',
+            'metric': 'Loss',
+            'color': 'blue',
+            'title': f'Train Loss - lr={lr} - batch_size={batch_size}',
+            'filename': 'loss-plot'
+        },
+        {
+            'log_file': 'metric.log',
+            'metric': eval_metric,
+            'color': eval_metric_color,
+            'title': f'Dev {eval_metric} - lr={lr} - batch_size={batch_size}',
+            'filename': f"dev-{eval_metric.lower().replace(' ', '-')}-plot"
+        }
+    ]

core/base.py

@@ -6,6 +6,7 @@
 from datasets.base import SimDatasetPartition
 from core.optim import Optimizer
 import common
+import visual_utils as vis
 
 
 class TrainingListener:
@@ -64,23 +65,42 @@ def __init__(self, loss_name: str, model: SimNet, loss_fn: nn.Module, partition:
     def _restore(self):
         if self.model_loader is not None:
             checkpoint = self.model_loader.restore(self.model, self.loss_fn, self.optim, self.loss_name)
-            epoch = checkpoint['epoch']
-            return checkpoint, epoch + 1
+            return checkpoint
         else:
-            return None, 1
-        
-    def train(self, epochs):
-        checkpoint, epoch = self._restore()
+            return None
+
+    def _create_plots(self, exp_path: str, plots: list):
+        print("Creating training plots before exiting...")
+        for plot in plots:
+            vis.visualize_logs(exp_path,
+                               log_file_name=plot['log_file'],
+                               metric_name=plot['metric'],
+                               color=plot['color'],
+                               title=plot['title'],
+                               plot_file_name=plot['filename'])
+        print("Done")
+
+    def _start_training(self, epochs):
+        checkpoint = self._restore()
 
         for cb in self.callbacks:
             cb.on_before_train(checkpoint)
 
-        for i in range(epoch, epoch+epochs):
+        for i in range(1, epochs + 1):
             self.train_epoch(i)
 
         for cb in self.callbacks:
             cb.on_after_train()
 
+    def train(self, epochs: int, exp_path: str, plots: list):
+        try:
+            self._start_training(epochs)
+            print("Training finished")
+            self._create_plots(exp_path, plots)
+        except KeyboardInterrupt:
+            print("Stopped by user")
+            self._create_plots(exp_path, plots)
+        
     def train_epoch(self, epoch):
         self.model.train()
 

datasets/semeval.py

@@ -7,6 +7,7 @@
 import torch
 
 from datasets.base import SimDataset, SimDatasetPartition
+from models import SimNet
 from sts.augmentation import SemEvalAugmentationStrategy, pad_sent_pair
 from sts import utils as sts
 

datasets/voxceleb.py

@@ -44,7 +44,7 @@ def __init__(self, batch_size: int, segment_size_millis: int):
         self.batch_size = batch_size
         self.segment_size_s = segment_size_millis / 1000
         self.nfeat = self.sample_rate * segment_size_millis // 1000
-        self.config = VoxCeleb1.config(self.segment_size_s)
+        self.config = self._create_config(self.segment_size_s)
         self.protocol = get_protocol(self.config.protocol_name, preprocessors=self.config.preprocessors)
         self.train_gen, self.dev_gen, self.test_gen = None, None, None
         print(f"[Segment Size: {self.segment_size_s}s]")
@@ -84,7 +84,7 @@ def _create_config(self, segment_size_sec: float):
                                                duration=segment_size_sec)
 
 
-class VoxCeleb2(VoxCelebDataset)
+class VoxCeleb2(VoxCelebDataset):
 
     def _create_config(self, segment_size_sec: float):
         return metrics.SpeakerValidationConfig(protocol_name='VoxCeleb.SpeakerVerification.VoxCeleb2',

distances.py

@@ -92,7 +92,7 @@ def to_sklearn_metric(self):
         return 'euclidean'
 
     def dist(self, x, y):
-        return torch.dist(x, y, p=2)
+        return torch.sum(torch.pow((x - y), 2), dim=1)
 
     def sqdist_sum(self, x, y):
         return (x - y).pow(2).sum()

losses/triplet.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
+import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
@@ -19,6 +20,9 @@ def triplets(self, y, distances):
         """
         raise NotImplementedError("a TripletSamplingStrategy should implement 'triplets'")
 
+    def filter(self, dist_pos, dist_neg):
+        return dist_pos, dist_neg
+
 
 class BatchAll(TripletSamplingStrategy):
     """
@@ -41,6 +45,20 @@ def triplets(self, y, distances):
         return anchors, positives, negatives
 
 
+class SemiHardNegative(TripletSamplingStrategy):
+
+    def __init__(self, m: float, deviation: float):
+        self.m = m
+        self.deviation = deviation
+
+    def filter(self, dist_pos, dist_neg):
+        keep_inds = []
+        for i in range(dist_neg.size(0)):
+            keep_inds.append(self.m >= dist_neg[i] or self.deviation >= dist_neg[i] - self.m)
+        keep_inds = torch.Tensor(keep_inds).float()
+        return keep_inds
+
+
 class HardestNegative(TripletSamplingStrategy):
     """
     Hardest negative strategy.
@@ -100,7 +118,7 @@ class TripletLoss(nn.Module):
     :param device: a device in which to run the computation
     :param margin: a margin value to separe classes
     :param distance: a distance object to measure between the samples
-    :param strategy: a TripletSamplingStrategy
+    :param sampling: a TripletSamplingStrategy
     """
 
     def __init__(self, device: str, margin: float, distance: Distance,
@@ -153,7 +171,10 @@ def forward(self, feat, logits, y):
             # Calculate the distances to positives and negatives for each anchor
             dpos = self.distance.dist(anchors, positives)
             dneg = self.distance.dist(anchors, negatives)
+            # keep_mask = self.sampling.filter(dpos, dneg).to(self.device)
+            # dpos = keep_mask * dpos
+            # dneg = keep_mask * dneg
 
         # Calculate the loss using the margin
-        loss = F.relu(dpos - dneg + self.margin)
+        loss = F.relu(torch.pow(dpos, 2) - torch.pow(dneg, 2) + self.margin)
         return loss.mean() if self.size_average else loss.sum()

metrics.py

@@ -195,10 +195,6 @@ def eval(self, model, partition: str = 'development'):
         # Returning 1-eer because the evaluator keeps track of the highest metric value
         return 1 - eer, y_pred, y_true
 
-    def on_before_train(self, checkpoint):
-        if checkpoint is not None:
-            self.best_metric = checkpoint['accuracy']
-
     def on_after_epoch(self, epoch, model, loss_fn, optim):
         if epoch % self.eval_interval == 0:
             metric_value, dists, y_true = self.eval(model.to_prediction_model(), self.partition)
@@ -259,10 +255,6 @@ def _eval(self, model):
         feat_test, y_test = np.concatenate(feat_test), np.concatenate(y_test)
         return feat_test, y_test
 
-    def on_before_train(self, checkpoint):
-        if checkpoint is not None:
-            self.best_metric = checkpoint['accuracy']
-
     def on_before_epoch(self, epoch):
         self.feat_train, self.y_train = [], []
 
@@ -339,10 +331,6 @@ def eval(self, model):
         y_test = np.concatenate(y_test)
         return phrases, feat_test, y_test
 
-    def on_before_train(self, checkpoint):
-        if checkpoint is not None:
-            self.best_metric = checkpoint['accuracy']
-
     def on_after_epoch(self, epoch, model, loss_fn, optim):
         _, feat_test, y_test = self.eval(model.to_prediction_model())
         metric_value = self.metric.get()
@@ -409,10 +397,6 @@ def eval(self, model):
         feat_test, y_test = np.concatenate(feat_test), np.concatenate(y_test)
         return phrases, feat_test, y_test
 
-    def on_before_train(self, checkpoint):
-        if checkpoint is not None:
-            self.best_metric = checkpoint['accuracy']
-
     def on_after_epoch(self, epoch, model, loss_fn, optim):
         phrases, feat_test, y_test = self.eval(model)
         metric_value = self.metric.get()