def test(ctx: dict):
ctx['load_best_model'] = True
ctx['net_wrapper'] = build_model(ctx)
ctx['num_mix'] = 1
dataset = MUSICMixDataset(get_ctx(ctx, 'list_test'), ctx, max_sample=get_ctx(ctx, 'num_test'), split='test')
ctx['loader_test'] = torch.utils.data.DataLoader(dataset, batch_size=get_ctx(ctx, 'batch_size'),
shuffle=False, num_workers=2, drop_last=False)
with torch.set_grad_enabled(False):
_test(ctx)
def predict_masks(audio_fname, frame_low_path):
ctx['load_best_model'] = True
ctx['net_wrapper'] = build_model(ctx)
ctx['num_mix'] = 1
sample = [[audio_fname, frame_low_path, len(os.listdir(frame_low_path))]]
dataset = MUSICMixDataset(sample, ctx, max_sample=1, split='regions')
loader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=1)
with torch.no_grad():
data = next(iter(loader))
output = ctx['net_wrapper'].forward(data, ctx, pixelwise=True)
return data, output
def train(ctx: dict):
ctx['net_wrapper'] = build_model(ctx)
ctx['optimizer'] = create_optimizer(
get_underlying_nets(get_ctx(ctx, 'net_wrapper')), ctx)
dataset_train = MUSICMixDataset(get_ctx(ctx, 'list_train'),
ctx,
split='train')
ctx['loader_train'] = torch.utils.data.DataLoader(
dataset_train,
batch_size=get_ctx(ctx, 'batch_size'),
shuffle=True,
num_workers=int(get_ctx(ctx, 'workers')),
drop_last=True)
ctx['epoch_iters'] = len(dataset_train) // get_ctx(ctx, 'batch_size')
print(f'1 Epoch = {get_ctx(ctx, "epoch_iters")} iters')
dataset_val = MUSICMixDataset(get_ctx(ctx, 'list_val'),
ctx,
max_sample=get_ctx(ctx, 'num_val'),
split='val')
ctx['loader_val'] = torch.utils.data.DataLoader(dataset_val,
batch_size=get_ctx(
ctx, 'batch_size'),
shuffle=False,
num_workers=2,
drop_last=False)
ctx['history'], from_epoch = init_history(ctx)
if get_ctx(ctx, 'continue_training') == '':
makedirs(get_ctx(ctx, 'path'), remove=True)
for epoch in range(from_epoch, get_ctx(ctx, 'num_epoch') + 1):
ctx['epoch'] = epoch
with torch.set_grad_enabled(True):
train_epoch(ctx)
with torch.set_grad_enabled(False):
if epoch % get_ctx(ctx, 'eval_epoch') == 0:
_evaluate(ctx)
checkpoint(ctx)
# drop learning rate
if epoch in get_ctx(ctx, 'lr_steps'):
adjust_learning_rate(ctx)
print('Training Done!')
def evaluate(ctx: dict):
ctx['load_best_model'] = True
ctx['net_wrapper'] = build_model(ctx)
dataset_val = MUSICMixDataset(get_ctx(ctx, 'list_val'), ctx, max_sample=get_ctx(ctx, 'num_val'), split='val')
ctx['loader_val'] = torch.utils.data.DataLoader(
dataset_val,
batch_size=get_ctx(ctx, 'batch_size'),
shuffle=False,
num_workers=2,
drop_last=False)
ctx['history'], _ = init_history(None)
ctx['epoch'] = 0
with torch.set_grad_enabled(False):
_evaluate(ctx)
def lr_range_test_for_part(ctx: dict, groups: List[int]):
ctx['net_wrapper'] = build_model(ctx)
ctx['optimizer'] = create_optimizer(get_underlying_nets(get_ctx(ctx, 'net_wrapper')), ctx)
dataset_train = MUSICMixDataset(get_ctx(ctx, 'list_train'), ctx, split='train')
ctx['loader_train'] = torch.utils.data.DataLoader(
dataset_train,
batch_size=get_ctx(ctx, 'batch_size'),
shuffle=True,
num_workers=int(get_ctx(ctx, 'workers')),
drop_last=True)
rates = []
losses = []
net_wrapper = get_ctx(ctx, 'net_wrapper')
optimizer = get_ctx(ctx, 'optimizer')
loader = get_ctx(ctx, 'loader_train')
total = 1000
min_lr = 1e-10
max_lr = 1e1
smooth_f = 0.05
for step, batch_data in tqdm(enumerate(loader), total=total):
if step == total:
break
it = step / total
lr = np.exp((1 - it) * np.log(min_lr) + it * np.log(max_lr))
set_lr(optimizer, lr, groups)
net_wrapper.zero_grad()
loss, _ = net_wrapper.forward(batch_data, ctx)
loss = loss.mean()
# backward
loss.backward()
optimizer.step()
if step > 0:
loss = smooth_f * loss + (1 - smooth_f) * losses[-1]
rates.append(lr)
losses.append(loss.item())
return np.array(rates), np.array(losses)
def regions(ctx: dict):
ctx['load_best_model'] = True
ctx['net_wrapper'] = build_model(ctx)
ctx['num_mix'] = 1
dataset = MUSICMixDataset(get_ctx(ctx, 'list_regions'), ctx, max_sample=get_ctx(ctx, 'num_regions'),
split='regions')
loader = torch.utils.data.DataLoader(dataset, batch_size=get_ctx(ctx, 'batch_size'), shuffle=True,
num_workers=1, drop_last=False)
makedirs(get_ctx(ctx, 'vis_regions'), remove=True)
cnt = 0
with torch.no_grad():
for data in tqdm(loader):
output = ctx['net_wrapper'].forward(data, ctx, pixelwise=True)
output_predictions(ctx, data, output)
cnt += len(data['audios'][0])