jittor.optim

这里是Jittor的优化器模块的API文档，您可以通过from jittor import optim来获取该模块。

class jittor.optim.Adam(params, lr, eps=1e-08, betas=(0.9, 0.999), weight_decay=0)[源代码]

Adam优化器，它结合了Momentum和RMSprop两种优化方法的主要观点，优化公式如下：

\[\begin{split}m_t &= b1 \times m_{t-1} + (1 - b1) \times g_t \\ v_t &= b2 \times v_{t-1} + (1 - b2) \times g_t^2 \\ p_t &= p_{t-1} - \frac{lr \times m_t} {\sqrt{v_t} + eps}\end{split}\]

其中，\(m_t\) 和 \(v_t\) 分别是梯度的一阶矩估计和二阶矩估计，\(b1\), \(b2\) 是估计的系数，\(lr\) 是学习率，\(eps\) 是用于数值稳定性的项，\(g_t\) 是梯度，\(p_t\) 是参数。

构造参数:

  • params (iterable): 待优化的参数，或定义了参数进行分组的字典。对于已经定义好的模型，可以用 model.parameters() 获取其参数

  • lr (float): 学习率

  • eps (float, 可选): 用于改善数值稳定性的项。默认值：1e-08

  • betas (Tuple[float, float], 可选): 用于计算梯度以及梯度平方的估计值的系数, 默认值: (0.9, 0.999)

  • weight_decay (float, 可选): 权重衰减（L2 penalty）。默认值: 0

step的执行参数：

  • loss (Var, 可选): 已经计算好的神经网络loss。默认值: None

  • retain_graph (bool, 可选): 是否保留计算图。默认值: False

代码示例: 
>>> optimizer = jt.optim.Adam(model.parameters(), lr, eps=1e-8, betas=(0.9, 0.999))
>>> optimizer.step(loss)
class jittor.optim.AdamW(params, lr, eps=1e-08, betas=(0.9, 0.999), weight_decay=0)[源代码]

AdamW 优化器，是 Adam 优化器的一个变体，其中对权重衰减进行了修正。

step 更新操作数学描述如下:

\[\begin{split}\begin{aligned} t &\leftarrow t + 1 \\ \widetilde{\eta_t} &\leftarrow \eta_t \times a^{t} \\ g_t &\leftarrow abla_{\theta} L_t(\theta) + \lambda \theta \\ m_t &\leftarrow \beta_{1_t} m_{t-1} + (1 - \beta_{1_t}) g_t \\ \widetilde{m_t} &\leftarrow \frac{m_t}{1 - \beta_{1_t}^t} \\ v_t &\leftarrow \beta_{2_t} v_{t-1} + (1 - \beta_{2_t}) g_t^2 \\ \widetilde{v_t} &\leftarrow \frac{v_t}{1 - \beta_{2_t}^t} \\ \theta &\leftarrow \theta - \widetilde{\eta_t} abla_{\theta} L_t(\theta) \\ \end{aligned}\end{split}\]

其中， \(\theta\) 是参数， \(g\) 是梯度， \(v\) 是梯度的平方， \(m\) 是梯度的指数移动平均值， \({m_t}\) 是偏差修正后的梯度的移动平均值, \({v_t}\) 是偏差修正后的梯度的平方的移动平均值， \(\eta\) 是学习率， \(\beta_{1_t}\) 和 \(\beta_{2_t}\) 是梯度和梯度平方动量项的系数， \(\lambda\) 是权重衰减系数。

参数:

  • params（iterable）：待优化参数的迭代器，或是定义了参数组的字典。

  • lr (float)：学习率。

  • eps (float)：为了增加数值计算的稳定性而加到分母里的项。默认值: 1e-8

  • betas (Tuple[float, float])：计算一阶动量和二阶动量的指数衰减率元组。默认值: (0.9, 0.999)

