transformer.py 24 KB
Newer Older
Hongkun Yu's avatar
Hongkun Yu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Keras-based transformer block layer."""
16
# pylint: disable=g-classes-have-attributes
Hongkun Yu's avatar
Hongkun Yu committed
17

Chen Chen's avatar
Chen Chen committed
18
import gin
Hongkun Yu's avatar
Hongkun Yu committed
19
20
21
import tensorflow as tf

from official.nlp.modeling.layers import attention
22
from official.nlp.modeling.layers import multi_channel_attention
23
from official.nlp.modeling.layers.util import tf_function_if_eager
Hongkun Yu's avatar
Hongkun Yu committed
24
25
26
27
28
29
30
31
32


@tf.keras.utils.register_keras_serializable(package="Text")
class Transformer(tf.keras.layers.Layer):
  """Transformer layer.

  This layer implements the Transformer from "Attention Is All You Need".
  (https://arxiv.org/abs/1706.03762).

33
  Arguments:
Hongkun Yu's avatar
Hongkun Yu committed
34
35
36
37
38
    num_attention_heads: Number of attention heads.
    intermediate_size: Size of the intermediate layer.
    intermediate_activation: Activation for the intermediate layer.
    dropout_rate: Dropout probability for the post-attention and output dropout.
    attention_dropout_rate: Dropout probability for within the attention layer.
39
40
    output_range: the sequence output range, [0, output_range) by slicing the
      target sequence. `None` means the target sequence is not sliced.
Hongkun Yu's avatar
Hongkun Yu committed
41
42
43
44
45
46
47
    kernel_initializer: Initializer for dense layer kernels.
    bias_initializer: Initializer for dense layer biases.
    kernel_regularizer: Regularizer for dense layer kernels.
    bias_regularizer: Regularizer for dense layer biases.
    activity_regularizer: Regularizer for dense layer activity.
    kernel_constraint: Constraint for dense layer kernels.
    bias_constraint: Constraint for dense layer kernels.
xinliupitt's avatar
xinliupitt committed
48
49
    use_bias: Whether to enable use_bias in attention layer. If set False,
      use_bias in attention layer is disabled.
xinliupitt's avatar
xinliupitt committed
50
    norm_first: Whether to normalize inputs to attention and intermediate dense
51
52
      layers. If set False, output of attention and intermediate dense layers is
      normalized.
xinliupitt's avatar
xinliupitt committed
53
    norm_epsilon: Epsilon value to initialize normalization layers.
54
    intermediate_dropout: Dropout probability for intermediate_dropout_layer.
xinliupitt's avatar
xinliupitt committed
55
56
    attention_initializer: Initializer for kernels of attention layers. If set
      `None`, attention layers use kernel_initializer as initializer for kernel.
Hongkun Yu's avatar
Hongkun Yu committed
57
58
59
60
61
62
63
64
  """

  def __init__(self,
               num_attention_heads,
               intermediate_size,
               intermediate_activation,
               dropout_rate=0.0,
               attention_dropout_rate=0.0,
65
               output_range=None,
Hongkun Yu's avatar
Hongkun Yu committed
66
67
68
69
70
71
72
               kernel_initializer="glorot_uniform",
               bias_initializer="zeros",
               kernel_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               bias_constraint=None,
xinliupitt's avatar
xinliupitt committed
73
74
75
               use_bias=True,
               norm_first=False,
               norm_epsilon=1e-12,
xinliupitt's avatar
xinliupitt committed
76
               intermediate_dropout=0.0,
xinliupitt's avatar
xinliupitt committed
77
               attention_initializer=None,
Hongkun Yu's avatar
Hongkun Yu committed
78
79
80
81
82
83
84
85
               **kwargs):
    super(Transformer, self).__init__(**kwargs)

    self._num_heads = num_attention_heads
    self._intermediate_size = intermediate_size
    self._intermediate_activation = intermediate_activation
    self._attention_dropout_rate = attention_dropout_rate
    self._dropout_rate = dropout_rate
86
    self._output_range = output_range
Hongkun Yu's avatar
Hongkun Yu committed
87
88
89
90
    self._kernel_initializer = tf.keras.initializers.get(kernel_initializer)
    self._bias_initializer = tf.keras.initializers.get(bias_initializer)
    self._kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
    self._bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
Hongkun Yu's avatar
Hongkun Yu committed
91
    self._activity_regularizer = tf.keras.regularizers.get(activity_regularizer)
Hongkun Yu's avatar
Hongkun Yu committed
92
93
    self._kernel_constraint = tf.keras.constraints.get(kernel_constraint)
    self._bias_constraint = tf.keras.constraints.get(bias_constraint)
xinliupitt's avatar
xinliupitt committed
94
95
96
    self._use_bias = use_bias
    self._norm_first = norm_first
    self._norm_epsilon = norm_epsilon
xinliupitt's avatar
xinliupitt committed
97
    self._intermediate_dropout = intermediate_dropout
xinliupitt's avatar
xinliupitt committed
98
    if attention_initializer:
xinliupitt's avatar
xinliupitt committed
99
100
      self._attention_initializer = tf.keras.initializers.get(
          attention_initializer)
xinliupitt's avatar
xinliupitt committed
101
102
    else:
      self._attention_initializer = self._kernel_initializer
Hongkun Yu's avatar
Hongkun Yu committed
103
104
105
106

  def build(self, input_shape):
    input_tensor = input_shape[0] if len(input_shape) == 2 else input_shape
    input_tensor_shape = tf.TensorShape(input_tensor)
107
    if len(input_tensor_shape.as_list()) != 3:
Hongkun Yu's avatar
Hongkun Yu committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
      raise ValueError("TransformerLayer expects a three-dimensional input of "
                       "shape [batch, sequence, width].")
    batch_size, sequence_length, hidden_size = input_tensor_shape

    if len(input_shape) == 2:
      mask_tensor_shape = tf.TensorShape(input_shape[1])
      expected_mask_tensor_shape = tf.TensorShape(
          [batch_size, sequence_length, sequence_length])
      if not expected_mask_tensor_shape.is_compatible_with(mask_tensor_shape):
        raise ValueError("When passing a mask tensor to TransformerLayer, the "
                         "mask tensor must be of shape [batch, "
                         "sequence_length, sequence_length] (here %s). Got a "
                         "mask tensor of shape %s." %
                         (expected_mask_tensor_shape, mask_tensor_shape))
    if hidden_size % self._num_heads != 0:
      raise ValueError(
          "The input size (%d) is not a multiple of the number of attention "
          "heads (%d)" % (hidden_size, self._num_heads))
    self._attention_head_size = int(hidden_size // self._num_heads)
127
    common_kwargs = dict(
Hongkun Yu's avatar
Hongkun Yu committed
128
129
130
131
132
        bias_initializer=self._bias_initializer,
        kernel_regularizer=self._kernel_regularizer,
        bias_regularizer=self._bias_regularizer,
        activity_regularizer=self._activity_regularizer,
        kernel_constraint=self._kernel_constraint,
133
        bias_constraint=self._bias_constraint)
134
    self._attention_layer = tf.keras.layers.MultiHeadAttention(
135
        num_heads=self._num_heads,
136
        key_dim=self._attention_head_size,
137
        dropout=self._attention_dropout_rate,
xinliupitt's avatar
xinliupitt committed
138
        use_bias=self._use_bias,
xinliupitt's avatar
xinliupitt committed
139
        kernel_initializer=self._attention_initializer,
140
141
        name="self_attention",
        **common_kwargs)
Hongkun Yu's avatar
Hongkun Yu committed
142
    self._attention_dropout = tf.keras.layers.Dropout(rate=self._dropout_rate)
Zongwei Zhou's avatar
Zongwei Zhou committed
143
144
    # Use float32 in layernorm for numeric stability.
    # It is probably safe in mixed_float16, but we haven't validated this yet.
Hongkun Yu's avatar
Hongkun Yu committed
145
146
    self._attention_layer_norm = (
        tf.keras.layers.LayerNormalization(
Chen Chen's avatar
Chen Chen committed
147
148
            name="self_attention_layer_norm",
            axis=-1,
xinliupitt's avatar
xinliupitt committed
149
            epsilon=self._norm_epsilon,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
150
            dtype=tf.float32))
151
152
153
154
    self._intermediate_dense = tf.keras.layers.experimental.EinsumDense(
        "abc,cd->abd",
        output_shape=(None, self._intermediate_size),
        bias_axes="d",
xinliupitt's avatar
xinliupitt committed
155
        kernel_initializer=self._kernel_initializer,
156
157
        name="intermediate",
        **common_kwargs)
158
159
160
161
162
163
    policy = tf.keras.mixed_precision.experimental.global_policy()
    if policy.name == "mixed_bfloat16":
      # bfloat16 causes BERT with the LAMB optimizer to not converge
      # as well, so we use float32.
      # TODO(b/154538392): Investigate this.
      policy = tf.float32
Chen Chen's avatar
Chen Chen committed
164
    self._intermediate_activation_layer = tf.keras.layers.Activation(
165
        self._intermediate_activation, dtype=policy)
166
167
    self._intermediate_dropout_layer = tf.keras.layers.Dropout(
        rate=self._intermediate_dropout)
168
169
170
171
172
    self._output_dense = tf.keras.layers.experimental.EinsumDense(
        "abc,cd->abd",
        output_shape=(None, hidden_size),
        bias_axes="d",
        name="output",
xinliupitt's avatar
xinliupitt committed
173
        kernel_initializer=self._kernel_initializer,
174
        **common_kwargs)
Hongkun Yu's avatar
Hongkun Yu committed
175
    self._output_dropout = tf.keras.layers.Dropout(rate=self._dropout_rate)
Zongwei Zhou's avatar
Zongwei Zhou committed
176
    # Use float32 in layernorm for numeric stability.
Hongkun Yu's avatar
Hongkun Yu committed
177
    self._output_layer_norm = tf.keras.layers.LayerNormalization(
178
179
180
        name="output_layer_norm",
        axis=-1,
        epsilon=self._norm_epsilon,
xinliupitt's avatar
xinliupitt committed
181
        dtype=tf.float32)
Hongkun Yu's avatar
Hongkun Yu committed
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196

    super(Transformer, self).build(input_shape)

  def get_config(self):
    config = {
        "num_attention_heads":
            self._num_heads,
        "intermediate_size":
            self._intermediate_size,
        "intermediate_activation":
            self._intermediate_activation,
        "dropout_rate":
            self._dropout_rate,
        "attention_dropout_rate":
            self._attention_dropout_rate,
197
198
        "output_range":
            self._output_range,
Hongkun Yu's avatar
Hongkun Yu committed
199
200
201
202
203
204
205
206
207
208
209
210
211
        "kernel_initializer":
            tf.keras.initializers.serialize(self._kernel_initializer),
        "bias_initializer":
            tf.keras.initializers.serialize(self._bias_initializer),
        "kernel_regularizer":
            tf.keras.regularizers.serialize(self._kernel_regularizer),
        "bias_regularizer":
            tf.keras.regularizers.serialize(self._bias_regularizer),
        "activity_regularizer":
            tf.keras.regularizers.serialize(self._activity_regularizer),
        "kernel_constraint":
            tf.keras.constraints.serialize(self._kernel_constraint),
        "bias_constraint":
xinliupitt's avatar
xinliupitt committed
212
213
214
215
216
217
            tf.keras.constraints.serialize(self._bias_constraint),
        "use_bias":
            self._use_bias,
        "norm_first":
            self._norm_first,
        "norm_epsilon":
xinliupitt's avatar
xinliupitt committed
218
219
            self._norm_epsilon,
        "intermediate_dropout":
xinliupitt's avatar
xinliupitt committed
220
221
            self._intermediate_dropout,
        "attention_initializer":
xinliupitt's avatar
xinliupitt committed
222
            tf.keras.initializers.serialize(self._attention_initializer)
Hongkun Yu's avatar
Hongkun Yu committed
223
224
225
226
227
228
229
230
231
232
    }
    base_config = super(Transformer, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def call(self, inputs):
    if isinstance(inputs, (list, tuple)) and len(inputs) == 2:
      input_tensor, attention_mask = inputs
    else:
      input_tensor, attention_mask = (inputs, None)

233
234
235
236
    if self._output_range:
      target_tensor = input_tensor[:, 0:self._output_range, :]
      attention_mask = attention_mask[:, 0:self._output_range, :]
    else:
xinliupitt's avatar
xinliupitt committed
237
238
239
      if self._norm_first:
        source_tensor = input_tensor
        input_tensor = self._attention_layer_norm(input_tensor)
240
      target_tensor = input_tensor
Hongkun Yu's avatar
Hongkun Yu committed
241

242
243
    attention_output = self._attention_layer(
        query=target_tensor, value=input_tensor, attention_mask=attention_mask)
244
    attention_output = self._attention_dropout(attention_output)
xinliupitt's avatar
xinliupitt committed
245
246
247
248
249
250
251
252
    if self._norm_first:
      attention_output = source_tensor + attention_output
    else:
      attention_output = self._attention_layer_norm(target_tensor +
                                                    attention_output)
    if self._norm_first:
      source_attention_output = attention_output
      attention_output = self._output_layer_norm(attention_output)
253
254
255
    intermediate_output = self._intermediate_dense(attention_output)
    intermediate_output = self._intermediate_activation_layer(
        intermediate_output)
256
    intermediate_output = self._intermediate_dropout_layer(intermediate_output)
257
258
259
260
261
262
    layer_output = self._output_dense(intermediate_output)
    layer_output = self._output_dropout(layer_output)
    # During mixed precision training, attention_output is from layer norm and
    # is always fp32 for now. Cast layer_output to fp32 for the subsequent
    # add.
    layer_output = tf.cast(layer_output, tf.float32)
xinliupitt's avatar
xinliupitt committed
263
264
265
266
    if self._norm_first:
      layer_output = source_attention_output + layer_output
    else:
      layer_output = self._output_layer_norm(layer_output + attention_output)
267
268

    return layer_output
269
270


Chen Chen's avatar
Chen Chen committed
271
272
@tf.keras.utils.register_keras_serializable(package="Text")
@gin.configurable
273
274
275
276
277
class CompiledTransformer(Transformer):

  @tf_function_if_eager(experimental_compile=True)
  def call(self, inputs):
    return super(CompiledTransformer, self).call(inputs)
278
279
280
281
282
283
284
285
286
287


@tf.keras.utils.register_keras_serializable(package="Text")
class TransformerDecoderLayer(tf.keras.layers.Layer):
  """Single transformer layer for decoder.

  It has three sub-layers:
  (1) a multi-head self-attention mechanism.
  (2) a encoder-decoder attention.
  (3) a positionwise fully connected feed-forward network.
Hongkun Yu's avatar
Hongkun Yu committed
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303

  Arguments:
    num_attention_heads: Number of attention heads.
    intermediate_size: Size of the intermediate layer.
    intermediate_activation: Activation for the intermediate layer.
    dropout_rate: Dropout probability for the post-attention and output dropout.
    attention_dropout_rate: Dropout probability for within the attention layer.
    multi_channel_cross_attention: Whether to use `MultiChannelAttention` for
      cross-attention between target sequences and source sequences.
    kernel_initializer: Initializer for dense layer kernels.
    bias_initializer: Initializer for dense layer biases.
    kernel_regularizer: Regularizer for dense layer kernels.
    bias_regularizer: Regularizer for dense layer biases.
    activity_regularizer: Regularizer for dense layer activity.
    kernel_constraint: Constraint for dense layer kernels.
    bias_constraint: Constraint for dense layer kernels.
xinliupitt's avatar
xinliupitt committed
304
305
    use_bias: Whether to enable use_bias in attention layer. If set False,
      use_bias in attention layer is disabled.
xinliupitt's avatar
xinliupitt committed
306
    norm_first: Whether to normalize inputs to attention and intermediate dense
307
308
      layers. If set False, output of attention and intermediate dense layers is
      normalized.
xinliupitt's avatar
xinliupitt committed
309
    norm_epsilon: Epsilon value to initialize normalization layers.
310
    intermediate_dropout: Dropout probability for intermediate_dropout_layer.
xinliupitt's avatar
xinliupitt committed
311
312
    attention_initializer: Initializer for kernels of attention layers. If set
      `None`, attention layers use kernel_initializer as initializer for kernel.
313
314
315
  """

  def __init__(self,
Hongkun Yu's avatar
Hongkun Yu committed
316
317
318
319
320
               num_attention_heads,
               intermediate_size,
               intermediate_activation,
               dropout_rate=0.0,
               attention_dropout_rate=0.0,
321
               multi_channel_cross_attention=False,
Hongkun Yu's avatar
Hongkun Yu committed
322
323
324
325
326
327
328
               kernel_initializer="glorot_uniform",
               bias_initializer="zeros",
               kernel_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               bias_constraint=None,
xinliupitt's avatar
xinliupitt committed
329
330
331
               use_bias=True,
               norm_first=False,
               norm_epsilon=1e-12,
xinliupitt's avatar
xinliupitt committed
332
               intermediate_dropout=0.0,
xinliupitt's avatar
xinliupitt committed
333
               attention_initializer=None,
334
335
336
337
338
339
               **kwargs):
    super(TransformerDecoderLayer, self).__init__(**kwargs)
    self.num_attention_heads = num_attention_heads
    self.intermediate_size = intermediate_size
    self.intermediate_activation = tf.keras.activations.get(
        intermediate_activation)
Hongkun Yu's avatar
Hongkun Yu committed
340
341
    self.dropout_rate = dropout_rate
    self.attention_dropout_rate = attention_dropout_rate
342
    self.multi_channel_cross_attention = multi_channel_cross_attention
Hongkun Yu's avatar
Hongkun Yu committed
343
344
345
346
347
348
349
    self._kernel_initializer = tf.keras.initializers.get(kernel_initializer)
    self._bias_initializer = tf.keras.initializers.get(bias_initializer)
    self._kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
    self._bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
    self._activity_regularizer = tf.keras.regularizers.get(activity_regularizer)
    self._kernel_constraint = tf.keras.constraints.get(kernel_constraint)
    self._bias_constraint = tf.keras.constraints.get(bias_constraint)
xinliupitt's avatar
xinliupitt committed
350
351
352
    self._use_bias = use_bias
    self._norm_first = norm_first
    self._norm_epsilon = norm_epsilon
xinliupitt's avatar
xinliupitt committed
353
    self._intermediate_dropout = intermediate_dropout
xinliupitt's avatar
xinliupitt committed
354
    if attention_initializer:
xinliupitt's avatar
xinliupitt committed
355
356
      self._attention_initializer = tf.keras.initializers.get(
          attention_initializer)
xinliupitt's avatar
xinliupitt committed
357
358
    else:
      self._attention_initializer = self._kernel_initializer
359
360
361
362
363
    if self.multi_channel_cross_attention:
      self._cross_attention_cls = multi_channel_attention.MultiChannelAttention
    else:
      self._cross_attention_cls = attention.MultiHeadAttention

Hongkun Yu's avatar
Hongkun Yu committed
364
365
  def build(self, input_shape):
    target_tensor_shape = tf.TensorShape(input_shape[0])
366
    if len(target_tensor_shape.as_list()) != 3:
Hongkun Yu's avatar
Hongkun Yu committed
367
368
369
370
      raise ValueError("TransformerLayer expects a three-dimensional input of "
                       "shape [batch, sequence, width].")
    hidden_size = target_tensor_shape[2]
    if hidden_size % self.num_attention_heads != 0:
371
372
      raise ValueError(
          "The hidden size (%d) is not a multiple of the number of attention "
Hongkun Yu's avatar
Hongkun Yu committed
373
374
          "heads (%d)" % (hidden_size, self.num_attention_heads))
    self.attention_head_size = int(hidden_size / self.num_attention_heads)
375
    common_kwargs = dict(
376
        bias_initializer=self._bias_initializer,
Hongkun Yu's avatar
Hongkun Yu committed
377
378
379
380
        kernel_regularizer=self._kernel_regularizer,
        bias_regularizer=self._bias_regularizer,
        activity_regularizer=self._activity_regularizer,
        kernel_constraint=self._kernel_constraint,
381
382
383
384
        bias_constraint=self._bias_constraint)
    # Self attention.
    self.self_attention = attention.CachedAttention(
        num_heads=self.num_attention_heads,
385
        key_dim=self.attention_head_size,
386
        dropout=self.attention_dropout_rate,
xinliupitt's avatar
xinliupitt committed
387
        use_bias=self._use_bias,
xinliupitt's avatar
xinliupitt committed
388
        kernel_initializer=self._attention_initializer,
389
390
391
392
393
394
        name="self_attention",
        **common_kwargs)
    self.self_attention_output_dense = tf.keras.layers.experimental.EinsumDense(
        "abc,cd->abd",
        output_shape=(None, hidden_size),
        bias_axes="d",
xinliupitt's avatar
xinliupitt committed
395
        kernel_initializer=self._kernel_initializer,
396
397
        name="output",
        **common_kwargs)
398
    self.self_attention_dropout = tf.keras.layers.Dropout(
Hongkun Yu's avatar
Hongkun Yu committed
399
        rate=self.dropout_rate)
400
401
    self.self_attention_layer_norm = (
        tf.keras.layers.LayerNormalization(
xinliupitt's avatar
xinliupitt committed
402
            name="self_attention_layer_norm",
403
404
            axis=-1,
            epsilon=self._norm_epsilon))
405
406
407
    # Encoder-decoder attention.
    self.encdec_attention = self._cross_attention_cls(
        num_heads=self.num_attention_heads,
408
        key_dim=self.attention_head_size,
Hongkun Yu's avatar
Hongkun Yu committed
409
410
        dropout=self.attention_dropout_rate,
        output_shape=hidden_size,
xinliupitt's avatar
xinliupitt committed
411
        use_bias=self._use_bias,
xinliupitt's avatar
xinliupitt committed
412
        kernel_initializer=self._attention_initializer,
413
414
        name="attention/encdec",
        **common_kwargs)
415
416

    self.encdec_attention_dropout = tf.keras.layers.Dropout(
Hongkun Yu's avatar
Hongkun Yu committed
417
        rate=self.dropout_rate)
418
419
    self.encdec_attention_layer_norm = (
        tf.keras.layers.LayerNormalization(
xinliupitt's avatar
xinliupitt committed
420
            name="attention/encdec_output_layer_norm",
421
422
            axis=-1,
            epsilon=self._norm_epsilon))
423
424

    # Feed-forward projection.
425
426
427
428
    self.intermediate_dense = tf.keras.layers.experimental.EinsumDense(
        "abc,cd->abd",
        output_shape=(None, self.intermediate_size),
        bias_axes="d",
xinliupitt's avatar
xinliupitt committed
429
        kernel_initializer=self._kernel_initializer,
430
431
        name="intermediate",
        **common_kwargs)
432
433
    self.intermediate_activation_layer = tf.keras.layers.Activation(
        self.intermediate_activation)
434
435
    self._intermediate_dropout_layer = tf.keras.layers.Dropout(
        rate=self._intermediate_dropout)
436
437
438
439
    self.output_dense = tf.keras.layers.experimental.EinsumDense(
        "abc,cd->abd",
        output_shape=(None, hidden_size),
        bias_axes="d",
xinliupitt's avatar
xinliupitt committed
440
        kernel_initializer=self._kernel_initializer,
441
442
        name="output",
        **common_kwargs)
Hongkun Yu's avatar
Hongkun Yu committed
443
    self.output_dropout = tf.keras.layers.Dropout(rate=self.dropout_rate)
444
    self.output_layer_norm = tf.keras.layers.LayerNormalization(
xinliupitt's avatar
xinliupitt committed
445
        name="output_layer_norm", axis=-1, epsilon=self._norm_epsilon)
446
447
    super(TransformerDecoderLayer, self).build(input_shape)

xinliupitt's avatar
xinliupitt committed
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
  def get_config(self):
    config = {
        "num_attention_heads":
            self.num_attention_heads,
        "intermediate_size":
            self.intermediate_size,
        "intermediate_activation":
            self.intermediate_activation,
        "dropout_rate":
            self.dropout_rate,
        "attention_dropout_rate":
            self.attention_dropout_rate,
        "multi_channel_cross_attention":
            self.multi_channel_cross_attention,
        "kernel_initializer":
            tf.keras.initializers.serialize(self._kernel_initializer),
        "bias_initializer":
            tf.keras.initializers.serialize(self._bias_initializer),
        "kernel_regularizer":
            tf.keras.regularizers.serialize(self._kernel_regularizer),
        "bias_regularizer":
            tf.keras.regularizers.serialize(self._bias_regularizer),
        "activity_regularizer":
            tf.keras.regularizers.serialize(self._activity_regularizer),
        "kernel_constraint":
            tf.keras.constraints.serialize(self._kernel_constraint),
        "bias_constraint":
            tf.keras.constraints.serialize(self._bias_constraint),
        "use_bias":
            self._use_bias,
        "norm_first":
            self._norm_first,
        "norm_epsilon":
xinliupitt's avatar
xinliupitt committed
481
482
            self._norm_epsilon,
        "intermediate_dropout":
xinliupitt's avatar
xinliupitt committed
483
484
            self._intermediate_dropout,
        "attention_initializer":
xinliupitt's avatar
xinliupitt committed
485
            tf.keras.initializers.serialize(self._attention_initializer)
xinliupitt's avatar
xinliupitt committed
486
487
488
489
    }
    base_config = super(TransformerDecoderLayer, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
  def common_layers_with_encoder(self):
    """Gets layer objects that can make a Transformer encoder block."""
    return [
        self.self_attention, self.self_attention_layer_norm,
        self.intermediate_dense, self.output_dense, self.output_layer_norm
    ]

  def call(self, inputs, cache=None, decode_loop_step=None):
    if self.multi_channel_cross_attention:
      if len(inputs) != 5:
        raise ValueError(
            "TransformerDecoderLayer must have 5 inputs, when it uses "
            "multi_channel_cross_attention. But it got: %d" % len(inputs))
    elif len(inputs) != 4:
      raise ValueError(
          "TransformerDecoderLayer must have 4 inputs, but it got: %d" %
          len(inputs))
    input_tensor, memory, attention_mask, self_attention_mask = inputs[:4]
xinliupitt's avatar
xinliupitt committed
508
509
510
    source_tensor = input_tensor
    if self._norm_first:
      input_tensor = self.self_attention_layer_norm(input_tensor)
511
    self_attention_output, cache = self.self_attention(
512
513
        query=input_tensor,
        value=input_tensor,
514
515
516
517
        attention_mask=self_attention_mask,
        cache=cache,
        decode_loop_step=decode_loop_step)
    self_attention_output = self.self_attention_dropout(self_attention_output)
xinliupitt's avatar
xinliupitt committed
518
519
520
521
522
523
524
525
526
    if self._norm_first:
      self_attention_output = source_tensor + self_attention_output
    else:
      self_attention_output = self.self_attention_layer_norm(
          input_tensor + self_attention_output)
    if self._norm_first:
      source_self_attention_output = self_attention_output
      self_attention_output = self.encdec_attention_layer_norm(
          self_attention_output)
527
528
529
530
    cross_attn_inputs = dict(
        query=self_attention_output,
        value=memory,
        attention_mask=attention_mask)
531
532
    if self.multi_channel_cross_attention:
      # Accesses the 5-th input tensor for the doc-attention probabilities.
533
534
      cross_attn_inputs["context_attention_weights"] = inputs[-1]
    attention_output = self.encdec_attention(**cross_attn_inputs)
535
    attention_output = self.encdec_attention_dropout(attention_output)
xinliupitt's avatar
xinliupitt committed
536
537
538
539
    if self._norm_first:
      attention_output = source_self_attention_output + attention_output
    else:
      attention_output = self.encdec_attention_layer_norm(
540
          self_attention_output + attention_output)
xinliupitt's avatar
xinliupitt committed
541
542
543
    if self._norm_first:
      source_attention_output = attention_output
      attention_output = self.output_layer_norm(attention_output)
544
545
546
547

    intermediate_output = self.intermediate_dense(attention_output)
    intermediate_output = self.intermediate_activation_layer(
        intermediate_output)
548
    intermediate_output = self._intermediate_dropout_layer(intermediate_output)
549
550
    layer_output = self.output_dense(intermediate_output)
    layer_output = self.output_dropout(layer_output)
xinliupitt's avatar
xinliupitt committed
551
552
553
554
    if self._norm_first:
      layer_output = source_attention_output + layer_output
    else:
      layer_output = self.output_layer_norm(layer_output + attention_output)
555
    return layer_output, cache