[TensorFlow] RaggedTensorで窓処理したい

# はじめに

TensorFlow 2.1以降で導入された、可変長のデータを表す [RaggedTensor](https://www.tensorflow.org/api_docs/python/tf/RaggedTensor) ですが、普通のTensorのノリで書こうとすると色々ハマります。
今回は信号処理などで使いそうな窓処理編。ある時間幅を持ったフレームを少しずつずらして、フレームの範囲に入った波形を抽出します。

# 検証環境

- Ubuntu 18.04
- Python 3.6.9
- TensorFlow 2.2.0 (CPU)

# やりたいこと

``x`` をバッチだと思って、各行のデータに対して短い区間の波形を切り出していきます。データの長さはバラバラです。
ここでは、フレーム幅を2として、切り出し位置を1ずつシフトしていきます。
``[3, 1, 4, 1]`` なら ``[[3, 1], [1, 4], [4, 1]]`` といった感じですね。

普通の ``Tensor`` に対しては [tf.signal.frame](https://www.tensorflow.org/api_docs/python/tf/signal/frame) という便利な関数がありますが、残念ながら ``RaggedTensor`` には使えません。

```python
x = tf.ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
print(tf.signal.frame(x, 2, 1)) # NG
# ValueError: TypeError: object of type 'RaggedTensor' has no len()
print(tf.signal.frame(x.to_tensor(), 2, 1)) # 動くけど余分な0がいっぱい出てくる
# tf.Tensor(
# [[[3 1]
#   [1 4]
#   [4 1]]
# 
#  [[0 0]
#   [0 0]
#   [0 0]]
# 
#  [[5 9]
#   [9 2]
#   [2 0]]
# 
#  [[6 0]
#   [0 0]
#   [0 0]]
# 
#  [[0 0]
#   [0 0]
#   [0 0]]], shape=(5, 3, 2), dtype=int32)
```

# 解決策

バッチ次元を外した (flattenした) 値を表す ``x.values`` と、``RaggedTensor`` から得られる各行の長さやオフセットをベースに考えます。

```python
print(x.values)        # 値を並べたTensor
# tf.Tensor([3 1 4 1 5 9 2 6], shape=(8,), dtype=int32)
print(x.row_starts())  # valuesにおける各行の開始インデックス（オフセット）
# tf.Tensor([0 4 4 7 8], shape=(5,), dtype=int64)
print(x.row_lengths()) # 各行の長さ
# tf.Tensor([4 0 3 1 0], shape=(5,), dtype=int64)
```

``x`` の各行ごとに、``x.values`` のどのインデックスから値を取ってくればよいかを考えます(*)。

- 0行目は ``[0, 1], [1, 2], [2, 3]``
- 1行目は空
- 2行目は ``[4, 5], [5, 6]``
- 3行目は空
- 4行目は空

まずは(*)の先頭インデックスを持つような``RaggedTensor`` を作ってやると

```python
s = x.row_starts()
e = s + x.row_lengths() - 1
r = tf.ragged.range(s, e)
print(r)
# <tf.RaggedTensor [[0, 1, 2], [], [4, 5], [], []]>
```

さらに、インデックスを1進めたものを結合すると、窓処理後の期待する結果について、``x.values`` のどこから値を持ってくればよいかが分かります。先の箇条書き(*)と対応する結果が得られています。

```python
ind = tf.stack([r, r+1], axis=2)
print(ind)
# <tf.RaggedTensor [[[0, 1], [1, 2], [2, 3]], [], [[4, 5], [5, 6]], [], []]>
```

あとは、``tf.gather()`` を使って、``ind`` に入ったインデックスに基づいて ``x.values`` から値を取ってくればOK。

```python
ret = tf.gather(x.values, ind)
print(ret)
# <tf.RaggedTensor [[[3, 1], [1, 4], [4, 1]], [], [[5, 9], [9, 2]], [], []]>
```

## フレーム長が3以上の場合

``e`` と ``ind`` の作り方が少し変わりますが、大筋は同じです。
``ind`` の作成にはブロードキャストを利用しています。そのために ``r[:, :, tf.newaxis]`` として末尾に長さ1の次元を追加しています。

```python
len_frame = 3
s = x.row_starts()
e = s + x.row_lengths() + 1 - len_frame
r = tf.ragged.range(s, e)
ind = r[:, :, tf.newaxis] + tf.range(0, len_frame, dtype="int64")
ret = tf.gather(x.values, ind)
print(ret)
# <tf.RaggedTensor [[[3, 1, 4], [1, 4, 1]], [], [[5, 9, 2]], [], []]>
```

もちろん ``len_frame = 2`` の場合でも使えます。

## フレームシフトが2以上の場合

``r`` の刻み幅を変えればOKです。

```python
len_frame = 2
shift_frame = 2
s = x.row_starts()
e = s + x.row_lengths() + 1 - len_frame
r = tf.ragged.range(s, e, shift_frame)
ind = r[:, :, tf.newaxis] + tf.range(0, len_frame, dtype="int64")
ret = tf.gather(x.values, ind)
print(ret)
# <tf.RaggedTensor [[[3, 1], [4, 1]], [], [[5, 9]], [], []]>
```

``shift_frame = 1`` でも大丈夫です。

## サンプルが多次元の場合

例えばステレオ音声でLとRの値がペアで格納されているような場合が該当します。

```python
x = tf.ragged.constant([[[3, 2], [1, 7], [4, 1], [1, 8]], [], [[5, 2], [9, 8], [2, 1]], [[6, 8]], []])
```

実は今までと全く同じ方法で動きます。

```python
len_frame = 2
shift_frame = 1
s = x.row_starts()
e = s + x.row_lengths() + 1 - len_frame
r = tf.ragged.range(s, e, shift_frame)
ind = r[:, :, tf.newaxis] + tf.range(0, len_frame, dtype="int64")
ret = tf.gather(x.values, ind)
print(ret)
# <tf.RaggedTensor [[[[3, 2], [1, 7]], [[1, 7], [4, 1]], [[4, 1], [1, 8]]], [], [[[5, 2], [9, 8]], [[9, 8], [2, 1]]], [], []]>
```

次元数が増えすぎて、パッと見ただけでは合っているかどうか分からなくなってきましたが、大丈夫なはず…。