Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/keras/src/layers/core/embedding.py: 24%
93 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Embedding layer."""
18import tensorflow.compat.v2 as tf
20from keras.src import backend
21from keras.src import constraints
22from keras.src import initializers
23from keras.src import regularizers
24from keras.src.dtensor import utils
25from keras.src.engine import base_layer_utils
26from keras.src.engine.base_layer import Layer
27from keras.src.utils import tf_utils
29# isort: off
30from tensorflow.python.util.tf_export import keras_export
33@keras_export("keras.layers.Embedding")
34class Embedding(Layer):
35 """Turns positive integers (indexes) into dense vectors of fixed size.
37 e.g. `[[4], [20]] -> [[0.25, 0.1], [0.6, -0.2]]`
39 This layer can only be used on positive integer inputs of a fixed range. The
40 `tf.keras.layers.TextVectorization`, `tf.keras.layers.StringLookup`,
41 and `tf.keras.layers.IntegerLookup` preprocessing layers can help prepare
42 inputs for an `Embedding` layer.
44 This layer accepts `tf.Tensor`, `tf.RaggedTensor` and `tf.SparseTensor`
45 input.
47 Example:
49 >>> model = tf.keras.Sequential()
50 >>> model.add(tf.keras.layers.Embedding(1000, 64, input_length=10))
51 >>> # The model will take as input an integer matrix of size (batch,
52 >>> # input_length), and the largest integer (i.e. word index) in the input
53 >>> # should be no larger than 999 (vocabulary size).
54 >>> # Now model.output_shape is (None, 10, 64), where `None` is the batch
55 >>> # dimension.
56 >>> input_array = np.random.randint(1000, size=(32, 10))
57 >>> model.compile('rmsprop', 'mse')
58 >>> output_array = model.predict(input_array)
59 >>> print(output_array.shape)
60 (32, 10, 64)
62 Args:
63 input_dim: Integer. Size of the vocabulary,
64 i.e. maximum integer index + 1.
65 output_dim: Integer. Dimension of the dense embedding.
66 embeddings_initializer: Initializer for the `embeddings`
67 matrix (see `keras.initializers`).
68 embeddings_regularizer: Regularizer function applied to
69 the `embeddings` matrix (see `keras.regularizers`).
70 embeddings_constraint: Constraint function applied to
71 the `embeddings` matrix (see `keras.constraints`).
72 mask_zero: Boolean, whether or not the input value 0 is a special
73 "padding" value that should be masked out. This is useful when using
74 recurrent layers which may take variable length input. If this is
75 `True`, then all subsequent layers in the model need to support masking
76 or an exception will be raised. If mask_zero is set to True, as a
77 consequence, index 0 cannot be used in the vocabulary (input_dim should
78 equal size of vocabulary + 1).
79 input_length: Length of input sequences, when it is constant.
80 This argument is required if you are going to connect
81 `Flatten` then `Dense` layers upstream
82 (without it, the shape of the dense outputs cannot be computed).
83 sparse: If True, calling this layer returns a `tf.SparseTensor`. If False,
84 the layer returns a dense `tf.Tensor`. For an entry with no features in
85 a sparse tensor (entry with value 0), the embedding vector of index 0 is
86 returned by default.
88 Input shape:
89 2D tensor with shape: `(batch_size, input_length)`.
91 Output shape:
92 3D tensor with shape: `(batch_size, input_length, output_dim)`.
94 **Note on variable placement:**
95 By default, if a GPU is available, the embedding matrix will be placed on
96 the GPU. This achieves the best performance, but it might cause issues:
98 - You may be using an optimizer that does not support sparse GPU kernels.
99 In this case you will see an error upon training your model.
100 - Your embedding matrix may be too large to fit on your GPU. In this case
101 you will see an Out Of Memory (OOM) error.
103 In such cases, you should place the embedding matrix on the CPU memory.
104 You can do so with a device scope, as such:
106 ```python
107 with tf.device('cpu:0'):
108 embedding_layer = Embedding(...)
109 embedding_layer.build()
110 ```
112 The pre-built `embedding_layer` instance can then be added to a `Sequential`
113 model (e.g. `model.add(embedding_layer)`), called in a Functional model
114 (e.g. `x = embedding_layer(x)`), or used in a subclassed model.
115 """
117 @utils.allow_initializer_layout
118 def __init__(
119 self,
120 input_dim,
121 output_dim,
122 embeddings_initializer="uniform",
123 embeddings_regularizer=None,
124 activity_regularizer=None,
125 embeddings_constraint=None,
126 mask_zero=False,
127 input_length=None,
128 sparse=False,
129 **kwargs,
130 ):
131 if "input_shape" not in kwargs:
132 if input_length:
133 kwargs["input_shape"] = (input_length,)
134 else:
135 kwargs["input_shape"] = (None,)
136 if input_dim <= 0 or output_dim <= 0:
137 raise ValueError(
138 "Both `input_dim` and `output_dim` should be positive, "
139 f"Received input_dim = {input_dim} "
140 f"and output_dim = {output_dim}"
141 )
142 if (
143 not base_layer_utils.v2_dtype_behavior_enabled()
144 and "dtype" not in kwargs
145 ):
146 # In TF1, the dtype defaults to the input dtype which is typically
147 # int32, so explicitly set it to floatx
148 kwargs["dtype"] = backend.floatx()
149 # We set autocast to False, as we do not want to cast floating- point
150 # inputs to self.dtype. In call(), we cast to int32, and casting to
151 # self.dtype before casting to int32 might cause the int32 values to be
152 # different due to a loss of precision.
153 kwargs["autocast"] = False
154 use_one_hot_matmul = kwargs.pop("use_one_hot_matmul", False)
155 super().__init__(**kwargs)
157 self.input_dim = input_dim
158 self.output_dim = output_dim
159 self.embeddings_initializer = initializers.get(embeddings_initializer)
160 self.embeddings_regularizer = regularizers.get(embeddings_regularizer)
161 self.activity_regularizer = regularizers.get(activity_regularizer)
162 self.embeddings_constraint = constraints.get(embeddings_constraint)
163 self.mask_zero = mask_zero
164 self.supports_masking = mask_zero
165 self.input_length = input_length
166 self.sparse = sparse
167 if self.sparse and self.mask_zero:
168 raise ValueError(
169 "`mask_zero` cannot be enabled when "
170 "`tf.keras.layers.Embedding` is used with `tf.SparseTensor` "
171 "input."
172 )
173 # Make this flag private and do not serialize it for now.
174 # It will be part of the public API after further testing.
175 self._use_one_hot_matmul = use_one_hot_matmul
177 @tf_utils.shape_type_conversion
178 def build(self, input_shape=None):
179 self.embeddings = self.add_weight(
180 shape=(self.input_dim, self.output_dim),
181 initializer=self.embeddings_initializer,
182 name="embeddings",
183 regularizer=self.embeddings_regularizer,
184 constraint=self.embeddings_constraint,
185 experimental_autocast=False,
186 )
187 self.built = True
189 def compute_mask(self, inputs, mask=None):
190 if not self.mask_zero:
191 return None
192 return tf.not_equal(inputs, 0)
194 @tf_utils.shape_type_conversion
195 def compute_output_shape(self, input_shape):
196 if self.input_length is None:
197 return input_shape + (self.output_dim,)
198 else:
199 # input_length can be tuple if input is 3D or higher
200 if isinstance(self.input_length, (list, tuple)):
201 in_lens = list(self.input_length)
202 else:
203 in_lens = [self.input_length]
204 if len(in_lens) != len(input_shape) - 1:
205 raise ValueError(
206 f'"input_length" is {self.input_length}, but received '
207 f"input has shape {input_shape}"
208 )
209 else:
210 for i, (s1, s2) in enumerate(zip(in_lens, input_shape[1:])):
211 if s1 is not None and s2 is not None and s1 != s2:
212 raise ValueError(
213 f'"input_length" is {self.input_length}, but '
214 f"received input has shape {input_shape}"
215 )
216 elif s1 is None:
217 in_lens[i] = s2
218 return (input_shape[0],) + tuple(in_lens) + (self.output_dim,)
220 def call(self, inputs):
221 dtype = backend.dtype(inputs)
222 if dtype != "int32" and dtype != "int64":
223 inputs = tf.cast(inputs, "int32")
224 if isinstance(inputs, tf.sparse.SparseTensor):
225 if self.sparse:
226 # get sparse embedding values
227 embedding_values = tf.nn.embedding_lookup(
228 params=self.embeddings, ids=inputs.values
229 )
230 embedding_values = tf.reshape(embedding_values, [-1])
231 # get sparse embedding indices
232 indices_values_embed_axis = tf.range(self.output_dim)
233 repeat_times = [inputs.indices.shape[0]]
234 indices_values_embed_axis = tf.expand_dims(
235 tf.tile(indices_values_embed_axis, repeat_times), -1
236 )
237 indices_values_embed_axis = tf.cast(
238 indices_values_embed_axis, dtype=tf.int64
239 )
240 current_indices = tf.repeat(
241 inputs.indices, [self.output_dim], axis=0
242 )
243 new_indices = tf.concat(
244 [current_indices, indices_values_embed_axis], 1
245 )
246 new_shape = tf.concat(
247 [tf.cast(inputs.shape, dtype=tf.int64), [self.output_dim]],
248 axis=-1,
249 )
250 out = tf.SparseTensor(
251 indices=new_indices,
252 values=embedding_values,
253 dense_shape=new_shape,
254 )
255 else:
256 sparse_inputs_expanded = tf.sparse.expand_dims(inputs, axis=-1)
257 out = tf.nn.safe_embedding_lookup_sparse(
258 embedding_weights=self.embeddings,
259 sparse_ids=sparse_inputs_expanded,
260 default_id=0,
261 )
262 elif self._use_one_hot_matmul:
263 # Note that we change the dtype of the one_hot to be same as the
264 # weight tensor, since the input data are usually ints, and weights
265 # are floats. The nn.embedding_lookup support ids as ints, but
266 # the one_hot matmul need both inputs and weights to be same dtype.
267 one_hot_data = tf.one_hot(
268 inputs, depth=self.input_dim, dtype=self.dtype
269 )
270 out = tf.matmul(one_hot_data, self.embeddings)
271 else:
272 out = tf.nn.embedding_lookup(self.embeddings, inputs)
274 if self.sparse and not isinstance(out, tf.SparseTensor):
275 out = tf.sparse.from_dense(out)
277 if (
278 self._dtype_policy.compute_dtype
279 != self._dtype_policy.variable_dtype
280 ):
281 # Instead of casting the variable as in most layers, cast the
282 # output, as this is mathematically equivalent but is faster.
283 out = tf.cast(out, self._dtype_policy.compute_dtype)
284 return out
286 def get_config(self):
287 config = {
288 "input_dim": self.input_dim,
289 "output_dim": self.output_dim,
290 "embeddings_initializer": initializers.serialize(
291 self.embeddings_initializer
292 ),
293 "embeddings_regularizer": regularizers.serialize(
294 self.embeddings_regularizer
295 ),
296 "activity_regularizer": regularizers.serialize(
297 self.activity_regularizer
298 ),
299 "embeddings_constraint": constraints.serialize(
300 self.embeddings_constraint
301 ),
302 "mask_zero": self.mask_zero,
303 "input_length": self.input_length,
304 }
305 base_config = super().get_config()
306 return dict(list(base_config.items()) + list(config.items()))