예제 #1
0
 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)
예제 #2
0
    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
예제 #3
0
 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))
예제 #4
0
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
예제 #5
0
    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))
예제 #6
0
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)
예제 #7
0
    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)
예제 #8
0
 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
예제 #9
0
    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
예제 #10
0
 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
예제 #11
0
    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)
예제 #12
0
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)
예제 #13
0
 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)