使用 @tf.function 重用嵌入层的权重矩阵

Question

不使用@tf.function，脚本可以完美运行

我想用它来加速训练，但是在我重复使用嵌入层的权重矩阵时会出错。

我认为错误是由 get_weights() 引起的，因为它将张量转换回 numpy

我尝试使用 tf.keras.layers.Dense 而不是重新使用嵌入的权重，并且效果很好。

class Example(tf.keras.Model):
    def __init__(self,):
        super(Example, self).__init__()
        self.embed_dim = embed_dim
        self.vocab_size = vocab_size
        self.embed = tf.keras.layers.Embedding(self.vocab_size, self.embed_dim)
        ...

    def call(self, inputs, trianing):
        ...
        embed_matrix = self.embed.get_weights()

        # a dense layer
        Vhid = tf.matmul(self.kernel, tf.transpose(embed_matrix[0]))
        pred_w = tf.matmul(pred, Vhid) + self.bias

在我的火车脚本中。 我做了

@tf.function
def train_step(x, y, training=None):

    with tf.GradientTape() as tape:
        pred = model(x, y, training)
        losses = compute_loss(y, pred)

    grads = tape.gradient(losses, model.trainable_variables)

    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    return losses

/home/thomas/projects/tf_convsent/models/.py:195 call  *
   embed_matrix = self.embed.get_weights()  # [vocab_size, 300]
/home/thomas/.conda/envs/tf2_p37/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:1177 get_weights
   return backend.batch_get_value(params)
/home/thomas/.conda/envs/tf2_p37/lib/python3.7/site-packages/tensorflow/python/keras/backend.py:3011 batch_get_value
   raise RuntimeError('Cannot get value inside Tensorflow graph function.')

RuntimeError: Cannot get value inside Tensorflow graph function.

Answer 1

找到最简单的解决方案，将训练速度提高 50%（122 小时到 ~65 小时）

只是改变

embed_matrix = self.embed.get_weights()

到

embed_matrix = self.embed.weights

会成功的。