transformer.py 18.7 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Hongkun Yu's avatar
Hongkun Yu committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Frederick Liu's avatar
Frederick Liu committed
14

Hongkun Yu's avatar
Hongkun Yu committed
15
"""Keras-based transformer block layer."""
16
# pylint: disable=g-classes-have-attributes
Hongkun Yu's avatar
Hongkun Yu committed
17

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

Scott Zhu's avatar
Scott Zhu committed
22
from official.modeling import tf_utils
Hongkun Yu's avatar
Hongkun Yu committed
23
from official.nlp.modeling.layers import attention
24
from official.nlp.modeling.layers import multi_channel_attention
25
from official.nlp.modeling.layers import transformer_encoder_block
26
from official.nlp.modeling.layers.util import tf_function_if_eager
Hongkun Yu's avatar
Hongkun Yu committed
27
28
29


@tf.keras.utils.register_keras_serializable(package="Text")
30
class Transformer(transformer_encoder_block.TransformerEncoderBlock):
Hongkun Yu's avatar
Hongkun Yu committed
31
32
33
34
35
  """Transformer layer.

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

36
37
38
  **Warning: this layer is deprecated. Please don't use it. Use the
  `TransformerEncoderBlock` layer instead.**

39
  Args:
Hongkun Yu's avatar
Hongkun Yu committed
40
41
42
43
44
    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.
45
46
    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
47
48
49
50
51
52
53
    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
54
55
    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
56
    norm_first: Whether to normalize inputs to attention and intermediate dense
57
58
      layers. If set False, output of attention and intermediate dense layers is
      normalized.
xinliupitt's avatar
xinliupitt committed
59
    norm_epsilon: Epsilon value to initialize normalization layers.
60
    intermediate_dropout: Dropout probability for intermediate_dropout_layer.
xinliupitt's avatar
xinliupitt committed
61
62
    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
63
64
65
66
67
68
69
70
  """

  def __init__(self,
               num_attention_heads,
               intermediate_size,
               intermediate_activation,
               dropout_rate=0.0,
               attention_dropout_rate=0.0,
71
               output_range=None,
Hongkun Yu's avatar
Hongkun Yu committed
72
73
74
75
76
77
78
               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
79
80
81
               use_bias=True,
               norm_first=False,
               norm_epsilon=1e-12,
xinliupitt's avatar
xinliupitt committed
82
               intermediate_dropout=0.0,
xinliupitt's avatar
xinliupitt committed
83
               attention_initializer=None,
Hongkun Yu's avatar
Hongkun Yu committed
84
               **kwargs):
85
    super().__init__(
Zhenyu Tan's avatar
Zhenyu Tan committed
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
        num_attention_heads=num_attention_heads,
        inner_dim=intermediate_size,
        inner_activation=intermediate_activation,
        output_dropout=dropout_rate,
        attention_dropout=attention_dropout_rate,
        output_range=output_range,
        kernel_initializer=kernel_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
        kernel_constraint=kernel_constraint,
        bias_constraint=bias_constraint,
        use_bias=use_bias,
        norm_first=norm_first,
        norm_epsilon=norm_epsilon,
        inner_dropout=intermediate_dropout,
        attention_initializer=attention_initializer,
        **kwargs)
105
106
    logging.warning("The `Transformer` layer is deprecated. Please directly "
                    "use `TransformerEncoderBlock`.")
107

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
108
109
110
111
112
113
114
115
116
  def get_config(self):
    return {
        "num_attention_heads":
            self._num_heads,
        "intermediate_size":
            self._inner_dim,
        "intermediate_activation":
            self._inner_activation,
        "dropout_rate":
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
117
            self._output_dropout_rate,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
        "attention_dropout_rate":
            self._attention_dropout_rate,
        "output_range":
            self._output_range,
        "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":
            self._norm_epsilon,
        "intermediate_dropout":
            self._inner_dropout,
        "attention_initializer":
            tf.keras.initializers.serialize(self._attention_initializer)
    }

148

Chen Chen's avatar
Chen Chen committed
149
150
@tf.keras.utils.register_keras_serializable(package="Text")
@gin.configurable
151
152
153
154
class CompiledTransformer(Transformer):

  @tf_function_if_eager(experimental_compile=True)
  def call(self, inputs):
155
    return super().call(inputs)
156
157
158


