def sub_model(self, model_input, vocab_size, num_mixtures=None,
l2_penalty=1e-8, sub_scope="", distill_labels=None,**unused_params):
num_mixtures = num_mixtures or FLAGS.moe_num_mixtures
class_size = 256
if distill_labels is not None:
class_input = slim.fully_connected(
distill_labels,
class_size,
activation_fn=tf.nn.relu,
weights_regularizer=slim.l2_regularizer(l2_penalty),
scope="class_inputs")
class_input = tf.nn.l2_normalize(class_input, dim=1)
model_input = tf.concat((model_input, class_input),axis=1)
gate_activations = slim.fully_connected(
model_input,
vocab_size * (num_mixtures + 1),
activation_fn=None,
biases_initializer=None,
weights_regularizer=slim.l2_regularizer(l2_penalty),
scope="gates-"+sub_scope)
expert_activations = slim.fully_connected(
model_input,
vocab_size * num_mixtures,
activation_fn=None,
weights_regularizer=slim.l2_regularizer(l2_penalty),
scope="experts-"+sub_scope)
gating_distribution = tf.nn.softmax(tf.reshape(
gate_activations,
[-1, num_mixtures + 1])) # (Batch * #Labels) x (num_mixtures + 1)
expert_distribution = tf.nn.sigmoid(tf.reshape(
expert_activations,
[-1, num_mixtures])) # (Batch * #Labels) x num_mixtures
final_probabilities_by_class_and_batch = tf.reduce_sum(
gating_distribution[:, :num_mixtures] * expert_distribution, 1)
final_probabilities = tf.reshape(final_probabilities_by_class_and_batch,
[-1, vocab_size])
return final_probabilities
评论列表
文章目录
