【Pytorch】損失関数の勾配を計算する

# はじめに

NNの学習時で計算する、勾配を計算する方法を調べたのでメモしておきます。

```math
\theta \leftarrow \theta + \eta \nabla_\theta L(\theta) 
```

通常は、これらはoptimizer が良しなに更新してくれるので、自分では learning rate をいじったりはしません。なので、何次元のパラメータかなども意識しません。が、それを見てみよう、というのが今回の趣旨です。

# Modelのgradient vector を計算

ネットワークモデルについては、pytorchのtuotrial に沿って行います。

https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html


### モデルの定義とパラメータの参照

```python: net.py
import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 1 input image channel, 6 output channels, 5x5 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5*5 from image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # If the size is a square, you can specify with a single number
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1) # flatten all dimensions except the batch dimension
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
```

このように、モデルを定義するようです。Linear や Conv2d では次元数を設定してあげると、中に重みパラメータをもったネットワークが作られます。

モデルのパラメータを列挙するには、nn.parameters() や nn.named_parameters() を使えます。
nn.parameters の説明は下記のリンクに。Tensorのsubclass で、nn.Module に追加すると自動的に参照されるようになるようです。data と requires_grad を持っています。データ数は```numel()``` で取得できます。

https://pytorch.org/docs/stable/generated/torch.nn.parameter.Parameter.html

```python
from mynet import Net

net = Net()
for p in net.parameters():
    print(f'size={p.size()} {p.requires_grad}')
```

各モジュールの名前とともに取得することもできます。named_paremters() で出力される順番は、定義した順番のようでした。（自分が動かした範囲では）

```python
for n, p in net.named_parameters():
    print(f'name={n} size={list(p.size())} ({p.numel()}) {p.requires_grad}')
```

出力：

> name=conv1.weight size=[6, 1, 5, 5] (150) True
name=conv1.bias size=[6] (6) True
name=fc1.weight size=[120, 400] (48000) True
name=fc1.bias size=[120] (120) True
name=fc2.weight size=[84, 120] (10080) True
name=fc2.bias size=[84] (84) True
name=fc3.weight size=[10, 84] (840) True
name=fc3.bias size=[10] (10) True
name=conv2.weight size=[16, 6, 5, 5] (2400) True
name=conv2.bias size=[16] (16) True


### モデルによる値の計算

入力を入れると出力を得られます。入力はMNIST を想定して 32 x 32 のサイズを推定しています。
ダミーのデータを入力すると、１０次元の値が出力されます。ここのところ、良く分かっていないのですが、
4D Tensor で nSamples x nChannels x Height x Widthを入力するようです。これはこの入力が最初に入力されるモデルConv2dの入力の仕様のようです。

```python
x0 = torch.randn(1, 1, 32, 32)
y = net(x0)
print(y) # # torch.Size([10])
```

出力はこんな感じに。

> tensor([[ 0.0110,  0.0724, -0.0437, -0.0916, -0.0550,  0.0797,  0.1916, -0.0253,
         -0.0806,  0.1785]], grad_fn=<AddmmBackward0>)

### 損失関数

この予測値と実際の正解値との差を、定義した損失関数で計算します。
ここでは既にあるMean Saure Error を使います。

```
target = torch.randn(10)  # torch.Size([10])
target = target.view(1, -1)  # torch.Size([1, 10])

criterion = nn.MSELoss()
loss = criterion(output, target)
print(loss)
```

ここでは view という関数を用いて、Tensor のサイズを ```torch.Size([10])```から```# torch.Size([1, 10])``` に変更しています。これは、loss 関数への入力として```y=net(x)``` の出力と同じ形にするためです。

出力:
> tensor(0.9695, grad_fn=<MseLossBackward0>

## 勾配の計算

### 直接gradient の値を取り出してベクトルのする方法

#### 損失と誤差逆伝搬法

ニューラルネットワークの出力の計算では、ネットワークでforward 計算を行い、隠れ層の値を計算します。隠れ層の値と、先程計算した損失関数の値から、微分値を計算できます。これが、back propagation で計算されます。

loss 関数から、逆にたどることを確認できます。もともとのモデルは

> input -> conv2d -> relu -> maxpool2d -> conv2d -> relu -> maxpool2d -> flatten -> linear -> relu -> linear -> relu -> linear -> MSELoss -> loss

となっていますが、

```
print(loss.grad_fn)  # MSELoss
print(loss.grad_fn.next_functions[0][0])  # Linear
print(loss.grad_fn.next_functions[0][0].next_functions[0][0])  # ReLU
```

で

> \<MseLossBackward0 object at 0x7f7c50a149d0>
\<AddmmBackward0 object at 0x7f7c50a149a0>
\<AccumulateGrad object at 0x7f7c50a149a0>

となっています。この逆伝搬で勾配を計算します。勾配計算をする前に、勾配の値を保持している変数をクリアします。
Back propagation の計算は、backward()で実行されます。

```
print('clear gradient variables.')
net.zero_grad()     # zeroes the gradient buffers of all parameters
print(f'before backward: conv1.bias.grad={net.conv1.bias.grad}')
loss.backward()
print(f'after backward: conv1.bias.grad={net.conv1.bias.grad}')
```

出力: 

> clear gradient variables.
before backward: conv1.bias.grad=None
after backward: conv1.bias.grad=tensor([ 0.0155,  0.0101, -0.0022, -0.0421,  0.0010, -0.0052])

この後、パラメータ更新を直接計算
```math
\vec{\theta} \leftarrow \vec{\theta} + \eta \cdot \nabla_{\theta} L(\theta)
```

を実行したかったら、下記でできます。
```python
learning_rate = 0.01
for f in net.parameters():
    f.data.sub_(f.grad.data * learning_rate)
```

#### 勾配の保持

全ての要素の勾配を結合させます。

```
grad_vec = []
for n, f in net.named_parameters():
    grad_vec.append(p.grad.data.view(-1))
grad_vec = torch.cat(grad_vec)
print(grad_vec)
print(grad_vec.shape)
```

出力: 

> tensor([ 0.0004, -0.0112,  0.0313,  ...,  0.0365, -0.1341, -0.2248])
torch.Size([61706])

パラメータ数は多いんですね。

# まとめ

とりあえず、ニューラルネットワークで訓練データに対する勾配の値を直接取り出すことができるようになった。これをどう使おうかな。
(2022/11/12)

### 追記
- tensorについて data.view(-1) として多次元配列をなくすところを修正 (2022/12/04)