代码示例:
>>> import jittor as jt
>>> model = jt.nn.Linear(10, 2)
>>> loss_fn = jt.nn.CrossEntropyLoss()
>>> optimizer = jt.optim.AdamW(params=model.parameters(), lr=0.1, eps=1e-8, betas=(0.9, 0.999))
>>> x = jt.randn([5, 10])
>>> y_true = jt.array([0, 1, 0, 1, 1])
>>> y_pred = model(x)
>>> loss = loss_fn(y_pred, y_true)
>>> optimizer.step(loss)
class jittor.optim.Adan(params, lr=0.001, betas=(0.98, 0.92, 0.99), eps=1e-08, weight_decay=0.0, max_grad_norm=0.0)[源代码]

实现Adan优化器。它在 Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models[J].arXiv preprint arXiv:2208.06677, 2022 中提出。

Adan是大多数DNN框架的高效优化器，计算负载比其他最新方法小约2倍，对训练设置和批处理大小具有稳健性，易于即插即用。

构造参数:

  • params (iterable): 要优化的参数的可迭代对象或者定义参数组的字典.

  • lr (float, 可选): 学习率。默认值: 1e-3

  • betas (Tuple[float, float, float], 可选): 用于一阶和二阶矩的系数。默认值: (0.98, 0.92, 0.99)

  • eps (float, 可选): 用于改善数值稳定性的分母中添加的项。默认值: 1e-8

  • weight_decay (float, 可选): 分离的权重衰减 (L2惩罚) 。默认值：0

  • max_grad_norm (float, 可选): 用于裁剪全局梯度范数的值，默认值: 0.0 （无裁剪）

step执行参数：

  • loss (Var, 可选): 已经计算好的神经网络loss。默认值: None

  • retain_graph (bool, 可选): 是否保留计算图。默认值: False

代码示例: 
>>> optimizer = jt.optim.Adan(model.parameters(), lr=0.001)
>>> optimizer.step(loss)
class jittor.optim.LRScheduler(optimizer, last_epoch=-1)[源代码]

学习率调度器的基类，用来根据训练的轮次调整学习率。

参数:

  • optimizer (Optimizer): 优化器，用来进行模型参数的优化。

  • last_epoch (int, 可选): 最后的轮次（epoch）。该值默认为 -1，代表使用优化器的学习率来初始化。否则，需要确保每个参数组中都有 initial_lr 这个值来进行学习率的初始化。

注意:

该类方法的使用，需要搭配其子类来使用，因为 get_lr() 方法需要在子类中被重写以定义如何根据轮次来调整学习率。

class jittor.optim.LambdaLR(optimizer, lr_lambda, last_epoch=-1)[源代码]

用于实现学习率调度功能的类，其根据使用者设置的学习率迭代公式（可用 lambda 表达式作为参数给出），对基本学习率进行调整。

参数:

  • optimizer (Optimizer): 已经初始化的优化器。

  • lr_lambda (callable or list): 一个函数或者一个函数列表，用以定义学习率策略. 对应公式为: \(lr = \text{init_lr} \times \text{lr_lambda(last_epoch)}\)

  • last_epoch (int, 可选): 最后一次迭代的 epoch 数。默认值: -1。

代码示例:
>>> from jittor.optim import SGD, LambdaLR
>>> optimizer = SGD(model.parameters(), lr=0.1)
>>> scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: 0.95 ** epoch)
>>> for epoch in range(100):
...     train(...)
...     validate(...)
...     scheduler.step()
class jittor.optim.Optimizer(params, lr, param_sync_iter=10000)[源代码]

优化器的基础类。这个基类可以用来实现各种优化算法，比如随机梯度下降等。

参数：

  • params (list): 模型参数。

  • lr (float): 学习速率。

  • param_sync_iter (int, optional): 参数同步的迭代次数。默认值：10000

代码示例：
>>> import jittor as jt
>>> from jittor.optim import Optimizer
>>> class MyOptimizer(Optimizer):
...     def __init__(self, params, lr):
...         super(MyOptimizer, self).__init__(params, lr)     
...     def step(self, loss):
...         self.zero_grad()
...         self.backward(loss)
...             for group in self.param_groups:
...             for param, grad in zip(group['params'], group['grads']):
...                 if not param.is_stop_grad():
...                     param.update(param - self.lr * grad)
... 
>>> x = jt.randn([2,3])
>>> optimizer = MyOptimizer([x], lr=0.1)
>>> loss = x.sum()
>>> optimizer.step(loss)
>>> print(x)
jt.Var([[-0.481019    0.01914055 -0.74143946]
        [ 0.33761212 -1.7029546  -0.8524694 ]], dtype=float32)
