def __init__(self, num_units, num_heads, ffn_inner_dim, dropout=0.1, attention_dropout=0.1, relu_dropout=0.1, **kwargs): """Initializes the layer. Args: num_units: The number of hidden units. num_heads: The number of heads in the multi-head attention. ffn_inner_dim: The number of units of the inner linear transformation in the feed forward layer. dropout: The probability to drop units from the outputs. attention_dropout: The probability to drop units from the attention. relu_dropout: The probability to drop units from the ReLU activation in the feed forward layer. kwargs: Additional layer arguments. """ super(_SelfAttentionEncoderLayer, self).__init__(**kwargs) self.self_attention = transformer.MultiHeadAttention( num_heads, num_units, dropout=attention_dropout) self.self_attention = common.LayerWrapper(self.self_attention, normalize_input=True, output_dropout=dropout, residual_connection=True) self.ffn = transformer.FeedForwardNetwork(ffn_inner_dim, num_units, dropout=relu_dropout) self.ffn = common.LayerWrapper(self.ffn, normalize_input=True, output_dropout=dropout, residual_connection=True)
def __init__(self, num_units, num_heads, ffn_inner_dim, num_sources=1, dropout=0.1, attention_dropout=0.1, ffn_dropout=0.1, ffn_activation=tf.nn.relu, **kwargs): """Initializes the layer. Args: num_units: The number of hidden units. num_heads: The number of heads in the multi-head attention. ffn_inner_dim: The number of units of the inner linear transformation in the feed forward layer. num_sources: The number of source contexts. dropout: The probability to drop units from the outputs. attention_dropout: The probability to drop units from the attention. ffn_dropout: The probability to drop units from the activation output in the feed forward layer. ffn_activation: The activation function to apply between the two linear transformations of the feed forward layer. **kwargs: Additional layer arguments. """ super(_SelfAttentionDecoderLayer, self).__init__(**kwargs) self.self_attention = transformer.MultiHeadAttention( num_heads, num_units, dropout=attention_dropout, name="masked_multi_head_attention") self.self_attention = transformer.TransformerLayerWrapper( self.self_attention, dropout, name="sub_layer_0") self.attention = [] for i in range(num_sources): attention = transformer.MultiHeadAttention( num_heads, num_units, dropout=attention_dropout, return_attention=num_sources == 1, name="multi_head_attention") attention = transformer.TransformerLayerWrapper( attention, dropout, name="sub_layer_%d" % (i + 1)) self.attention.append(attention) self.ffn = transformer.FeedForwardNetwork(ffn_inner_dim, num_units, dropout=ffn_dropout, activation=ffn_activation, name="feed_forward") self.ffn = transformer.TransformerLayerWrapper(self.ffn, dropout, name="sub_layer_%d" % (num_sources + 1))
def testFeedForwardNetwork(self): ffn = transformer.FeedForwardNetwork(20, 10) x = tf.random.uniform([4, 5, 10]) y = ffn(x) self.assertEqual(y.shape, x.shape)