def validation(model, args, lr, epoch, device):
dataloader, dataset = make_loader(
args.cv_list,
args.batch_size,
num_workers=args.num_threads,
)
model.eval()
loss_total = 0.
sisnr_total = 0.
num_batch = len(dataloader)
stime = time.time()
with torch.no_grad():
for idx, data in enumerate(dataloader):
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
outputs, wav = data_parallel(model, (inputs, ))
loss = model.loss(outputs, labels,mode='Mix')[0]
sisnr = model.loss(wav, labels, mode='SiSNR')
loss_total += loss.data.cpu()
sisnr_total += sisnr.data.cpu()
del loss, data, inputs, labels, wav, outputs
etime = time.time()
eplashed = (etime - stime) / num_batch
loss_total_avg = loss_total / num_batch
sisnr_total_avg = sisnr_total / num_batch
print('CROSSVAL AVG.LOSS | Epoch {:3d}/{:3d} '
'| lr {:.4e} | {:2.3f}s/batch| time {:2.1f}mins '
'| Mixloss {:2.4f} | SiSNR {:2.4f}'.format(epoch + 1, args.max_epoch, lr, eplashed,
(etime - stime)/60.0, loss_total_avg.item(), -sisnr_total_avg.item()))
sys.stdout.flush()
return loss_total_avg, sisnr_total_avg
def validation(model, args, lr, epoch, device):
dataloader, dataset = make_loader(args.cv_list,
args.batch_size,
num_workers=args.num_threads,
sample_rate=args.sample_rate,
segment_length=2,
processer=Processer(
sample_rate=args.sample_rate, ))
model.eval()
num_batch = len(dataloader)
stime = time.time()
all_moniter = Moniter('All', num_batch, 1)
sdr_moniter = Moniter('SDR', num_batch, 1)
speaker_moniter = Moniter('Speaker', num_batch, 1)
reg_moniter = Moniter('Reg', num_batch, 1)
num_batch = len(dataloader)
with torch.no_grad():
for idx, data in enumerate(dataloader):
inputs, labels, spkid = data
inputs = inputs.to(device)
labels = labels.to(device)
spkid = spkid.to(device)
est_wav, speaker_loss, reg_loss = data_parallel(
model, (inputs, spkid))
speaker_loss = torch.mean(speaker_loss)
reg_loss = torch.mean(reg_loss)
#gth_spec = data_parallel(model.stft, (labels))[0]
#loss = model.loss(est_spec, gth_spec, loss_mode='MSE')
sdr = model.loss(est_wav, labels, loss_mode='SDR')
all = sdr + 2 * speaker_loss + 0.3 * reg_loss
all_moniter(all)
sdr_moniter(sdr)
speaker_moniter(speaker_loss)
reg_moniter(reg_loss)
etime = time.time()
eplashed = time.time() - stime
log_str = 'CROSSVAL AVG.LOSS | Epoch {:3d}/{:3d} ' \
'| lr {:.4e} | {:2.3f}s/batch| time {:2.1f}mins |' \
' {:s} |'\
' {:s} |'\
' {:s} |'\
' {:s} |'\
''.format(
epoch + 1,
args.max_epoch,
lr,
eplashed,
(etime - stime)/60.0,
all_moniter.average(),
sdr_moniter.average(),
speaker_moniter.average(),
reg_moniter.average()
)
print(log_str)
sys.stdout.flush()
return all_moniter.ave_float(), sdr_moniter.ave_float()
def validation(model, args, lr, epoch, device):
dataloader, dataset = make_loader(args.cv_list,
args.batch_size,
8,
num_workers=args.num_threads,
processer=Processer())
model.eval()
loss_total = 0.0
sisnr_total = 0.0
num_batch = len(dataloader)
stime = time.time()
with torch.no_grad():
for idx, data in enumerate(dataloader):
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
est_spec, est_wav = data_parallel(model, (inputs, ))
#gth_spec = data_parallel(model.stft, (labels))
#loss = model.loss(est_spec, gth_spec, loss_mode='MSE')
sisnr = model.loss(est_wav, labels, loss_mode='SI-SNR')
loss = sisnr
loss_total += loss.data.cpu()
sisnr_total += sisnr.data.cpu()
etime = time.time()
eplashed = (etime - stime) / num_batch
loss_total_avg = loss_total / num_batch
sisnr_total_avg = sisnr_total / num_batch
print('CROSSVAL AVG.LOSS | Epoch {:3d}/{:3d} '
'| lr {:.4e} | {:2.3f}s/batch| time {:2.1f}mins '
'| loss {:2.6f} |'
'| SISNR {:2.4f} '.format(
epoch + 1,
args.max_epoch,
lr,
eplashed,
(etime - stime) / 60.0,
loss_total_avg,
-sisnr_total_avg,
))
sys.stdout.flush()
return loss_total_avg, sisnr_total_avg
def train(model, args, device, writer):
print('preparing data...')
dataloader, dataset = make_loader(
args.tr_list,
args.batch_size,
num_workers=args.num_threads,
)
print_freq = 100
num_batch = len(dataloader)
params = model.get_params(args.weight_decay)
optimizer = optim.Adam(params, lr=args.learn_rate)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, 'min', factor=0.5, patience=1, verbose=True)
if args.retrain:
start_epoch, step = reload_model(model, optimizer, args.exp_dir,
args.use_cuda)
else:
start_epoch, step = 0, 0
print('---------PRERUN-----------')
lr = get_learning_rate(optimizer)
print('(Initialization)')
val_loss, val_sisnr = validation(model, args, lr, -1, device)
writer.add_scalar('Loss/Train', val_loss, step)
writer.add_scalar('Loss/Cross-Validation', val_loss, step)
writer.add_scalar('SiSNR/Train', -val_sisnr, step)
writer.add_scalar('SiSNR/Cross-Validation', -val_sisnr, step)
for epoch in range(start_epoch, args.max_epoch):
torch.manual_seed(args.seed + epoch)
if args.use_cuda:
torch.cuda.manual_seed(args.seed + epoch)
model.train()
sisnr_total = 0.0
sisnr_print = 0.0
mix_loss_total = 0.0
mix_loss_print = 0.0
amp_loss_total = 0.0
amp_loss_print = 0.0
phase_loss_total = 0.0
phase_loss_print = 0.0
stime = time.time()
lr = get_learning_rate(optimizer)
for idx, data in enumerate(dataloader):
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
model.zero_grad()
outputs, wav = data_parallel(model, (inputs,))
loss = model.loss(outputs, labels, mode='Mix')
loss[0].backward()
nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad_norm)
optimizer.step()
step += 1
sisnr = model.loss(wav, labels, mode='SiSNR')
mix_loss_total += loss[0].data.cpu()
mix_loss_print += loss[0].data.cpu()
amp_loss_total += loss[1].data.cpu()
amp_loss_print += loss[1].data.cpu()
phase_loss_total += loss[2].data.cpu()
phase_loss_print += loss[2].data.cpu()
sisnr_print += sisnr.data.cpu()
sisnr_total += sisnr.data.cpu()
del outputs, labels, inputs, loss, wav
if (idx+1) % 1000 == 0:
save_checkpoint(model, optimizer, -1, step, args.exp_dir)
if (idx + 1) % print_freq == 0:
eplashed = time.time() - stime
speed_avg = eplashed / (idx+1)
mix_loss_print_avg = mix_loss_print / print_freq
amp_loss_print_avg = amp_loss_print / print_freq
phase_loss_print_avg = phase_loss_print / print_freq
sisnr_print_avg = sisnr_print / print_freq
print('Epoch {:3d}/{:3d} | batches {:5d}/{:5d} | lr {:1.4e} |'
'{:2.3f}s/batches '
'| Mixloss {:2.4f}'
'| AMPloss {:2.4f}'
'| Phaseloss {:2.4f}'
'| SiSNR {:2.4f}'
.format(
epoch, args.max_epoch, idx + 1, num_batch, lr,
speed_avg,
mix_loss_print_avg,
amp_loss_print_avg,
phase_loss_print_avg,
-sisnr_print_avg
))
sys.stdout.flush()
writer.add_scalar('SiSNR/Train', -sisnr_print_avg, step)
writer.add_scalar('Loss/Train', mix_loss_print_avg, step)
mix_loss_print = 0.
amp_loss_print = 0.
phase_loss_print = 0.
sisnr_print = 0.
eplashed = time.time() - stime
mix_loss_total_avg = mix_loss_total / num_batch
sisnr_total_avg = sisnr_total / num_batch
print(
'Training AVG.LOSS |'
' Epoch {:3d}/{:3d} | lr {:1.4e} |'
' {:2.3f}s/batch | time {:3.2f}mins |'
' Mixloss {:2.4f}'
' SiSNR {:2.4f}'
.format(
epoch + 1, args.max_epoch,
lr,
eplashed/num_batch,
eplashed/60.0,
mix_loss_total_avg,
-sisnr_total_avg
))
val_loss, val_sisnr = validation(model, args, lr, epoch, device)
writer.add_scalar('Loss/Cross-Validation', val_loss, step)
writer.add_scalar('SiSNR/Cross-Validation', -val_sisnr, step)
writer.add_scalar('learn_rate', lr, step)
if val_loss > scheduler.best:
print('Rejected !!! The best is {:2.6f}'.format(scheduler.best))
else:
save_checkpoint(model, optimizer, epoch + 1, step, args.exp_dir, mode='best_model')
scheduler.step(val_loss)
sys.stdout.flush()
stime = time.time()
def train(model, args, device, writer):
print('preparing data...')
dataloader, dataset = make_loader(args.tr_list,
args.batch_size,
num_workers=args.num_threads,
segment_length=1.5,
sample_rate=args.sample_rate,
processer=Processer(
sample_rate=args.sample_rate, ))
print_freq = 100
num_batch = len(dataloader)
params = model.get_params(args.weight_decay)
optimizer = optim.Adam(params, lr=args.learn_rate)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
'min',
factor=0.5,
patience=1,
verbose=True)
if args.retrain:
start_epoch, step = reload_model(model, optimizer, args.exp_dir,
args.use_cuda)
else:
start_epoch, step = 0, 0
print('---------PRERUN-----------')
lr = get_learning_rate(optimizer)
print('(Initialization)')
val_loss, val_sisnr = 30, 30. #validation(model, args, lr, -1, device)
writer.add_scalar('Loss/Train', val_loss, step)
writer.add_scalar('Loss/Cross-Validation', val_loss, step)
writer.add_scalar('SISNR/Train', -val_sisnr, step)
writer.add_scalar('SISNR/Cross-Validation', -val_sisnr, step)
for epoch in range(start_epoch, args.max_epoch):
torch.manual_seed(args.seed + epoch)
if args.use_cuda:
torch.cuda.manual_seed(args.seed + epoch)
model.train()
all_moniter = Moniter('All', num_batch, print_freq)
sdr_moniter = Moniter('SDR', num_batch, print_freq)
speaker_moniter = Moniter('Speaker', num_batch, print_freq)
reg_moniter = Moniter('Reg', num_batch, print_freq)
stime = time.time()
lr = get_learning_rate(optimizer)
for idx, data in enumerate(dataloader):
inputs, labels, spkid = data
inputs = inputs.to(device)
labels = labels.to(device)
spkid = spkid.to(device)
model.zero_grad()
est_wav, speaker_loss, reg_loss = data_parallel(
model, (inputs, spkid))
speaker_loss = torch.mean(speaker_loss)
reg_loss = torch.mean(reg_loss)
sdr = model.loss(est_wav, labels, loss_mode='SDR')
all = sdr + 2 * speaker_loss + 0.3 * reg_loss
all.backward()
nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad_norm)
optimizer.step()
all_moniter(all)
sdr_moniter(sdr)
speaker_moniter(speaker_loss)
reg_moniter(reg_loss)
step += 1
if (idx + 1) % 3000 == 0:
save_checkpoint(model, optimizer, -1, step, args.exp_dir)
#val_loss, val_sisnr= validation(model, args, lr, epoch, device)
#scheduler.step(val_loss)
#lr = get_learning_rate(optimizer)
if (idx + 1) % print_freq == 0:
eplashed = time.time() - stime
speed_avg = eplashed / (idx + 1)
log_str = 'Epoch {:3d}/{:3d} | batches {:5d}/{:5d} | lr {:1.4e} |'\
'{:2.3f}s/batches |' \
' {:s} |'\
' {:s} |'\
' {:s} |'\
' {:s} |'\
''.format(
epoch, args.max_epoch, idx + 1, num_batch, lr,
speed_avg,
all_moniter.recent(),
sdr_moniter.recent(),
speaker_moniter.recent(),
reg_moniter.recent()
)
writer.add_scalar('Loss/Train', all_moniter.rec_float(), step)
writer.add_scalar('SISNR/Train', -sdr_moniter.rec_float(),
step)
print(log_str)
all_moniter.reset()
sdr_moniter.reset()
speaker_moniter.reset()
reg_moniter.reset()
sys.stdout.flush()
eplashed = time.time() - stime
log_str = 'Training AVG.LOSS |' \
' Epoch {:3d}/{:3d} | lr {:1.4e} |' \
' {:2.3f}s/batch | time {:3.2f}mins |' \
' {:s} |'\
' {:s} |'\
' {:s} |'\
' {:s} |'\
''.format(
epoch + 1,
args.max_epoch,
lr,
eplashed/num_batch,
eplashed/60.0,
all_moniter.average(),
sdr_moniter.average(),
speaker_moniter.average(),
reg_moniter.average()
)
print(log_str)
val_loss, val_sisnr = validation(model, args, lr, epoch, device)
writer.add_scalar('Loss/Cross-Validation', val_loss, step)
writer.add_scalar('SISNR/Cross-Validation', -val_sisnr, step)
writer.add_scalar('learn_rate', lr, step)
if val_loss > scheduler.best:
print('Rejected !!! The best is {:2.6f}'.format(scheduler.best))
else:
save_checkpoint(model,
optimizer,
epoch + 1,
step,
args.exp_dir,
mode='best_model')
scheduler.step(val_loss)
sys.stdout.flush()
stime = time.time()
def train(model, args, device, writer):
print('preparing data...')
dataloader, dataset = make_loader(args.tr_list,
args.batch_size,
8,
num_workers=args.num_threads,
processer=Processer())
print_freq = 100
num_batch = len(dataloader)
params = model.get_params(args.weight_decay)
optimizer = optim.Adam(params, lr=args.learn_rate)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
'min',
factor=0.5,
patience=1,
verbose=True)
if args.retrain:
start_epoch, step = reload_model(model, optimizer, args.exp_dir,
args.use_cuda)
else:
start_epoch, step = 0, 0
print('---------PRERUN-----------')
lr = get_learning_rate(optimizer)
print('(Initialization)')
val_loss, val_sisnr = validation(model, args, lr, -1, device)
writer.add_scalar('Loss/Train', val_loss, step)
writer.add_scalar('Loss/Cross-Validation', val_loss, step)
writer.add_scalar('SISNR/Train', -val_sisnr, step)
writer.add_scalar('SISNR/Cross-Validation', -val_sisnr, step)
warmup_epoch = 6
warmup_lr = args.learn_rate / (4 * warmup_epoch)
for epoch in range(start_epoch, args.max_epoch):
torch.manual_seed(args.seed + epoch)
if args.use_cuda:
torch.cuda.manual_seed(args.seed + epoch)
model.train()
loss_total = 0.0
loss_print = 0.0
sisnr_total = 0.0
sisnr_print = 0.0
'''
if epoch == 0 and warmup_epoch > 0:
print('Use warmup stragery, and the lr is set to {:.5f}'.format(warmup_lr))
setup_lr(optimizer, warmup_lr)
warmup_lr *= 4*(epoch+1)
elif epoch == warmup_epoch:
print('The warmup was end, and the lr is set to {:.5f}'.format(args.learn_rate))
setup_lr(optimizer, args.learn_rate)
'''
stime = time.time()
lr = get_learning_rate(optimizer)
for idx, data in enumerate(dataloader):
torch.cuda.empty_cache()
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
model.zero_grad()
est_spec, est_wav = data_parallel(model, (inputs, ))
'''
if epoch > 8:
gth_spec, gth_wav = data_parallel(model, (labels,))
else:
gth_spec = data_parallel(model.stft, (labels))[0]
'''
#gth_spec = data_parallel(model.stft, (labels))
#loss = model.loss(est_spec, gth_spec, loss_mode='MSE')
#loss.backward()
sisnr = model.loss(est_wav, labels, loss_mode='SI-SNR')
sisnr.backward()
nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad_norm)
optimizer.step()
step += 1
#loss_total += loss.data.cpu()
#loss_print += loss.data.cpu()
sisnr_total += sisnr.data.cpu()
sisnr_print += sisnr.data.cpu()
loss_total = sisnr_total
loss_print = sisnr_print
del est_wav, est_spec
if (idx + 1) % 3000 == 0:
save_checkpoint(model, optimizer, -1, step, args.exp_dir)
if (idx + 1) % print_freq == 0:
eplashed = time.time() - stime
speed_avg = eplashed / (idx + 1)
loss_print_avg = loss_print / print_freq
sisnr_print_avg = sisnr_print / print_freq
print('Epoch {:3d}/{:3d} | batches {:5d}/{:5d} | lr {:1.4e} |'
'{:2.3f}s/batches | loss {:2.6f} |'
'SI-SNR {:2.4f} '.format(
epoch,
args.max_epoch,
idx + 1,
num_batch,
lr,
speed_avg,
loss_print_avg,
-sisnr_print_avg,
))
sys.stdout.flush()
writer.add_scalar('Loss/Train', loss_print_avg, step)
writer.add_scalar('SISNR/Train', -sisnr_print_avg, step)
loss_print = 0.0
sisnr_print = 0.0
eplashed = time.time() - stime
loss_total_avg = loss_total / num_batch
sisnr_total_avg = sisnr_total / num_batch
print('Training AVG.LOSS |'
' Epoch {:3d}/{:3d} | lr {:1.4e} |'
' {:2.3f}s/batch | time {:3.2f}mins |'
' loss {:2.6f} |'
' SISNR {:2.4f}|'.format(epoch + 1, args.max_epoch, lr,
eplashed / num_batch, eplashed / 60.0,
loss_total_avg.item(),
-sisnr_total_avg.item()))
val_loss, val_sisnr = validation(model, args, lr, epoch, device)
writer.add_scalar('Loss/Cross-Validation', val_loss, step)
writer.add_scalar('SISNR/Cross-Validation', -val_sisnr, step)
writer.add_scalar('learn_rate', lr, step)
if val_loss > scheduler.best:
print('Rejected !!! The best is {:2.6f}'.format(scheduler.best))
else:
save_checkpoint(model,
optimizer,
epoch + 1,
step,
args.exp_dir,
mode='best_model')
scheduler.step(val_loss)
sys.stdout.flush()
stime = time.time()