def main(): args = get_args() # No need gpu for model export os.environ['CUDA_VISIBLE_DEVICES'] = '-1' with open(args.config, 'r') as fin: configs = yaml.load(fin, Loader=yaml.FullLoader) model = init_model(configs) print(model) load_checkpoint(model, args.checkpoint) # Export jit torch script model if args.output_file: script_model = torch.jit.script(model) script_model.save(args.output_file) print('Export model successfully, see {}'.format(args.output_file)) # Export quantized jit torch script model if args.output_quant_file: quantized_model = torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtype=torch.qint8) print(quantized_model) script_quant_model = torch.jit.script(quantized_model) script_quant_model.save(args.output_quant_file) print('Export quantized model successfully, ' 'see {}'.format(args.output_quant_file))
def main(): torch.manual_seed(777) args = get_args() output_dir = args.output_dir os.system("mkdir -p " + output_dir) os.environ['CUDA_VISIBLE_DEVICES'] = '-1' with open(args.config, 'r') as fin: configs = yaml.load(fin, Loader=yaml.FullLoader) model = init_model(configs) load_checkpoint(model, args.checkpoint) model.eval() print(model) arguments = {} arguments['output_dir'] = output_dir arguments['batch'] = 1 arguments['chunk_size'] = args.chunk_size arguments['left_chunks'] = args.num_decoding_left_chunks arguments['reverse_weight'] = args.reverse_weight arguments['output_size'] = configs['encoder_conf']['output_size'] arguments['num_blocks'] = configs['encoder_conf']['num_blocks'] arguments['cnn_module_kernel'] = configs['encoder_conf'][ 'cnn_module_kernel'] arguments['head'] = configs['encoder_conf']['attention_heads'] arguments['feature_size'] = configs['input_dim'] arguments['vocab_size'] = configs['output_dim'] # NOTE(xcsong): if chunk_size == -1, hardcode to 67 arguments['decoding_window'] = (args.chunk_size - 1) * \ model.encoder.embed.subsampling_rate + \ model.encoder.embed.right_context + 1 if args.chunk_size > 0 else 67 arguments['encoder'] = configs['encoder'] arguments['decoder'] = configs['decoder'] arguments['subsampling_rate'] = model.subsampling_rate() arguments['right_context'] = model.right_context() arguments['sos_symbol'] = model.sos_symbol() arguments['eos_symbol'] = model.eos_symbol() arguments['is_bidirectional_decoder'] = 1 \ if model.is_bidirectional_decoder() else 0 # NOTE(xcsong): Please note that -1/-1 means non-streaming model! It is # not a [16/4 16/-1 16/0] all-in-one model and it should not be used in # streaming mode (i.e., setting chunk_size=16 in `decoder_main`). If you # want to use 16/-1 or any other streaming mode in `decoder_main`, # please export onnx in the same config. if arguments['left_chunks'] > 0: assert arguments['chunk_size'] > 0 # -1/4 not supported export_encoder(model, arguments) export_ctc(model, arguments) export_decoder(model, arguments)
test_collate_func = CollateFunc(**test_collate_conf, raw_wav=raw_wav) dataset_conf = configs.get('dataset_conf', {}) dataset_conf['batch_size'] = args.batch_size dataset_conf['batch_type'] = 'static' dataset_conf['sort'] = False test_dataset = AudioDataset(args.test_data, **dataset_conf, raw_wav=raw_wav) test_data_loader = DataLoader(test_dataset, collate_fn=test_collate_func, shuffle=False, batch_size=1, num_workers=0) # Init asr model from configs model = init_model(configs) # Load dict char_dict = {} with open(args.dict, 'r') as fin: for line in fin: arr = line.strip().split() assert len(arr) == 2 char_dict[int(arr[1])] = arr[0] eos = len(char_dict) - 1 load_checkpoint(model, args.checkpoint) use_cuda = args.gpu >= 0 and torch.cuda.is_available() device = torch.device('cuda' if use_cuda else 'cpu') model = model.to(device)
def main(): args = get_args() logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(message)s') os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu) if args.mode in ['ctc_prefix_beam_search', 'attention_rescoring' ] and args.batch_size > 1: logging.fatal( 'decoding mode {} must be running with batch_size == 1'.format( args.mode)) sys.exit(1) with open(args.config, 'r') as fin: configs = yaml.load(fin, Loader=yaml.FullLoader) if len(args.override_config) > 0: configs = override_config(configs, args.override_config) symbol_table = read_symbol_table(args.dict) test_conf = copy.deepcopy(configs['dataset_conf']) test_conf['filter_conf']['max_length'] = 102400 test_conf['filter_conf']['min_length'] = 0 test_conf['filter_conf']['token_max_length'] = 102400 test_conf['filter_conf']['token_min_length'] = 0 test_conf['filter_conf']['max_output_input_ratio'] = 102400 test_conf['filter_conf']['min_output_input_ratio'] = 0 test_conf['speed_perturb'] = False test_conf['spec_aug'] = False test_conf['spec_sub'] = False test_conf['shuffle'] = False test_conf['sort'] = False if 'fbank_conf' in test_conf: test_conf['fbank_conf']['dither'] = 0.0 elif 'mfcc_conf' in test_conf: test_conf['mfcc_conf']['dither'] = 0.0 test_conf['batch_conf']['batch_type'] = "static" test_conf['batch_conf']['batch_size'] = args.batch_size non_lang_syms = read_non_lang_symbols(args.non_lang_syms) test_dataset = Dataset(args.data_type, args.test_data, symbol_table, test_conf, args.bpe_model, non_lang_syms, partition=False) test_data_loader = DataLoader(test_dataset, batch_size=None, num_workers=0) # Init asr model from configs model = init_model(configs) # Load dict char_dict = {v: k for k, v in symbol_table.items()} eos = len(char_dict) - 1 load_checkpoint(model, args.checkpoint) use_cuda = args.gpu >= 0 and torch.cuda.is_available() device = torch.device('cuda' if use_cuda else 'cpu') model = model.to(device) model.eval() with torch.no_grad(), open(args.result_file, 'w') as fout: for batch_idx, batch in enumerate(test_data_loader): keys, feats, target, feats_lengths, target_lengths = batch feats = feats.to(device) target = target.to(device) feats_lengths = feats_lengths.to(device) target_lengths = target_lengths.to(device) if args.mode == 'attention': hyps, _ = model.recognize( feats, feats_lengths, beam_size=args.beam_size, decoding_chunk_size=args.decoding_chunk_size, num_decoding_left_chunks=args.num_decoding_left_chunks, simulate_streaming=args.simulate_streaming) hyps = [hyp.tolist() for hyp in hyps] elif args.mode == 'ctc_greedy_search': hyps, _ = model.ctc_greedy_search( feats, feats_lengths, decoding_chunk_size=args.decoding_chunk_size, num_decoding_left_chunks=args.num_decoding_left_chunks, simulate_streaming=args.simulate_streaming) elif args.mode == 'rnnt_greedy_search': assert (feats.size(0) == 1) assert 'predictor' in configs hyps = model.greedy_search( feats, feats_lengths, decoding_chunk_size=args.decoding_chunk_size, num_decoding_left_chunks=args.num_decoding_left_chunks, simulate_streaming=args.simulate_streaming) # ctc_prefix_beam_search and attention_rescoring only return one # result in List[int], change it to List[List[int]] for compatible # with other batch decoding mode elif args.mode == 'ctc_prefix_beam_search': assert (feats.size(0) == 1) hyp, _ = model.ctc_prefix_beam_search( feats, feats_lengths, args.beam_size, decoding_chunk_size=args.decoding_chunk_size, num_decoding_left_chunks=args.num_decoding_left_chunks, simulate_streaming=args.simulate_streaming) hyps = [hyp] elif args.mode == 'attention_rescoring': assert (feats.size(0) == 1) hyp, _ = model.attention_rescoring( feats, feats_lengths, args.beam_size, decoding_chunk_size=args.decoding_chunk_size, num_decoding_left_chunks=args.num_decoding_left_chunks, ctc_weight=args.ctc_weight, simulate_streaming=args.simulate_streaming, reverse_weight=args.reverse_weight) hyps = [hyp] for i, key in enumerate(keys): content = [] for w in hyps[i]: if w == eos: break content.append(char_dict[w]) logging.info('{} {}'.format(key, args.connect_symbol.join(content))) fout.write('{} {}\n'.format(key, args.connect_symbol.join(content)))