如何将 torch.scatter_add 更改为 tensorflow 函数

Question

我需要将代码 pytorch 转移到 tensorflow 这个pytorch代码在这里NADST

    encoded_context = ft['encoded_context2']
    encoded_in_domainslots = ft['encoded_in_domainslots2']
    self.pointer_attn(ft['out_states'], encoded_context, encoded_context, context_mask)
    pointer_attn = self.pointer_attn.attn.squeeze(1)
    p_vocab = F.softmax(vocab_attn, dim = -1)
    context_index = context.unsqueeze(1).expand_as(pointer_attn)
    p_context_ptr = torch.zeros(p_vocab.size()).cuda()
    p_context_ptr.scatter_add_(2, context_index, pointer_attn)

我想把代码"p_context_ptr.scatter_add_(2, context_index, pointer_attn)"改成tensorflow版本。

所以我使用了tensorflow函数的“tf.compat.v1.tensor_scatter_nd_add()”，但不是同一个操作torch scatter_add_()函数

我一直在尝试工作直到现在，但我的一些代码没有找到解决方案

def get_scatter_add(tensor, indices, updates):
    if indices.shape.rank > 2:
        tensor = tf.compat.v1.reshape(tensor, shape=[-1, tensor.shape[-1]])
        indices = tf.compat.v1.reshape(indices, shape=[-1, indices.shape[-1]])
        updates = tf.compat.v1.reshape(updates, shape=[-1, updates.shape[-1]])

    one_hot_index = tf.compat.v1.one_hot(indices=indices, depth=tensor.shape[-1])

    tile_update = tf.compat.v1.expand_dims(updates, axis=-1)
    updates = tf.compat.v1.to_float(one_hot_index) * tf.compat.v1.to_float(tile_update)
    indices = tf.compat.v1.expand_dims(indices, axis=-1)

    update = tensor.shape[indices.shape[-1]:]
    res = indices.shape[:-1] + update

    scatter = tf.compat.v1.tensor_scatter_nd_add(tensor, indices, updates)
    return scatter

但是，内存溢出，我的变量形状是 tensor.shape()->[1100, 19200], update.shape()->[1100, 900], updates.shape()->[1100 , 900]

这个问题怎么解决？？？

感谢您的回复

祝你有美好的一天！！！

Answer 1

我自己找到了解决办法

tensorflow tensor_scatter_nd_add 函数是一些问题向量维度被扩展为目标向量。 但除了一种情况外，与 torch scatter_add_ fucntion 的操作相同 这种情况：

import tensorflow as tf
indices = tf.constant([[4], [3], [1], [7]])
updates = tf.constant([9, 10, 11, 12])
tensor = tf.ones([8], dtype=tf.int32)
updated = tf.tensor_scatter_nd_add(tensor, indices, updates)
print(updated)

它只更新，张量一维和索引是等级 2 形状 所以我像上面这样改变形状

tensor.shape()->reshape[-1]
update.shape()->reshape[-1]
indices.shape()->reshape[-1, 1]

与上述情况相同，但我们需要更新索引操作，但如果我们有用于 DST 任务的指针生成器，因为张量是最后一维的词汇量，所以索引 + 词汇量下一批和 +vocab*2 下一批

所以它具有相同的操作 Torch scatter_add_

示例： 张量 = [35, 32, vocab_size]，索引 = [35, 32, 900]，更新 = [35, 32, 900]

手电筒： tensor.scatter_add_(2, 索引, 更新)

Tensorflow 案例： tensor = my_tensorflow_scatter_add(张量、索引、更新)

可变维度以上同样的操作案例

my_tensorflow_scatter_add 函数：

    def my_tensorflow_scatter_add(tensor, indices, updates):
    original_tensor = tensor
    # expand index value from vocab size
    indices = tf.compat.v1.reshape(indices, shape=[-1, tf.shape(indices)[-1]])
    indices_add = tf.compat.v1.expand_dims(tf.range(0, tf.shape(indices)[0], 1)*(tf.shape(tensor)[-1]), axis=-1)
    indices += indices_add

    # resize
    tensor = tf.compat.v1.reshape(tensor, shape=[-1])
    indices = tf.compat.v1.reshape(indices, shape=[-1, 1])
    updates = tf.compat.v1.reshape(updates, shape=[-1])

    #check_
    """
    update = tensor.shape[indices.shape[-1]:]
    res = indices.shape[:-1] + update
    """
    #same Torch scatter_add_
    scatter = tf.compat.v1.tensor_scatter_nd_add(tensor, indices, updates)
    scatter = tf.compat.v1.reshape(scatter, shape=[tf.shape(original_tensor)[0], tf.shape(original_tensor)[1], -1])
    return scatter

我解决了我的问题

Answer 2

不展平所有张量的替代解决方案。假设张量形状tensor = [35, 32, vocab_size], indices = [35, 32, 900], update = [35, 32, 900]（基于Proper usage of `tf.scatter_nd` in tensorflow-r1.2）：

def scatter_add(tensor, indices, updates):
    """

    Args:
        tensor: (seq_len, batch_size, vocab_size)
        indices: (seq_len, batch_size, dim)
        updates: (seq_len, batch_size, dim)

    Returns:
        (seq_len, batch_size, vocab_size)
    """
    seq_len, batch_size, dim = indices.shape
    # Create additional indices
    i1, i2 = tf.meshgrid(tf.range(seq_len),
                         tf.range(batch_size), indexing="ij")
    i1 = tf.tile(i1[:, :, tf.newaxis], [1, 1, dim])
    i2 = tf.tile(i2[:, :, tf.newaxis], [1, 1, dim])
    # Create final indices
    idx = tf.stack([i1, i2, indices], axis=-1)
    # Get scatter-added tensor
    scatter = tf.tensor_scatter_nd_add(tensor, idx, updates)
    return scatter