Tensorflow：テンソルの最後の次元に関数を適用する方法

Question

私は配列を持っています。この配列の最後の次元の値に基づいてマスクを作成したいと思います。 numpyのでは、私は何ができる：上記のコードで

import numpy as np 
room = np.array([ 
    [[0, 0, 1], [1, 0, 0], [1, 0, 0]], 
    [[1, 0, 0], [0, 0, 1], [1, 0, 0]], 
    [[1, 0, 0], [1, 0, 0], [0, 0, 1]] 
]) 
mask = np.apply_along_axis(lambda x: [1, 1, 1] if (x == [0, 0, 1]).all() else [0, 0, 0], axis=-1, arr=room) 
result = mask * room 
print(result) 

をroom私は、マスクを作成したかに基づいて（3、3、3）配列です。作成されたマスクは（3,3,3）配列でもあり、不要な要素をマスクするために他の配列と乗算するために使用されます。

しかし、私はTensorflowで同じことを達成することに問題があります。私は

room = tf.constant([ 
    [[0, 0, 1], [1, 0, 0], [1, 0, 0]], 
    [[1, 0, 0], [0, 0, 1], [1, 0, 0]], 
    [[1, 0, 0], [1, 0, 0], [0, 0, 1]] 
]) 
room = tf.reshape(room, shape=(9, -1)) 
mask = tf.map_fn(lambda x: [1, 1, 1] if x == [0, 0, 1] else [0, 0, 0], room) 

、次のコードを試してみましたが、それは次のエラーで終了：

ValueError: The two structures don't have the same number of elements. First structure: <dtype: 'int32'>, second structure: [0, 0, 0]. 

Answer 1

map_fnは、それがxの形状とは異なるかどう出力の形状を指定することができますパラメータdtypeを持っています。

ただし、これは問題ではありません。

Pythonの条件をテンソルフローの操作に混ぜると、pythonの操作（ifのような）がグラフの外側で実行され、代わりに、目的のopを実行するグ​​ラフを定義する必要があります。

のはあなたの問題を掘り下げるてみましょう：

1. あなたが最初の次元に沿ってroom変数アンロールしたい：map_fnは、そのためのOKです。
2. 現在の行（3要素のベクトル）が[0, 0, 1]に等しいかどうかを確認する必要があります。

これを行うには、テンソルフロー条件tf.cond(pred, true_fn, false_fn)を使用する必要があります。

predはスカラーでなければならないことに注意してください。したがって、現在のrowが目的の行と等しく、その結果を単一のスカラーに減らした場合は、テンソルフロー操作のみを使用して調べてみましょう。

定数値[1,1,1]を返します。それ以外の場合は[0,0,0]です。

mask = tf.map_fn(lambda row: tf.cond(
    tf.equal(
     tf.reduce_prod(tf.cast(tf.equal(row, tf.constant([0,0,1])), tf.int32)), 1), 
    lambda: tf.constant([1,1,1]), 
    lambda: tf.constant([0,0,0])), room)