def __init__(self,
nnet,
checkpoint="checkpoint",
optimizer="adam",
lr=1e-5,
momentum=0.9,
weight_decay=0,
clip_norm=None,
min_lr=0,
patience=1,
factor=0.5):
self.nnet = nnet
logger.info("Network structure:\n{}".format(self.nnet))
self.optimizer = create_optimizer(optimizer,
self.nnet.parameters(),
lr=lr,
momentum=momentum,
weight_decay=weight_decay)
self.scheduler = ReduceLROnPlateau(self.optimizer,
mode='min',
factor=factor,
patience=patience,
min_lr=min_lr,
verbose=True)
self.nnet.to(device)
self.checkpoint = checkpoint
self.num_spks = nnet.num_spks
self.clip_norm = clip_norm
if self.clip_norm:
logger.info("Clip gradient by 2-norm {}".format(clip_norm))
if not os.path.exists(checkpoint):
os.makedirs(checkpoint)
def validate(self, dataloader, epoch):
"""one epoch"""
self.TasNET.eval()
logger.info("Evaluating...")
tot_loss = 0
tot_batch = len(dataloader)
batch_indx = epoch * tot_batch
#print(tot_batch)
with torch.no_grad():
for mix_speech, speech1, speech2 in dataloader:
if torch.cuda.is_available():
mix_speech = mix_speech.cuda()
speech1 = speech1.cuda()
speech2 = speech2.cuda()
mix_speech = Variable(mix_speech)
speech1 = Variable(speech1)
speech2 = Variable(speech2)
output1, output2 = self.TasNET(mix_speech)
cur_loss = self.loss(output1, output2, speech1, speech2)
tot_loss += cur_loss.item()
#write summary
batch_indx += 1
self.writer.add_scalar('dev_loss', cur_loss, batch_indx)
return tot_loss / tot_batch, tot_batch
def rerun(self, train_set, dev_set, model_path, epoch_done):
self.TasNET.load_state_dict(torch.load(model_path))
init_loss, _ = self.validate(dev_set, epoch_done)
logger.info("Start training for {} epoches".format(self.num_epoches))
logger.info("Epoch {:2d}: dev loss ={:.4e}".format(0, init_loss))
torch.save(self.TasNET.state_dict(),
os.path.join(self.checkpoint, 'TasNET_0.pkl'))
for epoch in range(epoch_done, self.num_epoches + 1):
train_start = time.time()
train_loss, train_num_batch = self.train(train_set, epoch)
valid_start = time.time()
valid_loss, valid_num_batch = self.validate(dev_set, epoch)
valid_end = time.time()
self.scheduler.step(valid_loss)
logger.info("Epoch {:2d}: train loss = {:.4e}({:.2f}s/{:d}) |"
" dev loss= {:.4e}({:.2f}s/{:d})".format(
epoch, train_loss, valid_start - train_start,
train_num_batch, valid_loss,
valid_end - valid_start, valid_num_batch))
save_path = os.path.join(
self.checkpoint,
"TasNET_{:d}_trainloss_{:.4e}_valloss_{:.4e}.pkl".format(
epoch, train_loss, valid_loss))
torch.save(self.TasNET.state_dict(), save_path)
logger.info("Training for {} epoches done!".format(self.num_epoches))
def create_optimizer(optimizer, params, **kwargs):
supported_optimizer = {
'sgd': th.optim.SGD, # momentum, weight_decay, lr
'rmsprop': th.optim.RMSprop, # momentum, weight_decay, lr
'adam': th.optim.Adam # weight_decay, lr
# ...
}
if optimizer not in supported_optimizer:
raise ValueError('Unsupported optimizer {}'.format(optimizer))
if optimizer == 'adam':
del kwargs['momentum']
opt = supported_optimizer[optimizer](params, **kwargs)
logger.info('Create optimizer {}({})'.format(optimizer, kwargs))
return opt
def run(self, train_set, dev_set, num_epoches=20, start=1):
init_loss, _ = self.validate(dev_set)
logger.info("Epoch {:2d}: dev = {:.4f}".format(0, init_loss))
if (start == 1):
th.save(self.nnet.state_dict(),
os.path.join(self.checkpoint, 'epoch.0.pkl'))
else:
logger.info("training from {}epoch".format(start))
for epoch in range(start, num_epoches + 1):
on_train_start = time.time()
train_loss, train_num_batch = self.train(train_set)
on_valid_start = time.time()
valid_loss, valid_num_batch = self.validate(dev_set)
on_valid_end = time.time()
# scheduler learning rate
self.scheduler.step(valid_loss)
logger.info(
"Loss(time/mini-batch) - Epoch {:2d}: train = {:.4f}({:.2f}s/{:d}) |"
" dev = {:.4f}({:.2f}s/{:d})".format(
epoch, train_loss, on_valid_start - on_train_start,
train_num_batch, valid_loss, on_valid_end - on_valid_start,
valid_num_batch))
save_path = os.path.join(self.checkpoint,
'epoch.{:d}.pkl'.format(epoch))
th.save(self.nnet.state_dict(), save_path)
logger.info("Training for {} epoches done!".format(num_epoches))
def train(args):
debug = args.debug
logger.info(
"Start training in {} model".format('debug' if debug else 'normal'))
num_bins, config_dict = parse_yaml(args.config)
reader_conf = config_dict["spectrogram_reader"]
loader_conf = config_dict["dataloader"]
dcnnet_conf = config_dict["model"]
logger.info("Training with {}".format(
"IRM" if reader_conf["apply_abs"] else "PSM"))
batch_size = loader_conf["batch_size"]
logger.info(
"Training in {}".format("per utterance" if batch_size == 1 else
'{} utterance per batch'.format(batch_size)))
train_loader = uttloader(config_dict["train_scp_conf"]
if not debug else config_dict["debug_scp_conf"],
reader_conf,
loader_conf,
train=True)
valid_loader = uttloader(config_dict["valid_scp_conf"]
if not debug else config_dict["debug_scp_conf"],
reader_conf,
loader_conf,
train=False)
checkpoint = config_dict["trainer"]["checkpoint"]
logger.info("Training for {} epoches -> {}...".format(
args.num_epoches,
"default checkpoint" if checkpoint is None else checkpoint))
nnet = PITNet(num_bins, **dcnnet_conf)
trainer = PITrainer(nnet, **config_dict["trainer"])
trainer.run(train_loader, valid_loader, num_epoches=args.num_epoches)
def __init__(self,
chimera,
alpha=0.1,
checkpoint="checkpoint",
optimizer="adam",
lr=1e-5,
momentum=0.9,
weight_decay=0,
clip_norm=None,
num_spks=2):
self.nnet = chimera
self.nnet.load_state_dict(torch.load("model/chimera++0.975_107_trainloss_2.2019e-01_valloss_3.0721e-01.pkl"))
logger.info("Load trained model...")
self.alpha = alpha
logger.info("chimeraNet:\n{}".format(self.nnet))
if type(lr) is str:
lr = float(lr)
logger.info("Transfrom lr from str to float => {}".format(lr))
self.optimizer = torch.optim.Adam(
self.nnet.parameters(),
lr=lr,
weight_decay=weight_decay)
self.scheduler = ReduceLROnPlateau(self.optimizer, verbose=True)
self.nnet.to(device)
self.checkpoint = checkpoint
self.num_spks = num_spks
self.clip_norm = clip_norm
if self.clip_norm:
logger.info("Clip gradient by 2-norm {}".format(clip_norm))
if not os.path.exists(self.checkpoint):
os.makedirs(checkpoint)
def validate(self, dataloader):
self.nnet.eval()
logger.info("Evaluating...")
tot_loss = 0
num_batches = len(dataloader)
# do not need to keep gradient
with th.no_grad():
for mix_spect, tgt_index, vad_masks in dataloader:
mix_spect = mix_spect.cuda() if isinstance(
mix_spect, PackedSequence) else mix_spect.to(device)
tgt_index = tgt_index.to(device)
vad_masks = vad_masks.to(device)
# mix_spect = mix_spect * vad_masks
net_embed = self.nnet(mix_spect)
cur_loss = self.loss(net_embed, tgt_index, vad_masks)
tot_loss += cur_loss.item()
return tot_loss / num_batches, num_batches
def train(self, dataloader):
"""one epoch"""
self.nnet.train()
logger.info("Training...")
tot_loss1 = 0
tot_loss2 = 0
tot_loss = 0
tot_batch = len(dataloader)
# import pdb;pdb.set_trace()
for mix_true, inst_true, vocal_true, log_mag, mix_phase, tgt_index, vad_masks, pss_inst, pss_vocal in dataloader:
self.optimizer.zero_grad()
if torch.cuda.is_available():
log_mag = log_mag.cuda()
tgt_index = tgt_index.cuda()
vad_masks = vad_masks.cuda()
mix_phase = mix_phase.cuda()
inst_true = inst_true.cuda()
vocal_true = vocal_true.cuda()
# log_mag = log_mag.requires_grad_()
tgt_index = tgt_index.requires_grad_()
vad_masks = vad_masks.requires_grad_()
mix_phase = mix_phase.requires_grad_()
inst_true = inst_true.requires_grad_()
vocal_true = vocal_true.requires_grad_()
net_embed, net_irm = self.nnet(log_mag)
cur_loss1, cur_loss2, cur_loss = self.loss(net_embed, net_irm, tgt_index, vad_masks, log_mag, mix_phase,
inst_true, vocal_true)
tot_loss1 += cur_loss1.item()
tot_loss2 += cur_loss2.item()
tot_loss += cur_loss.item()
cur_loss.backward()
if self.clip_norm:
nn.utils.clip_grad_norm_(self.nnet.parameters(),
self.clip_norm)
self.optimizer.step()
return tot_loss / tot_batch, tot_batch, tot_loss1 / tot_batch, tot_loss2 / tot_batch
def train(self, dataloader):
self.nnet.train()
logger.info("Training...")
tot_loss = 0
num_batches = len(dataloader)
for mix_spect, tgt_index, vad_masks in dataloader:
self.optimizer.zero_grad()
mix_spect = mix_spect.cuda() if isinstance(
mix_spect, PackedSequence) else mix_spect.to(device)
tgt_index = tgt_index.to(device)
vad_masks = vad_masks.to(device)
# mix_spect = mix_spect * vad_masks
net_embed = self.nnet(mix_spect)
cur_loss = self.loss(net_embed, tgt_index, vad_masks)
tot_loss += cur_loss.item()
cur_loss.backward()
if self.clip_norm:
th.nn.utils.clip_grad_norm_(self.nnet.parameters(),
self.clip_norm)
self.optimizer.step()
return tot_loss / num_batches, num_batches
def __init__(self,
TasNET,
checkpoint="checkpoint",
optimizer="adam",
lr=1e-5,
momentum=0.9,
weight_decay=0,
num_epoches=20,
clip_norm=False,
sr=8000):
self.TasNET = TasNET
logger.info("TasNET:\n{}".format(self.TasNET))
if type(lr) is str:
lr = float(lr)
logger.info("Transfrom lr from str to float => {}".format(lr))
self.optimizer = torch.optim.Adam(self.TasNET.parameters(),
lr=lr,
weight_decay=weight_decay)
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer, 'min', factor=0.5, patience=3, verbose=True)
self.TasNET.to(device)
self.checkpoint = checkpoint
self.num_epoches = num_epoches
self.clip_norm = clip_norm
self.sr = sr
self.writer = SummaryWriter('./log')
if self.clip_norm:
logger.info("Clip gradient by 2-norm {}".format(clip_norm))
if not os.path.exists(self.checkpoint):
os.makedirs(checkpoint)
def train(args):
debug = args.debug
logger.info(
"Start training in {} model".format('debug' if debug else 'normal'))
num_bins, config_dict = parse_yaml(args.config)
reader_conf = config_dict["spectrogram_reader"]
loader_conf = config_dict["dataloader"]
dcnnet_conf = config_dict["dcnet"]
batch_size = loader_conf["batch_size"]
logger.info(
"Training in {}".format("per utterance" if batch_size == 1 else
'{} utterance per batch'.format(batch_size)))
train_loader = uttloader(
config_dict["train_scp_conf"]
if not debug else config_dict["debug_scp_conf"],
reader_conf,
loader_conf,
train=True)
valid_loader = uttloader(
config_dict["valid_scp_conf"]
if not debug else config_dict["debug_scp_conf"],
reader_conf,
loader_conf,
train=False)
checkpoint = config_dict["trainer"]["checkpoint"]
logger.info("Training for {} epoches -> {}...".format(
args.num_epoches, "default checkpoint"
if checkpoint is None else checkpoint))
dcnet = DCNet(num_bins, **dcnnet_conf)
trainer = Trainer(dcnet, **config_dict["trainer"])
trainer.run(train_loader, valid_loader, num_epoches=args.num_epoches)
def __init__(self,
dcnet,
checkpoint="checkpoint",
optimizer="adam",
lr=1e-5,
momentum=0.9,
weight_decay=0,
clip_norm=None,
num_spks=2):
self.nnet = dcnet
logger.info("DCNet:\n{}".format(self.nnet))
self.optimizer = create_optimizer(optimizer,
self.nnet.parameters(),
lr=lr,
momentum=momentum,
weight_decay=weight_decay)
self.nnet.to(device)
self.checkpoint = checkpoint
self.num_spks = num_spks
self.clip_norm = clip_norm
if self.clip_norm:
logger.info("Clip gradient by 2-norm {}".format(clip_norm))
if not os.path.exists(checkpoint):
os.makedirs(checkpoint)