TensorFlow 2：运行时错误：无法在稀疏变量上使用约束函数

Question

我用

embeddings_constraint

定义嵌入对象：

from tensorflow.keras.layers import Embedding
from tensorflow.keras.constraints import UnitNorm
. . .
emb = Embedding(input_dim, output_dim, name='embedding_name', embeddings_constraint=UnitNorm(axis=1))
. . .

稍后在代码中，当我想训练包含

emb

的模型时，我从函数

model.fit

:

中得到一个异常

RuntimeError: Cannot use a constraint function on a sparse variable.

但是，当我不对

emb

施加嵌入约束时，代码不会抛出错误。此外，我用 TF 1 尝试过，效果也很好（有或没有

embeddings_constraint

）。根据 GitHub 讨论，这似乎是 TF 2 错误，尽管没有提出可行的解决方案。

有什么想法可以解决这个问题吗？

Answer 1

此问题的解决方法是直接调用约束，如下所示：

from tensorflow.keras.layers import Embedding
from tensorflow.keras.constraints import UnitNorm
. . .
emb = Embedding(input_dim, output_dim, name='embedding_name')
norm_layer = UnitNorm(axis=1)
norm_embedding = norm_layer(emb(embedding_id_input))
. . .

Answer 2

我也点击了这个，添加到tensorflow github问题中，在tf 2.7和tf-nightly中尝试过，仍然看到错误，所以提出了新的keras问题：https://github.com/keras-team/keras/issues/ 15818.

该解决方法不会解决所有用例 - 它限制嵌入活动而不是整个嵌入矩阵。我确实有一个使用 keras 回调的解决方法 - 但我发现它减慢了训练速度 - 我相信回调正在处理无法在 GPU 上运行的 numpy 数组 - 但如果有用，这里是方法

import numpy as np
import tensorflow as tf


class ConstrainEmbeddings(tf.keras.callbacks.Callback):
    def __init__(self, min_norm, emb, eps=1e-9):
        super(ConstrainEmbeddings, self).__init__()
        self.min_norm = min_norm
        self.emb = emb
        self.eps = eps

    def on_batch_begin(self, *args, **kwargs):
        W = self.emb.get_weights()[0]
        norms = tf.maximum(self.eps, tf.norm(W, axis=1))
        delta = tf.expand_dims(tf.math.divide(self.min_norm, norms) - 1.0, 1)
        deltaW = tf.math.multiply(W, delta)
        constrainedW = W + tf.expand_dims(tf.cast(norms < self.min_norm, dtype=tf.float32), 1) * deltaW
        self.emb.set_weights([constrainedW])


def constrain_embeddings(use_keras):
    N = 10
    batch_size = 5
    data = {
        'X': np.random.randint(0, 10, N),
        'Y': np.random.randint(0, 2, N)
    }

    def get_labels(features):
        labels = features.pop('Y')
        return features, labels

    dset = tf.data.Dataset.from_tensor_slices(data).map(get_labels).batch(batch_size)

    inp = tf.keras.Input(shape=(1,), name='X', dtype='int64')
    constraint = tf.keras.constraints.MaxNorm(max_value=0.1) if use_keras else None

    emb = tf.keras.layers.Embedding(
        10, 3, input_length=1,
        embeddings_initializer=tf.keras.initializers.RandomUniform(minval=-1.0, maxval=1.0),
        embeddings_constraint=constraint
    )
    emb_out = emb(inp)
    out = tf.keras.layers.Dense(1)(emb_out)
    model = tf.keras.Model(inputs=inp, outputs=out)
    model.compile(optimizer='adam', loss=tf.keras.losses.binary_crossentropy)
    callbacks = [ConstrainEmbeddings(.3, emb)] if not use_keras else []
    model.fit(dset, epochs=10, callbacks=callbacks)


if __name__ == '__main__':
    print("tensorflow git version: ", tf.version.GIT_VERSION)
    constrain_embeddings(use_keras=False)

Answer 3

我将“if grad.shape!=var.shape:”添加到 OptimizerV2 类中代码的违规部分，在我的例子中，训练现在按预期运行：

        if isinstance(grad, tf.IndexedSlices):
            if var.constraint is not None:
                if grad.shape!=var.shape:
                    raise RuntimeError(
                        "Cannot use a constraint function on a sparse "
                        f"variable. Received: grad={grad}, "
                        f"var.constraint={var.constraint}."
                    )
            if "apply_state" in self._sparse_apply_args:
                apply_kwargs["apply_state"] = apply_state
            return self._resource_apply_sparse_duplicate_indices(
                grad.values, var, grad.indices, **apply_kwargs
            )