Tensorflow 2下怎么根据Index位置获取或设置Tensor的值呢？

大家知道在numpy下获取/设置array的值是很方便的，但是tensorflow下由于一些限制就没那么方便了。下面就对比一下。

注意，tensor是不可变的对象，所谓改变tensor的值是指用新的值替换旧的值，并返回一个新的tensor对象。

numpy下的操作

取值

# 1.
val_at_idx = some_array[idx]

# 2.
boolean_mask = [True, False, ...]
val_at_mask = some_array[boolean_mask]

设置值

some_array[idx] = new_val

Tensorflow下的操作

取值

单个值可以用tensor[idx]方法获得，其他情况需要使用gather、gather_nd或boolean_mask方法。

import tensorflow as tf
# 1.
val_at_idx = tf.gather(some_tensor, idx_array)

# 2.
boolean_mask = [True, False, ...]
val_at_idx = tf.boolean_mask(some_tensor, boolean_mask)

设置值

如果直接用tensor[idx]=new_val的方式会报下面的错误：

AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute 'assign'

设置值分两种情况，如果是根据值的条件来改变，那么优先使用tf.where方法，它会进行elementwise的操作，以下面的numpy操作作为例子：

>>> import numpy as np
>>> a = np.array( [1,2,3,4,5] )
>>> print(a)
[1 2 3 4 5]
>>> a_max = 3
>>> i = np.where( a > a_max)
>>> a[i] = a_max
>>> print(a)
[ 1  2  3  3  3]

tensorflow:

>>> a = tf.convert_to_tensor([1, 2, 3, 4, 5], dtype=tf.int64)
>>> print(a)
tf.Tensor([1 2 3 4 5], shape=(5,), dtype=int64)
>>> a_max = 3
>>> a = tf.where(a > a_max, a_max, a)
>>> print(a)
tf.Tensor([1 2 3 3 3], shape=(5,), dtype=int64)

第二种情况是改变指定index位置的值，比如像下面的操作：

idx = [1,3]
a[idx] = 0  <==== tensorflow出错

这种操作可以创建一个SparseTensor来代表对应Index的改变值，然后用tf.add操作来合并原Tensor和SparseTensor。

# 需要改变值的位置，注意位置要递增顺序，因为对SparseTensor的操作是如此假定的。如果不是递增的话需要调用tf.sparse.reorder来调整
idx_to_set = np.array([1, 3])
# 改变后的值
vals_to_set = np.array([0, 0])

# 取得现在位置的值
val_at_idx = tf.gather(a, idx_to_set)
# 取得要改变的差值，方便用tf.add进行合并
delta_vals_at_idx = tf.convert_to_tensor(vals_to_set) - val_at_idx

delta_sparse = tf.sparse.SparseTensor(indices=idx_to_set[:, tf.newaxis], values=delta_vals_at_idx, dense_shape=tf.shape(a, out_type=tf.int64))
a = tf.add(a, tf.sparse.to_dense(delta_sparse))
print(a)
# tf.Tensor([1 0 3 0 3], shape=(5,), dtype=int64)

谢谢你的阅读，如果能帮到你，欢迎点赞、关注！谢谢！