demucs/__main__.py

# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import json
import math
import os
import sys
import time
from dataclasses import dataclass, field

import torch as th
from torch import distributed, nn
from torch.nn.parallel.distributed import DistributedDataParallel

from .augment import FlipChannels, FlipSign, Remix, Scale, Shift
from .compressed import get_compressed_datasets
from .model import Demucs
from .parser import get_name, get_parser
from .raw import Rawset
from .repitch import RepitchedWrapper
from .pretrained import load_pretrained, SOURCES
from .tasnet import ConvTasNet
from .test import evaluate
from .train import train_model, validate_model
from .utils import (human_seconds, load_model, save_model, get_state,
                    save_state, sizeof_fmt, get_quantizer)
from .wav import get_wav_datasets, get_musdb_wav_datasets


@dataclass
class SavedState:
    metrics: list = field(default_factory=list)
    last_state: dict = None
    best_state: dict = None
    optimizer: dict = None


def main():
    parser = get_parser()
    args = parser.parse_args()
    name = get_name(parser, args)
    print(f"Experiment {name}")

    if args.musdb is None and args.rank == 0:
        print(
            "You must provide the path to the MusDB dataset with the --musdb flag. "
            "To download the MusDB dataset, see https://sigsep.github.io/datasets/musdb.html.",
            file=sys.stderr)
        sys.exit(1)

    eval_folder = args.evals / name
    eval_folder.mkdir(exist_ok=True, parents=True)
    args.logs.mkdir(exist_ok=True)
    metrics_path = args.logs / f"{name}.json"
    eval_folder.mkdir(exist_ok=True, parents=True)
    args.checkpoints.mkdir(exist_ok=True, parents=True)
    args.models.mkdir(exist_ok=True, parents=True)

    if args.device is None:
        device = "cpu"
        if th.cuda.is_available():
            device = "cuda"
    else:
        device = args.device

    th.manual_seed(args.seed)
    # Prevents too many threads to be started when running `museval` as it can be quite
    # inefficient on NUMA architectures.
    os.environ["OMP_NUM_THREADS"] = "1"
    os.environ["MKL_NUM_THREADS"] = "1"

    if args.world_size > 1:
        if device != "cuda" and args.rank == 0:
            print("Error: distributed training is only available with cuda device", file=sys.stderr)
            sys.exit(1)
        th.cuda.set_device(args.rank % th.cuda.device_count())
        distributed.init_process_group(backend="nccl",
                                       init_method="tcp://" + args.master,
                                       rank=args.rank,
                                       world_size=args.world_size)

    checkpoint = args.checkpoints / f"{name}.th"
    checkpoint_tmp = args.checkpoints / f"{name}.th.tmp"
    if args.restart and checkpoint.exists() and args.rank == 0:
        checkpoint.unlink()

    if args.test or args.test_pretrained:
        args.epochs = 1
        args.repeat = 0
        if args.test:
            model = load_model(args.models / args.test)
        else:
            model = load_pretrained(args.test_pretrained)
    elif args.tasnet:
        model = ConvTasNet(audio_channels=args.audio_channels,
                           samplerate=args.samplerate, X=args.X,
                           segment_length=4 * args.samples,
                           sources=SOURCES)
    else:
        model = Demucs(
            audio_channels=args.audio_channels,
            channels=args.channels,
            context=args.context,
            depth=args.depth,
            glu=args.glu,
            growth=args.growth,
            kernel_size=args.kernel_size,
            lstm_layers=args.lstm_layers,
            rescale=args.rescale,
            rewrite=args.rewrite,
            stride=args.conv_stride,
            resample=args.resample,
            normalize=args.normalize,
            samplerate=args.samplerate,
            segment_length=4 * args.samples,
            sources=SOURCES,
        )
    model.to(device)
    if args.init:
        model.load_state_dict(load_pretrained(args.init).state_dict())

    if args.show:
        print(model)
        size = sizeof_fmt(4 * sum(p.numel() for p in model.parameters()))
        print(f"Model size {size}")
        return

    try:
        saved = th.load(checkpoint, map_location='cpu')
    except IOError:
        saved = SavedState()

    optimizer = th.optim.Adam(model.parameters(), lr=args.lr)

    quantizer = None
    quantizer = get_quantizer(model, args, optimizer)

    if saved.last_state is not None:
        model.load_state_dict(saved.last_state, strict=False)
    if saved.optimizer is not None:
        optimizer.load_state_dict(saved.optimizer)

    model_name = f"{name}.th"
    if args.save_model:
        if args.rank == 0:
            model.to("cpu")
            model.load_state_dict(saved.best_state)
            save_model(model, quantizer, args, args.models / model_name)
        return
    elif args.save_state:
        model_name = f"{args.save_state}.th"
        if args.rank == 0:
            model.to("cpu")
            model.load_state_dict(saved.best_state)
            state = get_state(model, quantizer)
            save_state(state, args.models / model_name)
        return

    if args.rank == 0:
        done = args.logs / f"{name}.done"
        if done.exists():
            done.unlink()

    augment = [Shift(args.data_stride)]
    if args.augment:
        augment += [FlipSign(), FlipChannels(), Scale(),
                    Remix(group_size=args.remix_group_size)]
    augment = nn.Sequential(*augment).to(device)
    print("Agumentation pipeline:", augment)

    if args.mse:
        criterion = nn.MSELoss()
    else:
        criterion = nn.L1Loss()

    # Setting number of samples so that all convolution windows are full.
    # Prevents hard to debug mistake with the prediction being shifted compared
    # to the input mixture.
    samples = model.valid_length(args.samples)
    print(f"Number of training samples adjusted to {samples}")
    samples = samples + args.data_stride
    if args.repitch:
        # We need a bit more audio samples, to account for potential
        # tempo change.
        samples = math.ceil(samples / (1 - 0.01 * args.max_tempo))

    args.metadata.mkdir(exist_ok=True, parents=True)
    if args.raw:
        train_set = Rawset(args.raw / "train",
                           samples=samples,
                           channels=args.audio_channels,
                           streams=range(1, len(model.sources) + 1),
                           stride=args.data_stride)

        valid_set = Rawset(args.raw / "valid", channels=args.audio_channels)
    elif args.wav:
        train_set, valid_set = get_wav_datasets(args, samples, model.sources)
    elif args.is_wav:
        train_set, valid_set = get_musdb_wav_datasets(args, samples, model.sources)
    else:
        train_set, valid_set = get_compressed_datasets(args, samples)

    if args.repitch:
        train_set = RepitchedWrapper(
            train_set,
            proba=args.repitch,
            max_tempo=args.max_tempo)

    best_loss = float("inf")
    for epoch, metrics in enumerate(saved.metrics):
        print(f"Epoch {epoch:03d}: "
              f"train={metrics['train']:.8f} "
              f"valid={metrics['valid']:.8f} "
              f"best={metrics['best']:.4f} "
              f"ms={metrics.get('true_model_size', 0):.2f}MB "
              f"cms={metrics.get('compressed_model_size', 0):.2f}MB "
              f"duration={human_seconds(metrics['duration'])}")
        best_loss = metrics['best']

    if args.world_size > 1:
        dmodel = DistributedDataParallel(model,
                                         device_ids=[th.cuda.current_device()],
                                         output_device=th.cuda.current_device())
    else:
        dmodel = model

    for epoch in range(len(saved.metrics), args.epochs):
        begin = time.time()
        model.train()
        train_loss, model_size = train_model(
            epoch, train_set, dmodel, criterion, optimizer, augment,
            quantizer=quantizer,
            batch_size=args.batch_size,
            device=device,
            repeat=args.repeat,
            seed=args.seed,
            diffq=args.diffq,
            workers=args.workers,
            world_size=args.world_size)
        model.eval()
        valid_loss = validate_model(
            epoch, valid_set, model, criterion,
            device=device,
            rank=args.rank,
            split=args.split_valid,
            overlap=args.overlap,
            world_size=args.world_size)

        ms = 0
        cms = 0
        if quantizer and args.rank == 0:
            ms = quantizer.true_model_size()
            cms = quantizer.compressed_model_size(num_workers=min(40, args.world_size * 10))

        duration = time.time() - begin
        if valid_loss < best_loss and ms <= args.ms_target:
            best_loss = valid_loss
            saved.best_state = {
                key: value.to("cpu").clone()
                for key, value in model.state_dict().items()
            }

        saved.metrics.append({
            "train": train_loss,
            "valid": valid_loss,
            "best": best_loss,
            "duration": duration,
            "model_size": model_size,
            "true_model_size": ms,
            "compressed_model_size": cms,
        })
        if args.rank == 0:
            json.dump(saved.metrics, open(metrics_path, "w"))

        saved.last_state = model.state_dict()
        saved.optimizer = optimizer.state_dict()
        if args.rank == 0 and not args.test:
            th.save(saved, checkpoint_tmp)
            checkpoint_tmp.rename(checkpoint)

        print(f"Epoch {epoch:03d}: "
              f"train={train_loss:.8f} valid={valid_loss:.8f} best={best_loss:.4f} ms={ms:.2f}MB "
              f"cms={cms:.2f}MB "
              f"duration={human_seconds(duration)}")

    if args.world_size > 1:
        distributed.barrier()

    del dmodel
    model.load_state_dict(saved.best_state)
    if args.eval_cpu:
        device = "cpu"
        model.to(device)
    model.eval()
    evaluate(model, args.musdb, eval_folder,
             is_wav=args.is_wav,
             rank=args.rank,
             world_size=args.world_size,
             device=device,
             save=args.save,
             split=args.split_valid,
             shifts=args.shifts,
             overlap=args.overlap,
             workers=args.eval_workers)
    model.to("cpu")
    if args.rank == 0:
        if not (args.test or args.test_pretrained):
            save_model(model, quantizer, args, args.models / model_name)
        print("done")
        done.write_text("done")


if __name__ == "__main__":
    main()