遍历' tf.不允许使用张量

Question

我试图使用这个函数到一个@tf。函数decorato:

h和h2是形状为[3,3]

的张量

def fn(h,i):
    print(h[i])
    return h[i]
tensor = [fn(h,i) for i in tf.range(tf.cast(tf.shape(h)[0],tf.int32)) if  tf.reduce_all(tf.equal(h[i],h2[i])) ]
tf.print(tensor)

但是我得到这个错误:

main_coat_rds.py:139 train_step  *
        pseudo_label_1,images_discard_rede1=predict_aug_images(rede_2,rede_1,img_rede1_aug_1,img_rede1_aug_2,img_rede1_aug_3,img_rede1_aug_4,img_rede1_aug_5,img_rede1_aug_6,img_rede1_aug_7,img_rede1_aug_8,images_discard_rede1,Correct_labels)
    /vitor/codigo_noise_label/codigo_rds/utils_loss_function.py:289 predict_aug_images  *
        pred_match = [get_value_labels(all_predics,i) for i in tf.range(tf.cast(tf.shape(all_predics)[0],tf.int32)) if  tf.reduce_all(tf.equal(all_predics[i],all_predics_aj[i])) ]
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py:503 __iter__
        self._disallow_iteration()
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py:496 _disallow_iteration
        self._disallow_when_autograph_enabled("iterating over `tf.Tensor`")
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py:474 _disallow_when_autograph_enabled
        " indicate you are trying to use an unsupported feature.".format(task))
    OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

我能做的另一种方法是什么?

Answer 1

当你想根据一个条件选择张量的某些部分时，一个好的选择是使用tf.gather和tf.where的组合。

在这里，例如，要选择h和h2之间相等的行，您可以使用:

tf.gather_nd(h, tf.where(tf.reduce_all(tf.equal(h, h2),axis=1)))