float()をtensorに使うと自動微分ができなくなった

introduction

今回はPyTorchの内容。実験で深層学習を使いたいので、ちょうどいい機会なのでPyTorchを学んでいる。学んでいる本は以下の本。この本を学んでいる時に詰まったことを紹介する。
これ、めっちゃおすすめなのでぜひ読んでほしいで。
https://www.amazon.co.jp/最短コースでわかる-PyTorch-＆深層学習プログラミング-赤石-雅典-ebook/dp/B09G622WB6

疑問点

tensorを作成する際に、Numpy変数から作成しようとすると、データ型がfloat64になってしまう。そうすると何が問題かというと、機械学習を最終的にしたいのに、"nn.linear"という機械学習に用いるライブラリを利用できなくなってしまうのだ。このライブラリは基本float32のデータ型に対して作用するので、型変換をする必要がある。それを行えるのが、以下の2つの手法。

r1_np = np.array([1.0, 2, 3, 4, 5])

#手法1
r1 = torch.tensor(r1_np).float()

#手法2
r1_2 = torch.tensor(r1_np,dtype=torch.float32)

まあつまり、メソッドを使うのか、はたまたdtypeの引数として指定するのかの違い。一見どっちでもいいやないかって感じなんだけど、ここで困るのが自動微分。手法2は本にも書いてあるとおり上手くいくんだけど、手法1を用いて、以下のコードのように逆伝播するためのコードを作成して、自動微分を行うと、、、

#必要なライブラリ
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import japanize_matplotlib
from IPython.display import display
from torchviz import make_dot
import torch

# xをnumpy配列で定義
x_np = np.arange(-2, 2.1, 0.25)

# (1) 勾配計算用変数の定義
x = torch.tensor(x_np, requires_grad=True, 
    dtype=torch.float32)
# (2) テンソル変数をもちいた計算
y = 2 * x**2 + 2

z = y.sum()

# (3)計算グラフの可視化
g= make_dot(z, params={'x': x})
display(g)

# (4) 勾配計算
z.backward()

# (5) 勾配値の取得
print(x.grad)

以下のようにエラーが出てしまう。しかも出るのはprint(x.grad)の部分。

None
<ipython-input-42-dec9236f5aaf>:3: UserWarning: 
The .grad attribute of a Tensor that is not a leaf Tensor is being accessed.
Its .grad attribute won't be populated during autograd.backward().
If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor.
If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead.
See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at aten/src/ATen/core/TensorBody.h:486.)
  print(x.grad)

さて、なぜこのエラーが出るのだろうか。これが今回の疑問。

chatGPTにぶちこんでみる

一旦chatGPTにぶちこめばなにかわかるかもしれないので、一旦ぶちこんでみると、以下のように返ってくる。

警告メッセージには、Tensorオブジェクトの.grad属性にアクセスした際に発生する可能性がある警告に関する情報が含まれています。.grad属性は、Tensorオブジェクトの微分値を保持するために使用されます。ただし、この属性は、リーフテンソル（自分自身に対する微分を計算するために必要なテンソル）にのみ適用されるため、非リーフテンソルに対して.grad属性にアクセスした場合、警告が発生します。
この警告を回避するには、次のいずれかの方法を使用する必要があります。
a..retain_grad()を使用して、非リーフテンソルに対して.grad属性を保持するように指定する。
b.非リーフテンソルに誤ってアクセスしないようにする。
c.非リーフテンソルにアクセスする必要がある場合は、そのリーフテンソルを取得し、それにアクセスする。
具体的な警告メッセージには、どのTensorオブジェクトが問題の原因であるかについての情報も含まれています。これを使用して、問題のTensorオブジェクトを特定し、適切な対処を行うことができます。

つまりこれを要約すると、
非リーフテンソルに対して.grad属性にアクセスしちゃってるから上手く勾配とれねーよ！
ってことなんだろうな。

んで、まぁ「非リーフテンソル」って言葉と「.grad属性」って言葉が何なのかよくわからん。
そこでまずは、.grad属性から説明する。簡単に言えば"requires_grad=True"の部分で、このテンソルは勾配計算するために必要なテンソルですよ！って明示的にしめすもの。なので、ここがtrueになっていれば基本的にgrad属性はtrueである。
次にleaf tensorについて。lead tensorというのは、末端変数って言う意味で、leafは木の葉って意味からも納得行くのではないかなと思う。じゃあどこの末端か？って言われると、計算グラフの末端である。この本をやっている人は計算グラフについてもう理解しているだろうけど、指定ない人にして言うと、Pytorchでは以下のように計算過程がグラフになって保存されている。その一番上流(末端)にあるものがleaf tensor、他の言い方をするとleaf nodeやleaf variableなんてものだ。以下の図で言うならばWとBがleaf tensorになる。

