# 优化器和学习率适配
MindTorch优化器对标了PyTorch的优化器大部分功能，具体可以在[torch接口支持列表](SupportedList.md)中查看。本章节主要描述动态图和精度图模式下使用的一些差异。

## 优化器

### 1.`param_groups`中学习率的类型差异
出于对后续性能优化的考虑，MindTorch优化器成员param_groups中学习率lr的默认类型是MindSpore的Parameter类型，而非PyTorch定义的Number类型,
一般该差异会涉及两个场景的使用：打印和修改学习率。

#### 1.1 打印学习率

由于lr的类型不同，打印时为达到与PyTorch一样的效果，**需要将lr转换成number类型。**
以下面代码为例：

PyTorch代码：
```python
import torch
optimizer = torch.optim.SGD([torch.nn.Parameter(torch.tensor(2.))], lr=0.01)
print("lr is {}".format(optimizer.param_groups[0]['lr']))
```

对应的MindTorch代码：
```python
import mindtorch.torch as torch
optimizer = torch.optim.SGD([torch.nn.Parameter(torch.tensor(2.))], lr=0.01)
print("lr is {}".format(float(optimizer.param_groups[0]['lr'])))  # 通过`float(optimizer.param_groups[0]['lr'])`将其转换为`number`类型。
```

#### 1.2 修改学习率

- 动态图模式下，与PyTorch没有差异。因为虽然lr的默认类型为mindspore的Parameter，但是在实现上，支持了修改成number类型的功能。
- 但是在**静态图**模式下，只能使用`mindspore.ops.assign`的方式修改学习率。比如以下代码：

**静态图学习率修改例子：**
```python
import mindspore
import mindtorch as torch
optimizer = torch.optim.SGD([torch.nn.Parameter(torch.tensor(2.))], lr=0.01)
# optimizer.param_groups[0]['lr'] = 0.1                    # 静态图下不支持直接赋值的修改方法
mindspore.ops.assign(optimizer.param_groups[0]['lr'], 0.1) # 需要使用mindspore.ops.assign的方式修改对应值
```

### 2. `optimizer.step()`的入参差异
与PyTorch用法相同的微分方案正在开发中，当前调用optimizer.step时仍**需将梯度作为入参传入。**

比如以下例子：

PyTorch代码：
```python
import torch
...
net = Net()
loss = net(input)
loss.backward()
optimizer.step()
```

对应的MindTorch代码：
```python
import mindspore
import mindtorch.torch as torch
...
net = Net()
grad_fn = mindspore.ops.value_and_grad(net, None, optimizer.parameters)
grads = grad_fn(input) # 通过value_and_grad接口求得梯度grads
optimizer.step(grads)  # 将梯度grads，通过入参的方式传入到step函数中。
```

### 3. 自定义优化器
MindTorch提供了optim.Optimizer父类, 用户可继承使用实现自定义优化器。**但是**，也同样存在上述1、2小节中所述的差异(学习率类型不同、step函数入参不同)，用户在实现自定义优化器时请留意。

下面为一个动态图模式下的迁移例子：

PyTorch代码：
```python
import torch
class Ranger(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3, alpha=0.5, k=6):
        defaults = dict(lr=lr, alpha=alpha)
        super().__init__(params, defaults)
        self.k = k
    def __setstate__(self, state):
        print("set state called")
        super().__setstate__(state)
    def step(self, closure=None):
        loss = None
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                p_data_fp32 = p.data.float()
                state = self.state[p]
                state['step'] += 1
                p_data_fp32.add_(grad)
                p.data.copy_(p_data_fp32)
        return loss
```

MindTorch代码：
```python
import mindtorch.torch as torch
class Ranger(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3, alpha=0.5, k=6):
        defaults = dict(lr=lr, alpha=alpha)
        super().__init__(params, defaults)
        self.k = k
    def __setstate__(self, state):
        print("set state called")
        super().__setstate__(state)
    def step(self, grads, closure=None): # 需要新增grads作为函数入参，以便传入梯度
        loss = None
        i = -1                           # 声明一个索引，用来遍历grads入参
        for group in self.param_groups:
            for p in group['params']:
                i = i + 1                # 索引递增
                grad = grads[i]          # grad从入参中获取。如果对应Parameter没有参与求导，grad为0
                p_data_fp32 = p.data.float()
                state = self.state[p]
                state['step'] += 1
                p_data_fp32.add_(grad)
                p.data.copy_(p_data_fp32)
        return loss
```

## 学习率

### 1. 要求优化器中学习率为Parameter类型

出于静态图下性能优化的考虑，当前MindTorch封装的LRScheduler，暂不支持对学习率为number类型的优化器进行学习率调整，只支持Parameter类型。

