Skip to main content
 首页 » 编程设计

python之TensorFlow 中稀疏加法的 scatter_nd add() 示例

2025年01月19日30freeliver54

我在将 tf.scatter_nd_add() 应用于二维张量时遇到困难。该文档有点不清楚,没有包含稀疏更新的示例,但仅包含完整切片更新的示例。

我的情况如下:

  • 更新 - 形状为 [None, 6] 的 2D 张量
  • indices - 形状为 [None, 6]
  • 的二维张量
  • ref - 形状为 [None, 6]
  • 零点的二维变量

保证 updates、indices 和 ref 的第一个维度始终相等,但该维度的大小可以变化.我要执行的更新看起来像

for i, j: 
    k = indices[i][j] 
    ref[i][k] += updates[i][j]  

请注意,indices 包含重复项。 tf.scatter_nd_add(ref, indices, updates) 提示形状不匹配,我无法弄清楚我需要如何重组张量才能执行更新。

请您参考如下方法:

我想通了。 indices 中的每个 2D 条目实际上必须指定将在 ref 中更新的绝对位置。这意味着 indices 必须是 3D,然后非矢量化更新看起来像:

for i, j: 
    r, k = indices[i][j] 
    ref[r][k] += updates[i][j] 

在上面的问题中,恰好 r 总是等于 i。

这是具有不同形状的完整 Tensorflow 实现。为清楚起见,在以下示例中,col_indices 对应于原始问题中的 indices:

import tensorflow as tf 
import numpy as np 
 
updates     = tf.placeholder(dtype=tf.float32,  shape=[None, 6]) 
col_indices = tf.placeholder(dtype=tf.int32,    shape=[None, 6]) 
row_indices = tf.cumsum(tf.ones_like(col_indices), axis=0, exclusive=True) 
indices     = tf.concat([tf.expand_dims(row_indices, axis=-1),  
                         tf.expand_dims(col_indices, axis=-1)], axis=-1) 
 
tmp_var     = tf.Variable(0, trainable=False, dtype=tf.float32, validate_shape=False) 
ref         = tf.assign(tmp_var, tf.zeros_like(updates), validate_shape=False) 
# This makes sure that ref is always 0 before scatter_nd_add() runs 
with tf.control_dependencies([target_var]): 
  result = tf.scatter_nd_add(ref, indices, updates) 
 
sess = tf.Session() 
sess.run(tf.global_variables_initializer()) 
 
# Create example input data  
np_input = np.arange(0, 6, 1, dtype=np.int32) 
np_input = np.tile(np_input[None,:], [10, 1]) 
 
res = sess.run(result, feed_dict={updates: np_input, col_indices: np_input}) 
print(res)