Update things around training

- Update batchsize to 16
- Fix validataion dataset collate and preprocessing

Author: mthrok

examples/source_separation/README.md

@@ -9,7 +9,7 @@ This directory contains reference implementations for source separations. For th
 ### Overview
 
 To traing a model, you can use [`train.py`](./train.py). This script takes the form of
-`trin.py [parameters for distributed training] -- [parameters for model/training]`
+`train.py [parameters for distributed training] -- [parameters for model/training]`
 
     ```
     python train.py \
@@ -162,7 +162,7 @@ python -u \
   --sample-rate 8000 \
   --dataset-dir "${dataset_dir}" \
   --save-dir "${save_dir}" \
-  --batch-size $((32 / SLURM_NTASKS))
+  --batch-size $((16 / SLURM_NTASKS))
 ```
 
 </details>

examples/source_separation/conv_tasnet/README.md

@@ -13,7 +13,7 @@ For the usage, please checkout the [source separation README](../README.md).
 The default training/model configurations follows the best non-causal implementation from the paper. (causal configuration is not implemented.)
 
  - Sample rate: 8000 Hz
- - Batch size: total 32 over distributed training workers
+ - Batch size: total 16 over distributed training workers
  - Epochs: 100
  - Initial learning rate: 1e-3
  - Gradient clipping: maximum L2 norm of 5.0
@@ -31,30 +31,22 @@ The default training/model configurations follows the best non-causal implementa
  - The number of TCN convolution block layers (X): 8
  - The number of TCN convolution blocks (R): 3
 
-## Training
-
-The training takes about 5 mins per epoch with 8 V100 GPUs in a single node.
-
 ## Evaluation
 
 The following is the evaluation result of training the model on WSJ0-2mix and WSJ0-3mix datasets.
 
 ### wsj0-mix 2speakers
 
-|                   | SI-SNRi (dB) | SDRi (dB) | Epoch |
-|:-----------------:|-------------:|----------:|------:|
-| Reference         |         15.3 |      15.6 |       |
-| "min" Validation  |        12.63 |     12.63 |   100 |
-| "min" Evaluation  |        10.59 |     10.58 |   100 |
-| "max" Validation  |        12.72 |     12.72 |   100 |
-| "max" Evaluation  |        11.00 |     11.00 |   100 |
+|                    | SI-SNRi (dB) | SDRi (dB) | Epoch |
+|:------------------:|-------------:|----------:|------:|
+| Reference          |         15.3 |      15.6 |       |
+| Validation dataset |         13.3 |      13.3 |    86 |
+| Evaluation dataset |         11.3 |      11.3 |    86 |
 
 ### wsj0-mix 3speakers
 
-|                   | SI-SNRi (dB) | SDRi (dB) | Epoch |
-|:-----------------:|-------------:|----------:|------:|
-| Reference         |         12.7 |      13.1 |       |
-| "min" Validation  |        10.75 |     10.75 |   99  |
-| "min" Evaluation  |         8.39 |      8.38 |   99  |
-| "max" Validation  |        10.87 |     10.86 |   99  |
-| "max" Evaluation  |         8.23 |      8.20 |   99  |
+|                    | SI-SNRi (dB) | SDRi (dB) | Epoch |
+|:------------------:|-------------:|----------:|------:|
+| Reference          |         12.7 |      13.1 |       |
+| Validation dataset |         11.5 |      11.5 |    87 |
+| Evaluation dataset |          8.7 |       8.6 |    87 |

examples/source_separation/conv_tasnet/train.py

@@ -14,7 +14,7 @@
 
 
 def _parse_args(args):
-    default_batch_size = 32 // torch.distributed.get_world_size()
+    default_batch_size = 16 // torch.distributed.get_world_size()
 
     parser = argparse.ArgumentParser(description=__doc__,)
     parser.add_argument("--debug", action="store_true", help="Enable debug behavior.")
@@ -61,7 +61,7 @@ def _parse_args(args):
         "--batch-size",
         default=default_batch_size,
         type=int,
