【PyTorch】同じ層内のパラメータに別々の学習率を設定する方法

## 問題
1つのPyTorchモデルにて，パラメータ別に異なる学習率で学習したい層があるとします．具体的には
* 更新しない層
* 更新する層
    * 学習率を高めにして更新するパラメータ
    * 学習率をちょっと低めにして更新するパラメータ
    * ...などなど

とわかれている状況です．

この問題で難しいのは，**同じ層内に異なる学習率で学習したいパラメータがある**ことです．下の記事では，層別に異なる学習率で学習する方法について紹介されています．

https://noconocolib.hatenablog.com/entry/2019/01/16/000000

しかし，この方法では，同じ層内のパラメータに対して，別々の学習率を指定することができません．

同じ層内のパラメータに対して，別々の学習率を指定するにはどうすればいいでしょうか？

## 問題例
この問題例としては，LoRA+の実装があります．

https://qiita.com/kernelian/items/9897d6d45203c9139939

詳しくは上などの記事を参照されたいのですが，簡単に言うと
* **更新しない層**：元のモデル
* **更新する層**:LoRA
    * **学習率を高めにして更新するパラメータ**：LoRAのB
    * **学習率を低めにして更新するパラメータ**：LoRAのA

というような学習を行います．LoRA層のそれぞれのパラメータA，Bに対して，別々の学習率を指定するにはどうすればいいでしょうか？

## 実装方法を例を交えて説明する
### 状況設定
PyTorchで作業するため，まずは必要なパッケージをインポートします．
<details><summary>コード：必要なパッケージをインポート</summary>

```python
import torch
from torch import nn, optim, Tensor, gradient, tensor
```
</details>

次に示すような構造のモデル，`MyModel`を考えます．
```plaintext
MyModel(
  (t0): Target()
  (d0): Dummy()
  (t1): Target()
  (m0): Mother(
    (t): Target()
    (d): Dummy()
  )
  (t2): Target()
  (d1): Dummy()
)
```
`Target`層内にはパラメータが2つあります．それぞれ`a`，`b`とします．`Dummy`内にもまた別名のパラメータがあります．`Mother`は内部に`Target`と`Dummy`を入れ子として持つ層です．
<details><summary>コード：モデルの定義</summary>

```python
class Target(nn.Module):
    def __init__(self, a0:float, b0:float, *args, **kwargs ) -> None:
        super().__init__(*args, **kwargs)
        self.a = nn.Parameter(Tensor([a0]))
        self.b = nn.Parameter(Tensor([b0]))
        
        self.a.requires_grad = True
        self.b.requires_grad = True
    def forward(self, x:Tensor)->Tensor:
        return self.a*x + self.b

class Mother(nn.Module):
    def __init__(self, *args, **kwargs ) -> None:
        super().__init__(*args, **kwargs)
        self.t = Target(10., 20.)
        self.d = Dummy()
    def forward(self, x:Tensor)->Tensor:
        x = self.t(x)
        x = self.d(x)
        return x

class Dummy(nn.Module):
    def __init__(self, *args, **kwargs ) -> None:
        super().__init__(*args, **kwargs)
        self.c = nn.Parameter(Tensor([-1.]))
        self.d = nn.Parameter(Tensor([-2.]))
    def forward(self, x:Tensor)->Tensor:
        return self.c*x + self.d

class MyModel(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.t0 = Target(1.,2.)
        self.d0 = Dummy()
        self.t1 = Target(3.,4.)
        self.m0 = Mother()
        self.t2 = Target(5.,6.)
        self.d1 = Dummy()
    def forward(self, x:Tensor)->Tensor:
        x = self.t0(x)
        x = self.d0(x)
        x = self.t1(x)
        x = self.m0(x)
        x = self.t2(x)
        x = self.d1(x)
        return x
```
</details>

パラメータの更新方法を次のように定めます．
* **更新しない層**：`Dummy`
* **更新する層**:`Target`
    * **学習率を高めにして更新するパラメータ**：`a`
    * **学習率を低めにして更新するパラメータ**：`b`

