def _create_prior_net(self, subnet_unit: ReturnnNetwork, opts): prior_type = opts.get('type', 'zero') # fixed_ctx_vec_variants = ['zero', 'avg', 'train_avg_ctx', 'train_avg_enc', 'avg_zero', 'trained_vec'] if prior_type == 'zero': # set att context vector to zero prior_att_input = subnet_unit.add_eval_layer( 'zero_att', 'att', eval='tf.zeros_like(source(0))') elif prior_type == 'avg': # during search per utterance self.base_model.network.add_reduce_layer('encoder_mean', 'encoder', mode='mean', axes=['t' ]) # [B, enc-dim] prior_att_input = 'base:encoder_mean' elif prior_type == 'train_avg_ctx': # average all context vectors over training data prior_att_input = subnet_unit.add_constant_layer( 'train_avg_ctx', value=opts['data'], with_batch_dim=True, dtype='float32') elif prior_type == 'train_avg_enc': # average all encoder states over training data prior_att_input = subnet_unit.add_constant_layer( 'train_avg_enc', value=opts['data'], with_batch_dim=True, dtype='float32') elif prior_type == 'mini_lstm': # train a mini LM-like LSTM and use that as prior # example: lstmdim_100-l2_5e-05-recwd_0.0 n_out = 50 l2 = 0.0 recwd = 0.0 if opts.get('prefix_name', None): segs = opts['prefix_name'].split('-') for arg in segs: name, val = arg.split('_', 1) if name == 'lstmdim': n_out = int(val) elif name == 'l2': l2 = float(val) elif name == 'recwd': recwd = float(val) mini_lstm_inputs = opts.get('mini_lstm_inp', 'prev:target_embed').split('+') if len(mini_lstm_inputs) == 1: mini_lstm_inputs = mini_lstm_inputs[0] subnet_unit.add_rec_layer('mini_att_lstm', mini_lstm_inputs, n_out=n_out, l2=l2, rec_weight_dropout=recwd) prior_att_input = subnet_unit.add_linear_layer('mini_att', 'mini_att_lstm', activation=None, n_out=2048, l2=0.0001) elif prior_type == 'adaptive_ctx_vec': # \hat{c}_i = FF(h_i) num_layers = opts.get('num_layers', 3) dim = opts.get('dim', 512) act = opts.get('act', 'relu') x = 's' for i in range(num_layers): x = subnet_unit.add_linear_layer('adaptive_att_%d' % i, x, n_out=dim, **opts.get('att_opts', {})) x = subnet_unit.add_activation_layer('adaptive_att_%d_%s' % (i, act), x, activation=act) prior_att_input = subnet_unit.add_linear_layer('adaptive_att', x, n_out=2048, **opts.get( 'att_opts', {})) elif prior_type == 'trained_vec': prior_att_input = subnet_unit.add_variable_layer( 'trained_vec_att_var', shape=[2048], L2=0.0001) elif prior_type == 'avg_zero': self.base_model.network.add_reduce_layer('encoder_mean', 'encoder', mode='mean', axes=['t' ]) # [B, enc-dim] subnet_unit.add_eval_layer('zero_att', 'att', eval='tf.zeros_like(source(0))') return elif prior_type == 'density_ratio': assert 'lm_subnet' in opts and 'lm_model' in opts return self._add_density_ratio(subnet_unit, lm_subnet=opts['lm_subnet'], lm_model=opts['lm_model']) else: raise ValueError( '{} prior type is not supported'.format(prior_type)) if prior_type != 'mini_lstm': is_first_frame = subnet_unit.add_compare_layer('is_first_frame', source=':i', kind='equal', value=0) zero_att = subnet_unit.add_eval_layer( 'zero_att', 'att', eval='tf.zeros_like(source(0))') prev_att = subnet_unit.add_switch_layer('prev_att', condition=is_first_frame, true_from=zero_att, false_from=prior_att_input) else: prev_att = 'prev:' + prior_att_input assert prev_att is not None key_names = ['s', 'readout_in', 'readout', 'output_prob'] for key_name in key_names: d = copy.deepcopy(subnet_unit[key_name]) # update attention input new_sources = [] from_list = d['from'] if isinstance(from_list, str): from_list = [from_list] assert isinstance(from_list, list) for src in from_list: if 'att' in src: if src.split(':')[0] == 'prev': assert prev_att not in new_sources new_sources += [prev_att] else: new_sources += [prior_att_input] elif src in key_names: new_sources += [('prev:' if 'prev' in src else '') + 'prior_{}'.format(src.split(':')[-1])] else: new_sources += [src] d['from'] = new_sources subnet_unit['prior_{}'.format(key_name)] = d return 'prior_output_prob'
def add_decoder_subnetwork(self, subnet_unit: ReturnnNetwork): # target embedding if self.embed_dropout: # TODO: this is not a good approach. if i want to load a checkpoint from a trained model without embed dropout, # i would need to remap variable name target_embed to target_embed0 to load target_embed0/W subnet_unit.add_linear_layer('target_embed0', 'output', n_out=self.embed_dim, initial_output=0, with_bias=False) subnet_unit.add_dropout_layer('target_embed', 'target_embed0', dropout=self.embed_dropout, dropout_noise_shape={'*': None}) else: subnet_unit.add_linear_layer('target_embed', 'output', n_out=self.embed_dim, initial_output=0, with_bias=False) subnet_unit.add_compare_layer('end', source='output', value=0) # sentence end token # ------ attention location-awareness ------ # # conv-based if self.loc_conv_att_filter_size: assert self.loc_conv_att_num_channels pad_left = subnet_unit.add_pad_layer( 'feedback_pad_left', 'prev:att_weights', axes='s:0', padding=((self.loc_conv_att_filter_size - 1) // 2, 0), value=0) pad_right = subnet_unit.add_pad_layer( 'feedback_pad_right', pad_left, axes='s:0', padding=(0, (self.loc_conv_att_filter_size - 1) // 2), value=0) loc_att_conv = subnet_unit.add_conv_layer( 'loc_att_conv', pad_right, activation=None, with_bias=False, filter_size=(self.loc_conv_att_filter_size, ), padding='valid', n_out=self.loc_conv_att_num_channels, l2=self.l2) subnet_unit.add_linear_layer('weight_feedback', loc_att_conv, activation=None, with_bias=False, n_out=self.enc_key_dim) else: # additive subnet_unit.add_eval_layer( 'accum_att_weights', [ "prev:accum_att_weights", "att_weights", "base:inv_fertility" ], eval='source(0) + source(1) * source(2) * 0.5', out_type={ "dim": self.att_num_heads, "shape": (None, self.att_num_heads) }) subnet_unit.add_linear_layer('weight_feedback', 'prev:accum_att_weights', n_out=self.enc_key_dim, with_bias=False) subnet_unit.add_linear_layer('s_transformed', 's', n_out=self.enc_key_dim, with_bias=False) subnet_unit.add_combine_layer( 'energy_in', ['base:enc_ctx', 'weight_feedback', 's_transformed'], kind='add', n_out=self.enc_key_dim) subnet_unit.add_activation_layer('energy_tanh', 'energy_in', activation='tanh') subnet_unit.add_linear_layer('energy', 'energy_tanh', n_out=self.att_num_heads, with_bias=False) if self.att_dropout: subnet_unit.add_softmax_over_spatial_layer('att_weights0', 'energy') subnet_unit.add_dropout_layer('att_weights', 'att_weights0', dropout=self.att_dropout, dropout_noise_shape={'*': None}) else: if self.relax_att_scale: subnet_unit.add_softmax_over_spatial_layer( 'att_weights0', 'energy') subnet_unit.add_length_layer('encoder_len', 'base:encoder', dtype='float32') # [B] subnet_unit.add_eval_layer('scaled_encoder_len', source=['encoder_len'], eval='{} / source(0)'.format( self.relax_att_scale)) subnet_unit.add_eval_layer( 'att_weights', source=['att_weights0', 'scaled_encoder_len'], eval='{} * source(0) + source(1)'.format( 1 - self.relax_att_scale)) else: subnet_unit.add_softmax_over_spatial_layer( 'att_weights', 'energy') subnet_unit.add_generic_att_layer('att0', weights='att_weights', base='base:enc_value') subnet_unit.add_merge_dims_layer('att', 'att0', axes='except_batch') # LM-like component same as here https://arxiv.org/pdf/2001.07263.pdf lstm_lm_component = None if self.add_lstm_lm: lstm_lm_component = subnet_unit.add_rnn_cell_layer( 'lm_like_s', 'prev:target_embed', n_out=self.lstm_lm_dim, l2=self.l2) lstm_inputs = [] if lstm_lm_component: lstm_inputs += [lstm_lm_component] else: lstm_inputs += ['prev:target_embed'] lstm_inputs += ['prev:att'] if self.dec_state_no_label_ctx: lstm_inputs = ['prev:att'] # no label feedback # LSTM decoder if self.dec_zoneout: subnet_unit.add_rnn_cell_layer('s', lstm_inputs, n_out=self.dec_lstm_num_units, unit='zoneoutlstm', unit_opts={ 'zoneout_factor_cell': 0.15, 'zoneout_factor_output': 0.05 }) else: if self.rec_weight_dropout: # a rec layer with unit nativelstm2 is required to use rec_weight_dropout subnet_unit.add_rec_layer( 's', lstm_inputs, n_out=self.dec_lstm_num_units, l2=self.l2, rec_weight_dropout=self.rec_weight_dropout, unit='NativeLSTM2') else: subnet_unit.add_rnn_cell_layer('s', lstm_inputs, n_out=self.dec_lstm_num_units, l2=self.l2) # AM softmax output layer if self.dec_state_no_label_ctx and self.add_lstm_lm: subnet_unit.add_linear_layer( 'readout_in', ["lm_like_s", "prev:target_embed", "att"], n_out=self.dec_output_num_units) if self.add_no_label_ctx_s_to_output: subnet_unit.add_linear_layer( 'readout_in', ["lm_like_s", "s", "prev:target_embed", "att"], n_out=self.dec_output_num_units) else: subnet_unit.add_linear_layer('readout_in', ["s", "prev:target_embed", "att"], n_out=self.dec_output_num_units) if self.reduceout: subnet_unit.add_reduceout_layer('readout', 'readout_in') else: subnet_unit.add_copy_layer('readout', 'readout_in') if self.local_fusion_opts: output_prob = subnet_unit.add_softmax_layer('output_prob', 'readout', l2=self.l2, target=self.target, dropout=self.dropout) self._add_local_fusion(subnet_unit, am_output_prob=output_prob) elif self.mwer: # only MWER so CE is disabled output_prob = subnet_unit.add_softmax_layer('output_prob', 'readout', l2=self.l2, target=self.target, dropout=self.dropout) else: ce_loss_opts = {'label_smoothing': self.label_smoothing} if self.ce_loss_scale: ce_loss_opts['scale'] = self.ce_loss_scale output_prob = subnet_unit.add_softmax_layer('output_prob', 'readout', l2=self.l2, loss='ce', loss_opts=ce_loss_opts, target=self.target, dropout=self.dropout) # do not load the bias if self.remove_softmax_bias: subnet_unit['output_prob']['with_bias'] = False # for prior LM estimation prior_output_prob = None if self.prior_lm_opts: prior_output_prob = self._create_prior_net( subnet_unit, self.prior_lm_opts ) # this require preload_from_files in config # Beam search # only support shallow fusion for now if self.ext_lm_opts: self._add_external_LM(subnet_unit, output_prob, prior_output_prob) else: if self.coverage_term_scale: output_prob = subnet_unit.add_eval_layer( 'combo_output_prob', eval='safe_log(source(0)) + {} * source(1)'.format( self.coverage_term_scale), source=['output_prob', 'accum_coverage']) input_type = 'log_prob' else: output_prob = 'output_prob' input_type = None if self.length_norm: subnet_unit.add_choice_layer('output', output_prob, target=self.target, beam_size=self.beam_size, initial_output=0, input_type=input_type) else: subnet_unit.add_choice_layer( 'output', output_prob, target=self.target, beam_size=self.beam_size, initial_output=0, length_normalization=self.length_norm, input_type=input_type) if self.ilmt_opts: self._create_ilmt_net(subnet_unit) # recurrent subnetwork dec_output = self.network.add_subnet_rec_layer( 'output', unit=subnet_unit.get_net(), target=self.target, source=self.source) return dec_output
def add_decoder_subnetwork(self, subnet_unit: ReturnnNetwork): subnet_unit.add_compare_layer('end', source='output', value=0) # sentence end token # target embedding subnet_unit.add_linear_layer( 'target_embed0', 'output', n_out=self.embed_dim, initial_output=0, with_bias=False, l2=self.l2, forward_weights_init=self.embed_weight_init) subnet_unit.add_dropout_layer('target_embed', 'target_embed0', dropout=self.embed_dropout, dropout_noise_shape={'*': None}) # attention att = AttentionMechanism( enc_key_dim=self.enc_key_dim, att_num_heads=self.att_num_heads, att_dropout=self.att_dropout, l2=self.l2, loc_filter_size=self.loc_conv_att_filter_size, loc_num_channels=self.loc_conv_att_num_channels) subnet_unit.update(att.create()) # LM-like component same as here https://arxiv.org/pdf/2001.07263.pdf lstm_lm_component = None if self.add_lstm_lm: lstm_lm_component = subnet_unit.add_rnn_cell_layer( 'lm_like_s', 'prev:target_embed', n_out=self.lstm_lm_dim, l2=self.l2) lstm_inputs = [] if lstm_lm_component: lstm_inputs += [lstm_lm_component] else: lstm_inputs += ['prev:target_embed'] lstm_inputs += ['prev:att'] # LSTM decoder (or decoder state) if self.dec_zoneout: subnet_unit.add_rnn_cell_layer('s', lstm_inputs, n_out=self.dec_lstm_num_units, l2=self.l2, weights_init=self.lstm_weights_init, unit='zoneoutlstm', unit_opts={ 'zoneout_factor_cell': 0.15, 'zoneout_factor_output': 0.05 }) else: if self.rec_weight_dropout: # a rec layer with unit nativelstm2 is required to use rec_weight_dropout subnet_unit.add_rec_layer( 's', lstm_inputs, n_out=self.dec_lstm_num_units, l2=self.l2, unit='NativeLSTM2', rec_weight_dropout=self.rec_weight_dropout, weights_init=self.lstm_weights_init) else: subnet_unit.add_rnn_cell_layer( 's', lstm_inputs, n_out=self.dec_lstm_num_units, l2=self.l2, weights_init=self.lstm_weights_init) # ASR softmax output layer subnet_unit.add_linear_layer('readout_in', ["s", "prev:target_embed", "att"], n_out=self.dec_output_num_units, l2=self.l2) if self.reduceout: subnet_unit.add_reduceout_layer('readout', 'readout_in') else: subnet_unit.add_copy_layer('readout', 'readout_in') output_prob = subnet_unit.add_softmax_layer( 'output_prob', 'readout', l2=self.l2, loss='ce', loss_opts={'label_smoothing': self.label_smoothing}, target=self.target, dropout=self.dropout) subnet_unit.add_choice_layer('output', output_prob, target=self.target, beam_size=self.beam_size, initial_output=0) # recurrent subnetwork dec_output = self.network.add_subnet_rec_layer( 'output', unit=subnet_unit.get_net(), target=self.target, source=self.source) return dec_output