优化器和学习率适配
MindTorch优化器对标了PyTorch的优化器大部分功能，具体可以在torch接口支持列表中查看。本章节主要描述动态图和精度图模式下使用的一些差异。
优化器
1.param_groups中学习率的类型差异
出于对后续性能优化的考虑，MindTorch优化器成员param_groups中学习率lr的默认类型是MindSpore的Parameter类型，而非PyTorch定义的Number类型, 一般该差异会涉及两个场景的使用：打印和修改学习率。
1.1 打印学习率
由于lr的类型不同，打印时为达到与PyTorch一样的效果，需要将lr转换成number类型。 以下面代码为例：
PyTorch代码：
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代码：
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的方式修改学习率。比如以下代码：
静态图学习率修改例子：
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代码：
import torch
...
net = Net()
loss = net(input)
loss.backward()
optimizer.step()
对应的MindTorch代码：
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代码：
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代码：
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代码：
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代码：
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类型（只有这样优化器学习率才能在静态图下正常使用）。比如以下静态图适配例子：
例子一：
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()
例子二：
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的写法。
2.3 暂不支持静态图模式下编译加速
由于自定义LRScheduler存在静态图模式下无法支持的语法，因此不支持编译使用。可以将执行的代码移动到编译的范围外避免该问题。
比如以下例子：
原始代码：
import mindspore as ms
scheduler = MySched()
@ms.jit
def train_step():
...
scheduler.step() # 在ms.jit装饰器范围内，会进行编译，此时编译会报错语法不支持。
train_step()
适配后代码：
import mindspore as ms
scheduler = MySched()
@ms.jit
def train_step():
...
train_step()
scheduler.step() # 将函数调用移到ms.jit装饰器范围外，不参与编译，以动态图方式执行。