implement power low compression loss

maum-ai · Sep 1, 2019 · 0acfd8e · 0acfd8e
1 parent 36b59da
commit 0acfd8e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/config/default.yaml b/config/default.yaml
@@ -39,6 +39,7 @@ train:
     final: 0.05
   summary_interval: 1
   checkpoint_interval: 1000
+  complex_loss_ratio: 0.1  # the lambda in power-law compression loss computation
 ---
 log:
   chkpt_dir: 'chkpt'

diff --git a/utils/train.py b/utils/train.py
@@ -64,9 +64,16 @@ def train(args, pt_dir, chkpt_path, trainloader, testloader, writer, logger, hp,
                 mask = model(mixed_mag, dvec)
                 output = mixed_mag * mask
 
-                # output = torch.pow(torch.clamp(output, min=0.0), hp.audio.power)
-                # target_mag = torch.pow(torch.clamp(target_mag, min=0.0), hp.audio.power)
-                loss = criterion(output, target_mag)
+                # Power-law compression
+                magnitude_loss = criterion(
+                    torch.pow(torch.abs(output), hp.audio.power), 
+                    torch.pow(torch.abs(target_mag), hp.audio.power),
+                )
+                complex_loss = criterion(
+                    torch.pow(torch.clamp(output, min=0.0), hp.audio.power), 
+                    torch.pow(torch.clamp(target_mag, min=0.0), hp.audio.power),
+                )
+                loss = magnitude_loss + complex_loss * hp.train.complex_loss_ratio
 
                 optimizer.zero_grad()
                 loss.backward()