Pseudo Theory of Everything

データサイエンス初心者物理学徒の奮闘記

PyTorch Tutorials : 60 MINUTE BLITZ/自動微分モジュール

AUTOGRAD: AUTOMATIC DIFFERENTIATION

PyTorchをはじめとするニューラルネットワークのライブラリ(のおそらく全て)は自動微分の機能を持ち合わせています。PyTorchも例外ではなく、自動微分モジュール autograd が存在します。今回のチュートリアルはその autograd がどのように使われ、どのようなことができるのかに触れ、次記事でニューラルネットワークの学習に進みます。

1.1 Tensor

公式チュートリアルで最初に例として扱う式展開を見てからライブラリの使い方に入りましょう。

\(2 \times 2\) 行列 \({\boldsymbol x}\) を用意します。以降、行列の足を \(i, j\) で表現します。ここで以下のように各成分に \(2\) を足し、行列 \(y_{i, j}\) を作ります。

\begin{align} y_{i, j} = x_{i, j} + 2 \end{align}

さらに、行列 \({\boldsymbol y}\) を用いて行列 \({\boldsymbol z}\) を作り、その全成分の平均値をとったものを \({\rm out}\) とします。

\begin{align} z_{i, j} &= 3 y_{i, j}^2 = 3 ( x_{i, j} + 2 )^2\\ {\rm out} &= \frac{1}{4} \sum_{i,j = 1}^2 z_{i, j} = \frac{3}{4} \sum_{i,j = 1}^2 y_{i, j}^2 = \frac{3}{4} \sum_{i,j = 1}^2 ( x_{i, j} + 2 )^2 \end{align}

チュートリアルでは全ての成分が \(1\) の行列 \({\boldsymbol x}\) を用意し、出力 \({\rm out}\) は \(27\) となります。この章では、この \({\rm out}\) の \({\boldsymbol x}\) 微分をライブラリに任せます。実際には、PyTorchといえど解析式を与えるわけではなく、 \(x_{i, j} = 1\) での微分の数値解を与えてくれます。まず解析的には

\begin{align} \frac{d {\rm out}}{d z^{i, j}} &= \frac{1}{4}\\ \frac{\partial {\rm out}}{\partial y^{i, j}} &= \frac{3}{2} y_{i, j}\\ \frac{\partial {\rm out}}{\partial x^{i, j}} &= \frac{3}{2} ( x_{i, j} + 2 ) \end{align}

です。数値的には \(x_{i, j} = 1\) を代入し、

\begin{align} \left. \frac{\partial {\rm out}}{\partial x^{i, j}} \right|_{x_{i, j} = 1} = \frac{9}{2} \end{align}

チュートリアルがこの例を挙げていたのでこれをやりますが、この演算ベクトルやテンソルで演算する意味がないです。一通りチュートリアルを終えたら+αをやってみましょう。

早速PyTorchを使って計算してみます。まずはライブラリの import をします。

import torch

前回の記事とは異なりテンソルを定義する際、引数に requires_grad=True を渡し、行列 \({\boldsymbol x}\) を作成します(テンソルのデフォルト値では requires_grad=False になっています)1

x = torch.ones(2, 2, requires_grad=True)
print(x)

出力 :

tensor([[1., 1.],
        [1., 1.]], requires_grad=True)

上記のような定義をすることで後に自動微分操作が可能になります。続いて、行列 \(y\) を定義します。

y = x + 2
print(y)

出力 :

tensor([[3., 3.],
        [3., 3.]], grad_fn=<AddBackward0>)

出力には grad_fn という属性がつきました。 y.grad_fn を参照してみると

print(y.grad_fn)

出力 :

<AddBackward0 object at 0x125b8d8d0>

つまり、 yx偏微分するときの計算に使うオブジェクトが格納されています。今回の場合、2の加算ですので AddBackward0 が格納され、内部では偏微分を行う際に無視されるのでしょう。
では、 \({\boldsymbol y}\) を二乗しかつ3倍し、 \({\boldsymbol z}\) を定義します。さらに \({\boldsymbol z}\) の全成分の平均を取り \({\rm out}\) も同時に定義し、出力しましょう。

z = y * y * 3
print(z)

out = z.mean()
print(out)

出力 :

tensor([[27., 27.],
        [27., 27.]], grad_fn=<MulBackward0>)
tensor(27., grad_fn=<MeanBackward0>)

事前に計算した通り、 \({\rm out} = 27\) が算出されました。

1.2 微分