backward(loss, retain_graph=False)[源代码]

optimize.backward(loss) 用于累积多个step的梯度，可以如下使用：

原始源代码：
>>> n_iter = 10000
... batch_size = 100
... ...
>>> for i in range(n_iter):
...     ...
...     loss = calc_loss()
...     optimizer.step(loss)
累积版本：
>>> n_iter = 10000
... batch_size = 100
... accumulation_steps = 10
... n_iter *= accumulation_steps
... batch_size //= accumulation_steps
... ...
>>> for i in range(n_iter):
...     ...
...     loss = calc_loss()
...     # 如果损失是跨批次的平均值，我们需要除以 accumulation_steps
...     optimizer.backward(loss / accumulation_steps)
...     if (i+1) % accumulation_steps == 0:
...         optimizer.step()
clip_grad_norm(max_norm: float, norm_type: int = 2)[源代码]

剪切此优化器的梯度范数，范数是对所有梯度一起计算的。

参数：

max_norm (float or int): 梯度的最大范数 norm_type (int): 1-范数或2-范数

示例：
>>> a = jt.ones(2)
... opt = jt.optim.SGD([a], 0.1)
... loss = a*a
... opt.zero_grad()
... opt.backward(loss)
... print(opt.param_groups[0]['grads'][0].norm())
2.83
>>> opt.clip_grad_norm(0.01, 2)
... print(opt.param_groups[0]['grads'][0].norm())
0.01
>>> opt.step()
load_state_dict(state)[源代码]
post_step()[源代码]

在 step 之后执行的操作，比如更新参数值。

例子：

>>> class MyOptimizer(Optimizer):
...     def step(self, loss):
...         self.pre_step(loss)
...         ...
...         self.post_step()
pre_step(loss, retain_graph=False)[源代码]

在 step 之前执行的操作，比如计算梯度，retain_graph 等。

例子：

>>> class MyOptimizer(Optimizer):
...     def step(self, loss):
...         self.pre_step(loss)
...         ...
...         self.post_step()
class jittor.optim.RMSprop(params, lr=0.01, eps=1e-08, alpha=0.99)[源代码]

RMSprop 优化器类，用于在深度学习模型训练过程中更新模型参数以最小化损失函数，继承自 Optimizer 类。 RMSprop 是一种自适应学习率方法，为每个参数设置单独的学习率。它的公式包括累积平方梯度和参数更新的计算：

\[\begin{split}v(w, t) &= \alpha \cdot v(w, t-1) + (1 - \alpha) g(w, t)^2 \\ w &= w - \frac{lr \cdot g(w, t)}{\sqrt{v(w, t)} + \varepsilon}\end{split}\]

其中，\(g(w, t)\) 为 \(t\) 时刻参数 \(w\) 上的梯度，\(v\) 为累积平方梯度， \(\alpha\) 为平滑常数， \(\varepsilon\) 为防止零除项。

参数:

  • params (list): 被优化的模型的参数。

  • lr (float): 学习率，默认值：1e-2

  • eps (float): 添加到分母的术语以避免零除。默认值： 1e-8

  • alpha (float): 平滑常数。 默认值： 0.99

代码示例：
>>> import jittor as jt
... model = jt.nn.Linear(10, 2)
... loss_fn = jt.nn.CrossEntropyLoss()
... optimizer = jt.optim.RMSprop(model.parameters(), lr=0.1)
... x = jt.randn([5, 10])
... y_true = jt.array([0, 1, 0, 1, 1])
... y_pred = model(x)
... loss = loss_fn(y_pred, y_true)
... optimizer.step(loss)
class jittor.optim.SGD(params, lr, momentum=0, weight_decay=0, dampening=0, nesterov=False)[源代码]

随机梯度下降算法优化器的实现，包括学习率衰减、动量、权重衰减等功能。在 step 方法中，根据参数的梯度和当前的动量更新参数值，也支持添加新的参数组，并能够计算带有动量的更新值。

