将原矩阵按照索引矩阵的方式分割
alist=[[[[1,1,1],[2,2,2],[3,3,3]],[[4,4,4],[5,5,5],[6,6,6]]],[[[7,7,7],[8,8,8],[9,9,9]],[[10,10,10],[11,11,11],[12,12,12]]]]
aal=tf.reshape(alist, [-1, ])
t=tf.to_int32(aal <= 5)
valid_logits = tf.dynamic_partition(aal, t, num_partitions=2)
with tf.Session() as sess:
print(sess.run(valid_logits))
输出:
[array([ 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 11,
11, 12, 12, 12]), array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5])]