Tensorの軸を拡張する方法

ディープラーニングに関するプログラミングをしている際の形状変更の一種として、軸を拡張するというケースがあります。本記事では、Tensorの軸を拡張する方法を紹介します。

なお、Tensorの形状を変更する方法については「Tensorの形状を変更する方法」で説明しているのであわせて見ていただけるとよいかと思います。

tf.newaxisを使用して軸を追加する

Tensorの軸を拡張する一つの方法としては、tf.newaxisを使用する方法があります。以降でいくつかの例をみてみましょう。

例1：行ベクトルや列ベクトルにする

tf.newaxisを使って行ベクトル(1, n)や列ベクトル(n, 1)を作るような場合には、以下のように使用します。

import tensorflow as tf

tensor = tf.constant([1, 2, 3, 4, 5], dtype=tf.float16)
print(tensor, "\n")

# ===== Tensorの軸を拡張する tf.newaxis
# 先頭に軸を追加する
expanded_tensor_1 = tensor[tf.newaxis, :]
print(expanded_tensor_1, "\n")

# 末尾に軸を追加する
expanded_tensor_2 = tensor[:, tf.newaxis]
print(expanded_tensor_2)

tf.Tensor([1. 2. 3. 4. 5.], shape=(5,), dtype=float16) 

tf.Tensor([[1. 2. 3. 4. 5.]], shape=(1, 5), dtype=float16) 

tf.Tensor(
[[1.]
 [2.]
 [3.]
 [4.]
 [5.]], shape=(5, 1), dtype=float16)

tf.newaxisを使用した方法では、要素参照の[]で軸を追加したい部分にtf.newaxisを指定します。例えば、先頭に軸を追加したい場合にはtensor[tf.newaxis, :]といった形です。

例2：2階以上のTensorに軸を追加

tf.newaxisの使い方は、2階以上のTensorでも考え方は同じです。以下の例では、先頭、真ん中、末尾に軸を追加しています。

import tensorflow as tf

tensor_2d = tf.constant([[1, 2], [3, 4]], dtype=tf.float16)
print(tensor_2d, "\n")

# ===== Tensorの軸を拡張する tf.newaxis
# 先頭に軸を追加する
expanded_tensor_1 = tensor_2d[tf.newaxis, :, :]
print(expanded_tensor_1, "\n")

# 真ん中に軸を追加する
expanded_tensor_2 = tensor_2d[:, tf.newaxis, :]
print(expanded_tensor_2, "\n")

# 末尾に軸を追加する
expanded_tensor_3 = tensor_2d[:, :, tf.newaxis]
print(expanded_tensor_3)

tf.Tensor(
[[1. 2.]
 [3. 4.]], shape=(2, 2), dtype=float16) 

tf.Tensor(
[[[1. 2.]
  [3. 4.]]], shape=(1, 2, 2), dtype=float16) 

tf.Tensor(
[[[1. 2.]]

 [[3. 4.]]], shape=(2, 1, 2), dtype=float16) 

tf.Tensor(
[[[1.]
  [2.]]

 [[3.]
  [4.]]], shape=(2, 2, 1), dtype=float16)

上記結果のshapeを見るとそれぞれで先頭、真ん中、末尾に軸が追加できていることが分かるかと思います。

tf.expand_dimsを使用して軸を拡張する

Tensorの軸を拡張するもう一つの方法は、tf.expand_dimsを使用する方法です。以降で、tf.newaxisの場合と同じ例で見てみましょう。

例1：行ベクトルや列ベクトルにする

tf.expand_dimsを使って行ベクトル(1, n)や列ベクトル(n, 1)を作るような場合には、以下のように使用します。

import tensorflow as tf

tensor = tf.constant([1, 2, 3, 4, 5], dtype=tf.float16)
print(tensor, "\n")

# ===== Tensorの軸を拡張する tf.expand_dims
# 先頭に軸を追加する
expanded_tensor_1 = tf.expand_dims(tensor, axis=0)
print(expanded_tensor_1, "\n")

# 末尾に軸を追加する
expanded_tensor_2 = tf.expand_dims(tensor, axis=-1)
print(expanded_tensor_2)

tf.Tensor([1. 2. 3. 4. 5.], shape=(5,), dtype=float16) 

tf.Tensor([[1. 2. 3. 4. 5.]], shape=(1, 5), dtype=float16) 

tf.Tensor(
[[1.]
 [2.]
 [3.]
 [4.]
 [5.]], shape=(5, 1), dtype=float16)

tf.expand_dimsを使用する場合は、拡張する元のTensorと拡張したい軸をaxis引数で指定します。例えば、先頭に追加したい場合には「axis=0」とします。

上記の例では、末尾に追加する例では「axis=-1」としています。もちろん上記の例で「axis=1」としても同様の結果が得られますが、末尾に追加する場合には「axis=-1」で指定する場合も多いですので覚えておきましょう。

例2：2階以上のTensorに軸を追加

tf.expand_dimsの使い方は、2階以上のTensorでも考え方は同じです。以下の例では、先頭、真ん中、末尾に軸を追加しています。

import tensorflow as tf

tensor_2d = tf.constant([[1, 2], [3, 4]], dtype=tf.float16)
print(tensor_2d, "\n")

# ===== Tensorの軸を拡張する tf.expand_dims
# 先頭に軸を追加する
expanded_tensor_1 = tf.expand_dims(tensor_2d, axis=0)
print(expanded_tensor_1, "\n")

# 真ん中に軸を追加する
expanded_tensor_2 = tf.expand_dims(tensor_2d, axis=1)
print(expanded_tensor_2, "\n")

# 末尾に軸を追加する
expanded_tensor_3 = tf.expand_dims(tensor_2d, axis=-1)
print(expanded_tensor_3)

tf.Tensor(
[[1. 2.]
 [3. 4.]], shape=(2, 2), dtype=float16) 

tf.Tensor(
[[[1. 2.]
  [3. 4.]]], shape=(1, 2, 2), dtype=float16) 

tf.Tensor(
[[[1. 2.]]

 [[3. 4.]]], shape=(2, 1, 2), dtype=float16) 

tf.Tensor(
[[[1.]
  [2.]]

 [[3.]
  [4.]]], shape=(2, 2, 1), dtype=float16)

上記結果のshapeを見るとそれぞれで先頭、真ん中、末尾に軸が追加できていることが分かるかと思います。

まとめ

Googleによって開発されている機械学習ライブラリであるTensorFlowで、Tensorの軸を拡張する方法を解説しました。

軸を拡張する方法としては、tf.newaxisを使用する方法とtf.expand_dimsを使用する方法があります。データ操作で形状を揃えるとき等に軸追加をして拡張することがよくありますので、方法をしっかり覚えておきましょう。

あわせて読みたい

【Python Tech】プログラミングガイド