これまで構成した \({\rm out}\) を用いて、本章の本題である自動微分を実行していきます。
自動微分操作には .backward() を宣言することで自動微分結果を参照できるようになります。ただし、PyTorchの自動微分モジュールとしてスカラーの1次元テンソルの値のみで .backward() が実行できます。もし多次元のテンソルの状態で .backward() を宣言するとエラーを返します。後に述べますが計算コスト上PyTorchがそのようにハードコードしているのだと思います。

out.backward()

事前に計算した \({\rm out}\) の \(x\) 微分結果は以下のようにして参照できます。

print(x.grad)

出力 :

tensor([[4.5000, 4.5000],
        [4.5000, 4.5000]])

事前に計算したように \(9/2\) の値が計算されました。

ここで、正しくスカラー値 \({\rm out}\) の微分結果が得られたものの、なぜテンソルの自動微分ができないのかに触れます。結論は先ほども挙げたように計算コストのためと考えています。
今仮に微分されるスカラー値を \(l = g({\boldsymbol y})\) とし、 \({\boldsymbol y} = y_i, (i = 1,2, \ldots, m)\) の関数とします。また、ベクトル \({\boldsymbol y}\) も \({\boldsymbol y} = f({\boldsymbol x})\) と \({\boldsymbol x} = x_i, (i = 1,2, \ldots, n)\) の関数であるとし、 \(l\) の \({\boldsymbol x}\) 微分である \(\left( \frac{\partial l}{\partial x_1} , \ldots , \frac{\partial l}{\partial x_n} \right)^T\) を考えます。
本来一般的には

\begin{align} \begin{pmatrix} \frac{\partial l}{\partial x_1}\\ \vdots\\ \frac{\partial l}{\partial x_n} \end{pmatrix} & = \begin{pmatrix} \frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_m}{\partial x_1}\\ \vdots & \ddots & \vdots\\ \frac{\partial y_1}{\partial x_n} & \cdots & \frac{\partial y_m}{\partial x_n} \end{pmatrix} \begin{pmatrix} \frac{\partial l}{\partial y_1}\\ \vdots\\ \frac{\partial l}{\partial y_m} \end{pmatrix} = J^T \cdot v\\ J &\equiv \begin{pmatrix} \frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_1}{\partial x_n}\\ \vdots & \ddots & \vdots\\ \frac{\partial y_m}{\partial x_1} & \cdots & \frac{\partial y_m}{\partial x_n} \end{pmatrix}\\ v &\equiv \begin{pmatrix} \frac{\partial l}{\partial y_1}\\ \vdots\\ \frac{\partial l}{\partial y_m} \end{pmatrix} \end{align}

のように \({\boldsymbol y}\) の \({\boldsymbol x}\) に対する微分であるヤコビアンを噛ませることで算出されます。これまで実際にPyTorchの自動微分の例のように \(x \rightarrow y \rightarrow z \rightarrow {\rm out}\) とすると \({\rm out}\) の \(x\) 微分を知るためにヤコビアンを2つ構成し、そのテンソル積を行わないとならないため、メモリ・計算量共にコストが膨大になっていきます。
実際のニューラルネットバックプロパゲーションを思い浮かべれば十分なのですが、常に「スカラー(損失関数やユニットへ入力される値)をテンソル(基本的に重み行列)で微分する」形になっているためにヤコビアンを構成することなくPyTorchはハードコードによって自動微分が構成されています。
そのために、常にスカラー値の .backward() しか許されていません。
(一方、ヤコビアンそのものを必要とする場合も少なからずあるため、その場合ちょっと手間やコストがかかってしまうため、ヤコビアンを作る関数も作っておけば良いのにとも思います。)

1.3 その他

チュートリアル内で扱われている補足的な内容を本節では扱っていきます。

1.3.1 .requires_grad_( ... ) について

テンソルを構成する際 requires_grad=True を渡してあげることで自動微分の対象とすることができることを紹介しました。

この属性はテンソルを構成した後でも以下のようにして付与することはできます。例えば以下のような形で requires_grad=True をし忘れると当然

a = torch.randn(2, 2)
a = ((a * 3) / (a - 1))
print(a.requires_grad)

出力 :

False

となってしまいます。 .requires_grad_(True) で後からこの属性を変更できます。

a.requires_grad_(True)
print(a.requires_grad)

出力 :

True

これによって、これまでみてきたように

b = (a * a).sum() 
print(b.grad_fn) 
b.backward() 
print(a.grad)                                             

出力 :

<SumBackward0 object at 0x127525d50>
tensor([[74.9851,  0.4784],
        [ 8.2700,  7.5751]])

と自動微分が実行できます。

1.3.2 vector-Jacobian product

