はじめに
TensorFlow 2.1以降で導入された、可変長のデータを表す RaggedTensor ですが、普通のTensorのノリで書こうとすると色々ハマります。
今回はIndexing編。RaggedTensor から特定のインデックスを指定して値を取り出したりしてみます。慣れてくると複雑な操作もできるようになる…はず。
検証環境
- Ubuntu 18.04
- Python 3.6.9
- TensorFlow 2.2.0 (CPU)
Indexingの例
Indexing対象の RaggedTensor として、x が以下のように作られているとします。
x = tf.RaggedTensor.from_row_lengths(tf.range(15), tf.range(1, 6))
print(x)
# <tf.RaggedTensor [[0], [1, 2], [3, 4, 5], [6, 7, 8, 9], [10, 11, 12, 13, 14]]>
| 列番号 | 0 | 1 | 2 | 3 | 4 |
|---|---|---|---|---|---|
| 0行目 | 0 | ||||
| 1行目 | 1 | 2 | |||
| 2行目 | 3 | 4 | 5 | ||
| 3行目 | 6 | 7 | 8 | 9 | |
| 4行目 | 10 | 11 | 12 | 13 | 14 |
特定の行でスライシング
まずはある行を取り出す操作ですが、これは普通の Tensor と同じです。numpy.ndarray の感覚で考えてもよいでしょう。範囲指定した場合、**最初のインデックスを含み、最後のインデックスを含みません。**Python使いの方なら問題ないと思いますが。
print(x[2])
# tf.Tensor([3 4 5], shape=(3,), dtype=int32)
print(x[1:4])
# <tf.RaggedTensor [[1, 2], [3, 4, 5], [6, 7, 8, 9]]>
ただし numpy.ndarray と異なり、飛び飛びの行を指定するスライシングは使えないようです。
# ndarrayに対してはこれでできる
print(x.numpy()[[1, 3]])
# [array([1, 2], dtype=int32) array([6, 7, 8, 9], dtype=int32)]
# Tensor/RaggedTensorに対しては使えない
print(x[[1, 3]])
# InvalidArgumentError: slice index 3 of dimension 0 out of bounds. [Op:StridedSlice] name: strided_slice/
代わりにこちらでどうぞ。
# Tensor/RaggedTensorでFancy Indexing
print(tf.gather(x, [1, 3], axis=0))
# <tf.RaggedTensor [[1, 2], [6, 7, 8, 9]]>
固定の列インデックスでスライシング
次に固定の列インデックスでスライシングする例です。
普通の Tensor と違って、行によってそのインデックスの要素があるかないかが変わるので、単純に
print(x[:, 2])
# ValueError: Cannot index into an inner ragged dimension.
のようにはできないようになっています。範囲指定をすれば
print(x[:, 2:3])
# <tf.RaggedTensor [[], [], [5], [8], [12]]>
のように動きます。指定したインデックスが存在しない行に対しては [] となっています。
| 列番号 | 0 | 1 | 2 | 3 | 4 |
|---|---|---|---|---|---|
| 0行目 | 0 | ||||
| 1行目 | 1 | 2 | |||
| 2行目 | 3 | 4 | 5 | ||
| 3行目 | 6 | 7 | 8 | 9 | |
| 4行目 | 10 | 11 | 12 | 13 | 14 |
行ごとに異なる列インデックスを指定してスライシング
集めたい2次元のインデックスを列挙した Tensor がある場合は、tf.gather_nd() が使えます。
ind = tf.constant([[0, 0], [1, 1], [2, 0], [4, 3]])
# x の (0, 0), (1, 1), (2, 0), (4, 3) 要素を集めたい
print(tf.gather_nd(x, ind))
# tf.Tensor([ 0 2 3 13], shape=(4,), dtype=int32)
| 列番号 | 0 | 1 | 2 | 3 | 4 |
|---|---|---|---|---|---|
| 0行目 | 0 | ||||
| 1行目 | 1 | 2 | |||
| 2行目 | 3 | 4 | 5 | ||
| 3行目 | 6 | 7 | 8 | 9 | |
| 4行目 | 10 | 11 | 12 | 13 | 14 |
一方、行ごとに要素を1つずつ取ってくるのですが、それぞれ異なる列から取ってきたい、という場合もあると思います。
col = tf.constant([0, 0, 2, 1, 2])
# x の (0, 0), (1, 0), (2, 2), (3, 1), (4, 2) 要素を集めたい
# インデックスに行番号をつけてから、先程と同じ方法で
ind = tf.transpose(tf.stack([tf.range(tf.shape(col)[0]), col]))
print(tf.gather_nd(x, ind))
# tf.Tensor([ 0 1 5 7 12], shape=(5,), dtype=int32)
| 列番号 | 0 | 1 | 2 | 3 | 4 |
|---|---|---|---|---|---|
| 0行目 | 0 | ||||
| 1行目 | 1 | 2 | |||
| 2行目 | 3 | 4 | 5 | ||
| 3行目 | 6 | 7 | 8 | 9 | |
| 4行目 | 10 | 11 | 12 | 13 | 14 |
でもなんとなく遅そうな気がするので、もう少し賢い方法を考えてみると
print(tf.gather(x.values, x.row_starts() + col))
# tf.Tensor([ 0 1 5 7 12], shape=(5,), dtype=int32)
これでOKです。
x の値の実体は各行をつなげた（1次元少ない）Tensor （RaggedTensor ではありません）に入っていて、x.values にアクセスすると取得できます。また、x の形状を表すために各行の開始インデックス (x.row_starts()) の情報が保持されています。よって、このインデックスに指定したオフセットを加え、x.values に対してスライシングすればよいというわけです。
%timeit tf.gather_nd(x, tf.transpose(tf.stack([tf.range(tf.shape(col)[0]), col])))
# 739 µs ± 75.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit tf.gather(x.values, x.row_starts() + col)
# 124 µs ± 6.47 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
こっちのほうが速いですね(^_^)
このあたりの操作を極めようと思えば、公式ドキュメントを見るのが吉です。
列インデックスがRaggedTensorに入っている場合
先程の「値の実体が1次元の Tensor に入っている」ということを応用します。
col = tf.ragged.constant([[0], [], [0, 2], [1, 3], [2]])
# x の (0, 0), (2, 0), (2, 2), (3, 1), (3, 3), (4, 2) 要素を集めたい
# xの各行の開始インデックスを得る
row_starts = tf.cast(x.row_starts(), "int32")
# colの各成分が属する行番号を得て、xにおける開始インデックスに変換し、オフセットを加える
ind_flat = tf.gather(row_starts, col.value_rowids()) + col.values
ret = tf.gather(x.values, ind_flat)
print(ret)
# tf.Tensor([ 0 3 5 7 9 12], shape=(6,), dtype=int32)
| 列番号 | 0 | 1 | 2 | 3 | 4 |
|---|---|---|---|---|---|
| 0行目 | 0 | ||||
| 1行目 | 1 | 2 | |||
| 2行目 | 3 | 4 | 5 | ||
| 3行目 | 6 | 7 | 8 | 9 | |
| 4行目 | 10 | 11 | 12 | 13 | 14 |
元の行の情報を保存したい場合
先程の結果は値を並べた普通の Tensor になっていて、元の行の情報は失われていますが、行の情報を保存したい場合はどうすればよいでしょう。
Tensor に対して行の開始インデックスの情報を与えることで RaggedTensor を作ることができます。各行の長さは col と一致するはずなので、この開始インデックスを col.value_rowids() から持ってくればよいですね。
print(tf.RaggedTensor.from_value_rowids(ret, col.value_rowids()))
# <tf.RaggedTensor [[0], [], [3, 5], [7, 9], [12]]>
対象のRaggedTensorが3次元以上の場合
2次元以上のデータが時系列に並んでいる（バッチ次元を含めて、RaggedTensor としては3次元以上）場合でも、今までの方法がそのまま使えます。
x = tf.RaggedTensor.from_row_lengths(tf.reshape(tf.range(30), (15, 2)), tf.range(1, 6))
print(x)
# <tf.RaggedTensor [[[0, 1]], [[2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]], [[12, 13], [14, 15], [16, 17], [18, 19]], [[20, 21], [22, 23], [24, 25], [26, 27], [28, 29]]]>
この x の構造は以下のように解釈できます。
| 列番号 | 0 | 1 | 2 | 3 | 4 |
|---|---|---|---|---|---|
| 0行目 | [0, 1] | ||||
| 1行目 | [2, 3] | [4, 5] | |||
| 2行目 | [6, 7] | [8, 9] | [10, 11] | ||
| 3行目 | [12, 13] | [14, 15] | [16, 17] | [18, 19] | |
| 4行目 | [20, 21] | [22, 23] | [24, 25] | [26, 27] | [28, 29] |
あとはこれまでと全く同じです。ただし、返ってくる Tensor が2次元になっていることにご注意ください。
ind = tf.constant([[0, 0], [1, 1], [2, 0], [4, 3]])
# x の (0, 0), (1, 1), (2, 0), (4, 3) 要素を集めたい
print(tf.gather_nd(x, ind))
# tf.Tensor(
# [[ 0 1]
# [ 4 5]
# [ 6 7]
# [26 27]], shape=(4, 2), dtype=int32)
col = tf.constant([0, 0, 2, 1, 2])
# x の (0, 0), (1, 0), (2, 2), (3, 1), (4, 2) 要素を集めたい
print(tf.gather(x.values, x.row_starts() + col))
# tf.Tensor(
# [[ 0 1]
# [ 2 3]
# [10 11]
# [14 15]
# [24 25]], shape=(5, 2), dtype=int32)
col = tf.ragged.constant([[0], [], [0, 2], [1, 3], [2]])
# x の (0, 0), (2, 0), (2, 2), (3, 1), (3, 3), (4, 2) 要素を集めたい
# xの各行の開始インデックスを得る
row_starts = tf.cast(x.row_starts(), "int32")
# colの各成分が属する行番号を得て、xにおける開始インデックスに変換し、オフセットを加える
ind_flat = tf.gather(row_starts, col.value_rowids()) + col.values
ret = tf.gather(x.values, ind_flat)
print(ret)
# tf.Tensor(
# [[ 0 1]
# [ 6 7]
# [10 11]
# [14 15]
# [18 19]
# [24 25]], shape=(6, 2), dtype=int32)
# 元の行の情報を保存したい場合
print(tf.RaggedTensor.from_value_rowids(ret, col.value_rowids()))
# <tf.RaggedTensor [[[0, 1]], [], [[6, 7], [10, 11]], [[14, 15], [18, 19]], [[24, 25]]]>