现在,我的函数使用了argmax:
p = tf.stop_gradient(tf.argmax(prev, 1))
我尝试使用以下内容,但dimn不兼容:
p = tf.stop_gradient(tf.nn.top_k(prev, 2)[1])
raise ValueError("Linear is expecting 2D arguments: %s" % str(shapes))
ValueError: Linear is expecting 2D arguments: [[None, 2, 1024], [None, 1024]]
我的TF版本可能是0.5,这就是为什么top_k只有2个args。
答案 0 :(得分:3)
查看tf.nn.top_k()的文档。该函数返回值和索引。所以下面的内容应该有用。
values, indices = tf.nn.top_k(prev,2)
p = tf.stop_gradient(indices[1])