torchrec.optim

Torchrec 优化器

Torchrec 包含一个名为 KeyedOptimizer 的特殊优化器。KeyedOptimizer 公开了 state_dict，其中包含有意义的键 - 它允许在适当的位置加载 torch.tensor 和 ShardedTensor，并且它禁止将空状态加载到已初始化的 KeyedOptimizer 中，反之亦然。

它还包含 - 包裹 KeyedOptimizer 的几个模块，称为 CombinedOptimizer 和 OptimizerWrapper - 用于 RecSys 的优化器：例如，rowwise adagrad/adam 等

torchrec.optim.clipping

class torchrec.optim.clipping.GradientClipping(value)

基类：Enum

枚举。

NONE = 'none'

NORM = 'norm'

VALUE = 'value'

class torchrec.optim.clipping.GradientClippingOptimizer(optimizer: KeyedOptimizer, clipping: GradientClipping = GradientClipping.NONE, max_gradient: float = 0.1, norm_type: Union[float, str] = 2.0)

基类：OptimizerWrapper

在执行优化步骤之前剪裁梯度。

参数：

• optimizer (KeyedOptimizer) – 要包装的优化器

• clipping (GradientClipping) – 如何剪裁梯度

• max_gradient (float) – 剪裁的最大值

• norm_type (float 或 str) – 使用的 p 范数类型。可以是 'inf' 用于无穷范数。

step(closure: Optional[Any] = None) → None

执行单次优化步骤以更新参数。

参数：

closure (Callable) – 重新评估模型并返回损失的闭包。对于大多数优化器来说是可选的。

注意

除非另有说明，否则此函数不应修改参数的 .grad 字段。

torchrec.optim.fused

class torchrec.optim.fused.EmptyFusedOptimizer

基类：FusedOptimizer

融合优化器类，具有无操作步骤，没有可优化参数

step(closure: Optional[Any] = None) → None

执行单次优化步骤以更新参数。

参数：

closure (Callable) – 重新评估模型并返回损失的闭包。对于大多数优化器来说是可选的。

注意

除非另有说明，否则此函数不应修改参数的 .grad 字段。

zero_grad(set_to_none: bool = False) → None

重置所有经过优化的 torch.Tensor 的梯度。

参数：

set_to_none (bool) – 代替设置为零，将梯度设置为 None。这通常具有更低的内存占用，并且可以适度提高性能。但是，它会改变某些行为。例如：1. 当用户尝试访问梯度并在其上执行手动操作时，None 属性或 Tensor 满 0s 的行为将有所不同。2. 如果用户请求 zero_grad(set_to_none=True) 然后进行反向传播，则对于没有收到梯度的参数，.grad 保证为 None。3. torch.optim 优化器在梯度为 0 或 None 时具有不同的行为（在一种情况下，它会使用梯度 0 执行步骤，在另一种情况下，它会完全跳过步骤）。

class torchrec.optim.fused.FusedOptimizer(params: Mapping[str, Union[Tensor, ShardedTensor]], state: Mapping[Any, Any], param_groups: Collection[Mapping[str, Any]])

继承自: KeyedOptimizer, ABC

假设权重更新在反向传播过程中完成，因此 step() 是一个无操作函数。

abstract step(closure: Optional[Any] = None) → None

执行单次优化步骤以更新参数。

参数：

closure (Callable) – 重新评估模型并返回损失的闭包。对于大多数优化器来说是可选的。

注意

除非另有说明，否则此函数不应修改参数的 .grad 字段。

abstract zero_grad(set_to_none: bool = False) → None

重置所有经过优化的 torch.Tensor 的梯度。

参数：

set_to_none (bool) – 代替设置为零，将梯度设置为 None。这通常具有更低的内存占用，并且可以适度提高性能。但是，它会改变某些行为。例如：1. 当用户尝试访问梯度并在其上执行手动操作时，None 属性或 Tensor 满 0s 的行为将有所不同。2. 如果用户请求 zero_grad(set_to_none=True) 然后进行反向传播，则对于没有收到梯度的参数，.grad 保证为 None。3. torch.optim 优化器在梯度为 0 或 None 时具有不同的行为（在一种情况下，它会使用梯度 0 执行步骤，在另一种情况下，它会完全跳过步骤）。

class torchrec.optim.fused.FusedOptimizerModule

继承自: ABC

模块，在反向传播过程中执行权重更新。

abstract property fused_optimizer: KeyedOptimizer

torchrec.optim.keyed

class torchrec.optim.keyed.CombinedOptimizer(optims: List[Union[KeyedOptimizer, Tuple[str, KeyedOptimizer]]])

继承自: KeyedOptimizer

将多个 KeyedOptimizers 合并为一个。

用于将不同子模块的不同优化器组合在一起。

property optimizers: List[Tuple[str, KeyedOptimizer]]

property param_groups: Collection[Mapping[str, Any]]

property params: Mapping[str, Union[Tensor, ShardedTensor]]

post_load_state_dict() → None

static prepend_opt_key(name: str, opt_key: str) → str