@tf.keras.utils.register_keras_serializable(package="Text")
Hongkun Yu's avatar
Hongkun Yu committed
159
class TransformerDecoderBlock(tf.keras.layers.Layer):
160
161
162
163
164
165
  """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
166

167
  Args:
Hongkun Yu's avatar
Hongkun Yu committed
168
169
170
171
172
173
174
175
176
177
178
179
180
181
    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
182
183
    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
184
    norm_first: Whether to normalize inputs to attention and intermediate dense
185
186
      layers. If set False, output of attention and intermediate dense layers is
      normalized.
xinliupitt's avatar
xinliupitt committed
187
    norm_epsilon: Epsilon value to initialize normalization layers.
188
    intermediate_dropout: Dropout probability for intermediate_dropout_layer.
xinliupitt's avatar
xinliupitt committed
189
190
    attention_initializer: Initializer for kernels of attention layers. If set
      `None`, attention layers use kernel_initializer as initializer for kernel.
191
192
193
  """

  def __init__(self,
Hongkun Yu's avatar
Hongkun Yu committed
194
195
196
197
198
               num_attention_heads,
               intermediate_size,
               intermediate_activation,
               dropout_rate=0.0,
               attention_dropout_rate=0.0,
199
               multi_channel_cross_attention=False,
Hongkun Yu's avatar
Hongkun Yu committed
200
201
202
203
204
205
206
               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
207
208
209
               use_bias=True,
               norm_first=False,
               norm_epsilon=1e-12,
xinliupitt's avatar
xinliupitt committed
210
               intermediate_dropout=0.0,
xinliupitt's avatar
xinliupitt committed
211
               attention_initializer=None,
212
               **kwargs):
Hongkun Yu's avatar
Hongkun Yu committed
213
    super().__init__(**kwargs)
214
215
216
217
    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
218
219
    self.dropout_rate = dropout_rate
    self.attention_dropout_rate = attention_dropout_rate
220
    self.multi_channel_cross_attention = multi_channel_cross_attention
Hongkun Yu's avatar
Hongkun Yu committed
221
222
223
224
225
226
227
    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
228
229
230
    self._use_bias = use_bias
    self._norm_first = norm_first
    self._norm_epsilon = norm_epsilon
xinliupitt's avatar
xinliupitt committed
231
    self._intermediate_dropout = intermediate_dropout
xinliupitt's avatar
xinliupitt committed
232
    if attention_initializer:
xinliupitt's avatar
xinliupitt committed
233
234
      self._attention_initializer = tf.keras.initializers.get(
          attention_initializer)
xinliupitt's avatar
xinliupitt committed
235
    else:
Scott Zhu's avatar
Scott Zhu committed
236
237
      self._attention_initializer = tf_utils.clone_initializer(
          self._kernel_initializer)
238
239
240
241
242
    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
243
244
  def build(self, input_shape):
    target_tensor_shape = tf.TensorShape(input_shape[0])
245
    if len(target_tensor_shape.as_list()) != 3:
Hongkun Yu's avatar
Hongkun Yu committed
246
247
248
249
      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:
250
251
      raise ValueError(
          "The hidden size (%d) is not a multiple of the number of attention "
Hongkun Yu's avatar
Hongkun Yu committed
252
          "heads (%d)" % (hidden_size, self.num_attention_heads))
253
    self.attention_head_size = int(hidden_size) // self.num_attention_heads
254
    common_kwargs = dict(
Hongkun Yu's avatar
Hongkun Yu committed
255
256
257
258
        kernel_regularizer=self._kernel_regularizer,
        bias_regularizer=self._bias_regularizer,
        activity_regularizer=self._activity_regularizer,
        kernel_constraint=self._kernel_constraint,
259
260
261
262
        bias_constraint=self._bias_constraint)
    # Self attention.
    self.self_attention = attention.CachedAttention(
        num_heads=self.num_attention_heads,
263
        key_dim=self.attention_head_size,
264
        dropout=self.attention_dropout_rate,
xinliupitt's avatar
xinliupitt committed
265
        use_bias=self._use_bias,
Scott Zhu's avatar
Scott Zhu committed
266
267
268
        kernel_initializer=tf_utils.clone_initializer(
            self._attention_initializer),
        bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
269
270
        name="self_attention",
        **common_kwargs)
271
    self.self_attention_output_dense = tf.keras.layers.EinsumDense(
272
273
274
        "abc,cd->abd",
        output_shape=(None, hidden_size),
        bias_axes="d",
Scott Zhu's avatar
Scott Zhu committed
275
276
        kernel_initializer=tf_utils.clone_initializer(self._kernel_initializer),
        bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
277
278
        name="output",
        **common_kwargs)
279
    self.self_attention_dropout = tf.keras.layers.Dropout(
Hongkun Yu's avatar
Hongkun Yu committed
280
        rate=self.dropout_rate)
281
282
    self.self_attention_layer_norm = (
        tf.keras.layers.LayerNormalization(
xinliupitt's avatar
xinliupitt committed
283
            name="self_attention_layer_norm",
284
            axis=-1,
285
286
            epsilon=self._norm_epsilon,
            dtype="float32"))
287
288
289
    # Encoder-decoder attention.
    self.encdec_attention = self._cross_attention_cls(
        num_heads=self.num_attention_heads,
290
        key_dim=self.attention_head_size,
Hongkun Yu's avatar
Hongkun Yu committed
291
292
        dropout=self.attention_dropout_rate,
        output_shape=hidden_size,
xinliupitt's avatar
xinliupitt committed
293
        use_bias=self._use_bias,
Scott Zhu's avatar
Scott Zhu committed
294
295
296
        kernel_initializer=tf_utils.clone_initializer(
            self._attention_initializer),
        bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
297
298
        name="attention/encdec",
        **common_kwargs)
