mirror of
https://github.com/Anjok07/ultimatevocalremovergui.git
synced 2024-09-24 11:38:37 +02:00
224 lines
8.2 KiB
Python
224 lines
8.2 KiB
Python
import argparse
|
|
from datetime import datetime as dt
|
|
import gc
|
|
import json
|
|
import os
|
|
import random
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from lib import dataset
|
|
from lib import nets
|
|
from lib import spec_utils
|
|
|
|
|
|
def train_val_split(mix_dir, inst_dir, val_rate, val_filelist_json):
|
|
input_exts = ['.wav', '.m4a', '.3gp', '.oma', '.mp3', '.mp4']
|
|
X_list = sorted([
|
|
os.path.join(mix_dir, fname)
|
|
for fname in os.listdir(mix_dir)
|
|
if os.path.splitext(fname)[1] in input_exts])
|
|
y_list = sorted([
|
|
os.path.join(inst_dir, fname)
|
|
for fname in os.listdir(inst_dir)
|
|
if os.path.splitext(fname)[1] in input_exts])
|
|
|
|
filelist = list(zip(X_list, y_list))
|
|
random.shuffle(filelist)
|
|
|
|
val_filelist = []
|
|
if val_filelist_json is not None:
|
|
with open(val_filelist_json, 'r', encoding='utf8') as f:
|
|
val_filelist = json.load(f)
|
|
|
|
if len(val_filelist) == 0:
|
|
val_size = int(len(filelist) * val_rate)
|
|
train_filelist = filelist[:-val_size]
|
|
val_filelist = filelist[-val_size:]
|
|
else:
|
|
train_filelist = [
|
|
pair for pair in filelist
|
|
if list(pair) not in val_filelist]
|
|
|
|
return train_filelist, val_filelist
|
|
|
|
|
|
def train_inner_epoch(X_train, y_train, model, optimizer, batchsize, instance_loss):
|
|
sum_loss = 0
|
|
model.train()
|
|
aux_crit = nn.L1Loss()
|
|
criterion = nn.L1Loss(reduction='none')
|
|
perm = np.random.permutation(len(X_train))
|
|
for i in range(0, len(X_train), batchsize):
|
|
local_perm = perm[i: i + batchsize]
|
|
X_batch = torch.from_numpy(X_train[local_perm]).cpu()
|
|
y_batch = torch.from_numpy(y_train[local_perm]).cpu()
|
|
|
|
model.zero_grad()
|
|
mask, aux = model(X_batch)
|
|
|
|
aux_loss = aux_crit(X_batch * aux, y_batch)
|
|
X_batch = spec_utils.crop_center(mask, X_batch, False)
|
|
y_batch = spec_utils.crop_center(mask, y_batch, False)
|
|
abs_diff = criterion(X_batch * mask, y_batch)
|
|
|
|
loss = abs_diff.mean() * 0.9 + aux_loss * 0.1
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
abs_diff_np = abs_diff.detach().cpu().numpy()
|
|
instance_loss[local_perm] += abs_diff_np.mean(axis=(1, 2, 3))
|
|
sum_loss += float(loss.detach().cpu().numpy()) * len(X_batch)
|
|
|
|
return sum_loss / len(X_train)
|
|
|
|
|
|
def val_inner_epoch(dataloader, model):
|
|
sum_loss = 0
|
|
model.eval()
|
|
criterion = nn.L1Loss()
|
|
with torch.no_grad():
|
|
for X_batch, y_batch in dataloader:
|
|
X_batch = X_batch.cpu()
|
|
y_batch = y_batch.cpu()
|
|
mask = model.predict(X_batch)
|
|
X_batch = spec_utils.crop_center(mask, X_batch, False)
|
|
y_batch = spec_utils.crop_center(mask, y_batch, False)
|
|
|
|
loss = criterion(X_batch * mask, y_batch)
|
|
sum_loss += float(loss.detach().cpu().numpy()) * len(X_batch)
|
|
|
|
return sum_loss / len(dataloader.dataset)
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument('--gpu', '-g', type=int, default=-1)
|
|
p.add_argument('--seed', '-s', type=int, default=2019)
|
|
p.add_argument('--sr', '-r', type=int, default=44100)
|
|
p.add_argument('--hop_length', '-l', type=int, default=1024)
|
|
p.add_argument('--mixture_dataset', '-m', required=True)
|
|
p.add_argument('--instrumental_dataset', '-i', required=True)
|
|
p.add_argument('--learning_rate', type=float, default=0.001)
|
|
p.add_argument('--lr_min', type=float, default=0.0001)
|
|
p.add_argument('--lr_decay_factor', type=float, default=0.9)
|
|
p.add_argument('--lr_decay_patience', type=int, default=6)
|
|
p.add_argument('--batchsize', '-B', type=int, default=4)
|
|
p.add_argument('--cropsize', '-c', type=int, default=256)
|
|
p.add_argument('--val_rate', '-v', type=float, default=0.1)
|
|
p.add_argument('--val_filelist', '-V', type=str, default=None)
|
|
p.add_argument('--val_batchsize', '-b', type=int, default=4)
|
|
p.add_argument('--val_cropsize', '-C', type=int, default=512)
|
|
p.add_argument('--patches', '-p', type=int, default=16)
|
|
p.add_argument('--epoch', '-E', type=int, default=100)
|
|
p.add_argument('--inner_epoch', '-e', type=int, default=4)
|
|
p.add_argument('--oracle_rate', '-O', type=float, default=0)
|
|
p.add_argument('--oracle_drop_rate', '-o', type=float, default=0.5)
|
|
p.add_argument('--mixup_rate', '-M', type=float, default=0.0)
|
|
p.add_argument('--mixup_alpha', '-a', type=float, default=1.0)
|
|
p.add_argument('--pretrained_model', '-P', type=str, default=None)
|
|
p.add_argument('--debug', '-d', action='store_true')
|
|
args = p.parse_args()
|
|
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
timestamp = dt.now().strftime('%Y%m%d%H%M%S')
|
|
|
|
model = nets.CascadedASPPNet()
|
|
if args.pretrained_model is not None:
|
|
model.load_state_dict(torch.load(args.pretrained_model))
|
|
if args.gpu >= 0:
|
|
model.cuda()
|
|
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
|
|
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
|
optimizer,
|
|
factor=args.lr_decay_factor,
|
|
patience=args.lr_decay_patience,
|
|
min_lr=args.lr_min,
|
|
verbose=True)
|
|
|
|
train_filelist, val_filelist = train_val_split(
|
|
mix_dir=args.mixture_dataset,
|
|
inst_dir=args.instrumental_dataset,
|
|
val_rate=args.val_rate,
|
|
val_filelist_json=args.val_filelist)
|
|
|
|
if args.debug:
|
|
print('### DEBUG MODE')
|
|
train_filelist = train_filelist[:1]
|
|
val_filelist = val_filelist[:1]
|
|
|
|
with open('val_{}.json'.format(timestamp), 'w', encoding='utf8') as f:
|
|
json.dump(val_filelist, f, ensure_ascii=False)
|
|
|
|
for i, (X_fname, y_fname) in enumerate(val_filelist):
|
|
print(i + 1, os.path.basename(X_fname), os.path.basename(y_fname))
|
|
|
|
val_dataset = dataset.make_validation_set(
|
|
filelist=val_filelist,
|
|
cropsize=args.val_cropsize,
|
|
sr=args.sr,
|
|
hop_length=args.hop_length,
|
|
offset=model.offset)
|
|
val_dataloader = torch.utils.data.DataLoader(
|
|
dataset=val_dataset,
|
|
batch_size=args.val_batchsize,
|
|
shuffle=False,
|
|
num_workers=4)
|
|
|
|
log = []
|
|
oracle_X = None
|
|
oracle_y = None
|
|
best_loss = np.inf
|
|
for epoch in range(args.epoch):
|
|
X_train, y_train = dataset.make_training_set(
|
|
train_filelist, args.cropsize, args.patches, args.sr, args.hop_length, model.offset)
|
|
|
|
X_train, y_train = dataset.mixup_generator(
|
|
X_train, y_train, args.mixup_rate, args.mixup_alpha)
|
|
|
|
if oracle_X is not None and oracle_y is not None:
|
|
perm = np.random.permutation(len(oracle_X))
|
|
X_train[perm] = oracle_X
|
|
y_train[perm] = oracle_y
|
|
|
|
print('# epoch', epoch)
|
|
instance_loss = np.zeros(len(X_train), dtype=np.float32)
|
|
for inner_epoch in range(args.inner_epoch):
|
|
print(' * inner epoch {}'.format(inner_epoch))
|
|
train_loss = train_inner_epoch(
|
|
X_train, y_train, model, optimizer, args.batchsize, instance_loss)
|
|
val_loss = val_inner_epoch(val_dataloader, model)
|
|
|
|
print(' * training loss = {:.6f}, validation loss = {:.6f}'
|
|
.format(train_loss * 1000, val_loss * 1000))
|
|
|
|
scheduler.step(val_loss)
|
|
|
|
if val_loss < best_loss:
|
|
best_loss = val_loss
|
|
print(' * best validation loss')
|
|
model_path = 'models/model_iter{}.pth'.format(epoch)
|
|
torch.save(model.state_dict(), model_path)
|
|
|
|
log.append([train_loss, val_loss])
|
|
with open('log_{}.json'.format(timestamp), 'w', encoding='utf8') as f:
|
|
json.dump(log, f, ensure_ascii=False)
|
|
|
|
if args.oracle_rate > 0:
|
|
instance_loss /= args.inner_epoch
|
|
oracle_X, oracle_y, idx = dataset.get_oracle_data(
|
|
X_train, y_train, instance_loss, args.oracle_rate, args.oracle_drop_rate)
|
|
print(' * oracle loss = {:.6f}'.format(instance_loss[idx].mean()))
|
|
|
|
del X_train, y_train
|
|
gc.collect()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|