TensorFlowのテンソル計算で、二次元のユニークな値から三次元のテンソルを素早く作る方法

Question

やりたい事の概要を以下に書きます。
44で要素が0～3のtensor_1から、44*4のtensor_2を作りたい。
tensor_2の3軸目は、インデックスに対応したtensor_1の要素と一致している場合に1となります。

tensor_1
[[0, 2, 1, 3],
[1, 3, 0, 0],
[2, 1, 3, 2],
[0, 2, 3, 1]]

↓

tensor_2[:, :, 0]
[[1, 0, 0, 0],
[0, 0, 1, 1],
[0, 0, 0, 0],
[1, 0, 0, 0]]

tensor_2[:, :, 1]
[[0, 0, 1, 0],
[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 0, 1]]

tensor_2[:, :, 2]
[[0, 1, 0, 0],
[0, 0, 0, 0],
[1, 0, 0, 1],
[0, 1, 0, 0]]

tensor_2[:, :, 3]
[[0, 0, 0, 1],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 1, 0]]

上のは例であり、実際想定しているtenosr_1のシェイプは[1000, 1200]で、取りうる要素は0～65の整数です。
このtensor_2を作るのに、私は以下のコードを書きました。
whereでそれぞれの数字に対してのテンソルを作り、それらを3軸目でconcatして[1000, 1200, 66]のtenosr_2を作っています。

# shapeを取得
h = tf.shape(tensor_1)[0]
w = tf.shape(tensor_1)[1]
c = 66  # チャンネル数

# 0番目の要素を取得
tensor_2 = tf.where(tensor_1 == 0, 1, 0)
tensor_2 = tensor_2.reshape((h, w, 1))  # 軸を増やす

# ループで1～65の要素を加える
for i in range(1, c):
    x = tf.where(tensor_1 == i, 1, 0)
    x = tensor_2.reshape((h, w, 1))

    # 結合する
    tensor_2  = tf.concat([tensor_2 , x], axis=2)  # 3軸目で結合

これにより、望みのtensor_2を作ることができたのですが、予想以上に処理に時間がかかってしました。
処理に時間がかかる場所を調べた結果、それはconcatでした。
配列の末尾を追加するのは時間がかかると聞いたことがあるので、それが原因なのかなと。。。
上記のtensor_2を素早く作る方法をご存じの方がおられましたらご教示お願いします。

〇なぜこの変換が必要か
セグメンテーションの深層学習をする際に、アノテーションデータの容量を節約するためtensor_1の形式で保存しました。
tf.datasetの前処理で、アノテーション毎にチャンネルを作る必要があり、この変換が必要になりました。

〇試したこと
先にtensor_2をtf.zeros((h, w, c), dtype(tf.int32))として作っておき、そこに代入すれば早いのではないかと考えたのですが、テンソルはスライスに代入することができないので失敗しました。
(前処理なので、Variable変数にしてassignすることは出来ないとの理解です)

tensor_2 = tf.zeros((h, w, c), dtype(tf.int32))
for i in range(c):
    tensor_2[: ,:, i] = tf.where(tensor_1 == i, 1, 0)

TypeError: 'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment

〇環境
python3.8
tensorflow2.4

Answer 1

こういう感じでできます。

>>> print(tensor_1)
tf.Tensor(
[[0 2 1 3]
 [1 3 0 0]
 [2 1 3 2]
 [0 2 3 1]], shape=(4, 4), dtype=int32)
>>> print(tensor_i)
tf.Tensor([0 1 2 3], shape=(4,), dtype=int32)
>>> print(tf.where(tf.reshape(tensor_1, (1,4,4)) == tf.reshape(tensor_i, (4,1,1)), 1, 0))
tf.Tensor(
[[[1 0 0 0]
  [0 0 1 1]
  [0 0 0 0]
  [1 0 0 0]]

 [[0 0 1 0]
  [1 0 0 0]
  [0 1 0 0]
  [0 0 0 1]]

 [[0 1 0 0]
  [0 0 0 0]
  [1 0 0 1]
  [0 1 0 0]]

 [[0 0 0 1]
  [0 1 0 0]
  [0 0 1 0]
  [0 0 1 0]]], shape=(4, 4, 4), dtype=int32)

Comments

[Amanokawa]

ありがとうございます。
簡潔かつエレガントな回答とても助かりました。