save_param_groups(save: bool) → None

property state: Mapping[Tensor, Any]

step(closure: Optional[Any] = None) → None

执行单次优化步骤以更新参数。

参数：

closure (Callable) – 重新评估模型并返回损失的闭包。对于大多数优化器来说是可选的。

注意

除非另有说明，否则此函数不应修改参数的 .grad 字段。

zero_grad(set_to_none: bool = False) → None

重置所有经过优化的 torch.Tensor 的梯度。

参数：

set_to_none (bool) – 代替设置为零，将梯度设置为 None。这通常具有更低的内存占用，并且可以适度提高性能。但是，它会改变某些行为。例如：1. 当用户尝试访问梯度并在其上执行手动操作时，None 属性或 Tensor 满 0s 的行为将有所不同。2. 如果用户请求 zero_grad(set_to_none=True) 然后进行反向传播，则对于没有收到梯度的参数，.grad 保证为 None。3. torch.optim 优化器在梯度为 0 或 None 时具有不同的行为（在一种情况下，它会使用梯度 0 执行步骤，在另一种情况下，它会完全跳过步骤）。

class torchrec.optim.keyed.KeyedOptimizer(params: Mapping[str, Union[Tensor, ShardedTensor]], state: Mapping[Any, Any], param_groups: Collection[Mapping[str, Any]])

继承自: Optimizer

接收一个参数字典，并通过参数键公开 state_dict。

此实现比 torch.Optimizer 中的实现严格得多：它要求实现完全初始化其状态在第一次优化迭代期间，并且它禁止将空状态加载到已初始化的 KeyedOptimizer 中，反之亦然。

它也不公开 state_dict() 中的 param_groups 默认情况下，可以通过设置 save_param_groups 标志来启用旧行为。原因是，在分布式训练期间，并非所有参数都存在于所有排名上，我们通过其参数识别 param_group。此外，param_groups 通常在训练初始化期间重新设置，因此将它们作为状态的一部分保存几乎没有意义。

add_param_group(param_group: Any) → None

向 Optimizer 的 param_groups 中添加一个参数组。

当微调预训练网络时，这可能很有用，因为冻结层可以被设置为可训练，并在训练过程中添加到 Optimizer 中。

参数：

param_group (dict) – 指定应与组特定的优化选项一起优化的张量。

init_state(sparse_grad_parameter_names: Optional[Set[str]] = None) → None

运行一个虚拟的优化器步骤，这允许初始化优化器状态，通常是延迟的。这使我们能够从检查点中进行优化器状态的原位加载。

load_state_dict(state_dict: Mapping[str, Any]) → None

此实现比 torch.Optimizer 中的实现严格得多：它要求实现完全初始化其状态在第一次优化迭代期间，并且它禁止将空状态加载到已初始化的 KeyedOptimizer 中，反之亦然。

由于引入了严格性，它允许我们

• 对状态和 param_groups 进行兼容性检查，从而提高可用性

• 通过直接复制到状态张量中来避免状态重复，例如 optimizer.step() # 确保优化器已初始化 sd = optimizer.state_dict() load_checkpoint(sd) # 将状态直接复制到张量中，如果需要则重新分片 optimizer.load_state_dict(sd) # 替换 param_groups

post_load_state_dict() → None

save_param_groups(save: bool) → None

state_dict() → Dict[str, Any]

返回的状态和 param_groups 将包含参数键而不是 torch.Optimizer 中的参数索引。这允许实现诸如优化器重新分片之类的更高级的功能。

还可以处理遵循 PyTorch 有状态协议的类和支持的数据结构。

class torchrec.optim.keyed.KeyedOptimizerWrapper(params: Mapping[str, Union[Tensor, ShardedTensor]], optim_factory: Callable[[List[Union[Tensor, ShardedTensor]]], Optimizer])

继承自: KeyedOptimizer

接收一个参数字典，并通过参数键公开 state_dict。

方便的包装器，用于接收 optim_factory 可调用函数以创建 KeyedOptimizer

step(closure: Optional[Any] = None) → None

执行单次优化步骤以更新参数。

参数：

closure (Callable) – 重新评估模型并返回损失的闭包。对于大多数优化器来说是可选的。

注意

除非另有说明，否则此函数不应修改参数的 .grad 字段。

zero_grad(set_to_none: bool = False) → None

重置所有经过优化的 torch.Tensor 的梯度。

参数：

set_to_none (bool) – 代替设置为零，将梯度设置为 None。这通常具有更低的内存占用，并且可以适度提高性能。但是，它会改变某些行为。例如：1. 当用户尝试访问梯度并在其上执行手动操作时，None 属性或 Tensor 满 0s 的行为将有所不同。2. 如果用户请求 zero_grad(set_to_none=True) 然后进行反向传播，则对于没有收到梯度的参数，.grad 保证为 None。3. torch.optim 优化器在梯度为 0 或 None 时具有不同的行为（在一种情况下，它会使用梯度 0 执行步骤，在另一种情况下，它会完全跳过步骤）。

