# coding: utf-8 __author__ = 'Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/' __version__ = '1.0.3' # Read more here: # https://huggingface.co/docs/accelerate/index import argparse import soundfile as sf import numpy as np import time import glob from tqdm.auto import tqdm import os import torch import wandb import auraloss import torch.nn as nn from torch.optim import Adam, AdamW, SGD, RAdam, RMSprop from torch.utils.data import DataLoader from torch.optim.lr_scheduler import ReduceLROnPlateau import torch.nn.functional as F from accelerate import Accelerator from dataset import MSSDataset from utils import get_model_from_config, demix, sdr, prefer_target_instrument from train import masked_loss, manual_seed, load_not_compatible_weights import warnings warnings.filterwarnings("ignore") def valid(model, valid_loader, args, config, device, verbose=False): instruments = prefer_target_instrument(config) all_sdr = dict() for instr in instruments: all_sdr[instr] = [] all_mixtures_path = valid_loader if verbose: all_mixtures_path = tqdm(valid_loader) pbar_dict = {} for path_list in all_mixtures_path: path = path_list[0] mix, sr = sf.read(path) folder = os.path.dirname(path) res = demix(config, model, mix.T, device, model_type=args.model_type) # mix.T for instr in instruments: if instr != 'other' or config.training.other_fix is False: track, sr1 = sf.read(folder + '/{}.wav'.format(instr)) else: # other is actually instrumental track, sr1 = sf.read(folder + '/{}.wav'.format('vocals')) track = mix - track # sf.write("{}.wav".format(instr), res[instr].T, sr, subtype='FLOAT') references = np.expand_dims(track, axis=0) estimates = np.expand_dims(res[instr].T, axis=0) sdr_val = sdr(references, estimates)[0] single_val = torch.from_numpy(np.array([sdr_val])).to(device) all_sdr[instr].append(single_val) pbar_dict['sdr_{}'.format(instr)] = sdr_val if verbose: all_mixtures_path.set_postfix(pbar_dict) return all_sdr class MSSValidationDataset(torch.utils.data.Dataset): def __init__(self, args): all_mixtures_path = [] for valid_path in args.valid_path: part = sorted(glob.glob(valid_path + '/*/mixture.wav')) if len(part) == 0: print('No validation data found in: {}'.format(valid_path)) all_mixtures_path += part self.list_of_files = all_mixtures_path def __len__(self): return len(self.list_of_files) def __getitem__(self, index): return self.list_of_files[index] def train_model(args): accelerator = Accelerator() device = accelerator.device parser = argparse.ArgumentParser() parser.add_argument("--model_type", type=str, default='mdx23c', help="One of mdx23c, htdemucs, segm_models, mel_band_roformer, bs_roformer, swin_upernet, bandit") parser.add_argument("--config_path", type=str, help="path to config file") parser.add_argument("--start_check_point", type=str, default='', help="Initial checkpoint to start training") parser.add_argument("--results_path", type=str, help="path to folder where results will be stored (weights, metadata)") parser.add_argument("--data_path", nargs="+", type=str, help="Dataset data paths. You can provide several folders.") parser.add_argument("--dataset_type", type=int, default=1, help="Dataset type. Must be one of: 1, 2, 3 or 4. Details here: https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/docs/dataset_types.md") parser.add_argument("--valid_path", nargs="+", type=str, help="validation data paths. You can provide several folders.") parser.add_argument("--num_workers", type=int, default=0, help="dataloader num_workers") parser.add_argument("--pin_memory", type=bool, default=False, help="dataloader pin_memory") parser.add_argument("--seed", type=int, default=0, help="random seed") parser.add_argument("--device_ids", nargs='+', type=int, default=[0], help='list of gpu ids') parser.add_argument("--use_multistft_loss", action='store_true', help="Use MultiSTFT Loss (from auraloss package)") parser.add_argument("--use_mse_loss", action='store_true', help="Use default MSE loss") parser.add_argument("--use_l1_loss", action='store_true', help="Use L1 loss") parser.add_argument("--wandb_key", type=str, default='', help='wandb API Key') parser.add_argument("--pre_valid", action='store_true', help='Run validation before training') if args is None: args = parser.parse_args() else: args = parser.parse_args(args) manual_seed(args.seed + int(time.time())) # torch.backends.cudnn.benchmark = True torch.backends.cudnn.deterministic = False # Fix possible slow down with dilation convolutions torch.multiprocessing.set_start_method('spawn') model, config = get_model_from_config(args.model_type, args.config_path) accelerator.print("Instruments: {}".format(config.training.instruments)) os.makedirs(args.results_path, exist_ok=True) device_ids = args.device_ids batch_size = config.training.batch_size # wandb if accelerator.is_main_process and args.wandb_key is not None and args.wandb_key.strip() != '': wandb.login(key = args.wandb_key) wandb.init(project = 'msst-accelerate', config = { 'config': config, 'args': args, 'device_ids': device_ids, 'batch_size': batch_size }) else: wandb.init(mode = 'disabled') # Fix for num of steps config.training.num_steps *= accelerator.num_processes trainset = MSSDataset( config, args.data_path, batch_size=batch_size, metadata_path=os.path.join(args.results_path, 'metadata_{}.pkl'.format(args.dataset_type)), dataset_type=args.dataset_type, verbose=accelerator.is_main_process, ) train_loader = DataLoader( trainset, batch_size=batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=args.pin_memory ) validset = MSSValidationDataset(args) valid_dataset_length = len(validset) valid_loader = DataLoader( validset, batch_size=1, shuffle=False, ) valid_loader = accelerator.prepare(valid_loader) if args.start_check_point != '': accelerator.print('Start from checkpoint: {}'.format(args.start_check_point)) if 1: load_not_compatible_weights(model, args.start_check_point, verbose=False) else: model.load_state_dict( torch.load(args.start_check_point) ) optim_params = dict() if 'optimizer' in config: optim_params = dict(config['optimizer']) accelerator.print('Optimizer params from config:\n{}'.format(optim_params)) if config.training.optimizer == 'adam': optimizer = Adam(model.parameters(), lr=config.training.lr, **optim_params) elif config.training.optimizer == 'adamw': optimizer = AdamW(model.parameters(), lr=config.training.lr, **optim_params) elif config.training.optimizer == 'radam': optimizer = RAdam(model.parameters(), lr=config.training.lr, **optim_params) elif config.training.optimizer == 'rmsprop': optimizer = RMSprop(model.parameters(), lr=config.training.lr, **optim_params) elif config.training.optimizer == 'prodigy': from prodigyopt import Prodigy # you can choose weight decay value based on your problem, 0 by default # We recommend using lr=1.0 (default) for all networks. optimizer = Prodigy(model.parameters(), lr=config.training.lr, **optim_params) elif config.training.optimizer == 'adamw8bit': import bitsandbytes as bnb optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=config.training.lr, **optim_params) elif config.training.optimizer == 'sgd': accelerator.print('Use SGD optimizer') optimizer = SGD(model.parameters(), lr=config.training.lr, **optim_params) else: accelerator.print('Unknown optimizer: {}'.format(config.training.optimizer)) exit() if accelerator.is_main_process: print('Processes GPU: {}'.format(accelerator.num_processes)) print("Patience: {} Reduce factor: {} Batch size: {} Optimizer: {}".format( config.training.patience, config.training.reduce_factor, batch_size, config.training.optimizer, )) # Reduce LR if no SDR improvements for several epochs scheduler = ReduceLROnPlateau( optimizer, 'max', # patience=accelerator.num_processes * config.training.patience, # This is strange place... patience=config.training.patience, factor=config.training.reduce_factor ) if args.use_multistft_loss: try: loss_options = dict(config.loss_multistft) except: loss_options = dict() accelerator.print('Loss options: {}'.format(loss_options)) loss_multistft = auraloss.freq.MultiResolutionSTFTLoss( **loss_options ) model, optimizer, train_loader, scheduler = accelerator.prepare(model, optimizer, train_loader, scheduler) if args.pre_valid: sdr_list = valid(model, valid_loader, args, config, device, verbose=accelerator.is_main_process) sdr_list = accelerator.gather(sdr_list) accelerator.wait_for_everyone() # print(sdr_list) sdr_avg = 0.0 instruments = prefer_target_instrument(config) for instr in instruments: # print(sdr_list[instr]) sdr_data = torch.cat(sdr_list[instr], dim=0).cpu().numpy() sdr_val = sdr_data.mean() accelerator.print("Valid length: {}".format(valid_dataset_length)) accelerator.print("Instr SDR {}: {:.4f} Debug: {}".format(instr, sdr_val, len(sdr_data))) sdr_val = sdr_data[:valid_dataset_length].mean() accelerator.print("Instr SDR {}: {:.4f} Debug: {}".format(instr, sdr_val, len(sdr_data))) sdr_avg += sdr_val sdr_avg /= len(instruments) if len(instruments) > 1: accelerator.print('SDR Avg: {:.4f}'.format(sdr_avg)) sdr_list = None accelerator.print('Train for: {}'.format(config.training.num_epochs)) best_sdr = -100 for epoch in range(config.training.num_epochs): model.train().to(device) accelerator.print('Train epoch: {} Learning rate: {}'.format(epoch, optimizer.param_groups[0]['lr'])) loss_val = 0. total = 0 pbar = tqdm(train_loader, disable=not accelerator.is_main_process) for i, (batch, mixes) in enumerate(pbar): y = batch x = mixes if args.model_type in ['mel_band_roformer', 'bs_roformer']: # loss is computed in forward pass loss = model(x, y) else: y_ = model(x) if args.use_multistft_loss: y1_ = torch.reshape(y_, (y_.shape[0], y_.shape[1] * y_.shape[2], y_.shape[3])) y1 = torch.reshape(y, (y.shape[0], y.shape[1] * y.shape[2], y.shape[3])) loss = loss_multistft(y1_, y1) # We can use many losses at the same time if args.use_mse_loss: loss += 1000 * nn.MSELoss()(y1_, y1) if args.use_l1_loss: loss += 1000 * F.l1_loss(y1_, y1) elif args.use_mse_loss: loss = nn.MSELoss()(y_, y) elif args.use_l1_loss: loss = F.l1_loss(y_, y) else: loss = masked_loss( y_, y, q=config.training.q, coarse=config.training.coarse_loss_clip ) accelerator.backward(loss) if config.training.grad_clip: accelerator.clip_grad_norm_(model.parameters(), config.training.grad_clip) optimizer.step() optimizer.zero_grad() li = loss.item() loss_val += li total += 1 if accelerator.is_main_process: wandb.log({'loss': 100 * li, 'avg_loss': 100 * loss_val / (i + 1), 'total': total, 'loss_val': loss_val, 'i': i }) pbar.set_postfix({'loss': 100 * li, 'avg_loss': 100 * loss_val / (i + 1)}) if accelerator.is_main_process: print('Training loss: {:.6f}'.format(loss_val / total)) wandb.log({'train_loss': loss_val / total, 'epoch': epoch}) # Save last store_path = args.results_path + '/last_{}.ckpt'.format(args.model_type) accelerator.wait_for_everyone() if accelerator.is_main_process: unwrapped_model = accelerator.unwrap_model(model) accelerator.save(unwrapped_model.state_dict(), store_path) sdr_list = valid(model, valid_loader, args, config, device, verbose=accelerator.is_main_process) sdr_list = accelerator.gather(sdr_list) accelerator.wait_for_everyone() sdr_avg = 0.0 instruments = prefer_target_instrument(config) for instr in instruments: if accelerator.is_main_process and 0: print(sdr_list[instr]) sdr_data = torch.cat(sdr_list[instr], dim=0).cpu().numpy() # sdr_val = sdr_data.mean() sdr_val = sdr_data[:valid_dataset_length].mean() if accelerator.is_main_process: print("Instr SDR {}: {:.4f} Debug: {}".format(instr, sdr_val, len(sdr_data))) wandb.log({ f'{instr}_sdr': sdr_val }) sdr_avg += sdr_val sdr_avg /= len(instruments) if len(instruments) > 1: if accelerator.is_main_process: print('SDR Avg: {:.4f}'.format(sdr_avg)) wandb.log({'sdr_avg': sdr_avg, 'best_sdr': best_sdr}) if accelerator.is_main_process: if sdr_avg > best_sdr: store_path = args.results_path + '/model_{}_ep_{}_sdr_{:.4f}.ckpt'.format(args.model_type, epoch, sdr_avg) print('Store weights: {}'.format(store_path)) unwrapped_model = accelerator.unwrap_model(model) accelerator.save(unwrapped_model.state_dict(), store_path) best_sdr = sdr_avg scheduler.step(sdr_avg) sdr_list = None accelerator.wait_for_everyone() if __name__ == "__main__": train_model(None)