Tensorflow получает информацию об измерениях при использовании tf.map_fn () для перебора в тензор

Question

Скажем, у меня есть тензор ts формы [s1, s2, s3], и я хочу перейти к нему с tf.map_fn как:

tf.map_fn(lambda dim1:
    tf.map_fn(lambda dim2:
        do_sth(dim, idx1, idx2)
    ,dim)
,ts)

idx1 и idx2 выше - это индекс измерения 0 и измерения 1 из ts, в котором в данный момент находится do_sth(). Как я могу получить их? Я хочу получить это, как будто я делаю что-то вроде:

for idx1 in range(s1):
    for idx2 in range(s2):
        tensor = ts[idx1][idx2]
        do_sth(tensor, idx1, idx2)

Причина, по которой я не могу сделать это таким образом, заключается в том, что большую часть времени s1, s2, s3 являются неизвестными (то есть ts имеет форму (?, ?, t3) или аналогичную)

Возможно ли это?

Answer 1

Я бы посоветовал вам добавить дополнительное измерение в конце и заполнить его индексами перед выполнением команды.

Этот код (проверенный) добавляет индексы к значениям, умноженным на 10 и 100 для индексов X и Y соответственно:

from __future__ import print_function
import tensorflow as tf
import numpy as np

r = 3

a = tf.reshape( tf.constant( range( r * r ) ), ( r, r ) )
x = tf.tile( tf.cast( tf.lin_space( 0.0, r - 1, r )[ None, : ], tf.int32 ), [ r, 1 ] )
y = tf.tile( tf.cast( tf.lin_space( 0.0, r - 1, r )[ :, None ], tf.int32 ), [ 1, r ] )

b = tf.stack( [ a, x, y ], axis = -1 )

c = tf.map_fn( lambda y: tf.map_fn( lambda x: 
        x[ 0 ] + 10 * x[ 1 ] + 100 * x[ 2 ]
    , y ), b )

with tf.Session() as sess:
    res = sess.run( [ c ] )
    for x in res:
        print()
        print( x )

Выходы:

[[0 11 22]
[103 114 125]
[206 217 228]]