classification.py 4.12 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

15
"""Classification and regression network."""
16
# pylint: disable=g-classes-have-attributes
17
import collections
Hongkun Yu's avatar
Hongkun Yu committed
18
import tensorflow as tf
19
from tensorflow.python.util import deprecation
Hongkun Yu's avatar
Hongkun Yu committed
20
21
22


@tf.keras.utils.register_keras_serializable(package='Text')
23
class Classification(tf.keras.Model):
Hongkun Yu's avatar
Hongkun Yu committed
24
25
  """Classification network head for BERT modeling.

26
27
  This network implements a simple classifier head based on a dense layer. If
  num_classes is one, it can be considered as a regression problem.
Hongkun Yu's avatar
Hongkun Yu committed
28

29
30
31
  *Note* that the network is constructed by
  [Keras Functional API](https://keras.io/guides/functional_api/).

32
  Args:
Hongkun Yu's avatar
Hongkun Yu committed
33
    input_width: The innermost dimension of the input tensor to this network.
34
35
    num_classes: The number of classes that this network should classify to. If
      equal to 1, a regression problem is assumed.
Hongkun Yu's avatar
Hongkun Yu committed
36
    activation: The activation, if any, for the dense layer in this network.
Hongkun Yu's avatar
Hongkun Yu committed
37
38
    initializer: The initializer for the dense layer in this network. Defaults
      to a Glorot uniform initializer.
39
40
    output: The output style for this network. Can be either `logits` or
      `predictions`.
Hongkun Yu's avatar
Hongkun Yu committed
41
42
  """

43
44
  @deprecation.deprecated(None, 'Classification as a network is deprecated. '
                          'Please use the layers.ClassificationHead instead.')
Hongkun Yu's avatar
Hongkun Yu committed
45
46
47
48
49
50
51
52
53
54
  def __init__(self,
               input_width,
               num_classes,
               initializer='glorot_uniform',
               output='logits',
               **kwargs):

    cls_output = tf.keras.layers.Input(
        shape=(input_width,), name='cls_output', dtype=tf.float32)

55
    logits = tf.keras.layers.Dense(
Hongkun Yu's avatar
Hongkun Yu committed
56
57
58
59
60
        num_classes,
        activation=None,
        kernel_initializer=initializer,
        name='predictions/transform/logits')(
            cls_output)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
61

Hongkun Yu's avatar
Hongkun Yu committed
62
    if output == 'logits':
63
      output_tensors = logits
Hongkun Yu's avatar
Hongkun Yu committed
64
    elif output == 'predictions':
65
      policy = tf.keras.mixed_precision.global_policy()
66
67
68
69
70
71
      if policy.name == 'mixed_bfloat16':
        # b/158514794: bf16 is not stable with post-softmax cross-entropy.
        policy = tf.float32
      output_tensors = tf.keras.layers.Activation(
          tf.nn.log_softmax, dtype=policy)(
              logits)
Hongkun Yu's avatar
Hongkun Yu committed
72
73
74
75
76
    else:
      raise ValueError(
          ('Unknown `output` value "%s". `output` can be either "logits" or '
           '"predictions"') % output)

Hongkun Yu's avatar
Hongkun Yu committed
77
    super().__init__(
Hongkun Yu's avatar
Hongkun Yu committed
78
79
        inputs=[cls_output], outputs=output_tensors, **kwargs)

80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
    # b/164516224
    # Once we've created the network using the Functional API, we call
    # super().__init__ as though we were invoking the Functional API Model
    # constructor, resulting in this object having all the properties of a model
    # created using the Functional API. Once super().__init__ is called, we
    # can assign attributes to `self` - note that all `self` assignments are
    # below this line.
    config_dict = {
        'input_width': input_width,
        'num_classes': num_classes,
        'initializer': initializer,
        'output': output,
    }
    # We are storing the config dict as a namedtuple here to ensure checkpoint
    # compatibility with an earlier version of this model which did not track
    # the config dict attribute. TF does not track immutable attrs which
    # do not contain Trackables, so by creating a config namedtuple instead of
    # a dict we avoid tracking it.
    config_cls = collections.namedtuple('Config', config_dict.keys())
    self._config = config_cls(**config_dict)
    self.logits = logits

Hongkun Yu's avatar
Hongkun Yu committed
102
  def get_config(self):
103
    return dict(self._config._asdict())
Hongkun Yu's avatar
Hongkun Yu committed
104
105
106
107

  @classmethod
  def from_config(cls, config, custom_objects=None):
    return cls(**config)