-        help=f"Batch size. (default: {default_batch_size} ( == 32 // batch_size))",
+        help=f"Batch size. (default: {default_batch_size} (== 16 // world_size))",
     )
     group = parser.add_argument_group("Training Options")
     group.add_argument(
@@ -133,34 +133,48 @@ def _get_dataloader(dataset_type, dataset_dir, num_speakers, sample_rate, batch_
     train_dataset, valid_dataset, eval_dataset = dataset_utils.get_dataset(
         dataset_type, dataset_dir, num_speakers, sample_rate,
     )
-    collate_fn = dataset_utils.get_collate_fn(
-        dataset_type, sample_rate=sample_rate, duration=4
+    train_collate_fn = dataset_utils.get_collate_fn(
+        dataset_type, mode='train', sample_rate=sample_rate, duration=4
     )
 
+    test_collate_fn = dataset_utils.get_collate_fn(dataset_type, mode='test')
+
     train_loader = torch.utils.data.DataLoader(
         train_dataset,
         batch_size=batch_size,
         sampler=torch.utils.data.distributed.DistributedSampler(train_dataset),
-        collate_fn=collate_fn,
+        collate_fn=train_collate_fn,
         pin_memory=True,
     )
     valid_loader = torch.utils.data.DataLoader(
         valid_dataset,
         batch_size=batch_size,
         sampler=torch.utils.data.distributed.DistributedSampler(valid_dataset),
-        collate_fn=collate_fn,
+        collate_fn=test_collate_fn,
         pin_memory=True,
     )
     eval_loader = torch.utils.data.DataLoader(
         eval_dataset,
         batch_size=batch_size,
         sampler=torch.utils.data.distributed.DistributedSampler(eval_dataset),
-        collate_fn=collate_fn,
+        collate_fn=test_collate_fn,
         pin_memory=True,
     )
     return train_loader, valid_loader, eval_loader
 
 
+def _write_header(log_path, args):
+    rows = [
+        [f"# torch: {torch.__version__}", ],
+        [f"# torchaudio: {torchaudio.__version__}", ]
+    ]
+    rows.append(["# arguments"])
+    for key, item in vars(args).items():
+        rows.append([f"#   {key}: {item}"])
+
+    dist_utils.write_csv_on_master(log_path, *rows)
+
+
 def train(args):
     args = _parse_args(args)
     _LG.info("%s", args)
@@ -237,7 +251,7 @@ def train(args):
     )
 
     log_path = args.save_dir / f"log.csv"
-    dist_utils.write_csv_on_master(log_path, [f"# {args}", ])
+    _write_header(log_path, args)
     dist_utils.write_csv_on_master(
         log_path,
         [

examples/source_separation/conv_tasnet/trainer.py

@@ -134,20 +134,22 @@ def _test(self, loader):
         total_si_snri = 0.0
         total_sdri = 0.0
 
-        for batch in loader:
-            mixed = batch.mix.to(self.device)
-            sources = batch.src.to(self.device)
-
-            estimate = self.model(mixed)
-            si_snri, sdri = si_sdr_improvement(estimate, sources, mixed)
-            si_snri = si_snri.sum()
-            sdri = sdri.sum()
-
-            dist.all_reduce(si_snri, dist.ReduceOp.SUM)
-            dist.all_reduce(sdri, dist.ReduceOp.SUM)
-
-            total_si_snri += si_snri.item()
-            total_sdri += sdri.item()
+        for samples in loader:
+            # Due to the possible length difference, we run evaluation sample-wise
+            for sample in samples:
+                mixed = sample.mix.to(self.device)
+                sources = sample.src.to(self.device)
+
+                estimate = self.model(mixed)
+                si_snri, sdri = si_sdr_improvement(estimate, sources, mixed)
+                si_snri = si_snri.sum()
+                sdri = sdri.sum()
+
+                dist.all_reduce(si_snri, dist.ReduceOp.SUM)
+                dist.all_reduce(sdri, dist.ReduceOp.SUM)
+
+                total_si_snri += si_snri.item()
+                total_sdri += sdri.item()
 
             if self.debug:
                 break

examples/source_separation/train.py

@@ -83,6 +83,11 @@ def _parse_args(args=None):
             'access, using `"file://..."` protocol. (default: "env://")'
         ),
     )
+    group.add_argument(
+        "--random-seed",
+        type=int,
+        help="Set random seed value. (default: None)",
+    )
     parser.add_argument(
         "rest", nargs=argparse.REMAINDER, help="Model-specific arguments."
     )
@@ -118,6 +123,8 @@ def _main(cli_args):
             backend='nccl' if torch.cuda.is_available() else 'gloo',
             init_method=args.sync_protocol,
         )
+        if args.random_seed is not None:
+            torch.manual_seed(args.random_seed)
         if torch.cuda.is_available():
             torch.cuda.set_device(args.device_id)
             _LG.info("CUDA device set to %s", args.device_id)

examples/source_separation/utils/dist_utils.py

@@ -58,11 +58,12 @@ def save_on_master(path, obj):
         torch.save(obj, path)
 
 
-def write_csv_on_master(path, items):
+def write_csv_on_master(path, *rows):
     if dist.get_rank() == 0:
         with open(path, "a", newline="") as fileobj:
             writer = csv.writer(fileobj)
-            writer.writerow(items)
+            for row in rows:
+                writer.writerow(row)
 
 
 def synchronize_params(path, device, *modules):