transformer.py 23.9 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
    layer_output = self._output_dense(intermediate_output)
    layer_output = self._output_dropout(layer_output)
259

xinliupitt's avatar
xinliupitt committed
260
    if self._norm_first:
261
      return source_attention_output + layer_output
262

263
264
265
266
    # During mixed precision training, layer norm output is always fp32 for now.
    # Casts fp32 for the subsequent add.
    layer_output = tf.cast(layer_output, tf.float32)
    return self._output_layer_norm(layer_output + attention_output)
267
268


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

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


@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
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301

  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
302
303
    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
304
    norm_first: Whether to normalize inputs to attention and intermediate dense
305
306
      layers. If set False, output of attention and intermediate dense layers is
      normalized.
xinliupitt's avatar
xinliupitt committed
307
    norm_epsilon: Epsilon value to initialize normalization layers.
308
    intermediate_dropout: Dropout probability for intermediate_dropout_layer.
xinliupitt's avatar
xinliupitt committed
309
310
    attention_initializer: Initializer for kernels of attention layers. If set
      `None`, attention layers use kernel_initializer as initializer for kernel.
311
312
313
  """

  def __init__(self,
Hongkun Yu's avatar
Hongkun Yu committed
314
315
316
317
318
               num_attention_heads,
               intermediate_size,
               intermediate_activation,
               dropout_rate=0.0,
               attention_dropout_rate=0.0,
319
               multi_channel_cross_attention=False,
Hongkun Yu's avatar
Hongkun Yu committed
320
321
322
323
324
325
326
               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
327
328
329
               use_bias=True,
               norm_first=False,
               norm_epsilon=1e-12,
xinliupitt's avatar
xinliupitt committed
330
               intermediate_dropout=0.0,
xinliupitt's avatar
xinliupitt committed
331
               attention_initializer=None,
332
333
334
335
336
337
               **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
338
339
    self.dropout_rate = dropout_rate
    self.attention_dropout_rate = attention_dropout_rate
340
    self.multi_channel_cross_attention = multi_channel_cross_attention
Hongkun Yu's avatar
Hongkun Yu committed
341
342
343
344
345
346
347
    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
348
349
350
    self._use_bias = use_bias
    self._norm_first = norm_first
    self._norm_epsilon = norm_epsilon
xinliupitt's avatar
xinliupitt committed
351
    self._intermediate_dropout = intermediate_dropout
xinliupitt's avatar
xinliupitt committed
352
    if attention_initializer:
xinliupitt's avatar
xinliupitt committed
353
354
      self._attention_initializer = tf.keras.initializers.get(
          attention_initializer)
xinliupitt's avatar
xinliupitt committed
355
356
    else:
      self._attention_initializer = self._kernel_initializer
357
358
359
360
361
    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
362
363
  def build(self, input_shape):
    target_tensor_shape = tf.TensorShape(input_shape[0])
364
    if len(target_tensor_shape.as_list()) != 3:
Hongkun Yu's avatar
Hongkun Yu committed
365
366
367
368
      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:
369
370
      raise ValueError(
          "The hidden size (%d) is not a multiple of the number of attention "
Hongkun Yu's avatar
Hongkun Yu committed
371
372
          "heads (%d)" % (hidden_size, self.num_attention_heads))
    self.attention_head_size = int(hidden_size / self.num_attention_heads)
373
    common_kwargs = dict(
374
        bias_initializer=self._bias_initializer,
Hongkun Yu's avatar
Hongkun Yu committed
375
376
377
378
        kernel_regularizer=self._kernel_regularizer,
        bias_regularizer=self._bias_regularizer,
        activity_regularizer=self._activity_regularizer,
        kernel_constraint=self._kernel_constraint,
379
380
381
382
        bias_constraint=self._bias_constraint)
    # Self attention.
    self.self_attention = attention.CachedAttention(
        num_heads=self.num_attention_heads,
383
        key_dim=self.attention_head_size,
384
        dropout=self.attention_dropout_rate,
xinliupitt's avatar
xinliupitt committed
385
        use_bias=self._use_bias,
xinliupitt's avatar
xinliupitt committed
386
        kernel_initializer=self._attention_initializer,
387
388
389
390
391
392
        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
393
        kernel_initializer=self._kernel_initializer,
394
395
        name="output",
        **common_kwargs)
396
    self.self_attention_dropout = tf.keras.layers.Dropout(
Hongkun Yu's avatar
Hongkun Yu committed
397
        rate=self.dropout_rate)
398
399
    self.self_attention_layer_norm = (
        tf.keras.layers.LayerNormalization(
xinliupitt's avatar
xinliupitt committed
400
            name="self_attention_layer_norm",
401
402
            axis=-1,
            epsilon=self._norm_epsilon))
403
404
405
    # Encoder-decoder attention.
    self.encdec_attention = self._cross_attention_cls(
        num_heads=self.num_attention_heads,
406
        key_dim=self.attention_head_size,
Hongkun Yu's avatar
Hongkun Yu committed
407
408
        dropout=self.attention_dropout_rate,
        output_shape=hidden_size,
xinliupitt's avatar
xinliupitt committed
409
        use_bias=self._use_bias,
xinliupitt's avatar
xinliupitt committed
410
        kernel_initializer=self._attention_initializer,
411
412
        name="attention/encdec",
        **common_kwargs)
413
414

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

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

xinliupitt's avatar
xinliupitt committed
446
447
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
  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
479
480
            self._norm_epsilon,
        "intermediate_dropout":
xinliupitt's avatar
xinliupitt committed
481
482
            self._intermediate_dropout,
        "attention_initializer":
xinliupitt's avatar
xinliupitt committed
483
            tf.keras.initializers.serialize(self._attention_initializer)
xinliupitt's avatar
xinliupitt committed
484
485
486
487
    }
    base_config = super(TransformerDecoderLayer, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
  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
506
507
508
    source_tensor = input_tensor
    if self._norm_first:
      input_tensor = self.self_attention_layer_norm(input_tensor)
509
    self_attention_output, cache = self.self_attention(
510
511
        query=input_tensor,
        value=input_tensor,
512
513
514
515
        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
516
517
518
519
520
521
522
523
524
    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)
525
526
527
528
    cross_attn_inputs = dict(
        query=self_attention_output,
        value=memory,
        attention_mask=attention_mask)
529
530
    if self.multi_channel_cross_attention:
      # Accesses the 5-th input tensor for the doc-attention probabilities.
531
532
      cross_attn_inputs["context_attention_weights"] = inputs[-1]
    attention_output = self.encdec_attention(**cross_attn_inputs)
533
    attention_output = self.encdec_attention_dropout(attention_output)
xinliupitt's avatar
xinliupitt committed
534
535
536
537
    if self._norm_first:
      attention_output = source_self_attention_output + attention_output
    else:
      attention_output = self.encdec_attention_layer_norm(
538
          self_attention_output + attention_output)
xinliupitt's avatar
xinliupitt committed
539
540
541
    if self._norm_first:
      source_attention_output = attention_output
      attention_output = self.output_layer_norm(attention_output)
542
543
544
545

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