该算法使用 SGD 进行更新，其中加入了动量项、权重衰减项，并可以选择是否使用 Nesterov 加速。如果使用动量项，则更新过程会考虑历史的梯度信息，更新公式如下:

\[\begin{split}v &= \text{momentum} \times v + dp \times (1 - \text{dampening}) \\ dp &= p \times \text{weight_decay} + g\end{split}\]

如果使用 Nesterov 动量, 则更新过程会提前按历史方向进行一步预测，这样在某些问题上会获得更好的训练效果，更新公式如下:

\[p = p - (dp + \text{momentum} * v) * lr\]

否则，更新公式为：

\[p = p - v \times lr\]

以上 \(p,g,v\) 分别表示参数值、梯度以及参数更新量。

参数：

  • params(list): 待优化的参数或者已经定义好的参数组

  • lr(float): 学习率，用于控制参数更新的步长。

  • momentum(float, 可选): 动量因子。默认值： 0

  • weight_decay(float, 可选): 权重衰减系数。默认值： 0

  • dampening(float, 可选): 动量的抑制因子。默认值： 0

  • nesterov(bool, 可选): 是否采用 Nesterov 动量。默认值：False

代码示例：
>>> import jittor as jt
>>> model = jt.nn.Linear(10, 2)
>>> loss_fn = jt.nn.CrossEntropyLoss()
>>> optimizer = jt.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.001)
>>> x = jt.randn([5, 10])
>>> y_true = jt.array([0, 1, 0, 1, 1])
>>> y_pred = model(x)
>>> loss = loss_fn(y_pred, y_true)
>>> optimizer.step(loss)
jittor.optim.opt_grad(v: Var, opt: Optimizer)[源代码]

获取优化器中某个变量的梯度，

参数：

  • v(Var): 优化器中的变量

  • opt (Optimizer): 优化器

代码示例：
>>> model = Model()
>>> optimizer = SGD(model.parameters())
>>> ...
>>> optimizer.backward(loss)
>>> for p in model.parameters():
>>>     grad = p.opt_grad(optimizer)
返回值：

优化器中某个变量的梯度（Var）

以下是Jittor的学习率调度模块的API文档，学习率调度模块需要配合优化器使用，您可以通过from jittor import lr_scheduler来获取该模块。

class jittor.lr_scheduler.CosineAnnealingLR(optimizer, T_max, eta_min=0, last_epoch=-1)[源代码]

是用于实现余弦退火学习率调度策略的工具。该类实现了余弦退火学习率调度策略，其中第 \(t\) 步的学习率计算过程如下：

\[ \begin{align}\begin{aligned}\eta_t = \eta_{min} + \dfrac{1}{2}(\eta_{max} - \eta_{min})\left(1+cos\left(\dfrac{T_{cur}}{T_{max}}\pi\right)\right), &\quad T_{cur} = (2k+1)T_{max};\\\eta_{t+1} = \eta_{t} + \dfrac{1}{2}(\eta_{max} - \eta_{min})\left(1-cos\left(\dfrac{1}{T_{max}}\pi\right)\right), &\quad T_{cur} = (2k+1)T_{max}.\end{aligned}\end{align} \]

其中，\(\eta_{max}\) 是初始学习率，\(\eta_{min}\) 是学习率下限，\(T_{max}\) 是最大周期数， \(T_{cur}\) 是 SGDR 上一次重启后的周期（epoch）数， 通过 last_epoch 指定。特别地， last_epoch 为 ``-1``（默认）时，将学习率设置为优化器的初始学习率。

参数:

  • optimizer (Optimizer): 优化器对象，其学习率将被调整。

  • T_max (int): 用来计算学习率的最大周期数。

  • eta_min (float, 可选): 学习率下限，默认为 0

  • last_epoch (int, 可选): 最后一次迭代的周期数，默认为 -1

class jittor.lr_scheduler.ExponentialLR(optimizer, gamma, last_epoch=-1)[源代码]

此类实现指数衰减学习率优化策略。在每个步骤中，将学习率乘以一个给定的系数 gamma。当 last_epoch 为 -1 时，将学习率设置为优化器的初始学习率。

