Refactored source

app.py

@@ -6,7 +6,6 @@ import torch
 import torchaudio
 import gradio as gr
 
-from src.helpers import utils
 from src.training.dcc_tf import Net as Waveformer
 
 TARGETS = [
@@ -34,7 +33,7 @@ if not os.path.exists('default_ckpt.pt'):
 # Instantiate model
 params = utils.Params('default_config.json')
 model = Waveformer(**params.model_params)
-utils.load_checkpoint('default_ckpt.pt', model)
+model.load_state_dict(torch.load('default_ckpt.pt', map_location=torch.device('cpu')))
 model.eval()
 
 def waveformer(audio, label_choices):

src/helpers/utils.py

@@ -1,205 +0,0 @@
-"""A collection of useful helper functions"""
-
-import os
-import logging
-import json
-
-import torch
-from torch.profiler import profile, record_function, ProfilerActivity
-import pandas as pd
-from torchmetrics.functional import(
-    scale_invariant_signal_noise_ratio as si_snr,
-    signal_noise_ratio as snr,
-    signal_distortion_ratio as sdr,
-    scale_invariant_signal_distortion_ratio as si_sdr)
-import matplotlib.pyplot as plt
-
-class Params():
-    """Class that loads hyperparameters from a json file.
-    Example:
-    ```
-    params = Params(json_path)
-    print(params.learning_rate)
-    params.learning_rate = 0.5  # change the value of learning_rate in params
-    ```
-    """
-
-    def __init__(self, json_path):
-        with open(json_path) as f:
-            params = json.load(f)
-            self.__dict__.update(params)
-
-    def save(self, json_path):
-        with open(json_path, 'w') as f:
-            json.dump(self.__dict__, f, indent=4)
-
-    def update(self, json_path):
-        """Loads parameters from json file"""
-        with open(json_path) as f:
-            params = json.load(f)
-            self.__dict__.update(params)
-
-    @property
-    def dict(self):
-        """Gives dict-like access to Params instance by `params.dict['learning_rate']"""
-        return self.__dict__
-
-def save_graph(train_metrics, test_metrics, save_dir):
-    metrics = [snr, si_snr]
-    results = {'train_loss': train_metrics['loss'],
-               'test_loss' : test_metrics['loss']}
-
-    for m_fn in metrics:
-        results["train_"+m_fn.__name__] = train_metrics[m_fn.__name__]
-        results["test_"+m_fn.__name__] = test_metrics[m_fn.__name__]
-
-    results_pd = pd.DataFrame(results)
-
-    results_pd.to_csv(os.path.join(save_dir, 'results.csv'))
-
-    fig, temp_ax = plt.subplots(2, 3, figsize=(15,10))
-    axs=[]
-    for i in temp_ax:
-        for j in i:
-            axs.append(j)
-
-    x = range(len(train_metrics['loss']))
-    axs[0].plot(x, train_metrics['loss'], label='train')
-    axs[0].plot(x, test_metrics['loss'], label='test')
-    axs[0].set(ylabel='Loss')
-    axs[0].set(xlabel='Epoch')
-    axs[0].set_title('loss',fontweight='bold')
-    axs[0].legend()
-
-    for i in range(len(metrics)):
-        axs[i+1].plot(x, train_metrics[metrics[i].__name__], label='train')
-        axs[i+1].plot(x, test_metrics[metrics[i].__name__], label='test')
-        axs[i+1].set(xlabel='Epoch')
-        axs[i+1].set_title(metrics[i].__name__,fontweight='bold')
-        axs[i+1].legend()
-
-    plt.tight_layout()
-    plt.savefig(os.path.join(save_dir, 'results.png'))
-    plt.close(fig)
-
-def set_logger(log_path):
-    """Set the logger to log info in terminal and file `log_path`.
-    In general, it is useful to have a logger so that every output to the terminal is saved
-    in a permanent file. Here we save it to `model_dir/train.log`.
-    Example:
-    ```
-    logging.info("Starting training...")
-    ```
-    Args:
-        log_path: (string) where to log
-    """
-    logger = logging.getLogger()
-    logger.setLevel(logging.INFO)
-    logger.handlers.clear()
-
-    # Logging to a file
-    file_handler = logging.FileHandler(log_path)
-    file_handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s: %(message)s'))
-    logger.addHandler(file_handler)
-
-    # Logging to console
-    stream_handler = logging.StreamHandler()
-    stream_handler.setFormatter(logging.Formatter('%(message)s'))
-    logger.addHandler(stream_handler)
-
-def load_checkpoint(checkpoint, model, optim=None, lr_sched=None, data_parallel=False):
-    """Loads model parameters (state_dict) from file_path.
-
-    Args:
-        checkpoint: (string) filename which needs to be loaded
-        model: (torch.nn.Module) model for which the parameters are loaded
-        data_parallel: (bool) if the model is a data parallel model
-    """
-    if not os.path.exists(checkpoint):
-        raise("File doesn't exist {}".format(checkpoint))
-
-    state_dict = torch.load(checkpoint)
-
-    if data_parallel:
-        state_dict['model_state_dict'] = {
-            'module.' + k: state_dict['model_state_dict'][k]
-            for k in state_dict['model_state_dict'].keys()}
-    model.load_state_dict(state_dict['model_state_dict'])
-
-    if optim is not None:
-        optim.load_state_dict(state_dict['optim_state_dict'])
-
-    if lr_sched is not None:
-        lr_sched.load_state_dict(state_dict['lr_sched_state_dict'])
-
-    return state_dict['epoch'], state_dict['train_metrics'], \
-           state_dict['val_metrics']
-
-def save_checkpoint(checkpoint, epoch, model, optim=None, lr_sched=None,
-                    train_metrics=None, val_metrics=None, data_parallel=False):
-    """Saves model parameters (state_dict) to file_path.
-
-    Args:
-        checkpoint: (string) filename which needs to be loaded
-        model: (torch.nn.Module) model for which the parameters are loaded
-        data_parallel: (bool) if the model is a data parallel model
-    """
-    if os.path.exists(checkpoint):
-        raise("File already exists {}".format(checkpoint))
-
-    model_state_dict = model.state_dict()
-    if data_parallel:
-        model_state_dict = {
-            k.partition('module.')[2]:
-            model_state_dict[k] for k in model_state_dict.keys()}
-
-    optim_state_dict = None if not optim else optim.state_dict()
-    lr_sched_state_dict = None if not lr_sched else lr_sched.state_dict()
-
-    state_dict = {
-        'epoch': epoch,
-        'model_state_dict': model_state_dict,
-        'optim_state_dict': optim_state_dict,
-        'lr_sched_state_dict': lr_sched_state_dict,
-        'train_metrics': train_metrics,
-        'val_metrics': val_metrics
-    }
-
-    torch.save(state_dict, checkpoint)
-
-def model_size(model):
-    """
-    Returns size of the `model` in millions of parameters.
-    """
-    num_train_params = sum(
-        p.numel() for p in model.parameters() if p.requires_grad)
-    return num_train_params / 1e6
-
-def run_time(model, inputs, profiling=False):
-    """
-    Returns runtime of a model in ms.
-    """
-    # Warmup
-    for _ in range(100):
-        output = model(*inputs)
-
-    with profile(activities=[ProfilerActivity.CPU],
-                 record_shapes=True) as prof:
-        with record_function("model_inference"):
-            output = model(*inputs)
-
-    # Print profiling results
-    if profiling:
-        print(prof.key_averages().table(sort_by="self_cpu_time_total",
-                                        row_limit=20))
-
-    # Return runtime in ms
-    return prof.profiler.self_cpu_time_total / 1000
-
-def format_lr_info(optimizer):
-    lr_info = ""
-    for i, pg in enumerate(optimizer.param_groups):
-        lr_info += " {group %d: params=%.5fM lr=%.1E}" % (
-            i, sum([p.numel() for p in pg['params']]) / (1024 ** 2), pg['lr'])
-    return lr_info
-

src/training/eval.py

@@ -1,214 +0,0 @@
-"""
-Test script to evaluate the model.
-"""
-
-import argparse
-import importlib
-import multiprocessing
-import os, glob
-import logging
-
-import numpy as np
-import torch
-import pandas as pd
-import torch.nn as nn
-from torch.utils.tensorboard import SummaryWriter
-from torch.profiler import profile, record_function, ProfilerActivity
-from tqdm import tqdm  # pylint: disable=unused-import
-from torchmetrics.functional import(
-    scale_invariant_signal_noise_ratio as si_snr,
-    signal_noise_ratio as snr,
-    signal_distortion_ratio as sdr,
-    scale_invariant_signal_distortion_ratio as si_sdr)
-
-from src.helpers import utils
-from src.training.synthetic_dataset import FSDSoundScapesDataset, tensorboard_add_metrics
-from src.training.synthetic_dataset import tensorboard_add_sample
-
-def test_epoch(model: nn.Module, device: torch.device,
-               test_loader: torch.utils.data.dataloader.DataLoader,
-               n_items: int, loss_fn, metrics_fn,
-               profiling: bool = False, epoch: int = 0,
-               writer: SummaryWriter = None, data_params = None) -> float:
-    """
-    Evaluate the network.
-    """
-    model.eval()
-    metrics = {}
- 
-    with torch.no_grad():
-        for batch_idx, (mixed, label, gt) in \
-                enumerate(tqdm(test_loader, desc='Test', ncols=100)):
-            mixed = mixed.to(device)
-            label = label.to(device)
-            gt = gt.to(device)
-
-            # Run through the model
-            with profile(activities=[ProfilerActivity.CPU],
-                         record_shapes=True) as prof:
-                with record_function("model_inference"):
-                    output = model(mixed, label)
-            if profiling:
-                logging.info(
-                    prof.key_averages().table(sort_by="self_cpu_time_total",
-                                              row_limit=20))
-
-            # Compute loss
-            loss = loss_fn(output, gt)
-
-            # Compute metrics
-            metrics_batch = metrics_fn(mixed, output, gt)
-            metrics_batch['loss'] = [loss.item()]
-            metrics_batch['runtime'] = [prof.profiler.self_cpu_time_total/1000]
-            for k in metrics_batch.keys():
-                if not k in metrics:
-                    metrics[k] = metrics_batch[k]
-                else:
-                    metrics[k] += metrics_batch[k]
-
-            if writer is not None:
-                if batch_idx == 0:
-                    tensorboard_add_sample(
-                        writer, tag='Test',
-                        sample=(mixed[:8], label[:8], gt[:8], output[:8]),
-                        step=epoch, params=data_params)
-                tensorboard_add_metrics(
-                    writer, tag='Test', metrics=metrics_batch, label=label,
-                    step=epoch)
-
-            if n_items is not None and batch_idx == (n_items - 1):
-                break
-
-        avg_metrics = {k: np.mean(metrics[k]) for k in metrics.keys()}
-        avg_metrics_str = "Test:"
-        for m in avg_metrics.keys():
-            avg_metrics_str += ' %s=%.04f' % (m, avg_metrics[m])
-        logging.info(avg_metrics_str)
-
-        return avg_metrics
-
-def evaluate(network, args: argparse.Namespace):
-    """
-    Evaluate the model on a given dataset.
-    """
-
-    # Load dataset
-    data_test = FSDSoundScapesDataset(**args.test_data)
-    logging.info("Loaded test dataset at %s containing %d elements" %
-                 (args.test_data['input_dir'], len(data_test)))
- 
-    # Set up the device and workers.
-    use_cuda = args.use_cuda and torch.cuda.is_available()
-    if use_cuda:
-        gpu_ids = args.gpu_ids if args.gpu_ids is not None\
-                        else range(torch.cuda.device_count())
-        device_ids = [_ for _ in gpu_ids]
-        data_parallel = len(device_ids) > 1
-        device = 'cuda:%d' % device_ids[0]
-        torch.cuda.set_device(device_ids[0])
-        logging.info("Using CUDA devices: %s" % str(device_ids))
-    else:
-        data_parallel = False
-        device = torch.device('cpu')
-        logging.info("Using device: CPU")
-
-    # Set multiprocessing params
-    num_workers = min(multiprocessing.cpu_count(), args.n_workers)
-    kwargs = {
-        'num_workers': num_workers,
-        'pin_memory': True
-    } if use_cuda else {}
-
-    # Set up data loader
-    test_loader = torch.utils.data.DataLoader(data_test,
-                                              batch_size=args.eval_batch_size,
-                                              **kwargs)
-
-    # Set up model
-    model = network.Net(**args.model_params)
-    if use_cuda and data_parallel:
-        model = nn.DataParallel(model, device_ids=device_ids)
-        logging.info("Using data parallel model")
-    model.to(device)
-
-    # Load weights
-    if args.pretrain_path == "best":
-        ckpts = glob.glob(os.path.join(args.exp_dir, '*.pt'))
-        ckpts.sort(
-            key=lambda _: int(os.path.splitext(os.path.basename(_))[0]))
-        val_metrics = torch.load(ckpts[-1])['val_metrics'][args.base_metric]
-        best_epoch = max(range(len(val_metrics)), key=val_metrics.__getitem__)
-        args.pretrain_path = os.path.join(args.exp_dir, '%d.pt' % best_epoch)
-        logging.info(
-            "Found 'best' validation %s=%.02f at %s" %
-            (args.base_metric, val_metrics[best_epoch], args.pretrain_path))
-    if args.pretrain_path != "":
-        utils.load_checkpoint(
-            args.pretrain_path, model, data_parallel=data_parallel)
-        logging.info("Loaded pretrain weights from %s" % args.pretrain_path)
-
-    # Evaluate
-    try:
-        return test_epoch(
-            model, device, test_loader, args.n_items, network.loss,
-            network.metrics, args.profiling)
-    except KeyboardInterrupt:
-        print("Interrupted")
-    except Exception as _:  # pylint: disable=broad-except
-        import traceback  # pylint: disable=import-outside-toplevel
-        traceback.print_exc()
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    # Data Params
-    parser.add_argument('experiments', nargs='+', type=str,
-                        default=None,
-                        help="List of experiments to evaluate. "
-                        "Provide only one experiment when providing "
-                        "pretrained path. If pretrianed path is not "
-                        "provided, epoch with best validation metric "
-                        "is used for evaluation.")
-    parser.add_argument('--results', type=str, default="",
-                        help="Path to the CSV file to store results.")
-
-    # System params
-    parser.add_argument('--n_items', type=int, default=None,
-                        help="Number of items to test.")
-    parser.add_argument('--pretrain_path', type=str, default="best",
-                        help="Path to pretrained weights")
-    parser.add_argument('--profiling', dest='profiling', action='store_true',
-                        help="Enable or disable profiling.")
-    parser.add_argument('--use_cuda', dest='use_cuda', action='store_true',
-                        help="Whether to use cuda")
-    parser.add_argument('--gpu_ids', nargs='+', type=int, default=None,
-                        help="List of GPU ids used for training. "
-                        "Eg., --gpu_ids 2 4. All GPUs are used by default.")
-    args = parser.parse_args()
-
-    results = []
-
-    for exp_dir in args.experiments:
-        eval_args = argparse.Namespace(**vars(args))
-        eval_args.exp_dir = exp_dir
-
-        utils.set_logger(os.path.join(exp_dir, 'eval.log'))
-        logging.info("Evaluating %s ..." % exp_dir)
-
-        # Load model and training params
-        params = utils.Params(os.path.join(exp_dir, 'config.json'))
-        for k, v in params.__dict__.items():
-            vars(eval_args)[k] = v
-
-        network = importlib.import_module(eval_args.model)
-        logging.info("Imported the model from '%s'." % eval_args.model)
-
-        curr_res = evaluate(network, eval_args)
-        curr_res['experiment'] = os.path.basename(exp_dir)
-        results.append(curr_res)
-
-        del eval_args
-
-    if args.results != "":
-        print("Writing results to %s" % args.results)
-        pd.DataFrame(results).to_csv(args.results, index=False)

src/training/synthetic_dataset.py

@@ -1,168 +0,0 @@
-"""
-Torch dataset object for synthetically rendered spatial data.
-"""
-
-import os
-import json
-import random
-from pathlib import Path
-import logging
-
-import numpy as np
-import pandas as pd
-import matplotlib.pyplot as plt
-import scaper
-import torch
-import torchaudio
-import torchaudio.transforms as AT
-from random import randrange
-
-class FSDSoundScapesDataset(torch.utils.data.Dataset):  # type: ignore
-    """
-    Base class for FSD Sound Scapes dataset
-    """
-
-    _labels = [
-    "Acoustic_guitar", "Applause", "Bark", "Bass_drum",
-    "Burping_or_eructation", "Bus", "Cello", "Chime", "Clarinet",
-    "Computer_keyboard", "Cough", "Cowbell", "Double_bass",
-    "Drawer_open_or_close", "Electric_piano", "Fart", "Finger_snapping",
-    "Fireworks", "Flute", "Glockenspiel", "Gong", "Gunshot_or_gunfire",
-    "Harmonica", "Hi-hat", "Keys_jangling", "Knock", "Laughter", "Meow",
-    "Microwave_oven", "Oboe", "Saxophone", "Scissors", "Shatter",
-    "Snare_drum", "Squeak", "Tambourine", "Tearing", "Telephone",
-    "Trumpet", "Violin_or_fiddle", "Writing"]
-
-    def __init__(self, input_dir, dset='', sr=None,
-                 resample_rate=None, max_num_targets=1):
-        assert dset in ['train', 'val', 'test'], \
-            "`dset` must be one of ['train', 'val', 'test']"
-        self.dset = dset
-        self.max_num_targets = max_num_targets
-        self.fg_dir = os.path.join(input_dir, 'FSDKaggle2018/%s' % dset)
-        if dset in ['train', 'val']:
-            self.bg_dir = os.path.join(
-                input_dir,
-                'TAU-acoustic-sounds/'
-                'TAU-urban-acoustic-scenes-2019-development')
-        else:
-            self.bg_dir = os.path.join(
-                input_dir,
-                'TAU-acoustic-sounds/'
-                'TAU-urban-acoustic-scenes-2019-evaluation')
-        logging.info("Loading %s dataset: fg_dir=%s bg_dir=%s" %
-                     (dset, self.fg_dir, self.bg_dir))
-
-        self.samples = sorted(list(
-            Path(os.path.join(input_dir, 'jams', dset)).glob('[0-9]*')))
-
-        jamsfile = os.path.join(self.samples[0], 'mixture.jams')
-        _, jams, _, _ = scaper.generate_from_jams(
-            jamsfile, fg_path=self.fg_dir, bg_path=self.bg_dir)
-        _sr = jams['annotations'][0]['sandbox']['scaper']['sr']
-        assert _sr == sr, "Sampling rate provided does not match the data"
-
-        if resample_rate is not None:
-            self.resampler = AT.Resample(sr, resample_rate)
-            self.sr = resample_rate
-        else:
-            self.resampler = lambda a: a
-            self.sr = sr
-
-    def _get_label_vector(self, labels):
-        """
-        Generates a multi-hot vector corresponding to `labels`.
-        """
-        vector = torch.zeros(len(FSDSoundScapesDataset._labels))
-
-        for label in labels:
-            idx = FSDSoundScapesDataset._labels.index(label)
-            assert vector[idx] == 0, "Repeated labels"
-            vector[idx] = 1 
-
-        return vector
-
-    def __len__(self):
-        return len(self.samples)
-
-    def __getitem__(self, idx):
-        sample_path = self.samples[idx]
-        jamsfile = os.path.join(sample_path, 'mixture.jams')
-
-        mixture, jams, ann_list, event_audio_list = scaper.generate_from_jams(
-            jamsfile, fg_path=self.fg_dir, bg_path=self.bg_dir)
-        isolated_events = {}
-        for e, a in zip(ann_list, event_audio_list[1:]):
-            # 0th event is background
-            isolated_events[e[2]] = a
-        gt_events = list(pd.read_csv(
-            os.path.join(sample_path, 'gt_events.csv'), sep='\t')['label'])
-
-        mixture = torch.from_numpy(mixture).permute(1, 0)
-        mixture = self.resampler(mixture.to(torch.float))
-
-        if self.dset == 'train':
-            labels = random.sample(gt_events, randrange(1,self.max_num_targets+1))
-        elif self.dset == 'val':
-            labels = gt_events[:idx%self.max_num_targets+1]
-        elif self.dset == 'test':
-            labels = gt_events[:self.max_num_targets]
-        label_vector = self._get_label_vector(labels)
-
-        gt = torch.zeros_like(
-            torch.from_numpy(event_audio_list[1]).permute(1, 0))
-        for l in labels:
-            gt = gt + torch.from_numpy(isolated_events[l]).permute(1, 0)
-        gt = self.resampler(gt.to(torch.float))
-
-        return mixture, label_vector, gt #, jams
-
-def tensorboard_add_sample(writer, tag, sample, step, params):
-    """
-    Adds a sample of FSDSynthDataset to tensorboard.
-    """
-    if params['resample_rate'] is not None:
-        sr = params['resample_rate']
-    else:
-        sr = params['sr']
-    resample_rate = 16000 if sr > 16000 else sr
-
-    m, l, gt, o = sample
-    m, gt, o = (
-        torchaudio.functional.resample(_, sr, resample_rate).cpu()
-        for _ in (m, gt, o))
-
-    def _add_audio(a, audio_tag, axis, plt_title):
-        for i, ch in enumerate(a):
-            axis.plot(ch, label='mic %d' % i)
-            writer.add_audio(
-                '%s/mic %d' % (audio_tag, i), ch.unsqueeze(0), step, resample_rate)
-        axis.set_title(plt_title)
-        axis.legend()
-
-    for b in range(m.shape[0]):
-        label = []
-        for i in range(len(l[b, :])):
-            if l[b, i] == 1:
-                label.append(FSDSoundScapesDataset._labels[i])
-
-        # Add waveforms
-        rows = 3 # input, output, gt
-        fig = plt.figure(figsize=(10, 2 * rows))
-        axes = fig.subplots(rows, 1, sharex=True)
-        _add_audio(m[b], '%s/sample_%d/0_input' % (tag, b), axes[0], "Mixed")
-        _add_audio(o[b], '%s/sample_%d/1_output' % (tag, b), axes[1], "Output (%s)" % label)
-        _add_audio(gt[b], '%s/sample_%d/2_gt' % (tag, b), axes[2], "GT (%s)" % label)
-        writer.add_figure('%s/sample_%d/waveform' % (tag, b), fig, step)
-
-def tensorboard_add_metrics(writer, tag, metrics, label, step):
-    """
-    Add metrics to tensorboard.
-    """
-    vals = np.asarray(metrics['scale_invariant_signal_noise_ratio'])
-
-    writer.add_histogram('%s/%s' % (tag, 'SI-SNRi'), vals, step)
-
-    label_names = [FSDSoundScapesDataset._labels[torch.argmax(_)] for _ in label]
-    for l, v in zip(label_names, vals):
-        writer.add_histogram('%s/%s' % (tag, l), v, step)

src/training/train.py

@@ -1,311 +0,0 @@
-"""
-The main training script for training on synthetic data
-"""
-
-import argparse
-import multiprocessing
-import os
-import logging
-from pathlib import Path
-import random
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-from torch.utils.tensorboard import SummaryWriter
-from tqdm import tqdm  # pylint: disable=unused-import
-from torchmetrics.functional import(
-    scale_invariant_signal_noise_ratio as si_snr,
-    signal_noise_ratio as snr,
-    signal_distortion_ratio as sdr,
-    scale_invariant_signal_distortion_ratio as si_sdr)
-
-from src.helpers import utils
-from src.training.eval import test_epoch
-from src.training.synthetic_dataset import FSDSoundScapesDataset as Dataset
-from src.training.synthetic_dataset import tensorboard_add_sample
-
-def train_epoch(model: nn.Module, device: torch.device,
-                optimizer: optim.Optimizer,
-                train_loader: torch.utils.data.dataloader.DataLoader,
-                n_items: int, epoch: int = 0,
-                writer: SummaryWriter = None, data_params = None) -> float:
-
-    """
-    Train a single epoch.
-    """
-    # Set the model to training.
-    model.train()
-
-    # Training loop
-    losses = []
-    metrics = {}
-
-    with tqdm(total=len(train_loader), desc='Train', ncols=100) as t:
-        for batch_idx, (mixed, label, gt) in enumerate(train_loader):
-            mixed = mixed.to(device)
-            label = label.to(device)
-            gt = gt.to(device)
-
-            # Reset grad
-            optimizer.zero_grad()
-
-            # Run through the model
-            output = model(mixed, label)
-
-            # Compute loss
-            loss = network.loss(output, gt)
-
-            losses.append(loss.item())
-
-            # Backpropagation
-            loss.backward()
-
-            # Gradient clipping
-            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
-
-            # Update the weights
-            optimizer.step()
-
-            metrics_batch = network.metrics(mixed.detach(), output.detach(),
-                                            gt.detach())
-            for k in metrics_batch.keys():
-                if not k in metrics:
-                    metrics[k] = metrics_batch[k]
-                else:
-                    metrics[k] += metrics_batch[k]
-
-            if writer is not None and batch_idx == 0:
-                tensorboard_add_sample(
-                    writer, tag='Train',
-                    sample=(mixed.detach()[:8], label.detach()[:8],
-                            gt.detach()[:8], output.detach()[:8]),
-                    step=epoch, params=data_params)
-
-            # Show current loss in the progress meter
-            t.set_postfix(loss='%.05f'%loss.item())
-            t.update()
-
-            if n_items is not None and batch_idx == n_items:
-                break
-
-    avg_metrics = {k: np.mean(metrics[k]) for k in metrics.keys()}
-    avg_metrics['loss'] = np.mean(losses)
-    avg_metrics_str = "Train:"
-    for m in avg_metrics.keys():
-        avg_metrics_str += ' %s=%.04f' % (m, avg_metrics[m])
-    logging.info(avg_metrics_str)
-
-    return avg_metrics
-
-
-def train(args: argparse.Namespace):
-    """
-    Train the network.
-    """
-
-    # Load dataset
-    data_train = Dataset(**args.train_data)
-    logging.info("Loaded train dataset at %s containing %d elements" %
-                 (args.train_data['input_dir'], len(data_train)))
-    data_val = Dataset(**args.val_data)
-    logging.info("Loaded test dataset at %s containing %d elements" %
-                 (args.val_data['input_dir'], len(data_val)))
-
-    # Set up the device and workers.
-    use_cuda = args.use_cuda and torch.cuda.is_available()
-    if use_cuda:
-        gpu_ids = args.gpu_ids if args.gpu_ids is not None\
-                        else range(torch.cuda.device_count())
-        device_ids = [_ for _ in gpu_ids]
-        data_parallel = len(device_ids) > 1
-        device = 'cuda:%d' % device_ids[0]
-        torch.cuda.set_device(device_ids[0])
-        logging.info("Using CUDA devices: %s" % str(device_ids))
-    else:
-        data_parallel = False
-        device = torch.device('cpu')
-        logging.info("Using device: CPU")
-
-    # Set multiprocessing params
-    num_workers = min(multiprocessing.cpu_count(), args.n_workers)
-    kwargs = {
-        'num_workers': num_workers,
-        'pin_memory': True
-    } if use_cuda else {}
-
-    # Set up data loaders
-    #print(args.batch_size, args.eval_batch_size)
-    train_loader = torch.utils.data.DataLoader(data_train,
-                                               batch_size=args.batch_size,
-                                               shuffle=True, **kwargs)
-    val_loader = torch.utils.data.DataLoader(data_val,
-                                             batch_size=args.eval_batch_size,
-                                             **kwargs)
-
-    # Set up model
-    model = network.Net(**args.model_params)
-
-    # Add graph to tensorboard with example train samples
-    # _mixed, _label, _ = next(iter(val_loader))
-    # args.writer.add_graph(model, (_mixed, _label))
-
-    if use_cuda and data_parallel:
-        model = nn.DataParallel(model, device_ids=device_ids)
-        logging.info("Using data parallel model")
-    model.to(device)
-
-    # Set up the optimizer
-    logging.info("Initializing optimizer with %s" % str(args.optim))
-    optimizer = network.optimizer(model, **args.optim, data_parallel=data_parallel)
-    logging.info('Learning rates initialized to:' + utils.format_lr_info(optimizer))
-
-    lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(
-        optimizer, **args.lr_sched)
-    logging.info("Initialized LR scheduler with params: fix_lr_epochs=%d %s"
-                 % (args.fix_lr_epochs, str(args.lr_sched)))
-
-    base_metric = args.base_metric
-    train_metrics = {}
-    val_metrics = {}
-
-    # Load the model if `args.start_epoch` is greater than 0. This will load the
-    # model from epoch = `args.start_epoch - 1`
-    assert args.start_epoch >=0, "start_epoch must be greater than 0."
-    if args.start_epoch > 0:
-        checkpoint_path = os.path.join(args.exp_dir,
-                                       '%d.pt' % (args.start_epoch - 1))
-        _, train_metrics, val_metrics = utils.load_checkpoint(
-            checkpoint_path, model, optim=optimizer, lr_sched=lr_scheduler,
-            data_parallel=data_parallel)
-        logging.info("Loaded checkpoint from %s" % checkpoint_path)
-        logging.info("Learning rates restored to:" + utils.format_lr_info(optimizer))
-
-    # Training loop
-    try:
-        torch.autograd.set_detect_anomaly(args.detect_anomaly)
-        for epoch in range(args.start_epoch, args.epochs + 1):
-            logging.info("Epoch %d:" % epoch)
-            checkpoint_file = os.path.join(args.exp_dir, '%d.pt' % epoch)
-            assert not os.path.exists(checkpoint_file), \
-                "Checkpoint file %s already exists" % checkpoint_file
-            #print("---- begin trianivg")
-            curr_train_metrics = train_epoch(model, device, optimizer,
-                                             train_loader, args.n_train_items,
-                                             epoch=epoch, writer=args.writer,
-                                             data_params=args.train_data)
-            #raise KeyboardInterrupt
-            curr_test_metrics = test_epoch(model, device, val_loader,
-                                           args.n_test_items, network.loss,
-                                           network.metrics, epoch=epoch,
-                                           writer=args.writer,
-                                           data_params=args.val_data)
-            # LR scheduler
-            if epoch >= args.fix_lr_epochs:
-                lr_scheduler.step(curr_test_metrics[base_metric])
-                logging.info(
-                    "LR after scheduling step: %s" %
-                    [_['lr'] for _ in optimizer.param_groups])
-
-            # Write metrics to tensorboard
-            args.writer.add_scalars('Train', curr_train_metrics, epoch)
-            args.writer.add_scalars('Val', curr_test_metrics, epoch)
-            args.writer.flush()
-
-            for k in curr_train_metrics.keys():
-                if not k in train_metrics:
-                    train_metrics[k] = [curr_train_metrics[k]]
-                else:
-                    train_metrics[k].append(curr_train_metrics[k])
-
-            for k in curr_test_metrics.keys():
-                if not k in val_metrics:
-                    val_metrics[k] = [curr_test_metrics[k]]
-                else:
-                    val_metrics[k].append(curr_test_metrics[k])
-
-            if max(val_metrics[base_metric]) == val_metrics[base_metric][-1]:
-                logging.info("Found best validation %s!" % base_metric)
-
-            utils.save_checkpoint(
-                checkpoint_file, epoch, model, optimizer, lr_scheduler,
-                train_metrics, val_metrics, data_parallel)
-            logging.info("Saved checkpoint at %s" % checkpoint_file)
-
-            utils.save_graph(train_metrics, val_metrics, args.exp_dir)
-
-        return train_metrics, val_metrics
-
-
-    except KeyboardInterrupt:
-        print("Interrupted")
-    except Exception as _:  # pylint: disable=broad-except
-        import traceback  # pylint: disable=import-outside-toplevel
-        traceback.print_exc()
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    # Data Params
-    parser.add_argument('exp_dir', type=str,
-                        default='./experiments/fsd_mask_label_mult',
-                        help="Path to save checkpoints and logs.")
-
-    parser.add_argument('--n_train_items', type=int, default=None,
-                        help="Number of items to train on in each epoch")
-    parser.add_argument('--n_test_items', type=int, default=None,
-                        help="Number of items to test.")
-    parser.add_argument('--start_epoch', type=int, default=0,
-                        help="Start epoch")
-    parser.add_argument('--pretrain_path', type=str,
-                        help="Path to pretrained weights")
-    parser.add_argument('--use_cuda', dest='use_cuda', action='store_true',
-                        help="Whether to use cuda")
-    parser.add_argument('--gpu_ids', nargs='+', type=int, default=None,
-                        help="List of GPU ids used for training. "
-                        "Eg., --gpu_ids 2 4. All GPUs are used by default.")
-    parser.add_argument('--detect_anomaly', dest='detect_anomaly',
-                        action='store_true',
-                        help="Whether to use cuda")
-    parser.add_argument('--wandb', dest='wandb', action='store_true',
-                        help="Whether to sync tensorboard to wandb")
-
-    args = parser.parse_args()
-
-    # Set the random seed for reproducible experiments
-    torch.manual_seed(230)
-    random.seed(230)
-    np.random.seed(230)
-    if args.use_cuda:
-        torch.cuda.manual_seed(230)
-
-    # Set up checkpoints
-    if not os.path.exists(args.exp_dir):
-        os.makedirs(args.exp_dir)
-
-    utils.set_logger(os.path.join(args.exp_dir, 'train.log'))
-
-    # Load model and training params
-    params = utils.Params(os.path.join(args.exp_dir, 'config.json'))
-    for k, v in params.__dict__.items():
-        vars(args)[k] = v
-
-    # Initialize tensorboard writer
-    tensorboard_dir = os.path.join(args.exp_dir, 'tensorboard')
-    args.writer = SummaryWriter(tensorboard_dir, purge_step=args.start_epoch)
-    if args.wandb:
-        import wandb
-        wandb.init(
-            project='Semaudio', sync_tensorboard=True,
-            dir=tensorboard_dir, name=os.path.basename(args.exp_dir))
-
-    exec("import %s as network" % args.model)
-    logging.info("Imported the model from '%s'." % args.model)
-
-    train(args)
-
-    args.writer.close()
-    if args.wandb:
-        wandb.finish()