class Solver(BaseSolver): ''' Solver for training''' def __init__(self, config, paras, mode): super().__init__(config, paras, mode) # ToDo : support tr/eval on different corpus assert self.config['data']['corpus']['name'] == self.src_config[ 'data']['corpus']['name'] self.config['data']['corpus']['path'] = self.src_config['data'][ 'corpus']['path'] self.config['data']['corpus']['bucketing'] = False # The follow attribute should be identical to training config self.config['data']['audio'] = self.src_config['data']['audio'] self.config['data']['text'] = self.src_config['data']['text'] self.config['model'] = self.src_config['model'] # Output file self.output_file = str(self.ckpdir) + '_{}_{}.csv' # Override batch size for beam decoding self.greedy = self.config['decode']['beam_size'] == 1 if not self.greedy: self.config['data']['corpus']['batch_size'] = 1 else: # ToDo : implement greedy raise NotImplementedError def load_data(self): ''' Load data for training/validation, store tokenizer and input/output shape''' self.dv_set, self.tt_set, self.feat_dim, self.vocab_size, self.tokenizer, msg = \ load_dataset(self.paras.njobs, self.paras.gpu, self.paras.pin_memory, False, **self.config['data']) self.verbose(msg) def set_model(self): ''' Setup ASR model ''' # Model init_adadelta = True self.model = ASR(self.feat_dim, self.vocab_size, init_adadelta, **self.config['model']) # Plug-ins if ('emb' in self.config) and (self.config['emb']['enable']) \ and (self.config['emb']['fuse'] > 0): from src.plugin import EmbeddingRegularizer self.emb_decoder = EmbeddingRegularizer(self.tokenizer, self.model.dec_dim, **self.config['emb']) # Load target model in eval mode self.load_ckpt() # Beam decoder self.decoder = BeamDecoder(self.model.cpu(), self.emb_decoder, **self.config['decode']) self.verbose(self.decoder.create_msg()) del self.model del self.emb_decoder def exec(self): ''' Testing End-to-end ASR system ''' for s, ds in zip(['dev', 'test'], [self.dv_set, self.tt_set]): # Setup output self.cur_output_path = self.output_file.format(s, 'output') with open(self.cur_output_path, 'w') as f: f.write('idx\thyp\ttruth\n') if self.greedy: # Greedy decode self.verbose( 'Performing batch-wise greedy decoding on {} set, num of batch = {}.' .format(s, len(ds))) self.verbose('Results will be stored at {}'.format( self.cur_output_path)) else: # Additional output to store all beams self.cur_beam_path = self.output_file.format(s, 'beam') with open(self.cur_beam_path, 'w') as f: f.write('idx\tbeam\thyp\ttruth\n') self.verbose( 'Performing instance-wise beam decoding on {} set. (NOTE: use --njobs to speedup)' .format(s)) # Minimal function to pickle beam_decode_func = partial(beam_decode, model=copy.deepcopy(self.decoder), device=self.device) # Parallel beam decode results = Parallel(n_jobs=self.paras.njobs)( delayed(beam_decode_func)(data) for data in tqdm(ds)) self.verbose('Results/Beams will be stored at {} / {}.'.format( self.cur_output_path, self.cur_beam_path)) self.write_hyp(results, self.cur_output_path, self.cur_beam_path) self.verbose('All done !') def write_hyp(self, results, best_path, beam_path): '''Record decoding results''' for name, hyp_seqs, truth in tqdm(results): hyp_seqs = [self.tokenizer.decode(hyp) for hyp in hyp_seqs] truth = self.tokenizer.decode(truth) with open(best_path, 'a') as f: f.write('\t'.join([name, hyp_seqs[0], truth]) + '\n') if not self.greedy: with open(beam_path, 'a') as f: for b, hyp in enumerate(hyp_seqs): f.write('\t'.join([name, str(b), hyp, truth]) + '\n')
class Solver(BaseSolver): ''' Solver for training''' def __init__(self, config, paras, mode): super().__init__(config, paras, mode) # ToDo : support tr/eval on different corpus assert self.config['data']['corpus']['name'] == self.src_config['data']['corpus']['name'] self.config['data']['corpus']['path'] = self.src_config['data']['corpus']['path'] self.config['data']['corpus']['bucketing'] = False # The follow attribute should be identical to training config self.config['data']['audio'] = self.src_config['data']['audio'] self.config['data']['text'] = self.src_config['data']['text'] self.config['model'] = self.src_config['model'] # Output file self.output_file = str(self.ckpdir)+'_{}_{}.csv' # Override batch size for beam decoding self.greedy = self.config['decode']['beam_size'] == 1 if not self.greedy: self.config['data']['corpus']['batch_size'] = 1 else: # ToDo : implement greedy raise NotImplementedError def load_data(self): ''' Load data for training/validation, store tokenizer and input/output shape''' if self.paras.upstream is not None: print(f'[Solver] - using S3PRL {self.paras.upstream}') self.dv_set, self.tt_set, self.vocab_size, self.tokenizer, msg = \ load_wav_dataset(self.paras.njobs, self.paras.gpu, self.paras.pin_memory, False, **self.config['data']) self.upstream = torch.hub.load( 's3prl/s3prl', args.upstream, feature_selection = args.upstream_feature_selection, refresh = args.upstream_refresh, ckpt = args.upstream_ckpt, force_reload = True, ) self.feat_dim = self.upstream.get_output_dim() else: self.dv_set, self.tt_set, self.feat_dim, self.vocab_size, self.tokenizer, msg = \ load_dataset(self.paras.njobs, self.paras.gpu, self.paras.pin_memory, False, **self.config['data']) self.verbose(msg) def set_model(self): ''' Setup ASR model ''' # Model init_adadelta = self.src_config['hparas']['optimizer'] == 'Adadelta' self.model = ASR(self.feat_dim, self.vocab_size, init_adadelta, ** self.config['model']).to(self.device) # Plug-ins if ('emb' in self.config) and (self.config['emb']['enable']) \ and (self.config['emb']['fuse'] > 0): from src.plugin import EmbeddingRegularizer self.emb_decoder = EmbeddingRegularizer( self.tokenizer, self.model.dec_dim, **self.config['emb']) # Load target model in eval mode self.load_ckpt() # Beam decoder self.decoder = BeamDecoder( self.model.cpu(), self.emb_decoder, **self.config['decode']) self.verbose(self.decoder.create_msg()) del self.model del self.emb_decoder def exec(self): ''' Testing End-to-end ASR system ''' for s, ds in zip(['dev', 'test'], [self.dv_set, self.tt_set]): # Setup output self.cur_output_path = self.output_file.format(s, 'output') with open(self.cur_output_path, 'w') as f: f.write('idx\thyp\ttruth\n') if self.greedy: # Greedy decode self.verbose( 'Performing batch-wise greedy decoding on {} set, num of batch = {}.'.format(s, len(ds))) self.verbose('Results will be stored at {}'.format( self.cur_output_path)) else: # Additional output to store all beams self.cur_beam_path = self.output_file.format(s, 'beam') with open(self.cur_beam_path, 'w') as f: f.write('idx\tbeam\thyp\ttruth\n') self.verbose( 'Performing instance-wise beam decoding on {} set. (NOTE: use --njobs to speedup)'.format(s)) # Minimal function to pickle beam_decode_func = partial(beam_decode, model=copy.deepcopy( self.decoder), device=self.device) def handler(data): if self.paras.upstream is not None: # feat is raw waveform name, feat, feat_len, txt = data device = 'cpu' if self.paras.deterministic else self.device self.upstream.to(device) def to_device(feat): return [f.to(device) for f in feat] def extract_feature(feat): feat = self.upstream(to_device(feat)) return feat self.upstream.eval() with torch.no_grad(): feat = extract_feature(feat) feat_len = torch.LongTensor([len(f) for f in feat]) feat = pad_sequence(feat, batch_first=True) txt = pad_sequence(txt, batch_first=True) data = [name, feat, feat_len, txt] return data # Parallel beam decode results = Parallel(n_jobs=self.paras.njobs)( delayed(beam_decode_func)(handler(data)) for data in tqdm(ds)) self.verbose( 'Results/Beams will be stored at {} / {}.'.format(self.cur_output_path, self.cur_beam_path)) self.write_hyp(results, self.cur_output_path, self.cur_beam_path) self.verbose('All done !') def write_hyp(self, results, best_path, beam_path): '''Record decoding results''' for name, hyp_seqs, truth in tqdm(results): hyp_seqs = [self.tokenizer.decode(hyp) for hyp in hyp_seqs] truth = self.tokenizer.decode(truth) with open(best_path, 'a') as f: f.write('\t'.join([name, hyp_seqs[0], truth])+'\n') if not self.greedy: with open(beam_path, 'a') as f: for b, hyp in enumerate(hyp_seqs): f.write('\t'.join([name, str(b), hyp, truth])+'\n')
class Solver(BaseSolver): ''' Solver for training''' def __init__(self, config, paras, mode): super().__init__(config, paras, mode) # ToDo : support tr/eval on different corpus assert self.config['data']['corpus']['name'] == self.src_config[ 'data']['corpus']['name'] self.config['data']['corpus']['path'] = self.src_config['data'][ 'corpus']['path'] self.config['data']['corpus']['bucketing'] = False # The follow attribute should be identical to training config self.config['data']['audio'] = self.src_config['data']['audio'] self.config['data']['text'] = self.src_config['data']['text'] self.config['hparas'] = self.src_config['hparas'] self.config['model'] = self.src_config['model'] # Output file self.output_file = str(self.ckpdir) + '_{}_{}.csv' # Override batch size for beam decoding self.greedy = self.config['decode']['beam_size'] == 1 if not self.greedy: self.config['data']['corpus']['batch_size'] = 1 self.step = 0 def fetch_data(self, data): ''' Move data to device and compute text seq. length, For Greedy decoding only ( beam_decode & ctc_beam_decode otherwise)''' _, feat, feat_len, txt = data feat = feat.to(self.device) feat_len = feat_len.to(self.device) txt = txt.to(self.device) txt_len = torch.sum(txt != 0, dim=-1) return feat, feat_len, txt, txt_len def load_data(self): ''' Load data for training/validation, store tokenizer and input/output shape''' self.dv_set, self.tt_set, self.feat_dim, self.vocab_size, self.tokenizer, msg = \ load_dataset(self.paras.njobs, self.paras.gpu, self.paras.pin_memory, False, **self.config['data']) self.verbose(msg) def set_model(self): ''' Setup ASR model ''' # Model init_adadelta = self.config['hparas']['optimizer'] == 'Adadelta' self.model = ASR(self.feat_dim, self.vocab_size, init_adadelta, **self.config['model']).to(self.device) # Plug-ins if ('emb' in self.config) and (self.config['emb']['enable']) \ and (self.config['emb']['fuse'] > 0): from src.plugin import EmbeddingRegularizer self.emb_decoder = EmbeddingRegularizer(self.tokenizer, self.model.dec_dim, **self.config['emb']) # Load target model in eval mode self.load_ckpt() self.ctc_only = False if self.greedy: # Greedy decoding: attention-based if the ASR has a decoder, else use CTC self.decoder = copy.deepcopy(self.model).to(self.device) else: if (not self.model.enable_att) or self.config['decode'].get( 'ctc_weight', 0.0) == 1.0: # Pure CTC Beam Decoder assert self.config['decode']['beam_size'] <= self.config[ 'decode']['vocab_candidate'] self.decoder = CTCBeamDecoder( self.model.to(self.device), [1] + [r for r in range(3, self.vocab_size)], self.config['decode']['beam_size'], self.config['decode']['vocab_candidate'], lm_path=self.config['decode']['lm_path'], lm_config=self.config['decode']['lm_config'], lm_weight=self.config['decode']['lm_weight'], device=self.device) self.ctc_only = True else: # Joint CTC-Attention Beam Decoder self.decoder = BeamDecoder(self.model.cpu(), self.emb_decoder, **self.config['decode']) self.verbose(self.decoder.create_msg()) del self.model del self.emb_decoder def greedy_decode(self, dv_set): ''' Greedy Decoding ''' results = [] for i, data in enumerate(dv_set): self.progress('Valid step - {}/{}'.format(i + 1, len(dv_set))) # Fetch data feat, feat_len, txt, txt_len = self.fetch_data(data) # Forward model with torch.no_grad(): ctc_output, encode_len, att_output, att_align, dec_state = \ self.decoder( feat, feat_len, int(float(feat_len.max()) * self.config['decode']['max_len_ratio']), emb_decoder=self.emb_decoder) for j in range(len(txt)): idx = j + self.config['data']['corpus']['batch_size'] * i if att_output is not None: hyp_seqs = att_output[j].argmax(dim=-1).tolist() else: hyp_seqs = ctc_output[j].argmax(dim=-1).tolist() true_txt = txt[j] results.append((str(idx), [hyp_seqs], true_txt)) return results def exec(self): ''' Testing End-to-end ASR system ''' for s, ds in zip(['dev', 'test'], [self.dv_set, self.tt_set]): # Setup output self.cur_output_path = self.output_file.format(s, 'output') with open(self.cur_output_path, 'w', encoding='UTF-8') as f: f.write('idx\thyp\ttruth\n') if self.greedy: # Greedy decode self.verbose( 'Performing batch-wise greedy decoding on {} set, num of batch = {}.' .format(s, len(ds))) results = self.greedy_decode(ds) self.verbose('Results will be stored at {}'.format( self.cur_output_path)) self.write_hyp(results, self.cur_output_path, '-') elif self.ctc_only: # CTC beam decode # Additional output to store all beams self.cur_beam_path = self.output_file.format( s, 'beam-{}-{}'.format(self.config['decode']['beam_size'], self.config['decode']['lm_weight'])) with open(self.cur_beam_path, 'w') as f: f.write('idx\tbeam\thyp\ttruth\n') self.verbose( 'Performing instance-wise CTC beam decoding on {} set, num of batch = {}.' .format(s, len(ds))) # Minimal function to pickle ctc_beam_decode_func = partial(ctc_beam_decode, model=copy.deepcopy( self.decoder), device=self.device) # Parallel beam decode results = Parallel(n_jobs=self.paras.njobs)( delayed(ctc_beam_decode_func)(data) for data in tqdm(ds)) self.verbose('Results/Beams will be stored at {} / {}'.format( self.cur_output_path, self.cur_beam_path)) self.write_hyp(results, self.cur_output_path, self.cur_beam_path) else: # Joint CTC-Attention beam decode # Additional output to store all beams self.cur_beam_path = self.output_file.format( s, 'beam-{}-{}'.format(self.config['decode']['beam_size'], self.config['decode']['lm_weight'])) with open(self.cur_beam_path, 'w') as f: f.write('idx\tbeam\thyp\ttruth\n') self.verbose( 'Performing instance-wise beam decoding on {} set. (NOTE: use --njobs to speedup)' .format(s)) # Minimal function to pickle beam_decode_func = partial(beam_decode, model=copy.deepcopy(self.decoder), device=self.device) # Parallel beam decode results = Parallel(n_jobs=self.paras.njobs)( delayed(beam_decode_func)(data) for data in tqdm(ds)) self.verbose('Results/Beams will be stored at {} / {}.'.format( self.cur_output_path, self.cur_beam_path)) self.write_hyp(results, self.cur_output_path, self.cur_beam_path) self.verbose('All done !') def write_hyp(self, results, best_path, beam_path): '''Record decoding results''' if self.greedy: # Ignores repeated symbols if is decoded with CTC ignore_repeat = not self.decoder.enable_att else: ignore_repeat = False for name, hyp_seqs, truth in tqdm(results): if self.ctc_only and not self.greedy: new_hyp_seqs = [ self.tokenizer.decode(hyp, ignore_repeat=False) for hyp in hyp_seqs[:-1] ] hyp_seqs = new_hyp_seqs + [ self.tokenizer.decode(hyp_seqs[-1], ignore_repeat=True) ] else: hyp_seqs = [self.tokenizer.decode(hyp) for hyp in hyp_seqs] truth = self.tokenizer.decode(truth) with open(best_path, 'a') as f: if len(hyp_seqs[0]) == 0: # Set the sequence to a whitespace if it was empty hyp_seqs[0] = ' ' f.write('\t'.join([name, hyp_seqs[0], truth]) + '\n') if not self.greedy: with open(beam_path, 'a', encoding='UTF-8') as f: for b, hyp in enumerate(hyp_seqs): f.write('\t'.join([name, str(b), hyp, truth]) + '\n')