def _add_local_fusion(self, subnet: ReturnnNetwork, am_output_prob):
prefix_name = self.local_fusion_opts.get('prefix', 'local_fusion')
with_label_smoothing = self.local_fusion_opts.get(
'with_label_smoothing', False)
if self.local_fusion_opts['lm_type'] == 'n_gram':
lm_output_prob = subnet.add_kenlm_layer(
'{}_lm_output_prob'.format(prefix_name),
**self.local_fusion_opts['kenlm_opts'])
else:
lm_subnet = self.local_fusion_opts['lm_subnet']
lm_model = self.local_fusion_opts['lm_model']
vocab_size = self.local_fusion_opts['vocab_size']
# make sure all layers in LM subnet are not trainable
def make_non_trainable(d):
for v in d.values(): # layers
assert isinstance(v, dict)
v.update({'trainable': False})
# Add LM subnetwork.
lm_subnet_copy = copy.deepcopy(lm_subnet)
make_non_trainable(lm_subnet_copy)
lm_subnet_name = '{}_lm_output'.format(prefix_name)
subnet.add_subnetwork(lm_subnet_name, ['prev:output'],
subnetwork_net=lm_subnet_copy,
load_on_init=lm_model,
trainable=False,
n_out=vocab_size)
lm_output_prob = subnet.add_activation_layer(
'{}_lm_output_prob'.format(prefix_name),
lm_subnet_name,
activation='softmax',
target=self.target) # not in log-space
# define new loss criteria
eval_str = "self.network.get_config().typed_value('fusion_eval0_norm')(safe_log(source(0)), safe_log(source(1)))"
if self.local_fusion_opts['lm_type'] == 'n_gram':
eval_str = "self.network.get_config().typed_value('fusion_eval0_norm')(safe_log(source(0)), source(1))"
combo_output_log_prob = subnet.add_eval_layer(
'combo_output_log_prob', [am_output_prob, lm_output_prob],
eval=eval_str)
# local fusion criteria. Eq. (8) in the paper
if with_label_smoothing:
subnet.add_eval_layer(
'combo_output_prob',
combo_output_log_prob,
eval="tf.exp(source(0))",
target=self.target,
loss='ce',
loss_opts={'label_smoothing': self.label_smoothing})
else:
subnet.add_eval_layer('combo_output_prob',
combo_output_log_prob,
eval="tf.exp(source(0))",
target=self.target,
loss='ce')
subnet.add_choice_layer('output',
combo_output_log_prob,
target=self.target,
beam_size=self.beam_size,
initial_output=0,
input_type='log_prob')
def _add_external_LM(self,
subnet_unit: ReturnnNetwork,
am_output_prob,
prior_output_prob=None):
ext_lm_scale = self.ext_lm_opts[
'lm_scale'] if not self.trained_scales else 'lm_scale'
is_recurrent = self.ext_lm_opts.get('is_recurrent', False)
log_lm_prob = False # if lm_prob is already in log-space or not
if 'gram_lm' in self.ext_lm_opts['name']:
log_lm_prob = True # already in log-space
lm_output_prob = subnet_unit.add_kenlm_layer(
'lm_output_prob', **self.ext_lm_opts['kenlm_opts'])
elif is_recurrent:
ext_lm_subnet = self.ext_lm_opts['lm_subnet']
assert isinstance(ext_lm_subnet, dict)
lm_output_prob = self.ext_lm_opts['lm_output_prob_name']
ext_lm_subnet[lm_output_prob]['target'] = self.target
ext_lm_subnet[lm_output_prob][
'loss'] = None # TODO: is this needed?
subnet_unit.update(ext_lm_subnet) # just append
else:
ext_lm_subnet = self.ext_lm_opts['lm_subnet']
assert isinstance(ext_lm_subnet, dict)
ext_lm_model = self.ext_lm_opts['lm_model']
subnet_unit.add_subnetwork('lm_output',
'prev:output',
subnetwork_net=ext_lm_subnet,
load_on_init=ext_lm_model)
lm_output_prob = subnet_unit.add_activation_layer(
'lm_output_prob',
'lm_output',
activation='softmax',
target=self.target)
fusion_str = 'safe_log(source(0)) + {} * '.format(ext_lm_scale)
if log_lm_prob:
fusion_str += 'source(1)'
else:
fusion_str += 'safe_log(source(1))'
fusion_source = [am_output_prob, lm_output_prob]
if prior_output_prob:
fusion_source += [prior_output_prob]
prior_scale = self.prior_lm_opts[
'scale'] if not self.trained_scales else 'prior_scale'
fusion_str += ' - {} * safe_log(source(2))'.format(prior_scale)
if self.coverage_term_scale:
fusion_str += ' + {} * source({})'.format(self.coverage_term_scale,
len(fusion_source))
fusion_source += ['accum_coverage']
if self.trained_scales:
fusion_str = 'source(0) * safe_log(source(1)) + source(2) * safe_log(source(3))'
fusion_source = [
'am_scale', am_output_prob, 'lm_scale', lm_output_prob
]
if prior_output_prob:
fusion_str += ' - source(4) * safe_log(source(5))'
fusion_source += ['prior_scale', prior_output_prob]
subnet_unit.add_eval_layer('combo_output_prob',
source=fusion_source,
eval=fusion_str)
subnet_unit.add_choice_layer('output',
'combo_output_prob',
target=self.target,
beam_size=self.beam_size,
initial_output=0,
input_type='log_prob')