**如果需要对number类型学习率的优化器进行支持**，可以通过**自定义**的方式，在动态图下使用。比如下面例子，通过自定义一个StepLR，来实现对number类型学习率优化器的支持：

PyTorch代码：
```python
import torch
my_optimizer_with_lr_number = ...    #自定义的optimizer中lr为Number类型
step_lr = torch.optim.StepLR(my_optimizer_with_lr_number, arg.step_size)
step_lr.step()
```

相应的MindTorch代码：
```python
from mindtorch.torch.optim.lr_scheduler import LRScheduler
my_optimizer_with_lr_number = ...    #自定义的optimizer中lr为Number类型

# 重新定义StepLR，实现对应逻辑，适配`group['lr']`为number类型。
class StepLR(LRScheduler):
    def __init__(self, optimizer, step_size, gamma=0.1, last_epoch=-1, verbose=False):
        self.step_size = step_size
        self.gamma = gamma
        super().__init__(optimizer, last_epoch, verbose)

    def get_lr(self):
        if (self.last_epoch == 0) or (self.last_epoch % self.step_size != 0):
            return [group['lr'] for group in self.optimizer.param_groups]
        return [group['lr'] * self.gamma
                for group in self.optimizer.param_groups]   # 经过get_lr计算逻辑返回number类型学习率，父类step函数会自动识别number类型并正确执行。

    def _get_closed_form_lr(self):
        return [base_lr * self.gamma ** (self.last_epoch // self.step_size)
                for base_lr in self.base_lrs]

step_lr = StepLR(my_optimizer_with_lr_number, arg.step_size)
step_lr.step()
```

### 2. 自定义`LRSchduler`

#### 2.1 修改优化器学习率的方式

- 动态图下，修改方式与PyTorch一致。
- 但是在**静态图**下，需要使用`mindspore.ops.assign`的方式对优化器中学习率进行修改（即使LRSchduler不参与编译），以保证优化器中学习率一直是Parameter类型（只有这样优化器学习率才能在静态图下正常使用）。比如以下**静态图**适配例子：

*例子一：*

```python
class TransformerLrScheduler():
  def __init__(self, optimizer, d_model, warmup_steps, multiplier=5):
    self._optimizer = optimizer
    self.d_model = d_model
    self.warmup_steps = warmup_steps
    self.n_steps = 0
    self.multiplier = multiplier

  def step(self):
    self.n_steps += 1
    lr = self._get_lr()
    for param_group in self._optimizer.param_groups:
        # param_group['lr'] = lr                    # PyTorch原始用法，直接对优化器中学习率赋值
        mindspore.ops.assign(param_group['lr'] ,lr) # 适配成ops.assign的方式, 以支持静态图下对优化器中学习率的修改
  def _get_lr(self):
    return self.multiplier * (self.d_model ** -0.5) * min(self.n_steps ** (-0.5), self.n_steps * (self.warmup_steps ** (-1.5)))

scheduler = TransformerLrScheduler(optimizer, args.d_encoder, args.warmup_steps)
...
scheduler.step()
```

*例子二：*

```python
def adjust_learning_rate(optimizer, gamma, step):
    lr = args.lr * (gamma ** (step))
    for param_group in optimizer.param_groups:
        # param_group['lr'] = lr                    # PyTorch原始用法，直接对优化器中学习率赋值 
        mindspore.ops.assign(param_group['lr'], lr) # 适配成ops.assign的方式, 以支持静态图下对优化器中学习率的修改
```

#### 2.2 `state_dict`及`load_state_dict`函数的适配方法

由于优化器中学习率默认类型为Parameter，所以`base_lrs`,`_last_lr`默认也将会是Parameter类型。那么此时保存和恢复状态时，就需要做额外的类型装换。

- 如果不需要重写`state_dict`、`load_state_dict`方法，则无需额外操作，父类已经做好相应处理。
- **如果需要重写方法**，可以调用父类的接口来实现自动类型转换。可参考MindTorch [CyclicLR](https://openi.pcl.ac.cn/OpenI/MSAdapter/src/branch/master/mindtorch/torch/optim/lr_scheduler.py#CyclicLR)的写法。


#### 2.3 暂不支持静态图模式下编译加速
由于自定义LRScheduler存在静态图模式下无法支持的语法，因此**不支持编译使用**。可以将执行的代码**移动到编译的范围外**避免该问题。

比如以下例子：

原始代码：
```python
import mindspore as ms
scheduler = MySched()
@ms.jit
def train_step():
    ...
    scheduler.step()  # 在ms.jit装饰器范围内，会进行编译，此时编译会报错语法不支持。

train_step()
```

适配后代码：
```python
import mindspore as ms
scheduler = MySched()
@ms.jit
def train_step():
    ...

train_step()
scheduler.step() # 将函数调用移到ms.jit装饰器范围外，不参与编译，以动态图方式执行。
```