参数:

  • optimizer (Optimizer): 优化器对象，需要进行学习率衰减的优化器。

  • gamma (float): 在每个步骤中乘以学习率的系数。

  • last_epoch (int): 上一个周期（epoch）编号。默认值：-1

代码示例：
>>> optimizer = jt.optim.SGD(model.parameters(), lr=0.1)
>>> scheduler = ExponentialLR(optimizer, gamma=0.9)
>>> for epoch in range(100):
...     train(...)
...     validate(...)
...     scheduler.step()
class jittor.lr_scheduler.MultiStepLR(optimizer, milestones=[], gamma=0.1, last_epoch=-1)[源代码]

实现了一个常见的学习率调节策略，即在预设的多个阶段（milestones）对学习率进行衰减。

参数：

  • optimizer (Optimizer): 被包装的优化器实例。

  • milestones (list): 预设的学习率衰减阶段。默认值：[]

  • gamma (float): 学习率衰减系数，每次衰减时，学习率都会乘以这个系数。默认值：0.1

  • last_epoch (int): 上一个 epoch 的标记，默认值：-1。

代码示例：
>>> # 当 epoch 达到 30 和 80 时学习率乘以 0.1
>>> scheduler = MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)
>>> for epoch in range(100):
...     train(...)
...     validate(...)
...     scheduler.step()
class jittor.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, verbose=False, threshold=0.0001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-08)[源代码]

在训练深度学习模型时，如果训练过程中的损失（Loss）一直不再降低则表明已进入平台期，将学习率 (Learning Rate) 降低可以有助于模型跳出局部最小值点。 ReduceLROnPlateau 类提供了这样一种策略：当损失在一定期限 (patience) 内没有降低时，降低学习率。在一定的慢冷期 (cooldown) 内，损失的变化不再触发学习率的减小。

参数：

  • optimizer (Optimizer): 要使用的优化器.

  • mode (str): 有两种模式 'min' 和 'max'。默认为 'min'，当损失不再降低时减小学习率。如果设置为 'max'，当损失不再增加时减小学习率。

  • factor (float): 学习率每次降低的倍数。默认为 0.1

  • patience (int): 当损失在这么多次迭代后没有下降时，学习率进行调整。默认为 10

  • verbose (bool): 设为 True 会在每次更新学习率时打印信息。默认为 False

  • threshold (float): 判断损失是否不再减少的阈值。只有损失减少的幅度大于该值才算作损失的减小。默认为 0.0001

  • threshold_mode (str): 损失阈值模式，'rel' 或 'abs' 。在 'rel' 模式下，只要相较于上一次的损失减少了相对于上一次损失的 threshold 个百分比即算作损失的减小；在 ‘abs’ 模式下，只有 相较于上一次的损失减少了 threshold 即算作损失的减小。默认为 'rel'

  • cooldown (int): 每次降低学习率后，会有一段慢冷期，在该期间内的损失变化不再触发学习率的减小。默认为 0

  • min_lr (float or list): 学习率的下限，可为单一浮点数或浮点数的列表。如果是浮点数，则所有参数组的学习率下限都是这个值，如果是列表，则每个参数组的学习率下限是列表中的对应值。默认为 0

  • eps (float): 学习率的最小减少值。如果新旧学习率之间的差异小于 eps ，则忽略这次更新。默认为 1e-8

代码示例：
>>> import jittor as jt
>>> optimizer = jt.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10)
>>> for epoch in range(100):
...     train(...)
...     loss = validate(...)
...     scheduler.step(loss)
class jittor.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1)[源代码]

实现了对学习率的周期性衰减调度，每达设定的步数时，对学习率进行一次衰减。

参数:

  • optimizer (Optimizer) : 使用的优化器。

  • step_size (int) : 设定的步数，即每过多少轮后，进行一次衰减。

  • gamma (float, 可选) : 学习率衰减的系数, 默认为 0.1

  • last_epoch (int, 可选) : 上个周期（epoch）的编号，即初始的步数。默认为 -1 ，表示从第一个周期开始调整。

代码示例：
>>> scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
>>> for epoch in range(100):
...     train(...)
...     validate(...)
...     scheduler.step()