YOLOv9-0.1部分代码阅读笔记-lion.py
lion.py
utils\lion.py
目录
lion.py
1.所需的库和模块
2.class Lion(Optimizer):
1.所需的库和模块
# Lion优化器是一种新型的神经网络优化算法，由Google Brain团队通过遗传算法发现，全称为EvoLved SIgn MOmeNtum，意为“进化的符号动量”。以下是Lion优化器的一些主要特点和优势 ：
# 性能表现 ：Lion优化器在多项任务中的表现超过了目前广泛使用的Adam和AdamW优化器。
# 学习率设置 ：Lion的学习率通常需要比AdamW小3-10倍。例如，如果之前使用AdamW时的学习率是1e-3，那么使用Lion时可能需要将学习率调整到1e-4左右。
# 权重衰减 ：由于有效权重衰减是学习率与λ的乘积，因此Lion使用的解耦权重衰减λ值通常需要比AdamW大3-10倍，以保持相似的强度。
# 学习率调度 ：研究人员发现，使用余弦衰减调度训练ViT模型时，Lion的性能增益比使用倒数平方根调度更大。
# 实际应用表现 ：在语言建模任务中，当参数设置合适时，Lion通常能够取得优于Adam的效果。在大规模文本到图像训练中也有不错的表现，但可能需要一些调优工作。Lion对批量大小和数据量/增强的敏感度较高，建议在批量大小大于或等于64的情况下使用。
# 内存和效率优势 ：与AdamW和各种自适应优化器相比，Lion只需要动量并利用符号操作来计算更新，并且将额外的内存占用减半。这在训练大模型或大批量时很有用。Lion在实验中具有更快的运行时间（steps/sec），通常比AdamW和Adafactor提速2-15%，具体取决于任务、代码库和硬件。
# 参数设置 ：Lion的性能仍然高度依赖于具体任务和参数设置，研究人员和工程师们需要在实际应用中进行更多的实验和调优。
# Lion优化器作为一种新兴的神经网络优化算法，展现出了巨大的潜力，它在某些任务中能够以更少的计算资源达到与现有优化器相当甚至更好的效果。然而，Lion的性能仍然高度依赖于具体任务和参数设置，需要在实际应用中进行更多的实验和调优。
# Lion 优化器的 PyTorch 实现。
"""PyTorch implementation of the Lion optimizer."""
import torch
from torch.optim.optimizer import Optimizer2.class Lion(Optimizer):
# 这段代码是一个PyTorch优化器类的实现，它定义了一个名为 Lion 的优化器。这个类继承自PyTorch的 Optimizer 基类，并实现了其 __init__ 和 step 方法。这个优化器的实现类似于 Adam 优化器，但是它在更新权重时使用了梯度的符号（ torch.sign ）而不是梯度本身。
# 定义了一个名为 Lion 的新类，它继承自PyTorch的 Optimizer 基类。
class Lion(Optimizer):
# 实现 Lion 算法。
r"""Implements Lion algorithm."""
# 这段代码是 Lion 优化器类的构造函数 __init__ ，它负责初始化优化器的参数。
# 这行定义了 Lion 类的构造函数 __init__ 。它接受四个参数
# 1.self : 指向类的实例，总是第一个参数。
# 2.params : 一个参数列表，包含了需要优化的参数。
# 3.lr : 学习率，默认值为 1e-4 。
# 4.betas : 一个包含两个元素的元组，分别代表一阶矩（ beta1 ）和二阶矩（ beta2 ）的衰减率，默认为 (0.9, 0.99) 。
# 5.weight_decay : 权重衰减率，默认值为 0.0 。
def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0):
# 初始化超参数。
"""Initialize the hyperparameters.
Args:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-4)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.99))
weight_decay (float, optional): weight decay coefficient (default: 0)
"""
# 检查学习率 lr 是否大于等于0，如果不是，则抛出 ValueError 异常，并显示无效的学习率值。
if not 0.0 <= lr:
raise ValueError('Invalid learning rate: {}'.format(lr)) # 无效的学习率：{} 。
# 检查 betas 元组中的第一个值（ beta1 ）是否在闭区间 [0.0, 1.0) 内，如果不是，则抛出 ValueError 异常，并显示无效的 beta1 值。
if not 0.0 <= betas[0] < 1.0:
raise ValueError('Invalid beta parameter at index 0: {}'.format(betas[0])) # 索引 0 处的 beta 参数无效：{} 。
# 检查 betas 元组中的第二个值（ beta2 ）是否在闭区间 [0.0, 1.0) 内，如果不是，则抛出 ValueError 异常，并显示无效的 beta2 值。
if not 0.0 <= betas[1] < 1.0:
raise ValueError('Invalid beta parameter at index 1: {}'.format(betas[1])) # 索引 1 处的 beta 参数无效：{} 。
# 创建一个字典 defaults ，存储优化器的默认参数，包括 学习率 lr 、 betas 和 权重衰减率 weight_decay 。
defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
# 在 PyTorch 中自定义优化器时，当你调用父类 Optimizer 类的 __init__ 方法，你需要传入以下参数 ：
# params ：这是一个参数组（通常是一个列表），包含了需要优化的参数（例如模型的权重和偏置）。每个参数组可以是一个字典，其中包含了参数和对应的超参数设置。
# defaults ：这是一个字典，包含了优化器的默认参数设置。这些设置将被用于所有参数组，除非在参数组中被明确覆盖。常见的默认参数包括学习率 lr 、动量 momentum 、权重衰减 weight_decay 等。
# 以下是一个自定义优化器构造函数的示例，展示了如何调用父类的 __init__ 方法 ：
# import torch
# class CustomOptimizer(torch.optim.Optimizer):
# def __init__(self, params, lr=1e-3, momentum=0, weight_decay=0):
# # 默认参数设置
# defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay)
# # 调用父类的构造函数
# super(CustomOptimizer, self).__init__(params, defaults)
# # 自定义优化器的其他初始化代码...
# 在这个例子中， params 是需要优化的参数列表，而 defaults 是一个字典，包含了学习率 lr 、动量 momentum 和权重衰减 weight_decay 等默认参数设置。这些默认设置将被应用到所有参数组，除非在参数组中被特别指定。通过这种方式，自定义优化器可以灵活地为不同的参数设置不同的优化参数。
# 在 PyTorch 的 Optimizer 类中， defaults 字典应该包含所有优化器的默认超参数。这些超参数将作为参数组的默认值，如果某个参数组中没有指定对应的键值对，就会使用 defaults 中的值。以下是一些常见的键值对 ：
# lr (float) ：学习率，用于控制每次迭代中参数更新的步长。
# momentum (float, 可选) ：动量项，用于加速 SGD 优化器中的梯度下降。
# betas (tuple of floats, 可选) ：对于像 Adam 这样的优化器， betas 是一个包含两个值的元组，分别用于计算梯度的一阶和二阶矩估计。
# eps (float, 可选) ：用于防止除以零的小值，通常在 Adam 优化器中使用。
# weight_decay (float) ：权重衰减，用于 L2 正则化。
# alpha (float, 可选) ：在某些优化器中， alpha 可以用于控制学习率的缩放。
# amsgrad (bool, 可选) ：仅在 Adam 优化器中使用，用于启用 AMSGrad 变体。
# nesterov (bool, 可选) ：对于使用 Nesterov 加速梯度的 SGD 优化器，这个参数用于控制是否启用 Nesterov 动量。
# 这些键值对构成了优化器的默认设置，它们可以被覆盖或扩展，以适应不同的训练需求。自定义优化器时，可以在 defaults 字典中包含任何需要的超参数，并在优化器的 step 方法中使用这些参数来控制参数的更新逻辑。
# 以下是一个包含这些键值对的 defaults 字典示例 ：
# defaults = dict(
# lr=0.01, # 学习率
# momentum=0.9, # 动量
# betas=(0.9, 0.999), # Adam 优化器的两个 beta 参数
# eps=1e-08, # 防止除以零的小值
# weight_decay=0, # 权重衰减
# alpha=0.99, # 某些优化器中的 alpha 参数
# amsgrad=False, # Adam 优化器的 AMSGrad 变体
# nesterov=False # Nesterov 动量
# )
# 在实际使用中，只需要包含那些优化器实际使用的参数。不是所有的优化器都需要上述所有的参数。例如，如果正在实现一个简单的 SGD 优化器，你可能只需要 lr 和 momentum 参数。
# 调用父类 Optimizer 的构造函数，将 params 和 defaults 传递给它，以完成类的初始化。 super() 函数用于调用父类的方法。这里，它用于确保 Lion 类正确地继承和初始化 Optimizer 类。
super().__init__(params, defaults)
# __init__ 方法的主要作用是初始化 Lion 优化器的参数，并进行参数有效性检查。它设置了学习率、两个beta值（用于计算梯度的指数移动平均）和权重衰减率。这些参数对于优化器的行为至关重要。如果任何参数不在预期的范围内，方法会抛出异常，以确保优化器能够以正确的配置运行。最后，通过调用父类的构造函数， Lion 优化器被正式注册为PyTorch优化器，准备好用于后续的参数更新步骤。
# 这段代码定义了 Lion 优化器的 step 方法，它是执行单步优化的核心函数。
# @torch.no_grad()
# @torch.no_grad() 是 PyTorch 中的一个装饰器，用于包裹一个函数，使得在该函数内部执行的所有操作都不会跟踪梯度，也就是说，这些操作不会计算梯度，也不会消耗计算图。
# 这通常用于推理（inference）或者评估（evaluation）阶段，因为在这些阶段我们不需要进行反向传播，因此不需要计算梯度，这样可以减少内存消耗并提高计算效率。
# 说明 ：
# 当你使用 @torch.no_grad() 装饰器时，它会暂时将 PyTorch 设置为评估模式，在这个模式下，所有的 Variable 对象都会禁用梯度计算。
# 装饰器的作用域仅限于它所装饰的函数内部。
# 如果你需要在代码的某个特定区域临时禁用梯度计算，而不是整个函数，可以使用 torch.no_grad() 作为一个上下文管理器： with torch.no_grad(): # 在这个代码块中，梯度计算被禁用 result = model(input_data)
# 在示例中， evaluate_model 函数被 @torch.no_grad() 装饰，这意味着在函数内部执行的所有操作都不会跟踪梯度。这在模型评估时非常有用，因为评估时通常不需要进行梯度更新。
# 一个装饰器，用于指示 PyTorch 在执行 step 方法时不要计算梯度。这通常用于优化步骤，因为在优化过程中不需要进行梯度计算。
@torch.no_grad()
# 定义 step 方法，它接受一个可选参数。
# 1.closure ：这个参数是一个可调用对象，当提供时，它会被用来计算模型的损失值。
def step(self, closure=None):
# 执行单个优化步骤。
"""Performs a single optimization step.
Args:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
Returns:
the loss.
"""
# 初始化 loss 变量为 None ，用于存储如果提供了 closure 时计算的损失值。
loss = None
# 检查是否提供了 closure 参数。
if closure is not None:
# torch.enable_grad()
# 在 PyTorch 中， torch.enable_grad() 函数用于开启梯度计算。这个函数通常与 torch.no_grad() 配合使用，用于在特定的代码块中启用或禁用梯度计算。
# torch.enable_grad() 函数实际上是 torch.set_grad_enabled() 函数的一个便捷调用，它将梯度计算的状态设置为 True 。这意味着在调用 torch.enable_grad() 的代码块中，所有的操作都会跟踪梯度，以便后续可以进行反向传播。
# torch.enable_grad() 通常用于以下情况 ：
# 评估模型时计算特定的梯度 ：在模型评估阶段，你可能需要计算某些特定部分的梯度，而不是整个模型的梯度。
# 动态控制梯度计算 ：在某些复杂的操作中，你可能需要动态地开启或关闭梯度计算，以优化性能或内存使用。
# 调试和开发 ：在开发和调试模型时，你可能需要在某些特定的代码路径上启用梯度计算，以便检查模型的行为。
# 请注意， torch.enable_grad() 应该在 torch.no_grad() 的上下文管理器之外使用，或者在全局梯度计算被禁用的情况下使用，以确保梯度计算被正确地启用。
# 如果提供了 closure ，则使用 torch.enable_grad() 装饰器来启用梯度计算，因为 closure 需要计算梯度。
with torch.enable_grad():
# 调用 closure 函数并将其返回值赋给 loss 变量。
loss = closure()
# 在 PyTorch 中，优化器（Optimizer）的主要参数通常包括以下几个 ：
# params ：这是优化器构造函数中必须提供的参数，它是一个张量（Tensor）列表，表示需要优化的模型参数。
# 学习率（ lr ） ：这是控制优化过程中参数更新步长的参数。大多数优化器都有一个默认的学习率值，但用户可以在构造函数中指定不同的值。
# 权重衰减（ weight_decay ） ：用于正则化以防止过拟合，它将参数的平方乘以一个系数加到损失函数上，从而惩罚大的权重值。
# 动量参数（ betas 或 momentum ） ：某些优化器（如 Adam、SGD with momentum）使用动量参数来加速梯度下降。动量参数通常表示为一个或两个值的元组，用于控制一阶和二阶矩估计。
# eps ：用于防止分母为零的小值，在某些优化器（如 Adam）中，它会加到分母上以保证数值稳定性。
# alpha ：在某些优化器中， alpha 可以用于控制学习率的缩放。
# schedule ：对于某些优化器，可能包含一个学习率调度器（Learning Rate Scheduler），用于在训练过程中动态调整学习率。
# nesterov ：对于使用 Nesterov 加速梯度的优化器（如 SGD with Nesterov），这个参数用于控制是否启用 Nesterov 动量。
# amsgrad ：在 Adam 优化器中， amsgrad 是一个可选参数，用于启用 AMSGrad 变体，它有助于处理非平稳目标的问题。
# lr_scheduler ：学习率调度器，用于在训练的不同阶段调整学习率。它可以是一个独立的组件，也可以在优化器中直接指定。
# 这些参数的具体含义和作用会根据优化器的不同而有所差异。例如，SGD 优化器可能只使用 lr 和 momentum ，而 Adam 优化器会使用 lr 、 betas 和 eps 。自定义优化器时，可以根据自己的需求添加或修改这些参数。
# 优化器本身并不包含可学习参数，而是包含控制参数更新过程的超参数。
# 优化器的这些超参数不是在训练过程中学习的，而是在训练开始之前设置的，并且可以根据训练的需要进行调整。模型的参数（例如神经网络的权重和偏置）是在训练过程中通过优化器更新的，这些才是可学习参数。优化器的状态（ state ）可能会存储一些中间计算结果，如动量项或梯度的平方的累积值，但这些也不是可学习参数，它们仅用于辅助参数更新的计算。
# 遍历优化器中的所有参数组。
for group in self.param_groups:
# 遍历当前参数组中的所有参数。
for p in group['params']:
# 检查参数 p 是否有梯度。
if p.grad is None:
# 如果没有梯度，则跳过当前参数，继续下一个参数的迭代。
continue
# Perform stepweight decay 执行权重衰减步骤。
# torch.mul(input, other, *, out=None) -> Tensor
# 在 PyTorch 中， torch.mul() 函数用于执行两个张量的元素乘法操作。这个函数会返回一个新的张量，其中包含两个输入张量对应元素相乘的结果。与 mul_() 不同， torch.mul() 不会就地修改输入张量，而是返回一个新的张量。
# 参数 ：
# input (Tensor) ：第一个输入张量。
# other (Tensor 或数值类型) ：第二个输入张量或数值，将与 input 相乘。
# out (Tensor, 可选) ：输出张量，用于存储结果。如果提供，其形状和数据类型必须与结果张量匹配。
# 返回值 ：
# 一个新张量，包含 input 和 other 相乘的结果。
# torch.mul() 是一个非常基础且常用的操作，它在各种张量计算和神经网络的前向传播中扮演着重要角色。例如，在卷积神经网络中，卷积操作就是一种特殊的乘法操作。此外， torch.mul() 也常用于实现学习率的调整、权重的衰减等操作。
# 将参数 p 的数据乘以 (1 - lr * weight_decay) ，实现权重衰减。
p.data.mul_(1 - group['lr'] * group['weight_decay'])
# 获取参数 p 的梯度。
grad = p.grad
# 在 PyTorch 的优化器中，每个参数的状态（ state ）是一个字典，它存储了优化器在优化过程中的中间状态信息。这些信息对于某些优化器算法来说是必要的，因为它们需要在不同的迭代步骤之间保持一些统计数据。
# 优化器的状态（ state ）中可能包含的参数因优化器的不同而异，但以下是一些常见的条目 ：
# step (int) ：表示当前的迭代步数或更新次数。
# exp_avg (Tensor) ：存储梯度的指数加权平均值，常用于具有动量的优化器（如 SGD with momentum, Adam）。
# exp_avg_sq (Tensor) ：存储梯度平方的指数加权平均值，用于计算自适应学习率（如 Adam 优化器）。
# exp_inf_norm (Tensor) ：在某些优化器中，存储梯度的无穷范数的指数加权平均值。
# max_exp_avg_sq (Tensor) ：在 Adam 优化器的 AMSGrad 变体中使用，存储到目前为止遇到的最大的 exp_avg_sq 。
# buffer (Tensor 或 list of Tensors) ：用于存储额外的中间状态信息，具体取决于优化器的实现。
# bits (Tensor) ：在某些优化器中，存储参数的位表示，用于量化感知优化。
# sign_buffer (Tensor) ：用于存储梯度符号的缓冲区，用于某些优化器的更新规则。
# inertia (Tensor) ：在某些优化器中，存储用于模拟惯性的中间状态。
# s (Tensor) ：在某些优化器中，存储梯度的累积值。
# 这些状态信息通常在优化器的 step 方法中更新，并在每次迭代中用于计算参数的更新。需要注意的是，不是所有的优化器都会使用上述所有状态信息，而且不同的优化器可能会有自己特定的状态信息。例如，SGD 优化器可能只需要 step 和 momentum 相关的信息，而 Adam 优化器则需要 exp_avg 和 exp_avg_sq 。
# torch.optim.Optimizer.step()
# 在 PyTorch 中， step() 函数是优化器（Optimizer）类的一个方法，用于执行单步参数更新。这个方法会根据优化器内部的参数和学习率调度器来更新模型的参数。
# 参数说明 ：无参数。
# 返回值 ：无返回值。
# step() 方法通常在每次迭代后调用，它会对所有之前通过 zero_grad() 方法清零的参数进行梯度更新。在调用 step() 之前，需要确保已经计算了模型的梯度（通过反向传播）并且清零了之前的梯度。
# 需要注意的是， step() 方法不应该在没有梯度的情况下被调用，因为这可能会导致错误或未定义的行为。
# 此外，如果你使用了学习率调度器（LR Scheduler），通常在 step() 调用之后立即更新学习率 ：
# scheduler.step() # 在每个 epoch 结束时更新学习率。
# 这里的 scheduler.step() 会根据学习率调度器的策略更新学习率，为下一个 epoch 的训练做好准备。
# 获取与参数 p 相关的状态字典。
state = self.state[p]
# State initialization 状态初始化。
# 检查状态字典是否为空。
if len(state) == 0:
# Exponential moving average of gradient values 梯度值的指数移动平均。
# torch.zeros_like(input, dtype=None, layout=torch.strided, device=None, requires_grad=False)
# 在 PyTorch 中， torch.zeros_like() 函数用于创建一个与给定张量 input 形状相同、数据类型相同的新张量，并且新张量中的所有元素都初始化为0。
# 参数 ：
# input (Tensor) ：一个已有的张量，新创建的张量将具有相同的形状和数据类型。
# dtype (torch.dtype, 可选) ：新张量的数据类型，默认为 None ，这意味着新张量将具有与 input 相同的数据类型。
# layout (torch.layout, 可选) ：新张量的布局，默认为 torch.strided 。
# device (torch.device, 可选) ：新张量所在的设备，默认为 None ，这意味着新张量将被创建在与 input 相同的设备上。
# requires_grad (bool, 可选) ：新张量是否需要梯度计算，默认为 False 。
# 返回值 ：
# 一个与 input 形状和数据类型相同，所有元素都为0的新张量。
# 在优化算法的上下文中， state['exp_avg'] 通常指的是“指数移动平均”（Exponential Moving Average）的缩写，它是一种计算运行平均值的方法，特别适用于需要平滑处理时间序列数据的场景。在优化算法中， exp_avg 通常用来存储梯度的指数移动平均值。
# 在许多优化算法中，如 Adam 优化器，会使用两个主要的指数移动平均值 ：
# 梯度的指数移动平均（ exp_avg ） ：这个平均值用来估计参数的一阶矩（即均值），它可以帮助优化器调整步长，并减少梯度的噪声影响。
# 梯度平方的指数移动平均（ exp_avg_sq ） ：这个平均值用来估计参数的二阶矩（即未中心化的方差），它可以帮助优化器调整不同参数的步长，以适应不同参数的更新频率。
# state['exp_avg'] 中的 state 是一个字典，它存储了与每个参数相关的优化状态信息。每个参数可能有自己的 exp_avg 值，因此 state 字典会为每个参数维护一个条目。
# 在优化过程中，每次迭代都会更新这些指数移动平均值，以反映最新的梯度信息。
# 以下是在 Adam 优化器中如何初始化和更新 exp_avg 的一个简单示例：
# for p in group['params']:
# if p.grad is not None:
# grad = p.grad.data
# if 'exp_avg' not in self.state[p]:
# self.state[p]['exp_avg'] = torch.zeros_like(p.data)
# self.state[p]['exp_avg'].mul_(beta1).add_(1 - beta1, grad)
# 在这个示例中 ：
# beta1 是指数移动平均的衰减率。
# self.state[p]['exp_avg'] 是存储参数 p 的梯度指数移动平均值的张量。
# mul_(beta1) 将当前的 exp_avg 值乘以 beta1 。
# add_(1 - beta1, grad) 将更新项 (1 - beta1) * grad 加到 exp_avg 上，以反映最新的梯度信息。
# 通过这种方式， exp_avg 能够捕捉到梯度的历史信息，并为优化器提供一种平滑的梯度估计，这有助于提高训练的稳定性和收敛速度。
# 如果状态字典为空，则初始化 exp_avg （指数移动平均）为与 p 形状相同且所有元素为0的张量。
state['exp_avg'] = torch.zeros_like(p)
# 获取 exp_avg 。
exp_avg = state['exp_avg']
# 从参数组中获取 beta1 和 beta2 。
beta1, beta2 = group['betas']
# Weight update 权重更新。
# 计算更新项，它是 exp_avg 乘以 beta1 加上梯度 grad 乘以 (1 - beta1) 。
update = exp_avg * beta1 + grad * (1 - beta1)
# torch.sign(input, *, out=None) -> Tensor
# 在 PyTorch 中， torch.sign() 函数用于计算输入张量中每个元素的符号。符号函数返回一个与输入形状相同的新张量，其中每个元素的值是 ： -1 如果输入元素是负数， 0 如果输入元素是零， 1 如果输入元素是正数。
# 参数 ：
# input (Tensor) ：输入张量，可以是任意数值类型的张量。
# out (Tensor, 可选) ：输出张量，用于存储结果。如果提供，其形状和数据类型必须与输入张量匹配。
# 返回值 ：
# 一个与输入张量形状相同，元素为 -1 、 0 或 1 的新张量。
# torch.add(input, other, *, alpha=1, out=None) -> Tensor
# torch.add() 函数在 PyTorch 中用于执行两个张量的加法操作。这个函数会返回一个新的张量，其中包含两个输入张量对应元素相加的结果。与 add_() 不同， torch.add() 不会就地修改输入张量，而是返回一个新的张量。
# 参数 :
# input (Tensor) ：第一个输入张量。
# other (Tensor 或数值类型) ：第二个输入张量或数值，将与 input 相加。
# alpha (float, 可选) ：标量乘数，用于 other ， alpha * other 会被加到 input 上，默认为 1 。
# out (Tensor, 可选) ：输出张量，用于存储结果。如果提供，其形状和数据类型必须与结果张量匹配。
# 返回值 ：
# 一个新张量，包含 input 和 other 相加的结果。
# torch.add() 是一个非常基础且常用的操作，它在各种张量计算和神经网络的前向传播中扮演着重要角色。
# add_() 是一个就地（in-place）操作函数，用于将一个值或张量加到原张量上，并更新原张量。这个函数是 torch.add() 的就地版本，意味着它直接修改调用它的张量，而不是返回一个新的张量。
# 就地操作可以节省内存，因为它们不需要分配新的内存来存储结果，但这也意味着原始数据会被覆盖。因此，如果你需要保留原始数据，应该先复制原张量，或者使用非就地版本的函数，如 torch.add() 。
# 使用 update 的符号和学习率更新参数 p 。这里使用 add_ 方法将 torch.sign(update) 乘以 -group['lr'] 加到参数 p 上。
p.add_(torch.sign(update), alpha=-group['lr'])
# Decay the momentum running average coefficient 衰减动量运行平均系数。
# 更新 exp_avg ，首先乘以 beta2 ，然后加上 grad 乘以 (1 - beta2) 。
exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2)
# 返回计算的损失值。
return loss
# step 方法首先检查是否提供了 closure 来计算损失，然后对每个参数组中的每个参数执行权重衰减和更新步骤。更新步骤包括计算指数移动平均的梯度值，然后使用这个值和梯度的符号来更新参数。这个方法是优化器的核心，负责在每次迭代中调整模型的参数。
# 这个 Lion 优化器的实现是一个简化版本，它结合了动量和权重衰减的概念，但更新规则与标准的Adam或SGD不同，因为它使用了 torch.sign(update) 来确定更新的方向。这种实现可能需要进一步的调整和测试以确保其性能和稳定性。
# 这个优化器的更新规则与 Adam 优化器不同，因为它使用了梯度的符号而不是梯度本身，这可能会导致不同的优化行为。此外，代码中没有对 exp_avg 进行初始化为零以外的值，这可能会导致优化器在第一次迭代时表现不佳。通常，我们会使用小的随机值来初始化动量项，以避免在训练初期对参数更新产生过大的影响。
# 注意事项 ：
# 这段代码中的 Lion 优化器实现与之前描述的 Lion 优化器的特性有所不同。特别是，它使用了类似于Adam优化器的指数移动平均和权重更新策略，但更新规则有所不同（使用 torch.sign(update) ）。
# 代码中的 torch.no_grad() 装饰器确保在优化步骤中不计算梯度，这有助于减少内存消耗和计算成本。
# torch.enable_grad() 在计算 closure 时启用梯度计算。
# 这段代码提供了一个基本的 Lion 优化器实现，但可能需要进一步调整和测试以确保其性能和稳定性。
# 在 PyTorch 中自定义优化器时，通常会继承自 torch.optim.Optimizer 类，并实现以下方法 ：
# __init__ 方法 ：
# 这是类的构造函数，你需要在这里初始化优化器的状态，包括参数组（ param_groups ）和任何需要的超参数。通常，会在这里调用父类的 __init__ 方法，并传递参数组。
# step 方法 ：
# 这是执行单步参数更新的核心方法。在这个方法中，需要实现参数的更新逻辑，包括计算梯度、应用更新规则等。这个方法通常在每次迭代后被调用。
# zero_grad 方法（可选） ：
# 这个方法用于清除所有参数的梯度。虽然这不是必须实现的，但通常建议提供这个方法，以便在每次迭代开始时清除梯度，避免梯度累积。
# 除了这些方法， Optimizer 类还提供了一些其他方法和属性，它们可以被重写或扩展，以实现更复杂的优化逻辑 ：
# add_param_group 方法 ：用于在优化过程中动态添加参数组。
# state 和 param_groups 属性 ：这些属性存储了优化器的状态和参数组信息，可以在自定义优化器中使用或修改它们。
# load_state_dict 和 state_dict 方法 ：用于保存和加载优化器的状态。
# 以下是一个简单的自定义优化器的框架，展示了必须实现的方法 ：
'''
import torch
class CustomOptimizer(torch.optim.Optimizer):
def __init__(self, params, lr=1e-3):
defaults = dict(lr=lr)
super(CustomOptimizer, self).__init__(params, defaults)
def step(self, closure=None):
# 清除梯度（可选）
for group in self.param_groups:
for p in group['params']:
if p.grad is not None:
p.grad.detach_()
p.grad.zero_()
# 执行参数更新
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
# 在这里实现参数更新逻辑
grad = p.grad
state = self.state[p]
# 状态初始化（可选）
if len(state) == 0:
state['step'] = 0
# 可以在这里初始化更多的状态变量
state['step'] += 1
# 参数更新逻辑
p.data.add_(-group['lr'] * grad)
def zero_grad(self):
for group in self.param_groups:
for p in group['params']:
if p.grad is not None:
p.grad.detach_()
p.grad.zero_()
# 使用自定义优化器
model = ...
optimizer = CustomOptimizer(model.parameters(), lr=0.01)
'''
# 在这个例子中， CustomOptimizer 类实现了 __init__ 和 step 方法，并且提供了一个可选的 zero_grad 方法。这是自定义优化器的基本结构，你可以根据需要扩展更多的功能。