def testEvolvedTransformerDecoderLayerExtendStep(self): with self.session(use_gpu=True) as sess: np.random.seed(6348575) depth = 4 p = GPipeEvolvedTransformerDecoderLayer.Params() p.name = 'gpipe_evolved_transformer_decoder' p.source_dim = depth p.has_aux_atten = True p.mask_self_atten = True p.tr_double_heads_atten_tpl.num_attention_heads = 2 p.tr_atten_tpl.num_attention_heads = 2 p.transformer_tpl.tr_atten_tpl.num_attention_heads = 2 et_decoder = GPipeEvolvedTransformerDecoderLayer(p) (source_vecs, _, aux_vecs, aux_paddings) = self._testInputs(depth=depth) source_padding = tf.zeros([5, 2]) h1 = et_decoder.FPropDefaultTheta( aux_vecs, aux_paddings, source_vecs, source_padding, None, None, None, None, )[2] h2 = [] double_head_attention_states = py_utils.NestedMap( key=tf.zeros([0, 2, 4]), value=tf.zeros([0, 2, 4])) transformer_layer_states = py_utils.NestedMap( key=tf.zeros([0, 2, 4]), value=tf.zeros([0, 2, 4])) branched_convs_input = tf.zeros([0, 2, 4]) prefix_states = py_utils.NestedMap( double_head_attention_states=double_head_attention_states, transformer_layer_states=transformer_layer_states, branched_convs_input=branched_convs_input) for i in range(5): h, _, prefix_states = et_decoder.ExtendStep( et_decoder.theta, source_vecs[i, :, :], prefix_states, aux_vecs, aux_paddings) h2.append(h) h2 = tf.stack(h2) tf.global_variables_initializer().run() h1_v, h2_v = sess.run([h1, h2]) self.assertAllClose(h1_v, h2_v)
def testEvolvedTransformerDecoderLayerFProp(self): with self.session(use_gpu=True) as sess: np.random.seed(6348575) depth = 4 p = GPipeEvolvedTransformerDecoderLayer.Params() p.name = 'gpipe_evolved_transformer_decoder' p.source_dim = depth p.has_aux_atten = True p.mask_self_atten = True p.tr_double_heads_atten_tpl.num_attention_heads = 2 p.tr_atten_tpl.num_attention_heads = 2 p.transformer_tpl.tr_atten_tpl.num_attention_heads = 2 transformer = GPipeEvolvedTransformerDecoderLayer(p) (source_vecs, source_padding, aux_vecs, aux_paddings, _, _) = self._testInputs(depth=depth) output = transformer.FPropDefaultTheta(aux_vecs, aux_paddings, source_vecs, source_padding, None, None, None, None, None, None) h = output[0] tf.global_variables_initializer().run() actual_layer_output = sess.run([h])[0] tf.logging.info(np.array_repr(actual_layer_output)) # pylint: disable=bad-whitespace # pyformat: disable expected_layer_output = [ [[0.5904724, 0.05267439, 0.89581013, 0.63010913], [0.79584485, 0.07670615, 0.40381077, 0.26504567]], [[0.35448784, 0.28477612, 0.05394353, 0.06531866], [0.44413447, 0.81940264, 0.98786688, 0.35846332]], [[0.66811442, 0.07942203, 0.56781054, 0.83598584], [0.45858502, 0.44949403, 0.06522893, 0.10947803]], [[0.58166796, 0.94657594, 0.17643142, 0.02062288], [0.40596515, 0.01996579, 0.93727112, 0.97478259]], [[0.34873158, 0.0095871, 0.34063059, 0.64620447], [0.70584863, 0.69263214, 0.38247514, 0.28985959]], [[0.66496903, 0.20383522, 0.35497066, 0.66646087], [0.0787568, 0.26172587, 0.23034802, 0.88751978]], [[0.68153989, 0.81061888, 0.90142977, 0.87612331], [0.15129775, 0.56084079, 0.87029755, 0.37908044]] ] # pyformat: enable # pylint: enable=bad-whitespace self.assertAllClose(expected_layer_output, actual_layer_output)