def validate(self): kaldi.InstantiateKaldiCuda() chain_opts = self.chain_opts den_fst_path = os.path.join(chain_opts.dir, "den.fst") # load model model = self.initialize_model() self.load_base_model(model) model.eval() training_opts = kaldi.chain.CreateChainTrainingOptions( chain_opts.l2_regularize, chain_opts.out_of_range_regularize, chain_opts.leaky_hmm_coefficient, chain_opts.xent_regularize, ) compute_chain_objf( model, chain_opts.egs, den_fst_path, training_opts, minibatch_size="1:64", left_context=chain_opts.context, right_context=chain_opts.context, use_ivector=True if self.chain_opts.ivector_dir else False)
def validate(self): kaldi.InstantiateKaldiCuda() chain_opts = self.chain_opts den_fst_path = os.path.join(chain_opts.dir, "den.fst") # load model model = self.Net(self.chain_opts.feat_dim, self.chain_opts.output_dim) model.load_state_dict(torch.load(chain_opts.base_model)) model.eval() training_opts = kaldi.chain.CreateChainTrainingOptions( chain_opts.l2_regularize, chain_opts.out_of_range_regularize, chain_opts.leaky_hmm_coefficient, chain_opts.xent_regularize, ) compute_chain_objf( model, chain_opts.egs, den_fst_path, training_opts, minibatch_size="1:64", left_context=chain_opts.context, right_context=chain_opts.context, )
def train(self): """Run one iteration of LF-MMI training This is called by >>> self.train() It will probably be renamed as self.fit() since this seems to be the standard way other libraries call the training function. """ kaldi.InstantiateKaldiCuda() chain_opts = self.chain_opts lr = chain_opts.lr den_fst_path = os.path.join(chain_opts.dir, "den.fst") # load model model = self.initialize_model() self.load_base_model(model) training_opts = kaldi.chain.CreateChainTrainingOptions( chain_opts.l2_regularize, chain_opts.out_of_range_regularize, chain_opts.leaky_hmm_coefficient, chain_opts.xent_regularize, ) context = chain_opts.context model = model.cuda() optimizer = self.get_optimizer( model, lr=chain_opts.lr, weight_decay=chain_opts.l2_regularize_factor) new_model = train_lfmmi_one_iter( model, chain_opts.egs, den_fst_path, training_opts, chain_opts.feat_dim, minibatch_size=chain_opts.minibatch_size, left_context=context, right_context=context, lr=chain_opts.lr, weight_decay=chain_opts.l2_regularize_factor, frame_shift=chain_opts.frame_shift, optimizer=optimizer, e2e=True) torch.save(new_model.state_dict(), chain_opts.new_model)
def combine_final_model(self): """Implements Kaldi-style model ensembling""" kaldi.InstantiateKaldiCuda() chain_opts = self.chain_opts den_fst_path = os.path.join(chain_opts.dir, "den.fst") base_models = chain_opts.base_model.split(',') assert len(base_models) > 0 training_opts = kaldi.chain.CreateChainTrainingOptions( chain_opts.l2_regularize, chain_opts.out_of_range_regularize, chain_opts.leaky_hmm_coefficient, chain_opts.xent_regularize, ) moving_average = self.Net(self.chain_opts.feat_dim, self.chain_opts.output_dim) best_mdl = self.Net(self.chain_opts.feat_dim, self.chain_opts.output_dim) moving_average.load_state_dict(torch.load(base_models[0])) moving_average.cuda() best_mdl = moving_average compute_objf = lambda mdl: compute_chain_objf( mdl, chain_opts.egs, den_fst_path, training_opts, minibatch_size="1:64", # TODO: this should come from a config left_context=chain_opts.context, right_context=chain_opts.context, frame_shift=chain_opts.frame_shift, ) _, init_objf = compute_objf(moving_average) best_objf = init_objf model_acc = dict(moving_average.named_parameters()) num_accumulated = torch.Tensor([1.0]).reshape(1).cuda() best_num_to_combine = 1 for mdl_name in base_models[1:]: this_mdl = self.Net(self.chain_opts.feat_dim, self.chain_opts.output_dim) logging.info("Combining model {}".format(mdl_name)) this_mdl.load_state_dict(torch.load(mdl_name)) this_mdl = this_mdl.cuda() # TODO(srikanth): check why is this even necessary moving_average.cuda() num_accumulated += 1. for name, params in this_mdl.named_parameters(): model_acc[name].data.mul_( (num_accumulated - 1.) / (num_accumulated)) model_acc[name].data.add_( params.data.mul_(1. / num_accumulated)) _, this_objf = compute_objf(moving_average) if this_objf > best_objf: best_objf = this_objf best_mdl = moving_average best_num_to_combine = num_accumulated.clone().detach() logging.info("Found best model") else: logging.info("Won't update best model") logging.info("Combined {} models".format(best_num_to_combine.cpu())) logging.info("Initial objf = {}, Final objf = {}".format( init_objf, best_objf)) best_mdl.cpu() torch.save(best_mdl.state_dict(), chain_opts.new_model) return self
def train_lfmmi_one_iter( model, egs_file, den_fst_path, training_opts, feat_dim, minibatch_size="64", use_gpu=True, lr=0.0001, weight_decay=0.25, frame_shift=0, left_context=0, right_context=0, print_interval=10, frame_subsampling_factor=3, ): """Run one iteration of LF-MMI training The function loads the latest model, takes a list of egs, path to denominator fst and runs through the merged egs for one iteration of training. This is similar to how one iteration of training is completed in Kaldi. Args: model: Path to pytorch model (.pt file) egs_file: scp or ark file (a string), should be prefix accordingly just like Kaldi den_fst_path: path to den.fst file training_opts: options of type ChainTrainingOpts feat_dim: dimension of features (e.g. 40 for MFCC hires features) minibatch_size: a string of minibatch sizes separated by commas. E.g "64" or "128,64" use_gpu: a boolean to set or unset the use of GPUs while training lr: learning rate frame_shift: an integer (usually 0, 1, or 2) used to shift the training features print_interval: the interval (a positive integer) to print the loss value Returns: updated model in CPU """ # this is required to make sure Kaldi uses GPU kaldi.InstantiateKaldiCuda() if training_opts is None: training_opts = kaldi.chain.CreateChainTrainingOptionsDefault() den_graph = kaldi.chain.LoadDenominatorGraph(den_fst_path, model.output_dim) criterion = KaldiChainObjfFunction.apply if use_gpu: model = model.cuda() optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=weight_decay) acc_sum = torch.tensor(0., requires_grad=False) for mb_id, merged_egs in enumerate( prepare_minibatch(egs_file, minibatch_size)): chunk_size = kaldi.chain.GetFramesPerSequence( merged_egs) * frame_subsampling_factor features = kaldi.chain.GetFeaturesFromEgs(merged_egs) features = features[:, frame_shift:frame_shift + chunk_size + left_context + right_context, :] features = features.cuda() output, xent_output = model(features) sup = kaldi.chain.GetSupervisionFromEgs(merged_egs) deriv = criterion(training_opts, den_graph, sup, output, xent_output) acc_sum.add_(deriv[0]) if mb_id > 0 and mb_id % print_interval == 0: logging.info("Overall objf={}\n".format(acc_sum / print_interval)) acc_sum.zero_() optimizer.zero_grad() deriv.backward() clip_grad_value_(model.parameters(), 5.0) optimizer.step() model = model.cpu() return model