"sgl-kernel/pyproject_rocm.toml" did not exist on "a53454c55e222fcc7375676b665b7e1276464170"
nn_layers.py 4.52 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Hye Yoon's avatar
Hye Yoon 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.
14

Hye Yoon's avatar
Hye Yoon committed
15
"""Contains model definitions."""
Yeqing Li's avatar
Yeqing Li committed
16
from typing import Any, Dict, Optional
Hye Yoon's avatar
Hye Yoon committed
17
18

import tensorflow as tf
Yeqing Li's avatar
Yeqing Li committed
19
from official.projects.yt8m.modeling import yt8m_model_utils as utils
Hye Yoon's avatar
Hye Yoon committed
20
21
22
23
24
25
26

layers = tf.keras.layers


class LogisticModel():
  """Logistic model with L2 regularization."""

27
  def create_model(self, model_input, vocab_size, l2_penalty=1e-8):
Hye Yoon's avatar
Hye Yoon committed
28
    """Creates a logistic model.
29

Hye Yoon's avatar
Hye Yoon committed
30
31
32
    Args:
      model_input: 'batch' x 'num_features' matrix of input features.
      vocab_size: The number of classes in the dataset.
33
34
      l2_penalty: L2 weight regularization ratio.

Hye Yoon's avatar
Hye Yoon committed
35
36
37
38
39
40
41
42
    Returns:
      A dictionary with a tensor containing the probability predictions of the
      model in the 'predictions' key. The dimensions of the tensor are
      batch_size x num_classes.
    """
    output = layers.Dense(
        vocab_size,
        activation=tf.nn.sigmoid,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
43
        kernel_regularizer=tf.keras.regularizers.l2(l2_penalty))(
44
            model_input)
Hye Yoon's avatar
Hye Yoon committed
45
46
47
48
49
50
51
52
53
    return {"predictions": output}


class MoeModel():
  """A softmax over a mixture of logistic models (with L2 regularization)."""

  def create_model(self,
                   model_input,
                   vocab_size,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
54
55
56
57
58
59
                   num_mixtures: int = 2,
                   use_input_context_gate: bool = False,
                   use_output_context_gate: bool = False,
                   normalizer_fn=None,
                   normalizer_params: Optional[Dict[str, Any]] = None,
                   l2_penalty: float = 1e-5):
Hye Yoon's avatar
Hye Yoon committed
60
    """Creates a Mixture of (Logistic) Experts model.
61
62
63
64
65
66
67
68
69

     The model consists of a per-class softmax distribution over a
     configurable number of logistic classifiers. One of the classifiers
     in the mixture is not trained, and always predicts 0.
    Args:
      model_input: 'batch_size' x 'num_features' matrix of input features.
      vocab_size: The number of classes in the dataset.
      num_mixtures: The number of mixtures (excluding a dummy 'expert' that
        always predicts the non-existence of an entity).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
70
71
72
73
      use_input_context_gate: if True apply context gate layer to the input.
      use_output_context_gate: if True apply context gate layer to the output.
      normalizer_fn: normalization op constructor (e.g. batch norm).
      normalizer_params: parameters to the `normalizer_fn`.
74
75
76
77
78
79
80
81
      l2_penalty: How much to penalize the squared magnitudes of parameter
        values.

    Returns:
      A dictionary with a tensor containing the probability predictions
      of the model in the 'predictions' key. The dimensions of the tensor
      are batch_size x num_classes.
    """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
82
83
84
85
86
87
    if use_input_context_gate:
      model_input = utils.context_gate(
          model_input,
          normalizer_fn=normalizer_fn,
          normalizer_params=normalizer_params,
      )
Hye Yoon's avatar
Hye Yoon committed
88
89
90
91
92

    gate_activations = layers.Dense(
        vocab_size * (num_mixtures + 1),
        activation=None,
        bias_initializer=None,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
93
        kernel_regularizer=tf.keras.regularizers.l2(l2_penalty))(
94
            model_input)
Hye Yoon's avatar
Hye Yoon committed
95
96
97
    expert_activations = layers.Dense(
        vocab_size * num_mixtures,
        activation=None,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
98
        kernel_regularizer=tf.keras.regularizers.l2(l2_penalty))(
99
            model_input)
Hye Yoon's avatar
Hye Yoon committed
100
101
102
103
104
105
106
107
108
109
110
111
112

    gating_distribution = tf.nn.softmax(
        tf.reshape(
            gate_activations,
            [-1, num_mixtures + 1]))  # (Batch * #Labels) x (num_mixtures + 1)
    expert_distribution = tf.nn.sigmoid(
        tf.reshape(expert_activations,
                   [-1, num_mixtures]))  # (Batch * #Labels) x num_mixtures

    final_probabilities_by_class_and_batch = tf.reduce_sum(
        gating_distribution[:, :num_mixtures] * expert_distribution, 1)
    final_probabilities = tf.reshape(final_probabilities_by_class_and_batch,
                                     [-1, vocab_size])
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
113
114
115
116
117
118
    if use_output_context_gate:
      final_probabilities = utils.context_gate(
          final_probabilities,
          normalizer_fn=normalizer_fn,
          normalizer_params=normalizer_params,
      )
Hye Yoon's avatar
Hye Yoon committed
119
    return {"predictions": final_probabilities}