ディープラーニングフレームワークで可変長の入力を扱うときのTips

# はじめに

自然言語処理などで可変長の行列を使うときのパターンがいくつかある。毎回再実装しているような気がするので備忘録としてまとめる。

本記事では自分がよく使うChainerとTensorflowの実装をのせる。
(注: プロダクションのコードをコピペしたのではなく、このポストのために1から再実装したので、テストなどされていません。)

## Chainerに関する留意点

Chainerでは、可変長は下記で扱うように `Variable` + `length` で管理するのではなく、`Variable` のリストとして管理することを推奨しているように思う。具体的には[`L.NStepLSTM`](https://docs.chainer.org/en/stable/reference/generated/chainer.links.NStepLSTM.html)や[`F.pad_sequence`](https://docs.chainer.org/en/stable/reference/generated/chainer.functions.pad_sequence.html)などである。 

## 注記

下記記載のコードは、それぞれ下記のimportがされていることを前提にしている。

```python:Chainer
import chainer
import chainer.functions as F
import numpy as np
```

```python:Tensorflow
import tensorflow as tf
import numpy as np


sess = tf.InteractiveSession()
```

# 本文

## Padding

多くのディープラーニングフレームワークはGPUやCPUの並列計算を活用するために可変長の行列の計算を直接的にはサポートしていない。そこで、最大長の行列にあわせて、系列長外の部分を適当な値で埋めるパディングを行う。

なお、この部分はディープラーニングフレームワークではなくデータ作成の段階で自分でやってしまうことが多い。

```python
X = [np.array([1, 2]),
     np.array([11, 12, 13, 14]),
     np.array([21])]

# 単語IDを扱うことを想定してint32で
x = np.zeros([3, 4], dtype=np.int32)

for i, xi in enumerate(X):
    x[i, :len(xi)] = xi[:]

print x
# [[ 1  2  0  0]
#  [11 12 13 14]
#  [21  0  0  0]]
```

なお、Chainerの`L.EmbedId`を使うときは、0埋めではなく-1埋めにし、`L.EmbedId(..., ignore_label=-1)`を使うと良い。


## Masking

Sum poolingなどをするときに、[上記padding](#Padding)によって生まれた系列長外の部分を0でマスキングする(ただし、あまり[マスキングを過信しないこと](#Maskを過信しない)。こういった計算は`where`の計算で実現できる。

![mask.png](https://qiita-image-store.s3.amazonaws.com/0/152362/d11fed44-5372-2eed-9ad6-bcb47f19851d.png)
(Trueの場合は左辺値、Falseの場合は右辺値を採用することでマスキングとして機能する)

この処理をステップバイステップで書くと:

```python:Chainer
x = chainer.Variable(np.arange(1, 7).reshape(2, 3))
print x
# variable([[1 2 3]
#           [4 5 6]])

length = np.array([3, 2], dtype=np.int32)
print length
# [3 2]

xp = chainer.cuda.get_array_module(x.data)
mask = xp.tile(xp.arange(x.shape[-1]).reshape(1, -1), (x.shape[0], 1))
print mask
# [[0 1 2]
#  [0 1 2]]

mask = mask < length.reshape(-1, 1)
print mask
# [[ True  True  True]
#  [ True  True False]]

padding = xp.zeros(x.shape, dtype=x.dtype)
print padding
# [[0 0 0]
#  [0 0 0]]

z = F.where(mask, x, padding)
print z
# variable([[1 2 3]
#           [4 5 0]])

```

Tensorflowでは`sequence_mask`が便利。

```python:Tensorflow
x = tf.constant(np.arange(1, 7).reshape(2, 3).astype(np.float32))
length = tf.constant(np.array([3, 2], dtype=np.int32))

mask = tf.sequence_mask(length, tf.shape(x)[-1])
padding = tf.fill(tf.shape(x), 0.0)
z = tf.where(mask, x, padding)
print z.eval()
# [[ 1.  2.  3.]
#  [ 4.  5.  0.]]
```

Chainer版 (というよりnumpy版) `sequence_mask`

```python:Chainer
def sequence_mask(length, max_num=None):
    xp = chainer.cuda.get_array_module(length.data)
    if max_num is None:
        max_num = xp.max(length)
    # create permutation on (length.ndim + 1) dimension
    perms = xp.arange(max_num).reshape([1] * length.ndim + [-1])
    length = length.reshape([1] * (length.ndim - 1) + [-1] + [1])
    return perms < length
```

## Reshape

ディープラーニングでは`ミニバッチサイズ×特徴量`のランク2行列を扱うことが多いので、多くのフレームワークにはそのような行列を入力とした関数が多く用意されている。それらの関数の恩恵を享受するために、`ミニバッチ×シークエンス長さ×特徴量`の行列を`(ミニバッチサイズ*シークエンス長さ)×特徴量`のランク2行列に変換して処理をする。

![reshape_1.png](https://qiita-image-store.s3.amazonaws.com/0/152362/011ad625-4e7d-84bb-4823-06f3b9e92db2.png)

ただし、これだと比較的行列が疎の場合に余計な処理がもったいない。インデキシングをがんばることで処理を減らすことができる。
（試していないが、行列が疎でない場合はメモリの再確保で逆に時間がかかるかもしれないので注意）

![reshape_2.png](https://qiita-image-store.s3.amazonaws.com/0/152362/1443c289-1d41-843b-ae9f-f476c9ce9da6.png)


Tensorflowの場合は下記の処理によってこのような処理を実現できる。

![reshape_4.png](https://qiita-image-store.s3.amazonaws.com/0/152362/b82753eb-8d69-c999-a56b-84240b784ab9.png)


```python:Chainer
# WARNING: I have not checked it in case of rank != 3

x = chainer.Variable(np.arange(18).astype(np.float32).reshape(3, 3, 2))
length = np.array([2, 3, 1], dtype=np.int32)
w = chainer.Variable(np.ones([2, 3], dtype=np.float32))

# sequence_mask は前述
mask = sequence_mask(length, x.shape[length.ndim])
print mask
# [[ True  True False]
#  [ True  True  True]
#  [ True False False]]

x_reshaped = F.get_item(x, mask)
print x_reshaped
# [[  0.   1.]
#  [  2.   3.]
#  [  6.   7.]
#  [  8.   9.]
#  [ 10.  11.]
#  [ 12.  13.]]

y_reshaped = F.matmul(x_reshaped, w)
print y_reshaped
# [[  1.   1.   1.]
#  [  5.   5.   5.]
#  [ 13.  13.  13.]
#  [ 17.  17.  17.]
#  [ 21.  21.  21.]
#  [ 25.  25.  25.]]

pad_shape = [[0, 0] for _ in xrange(y_reshaped.ndim)]
pad_shape[length.ndim - 1][1] = 1
y_reshaped = F.pad(y_reshaped, pad_shape, 'constant', constant_values=0.)
print y_reshaped
# variable([[  1.,   1.,   1.],
#           [  5.,   5.,   5.],
#           [ 13.,  13.,  13.],
#           [ 17.,  17.,  17.],
#           [ 21.,  21.,  21.],
#           [ 25.,  25.,  25.],
#           [  0.,   0.,   0.]])


idx_size = np.prod(mask.shape)
inv_idx = np.ones([idx_size], dtype=np.int32) * -1
inv_idx[np.nonzero(mask.flat)[0]] = np.arange(x_reshaped.shape[0]).astype(np.int32)
print inv_idx
# [ 0  1 -1  2  3  4  5 -1 -1]

y = F.reshape(F.get_item(y_reshaped, inv_idx), list(x.shape[:length.ndim + 1]) + [-1])
print y
# [[[  1.   1.   1.]
#   [  5.   5.   5.]
#   [  0.   0.   0.]]
# 
#  [[ 13.  13.  13.]
#   [ 17.  17.  17.]
#   [ 21.  21.  21.]]
# 
#  [[ 25.  25.  25.]
#   [  0.   0.   0.]
#   [  0.   0.   0.]]]
```



Tensorflowの場合は下記の処理によってこのような処理を実現できる。

![reshape_3.png](https://qiita-image-store.s3.amazonaws.com/0/152362/5c9c2451-b23a-f917-2adc-9840ce0be50b.png)

```python:Tensorflow
# WARNING: I have not checked it in case of rank != 3
x = tf.constant(np.arange(18).astype(np.float32).reshape(3, 3, 2))
length = tf.constant(np.array([2, 3, 1], dtype=np.int32))
w = tf.constant(np.ones([2, 3], dtype=np.float32))

mask = tf.sequence_mask(length, tf.shape(x)[tf.rank(length)])
print mask.eval()
# [[ True  True False]
#  [ True  True  True]
#  [ True False False]]

x_reshaped = tf.boolean_mask(x, mask)
print x_reshaped.eval()
# [[  0.   1.]
#  [  2.   3.]
#  [  6.   7.]
#  [  8.   9.]
#  [ 10.  11.]
#  [ 12.  13.]]

y_reshaped = tf.matmul(x_reshaped, w)
print y_reshaped.eval()
# [[  1.   1.   1.]
#  [  5.   5.   5.]
#  [ 13.  13.  13.]
#  [ 17.  17.  17.]
#  [ 21.  21.  21.]
#  [ 25.  25.  25.]]

idx = tf.to_int32(tf.where(mask))
print idx.eval()
# [[0 0]
#  [0 1]
#  [1 0]
#  [1 1]
#  [1 2]
#  [2 0]]

shape = tf.concat([tf.shape(x)[:-1], tf.shape(y_reshaped)[-1:]], 0)
print shape.eval()
# [3 3 3]

y = tf.scatter_nd(idx, y_reshaped, shape)
print y.eval()
# [[[  1.   1.   1.]
#   [  5.   5.   5.]
#   [  0.   0.   0.]]
# 
#  [[ 13.  13.  13.]
#   [ 17.  17.  17.]
#   [ 21.  21.  21.]]
# 
#  [[ 25.  25.  25.]
#   [  0.   0.   0.]
#   [  0.   0.   0.]]]
```

## Softmaxの実装

与えられた行列の最外次元にsoftmaxを行うことを考える。このようなシチュエーションは[ListNetのPermutation probability distribution](http://qiita.com/koreyou/items/a69750696fd0b9d88608)や、アテンションの計算で生じる。

Softmaxの式
$$
y_i = \frac{exp(x_i)}{\sum_jexp({x_j})}
$$

```python
x = np.random.random([2, 3]).astype(np.float32)
# array([[ 0.44715771,  0.85983515,  0.08915455],
#        [ 0.02465274,  0.63411605,  0.01340247]], dtype=float32)

length = np.array([3, 2], dtype=np.int32)
```
下記図のように青い領域だけを使ってSoftmaxを計算したい。

![masked_softmax.png](https://qiita-image-store.s3.amazonaws.com/0/152362/a2339b0b-6cc6-eeaf-a5bb-e090c6bb0e34.png)

ちなみに事前/事後にマスクをかけるのはダメ。

```python:Chainer
# ダメな例1
x_ = np.copy(x)
x_[1, 2] = 0.
print F.softmax(x_)
# variable([[ 0.31153342,  0.47068265,  0.21778394],
#           [ 0.26211682,  0.48214924,  0.25573397]])

# ダメな例2
y = F.softmax(x)
y[1, 2] = 0.
print y
# variable([[ 0.31153342,  0.47068265,  0.21778394],
#           [ 0.26121548,  0.48049128,  0.0       ]])
# 2行目の合計が1.0になっていないことから明らかにだめ
```

理由は非常に簡単で、例1は $exp(0.258) \neq 0$のため。例2では、 `x[2,1]`が分母の計算に影響してしまっている。

Softmaxの計算では$exp(-inf) = 0$であることを利用してマスキングを行う。

```python:Chainer
def masked_softmax(x, length):
    """
    Softmax operation on the ourter-most dimenstion of x.

    Args:
         x (chainer.Variable): Values to be passed to softmax
         length (numpy.ndarray or cupy.ndarray):
             Number of items in the outer-most dimension of x
    """
    assert x.ndim - 1 == length.ndim
    xp = chainer.cuda.get_array_module(x.data)
    x_shape = x.shape
    x = F.reshape(x, (-1, x_shape[-1]))
    # mask: (B, T)
    mask = xp.tile(xp.arange(x.shape[-1]).reshape(1, -1), (x.shape[0], 1))
    mask = mask < length.reshape(-1, 1)
    padding = xp.ones(x.shape, dtype=x.dtype) * -np.inf
    z = F.where(mask, x, padding)
    return F.reshape(F.softmax(z), x_shape)


print masked_softmax(chainer.Variable(x), length)
# variable([[ 0.31153342,  0.47068265,  0.21778394],
#           [ 0.35218161,  0.64781839,  0.        ]])
```

```python:Tensorflow
def masked_softmax(x, length):
    """
    Softmax operation on the ourter-most dimenstion of x.

    Args:
         x (tf.Tensor): Values to be passed to softmax
         length (tf.Tensor): Number of items in the outer-most dimension of x
    """
    mask = tf.sequence_mask(length, tf.shape(x)[-1])
    padding = tf.fill(tf.shape(x), -np.inf)
    z = tf.where(mask, x, padding)
    return tf.nn.softmax(z, dim=-1)


print masked_softmax(
    tf.constant(x),
    tf.constant(length)).eval()
# [[ 0.31153342,  0.47068265,  0.21778394],
#  [ 0.35218161,  0.64781839,  0.        ]]
```

# Appendix:

## Maskを過信しない

ディープラーニングフレームワークでは0除算が生じると、`where`を使ったとしても勾配が`inf`になる[仕様](https://github.com/tensorflow/tensorflow/issues/2540)がある。なので、「不安定な計算をしてもマスクすればいいや」は通用しない。

次の式のようなネットワークがある。

```math
e = f_0(x) \\
w = f_1(e)
```

これをchainルールで表現すると下記になる。
$$
\frac{\partial w}{\partial x} = \frac{\partial w}{\partial e}\frac{\partial e}{\partial x}
$$

さて、これは自動微分では下記のように（ざっくりとだが）実現される。
```
x.grad = e.grad * g(f_0, e, x)
```
ここで、 `g(f_0, e, x)`は $f_0$とその入出力から表現される偏微分である。つまり、上段の式からどのような微分値`e.grad`が来たとしても、式$f_0$の偏微分値が`inf`や`nan`であれば`x.grad`もまた`inf`や`nan`になってしまう。試しにChainerとTensorflowでこれを試すと、

```python:Tensorflow
sess = tf.InteractiveSession()

x = tf.constant(0.0)

t = x
e = 1. / x
w = tf.where(True, t, e)

print w.eval()  # 0.0
print tf.gradients(w, x)[0].eval()  # nan
```

```python:Chainer
x = chainer.Variable(np.array([0.0], dtype=np.float32))
t = x
e = 1. / x
w = chainer.functions.where(np.array([True]), t, e)

w.grad = np.array([1.0], np.float32)
w.backward(retain_grad=True)

print w  # 0.
print x.grad  # nan
```