g=tf.random.normal([1000,627,16,1])
ind1=tf.constant(np.random.randint(100,size=(251,128,128,3)))
ind2=tf.constant(np.random.randint(100,size=(150,128,128,3)))
ind3=tf.constant(np.random.randint(100,size=(100,128,128,3)))
f1=tf.gather_nd(g[0:251]*1.0,ind1)
f2=tf.gather_nd(g[30:150+30]*1.0,ind2)
f3=tf.gather_nd(g[120:120+100]*1.0,ind3)
f=tf.concat([f1,f2,f3],0)
print('debug')
ind1,ind2,ind3 可以合并为一个愤怒的张量。但是用于每个 ind 的 g 的形状仍然不统一。我知道解决这个问题的一种方法是将 ind1、ind2 和 ind3 扩展到 g 的相同形状。但是这种方法会消耗过多的 GPU 内存。
有没有一种方法可以通过使用不规则张量在一个函数中实现上述代码，例如：ind=tf.ragged.constant([ind1,ind2,ind3]) f=tf.gather_nd(g,ind)