Tensorflow(Keras)のNormalizationで多次元配列の平均を取るのにaxisの定義がNumpyのmeanなどと違うのにハマる…タプルを使わないとだめなのか。

もう世の中から100周回遅れくらいだろうけれど機械学習で遊び始めた。とりあえずTensorflow（Keras）を使ってみている。

で、早速ハマったところがある（が、ググってもあまりいいのが出てこない）ので備忘録としてここに記す。

データとしては2次元配列で与えられていて、列ごとに標準化したい。検証用のデータを、

data = np.arange(24).reshape(4,3,2)

とする。4つの3x2行列があることに相当。

array([[[ 0, 1],
[ 2, 3],
[ 4, 5]],
[[ 6, 7],
[ 8, 9],
[10, 11]],
[[12, 13],
[14, 15],
[16, 17]],
[[18, 19],
[20, 21],
[22, 23]]])

numpyで平均をとるのには、まずaxis設定なしなら

np.mean(data)

11.5

になる。

axis = 0なら3x2行列になる。これが最終的に欲しい平均。

np.mean(data, axis = 0)

array([[ 9., 10.],
[11., 12.],
[13., 14.]])

axis =1なら4x2行列。

np.mean(data, axis = 1)

array([[ 2., 3.],
[ 8., 9.],
[14., 15.],
[20., 21.]])

axis =2またはaxis = -1なら4x3行列になる。

np.mean(data, axis = -1)

array([[ 0.5, 2.5, 4.5],
[ 6.5, 8.5, 10.5],
[12.5, 14.5, 16.5],
[18.5, 20.5, 22.5]])

さて問題はここから。Tensorflow.Keras.layers.Normalizationを使うとき。

axis = Noneなら

norm1 = tf.keras.layers.Normalization(axis = None)

norm1.adapt(data)

norm1.mean

<tf.Tensor: shape=(1, 1, 1), dtype=float32, numpy=array([[[11.5]]], dtype=float32)>

なんでこれは分かる。

axis=0なら

norm2 = tf.keras.layers.Normalization(axis = 0)

norm2.adapt(data)

norm2.mean

ValueError: All `axis` values to be kept must have known shape. Got axis: (0,), input shape: [None, 3, 2], with unknown axis at index: 0

というエラー。ああ最初の添え字はNoneになるからか…え？じゃあmeanのaxis=0に相当するのはどうすんの？

とりあえず別のaxisを試そう。axis=1のときは？

norm3 = tf.keras.layers.Normalization(axis = 1)

norm3.adapt(data)

norm3.mean

<tf.Tensor: shape=(1, 3, 1), dtype=float32, numpy=
array([[[ 9.5],
[11.5],
[13.5]]], dtype=float32)>

ああ、こうなるのか。1x3x1になる。

axis=-1(デフォルト）or axis=2なら

norm4 = tf.keras.layers.Normalization(axis = -1)

norm4.adapt(data)

norm4.mean

<tf.Tensor: shape=(1, 1, 2), dtype=float32, numpy=array([[[11., 12.]]], dtype=float32)>

1x1x2になる。

じゃあどうするんだろう、、、とドキュメントをちゃんと読む。

https://www.tensorflow.org/api_docs/python/tf/keras/layers/Normalization

Integer, tuple of integers, or None. The axis or axes that should have a separate mean and variance for each index in the shape. 

と書いてある。あ、タプルを受け取るのか！じゃあ、

norm5 = tf.keras.layers.Normalization(axis = (1,2))

norm5.adapt(data)

norm5.mean

<tf.Tensor: shape=(1, 3, 2), dtype=float32, numpy=
array([[[ 9., 10.],
[11., 12.],
[13., 14.]]], dtype=float32)>

これだ！やっとやり方が分かった…