Укладка Tensorflow или Numpy meshgrid с произвольным количеством штабелируемых элементов - PullRequest
0 голосов
/ 20 июня 2020

Я хочу создать специальную c мне sh сетку. Взгляните на следующий код, который работает правильно для лучшего объяснения (я использую сетку тензорного потока):

myDtype = tf.float32
def grid(central, d):
    b_sz = tf.shape(d)[0]
    y_sz = tf.shape(d)[1]
    x_sz = tf.shape(d)[2]
    u_sz = tf.shape(d)[3]
    v_sz = tf.shape(d)[4]
    
    central = tf.expand_dims(tf.expand_dims(central, 3), 4)
                                            

    b_vals = tf.cast(tf.range(b_sz),myDtype)
    v_vals = tf.math.subtract(tf.cast(tf.range(v_sz),myDtype), tf.cast(v_sz,myDtype)/2.0)
    u_vals = tf.math.subtract(tf.cast(tf.range(u_sz),myDtype), tf.cast(u_sz,myDtype)/2.0)
    y_vals = tf.cast(tf.range(y_sz),myDtype)
    x_vals = tf.cast(tf.range(x_sz),myDtype)

    b, y, x, v, u = tf.meshgrid(b_vals, y_vals, x_vals, v_vals, u_vals, indexing='ij')
    return b, y, x, v, u

следующий пример:

c = tf.random.normal([5,128,128])
r = tf.random.normal([5,128,128,7,7])
ret = grid(c,r)
ret[0].shape

дает форму возвращаемых значений of: TensorShape([5, 128, 128, 7, 7]).

Теперь я хочу изменить функцию на следующую:

def grid2(central, d, u_vector,v_vector):
    b_sz = tf.shape(d)[0]
    y_sz = tf.shape(d)[1]
    x_sz = tf.shape(d)[2]
    
    central = tf.expand_dims(tf.expand_dims(central, 3), 4)
                                            

    b_vals = tf.cast(tf.range(b_sz),myDtype)
    y_vals = tf.cast(tf.range(y_sz),myDtype)
    x_vals = tf.cast(tf.range(x_sz),myDtype)

    _, _, _, v, u = tf.meshgrid(b_vals, y_vals, x_vals, v_vector, u_vector, indexing='ij')
    return v, u

В следующем примере:

c = tf.random.normal([5,128,128])
r = tf.random.normal([5,128,128,7,7])
u_vec = tf.random.normal([5,7])
v_vec = tf.random.normal([5,7])
ret2=grid2(c,r,u_vec,v_vec)
ret2[0].shape

возвращает: TensorShape([5, 128, 128, 35, 35])

Вот в чем проблема. Я хочу, чтобы grid2 имел ту же измененную форму, что и grid. Я действительно хочу иметь функцию, которая может дать мне такой результат:

rets = []
for i in range(u_vec.shape[0]):
    rets.append(grid2(tf.expand_dims(c[i],axis=0),tf.expand_dims(r[i],axis=0),u_vec[i],v_vec[i]))
v = tf.stack([tf.squeeze(rets[0][0]),tf.squeeze(rets[1][0]),tf.squeeze(rets[2][0]),tf.squeeze(rets[3][0]),tf.squeeze(rets[4][0])], axis=0)
v.shape

как grid1: TensorShape([5, 128, 128, 7, 7])

Я, вероятно, плохо понял концепцию меня sh grid но, пожалуйста, помогите мне найти решение. Обратите внимание, что последний результат me sh отличается для каждого тензорного среза, нарезанного по оси = 0. Как я могу получить результат, аналогичный v (создан для l oop) без использования al oop и с возможностью обобщения на произвольной оси ввода = 0 размеров?

Пожалуйста, укажите, если любая часть вопроса расплывчата, поэтому я могу отредактировать ее как можно скорее.

...