class torchrec.optim.keyed.OptimizerWrapper(optimizer: KeyedOptimizer)

继承自: KeyedOptimizer

包装器，它接收一个 KeyedOptimizer 并且是 KeyedOptimizer

用于像 GradientClippingOptimizer 和 WarmupOptimizer 这样的优化器的子类

add_param_group(param_group: Any) → None

向 Optimizer 的 param_groups 中添加一个参数组。

当微调预训练网络时，这可能很有用，因为冻结层可以被设置为可训练，并在训练过程中添加到 Optimizer 中。

参数：

param_group (dict) – 指定应与组特定的优化选项一起优化的张量。

load_state_dict(state_dict: Mapping[str, Any]) → None

此实现比 torch.Optimizer 中的实现严格得多：它要求实现完全初始化其状态在第一次优化迭代期间，并且它禁止将空状态加载到已初始化的 KeyedOptimizer 中，反之亦然。

由于引入了严格性，它允许我们

• 对状态和 param_groups 进行兼容性检查，从而提高可用性

• 通过直接复制到状态张量中来避免状态重复，例如 optimizer.step() # 确保优化器已初始化 sd = optimizer.state_dict() load_checkpoint(sd) # 将状态直接复制到张量中，如果需要则重新分片 optimizer.load_state_dict(sd) # 替换 param_groups

post_load_state_dict() → None

save_param_groups(save: bool) → None

state_dict() → Dict[str, Any]

返回的状态和 param_groups 将包含参数键而不是 torch.Optimizer 中的参数索引。这允许实现诸如优化器重新分片之类的更高级的功能。

还可以处理遵循 PyTorch 有状态协议的类和支持的数据结构。

step(closure: Optional[Any] = None) → None

执行单次优化步骤以更新参数。

参数：

closure (Callable) – 重新评估模型并返回损失的闭包。对于大多数优化器来说是可选的。

注意

除非另有说明，否则此函数不应修改参数的 .grad 字段。

zero_grad(set_to_none: bool = False) → None

重置所有经过优化的 torch.Tensor 的梯度。

参数：

set_to_none (bool) – 代替设置为零，将梯度设置为 None。这通常具有更低的内存占用，并且可以适度提高性能。但是，它会改变某些行为。例如：1. 当用户尝试访问梯度并在其上执行手动操作时，None 属性或 Tensor 满 0s 的行为将有所不同。2. 如果用户请求 zero_grad(set_to_none=True) 然后进行反向传播，则对于没有收到梯度的参数，.grad 保证为 None。3. torch.optim 优化器在梯度为 0 或 None 时具有不同的行为（在一种情况下，它会使用梯度 0 执行步骤，在另一种情况下，它会完全跳过步骤）。

torchrec.optim.warmup

class torchrec.optim.warmup.WarmupOptimizer(optimizer: KeyedOptimizer, stages: List[WarmupStage], lr: float = 0.1, lr_param: str = 'lr', param_name: str = '__warmup')

基类：OptimizerWrapper

根据计划调整学习率。

参数：

• optimizer (KeyedOptimizer) – 要包装的优化器

• stages (List[WarmupStage]) – 要经历的阶段

• lr (float) – 初始学习率

• lr_param (str) – 参数组中的学习率参数。

• param_name – 用于保存预热状态的伪参数的名称。

post_load_state_dict() → None

step(closure: Optional[Any] = None) → None

执行单次优化步骤以更新参数。

参数：

closure (Callable) – 重新评估模型并返回损失的闭包。对于大多数优化器来说是可选的。

注意

除非另有说明，否则此函数不应修改参数的 .grad 字段。

class torchrec.optim.warmup.WarmupPolicy(value)

基类：Enum

枚举。

CONSTANT = 'constant'

COSINE_ANNEALING_WARM_RESTARTS = 'cosine_annealing_warm_restarts'

INVSQRT = 'inv_sqrt'

LINEAR = 'linear'

NONE = 'none'

POLY = 'poly'

STEP = 'step'

class torchrec.optim.warmup.WarmupStage(policy: torchrec.optim.warmup.WarmupPolicy = <WarmupPolicy.LINEAR: 'linear'>, max_iters: int = 1, value: float = 1.0, lr_scale: float = 1.0, decay_iters: int = -1, sgdr_period: int = 1)

Bases: object

decay_iters: int = -1

lr_scale: float = 1.0

max_iters: int = 1

policy: WarmupPolicy = 'linear'

sgdr_period: int = 1

value: float = 1.0

模块内容

Torchrec 优化器

Torchrec 包含一个名为 KeyedOptimizer 的特殊优化器。KeyedOptimizer 公开了 state_dict，其中包含有意义的键 - 它允许在适当的位置加载 torch.tensor 和 ShardedTensor，并且它禁止将空状态加载到已初始化的 KeyedOptimizer 中，反之亦然。

它还包含 - 包裹 KeyedOptimizer 的几个模块，称为 CombinedOptimizer 和 OptimizerWrapper - 用于 RecSys 的优化器：例如，rowwise adagrad/adam 等