ヤコビアンとベクトルの積をvector-Jacobian productと呼びます(正確な定義は知りません)。微分の節で挙げたように、自動微分.backward()スカラー値にしか渡すことができません。

x = torch.randn(3, requires_grad=True) 
print(x) 
 
y = x * 2 
count = 0 
while y.data.norm() < 1000: 
    y = y * 2 
    count += 1 

print(y, count)

出力 :

tensor([ 0.9203, -0.0432,  0.9675], requires_grad=True)
tensor([942.3748, -44.2408, 990.6782], grad_fn=<MulBackward0>) 9

上記の例では \({\boldsymbol x}\) を \({\boldsymbol \mu} = (0, 0, 0)^T\),  {\boldsymbol \sigma} = {\rm diag} (1, 1, 1)正規分布に従うように3次元のベクトルを生成し、その \(L_2\) ノルムが1,000を超えるまで2倍する操作を続けています。
このとき yスカラー値ではないため、 x微分を得ることができません。

y.backward()

出力 :

--------------------------------------------------------------------
RuntimeError                       Traceback (most recent call last)
<ipython-input-149-ab75bb780f4c> in <module>
----> 1 y.backward()

~/.pyenv/versions/3.7.5/lib/python3.7/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    164                 products. Defaults to ``False``.
    165         """
--> 166         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    167 
    168     def register_hook(self, hook):

~/.pyenv/versions/3.7.5/lib/python3.7/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     91         grad_tensors = list(grad_tensors)
     92 
---> 93     grad_tensors = _make_grads(tensors, grad_tensors)
     94     if retain_graph is None:
     95         retain_graph = create_graph

~/.pyenv/versions/3.7.5/lib/python3.7/site-packages/torch/autograd/__init__.py in _make_grads(outputs, grads)
     32             if out.requires_grad:
     33                 if out.numel() != 1:
---> 34                     raise RuntimeError("grad can be implicitly created only for scalar outputs")
     35                 new_grads.append(torch.ones_like(out))
     36             else:

RuntimeError: grad can be implicitly created only for scalar outputs

この yx 微分を求める場合、適当なベクトル値を用意し、 .backward() に渡すことで x 微分が得られます。

v = torch.tensor([0.1, 1.0, 0.0001], dtype=torch.float)
y.backward(v)

print(x.grad)

出力 :

tensor([1.0240e+02, 1.0240e+03, 1.0240e-01])

上記の例では定数 torch.tensor([0.1, 1.0, 0.0001])y内積を取ることでスカラーにし(コード上では .backward(v) でそれが実現できる)、微分係数を求めることをしました。
当然ですが、これまで見てきたように sum を使うことでも y における x の係数 \(1,024(=2^{10})\) が計算されます。

x = torch.randn(3, requires_grad=True) 
 
y = x * 2 
while y.data.norm() < 1000: 
    y = y * 2 
 
y = y.sum() 
y.backward() 
print(x.grad)

出力 :

tensor([1024., 1024., 1024.])

1.3.3 自動微分を止める

これまでの例で見た通り、 requires_grad=True とすると、それ以後の演算を全ての自動微分を算出する操作となっていました。ここで紹介するのは自動微分を明示的にさせない操作です(正直メモリの節約とかには使えるんでしょうが、簡単な例が思いつきません)。
まずは with torch.no_grad() を利用することで微分の追跡から省くことができます。

print(x.requires_grad)
print((x ** 2).requires_grad)

with torch.no_grad():
    print((x ** 2).requires_grad)

出力 :

True
True
False

上記のように with torch.no_grad() 句内では演算を行っても requires_grad 属性は False になります。
同様に .detach() を使うことで成分は同じであるものの、 requires_grad=Falseテンソルをコピーできます。

print(x.requires_grad)
y = x.detach()
print(y.requires_grad)
print(x.eq(y).all())

出力 :

True
False
tensor(True)

最後の eq()x, y の成分が同じか否かを判断する関数になります。

1.3.4 中間変数の勾配

以上で公式のチュートリアルは終わりなのですが、加えて中間変数の勾配を見てみましょう。

最初の自動微分の例では \(x \rightarrow y \rightarrow z \rightarrow {\rm out}\) と変数変換していきましたが、実は中間変数にあたる \(y, z\) の微分を見ることはできません。試してみましょう。

x = torch.ones(2, 2, requires_grad=True)
y = x + 2
z = y * y * 3
out = z.mean()

out.backward()
print(x.grad)
print(y.grad)
print(z.grad)

出力 :

tensor([[4.5000, 4.5000],
        [4.5000, 4.5000]])
None
None

となってしまいます。これの解決の方法は .retain_grad() です。

x = torch.ones(2, 2, requires_grad=True)
y = x + 2
z = y * y * 3
out = z.mean()

y.retain_grad()
z.retain_grad()
out.backward()
print(x.grad)
print(y.grad)
print(z.grad)

出力 :

tensor([[4.5000, 4.5000],
        [4.5000, 4.5000]])
tensor([[4.5000, 4.5000],
        [4.5000, 4.5000]])
tensor([[0.2500, 0.2500],
        [0.2500, 0.2500]])

1.3.5 高階微分

高階微分もこれまでの知識では実行できません。

\begin{align} {\boldsymbol x} &= (1, 2)^T\\ y &= \sum_i x_i^3 \end{align}

を準備し、 \({\boldsymbol x}\) で1, 2, 3階微分をします。あらかじめ解くと

\begin{align} \left. \frac{d y}{d x^i} \right|_{{\boldsymbol x} = (1, 2)^T} &= 3 x_i^2|_{{\boldsymbol x} = (1, 2)^T} = (3, 12)^T\\ \left. \frac{d^2 y}{d x^{i2}} \right|_{{\boldsymbol x} = (1, 2)^T} &= 6 x_i|_{{\boldsymbol x} = (1, 2)^T} = (6, 12)^T\\ \left. \frac{d^3 y}{d x^{i3}} \right|_{{\boldsymbol x} = (1, 2)^T} &= 6 = (6, 6)^T \end{align}
x = torch.tensor([1.0, 2.0], requires_grad=True) 
y = x**3 
y = y.sum() 
y.backward() 
print(x.grad) 
print(x.grad.grad) # これが意図としたことなのかわかりませんが

出力 :

tensor([ 3., 12.])
None

例えば上記のように2階微分を知りたい場合、 torch.autograd.grad() を使います。

x = torch.tensor([1.0, 2.0], requires_grad=True) 
y = x**3 
y = y.sum() 
grad = torch.autograd.grad(outputs=y, inputs=x, create_graph=True)
print(grad)  # タプルを返す
grad[0].sum().backward() 
print(x.grad)

出力 :

(tensor([ 3., 12.], grad_fn=<MulBackward0>),)
tensor([ 6., 12.])

上記の通り、2階微分が求まりました。見てわかるように torch.autograd.grad 関数を使うことで1階微分の勾配が取得できます。 create_graph=True をすることで計算グラフが構築され、高階微分が可能になります。

同じ要領で3階微分は以下の通りです。

x = torch.tensor([1.0, 2.0], requires_grad=True)  
y = x**3  
y = y.sum()  
grad1 = torch.autograd.grad(outputs=y, inputs=x, create_graph=True) 
grad1 = grad1[0].sum() 
grad2 = torch.autograd.grad(outputs=grad1, inputs=x, create_graph=True) 
grad2[0].sum().backward()  
print(x.grad)

出力 :

tensor([6., 6.])

ここで torch.autograd.grad がタプルを返す件ですが、この関数は通常複数の変数の微分を扱うため、その都度タプルの成分が増えるようです。試しに

\begin{align} f(x, y) = (x + 2y)^3 \end{align}

を作り、 \(x, y\) で微分します。これまでは \(x, y\) もベクトルとして扱っていたのですが、本節では処理が多くなり面倒なので1次元のテンソル微分をします。

\begin{align} \left. \frac{\partial^2 f(x, y)}{\partial x \partial y} \right|_{x=2, y=3} = \left. \frac{\partial^2 f(x, y)}{\partial y \partial x} \right|_{x=2, y=3} = 12 (x + 2y) |_{x=2, y=3} = 96 \end{align}
x = torch.tensor([2.], requires_grad=True) 
y = torch.tensor([3.], requires_grad=True) 
f = (x + 2 * y) ** 3 
grad = torch.autograd.grad(outputs=f, inputs=(x, y), create_graph=True) 
print(grad)

出力 :

(tensor([192.], grad_fn=<MulBackward0>), tensor([384.], grad_fn=<MulBackward0>))

それぞれ第一成分が \(x\) の1階微分, 第二成分が \(y\) の1階微分を指します。つまり、

grad[0].backward() 
print(y.grad)

出力 :

tensor([96.])

を返し、

grad[1].backward() 
print(x.grad) 

出力 :

tensor([96.])

になります。ここで上記の .backward() を連続では使えないことに注意してください。片方を実行したら、再度 grad を再定義してください。

Footnotes:

1

試したところ浮動小数点型の torch.half, torch.float, torch.double しか requires_grad=True にすることができません。