### 実装
#### ❶層の取り出し
パラメータ`a`，`b`を取り出すためには，まずはそれぞれが属する層，つまり`Target`層を取り出してくる必要があります．
次のコードでは，`Target`のみを取り出し，`Target`だけの配列を作ります．
```python
target_list = []
for module in model.modules(): 
    if isinstance(module, Target):
        print("✅", end="")
        target_list.append(module)
    else:
        print("🟥", end="")
    print(module)
target_list
```
`model.modules`メソッドは，自分のモジュール内の子モジュールを再帰的に挙げるものです．
```plaintext:結果
🟥MyModel(
  (t0): Target()
  (d0): Dummy()
  (t1): Target()
  (m0): TargetMother(
    (t): Target()
    (d): Dummy()
  )
  (t2): Target()
  (d1): Dummy()
)
✅Target()
🟥Dummy()
✅Target()
🟥TargetMother(
  (t): Target()
  (d): Dummy()
)
✅Target()
🟥Dummy()
✅Target()
🟥Dummy()

[Target(), Target(), Target(), Target()]
```
`Target`だけうまく選べていることが分かります．

#### ❷パラメータの取り出し
次に，`Target`層だけの配列から，`a`だけの配列と`b`だけの配列を作ります．
```python
a_param_list = []
b_param_list = []
for target in target_list:
    a_param_list.append(target.a)
    b_param_list.append(target.b)
a_param_list, b_param_list
```

<details><summary>結果</summary>

```plaintext
([Parameter containing:
  tensor([1.], requires_grad=True),
  Parameter containing:
  tensor([3.], requires_grad=True),
  Parameter containing:
  tensor([10.], requires_grad=True),
  Parameter containing:
  tensor([5.], requires_grad=True)],
 [Parameter containing:
  tensor([2.], requires_grad=True),
  Parameter containing:
  tensor([4.], requires_grad=True),
  Parameter containing:
  tensor([20.], requires_grad=True),
  Parameter containing:
  tensor([6.], requires_grad=True)])
```
</details>
結果を見ると，`a`と`b`だけが取り出されていることが分かります．

#### ❸最適化アルゴリズムの起動
最後に，PyTorchの機能を使って，Optimizerを起動します．
```python
optimizer = optim.SGD([
    {"params":a_param_list, "lr":1.},
    {"params":b_param_list, "lr":0.0001}
])
```
このように書くと，`a_param_list`内のパラメータに対しては学習率を1に，`b_param_list`内のパラメータに対しては学習率を0.0001にして学習が行われます．また，ここに出てこなかった`Dummy`層などのパラメータは更新されません．

#### ❹学習過程
試しに簡単な例で学習を行い，パラメータ更新前と後でどのパラメータがどのくらい変わったかを見てみます．
<details><summary>コードと結果</summary>

```python
print("Before")
for name, param in model.named_parameters():
    print(f"{name:<10s}", param.detach().numpy()[0])

x = tensor([3.0]) # 入力データ
y = model(x)
y # 損失関数

y.backward() # dy/dθ，それぞれの微分指定されたパラメータθ を計算する
optimizer.step() # パラメータを更新

print("After")
for name, param in model.named_parameters():
    print(f"{name:<10s}", param.detach().numpy()[0])

```
```plaintext
Before
t0.a       1.0
t0.b       2.0
d0.c       -1.0
d0.d       -2.0
t1.a       3.0
t1.b       4.0
m0.t.a     10.0
m0.t.b     20.0
m0.d.c     -1.0
m0.d.d     -2.0
t2.a       5.0
t2.b       6.0
d1.c       -1.0
d1.d       -2.0
After
t0.a       451.0
t0.b       2.015
d0.c       -1.0
d0.d       -2.0
t1.a       353.0
t1.b       3.995
m0.t.a     95.0
m0.t.b     19.9995
m0.d.c     -1.0
m0.d.d     -2.0
t2.a       153.0
t2.b       6.0001
d1.c       -1.0
d1.d       -2.0
```
</details>
結果を見ると，$a$が大きく更新されていて，$b$が小さく更新されていることが分かります．ちなみに，`Dummy`層にあるパラメータ$c,d$に関しては変化していないことが分かります．

`Target`層内にあるパラメータ`a`，`b`に対して，別々の学習率を設定して学習を進めることができました．