leaf tensorであるならば、そのテンソルはちゃんと求めたい損失関数までの伝播の道のりが出来ているわけだ。そしてそれはgrad属性がtrueということにほかならない。僕らが"requires_grad=True"って打った時点で、その変数はある計算グラフにおけるleaf tensorになってたってことだ。
もしもそれがleaf tensorであるかどうかの確認をしたければ、以下のコードを打つといい。

print(x.is_leaf)

さて、つまりさっきの文が言いたいのは、xがfloat()によってleaf tensorじゃなくなっている、いいかえるなら"requires_grad=false"になってるってこと。
え？
ってならん？だっておれらちゃんと"requires_grad=true"って打ったぜ、、、？まぁ試しにis_leafで確認してみるか、、、

x = torch.tensor(x_np, requires_grad=True).float()
print(x.is_leaf)
>>> False

Falseなっとるやないか！！！なんでや！！！

ってことで、色々調べてみた。

解説

なんとなく意味はわかった。まずは本の引用から。

xのデータ型をfloat32に変換するため、いつものfloat関数を呼び出すのではなく、ここではdtyprパラメータを指定しています。float関数を呼び出すと、計算グラフとしてはcopy関数の呼び出しが入り、xが勾配値計算が可能なリーフノードではなくなるためです。

これによると、float関数を呼び出すとcopy関数の役割を果たすことになっているらしい。なんでこんなめんどくさい機能をしているかというと、python特有のコンテナデータ型という性質に基づいているのかもしれない。コンテナデータ型であることから、numpyのデータ型からtensorに変換した際、numpyの方を変換した際にtensorの値も変わってしまうのだ。これは本で扱っているが、もしも本を使っていないなら以下のコードを見て理解してくれ。

import torch

# x1: shape=[5] となるすべて値が1テンソル
x1 = torch.ones(5)

# x2 x1から生成したNumPy
x2 = x1.data.numpy()

# x1の値を変更
x1[1] = -1

# 連動してx2の値も変わる
print(x1)
print(x2)

つまり、データ型変換すると、numpyのデータに対してまで影響を与えないといけないので、まずはcopyをとって対象を絞る事によってそれを考慮しないことにしている。

実際、xを作る時点でxを確認してみると、grad_fnにcopyのレイヤーが追加されていることがわかる。

x = torch.tensor(x_np, requires_grad=True).float()
print(x)
>>>tensor([-2.0000, -1.7500, -1.5000, -1.2500, -1.0000, -0.7500, -0.5000, -0.2500,
         0.0000,  0.2500,  0.5000,  0.7500,  1.0000,  1.2500,  1.5000,  1.7500,
         2.0000], grad_fn=<ToCopyBackward0>)

このレイヤーがあることによって、xが勾配計算ができないleaf tensorとなっているようだ。

じゃあ手法2で全部やれば？

そんなん言うなら手法2で全部やればいいやんけ！って思うだろうけど、単純にタイピング量が多いのでできるだけ使いたくはない。そんで、そもそも初期値の設定に対してnumpyからの変換とか全然使わない。実はこのcopyの呼び出しが起こるのは、numpyからの変換の時に特異的に起こるものなのだ。
本を先に進んだ人ならわかると思うけど、普通にleaf tensorに対して普通にfloat()を使っている。んで、それらに対しては、grad_fnは表示されない。

W = torch.tensor(1.0, requires_grad=True).float()
B = torch.tensor(1.0, requires_grad=True).float()
print(W,B)
>>>tensor(1., requires_grad=True) tensor(1., requires_grad=True)

そう、tensor作成の対象がpythonのリスト型や、スカラーなどに関してはcopyレイヤーは作られることはないし、そもそもこの使い方がメインになってくる。だからこそ、基本的にはfloat()に対して用いてよいということだ。
勾配を単純な一次の微分に用いたいなら、手法2でやらないといけないというわけだな(本を参照、別に重要な内容ではない。

計算グラフはどうなるのか

一応leaf tensorではなくても計算グラフはかけているようだ。

でもこれをみてみるとleaf tensor(x)への参照がなくなっている。ToCopyBackward0が追加されることで、逆伝播で計算できなくなってしまい、この計算グラフ自体に意味がなくなってしまう。計算グラフを枝と捉えるのなら、途中が壊死してしまってleafが枯れてしまうことに相当するのだろう。