def __init__(self, args, save_path=None, idx2token=None): super(ModelBase, self).__init__() self.save_path = save_path # for encoder, decoder self.input_type = args.input_type self.input_dim = args.input_dim self.enc_type = args.enc_type self.enc_n_units = args.enc_n_units if args.enc_type in ['blstm', 'bgru', 'conv_blstm', 'conv_bgru']: self.enc_n_units *= 2 self.dec_type = args.dec_type # for OOV resolution self.enc_n_layers = args.enc_n_layers self.enc_n_layers_sub1 = args.enc_n_layers_sub1 self.subsample = [int(s) for s in args.subsample.split('_')] # for decoder self.vocab = args.vocab self.vocab_sub1 = args.vocab_sub1 self.vocab_sub2 = args.vocab_sub2 self.blank = 0 self.unk = 1 self.eos = 2 self.pad = 3 # NOTE: reserved in advance # for the sub tasks self.main_weight = 1 - args.sub1_weight - args.sub2_weight self.sub1_weight = args.sub1_weight self.sub2_weight = args.sub2_weight self.mtl_per_batch = args.mtl_per_batch self.task_specific_layer = args.task_specific_layer # for CTC self.ctc_weight = min(args.ctc_weight, self.main_weight) self.ctc_weight_sub1 = min(args.ctc_weight_sub1, self.sub1_weight) self.ctc_weight_sub2 = min(args.ctc_weight_sub2, self.sub2_weight) # for backward decoder self.bwd_weight = min(args.bwd_weight, self.main_weight) self.fwd_weight = self.main_weight - self.bwd_weight - self.ctc_weight self.fwd_weight_sub1 = self.sub1_weight - self.ctc_weight_sub1 self.fwd_weight_sub2 = self.sub2_weight - self.ctc_weight_sub2 # for MBR self.mbr_training = args.mbr_training self.recog_params = vars(args) self.idx2token = idx2token # Feature extraction self.gaussian_noise = args.gaussian_noise self.n_stacks = args.n_stacks self.n_skips = args.n_skips self.n_splices = args.n_splices self.use_specaug = args.n_freq_masks > 0 or args.n_time_masks > 0 self.specaug = None self.flip_time_prob = args.flip_time_prob self.flip_freq_prob = args.flip_freq_prob self.weight_noise = args.weight_noise if self.use_specaug: assert args.n_stacks == 1 and args.n_skips == 1 assert args.n_splices == 1 self.specaug = SpecAugment(F=args.freq_width, T=args.time_width, n_freq_masks=args.n_freq_masks, n_time_masks=args.n_time_masks, p=args.time_width_upper) # Frontend self.ssn = None if args.sequence_summary_network: assert args.input_type == 'speech' self.ssn = SequenceSummaryNetwork(args.input_dim, n_units=512, n_layers=3, bottleneck_dim=100, dropout=0, param_init=args.param_init) # Encoder self.enc = build_encoder(args) if args.freeze_encoder: for p in self.enc.parameters(): p.requires_grad = False # main task external_lm = None directions = [] if self.fwd_weight > 0 or (self.bwd_weight == 0 and self.ctc_weight > 0): directions.append('fwd') if self.bwd_weight > 0: directions.append('bwd') for dir in directions: # Load the LM for LM fusion and decoder initialization if args.external_lm and dir == 'fwd': external_lm = RNNLM(args.lm_conf) load_checkpoint(external_lm, args.external_lm) # freeze LM parameters for n, p in external_lm.named_parameters(): p.requires_grad = False # Decoder special_symbols = { 'blank': self.blank, 'unk': self.unk, 'eos': self.eos, 'pad': self.pad, } dec = build_decoder( args, special_symbols, self.enc.output_dim, args.vocab, self.ctc_weight, args.ctc_fc_list, self.main_weight - self.bwd_weight if dir == 'fwd' else self.bwd_weight, external_lm) setattr(self, 'dec_' + dir, dec) # sub task for sub in ['sub1', 'sub2']: if getattr(self, sub + '_weight') > 0: dec_sub = build_decoder(args, special_symbols, self.enc.output_dim, getattr(self, 'vocab_' + sub), getattr(self, 'ctc_weight_' + sub), getattr(args, 'ctc_fc_list_' + sub), getattr(self, sub + '_weight'), external_lm) setattr(self, 'dec_fwd_' + sub, dec_sub) if args.input_type == 'text': if args.vocab == args.vocab_sub1: # Share the embedding layer between input and output self.embed = dec.embed else: self.embed = nn.Embedding(args.vocab_sub1, args.emb_dim, padding_idx=self.pad) self.dropout_emb = nn.Dropout(p=args.dropout_emb) # Recurrent weights are orthogonalized if args.rec_weight_orthogonal: self.reset_parameters(args.param_init, dist='orthogonal', keys=['rnn', 'weight']) # Initialize bias in forget gate with 1 # self.init_forget_gate_bias_with_one() # Fix all parameters except for the gating parts in deep fusion if args.lm_fusion == 'deep' and external_lm is not None: for n, p in self.named_parameters(): if 'output' in n or 'output_bn' in n or 'linear' in n: p.requires_grad = True else: p.requires_grad = False
def __init__(self, args, save_path=None, idx2token=None): super(ModelBase, self).__init__() self.save_path = save_path # for encoder, decoder self.input_type = args.input_type self.input_dim = args.input_dim self.enc_type = args.enc_type self.dec_type = args.dec_type # for OOV resolution self.enc_n_layers = args.enc_n_layers self.enc_n_layers_sub1 = args.enc_n_layers_sub1 self.subsample = [int(s) for s in args.subsample.split('_')] # for decoder self.vocab = args.vocab self.vocab_sub1 = args.vocab_sub1 self.vocab_sub2 = args.vocab_sub2 self.blank = 0 self.unk = 1 self.eos = 2 self.pad = 3 # NOTE: reserved in advance # for the sub tasks self.main_weight = args.total_weight - args.sub1_weight - args.sub2_weight self.sub1_weight = args.sub1_weight self.sub2_weight = args.sub2_weight self.mtl_per_batch = args.mtl_per_batch self.task_specific_layer = args.task_specific_layer # for CTC self.ctc_weight = min(args.ctc_weight, self.main_weight) self.ctc_weight_sub1 = min(args.ctc_weight_sub1, self.sub1_weight) self.ctc_weight_sub2 = min(args.ctc_weight_sub2, self.sub2_weight) # for backward decoder self.bwd_weight = min(args.bwd_weight, self.main_weight) self.fwd_weight = self.main_weight - self.bwd_weight - self.ctc_weight self.fwd_weight_sub1 = self.sub1_weight - self.ctc_weight_sub1 self.fwd_weight_sub2 = self.sub2_weight - self.ctc_weight_sub2 # for MBR self.mbr_training = args.mbr_training self.recog_params = vars(args) self.idx2token = idx2token # for discourse-aware model self.utt_id_prev = None # Feature extraction self.input_noise_std = args.input_noise_std self.n_stacks = args.n_stacks self.n_skips = args.n_skips self.n_splices = args.n_splices self.weight_noise_std = args.weight_noise_std self.specaug = None if args.n_freq_masks > 0 or args.n_time_masks > 0: assert args.n_stacks == 1 and args.n_skips == 1 assert args.n_splices == 1 self.specaug = SpecAugment( F=args.freq_width, T=args.time_width, n_freq_masks=args.n_freq_masks, n_time_masks=args.n_time_masks, p=args.time_width_upper, adaptive_number_ratio=args.adaptive_number_ratio, adaptive_size_ratio=args.adaptive_size_ratio, max_n_time_masks=args.max_n_time_masks) # Frontend self.ssn = None if args.sequence_summary_network: assert args.input_type == 'speech' self.ssn = SequenceSummaryNetwork(args.input_dim, n_units=512, n_layers=3, bottleneck_dim=100, dropout=0, param_init=args.param_init) # Encoder self.enc = build_encoder(args) if args.freeze_encoder: for n, p in self.enc.named_parameters(): if 'bridge' in n or 'sub1' in n: continue p.requires_grad = False logger.info('freeze %s' % n) special_symbols = { 'blank': self.blank, 'unk': self.unk, 'eos': self.eos, 'pad': self.pad, } # main task external_lm = None directions = [] if self.fwd_weight > 0 or (self.bwd_weight == 0 and self.ctc_weight > 0): directions.append('fwd') if self.bwd_weight > 0: directions.append('bwd') for dir in directions: # Load the LM for LM fusion and decoder initialization if args.external_lm and dir == 'fwd': external_lm = RNNLM(args.lm_conf) load_checkpoint(args.external_lm, external_lm) # freeze LM parameters for n, p in external_lm.named_parameters(): p.requires_grad = False # Decoder dec = build_decoder( args, special_symbols, self.enc.output_dim, args.vocab, self.ctc_weight, self.main_weight - self.bwd_weight if dir == 'fwd' else self.bwd_weight, external_lm) setattr(self, 'dec_' + dir, dec) # sub task for sub in ['sub1', 'sub2']: if getattr(self, sub + '_weight') > 0: args_sub = copy.deepcopy(args) if hasattr(args, 'dec_config_' + sub): for k, v in getattr(args, 'dec_config_' + sub).items(): setattr(args_sub, k, v) # NOTE: Other parameters are the same as the main decoder dec_sub = build_decoder(args_sub, special_symbols, getattr(self.enc, 'output_dim_' + sub), getattr(self, 'vocab_' + sub), getattr(self, 'ctc_weight_' + sub), getattr(self, sub + '_weight'), external_lm) setattr(self, 'dec_fwd_' + sub, dec_sub) if args.input_type == 'text': if args.vocab == args.vocab_sub1: # Share the embedding layer between input and output self.embed = dec.embed else: self.embed = nn.Embedding(args.vocab_sub1, args.emb_dim, padding_idx=self.pad) self.dropout_emb = nn.Dropout(p=args.dropout_emb) # Initialize bias in forget gate with 1 # self.init_forget_gate_bias_with_one() # Fix all parameters except for the gating parts in deep fusion if args.lm_fusion == 'deep' and external_lm is not None: for n, p in self.named_parameters(): if 'output' in n or 'output_bn' in n or 'linear' in n: p.requires_grad = True else: p.requires_grad = False