299
300

    self.encdec_attention_dropout = tf.keras.layers.Dropout(
Hongkun Yu's avatar
Hongkun Yu committed
301
        rate=self.dropout_rate)
302
303
    self.encdec_attention_layer_norm = (
        tf.keras.layers.LayerNormalization(
xinliupitt's avatar
xinliupitt committed
304
            name="attention/encdec_output_layer_norm",
305
            axis=-1,
306
307
            epsilon=self._norm_epsilon,
            dtype="float32"))
308
309

    # Feed-forward projection.
310
    self.intermediate_dense = tf.keras.layers.EinsumDense(
311
312
313
        "abc,cd->abd",
        output_shape=(None, self.intermediate_size),
        bias_axes="d",
Scott Zhu's avatar
Scott Zhu committed
314
315
        kernel_initializer=tf_utils.clone_initializer(self._kernel_initializer),
        bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
316
317
        name="intermediate",
        **common_kwargs)
318
319
    self.intermediate_activation_layer = tf.keras.layers.Activation(
        self.intermediate_activation)
320
321
    self._intermediate_dropout_layer = tf.keras.layers.Dropout(
        rate=self._intermediate_dropout)
322
    self.output_dense = tf.keras.layers.EinsumDense(
323
324
325
        "abc,cd->abd",
        output_shape=(None, hidden_size),
        bias_axes="d",
Scott Zhu's avatar
Scott Zhu committed
326
327
        kernel_initializer=tf_utils.clone_initializer(self._kernel_initializer),
        bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
328
329
        name="output",
        **common_kwargs)
Hongkun Yu's avatar
Hongkun Yu committed
330
    self.output_dropout = tf.keras.layers.Dropout(rate=self.dropout_rate)
331
    self.output_layer_norm = tf.keras.layers.LayerNormalization(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
332
333
334
335
        name="output_layer_norm",
        axis=-1,
        epsilon=self._norm_epsilon,
        dtype="float32")
Hongkun Yu's avatar
Hongkun Yu committed
336
    super().build(input_shape)
337

xinliupitt's avatar
xinliupitt committed
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
  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
371
372
            self._norm_epsilon,
        "intermediate_dropout":
xinliupitt's avatar
xinliupitt committed
373
374
            self._intermediate_dropout,
        "attention_initializer":
xinliupitt's avatar
xinliupitt committed
375
            tf.keras.initializers.serialize(self._attention_initializer)
xinliupitt's avatar
xinliupitt committed
376
    }
Hongkun Yu's avatar
Hongkun Yu committed
377
    base_config = super().get_config()
xinliupitt's avatar
xinliupitt committed
378
379
    return dict(list(base_config.items()) + list(config.items()))

380
381
382
383
384
385
386
  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
    ]

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
387
388
389
  def call(self, inputs, cache=None, decode_loop_step=None):
    if self.multi_channel_cross_attention:
      if len(inputs) != 5:
390
        raise ValueError(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
391
            "TransformerDecoderBlock must have 5 inputs, when it uses "
392
            "multi_channel_cross_attention. But it got: %d" % len(inputs))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
393
394
395
396
397
    elif len(inputs) != 4:
      raise ValueError(
          "TransformerDecoderBlock 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
398
399
400
    source_tensor = input_tensor
    if self._norm_first:
      input_tensor = self.self_attention_layer_norm(input_tensor)
401
    self_attention_output, cache = self.self_attention(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
402
        query=input_tensor,
403
        value=input_tensor,
404
405
406
407
        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
408
409
410
411
412
413
414
415
416
    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)
417
    cross_attn_inputs = dict(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
418
        query=self_attention_output,
419
420
        value=memory,
        attention_mask=attention_mask)
421
422
    if self.multi_channel_cross_attention:
      # Accesses the 5-th input tensor for the doc-attention probabilities.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
423
      cross_attn_inputs["context_attention_weights"] = inputs[-1]
424
    attention_output = self.encdec_attention(**cross_attn_inputs)
425
    attention_output = self.encdec_attention_dropout(attention_output)
xinliupitt's avatar
xinliupitt committed
426
427
428
429
    if self._norm_first:
      attention_output = source_self_attention_output + attention_output
    else:
      attention_output = self.encdec_attention_layer_norm(
430
          self_attention_output + attention_output)
xinliupitt's avatar
xinliupitt committed
431
432
433
    if self._norm_first:
      source_attention_output = attention_output
      attention_output = self.output_layer_norm(attention_output)
434
435
436
437

    intermediate_output = self.intermediate_dense(attention_output)
    intermediate_output = self.intermediate_activation_layer(
        intermediate_output)
438
    intermediate_output = self._intermediate_dropout_layer(intermediate_output)
439
440
    layer_output = self.output_dense(intermediate_output)
    layer_output = self.output_dropout(layer_output)
xinliupitt's avatar
xinliupitt committed
441
442
443
444
    if self._norm_first:
      layer_output = source_attention_output + layer_output
    else:
      layer_output = self.output_layer_norm(layer_output + attention_output)
445
    return layer_output, cache