此文件包含用于确定性和概率性 PyTorch Lightning 模块的抽象类

class darts.models.forecasting.pl_forecasting_module.PLDualCovariatesModule(input_chunk_length, output_chunk_length, output_chunk_shift=0, train_sample_shape=None, loss_fn=MSELoss(), torch_metrics=None, likelihood=None, optimizer_cls=<class 'torch.optim.adam.Adam'>, optimizer_kwargs=None, lr_scheduler_cls=None, lr_scheduler_kwargs=None, use_reversible_instance_norm=False)

基类：PLForecastingModule, ABC

基于 PyTorch Lightning 的预测模块。

这个类旨在通过继承来创建一个新的基于 PyTorch Lightning 的预测模块。当子类化这个类时，请确保添加以下具有给定签名的方法：

• PLForecastingModule.__init__()

• PLForecastingModule.forward()

• PLForecastingModule._produce_train_output()

• PLForecastingModule._get_batch_prediction()

在子类 MyModel 的 __init__() 函数中调用 super(MyModel, self).__init__(**kwargs)，其中 kwargs 是 PLForecastingModule 的参数。

参数

• input_chunk_length (int) – 过去作为模型输入的时间步数（每个块）。适用于目标序列，以及过去和/或未来协变量（如果模型支持）。

• output_chunk_length (int) – 内部模型一次预测的时间步数（每个块）。此外，模型输入中使用的未来协变量的未来值数量（如果模型支持未来协变量）。这与 predict() 中使用的预测范围 n 不同，后者是使用一次性或自回归预测生成的所需预测点数。设置 n <= output_chunk_length 可以防止自回归。当协变量不足以延伸到未来，或禁止模型使用过去和/或未来协变量的未来值进行预测时（取决于模型的协变量支持），这很有用。

• train_sample_shape (Optional[Tuple, None]) – 模型的输入形状，用于在不调用 fit_from_dataset 的情况下实例化模型，并对用于重新训练或预测的新训练/推理数据集进行完整性检查。

• loss_fn (_Loss) – 用于训练的 PyTorch 损失函数。如果指定了 likelihood 参数，则此参数将被概率模型忽略。默认值：torch.nn.MSELoss()。

• torch_metrics (Union[Metric, MetricCollection, None]) – 用于评估的 torch 指标或 MetricCollection。可用指标的完整列表可以在 https://torchmetrics.readthedocs.io/en/latest/ 找到。默认值：None。

• likelihood (Optional[Likelihood, None]) – Darts 的 似然模型 之一，用于概率预测。默认值：None。

• optimizer_cls (Optimizer) – 要使用的 PyTorch 优化器类。默认值：torch.optim.Adam。

• optimizer_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 优化器的键值参数（例如，{'lr': 1e-3} 用于指定学习率）。否则将使用所选 optimizer_cls 的默认值。默认值：None。

• lr_scheduler_cls (Optional[_LRScheduler, None]) – 可选地，指定要使用的 PyTorch 学习率调度器类。指定 None 对应于使用恒定的学习率。默认值：None。

• lr_scheduler_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 学习率调度器的键值参数。默认值：None。

• use_reversible_instance_norm (bool) – 是否使用可逆实例归一化 RINorm 来对抗分布偏移，如 [1] 所示。它仅应用于目标序列的特征，而不应用于协变量。

引用

1

T. Kim et al. “Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift”, https://openreview.net/forum?id=cGDAkQo1C0p

属性

automatic_optimization	如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。
current_epoch	Trainer 中的当前时期，如果未附加则为 0。
device_mesh	像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。
example_input_array	示例输入数组是模块在 forward() 方法中可以消费的规范。
global_rank	当前进程在所有节点和设备中的索引。
global_step	所有epoch中看到的总训练批次。
hparams	使用 save_hyperparameters() 保存的超参数集合。
hparams_initial	使用 save_hyperparameters() 保存的超参数集合。
local_rank	当前进程在单个节点中的索引。
logger	Trainer 中对 logger 对象的引用。
loggers	参考 Trainer 中的日志记录器列表。
on_gpu	如果该模型当前位于GPU上，则返回 True。
output_chunk_length	模型一次预测的时间步数。
strict_loading	确定 Lightning 如何使用 .load_state_dict(..., strict=model.strict_loading) 加载此模型。

设备
dtype
已训练的轮数
fabric
supports_probabilistic_prediction
训练者

方法

add_module(name, module)	将一个子模块添加到当前模块。
all_gather(data[, group, sync_grads])	从多个进程中收集张量或张量集合。
apply(fn)	递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。
backward(loss, *args, **kwargs)	在 training_step() 中返回的损失上调用以执行反向传播。
bfloat16()	将所有浮点参数和缓冲区转换为 bfloat16 数据类型。
buffers([recurse])	返回一个遍历模块缓冲区的迭代器。
children()	返回一个遍历直接子模块的迭代器。
clip_gradients(optimizer[, ...])	内部处理梯度裁剪。
compile(*args, **kwargs)	使用 torch.compile() 编译此模块的 forward。
configure_callbacks()	配置模型特定的回调。
configure_gradient_clipping(optimizer[, ...])	对优化器参数执行梯度裁剪。
configure_model()	在策略和精度感知的环境中创建模块的钩子。
configure_optimizers()	配置优化器和学习率调度器以优化模型。
configure_sharded_model()	已弃用。
configure_torch_metrics(torch_metrics)	处理 torch_metrics 参数。
cpu()	参见 torch.nn.Module.cpu()。
cuda([device])	将所有模型参数和缓冲区移动到GPU。
double()	参见 torch.nn.Module.double()。
eval()	将模块设置为评估模式。
extra_repr()	设置模块的额外表示。
float()	参见 torch.nn.Module.float()。
forward(*args, **kwargs)	与 torch.nn.Module.forward() 相同。
freeze()	冻结所有参数以进行推理。
get_buffer(target)	如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。
get_extra_state()	返回要包含在模块的 state_dict 中的任何额外状态。
get_parameter(target)	如果存在，则返回由 target 给出的参数，否则抛出错误。
get_submodule(target)	如果存在，返回由 target 指定的子模块，否则抛出错误。
half()	参见 torch.nn.Module.half()。
ipu([device])	将所有模型参数和缓冲区移动到 IPU。
load_from_checkpoint(checkpoint_path[, ...])	从检查点加载模型的主要方式。
load_state_dict(state_dict[, strict, assign])	将参数和缓冲区从 state_dict 复制到此模块及其后代中。
log(name, value[, prog_bar, logger, ...])	记录一个键值对。
log_dict(dictionary[, prog_bar, logger, ...])	一次性记录一个值的字典。
lr_scheduler_step(scheduler, metric)	重写此方法以调整 Trainer 调用每个调度器的默认方式。
lr_schedulers()	返回在训练期间使用的学习率调度器。
manual_backward(loss, *args, **kwargs)	在进行手动优化时，直接从 training_step() 调用此方法。
modules()	返回网络中所有模块的迭代器。
named_buffers([prefix, recurse, ...])	返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。
named_children()	返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。
named_modules([memo, prefix, remove_duplicate])	返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。
named_parameters([prefix, recurse, ...])	返回一个模块参数的迭代器，同时生成参数的名称和参数本身。
on_after_backward()	在 loss.backward() 之后和优化器步进之前调用。
on_after_batch_transfer(batch, dataloader_idx)	覆盖此方法以在批次转移到设备后对其进行批量增强或修改。
on_before_backward(loss)	在 loss.backward() 之前调用。
on_before_batch_transfer(batch, dataloader_idx)	在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。
on_before_optimizer_step(optimizer)	在 optimizer.step() 之前调用。
on_before_zero_grad(optimizer)	在 training_step() 之后和 optimizer.zero_grad() 之前调用。
on_fit_end()	在拟合的最后调用。
on_fit_start()	在拟合开始时调用。
on_load_checkpoint(checkpoint)	由 Lightning 调用以恢复您的模型。
on_predict_batch_end(outputs, batch, batch_idx)	在批次之后调用预测循环。
on_predict_batch_start(batch, batch_idx[, ...])	在预测循环中，在该批次发生任何事情之前调用。
on_predict_end()	在预测结束时调用。
on_predict_epoch_end()	在预测结束时调用。
on_predict_epoch_start()	在预测开始时调用。
on_predict_model_eval()	在预测循环开始时调用。
on_predict_start()	在预测开始时调用。
on_save_checkpoint(checkpoint)	当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。
on_test_batch_end(outputs, batch, batch_idx)	在批次之后在测试循环中调用。
on_test_batch_start(batch, batch_idx[, ...])	在测试循环中，在该批次发生任何事情之前调用。
on_test_end()	在测试结束时调用。
on_test_epoch_end()	在每个epoch的测试循环结束时调用。
on_test_epoch_start()	在epoch开始时在测试循环中调用。
on_test_model_eval()	在测试循环开始时调用。
on_test_model_train()	在测试循环结束时调用。
on_test_start()	在测试开始时调用。
on_train_batch_end(outputs, batch, batch_idx)	在训练循环中批次处理后调用。
on_train_batch_start(batch, batch_idx)	在训练循环中，在该批次发生任何事情之前调用。
on_train_end()	在训练结束前调用，在记录器实验关闭之前。
on_train_epoch_end()	在训练循环的每个epoch结束时调用。
on_train_epoch_start()	在训练循环中，在epoch开始时调用。
on_train_start()	在完整性检查后，训练开始时调用。
on_validation_batch_end(outputs, batch, ...)	在批次验证循环后调用。
on_validation_batch_start(batch, batch_idx)	在验证循环中，对该批次进行任何操作之前调用。
on_validation_end()	在验证结束时调用。
on_validation_epoch_end()	在epoch结束时的验证循环中调用。
on_validation_epoch_start()	在epoch开始时验证循环中调用。
on_validation_model_eval()	在验证循环开始时调用。
on_validation_model_train()	当验证循环结束时调用。
on_validation_model_zero_grad()	在进入验证循环之前，由训练循环调用来释放梯度。
on_validation_start()	在验证开始时调用。
optimizer_step(epoch, batch_idx, optimizer)	重写此方法以调整 Trainer 调用优化器的默认方式。
optimizer_zero_grad(epoch, batch_idx, optimizer)	重写此方法以改变 optimizer.zero_grad() 的默认行为。
optimizers([use_pl_optimizer])	返回训练期间使用的优化器。
parameters([recurse])	返回一个遍历模块参数的迭代器。
predict_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。
predict_step(batch, batch_idx[, dataloader_idx])	执行预测步骤
prepare_data()	使用此功能下载并准备数据。
print(*args, **kwargs)	仅从进程 0 打印。
register_backward_hook(hook)	在模块上注册一个反向钩子。
register_buffer(name, tensor[, persistent])	向模块添加一个缓冲区。
register_forward_hook(hook, *[, prepend, ...])	在模块上注册一个前向钩子。
register_forward_pre_hook(hook, *[, ...])	在模块上注册一个前向预钩子。
register_full_backward_hook(hook[, prepend])	在模块上注册一个反向钩子。
register_full_backward_pre_hook(hook[, prepend])	在模块上注册一个反向预钩子。
register_load_state_dict_post_hook(hook)	注册一个在模块的 load_state_dict 调用后运行的后置钩子。
register_module(name, module)	别名 add_module()。
register_parameter(name, param)	向模块添加一个参数。
register_state_dict_pre_hook(hook)	为 state_dict() 方法注册一个预钩子。
requires_grad_([requires_grad])	更改是否应在此模块的参数上记录自动梯度操作。
save_hyperparameters(*args[, ignore, frame, ...])	将参数保存到 hparams 属性中。
set_extra_state(state)	设置包含在加载的 state_dict 中的额外状态。
set_predict_parameters(n, num_samples, ...)	在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。
setup(stage)	在 fit（训练 + 验证）、验证、测试或预测开始时调用。
share_memory()	参见 torch.Tensor.share_memory_()。
state_dict(*args[, destination, prefix, ...])	返回一个包含模块整体状态引用的字典。
teardown(stage)	在 fit（训练 + 验证）、验证、测试或预测结束时调用。
test_dataloader()	一个可迭代对象或多个可迭代对象的集合，指定测试样本。
test_step(*args, **kwargs)	对测试集中的单个批次数据进行操作。
to(*args, **kwargs)	参见 torch.nn.Module.to()。
to_dtype(dtype)	将模块精度（默认为 float32）转换为另一种精度。
to_empty(*, device[, recurse])	将参数和缓冲区移动到指定设备，而不复制存储。
to_onnx(file_path[, input_sample])	将模型保存为 ONNX 格式。
to_torchscript([file_path, method, ...])	默认情况下，将整个模型编译为 ScriptModule。
toggle_optimizer(optimizer)	确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。
train([mode])	设置模块为训练模式。
train_dataloader()	指定训练样本的可迭代对象或可迭代对象的集合。
training_step(train_batch, batch_idx)	执行训练步骤
transfer_batch_to_device(batch, device, ...)	如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。
type(dst_type)	参见 torch.nn.Module.type()。
unfreeze()	解冻所有参数以进行训练。
untoggle_optimizer(optimizer)	重置通过 toggle_optimizer() 切换的所需梯度的状态。
val_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。
validation_step(val_batch, batch_idx)	执行验证步骤
xpu([device])	将所有模型参数和缓冲区移动到XPU。
zero_grad([set_to_none])	重置所有模型参数的梯度。

__call__
set_mc_dropout

CHECKPOINT_HYPER_PARAMS_KEY = 'hyper_parameters'

CHECKPOINT_HYPER_PARAMS_NAME = 'hparams_name'

CHECKPOINT_HYPER_PARAMS_TYPE = 'hparams_type'

T_destination

TypeVar(‘T_destination’, bound=Dict[str, Any]) 的别名

add_module(name, module)

将一个子模块添加到当前模块。

可以通过给定的名称将模块作为属性访问。

参数

• name (str) – 子模块的名称。可以使用给定的名称从此模块访问子模块。

• module (Module) – 要添加到模块中的子模块。

返回类型

None

all_gather(data, group=None, sync_grads=False)

从多个进程中收集张量或张量集合。

此方法需要在所有进程上调用，并且所有进程上的张量需要具有相同的形状，否则程序将永远停滞。

参数

• data (Union[Tensor, Dict, List, Tuple]) – int, float, 形状为 (batch, …) 的张量，或可能是嵌套的集合。

• group (Optional[Any, None]) – 要从中收集结果的进程组。默认为所有进程（world）

• sync_grads (bool) – 允许用户为 all_gather 操作同步梯度的标志

返回类型

Union[Tensor, Dict, List, Tuple]

返回

一个形状为 (world_size, batch, …) 的张量，或者如果输入是一个集合，输出也将是一个具有这种形状张量的集合。对于 world_size 为 1 的特殊情况，张量不会增加额外的维度。

apply(fn)

递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。

典型用途包括初始化模型的参数（另见 nn-init-doc）。

参数

fn (Module -> None) – 要应用于每个子模块的函数

返回

自己

返回类型

Module

示例:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

property automatic_optimization: bool

如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。

返回类型

bool

backward(loss, *args, **kwargs)

在 training_step() 中返回的损失上调用以执行反向传播。如果需要，请使用您自己的实现覆盖此钩子。

参数

loss (Tensor) – 由 training_step() 返回的损失张量。如果使用梯度累积，这里的损失值是归一化的（按 1 / 累积步数缩放）。

示例:

def backward(self, loss):
    loss.backward()

返回类型

None

bfloat16()

将所有浮点参数和缓冲区转换为 bfloat16 数据类型。

备注

此方法就地修改模块。

返回

自己

返回类型

Module

buffers(recurse=True)

返回一个遍历模块缓冲区的迭代器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。

生成器

torch.Tensor – 模块缓冲区

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Tensor]

call_super_init: bool = False

children()

返回一个遍历直接子模块的迭代器。

生成器

模块 – 子模块

返回类型

Iterator[Module]

clip_gradients(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

内部处理梯度裁剪。

备注

• 不要重写此方法。如果你想自定义梯度裁剪，考虑使用 configure_gradient_clipping() 方法。

• 对于手动优化（self.automatic_optimization = False），如果你想使用梯度裁剪，可以在训练步骤中手动调用 self.clip_gradients(opt, gradient_clip_val=0.5, gradient_clip_algorithm="norm")。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。

• gradient_clip_algorithm (Optional[str, None]) – 使用的梯度裁剪算法。传递 gradient_clip_algorithm="value" 按值裁剪，传递 gradient_clip_algorithm="norm" 按范数裁剪。

返回类型

None

compile(*args, **kwargs)

使用 torch.compile() 编译此模块的 forward。

此模块的 __call__ 方法被编译，所有参数都按原样传递给 torch.compile()。

有关此函数的参数详情，请参阅 torch.compile()。

configure_callbacks()

配置模型特定的回调。当模型被附加时，例如当调用 .fit() 或 .test() 时，这里返回的回调列表或回调将与传递给 Trainer 的 callbacks 参数的回调列表合并。如果这里返回的回调与 Trainer 的回调列表中已有的一个或多个回调类型相同，它将优先并替换它们。此外，Lightning 将确保 ModelCheckpoint 回调最后运行。

返回类型

Union[Sequence[Callback], Callback]

返回

一个回调函数或回调函数列表，这些回调函数将扩展 Trainer 中的回调函数列表。

示例:

def configure_callbacks(self):
    early_stop = EarlyStopping(monitor="val_acc", mode="max")
    checkpoint = ModelCheckpoint(monitor="val_loss")
    return [early_stop, checkpoint]

configure_gradient_clipping(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

对优化器参数执行梯度裁剪。在调用 optimizer_step() 之前执行。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。默认情况下，传递给 Trainer 的值将在此处可用。

• gradient_clip_algorithm (Optional[str, None]) – 要使用的梯度裁剪算法。默认情况下，传递给 Trainer 的值将在此处可用。

示例:

def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm):
    # Implement your own custom logic to clip gradients
    # You can call `self.clip_gradients` with your settings:
    self.clip_gradients(
        optimizer,
        gradient_clip_val=gradient_clip_val,
        gradient_clip_algorithm=gradient_clip_algorithm
    )

返回类型

None

configure_model()

在策略和精度感知的环境中创建模块的钩子。

这在使用分片策略（FSDP 和 DeepSpeed）时特别有用，我们希望立即分片模型以节省内存和初始化时间。对于非分片策略，您可以选择覆盖此钩子或在 init_module() 上下文管理器下初始化您的模型。

此钩子在每个 fit/val/test/predict 阶段都会被调用，因此请确保此钩子的实现是 幂等的，即，在第一次调用钩子后，后续的调用应该是无操作的。

返回类型

None

configure_optimizers()

配置优化器和学习率调度器以优化模型。

configure_sharded_model()

已弃用。

使用 configure_model() 代替。

返回类型

None

static configure_torch_metrics(torch_metrics)

处理 torch_metrics 参数。

返回类型

MetricCollection

cpu()

参见 torch.nn.Module.cpu()。

返回类型

Self

cuda(device=None)

将所有模型参数和缓冲区移动到GPU。这也会使相关的参数和缓冲区成为不同的对象。因此，如果模块在优化期间将驻留在GPU上，则应在构造优化器之前调用此方法。

参数

device (Union[device, int, None]) – 如果指定，所有参数将被复制到该设备。如果为 None，将使用当前的 CUDA 设备索引。

返回

自己

返回类型

Module

property current_epoch: int

Trainer 中的当前时期，如果未附加则为 0。

返回类型

int

property device: device

返回类型

device

property device_mesh: Optional[DeviceMesh]

像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。

返回类型

Optional[ForwardRef, None]

double()

参见 torch.nn.Module.double()。

返回类型

Self

property dtype: Union[str, dtype]

返回类型

Union[str, dtype]

dump_patches: bool = False

property epochs_trained

eval()

将模块设置为评估模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

这等同于 self.train(False)。

请参阅 locally-disable-grad-doc 以比较 .eval() 和几个可能与之混淆的类似机制。

返回

自己

返回类型

Module

property example_input_array: Optional[Union[Tensor, Tuple, Dict]]

示例输入数组是模块在 forward() 方法中可以使用的规范。返回类型解释如下：

• 单个张量：假设模型接受一个参数，即 model.forward(model.example_input_array)

• 元组：输入数组应被解释为位置参数的序列，即 model.forward(*model.example_input_array)

• 字典：输入数组表示命名关键字参数，即 model.forward(**model.example_input_array)

返回类型

Union[Tensor, Tuple, Dict, None]

extra_repr()

设置模块的额外表示。

要打印自定义的额外信息，您应该在您自己的模块中重新实现此方法。单行和多行字符串都是可接受的。

返回类型

str

property fabric: Optional[Fabric]

返回类型

Optional[Fabric, None]

float()

参见 torch.nn.Module.float()。

返回类型

Self

abstract forward(*args, **kwargs)

与 torch.nn.Module.forward() 相同。

参数

• *args – 无论你决定传入 forward 方法的内容。

• **kwargs – 关键字参数也是可能的。

返回类型

Any

返回

你的模型的输出

freeze()

冻结所有参数以进行推理。

示例:

model = MyLightningModule(...)
model.freeze()

返回类型

None

get_buffer(target)

如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的缓冲区的完全限定字符串名称。（有关如何指定完全限定字符串，请参见 get_submodule。）

返回

target 引用的缓冲区

返回类型

torch.Tensor

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是缓冲区的内容

get_extra_state()

返回要包含在模块的 state_dict 中的任何额外状态。

如果需要存储额外状态，请在你的模块中实现此功能以及相应的 set_extra_state()。此函数在构建模块的 state_dict() 时被调用。

请注意，额外的状态应该是可pickle的，以确保state_dict的序列化工作。我们仅对序列化张量提供向后兼容性保证；如果其他对象的序列化pickle形式发生变化，它们可能会破坏向后兼容性。

返回

模块的 state_dict 中存储的任何额外状态

返回类型

object

get_parameter(target)

如果存在，则返回由 target 给出的参数，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的参数的全限定字符串名称。（关于如何指定全限定字符串，请参见 get_submodule。）

返回

target 引用的参数

返回类型

torch.nn.Parameter

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Parameter 的内容

get_submodule(target)

如果存在，返回由 target 指定的子模块，否则抛出错误。

例如，假设你有一个 nn.Module A ，它看起来像这样：

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(该图展示了一个 nn.Module A。A 有一个嵌套的子模块 net_b，它本身有两个子模块 net_c 和 linear。net_c 然后有一个子模块 conv。)

要检查我们是否拥有 linear 子模块，我们会调用 get_submodule("net_b.linear")。要检查我们是否拥有 conv 子模块，我们会调用 get_submodule("net_b.net_c.conv")。

get_submodule 的运行时间受限于 target 中的模块嵌套程度。对 named_modules 的查询可以实现相同的结果，但它是 O(N) 的，其中 N 是传递模块的数量。因此，对于检查某个子模块是否存在的简单检查，应始终使用 get_submodule。

参数

target (str) – 要查找的子模块的完全限定字符串名称。（请参见上面的示例，了解如何指定完全限定的字符串。）

返回

由 target 引用的子模块

返回类型

torch.nn.Module

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Module 的内容

property global_rank: int

当前进程在所有节点和设备中的索引。

返回类型

int

property global_step: int

所有epoch中看到的总训练批次。

如果没有连接训练器，此属性为 0。

返回类型

int

half()

参见 torch.nn.Module.half()。

返回类型

Self

property hparams: Union[AttributeDict, MutableMapping]

通过 save_hyperparameters() 保存的超参数集合。用户可以对其进行修改。对于初始超参数的冻结集，请使用 hparams_initial。

返回类型

Union[AttributeDict, MutableMapping]

返回

可变超参数字典

property hparams_initial: AttributeDict

通过 save_hyperparameters() 保存的超参数集合。这些内容是只读的。可以通过 hparams 手动更新保存的超参数。

返回

不可变的初始超参数

返回类型

AttributeDict

ipu(device=None)

将所有模型参数和缓冲区移动到 IPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在IPU上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs)

从检查点加载模型的主要方式。当Lightning保存检查点时，它会将传递给 __init__ 的参数存储在检查点中的 "hyper_parameters" 下。

通过 **kwargs 指定的任何参数将覆盖存储在 "hyper_parameters" 中的参数。

参数

• checkpoint_path (Union[str, Path, IO]) – 检查点的路径。这也可以是一个URL，或类似文件的对象。

• map_location (Union[device, str, int, Callable[[UntypedStorage, str], Optional[UntypedStorage, None]], Dict[Union[device, str, int], Union[device, str, int]], None]) – 如果你的检查点保存了一个GPU模型，而你现在在CPU上或不同数量的GPU上加载，使用这个来映射到新的设置。其行为与 torch.load() 相同。

• hparams_file (Union[str, Path, None]) –

  可选路径到一个 .yaml 或 .csv 文件，其层次结构如本示例所示:

  drop_prob: 0.2
  dataloader:
      batch_size: 32
  

  您很可能不需要这个，因为 Lightning 总是会将超参数保存到检查点。然而，如果您的检查点权重没有保存超参数，请使用此方法传入一个包含您希望使用的超参数的 .yaml 文件。这些将被转换为 dict 并传递到您的 LightningModule 中使用。

  如果你的模型的 hparams 参数是 Namespace 并且 .yaml 文件具有层次结构，你需要重构你的模型以将 hparams 视为 dict。

• strict (Optional[bool, None]) – 是否严格要求 checkpoint_path 中的键与该模块状态字典返回的键匹配。默认为 True ，除非设置了 LightningModule.strict_loading ，在这种情况下，它默认为 LightningModule.strict_loading 的值。

• **kwargs – 初始化模型所需的任何额外关键字参数。也可以用来覆盖已保存的超参数值。

返回类型

Self

返回

LightningModule 实例，加载了权重和超参数（如果可用）。

备注

load_from_checkpoint 是一个 类 方法。你应该使用你的 LightningModule 类 来调用它，而不是 LightningModule 实例，否则会引发 TypeError。

备注

为了确保所有层都能从检查点加载，如果您的 LightningModule 中重写了此钩子，此函数将在实例化模型后直接调用 configure_model()。但是，请注意 load_from_checkpoint 不支持加载分片检查点，如果模型太大，您可能会耗尽内存。在这种情况下，请考虑通过 .fit(ckpt_path=...) 通过 Trainer 加载。

示例:

# load weights without mapping ...
model = MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')

# or load weights mapping all weights from GPU 1 to GPU 0 ...
map_location = {'cuda:1':'cuda:0'}
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    map_location=map_location
)

# or load weights and hyperparameters from separate files.
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    hparams_file='/path/to/hparams_file.yaml'
)

# override some of the params with new values
model = MyLightningModule.load_from_checkpoint(
    PATH,
    num_layers=128,
    pretrained_ckpt_path=NEW_PATH,
)

# predict
pretrained_model.eval()
pretrained_model.freeze()
y_hat = pretrained_model(x)

load_state_dict(state_dict, strict=True, assign=False)

将参数和缓冲区从 state_dict 复制到此模块及其后代中。

如果 strict 是 True，那么 state_dict 的键必须与该模块的 state_dict() 函数返回的键完全匹配。

警告

如果 assign 是 True，优化器必须在调用 load_state_dict 之后创建，除非 get_swap_module_params_on_conversion() 是 True。

参数

• state_dict (dict) – 包含参数和持久缓冲区的字典。

• strict (bool, optional) – 是否严格要求 state_dict 中的键与该模块的 state_dict() 函数返回的键匹配。默认值：True

• assign (bool, optional) – 当 False 时，当前模块中张量的属性会被保留，而当 True 时，状态字典中张量的属性会被保留。唯一的例外是 Parameter 的 requires_grad 字段，该字段的值会从模块中保留。默认值：False

返回

• missing_keys 是一个包含任何预期键的 str 列表

  此模块中存在但未在提供的 state_dict 中出现的。

• unexpected_keys 是一个包含未预料到的键的 str 列表。

  此模块预期但未在提供的 state_dict 中出现。

返回类型

带有 missing_keys 和 unexpected_keys 字段的 NamedTuple

备注

如果一个参数或缓冲区被注册为 None ，并且其对应的键存在于 state_dict 中，load_state_dict() 将引发一个 RuntimeError 。

property local_rank: int

当前进程在单个节点中的索引。

返回类型

int

log(name, value, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, metric_attribute=None, rank_zero_only=False)

记录一个键值对。

示例:

self.log('train_loss', loss)

每个钩子的默认行为在此处有文档说明：extensions/logging:自动日志记录。

参数

• name (str) – 日志的键。如果使用DDP或其他分布式策略，必须在所有进程中保持一致。

• value (Union[Metric, Tensor, int, float]) – 要记录的值。可以是 float、Tensor 或 Metric。

• prog_bar (bool) – 如果 True 则记录到进度条。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• on_epoch (Optional[bool, None]) – 如果 True ，记录每个epoch累积的指标。默认值由钩子决定。详情请参见 extensions/logging:自动日志记录 。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图表。

• sync_dist (bool) – 如果 True，将在设备间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的DDP组。

• add_dataloader_idx (bool) – 如果 True，将当前数据加载器的索引附加到名称（当使用多个数据加载器时）。如果为 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的 batch_size。这将直接从加载的批次中推断出来，但对于某些数据结构，您可能需要显式提供它。

• metric_attribute (Optional[str, None]) – 要恢复指标状态，Lightning 需要模型中 torchmetrics.Metric 的引用。如果它是模型属性，则会自动找到。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

log_dict(dictionary, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, rank_zero_only=False)

一次性记录一个值的字典。

示例:

values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n}
self.log_dict(values)

参数

• dictionary (Union[Mapping[str, Union[Metric, Tensor, int, float]], MetricCollection]) – 键值对。如果使用DDP或其他分布式策略，键必须在所有进程中保持一致。值可以是 float 、 Tensor 、 Metric 或 MetricCollection 。

• prog_bar (bool) – 如果 True 则记录到进度基础。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。None 自动记录训练步骤，但不记录验证/测试步骤。默认值由钩子决定。详情请参阅 extensions/logging:自动记录。

• on_epoch (Optional[bool, None]) – 如果 True 记录累积的 epoch 指标。None 自动记录 val/test 步骤的日志，但不包括 training_step。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图形

• sync_dist (bool) – 如果 True，将在GPU/TPU之间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的 ddp 组。

• add_dataloader_idx (bool) – 如果 True，则将当前数据加载器的索引附加到名称（在使用多个数据加载器时）。如果 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的批次大小。这将直接从加载的批次中推断出来，但某些数据结构可能需要显式提供它。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

property logger: Optional[Union[Logger, Logger]]

Trainer 中对 logger 对象的引用。

返回类型

Union[Logger, Logger, None]

property loggers: Union[List[Logger], List[Logger]]

参考 Trainer 中的日志记录器列表。

返回类型

Union[List[Logger], List[Logger]]

lr_scheduler_step(scheduler, metric)

重写此方法以调整 Trainer 调用每个调度器的默认方式。默认情况下，Lightning 会根据每个调度器的 interval 调用 step() 方法，如示例所示。

参数

• scheduler (Union[LRScheduler, ReduceLROnPlateau]) – 学习率调度器。

• metric (Optional[Any, None]) – 用于 ReduceLROnPlateau 等调度器的监视器值。

示例:

# DEFAULT
def lr_scheduler_step(self, scheduler, metric):
    if metric is None:
        scheduler.step()
    else:
        scheduler.step(metric)

# Alternative way to update schedulers if it requires an epoch value
def lr_scheduler_step(self, scheduler, metric):
    scheduler.step(epoch=self.current_epoch)

返回类型

None

lr_schedulers()

返回在训练期间使用的学习率调度器。对于手动优化很有用。

返回类型

Union[None, List[Union[LRScheduler, ReduceLROnPlateau]], LRScheduler, ReduceLROnPlateau]

返回

一个调度器，或者在存在多个调度器的情况下是一个调度器列表，如果在 configure_optimizers() 中没有返回调度器，则为 None。

manual_backward(loss, *args, **kwargs)

当你手动进行优化时，请直接从你的 training_step() 调用此方法。通过使用此方法，Lightning 可以确保在使用混合精度时应用所有适当的缩放。

更多示例请参见 手动优化。

示例:

def training_step(...):
    opt = self.optimizers()
    loss = ...
    opt.zero_grad()
    # automatically applies scaling, etc...
    self.manual_backward(loss)
    opt.step()

参数

• loss (Tensor) – 要计算梯度的张量。必须附加一个图。

• *args – 要转发给 backward() 的附加位置参数

• **kwargs – 要转发给 backward() 的额外关键字参数

返回类型

None

modules()

返回网络中所有模块的迭代器。

生成器

模块 – 网络中的一个模块

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

返回类型

Iterator[Module]

named_buffers(prefix='', recurse=True, remove_duplicate=True)

返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。

参数

• prefix (str) – 要添加到所有缓冲区名称的前缀。

• recurse (bool, optional) – 如果为 True，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。默认为 True。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的缓冲区。默认为 True。

生成器

(str, torch.Tensor) – 包含名称和缓冲区的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

返回类型

Iterator[Tuple[str, Tensor]]

named_children()

返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。

生成器

(str, Module) – 包含名称和子模块的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

返回类型

Iterator[Tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)

返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。

参数

• memo (Optional[Set[Module], None]) – 一个备忘录，用于存储已添加到结果中的模块集合

• prefix (str) – 一个将添加到模块名称前的前缀

• remove_duplicate (bool) – 是否在结果中移除重复的模块实例

生成器

(str, Module) – 名称和模块的元组

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)

返回一个模块参数的迭代器，同时生成参数的名称和参数本身。

参数

• prefix (str) – 要添加到所有参数名称前的前缀。

• recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的参数。默认为 True。

生成器

(str, Parameter) – 包含名称和参数的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

返回类型

Iterator[Tuple[str, Parameter]]

on_after_backward()

在 loss.backward() 之后和优化器步进之前调用。

备注

如果使用原生 AMP，梯度在此刻将不会被缩放。如果你需要未缩放的梯度，请使用 on_before_optimizer_step。

返回类型

None

on_after_batch_transfer(batch, dataloader_idx)

覆盖此方法以在批次转移到设备后对其进行批量增强或修改。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_after_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = gpu_transforms(batch['x'])
    return batch

参见

• on_before_batch_transfer()

• transfer_batch_to_device()

on_before_backward(loss)

在 loss.backward() 之前调用。

参数

loss (Tensor) – 损失除以梯度累积的批次数量，如果使用AMP则进行缩放。

返回类型

None

on_before_batch_transfer(batch, dataloader_idx)

在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_before_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = transforms(batch['x'])
    return batch

参见

• on_after_batch_transfer()

• transfer_batch_to_device()

on_before_optimizer_step(optimizer)

在 optimizer.step() 之前调用。

如果使用梯度累积，钩子会在梯度累积完成后被调用。参见：:paramref:`~pytorch_lightning.trainer.trainer.Trainer.accumulate_grad_batches`。

如果使用AMP，损失将在调用此钩子之前未缩放。有关梯度缩放的更多信息，请参阅这些 文档。

如果裁剪梯度，梯度尚未被裁剪。

参数

optimizer (Optimizer) – 当前使用的优化器。

示例:

def on_before_optimizer_step(self, optimizer):
    # example to inspect gradient information in tensorboard
    if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
        for k, v in self.named_parameters():
            self.logger.experiment.add_histogram(
                tag=k, values=v.grad, global_step=self.trainer.global_step
            )

返回类型

None

on_before_zero_grad(optimizer)

在 training_step() 之后和 optimizer.zero_grad() 之前调用。

在训练循环中，调用优化器步骤之后和清零梯度之前。这是一个检查更新后权重信息的好地方。

这是调用它的地方:

for optimizer in optimizers:
    out = training_step(...)

    model.on_before_zero_grad(optimizer) # < ---- called here
    optimizer.zero_grad()

    backward()

参数

optimizer (Optimizer) – 应该被置零的梯度优化器。

返回类型

None

on_fit_end()

在拟合的最后调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

on_fit_start()

在拟合开始时调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

property on_gpu: bool

如果该模型当前位于GPU上，则返回 True。

在LightningModule周围设置标志以区分CPU和GPU行为非常有用。

返回类型

bool

on_load_checkpoint(checkpoint)

由 Lightning 调用以恢复您的模型。如果您使用 on_save_checkpoint() 保存了某些内容，这是您恢复这些内容的机会。

参数

checkpoint (Dict[str, Any]) – 加载的检查点

示例:

def on_load_checkpoint(self, checkpoint):
    # 99% of the time you don't need to implement this method
    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']

备注

Lightning 自动恢复全局步数、轮次和训练状态，包括 amp 缩放。您无需恢复任何与训练相关的内容。

返回类型

None

on_predict_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后调用预测循环。

参数

• outputs (Optional[Any, None]) – predict_step(x) 的输出

• batch (Any) – 由预测 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_batch_start(batch, batch_idx, dataloader_idx=0)

在预测循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_start()

在预测开始时调用。

返回类型

None

on_predict_model_eval()

在预测循环开始时调用。

默认情况下，预测循环在开始之前会在 LightningModule 上调用 .eval()。重写此钩子以更改行为。

返回类型

None

on_predict_start()

在预测开始时调用。

返回类型

None

on_save_checkpoint(checkpoint)

当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。

参数

checkpoint (Dict[str, Any]) – 在转储到文件之前的完整检查点字典。此钩子的实现可以将额外数据插入此字典中。

示例:

def on_save_checkpoint(self, checkpoint):
    # 99% of use cases you don't need to implement this method
    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object

备注

Lightning 保存了训练的所有方面（如 epoch、全局步数等），包括 amp 缩放。您无需存储任何关于训练的信息。

返回类型

None

on_test_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后在测试循环中调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – test_step(x) 的输出

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_batch_start(batch, batch_idx, dataloader_idx=0)

在测试循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_end()

在测试结束时调用。

返回类型

None

on_test_epoch_end()

在每个epoch的测试循环结束时调用。

返回类型

None

on_test_epoch_start()

在epoch开始时在测试循环中调用。

返回类型

None

on_test_model_eval()

在测试循环开始时调用。

默认情况下，测试循环在开始前会调用 .eval() 方法在 LightningModule 上。重写此钩子以改变行为。另请参见 on_test_model_train()。

返回类型

None

on_test_model_train()

在测试循环结束时调用。

测试循环默认会将 LightningModule 的 training 模式恢复到开始测试之前的状态。覆盖此钩子以更改行为。另请参见 on_test_model_eval()。

返回类型

None

on_test_start()

在测试开始时调用。

返回类型

None

on_train_batch_end(outputs, batch, batch_idx)

在训练循环中批次处理后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – 训练步骤(x)的输出

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

备注

这里的 outputs["loss"] 值将是相对于 accumulate_grad_batches 的损失归一化值，该损失从 training_step 返回。

返回类型

None

on_train_batch_start(batch, batch_idx)

在训练循环中，在该批次发生任何事情之前调用。

如果你在这里返回 -1，你将跳过当前轮次的剩余训练。

参数

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

返回类型

Optional[int, None]

on_train_end()

在训练结束前调用，在记录器实验关闭之前。

返回类型

None

on_train_epoch_end()

在训练循环的每个epoch结束时调用。

要在每个epoch结束时访问所有批处理输出，您可以将步骤输出缓存为 LightningModule 的一个属性，并在此钩子中访问它们：

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # free up the memory
        self.training_step_outputs.clear()

on_train_epoch_start()

在训练循环中，在epoch开始时调用。

返回类型

None

on_train_start()

在完整性检查后，训练开始时调用。

返回类型

None

on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次验证循环后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – validation_step(x) 的输出

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_batch_start(batch, batch_idx, dataloader_idx=0)

在验证循环中，对该批次进行任何操作之前调用。

参数

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_end()

在验证结束时调用。

返回类型

None

on_validation_epoch_end()

在epoch结束时的验证循环中调用。

on_validation_epoch_start()

在epoch开始时验证循环中调用。

返回类型

None

on_validation_model_eval()

在验证循环开始时调用。

默认情况下，验证循环在开始前会在 LightningModule 上调用 .eval()。覆盖此钩子以更改行为。另请参见 on_validation_model_train()。

返回类型

None

on_validation_model_train()

当验证循环结束时调用。

默认情况下，验证循环会将 LightningModule 的 training 模式恢复到开始验证之前的状态。重写此钩子以更改行为。另请参见 on_validation_model_eval()。

返回类型

None

on_validation_model_zero_grad()

在进入验证循环之前，由训练循环调用来释放梯度。

返回类型

None

on_validation_start()

在验证开始时调用。

返回类型

None

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)

重写此方法以调整 Trainer 调用优化器的默认方式。

默认情况下，Lightning 会调用 step() 和 zero_grad()，如示例所示。当 Trainer(accumulate_grad_batches != 1) 时，在累积阶段不会调用此方法（和 zero_grad()）。手动优化时，重写此钩子没有好处。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Union[Optimizer, LightningOptimizer]) – 一个 PyTorch 优化器

• optimizer_closure (Optional[Callable[[], Any], None]) – 优化器闭包。此闭包必须执行，因为它包含了 training_step()、optimizer.zero_grad() 和 backward() 的调用。

示例:

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    # Add your custom logic to run directly before `optimizer.step()`

    optimizer.step(closure=optimizer_closure)

    # Add your custom logic to run directly after `optimizer.step()`

返回类型

None

optimizer_zero_grad(epoch, batch_idx, optimizer)

重写此方法以改变 optimizer.zero_grad() 的默认行为。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Optimizer) – 一个 PyTorch 优化器

示例:

# DEFAULT
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad()

# Set gradients to `None` instead of zero to improve performance (not required on `torch>=2.0.0`).
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad(set_to_none=True)

参见 torch.optim.Optimizer.zero_grad() 以了解上述示例的解释。

返回类型

None

optimizers(use_pl_optimizer=True)

返回训练期间使用的优化器。对于手动优化很有用。

参数

use_pl_optimizer (bool) – 如果 True，将会把优化器包装在 LightningOptimizer 中，以自动处理精度、分析和步数调用的计数，以进行适当的日志记录和检查点保存。它特别包装了 step 方法，不支持没有此方法的自定义优化器。

返回类型

Union[Optimizer, LightningOptimizer, _FabricOptimizer, List[Optimizer], List[LightningOptimizer], List[_FabricOptimizer]]

返回

一个单一的优化器，或者在存在多个优化器的情况下，一个优化器列表。

property output_chunk_length: Optional[int]

模型一次预测的时间步数。

返回类型

Optional[int, None]

parameters(recurse=True)

返回一个遍历模块参数的迭代器。

这通常传递给一个优化器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

生成器

参数 – 模块参数

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Parameter]

predict_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。

有关多个数据加载器的更多信息，请参阅此 部分。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• predict()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

返回

一个 torch.utils.data.DataLoader 或它们的序列，指定预测样本。

predict_step(batch, batch_idx, dataloader_idx=None)

执行预测步骤

批处理

Darts 的 InferenceDataset 的输出 - 包含 (过去的目标, 过去的协变量, 历史未来协变量, 未来协变量, 未来过去协变量, 输入时间序列, 预测开始时间步) 的元组

batch_idx

当前批次的批次索引

dataloader_idx

数据加载器索引

返回类型

Sequence[TimeSeries]

prepare_data()

使用此方法下载和准备数据。使用多个进程（分布式设置）下载和保存数据将导致数据损坏。Lightning 确保此方法仅在一个进程内调用，因此您可以安全地在其中添加下载逻辑。

警告

不要将状态设置到模型中（使用 setup 代替），因为这不会在每个设备上被调用

示例:

def prepare_data(self):
    # good
    download_data()
    tokenize()
    etc()

    # bad
    self.split = data_split
    self.some_state = some_other_state()

在分布式环境中，prepare_data 可以通过两种方式调用（使用 prepare_data_per_node）

1. 每个节点一次。这是默认设置，仅在 LOCAL_RANK=0 时调用。

2. 总共一次。仅在 GLOBAL_RANK=0 时调用。

示例:

# DEFAULT
# called once per node on LOCAL_RANK=0 of that node
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = True

# call on GLOBAL_RANK=0 (great for shared file systems)
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = False

在请求数据加载器之前调用此方法：

model.prepare_data()
initialize_distributed()
model.setup(stage)
model.train_dataloader()
model.val_dataloader()
model.test_dataloader()
model.predict_dataloader()

返回类型

None

print(*args, **kwargs)

仅从进程 0 打印。在任何分布式模式下使用此功能以仅记录一次。

参数

• *args – 要打印的内容。与Python内置的print函数相同。

• **kwargs – 与Python内置的print函数相同。

示例:

def forward(self, x):
    self.print(x, 'in forward')

返回类型

None

register_backward_hook(hook)

在模块上注册一个反向钩子。

此函数已被 register_full_backward_hook() 取代，并且此函数的行为将在未来版本中发生变化。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_buffer(name, tensor, persistent=True)

向模块添加一个缓冲区。

这通常用于注册一个不应被视为模型参数的缓冲区。例如，BatchNorm 的 running_mean 不是参数，而是模块状态的一部分。缓冲区默认是持久的，并将与参数一起保存。可以通过将 persistent 设置为 False 来改变此行为。持久缓冲区和非持久缓冲区之间的唯一区别是后者不会成为此模块的 state_dict 的一部分。

缓冲区可以通过给定的名称作为属性访问。

参数

• name (str) – 缓冲区的名称。可以通过给定的名称从此模块访问缓冲区。

• tensor (Tensor or None) – 要注册的缓冲区。如果为 None ，则运行在缓冲区上的操作，例如 cuda ，将被忽略。如果为 None ，该缓冲区将 不 包含在模块的 state_dict 中。

• persistent (bool) – 缓冲区是否是此模块的 state_dict 的一部分。

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

返回类型

None

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

在模块上注册一个前向钩子。

钩子将在每次 forward() 计算输出后被调用。

如果 with_kwargs 是 False 或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward 。钩子可以修改输出。它可以就地修改输入，但这不会影响前向传播，因为这是在调用 forward() 之后调用的。钩子应具有以下签名:

hook(module, args, output) -> None or modified output

如果 with_kwargs 是 True，前向钩子将被传递给前向函数的 kwargs，并期望返回可能被修改的输出。钩子应具有以下签名:

hook(module, args, kwargs, output) -> None or modified output

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果 True，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之后触发。请注意，使用 register_module_forward_hook() 注册的全局 forward 钩子将在通过此方法注册的所有钩子之前触发。默认值：False

• with_kwargs (bool) – 如果 True，hook 将会传递给 forward 函数的 kwargs。默认值：False

• always_call (bool) – 如果 True ，无论在调用模块时是否引发异常，都将运行 hook 。默认值： False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

在模块上注册一个前向预钩子。

钩子将在每次调用 forward() 之前被调用。

如果 with_kwargs 为假或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward。钩子可以修改输入。用户可以在钩子中返回一个元组或单个修改后的值。如果返回单个值（除非该值已经是元组），我们将把该值包装成一个元组。钩子应具有以下签名:

hook(module, args) -> None or modified input

如果 with_kwargs 为真，前向预钩子将传递给前向函数的 kwargs。如果钩子修改了输入，则应返回 args 和 kwargs。钩子应具有以下签名:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之后触发。请注意，通过 register_module_forward_pre_hook() 注册的全局 forward_pre 钩子将在所有通过此方法注册的钩子之前触发。默认值：False

• with_kwargs (bool) – 如果为真，hook 将被传递给 forward 函数的 kwargs。默认值：False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)

在模块上注册一个反向钩子。

每次计算与模块相关的梯度时，钩子都会被调用，即当且仅当计算与模块输出相关的梯度时，钩子才会执行。钩子应具有以下签名:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input 和 grad_output 是包含相对于输入和输出的梯度的元组。钩子不应修改其参数，但可以选择返回一个新的相对于输入的梯度，该梯度将用于替代 grad_input 在后续计算中。grad_input 将仅对应于作为位置参数给出的输入，所有 kwarg 参数都将被忽略。grad_input 和 grad_output 中的条目将为 None 用于所有非 Tensor 参数。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入或输出，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之后触发。请注意，通过 register_module_full_backward_hook() 注册的全局 backward 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

在模块上注册一个反向预钩子。

每次计算模块的梯度时，都会调用该钩子。钩子应具有以下签名:

hook(module, grad_output) -> tuple[Tensor] or None

The grad_output 是一个元组。钩子不应修改其参数，但它可以选择返回一个新的相对于输出的梯度，该梯度将用于替代 grad_output 在后续计算中。grad_output 中的条目对于所有非 Tensor 参数将为 None。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之后触发。请注意，通过 register_module_full_backward_pre_hook() 注册的全局 backward_pre 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)

注册一个在模块的 load_state_dict 调用后运行的后置钩子。

它应该有以下签名:

hook(module, incompatible_keys) -> None

module 参数是当前注册此钩子的模块，而 incompatible_keys 参数是一个由属性 missing_keys 和 unexpected_keys 组成的 NamedTuple。missing_keys 是一个包含缺失键的 str 的 list，而 unexpected_keys 是一个包含意外键的 str 的 list。

如果需要，可以就地修改给定的 incompatible_keys。

请注意，当使用 strict=True 调用 load_state_dict() 时执行的检查会受到钩子对 missing_keys 或 unexpected_keys 所做修改的影响，这是预期的。对任一组键的添加将导致在 strict=True 时抛出错误，而清除掉缺失和意外的键将避免错误。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_module(name, module)

别名 add_module()。

返回类型

None

register_parameter(name, param)

向模块添加一个参数。

该参数可以通过给定的名称作为属性访问。

参数

• name (str) – 参数的名称。可以使用给定的名称从此模块访问该参数。

• param (Parameter or None) – 要添加到模块的参数。如果为 None ，则运行在参数上的操作，例如 cuda ，将被忽略。如果为 None ，该参数将 不 包含在模块的 state_dict 中。

返回类型

None

register_state_dict_pre_hook(hook)

为 state_dict() 方法注册一个预钩子。

这些钩子将在调用 self 上的 state_dict 之前，使用参数 self、prefix 和 keep_vars 进行调用。注册的钩子可以用于在调用 state_dict 之前执行预处理。

requires_grad_(requires_grad=True)

更改是否应在此模块的参数上记录自动梯度操作。

此方法就地设置参数的 requires_grad 属性。

这种方法有助于冻结模块的一部分以进行微调或单独训练模型的一部分（例如，GAN训练）。

请参阅 locally-disable-grad-doc 以比较 .requires_grad_() 和几个可能与之混淆的类似机制。

参数

requires_grad (bool) – 是否应在此模块的参数上记录自动求导操作。默认值：True。

返回

自己

返回类型

Module

save_hyperparameters(*args, ignore=None, frame=None, logger=True)

将参数保存到 hparams 属性中。

参数

• args (Any) – dict、NameSpace 或 OmegaConf 的单个对象，或来自类 __init__ 的字符串名称或参数

• ignore (Union[Sequence[str], str, None]) – 一个参数名称或从类 __init__ 中忽略的参数名称列表

• frame (Optional[frame, None]) – 一个框架对象。默认是 None

• logger (bool) – 是否将超参数发送到记录器。默认值：True

示例:

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # manually assign arguments
...         self.save_hyperparameters('arg1', 'arg3')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class AutomaticArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # equivalent automatic
...         self.save_hyperparameters()
...     def forward(self, *args, **kwargs):
...         ...
>>> model = AutomaticArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg2": abc
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class SingleArgModel(HyperparametersMixin):
...     def __init__(self, params):
...         super().__init__()
...         # manually assign single argument
...         self.save_hyperparameters(params)
...     def forward(self, *args, **kwargs):
...         ...
>>> model = SingleArgModel(Namespace(p1=1, p2='abc', p3=3.14))
>>> model.hparams
"p1": 1
"p2": abc
"p3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # pass argument(s) to ignore as a string or in a list
...         self.save_hyperparameters(ignore='arg2')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

返回类型

None

set_extra_state(state)

设置包含在加载的 state_dict 中的额外状态。

此函数从 load_state_dict() 调用，以处理在 state_dict 中找到的任何额外状态。如果你的模块需要在 state_dict 中存储额外状态，请实现此函数以及相应的 get_extra_state()。

参数

state (dict) – 来自 state_dict 的额外状态

返回类型

None

set_mc_dropout(active)

set_predict_parameters(n, num_samples, roll_size, batch_size, n_jobs, predict_likelihood_parameters, mc_dropout)

在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。

返回类型

None

setup(stage)

在 fit（训练 + 验证）、验证、测试或预测开始时调用。当你需要动态构建模型或调整模型时，这是一个很好的钩子。在使用 DDP 时，此钩子会在每个进程上调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

示例:

class LitModel(...):
    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

        # don't do this
        self.something = else

    def setup(self, stage):
        data = load_data(...)
        self.l1 = nn.Linear(28, data.num_classes)

返回类型

None

share_memory()

参见 torch.Tensor.share_memory_()。

返回类型

~T

state_dict(*args, destination=None, prefix='', keep_vars=False)

返回一个包含模块整体状态引用的字典。

参数和持久缓冲区（例如运行平均值）都包含在内。键是对应的参数和缓冲区名称。设置为 None 的参数和缓冲区不包含在内。

备注

返回的对象是一个浅拷贝。它包含对模块参数和缓冲区的引用。

警告

目前，state_dict() 也按顺序接受 destination、prefix 和 keep_vars 的位置参数。然而，这正在被弃用，未来版本将强制使用关键字参数。

警告

请避免使用参数 destination，因为它不是为最终用户设计的。

参数

• destination (dict, optional) – 如果提供了，模块的状态将被更新到字典中，并返回相同的对象。否则，将创建并返回一个 OrderedDict。默认值：None。

• prefix (str, optional) – 在 state_dict 中组合键时添加到参数和缓冲区名称的前缀。默认值：''。

• keep_vars (bool, optional) – 默认情况下，状态字典中返回的 Tensor 是与自动求导分离的。如果设置为 True，则不会执行分离操作。默认值：False。

返回

包含模块整个状态的字典

返回类型

dict

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

property strict_loading: bool

确定 Lightning 如何使用 .load_state_dict(…, strict=model.strict_loading) 加载此模型。

返回类型

bool

property supports_probabilistic_prediction: bool

返回类型

bool

teardown(stage)

在 fit（训练 + 验证）、验证、测试或预测结束时调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

返回类型

None

test_dataloader()

一个可迭代对象或多个可迭代对象的集合，指定测试样本。

有关多个数据加载器的更多信息，请参阅此 部分。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• test()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要测试数据集和 test_step() ，你不需要实现这个方法。

返回类型

Any

test_step(*args, **kwargs)

对测试集中的单个批次数据进行操作。在这一步中，您通常会生成示例或计算任何感兴趣的内容，例如准确性。

参数

• batch – 你的数据可迭代对象的输出，通常是一个 DataLoader。

• batch_idx – 这一批的索引。

• dataloader_idx – 生成此批次的 dataloader 的索引。（仅在使用多个 dataloader 时）

返回类型

Union[Tensor, Mapping[str, Any], None]

返回

• Tensor - 损失张量

• dict - 一个字典。可以包含任何键，但必须包含键 'loss'。

• None - 跳到下一个批次。

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...

# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

示例:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

如果你传入多个测试数据加载器，test_step() 将会有一个额外的参数。我们建议将默认值设置为0，这样你可以快速在单个和多个数据加载器之间切换。

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

备注

如果你不需要测试，你就不需要实现这个方法。

备注

当调用 test_step() 时，模型已处于评估模式，并且已禁用 PyTorch 梯度。在测试轮次结束时，模型返回到训练模式，并且启用梯度。

to(*args, **kwargs)

参见 torch.nn.Module.to()。

返回类型

Self

to_dtype(dtype)

将模块精度（默认为 float32）转换为另一种精度。

to_empty(*, device, recurse=True)

将参数和缓冲区移动到指定设备，而不复制存储。

参数

• device (torch.device) – 此模块中参数和缓冲区的所需设备。

• recurse (bool) – 子模块的参数和缓冲区是否应递归地移动到指定设备。

返回

自己

返回类型

Module

to_onnx(file_path, input_sample=None, **kwargs)

将模型保存为 ONNX 格式。

参数

• file_path (Union[str, Path]) – 应保存 onnx 模型的文件路径。

• input_sample (Optional[Any, None]) – 用于跟踪的输入。默认值：None（使用 self.example_input_array）

• **kwargs – 将被传递给 torch.onnx.export 函数。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)

model = SimpleModel()
input_sample = torch.randn(1, 64)
model.to_onnx("export.onnx", input_sample, export_params=True)

返回类型

None

to_torchscript(file_path=None, method='script', example_inputs=None, **kwargs)

默认情况下，将整个模型编译为 ScriptModule。如果你想使用追踪，请提供参数 method='trace' 并确保提供了 example_inputs 参数，或者模型设置了 example_input_array。如果你想自定义被脚本的模块，你应该重写这个方法。如果你想要返回多个模块，我们建议使用字典。

参数

• file_path (Union[str, Path, None]) – 保存 torchscript 的路径。默认：None（不保存文件）。

• method (Optional[str, None]) – 是否使用 TorchScript 的 script 或 trace 方法。默认值：’script’

• example_inputs (Optional[Any, None]) – 当方法设置为’trace’时，用于进行追踪的输入。默认值：None（使用 example_input_array）

• **kwargs – 将传递给 torch.jit.script() 或 torch.jit.trace() 函数的额外参数。

备注

• 需要实现 forward() 方法。

• 导出的脚本将被设置为评估模式。

• 建议您安装最新支持版本的 PyTorch 以无限制地使用此功能。另请参阅 torch.jit 文档以了解支持的功能。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

model = SimpleModel()
model.to_torchscript(file_path="model.pt")

torch.jit.save(model.to_torchscript(
    file_path="model_trace.pt", method='trace', example_inputs=torch.randn(1, 64))
)

返回类型

Union[ScriptModule, Dict[str, ScriptModule]]

返回

无论 file_path 是否定义，此 LightningModule 都可以作为 torchscript 使用。

toggle_optimizer(optimizer)

确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。

它与 untoggle_optimizer() 一起工作，以确保 param_requires_grad_state 被正确重置。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要切换的优化器。

返回类型

None

train(mode=True)

设置模块为训练模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

参数

mode (bool) – 是否设置训练模式（True）或评估模式（False）。默认值：True。

返回

自己

返回类型

Module

train_dataloader()

指定训练样本的可迭代对象或可迭代对象的集合。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• fit()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

property trainer: Trainer

返回类型

Trainer

training: bool

training_step(train_batch, batch_idx)

执行训练步骤

返回类型

Tensor

transfer_batch_to_device(batch, device, dataloader_idx)

如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。

以下列出的数据类型（以及它们的任意嵌套）都是开箱即用的：

• torch.Tensor 或任何实现了 .to(…) 的对象

• 列表

• 字典

• 元组

对于其他情况，你需要定义数据如何移动到目标设备（CPU、GPU、TPU等）。

备注

此钩子应仅传输数据，不应修改数据，也不应将数据移动到除作为参数传入的设备之外的任何其他设备（除非你知道自己在做什么）。要检查此钩子的当前执行状态，可以使用 self.trainer.training/testing/validating/predicting，以便可以根据你的需求添加不同的逻辑。

参数

• batch (Any) – 一批需要转移到新设备的数据。

• device (device) – 在 PyTorch 中定义的目标设备。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

对新设备上数据的引用。

示例:

def transfer_batch_to_device(self, batch, device, dataloader_idx):
    if isinstance(batch, CustomBatch):
        # move all tensors in your custom data structure to the device
        batch.samples = batch.samples.to(device)
        batch.targets = batch.targets.to(device)
    elif dataloader_idx == 0:
        # skip device transfer for the first dataloader or anything you wish
        pass
    else:
        batch = super().transfer_batch_to_device(batch, device, dataloader_idx)
    return batch

参见

• move_data_to_device()

• apply_to_collection()

type(dst_type)

参见 torch.nn.Module.type()。

返回类型

Self

unfreeze()

解冻所有参数以进行训练。

model = MyLightningModule(...)
model.unfreeze()

返回类型

None

untoggle_optimizer(optimizer)

重置通过 toggle_optimizer() 切换的所需梯度的状态。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要取消选择的优化器。

返回类型

None

val_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• fit()

• validate()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要验证数据集和 validation_step()，你不需要实现这个方法。

返回类型

Any

validation_step(val_batch, batch_idx)

执行验证步骤

返回类型

Tensor

xpu(device=None)

将所有模型参数和缓冲区移动到XPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在 XPU 上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

zero_grad(set_to_none=True)

重置所有模型参数的梯度。

更多上下文请参见 torch.optim.Optimizer 下的类似函数。

参数

set_to_none (bool) – 而不是设置为零，将梯度设置为 None。详情请参阅 torch.optim.Optimizer.zero_grad()。

返回类型

None

class darts.models.forecasting.pl_forecasting_module.PLForecastingModule(input_chunk_length, output_chunk_length, output_chunk_shift=0, train_sample_shape=None, loss_fn=MSELoss(), torch_metrics=None, likelihood=None, optimizer_cls=<class 'torch.optim.adam.Adam'>, optimizer_kwargs=None, lr_scheduler_cls=None, lr_scheduler_kwargs=None, use_reversible_instance_norm=False)

基类：LightningModule, ABC

基于 PyTorch Lightning 的预测模块。

这个类旨在通过继承来创建一个新的基于 PyTorch Lightning 的预测模块。当子类化这个类时，请确保添加以下具有给定签名的方法：

• PLForecastingModule.__init__()

• PLForecastingModule.forward()

• PLForecastingModule._produce_train_output()

• PLForecastingModule._get_batch_prediction()

在子类 MyModel 的 __init__() 函数中调用 super(MyModel, self).__init__(**kwargs)，其中 kwargs 是 PLForecastingModule 的参数。

参数

• input_chunk_length (int) – 过去作为模型输入的时间步数（每个块）。适用于目标序列，以及过去和/或未来协变量（如果模型支持）。

• output_chunk_length (int) – 内部模型一次预测的时间步数（每个块）。此外，模型输入中使用的未来协变量的未来值数量（如果模型支持未来协变量）。这与 predict() 中使用的预测范围 n 不同，后者是使用一次性或自回归预测生成的所需预测点数。设置 n <= output_chunk_length 可以防止自回归。当协变量不足以延伸到未来，或禁止模型使用过去和/或未来协变量的未来值进行预测时（取决于模型的协变量支持），这很有用。

• train_sample_shape (Optional[Tuple, None]) – 模型的输入形状，用于在不调用 fit_from_dataset 的情况下实例化模型，并对用于重新训练或预测的新训练/推理数据集进行完整性检查。

• loss_fn (_Loss) – 用于训练的 PyTorch 损失函数。如果指定了 likelihood 参数，则此参数将被概率模型忽略。默认值：torch.nn.MSELoss()。

• torch_metrics (Union[Metric, MetricCollection, None]) – 用于评估的 torch 指标或 MetricCollection。可用指标的完整列表可以在 https://torchmetrics.readthedocs.io/en/latest/ 找到。默认值：None。

• likelihood (Optional[Likelihood, None]) – Darts 的 似然模型 之一，用于概率预测。默认值：None。

• optimizer_cls (Optimizer) – 要使用的 PyTorch 优化器类。默认值：torch.optim.Adam。

• optimizer_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 优化器的键值参数（例如，{'lr': 1e-3} 用于指定学习率）。否则将使用所选 optimizer_cls 的默认值。默认值：None。

• lr_scheduler_cls (Optional[_LRScheduler, None]) – 可选地，指定要使用的 PyTorch 学习率调度器类。指定 None 对应于使用恒定的学习率。默认值：None。

• lr_scheduler_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 学习率调度器的键值参数。默认值：None。

• use_reversible_instance_norm (bool) – 是否使用可逆实例归一化 RINorm 来对抗分布偏移，如 [1] 所示。它仅应用于目标序列的特征，而不应用于协变量。

引用

1

T. Kim et al. “Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift”, https://openreview.net/forum?id=cGDAkQo1C0p

属性

automatic_optimization	如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。
current_epoch	Trainer 中的当前时期，如果未附加则为 0。
device_mesh	像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。
example_input_array	示例输入数组是模块在 forward() 方法中可以消费的规范。
global_rank	当前进程在所有节点和设备中的索引。
global_step	所有epoch中看到的总训练批次。
hparams	使用 save_hyperparameters() 保存的超参数集合。
hparams_initial	使用 save_hyperparameters() 保存的超参数集合。
local_rank	当前进程在单个节点中的索引。
logger	Trainer 中对 logger 对象的引用。
loggers	参考 Trainer 中的日志记录器列表。
on_gpu	如果该模型当前位于GPU上，则返回 True。
output_chunk_length	模型一次预测的时间步数。
strict_loading	确定 Lightning 如何使用 .load_state_dict(..., strict=model.strict_loading) 加载此模型。

设备
dtype
已训练的轮数
fabric
supports_probabilistic_prediction
训练者

方法

add_module(name, module)	将一个子模块添加到当前模块。
all_gather(data[, group, sync_grads])	从多个进程中收集张量或张量集合。
apply(fn)	递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。
backward(loss, *args, **kwargs)	在 training_step() 中返回的损失上调用以执行反向传播。
bfloat16()	将所有浮点参数和缓冲区转换为 bfloat16 数据类型。
buffers([recurse])	返回一个遍历模块缓冲区的迭代器。
children()	返回一个遍历直接子模块的迭代器。
clip_gradients(optimizer[, ...])	内部处理梯度裁剪。
compile(*args, **kwargs)	使用 torch.compile() 编译此模块的 forward。
configure_callbacks()	配置模型特定的回调。
configure_gradient_clipping(optimizer[, ...])	对优化器参数执行梯度裁剪。
configure_model()	在策略和精度感知的环境中创建模块的钩子。
configure_optimizers()	配置优化器和学习率调度器以优化模型。
configure_sharded_model()	已弃用。
configure_torch_metrics(torch_metrics)	处理 torch_metrics 参数。
cpu()	参见 torch.nn.Module.cpu()。
cuda([device])	将所有模型参数和缓冲区移动到GPU。
double()	参见 torch.nn.Module.double()。
eval()	将模块设置为评估模式。
extra_repr()	设置模块的额外表示。
float()	参见 torch.nn.Module.float()。
forward(*args, **kwargs)	与 torch.nn.Module.forward() 相同。
freeze()	冻结所有参数以进行推理。
get_buffer(target)	如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。
get_extra_state()	返回要包含在模块的 state_dict 中的任何额外状态。
get_parameter(target)	如果存在，则返回由 target 给出的参数，否则抛出错误。
get_submodule(target)	如果存在，返回由 target 指定的子模块，否则抛出错误。
half()	参见 torch.nn.Module.half()。
ipu([device])	将所有模型参数和缓冲区移动到 IPU。
load_from_checkpoint(checkpoint_path[, ...])	从检查点加载模型的主要方式。
load_state_dict(state_dict[, strict, assign])	将参数和缓冲区从 state_dict 复制到此模块及其后代中。
log(name, value[, prog_bar, logger, ...])	记录一个键值对。
log_dict(dictionary[, prog_bar, logger, ...])	一次性记录一个值的字典。
lr_scheduler_step(scheduler, metric)	重写此方法以调整 Trainer 调用每个调度器的默认方式。
lr_schedulers()	返回在训练期间使用的学习率调度器。
manual_backward(loss, *args, **kwargs)	在进行手动优化时，直接从 training_step() 调用此方法。
modules()	返回网络中所有模块的迭代器。
named_buffers([prefix, recurse, ...])	返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。
named_children()	返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。
named_modules([memo, prefix, remove_duplicate])	返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。
named_parameters([prefix, recurse, ...])	返回一个模块参数的迭代器，同时生成参数的名称和参数本身。
on_after_backward()	在 loss.backward() 之后和优化器步进之前调用。
on_after_batch_transfer(batch, dataloader_idx)	覆盖此方法以在批次转移到设备后对其进行批量增强或修改。
on_before_backward(loss)	在 loss.backward() 之前调用。
on_before_batch_transfer(batch, dataloader_idx)	在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。
on_before_optimizer_step(optimizer)	在 optimizer.step() 之前调用。
on_before_zero_grad(optimizer)	在 training_step() 之后和 optimizer.zero_grad() 之前调用。
on_fit_end()	在拟合的最后调用。
on_fit_start()	在拟合开始时调用。
on_load_checkpoint(checkpoint)	由 Lightning 调用以恢复您的模型。
on_predict_batch_end(outputs, batch, batch_idx)	在批次之后调用预测循环。
on_predict_batch_start(batch, batch_idx[, ...])	在预测循环中，在该批次发生任何事情之前调用。
on_predict_end()	在预测结束时调用。
on_predict_epoch_end()	在预测结束时调用。
on_predict_epoch_start()	在预测开始时调用。
on_predict_model_eval()	在预测循环开始时调用。
on_predict_start()	在预测开始时调用。
on_save_checkpoint(checkpoint)	当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。
on_test_batch_end(outputs, batch, batch_idx)	在批次之后在测试循环中调用。
on_test_batch_start(batch, batch_idx[, ...])	在测试循环中，在该批次发生任何事情之前调用。
on_test_end()	在测试结束时调用。
on_test_epoch_end()	在每个epoch的测试循环结束时调用。
on_test_epoch_start()	在epoch开始时在测试循环中调用。
on_test_model_eval()	在测试循环开始时调用。
on_test_model_train()	在测试循环结束时调用。
on_test_start()	在测试开始时调用。
on_train_batch_end(outputs, batch, batch_idx)	在训练循环中批次处理后调用。
on_train_batch_start(batch, batch_idx)	在训练循环中，在该批次发生任何事情之前调用。
on_train_end()	在训练结束前调用，在记录器实验关闭之前。
on_train_epoch_end()	在训练循环的每个epoch结束时调用。
on_train_epoch_start()	在训练循环中，在epoch开始时调用。
on_train_start()	在完整性检查后，训练开始时调用。
on_validation_batch_end(outputs, batch, ...)	在批次验证循环后调用。
on_validation_batch_start(batch, batch_idx)	在验证循环中，对该批次进行任何操作之前调用。
on_validation_end()	在验证结束时调用。
on_validation_epoch_end()	在epoch结束时的验证循环中调用。
on_validation_epoch_start()	在epoch开始时验证循环中调用。
on_validation_model_eval()	在验证循环开始时调用。
on_validation_model_train()	当验证循环结束时调用。
on_validation_model_zero_grad()	在进入验证循环之前，由训练循环调用来释放梯度。
on_validation_start()	在验证开始时调用。
optimizer_step(epoch, batch_idx, optimizer)	重写此方法以调整 Trainer 调用优化器的默认方式。
optimizer_zero_grad(epoch, batch_idx, optimizer)	重写此方法以改变 optimizer.zero_grad() 的默认行为。
optimizers([use_pl_optimizer])	返回训练期间使用的优化器。
parameters([recurse])	返回一个遍历模块参数的迭代器。
predict_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。
predict_step(batch, batch_idx[, dataloader_idx])	执行预测步骤
prepare_data()	使用此功能下载并准备数据。
print(*args, **kwargs)	仅从进程 0 打印。
register_backward_hook(hook)	在模块上注册一个反向钩子。
register_buffer(name, tensor[, persistent])	向模块添加一个缓冲区。
register_forward_hook(hook, *[, prepend, ...])	在模块上注册一个前向钩子。
register_forward_pre_hook(hook, *[, ...])	在模块上注册一个前向预钩子。
register_full_backward_hook(hook[, prepend])	在模块上注册一个反向钩子。
register_full_backward_pre_hook(hook[, prepend])	在模块上注册一个反向预钩子。
register_load_state_dict_post_hook(hook)	注册一个在模块的 load_state_dict 调用后运行的后置钩子。
register_module(name, module)	别名 add_module()。
register_parameter(name, param)	向模块添加一个参数。
register_state_dict_pre_hook(hook)	为 state_dict() 方法注册一个预钩子。
requires_grad_([requires_grad])	更改是否应在此模块的参数上记录自动梯度操作。
save_hyperparameters(*args[, ignore, frame, ...])	将参数保存到 hparams 属性中。
set_extra_state(state)	设置包含在加载的 state_dict 中的额外状态。
set_predict_parameters(n, num_samples, ...)	在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。
setup(stage)	在 fit（训练 + 验证）、验证、测试或预测开始时调用。
share_memory()	参见 torch.Tensor.share_memory_()。
state_dict(*args[, destination, prefix, ...])	返回一个包含模块整体状态引用的字典。
teardown(stage)	在 fit（训练 + 验证）、验证、测试或预测结束时调用。
test_dataloader()	一个可迭代对象或多个可迭代对象的集合，指定测试样本。
test_step(*args, **kwargs)	对测试集中的单个批次数据进行操作。
to(*args, **kwargs)	参见 torch.nn.Module.to()。
to_dtype(dtype)	将模块精度（默认为 float32）转换为另一种精度。
to_empty(*, device[, recurse])	将参数和缓冲区移动到指定设备，而不复制存储。
to_onnx(file_path[, input_sample])	将模型保存为 ONNX 格式。
to_torchscript([file_path, method, ...])	默认情况下，将整个模型编译为 ScriptModule。
toggle_optimizer(optimizer)	确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。
train([mode])	设置模块为训练模式。
train_dataloader()	指定训练样本的可迭代对象或可迭代对象的集合。
training_step(train_batch, batch_idx)	执行训练步骤
transfer_batch_to_device(batch, device, ...)	如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。
type(dst_type)	参见 torch.nn.Module.type()。
unfreeze()	解冻所有参数以进行训练。
untoggle_optimizer(optimizer)	重置通过 toggle_optimizer() 切换的所需梯度的状态。
val_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。
validation_step(val_batch, batch_idx)	执行验证步骤
xpu([device])	将所有模型参数和缓冲区移动到XPU。
zero_grad([set_to_none])	重置所有模型参数的梯度。

__call__
set_mc_dropout

CHECKPOINT_HYPER_PARAMS_KEY = 'hyper_parameters'

CHECKPOINT_HYPER_PARAMS_NAME = 'hparams_name'

CHECKPOINT_HYPER_PARAMS_TYPE = 'hparams_type'

T_destination

TypeVar(‘T_destination’, bound=Dict[str, Any]) 的别名

add_module(name, module)

将一个子模块添加到当前模块。

可以通过给定的名称将模块作为属性访问。

参数

• name (str) – 子模块的名称。可以使用给定的名称从此模块访问子模块。

• module (Module) – 要添加到模块中的子模块。

返回类型

None

all_gather(data, group=None, sync_grads=False)

从多个进程中收集张量或张量集合。

此方法需要在所有进程上调用，并且所有进程上的张量需要具有相同的形状，否则程序将永远停滞。

参数

• data (Union[Tensor, Dict, List, Tuple]) – int, float, 形状为 (batch, …) 的张量，或可能是嵌套的集合。

• group (Optional[Any, None]) – 要从中收集结果的进程组。默认为所有进程（world）

• sync_grads (bool) – 允许用户为 all_gather 操作同步梯度的标志

返回类型

Union[Tensor, Dict, List, Tuple]

返回

一个形状为 (world_size, batch, …) 的张量，或者如果输入是一个集合，输出也将是一个具有这种形状张量的集合。对于 world_size 为 1 的特殊情况，张量不会增加额外的维度。

apply(fn)

递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。

典型用途包括初始化模型的参数（另见 nn-init-doc）。

参数

fn (Module -> None) – 要应用于每个子模块的函数

返回

自己

返回类型

Module

示例:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

property automatic_optimization: bool

如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。

返回类型

bool

backward(loss, *args, **kwargs)

在 training_step() 中返回的损失上调用以执行反向传播。如果需要，请使用您自己的实现覆盖此钩子。

参数

loss (Tensor) – 由 training_step() 返回的损失张量。如果使用梯度累积，这里的损失值是归一化的（按 1 / 累积步数缩放）。

示例:

def backward(self, loss):
    loss.backward()

返回类型

None

bfloat16()

将所有浮点参数和缓冲区转换为 bfloat16 数据类型。

备注

此方法就地修改模块。

返回

自己

返回类型

Module

buffers(recurse=True)

返回一个遍历模块缓冲区的迭代器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。

生成器

torch.Tensor – 模块缓冲区

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Tensor]

call_super_init: bool = False

children()

返回一个遍历直接子模块的迭代器。

生成器

模块 – 子模块

返回类型

Iterator[Module]

clip_gradients(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

内部处理梯度裁剪。

备注

• 不要重写此方法。如果你想自定义梯度裁剪，考虑使用 configure_gradient_clipping() 方法。

• 对于手动优化（self.automatic_optimization = False），如果你想使用梯度裁剪，可以在训练步骤中手动调用 self.clip_gradients(opt, gradient_clip_val=0.5, gradient_clip_algorithm="norm")。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。

• gradient_clip_algorithm (Optional[str, None]) – 使用的梯度裁剪算法。传递 gradient_clip_algorithm="value" 按值裁剪，传递 gradient_clip_algorithm="norm" 按范数裁剪。

返回类型

None

compile(*args, **kwargs)

使用 torch.compile() 编译此模块的 forward。

此模块的 __call__ 方法被编译，所有参数都按原样传递给 torch.compile()。

有关此函数的参数详情，请参阅 torch.compile()。

configure_callbacks()

配置模型特定的回调。当模型被附加时，例如当调用 .fit() 或 .test() 时，这里返回的回调列表或回调将与传递给 Trainer 的 callbacks 参数的回调列表合并。如果这里返回的回调与 Trainer 的回调列表中已有的一个或多个回调类型相同，它将优先并替换它们。此外，Lightning 将确保 ModelCheckpoint 回调最后运行。

返回类型

Union[Sequence[Callback], Callback]

返回

一个回调函数或回调函数列表，这些回调函数将扩展 Trainer 中的回调函数列表。

示例:

def configure_callbacks(self):
    early_stop = EarlyStopping(monitor="val_acc", mode="max")
    checkpoint = ModelCheckpoint(monitor="val_loss")
    return [early_stop, checkpoint]

configure_gradient_clipping(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

对优化器参数执行梯度裁剪。在调用 optimizer_step() 之前执行。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。默认情况下，传递给 Trainer 的值将在此处可用。

• gradient_clip_algorithm (Optional[str, None]) – 要使用的梯度裁剪算法。默认情况下，传递给 Trainer 的值将在此处可用。

示例:

def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm):
    # Implement your own custom logic to clip gradients
    # You can call `self.clip_gradients` with your settings:
    self.clip_gradients(
        optimizer,
        gradient_clip_val=gradient_clip_val,
        gradient_clip_algorithm=gradient_clip_algorithm
    )

返回类型

None

configure_model()

在策略和精度感知的环境中创建模块的钩子。

这在使用分片策略（FSDP 和 DeepSpeed）时特别有用，我们希望立即分片模型以节省内存和初始化时间。对于非分片策略，您可以选择覆盖此钩子或在 init_module() 上下文管理器下初始化您的模型。

此钩子在每个 fit/val/test/predict 阶段都会被调用，因此请确保此钩子的实现是 幂等的，即，在第一次调用钩子后，后续的调用应该是无操作的。

返回类型

None

configure_optimizers()

配置优化器和学习率调度器以优化模型。

configure_sharded_model()

已弃用。

使用 configure_model() 代替。

返回类型

None

static configure_torch_metrics(torch_metrics)

处理 torch_metrics 参数。

返回类型

MetricCollection

cpu()

参见 torch.nn.Module.cpu()。

返回类型

Self

cuda(device=None)

将所有模型参数和缓冲区移动到GPU。这也会使相关的参数和缓冲区成为不同的对象。因此，如果模块在优化期间将驻留在GPU上，则应在构造优化器之前调用此方法。

参数

device (Union[device, int, None]) – 如果指定，所有参数将被复制到该设备。如果为 None，将使用当前的 CUDA 设备索引。

返回

自己

返回类型

Module

property current_epoch: int

Trainer 中的当前时期，如果未附加则为 0。

返回类型

int

property device: device

返回类型

device

property device_mesh: Optional[DeviceMesh]

像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。

返回类型

Optional[ForwardRef, None]

double()

参见 torch.nn.Module.double()。

返回类型

Self

property dtype: Union[str, dtype]

返回类型

Union[str, dtype]

dump_patches: bool = False

property epochs_trained

eval()

将模块设置为评估模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

这等同于 self.train(False)。

请参阅 locally-disable-grad-doc 以比较 .eval() 和几个可能与之混淆的类似机制。

返回

自己

返回类型

Module

property example_input_array: Optional[Union[Tensor, Tuple, Dict]]

示例输入数组是模块在 forward() 方法中可以使用的规范。返回类型解释如下：

• 单个张量：假设模型接受一个参数，即 model.forward(model.example_input_array)

• 元组：输入数组应被解释为位置参数的序列，即 model.forward(*model.example_input_array)

• 字典：输入数组表示命名关键字参数，即 model.forward(**model.example_input_array)

返回类型

Union[Tensor, Tuple, Dict, None]

extra_repr()

设置模块的额外表示。

要打印自定义的额外信息，您应该在您自己的模块中重新实现此方法。单行和多行字符串都是可接受的。

返回类型

str

property fabric: Optional[Fabric]

返回类型

Optional[Fabric, None]

float()

参见 torch.nn.Module.float()。

返回类型

Self

abstract forward(*args, **kwargs)

与 torch.nn.Module.forward() 相同。

参数

• *args – 无论你决定传入 forward 方法的内容。

• **kwargs – 关键字参数也是可能的。

返回类型

Any

返回

你的模型的输出

freeze()

冻结所有参数以进行推理。

示例:

model = MyLightningModule(...)
model.freeze()

返回类型

None

get_buffer(target)

如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的缓冲区的完全限定字符串名称。（有关如何指定完全限定字符串，请参见 get_submodule。）

返回

target 引用的缓冲区

返回类型

torch.Tensor

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是缓冲区的内容

get_extra_state()

返回要包含在模块的 state_dict 中的任何额外状态。

如果需要存储额外状态，请在你的模块中实现此功能以及相应的 set_extra_state()。此函数在构建模块的 state_dict() 时被调用。

请注意，额外的状态应该是可pickle的，以确保state_dict的序列化工作。我们仅对序列化张量提供向后兼容性保证；如果其他对象的序列化pickle形式发生变化，它们可能会破坏向后兼容性。

返回

模块的 state_dict 中存储的任何额外状态

返回类型

object

get_parameter(target)

如果存在，则返回由 target 给出的参数，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的参数的全限定字符串名称。（关于如何指定全限定字符串，请参见 get_submodule。）

返回

target 引用的参数

返回类型

torch.nn.Parameter

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Parameter 的内容

get_submodule(target)

如果存在，返回由 target 指定的子模块，否则抛出错误。

例如，假设你有一个 nn.Module A ，它看起来像这样：

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(该图展示了一个 nn.Module A。A 有一个嵌套的子模块 net_b，它本身有两个子模块 net_c 和 linear。net_c 然后有一个子模块 conv。)

要检查我们是否拥有 linear 子模块，我们会调用 get_submodule("net_b.linear")。要检查我们是否拥有 conv 子模块，我们会调用 get_submodule("net_b.net_c.conv")。

get_submodule 的运行时间受限于 target 中的模块嵌套程度。对 named_modules 的查询可以实现相同的结果，但它是 O(N) 的，其中 N 是传递模块的数量。因此，对于检查某个子模块是否存在的简单检查，应始终使用 get_submodule。

参数

target (str) – 要查找的子模块的完全限定字符串名称。（请参见上面的示例，了解如何指定完全限定的字符串。）

返回

由 target 引用的子模块

返回类型

torch.nn.Module

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Module 的内容

property global_rank: int

当前进程在所有节点和设备中的索引。

返回类型

int

property global_step: int

所有epoch中看到的总训练批次。

如果没有连接训练器，此属性为 0。

返回类型

int

half()

参见 torch.nn.Module.half()。

返回类型

Self

property hparams: Union[AttributeDict, MutableMapping]

通过 save_hyperparameters() 保存的超参数集合。用户可以对其进行修改。对于初始超参数的冻结集，请使用 hparams_initial。

返回类型

Union[AttributeDict, MutableMapping]

返回

可变超参数字典

property hparams_initial: AttributeDict

通过 save_hyperparameters() 保存的超参数集合。这些内容是只读的。可以通过 hparams 手动更新保存的超参数。

返回

不可变的初始超参数

返回类型

AttributeDict

ipu(device=None)

将所有模型参数和缓冲区移动到 IPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在IPU上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs)

从检查点加载模型的主要方式。当Lightning保存检查点时，它会将传递给 __init__ 的参数存储在检查点中的 "hyper_parameters" 下。

通过 **kwargs 指定的任何参数将覆盖存储在 "hyper_parameters" 中的参数。

参数

• checkpoint_path (Union[str, Path, IO]) – 检查点的路径。这也可以是一个URL，或类似文件的对象。

• map_location (Union[device, str, int, Callable[[UntypedStorage, str], Optional[UntypedStorage, None]], Dict[Union[device, str, int], Union[device, str, int]], None]) – 如果你的检查点保存了一个GPU模型，而你现在在CPU上或不同数量的GPU上加载，使用这个来映射到新的设置。其行为与 torch.load() 相同。

• hparams_file (Union[str, Path, None]) –

  可选路径到一个 .yaml 或 .csv 文件，其层次结构如本示例所示:

  drop_prob: 0.2
  dataloader:
      batch_size: 32
  

  您很可能不需要这个，因为 Lightning 总是会将超参数保存到检查点。然而，如果您的检查点权重没有保存超参数，请使用此方法传入一个包含您希望使用的超参数的 .yaml 文件。这些将被转换为 dict 并传递到您的 LightningModule 中使用。

  如果你的模型的 hparams 参数是 Namespace 并且 .yaml 文件具有层次结构，你需要重构你的模型以将 hparams 视为 dict。

• strict (Optional[bool, None]) – 是否严格要求 checkpoint_path 中的键与该模块状态字典返回的键匹配。默认为 True ，除非设置了 LightningModule.strict_loading ，在这种情况下，它默认为 LightningModule.strict_loading 的值。

• **kwargs – 初始化模型所需的任何额外关键字参数。也可以用来覆盖已保存的超参数值。

返回类型

Self

返回

LightningModule 实例，加载了权重和超参数（如果可用）。

备注

load_from_checkpoint 是一个 类 方法。你应该使用你的 LightningModule 类 来调用它，而不是 LightningModule 实例，否则会引发 TypeError。

备注

为了确保所有层都能从检查点加载，如果您的 LightningModule 中重写了此钩子，此函数将在实例化模型后直接调用 configure_model()。但是，请注意 load_from_checkpoint 不支持加载分片检查点，如果模型太大，您可能会耗尽内存。在这种情况下，请考虑通过 .fit(ckpt_path=...) 通过 Trainer 加载。

示例:

# load weights without mapping ...
model = MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')

# or load weights mapping all weights from GPU 1 to GPU 0 ...
map_location = {'cuda:1':'cuda:0'}
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    map_location=map_location
)

# or load weights and hyperparameters from separate files.
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    hparams_file='/path/to/hparams_file.yaml'
)

# override some of the params with new values
model = MyLightningModule.load_from_checkpoint(
    PATH,
    num_layers=128,
    pretrained_ckpt_path=NEW_PATH,
)

# predict
pretrained_model.eval()
pretrained_model.freeze()
y_hat = pretrained_model(x)

load_state_dict(state_dict, strict=True, assign=False)

将参数和缓冲区从 state_dict 复制到此模块及其后代中。

如果 strict 是 True，那么 state_dict 的键必须与该模块的 state_dict() 函数返回的键完全匹配。

警告

如果 assign 是 True，优化器必须在调用 load_state_dict 之后创建，除非 get_swap_module_params_on_conversion() 是 True。

参数

• state_dict (dict) – 包含参数和持久缓冲区的字典。

• strict (bool, optional) – 是否严格要求 state_dict 中的键与该模块的 state_dict() 函数返回的键匹配。默认值：True

• assign (bool, optional) – 当 False 时，当前模块中张量的属性会被保留，而当 True 时，状态字典中张量的属性会被保留。唯一的例外是 Parameter 的 requires_grad 字段，该字段的值会从模块中保留。默认值：False

返回

• missing_keys 是一个包含任何预期键的 str 列表

  此模块中存在但未在提供的 state_dict 中出现的。

• unexpected_keys 是一个包含未预料到的键的 str 列表。

  此模块预期但未在提供的 state_dict 中出现。

返回类型

带有 missing_keys 和 unexpected_keys 字段的 NamedTuple

备注

如果一个参数或缓冲区被注册为 None ，并且其对应的键存在于 state_dict 中，load_state_dict() 将引发一个 RuntimeError 。

property local_rank: int

当前进程在单个节点中的索引。

返回类型

int

log(name, value, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, metric_attribute=None, rank_zero_only=False)

记录一个键值对。

示例:

self.log('train_loss', loss)

每个钩子的默认行为在此处有文档说明：extensions/logging:自动日志记录。

参数

• name (str) – 日志的键。如果使用DDP或其他分布式策略，必须在所有进程中保持一致。

• value (Union[Metric, Tensor, int, float]) – 要记录的值。可以是 float、Tensor 或 Metric。

• prog_bar (bool) – 如果 True 则记录到进度条。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• on_epoch (Optional[bool, None]) – 如果 True ，记录每个epoch累积的指标。默认值由钩子决定。详情请参见 extensions/logging:自动日志记录 。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图表。

• sync_dist (bool) – 如果 True，将在设备间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的DDP组。

• add_dataloader_idx (bool) – 如果 True，将当前数据加载器的索引附加到名称（当使用多个数据加载器时）。如果为 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的 batch_size。这将直接从加载的批次中推断出来，但对于某些数据结构，您可能需要显式提供它。

• metric_attribute (Optional[str, None]) – 要恢复指标状态，Lightning 需要模型中 torchmetrics.Metric 的引用。如果它是模型属性，则会自动找到。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

log_dict(dictionary, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, rank_zero_only=False)

一次性记录一个值的字典。

示例:

values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n}
self.log_dict(values)

参数

• dictionary (Union[Mapping[str, Union[Metric, Tensor, int, float]], MetricCollection]) – 键值对。如果使用DDP或其他分布式策略，键必须在所有进程中保持一致。值可以是 float 、 Tensor 、 Metric 或 MetricCollection 。

• prog_bar (bool) – 如果 True 则记录到进度基础。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。None 自动记录训练步骤，但不记录验证/测试步骤。默认值由钩子决定。详情请参阅 extensions/logging:自动记录。

• on_epoch (Optional[bool, None]) – 如果 True 记录累积的 epoch 指标。None 自动记录 val/test 步骤的日志，但不包括 training_step。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图形

• sync_dist (bool) – 如果 True，将在GPU/TPU之间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的 ddp 组。

• add_dataloader_idx (bool) – 如果 True，则将当前数据加载器的索引附加到名称（在使用多个数据加载器时）。如果 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的批次大小。这将直接从加载的批次中推断出来，但某些数据结构可能需要显式提供它。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

property logger: Optional[Union[Logger, Logger]]

Trainer 中对 logger 对象的引用。

返回类型

Union[Logger, Logger, None]

property loggers: Union[List[Logger], List[Logger]]

参考 Trainer 中的日志记录器列表。

返回类型

Union[List[Logger], List[Logger]]

lr_scheduler_step(scheduler, metric)

重写此方法以调整 Trainer 调用每个调度器的默认方式。默认情况下，Lightning 会根据每个调度器的 interval 调用 step() 方法，如示例所示。

参数

• scheduler (Union[LRScheduler, ReduceLROnPlateau]) – 学习率调度器。

• metric (Optional[Any, None]) – 用于 ReduceLROnPlateau 等调度器的监视器值。

示例:

# DEFAULT
def lr_scheduler_step(self, scheduler, metric):
    if metric is None:
        scheduler.step()
    else:
        scheduler.step(metric)

# Alternative way to update schedulers if it requires an epoch value
def lr_scheduler_step(self, scheduler, metric):
    scheduler.step(epoch=self.current_epoch)

返回类型

None

lr_schedulers()

返回在训练期间使用的学习率调度器。对于手动优化很有用。

返回类型

Union[None, List[Union[LRScheduler, ReduceLROnPlateau]], LRScheduler, ReduceLROnPlateau]

返回

一个调度器，或者在存在多个调度器的情况下是一个调度器列表，如果在 configure_optimizers() 中没有返回调度器，则为 None。

manual_backward(loss, *args, **kwargs)

当你手动进行优化时，请直接从你的 training_step() 调用此方法。通过使用此方法，Lightning 可以确保在使用混合精度时应用所有适当的缩放。

更多示例请参见 手动优化。

示例:

def training_step(...):
    opt = self.optimizers()
    loss = ...
    opt.zero_grad()
    # automatically applies scaling, etc...
    self.manual_backward(loss)
    opt.step()

参数

• loss (Tensor) – 要计算梯度的张量。必须附加一个图。

• *args – 要转发给 backward() 的附加位置参数

• **kwargs – 要转发给 backward() 的额外关键字参数

返回类型

None

modules()

返回网络中所有模块的迭代器。

生成器

模块 – 网络中的一个模块

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

返回类型

Iterator[Module]

named_buffers(prefix='', recurse=True, remove_duplicate=True)

返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。

参数

• prefix (str) – 要添加到所有缓冲区名称的前缀。

• recurse (bool, optional) – 如果为 True，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。默认为 True。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的缓冲区。默认为 True。

生成器

(str, torch.Tensor) – 包含名称和缓冲区的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

返回类型

Iterator[Tuple[str, Tensor]]

named_children()

返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。

生成器

(str, Module) – 包含名称和子模块的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

返回类型

Iterator[Tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)

返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。

参数

• memo (Optional[Set[Module], None]) – 一个备忘录，用于存储已添加到结果中的模块集合

• prefix (str) – 一个将添加到模块名称前的前缀

• remove_duplicate (bool) – 是否在结果中移除重复的模块实例

生成器

(str, Module) – 名称和模块的元组

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)

返回一个模块参数的迭代器，同时生成参数的名称和参数本身。

参数

• prefix (str) – 要添加到所有参数名称前的前缀。

• recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的参数。默认为 True。

生成器

(str, Parameter) – 包含名称和参数的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

返回类型

Iterator[Tuple[str, Parameter]]

on_after_backward()

在 loss.backward() 之后和优化器步进之前调用。

备注

如果使用原生 AMP，梯度在此刻将不会被缩放。如果你需要未缩放的梯度，请使用 on_before_optimizer_step。

返回类型

None

on_after_batch_transfer(batch, dataloader_idx)

覆盖此方法以在批次转移到设备后对其进行批量增强或修改。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_after_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = gpu_transforms(batch['x'])
    return batch

参见

• on_before_batch_transfer()

• transfer_batch_to_device()

on_before_backward(loss)

在 loss.backward() 之前调用。

参数

loss (Tensor) – 损失除以梯度累积的批次数量，如果使用AMP则进行缩放。

返回类型

None

on_before_batch_transfer(batch, dataloader_idx)

在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_before_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = transforms(batch['x'])
    return batch

参见

• on_after_batch_transfer()

• transfer_batch_to_device()

on_before_optimizer_step(optimizer)

在 optimizer.step() 之前调用。

如果使用梯度累积，钩子会在梯度累积完成后被调用。参见：:paramref:`~pytorch_lightning.trainer.trainer.Trainer.accumulate_grad_batches`。

如果使用AMP，损失将在调用此钩子之前未缩放。有关梯度缩放的更多信息，请参阅这些 文档。

如果裁剪梯度，梯度尚未被裁剪。

参数

optimizer (Optimizer) – 当前使用的优化器。

示例:

def on_before_optimizer_step(self, optimizer):
    # example to inspect gradient information in tensorboard
    if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
        for k, v in self.named_parameters():
            self.logger.experiment.add_histogram(
                tag=k, values=v.grad, global_step=self.trainer.global_step
            )

返回类型

None

on_before_zero_grad(optimizer)

在 training_step() 之后和 optimizer.zero_grad() 之前调用。

在训练循环中，调用优化器步骤之后和清零梯度之前。这是一个检查更新后权重信息的好地方。

这是调用它的地方:

for optimizer in optimizers:
    out = training_step(...)

    model.on_before_zero_grad(optimizer) # < ---- called here
    optimizer.zero_grad()

    backward()

参数

optimizer (Optimizer) – 应该被置零的梯度优化器。

返回类型

None

on_fit_end()

在拟合的最后调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

on_fit_start()

在拟合开始时调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

property on_gpu: bool

如果该模型当前位于GPU上，则返回 True。

在LightningModule周围设置标志以区分CPU和GPU行为非常有用。

返回类型

bool

on_load_checkpoint(checkpoint)

由 Lightning 调用以恢复您的模型。如果您使用 on_save_checkpoint() 保存了某些内容，这是您恢复这些内容的机会。

参数

checkpoint (Dict[str, Any]) – 加载的检查点

示例:

def on_load_checkpoint(self, checkpoint):
    # 99% of the time you don't need to implement this method
    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']

备注

Lightning 自动恢复全局步数、轮次和训练状态，包括 amp 缩放。您无需恢复任何与训练相关的内容。

返回类型

None

on_predict_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后调用预测循环。

参数

• outputs (Optional[Any, None]) – predict_step(x) 的输出

• batch (Any) – 由预测 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_batch_start(batch, batch_idx, dataloader_idx=0)

在预测循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_start()

在预测开始时调用。

返回类型

None

on_predict_model_eval()

在预测循环开始时调用。

默认情况下，预测循环在开始之前会在 LightningModule 上调用 .eval()。重写此钩子以更改行为。

返回类型

None

on_predict_start()

在预测开始时调用。

返回类型

None

on_save_checkpoint(checkpoint)

当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。

参数

checkpoint (Dict[str, Any]) – 在转储到文件之前的完整检查点字典。此钩子的实现可以将额外数据插入此字典中。

示例:

def on_save_checkpoint(self, checkpoint):
    # 99% of use cases you don't need to implement this method
    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object

备注

Lightning 保存了训练的所有方面（如 epoch、全局步数等），包括 amp 缩放。您无需存储任何关于训练的信息。

返回类型

None

on_test_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后在测试循环中调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – test_step(x) 的输出

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_batch_start(batch, batch_idx, dataloader_idx=0)

在测试循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_end()

在测试结束时调用。

返回类型

None

on_test_epoch_end()

在每个epoch的测试循环结束时调用。

返回类型

None

on_test_epoch_start()

在epoch开始时在测试循环中调用。

返回类型

None

on_test_model_eval()

在测试循环开始时调用。

默认情况下，测试循环在开始前会调用 .eval() 方法在 LightningModule 上。重写此钩子以改变行为。另请参见 on_test_model_train()。

返回类型

None

on_test_model_train()

在测试循环结束时调用。

测试循环默认会将 LightningModule 的 training 模式恢复到开始测试之前的状态。覆盖此钩子以更改行为。另请参见 on_test_model_eval()。

返回类型

None

on_test_start()

在测试开始时调用。

返回类型

None

on_train_batch_end(outputs, batch, batch_idx)

在训练循环中批次处理后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – 训练步骤(x)的输出

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

备注

这里的 outputs["loss"] 值将是相对于 accumulate_grad_batches 的损失归一化值，该损失从 training_step 返回。

返回类型

None

on_train_batch_start(batch, batch_idx)

在训练循环中，在该批次发生任何事情之前调用。

如果你在这里返回 -1，你将跳过当前轮次的剩余训练。

参数

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

返回类型

Optional[int, None]

on_train_end()

在训练结束前调用，在记录器实验关闭之前。

返回类型

None

on_train_epoch_end()

在训练循环的每个epoch结束时调用。

要在每个epoch结束时访问所有批处理输出，您可以将步骤输出缓存为 LightningModule 的一个属性，并在此钩子中访问它们：

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # free up the memory
        self.training_step_outputs.clear()

on_train_epoch_start()

在训练循环中，在epoch开始时调用。

返回类型

None

on_train_start()

在完整性检查后，训练开始时调用。

返回类型

None

on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次验证循环后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – validation_step(x) 的输出

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_batch_start(batch, batch_idx, dataloader_idx=0)

在验证循环中，对该批次进行任何操作之前调用。

参数

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_end()

在验证结束时调用。

返回类型

None

on_validation_epoch_end()

在epoch结束时的验证循环中调用。

on_validation_epoch_start()

在epoch开始时验证循环中调用。

返回类型

None

on_validation_model_eval()

在验证循环开始时调用。

默认情况下，验证循环在开始前会在 LightningModule 上调用 .eval()。覆盖此钩子以更改行为。另请参见 on_validation_model_train()。

返回类型

None

on_validation_model_train()

当验证循环结束时调用。

默认情况下，验证循环会将 LightningModule 的 training 模式恢复到开始验证之前的状态。重写此钩子以更改行为。另请参见 on_validation_model_eval()。

返回类型

None

on_validation_model_zero_grad()

在进入验证循环之前，由训练循环调用来释放梯度。

返回类型

None

on_validation_start()

在验证开始时调用。

返回类型

None

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)

重写此方法以调整 Trainer 调用优化器的默认方式。

默认情况下，Lightning 会调用 step() 和 zero_grad()，如示例所示。当 Trainer(accumulate_grad_batches != 1) 时，在累积阶段不会调用此方法（和 zero_grad()）。手动优化时，重写此钩子没有好处。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Union[Optimizer, LightningOptimizer]) – 一个 PyTorch 优化器

• optimizer_closure (Optional[Callable[[], Any], None]) – 优化器闭包。此闭包必须执行，因为它包含了 training_step()、optimizer.zero_grad() 和 backward() 的调用。

示例:

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    # Add your custom logic to run directly before `optimizer.step()`

    optimizer.step(closure=optimizer_closure)

    # Add your custom logic to run directly after `optimizer.step()`

返回类型

None

optimizer_zero_grad(epoch, batch_idx, optimizer)

重写此方法以改变 optimizer.zero_grad() 的默认行为。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Optimizer) – 一个 PyTorch 优化器

示例:

# DEFAULT
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad()

# Set gradients to `None` instead of zero to improve performance (not required on `torch>=2.0.0`).
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad(set_to_none=True)

参见 torch.optim.Optimizer.zero_grad() 以了解上述示例的解释。

返回类型

None

optimizers(use_pl_optimizer=True)

返回训练期间使用的优化器。对于手动优化很有用。

参数

use_pl_optimizer (bool) – 如果 True，将会把优化器包装在 LightningOptimizer 中，以自动处理精度、分析和步数调用的计数，以进行适当的日志记录和检查点保存。它特别包装了 step 方法，不支持没有此方法的自定义优化器。

返回类型

Union[Optimizer, LightningOptimizer, _FabricOptimizer, List[Optimizer], List[LightningOptimizer], List[_FabricOptimizer]]

返回

一个单一的优化器，或者在存在多个优化器的情况下，一个优化器列表。

property output_chunk_length: Optional[int]

模型一次预测的时间步数。

返回类型

Optional[int, None]

parameters(recurse=True)

返回一个遍历模块参数的迭代器。

这通常传递给一个优化器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

生成器

参数 – 模块参数

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Parameter]

predict_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。

有关多个数据加载器的更多信息，请参阅此 部分。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• predict()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

返回

一个 torch.utils.data.DataLoader 或它们的序列，指定预测样本。

predict_step(batch, batch_idx, dataloader_idx=None)

执行预测步骤

批处理

Darts 的 InferenceDataset 的输出 - 包含 (过去的目标, 过去的协变量, 历史未来协变量, 未来协变量, 未来过去协变量, 输入时间序列, 预测开始时间步) 的元组

batch_idx

当前批次的批次索引

dataloader_idx

数据加载器索引

返回类型

Sequence[TimeSeries]

prepare_data()

使用此方法下载和准备数据。使用多个进程（分布式设置）下载和保存数据将导致数据损坏。Lightning 确保此方法仅在一个进程内调用，因此您可以安全地在其中添加下载逻辑。

警告

不要将状态设置到模型中（使用 setup 代替），因为这不会在每个设备上被调用

示例:

def prepare_data(self):
    # good
    download_data()
    tokenize()
    etc()

    # bad
    self.split = data_split
    self.some_state = some_other_state()

在分布式环境中，prepare_data 可以通过两种方式调用（使用 prepare_data_per_node）

1. 每个节点一次。这是默认设置，仅在 LOCAL_RANK=0 时调用。

2. 总共一次。仅在 GLOBAL_RANK=0 时调用。

示例:

# DEFAULT
# called once per node on LOCAL_RANK=0 of that node
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = True

# call on GLOBAL_RANK=0 (great for shared file systems)
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = False

在请求数据加载器之前调用此方法：

model.prepare_data()
initialize_distributed()
model.setup(stage)
model.train_dataloader()
model.val_dataloader()
model.test_dataloader()
model.predict_dataloader()

返回类型

None

print(*args, **kwargs)

仅从进程 0 打印。在任何分布式模式下使用此功能以仅记录一次。

参数

• *args – 要打印的内容。与Python内置的print函数相同。

• **kwargs – 与Python内置的print函数相同。

示例:

def forward(self, x):
    self.print(x, 'in forward')

返回类型

None

register_backward_hook(hook)

在模块上注册一个反向钩子。

此函数已被 register_full_backward_hook() 取代，并且此函数的行为将在未来版本中发生变化。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_buffer(name, tensor, persistent=True)

向模块添加一个缓冲区。

这通常用于注册一个不应被视为模型参数的缓冲区。例如，BatchNorm 的 running_mean 不是参数，而是模块状态的一部分。缓冲区默认是持久的，并将与参数一起保存。可以通过将 persistent 设置为 False 来改变此行为。持久缓冲区和非持久缓冲区之间的唯一区别是后者不会成为此模块的 state_dict 的一部分。

缓冲区可以通过给定的名称作为属性访问。

参数

• name (str) – 缓冲区的名称。可以通过给定的名称从此模块访问缓冲区。

• tensor (Tensor or None) – 要注册的缓冲区。如果为 None ，则运行在缓冲区上的操作，例如 cuda ，将被忽略。如果为 None ，该缓冲区将 不 包含在模块的 state_dict 中。

• persistent (bool) – 缓冲区是否是此模块的 state_dict 的一部分。

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

返回类型

None

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

在模块上注册一个前向钩子。

钩子将在每次 forward() 计算输出后被调用。

如果 with_kwargs 是 False 或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward 。钩子可以修改输出。它可以就地修改输入，但这不会影响前向传播，因为这是在调用 forward() 之后调用的。钩子应具有以下签名:

hook(module, args, output) -> None or modified output

如果 with_kwargs 是 True，前向钩子将被传递给前向函数的 kwargs，并期望返回可能被修改的输出。钩子应具有以下签名:

hook(module, args, kwargs, output) -> None or modified output

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果 True，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之后触发。请注意，使用 register_module_forward_hook() 注册的全局 forward 钩子将在通过此方法注册的所有钩子之前触发。默认值：False

• with_kwargs (bool) – 如果 True，hook 将会传递给 forward 函数的 kwargs。默认值：False

• always_call (bool) – 如果 True ，无论在调用模块时是否引发异常，都将运行 hook 。默认值： False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

在模块上注册一个前向预钩子。

钩子将在每次调用 forward() 之前被调用。

如果 with_kwargs 为假或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward。钩子可以修改输入。用户可以在钩子中返回一个元组或单个修改后的值。如果返回单个值（除非该值已经是元组），我们将把该值包装成一个元组。钩子应具有以下签名:

hook(module, args) -> None or modified input

如果 with_kwargs 为真，前向预钩子将传递给前向函数的 kwargs。如果钩子修改了输入，则应返回 args 和 kwargs。钩子应具有以下签名:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之后触发。请注意，通过 register_module_forward_pre_hook() 注册的全局 forward_pre 钩子将在所有通过此方法注册的钩子之前触发。默认值：False

• with_kwargs (bool) – 如果为真，hook 将被传递给 forward 函数的 kwargs。默认值：False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)

在模块上注册一个反向钩子。

每次计算与模块相关的梯度时，钩子都会被调用，即当且仅当计算与模块输出相关的梯度时，钩子才会执行。钩子应具有以下签名:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input 和 grad_output 是包含相对于输入和输出的梯度的元组。钩子不应修改其参数，但可以选择返回一个新的相对于输入的梯度，该梯度将用于替代 grad_input 在后续计算中。grad_input 将仅对应于作为位置参数给出的输入，所有 kwarg 参数都将被忽略。grad_input 和 grad_output 中的条目将为 None 用于所有非 Tensor 参数。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入或输出，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之后触发。请注意，通过 register_module_full_backward_hook() 注册的全局 backward 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

在模块上注册一个反向预钩子。

每次计算模块的梯度时，都会调用该钩子。钩子应具有以下签名:

hook(module, grad_output) -> tuple[Tensor] or None

The grad_output 是一个元组。钩子不应修改其参数，但它可以选择返回一个新的相对于输出的梯度，该梯度将用于替代 grad_output 在后续计算中。grad_output 中的条目对于所有非 Tensor 参数将为 None。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之后触发。请注意，通过 register_module_full_backward_pre_hook() 注册的全局 backward_pre 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)

注册一个在模块的 load_state_dict 调用后运行的后置钩子。

它应该有以下签名:

hook(module, incompatible_keys) -> None

module 参数是当前注册此钩子的模块，而 incompatible_keys 参数是一个由属性 missing_keys 和 unexpected_keys 组成的 NamedTuple。missing_keys 是一个包含缺失键的 str 的 list，而 unexpected_keys 是一个包含意外键的 str 的 list。

如果需要，可以就地修改给定的 incompatible_keys。

请注意，当使用 strict=True 调用 load_state_dict() 时执行的检查会受到钩子对 missing_keys 或 unexpected_keys 所做修改的影响，这是预期的。对任一组键的添加将导致在 strict=True 时抛出错误，而清除掉缺失和意外的键将避免错误。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_module(name, module)

别名 add_module()。

返回类型

None

register_parameter(name, param)

向模块添加一个参数。

该参数可以通过给定的名称作为属性访问。

参数

• name (str) – 参数的名称。可以使用给定的名称从此模块访问该参数。

• param (Parameter or None) – 要添加到模块的参数。如果为 None ，则运行在参数上的操作，例如 cuda ，将被忽略。如果为 None ，该参数将 不 包含在模块的 state_dict 中。

返回类型

None

register_state_dict_pre_hook(hook)

为 state_dict() 方法注册一个预钩子。

这些钩子将在调用 self 上的 state_dict 之前，使用参数 self、prefix 和 keep_vars 进行调用。注册的钩子可以用于在调用 state_dict 之前执行预处理。

requires_grad_(requires_grad=True)

更改是否应在此模块的参数上记录自动梯度操作。

此方法就地设置参数的 requires_grad 属性。

这种方法有助于冻结模块的一部分以进行微调或单独训练模型的一部分（例如，GAN训练）。

请参阅 locally-disable-grad-doc 以比较 .requires_grad_() 和几个可能与之混淆的类似机制。

参数

requires_grad (bool) – 是否应在此模块的参数上记录自动求导操作。默认值：True。

返回

自己

返回类型

Module

save_hyperparameters(*args, ignore=None, frame=None, logger=True)

将参数保存到 hparams 属性中。

参数

• args (Any) – dict、NameSpace 或 OmegaConf 的单个对象，或来自类 __init__ 的字符串名称或参数

• ignore (Union[Sequence[str], str, None]) – 一个参数名称或从类 __init__ 中忽略的参数名称列表

• frame (Optional[frame, None]) – 一个框架对象。默认是 None

• logger (bool) – 是否将超参数发送到记录器。默认值：True

示例:

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # manually assign arguments
...         self.save_hyperparameters('arg1', 'arg3')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class AutomaticArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # equivalent automatic
...         self.save_hyperparameters()
...     def forward(self, *args, **kwargs):
...         ...
>>> model = AutomaticArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg2": abc
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class SingleArgModel(HyperparametersMixin):
...     def __init__(self, params):
...         super().__init__()
...         # manually assign single argument
...         self.save_hyperparameters(params)
...     def forward(self, *args, **kwargs):
...         ...
>>> model = SingleArgModel(Namespace(p1=1, p2='abc', p3=3.14))
>>> model.hparams
"p1": 1
"p2": abc
"p3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # pass argument(s) to ignore as a string or in a list
...         self.save_hyperparameters(ignore='arg2')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

返回类型

None

set_extra_state(state)

设置包含在加载的 state_dict 中的额外状态。

此函数从 load_state_dict() 调用，以处理在 state_dict 中找到的任何额外状态。如果你的模块需要在 state_dict 中存储额外状态，请实现此函数以及相应的 get_extra_state()。

参数

state (dict) – 来自 state_dict 的额外状态

返回类型

None

set_mc_dropout(active)

set_predict_parameters(n, num_samples, roll_size, batch_size, n_jobs, predict_likelihood_parameters, mc_dropout)

在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。

返回类型

None

setup(stage)

在 fit（训练 + 验证）、验证、测试或预测开始时调用。当你需要动态构建模型或调整模型时，这是一个很好的钩子。在使用 DDP 时，此钩子会在每个进程上调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

示例:

class LitModel(...):
    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

        # don't do this
        self.something = else

    def setup(self, stage):
        data = load_data(...)
        self.l1 = nn.Linear(28, data.num_classes)

返回类型

None

share_memory()

参见 torch.Tensor.share_memory_()。

返回类型

~T

state_dict(*args, destination=None, prefix='', keep_vars=False)

返回一个包含模块整体状态引用的字典。

参数和持久缓冲区（例如运行平均值）都包含在内。键是对应的参数和缓冲区名称。设置为 None 的参数和缓冲区不包含在内。

备注

返回的对象是一个浅拷贝。它包含对模块参数和缓冲区的引用。

警告

目前，state_dict() 也按顺序接受 destination、prefix 和 keep_vars 的位置参数。然而，这正在被弃用，未来版本将强制使用关键字参数。

警告

请避免使用参数 destination，因为它不是为最终用户设计的。

参数

• destination (dict, optional) – 如果提供了，模块的状态将被更新到字典中，并返回相同的对象。否则，将创建并返回一个 OrderedDict。默认值：None。

• prefix (str, optional) – 在 state_dict 中组合键时添加到参数和缓冲区名称的前缀。默认值：''。

• keep_vars (bool, optional) – 默认情况下，状态字典中返回的 Tensor 是与自动求导分离的。如果设置为 True，则不会执行分离操作。默认值：False。

返回

包含模块整个状态的字典

返回类型

dict

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

property strict_loading: bool

确定 Lightning 如何使用 .load_state_dict(…, strict=model.strict_loading) 加载此模型。

返回类型

bool

property supports_probabilistic_prediction: bool

返回类型

bool

teardown(stage)

在 fit（训练 + 验证）、验证、测试或预测结束时调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

返回类型

None

test_dataloader()

一个可迭代对象或多个可迭代对象的集合，指定测试样本。

有关多个数据加载器的更多信息，请参阅此 部分。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• test()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要测试数据集和 test_step() ，你不需要实现这个方法。

返回类型

Any

test_step(*args, **kwargs)

对测试集中的单个批次数据进行操作。在这一步中，您通常会生成示例或计算任何感兴趣的内容，例如准确性。

参数

• batch – 你的数据可迭代对象的输出，通常是一个 DataLoader。

• batch_idx – 这一批的索引。

• dataloader_idx – 生成此批次的 dataloader 的索引。（仅在使用多个 dataloader 时）

返回类型

Union[Tensor, Mapping[str, Any], None]

返回

• Tensor - 损失张量

• dict - 一个字典。可以包含任何键，但必须包含键 'loss'。

• None - 跳到下一个批次。

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...

# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

示例:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

如果你传入多个测试数据加载器，test_step() 将会有一个额外的参数。我们建议将默认值设置为0，这样你可以快速在单个和多个数据加载器之间切换。

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

备注

如果你不需要测试，你就不需要实现这个方法。

备注

当调用 test_step() 时，模型已处于评估模式，并且已禁用 PyTorch 梯度。在测试轮次结束时，模型返回到训练模式，并且启用梯度。

to(*args, **kwargs)

参见 torch.nn.Module.to()。

返回类型

Self

to_dtype(dtype)

将模块精度（默认为 float32）转换为另一种精度。

to_empty(*, device, recurse=True)

将参数和缓冲区移动到指定设备，而不复制存储。

参数

• device (torch.device) – 此模块中参数和缓冲区的所需设备。

• recurse (bool) – 子模块的参数和缓冲区是否应递归地移动到指定设备。

返回

自己

返回类型

Module

to_onnx(file_path, input_sample=None, **kwargs)

将模型保存为 ONNX 格式。

参数

• file_path (Union[str, Path]) – 应保存 onnx 模型的文件路径。

• input_sample (Optional[Any, None]) – 用于跟踪的输入。默认值：None（使用 self.example_input_array）

• **kwargs – 将被传递给 torch.onnx.export 函数。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)

model = SimpleModel()
input_sample = torch.randn(1, 64)
model.to_onnx("export.onnx", input_sample, export_params=True)

返回类型

None

to_torchscript(file_path=None, method='script', example_inputs=None, **kwargs)

默认情况下，将整个模型编译为 ScriptModule。如果你想使用追踪，请提供参数 method='trace' 并确保提供了 example_inputs 参数，或者模型设置了 example_input_array。如果你想自定义被脚本的模块，你应该重写这个方法。如果你想要返回多个模块，我们建议使用字典。

参数

• file_path (Union[str, Path, None]) – 保存 torchscript 的路径。默认：None（不保存文件）。

• method (Optional[str, None]) – 是否使用 TorchScript 的 script 或 trace 方法。默认值：’script’

• example_inputs (Optional[Any, None]) – 当方法设置为’trace’时，用于进行追踪的输入。默认值：None（使用 example_input_array）

• **kwargs – 将传递给 torch.jit.script() 或 torch.jit.trace() 函数的额外参数。

备注

• 需要实现 forward() 方法。

• 导出的脚本将被设置为评估模式。

• 建议您安装最新支持版本的 PyTorch 以无限制地使用此功能。另请参阅 torch.jit 文档以了解支持的功能。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

model = SimpleModel()
model.to_torchscript(file_path="model.pt")

torch.jit.save(model.to_torchscript(
    file_path="model_trace.pt", method='trace', example_inputs=torch.randn(1, 64))
)

返回类型

Union[ScriptModule, Dict[str, ScriptModule]]

返回

无论 file_path 是否定义，此 LightningModule 都可以作为 torchscript 使用。

toggle_optimizer(optimizer)

确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。

它与 untoggle_optimizer() 一起工作，以确保 param_requires_grad_state 被正确重置。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要切换的优化器。

返回类型

None

train(mode=True)

设置模块为训练模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

参数

mode (bool) – 是否设置训练模式（True）或评估模式（False）。默认值：True。

返回

自己

返回类型

Module

train_dataloader()

指定训练样本的可迭代对象或可迭代对象的集合。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• fit()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

property trainer: Trainer

返回类型

Trainer

training: bool

training_step(train_batch, batch_idx)

执行训练步骤

返回类型

Tensor

transfer_batch_to_device(batch, device, dataloader_idx)

如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。

以下列出的数据类型（以及它们的任意嵌套）都是开箱即用的：

• torch.Tensor 或任何实现了 .to(…) 的对象

• 列表

• 字典

• 元组

对于其他情况，你需要定义数据如何移动到目标设备（CPU、GPU、TPU等）。

备注

此钩子应仅传输数据，不应修改数据，也不应将数据移动到除作为参数传入的设备之外的任何其他设备（除非你知道自己在做什么）。要检查此钩子的当前执行状态，可以使用 self.trainer.training/testing/validating/predicting，以便可以根据你的需求添加不同的逻辑。

参数

• batch (Any) – 一批需要转移到新设备的数据。

• device (device) – 在 PyTorch 中定义的目标设备。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

对新设备上数据的引用。

示例:

def transfer_batch_to_device(self, batch, device, dataloader_idx):
    if isinstance(batch, CustomBatch):
        # move all tensors in your custom data structure to the device
        batch.samples = batch.samples.to(device)
        batch.targets = batch.targets.to(device)
    elif dataloader_idx == 0:
        # skip device transfer for the first dataloader or anything you wish
        pass
    else:
        batch = super().transfer_batch_to_device(batch, device, dataloader_idx)
    return batch

参见

• move_data_to_device()

• apply_to_collection()

type(dst_type)

参见 torch.nn.Module.type()。

返回类型

Self

unfreeze()

解冻所有参数以进行训练。

model = MyLightningModule(...)
model.unfreeze()

返回类型

None

untoggle_optimizer(optimizer)

重置通过 toggle_optimizer() 切换的所需梯度的状态。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要取消选择的优化器。

返回类型

None

val_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• fit()

• validate()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要验证数据集和 validation_step()，你不需要实现这个方法。

返回类型

Any

validation_step(val_batch, batch_idx)

执行验证步骤

返回类型

Tensor

xpu(device=None)

将所有模型参数和缓冲区移动到XPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在 XPU 上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

zero_grad(set_to_none=True)

重置所有模型参数的梯度。

更多上下文请参见 torch.optim.Optimizer 下的类似函数。

参数

set_to_none (bool) – 而不是设置为零，将梯度设置为 None。详情请参阅 torch.optim.Optimizer.zero_grad()。

返回类型

None

class darts.models.forecasting.pl_forecasting_module.PLFutureCovariatesModule(input_chunk_length, output_chunk_length, output_chunk_shift=0, train_sample_shape=None, loss_fn=MSELoss(), torch_metrics=None, likelihood=None, optimizer_cls=<class 'torch.optim.adam.Adam'>, optimizer_kwargs=None, lr_scheduler_cls=None, lr_scheduler_kwargs=None, use_reversible_instance_norm=False)

基类：PLForecastingModule, ABC

基于 PyTorch Lightning 的预测模块。

这个类旨在通过继承来创建一个新的基于 PyTorch Lightning 的预测模块。当子类化这个类时，请确保添加以下具有给定签名的方法：

• PLForecastingModule.__init__()

• PLForecastingModule.forward()

• PLForecastingModule._produce_train_output()

• PLForecastingModule._get_batch_prediction()

在子类 MyModel 的 __init__() 函数中调用 super(MyModel, self).__init__(**kwargs)，其中 kwargs 是 PLForecastingModule 的参数。

参数

• input_chunk_length (int) – 过去作为模型输入的时间步数（每个块）。适用于目标序列，以及过去和/或未来协变量（如果模型支持）。

• output_chunk_length (int) – 内部模型一次预测的时间步数（每个块）。此外，模型输入中使用的未来协变量的未来值数量（如果模型支持未来协变量）。这与 predict() 中使用的预测范围 n 不同，后者是使用一次性或自回归预测生成的所需预测点数。设置 n <= output_chunk_length 可以防止自回归。当协变量不足以延伸到未来，或禁止模型使用过去和/或未来协变量的未来值进行预测时（取决于模型的协变量支持），这很有用。

• train_sample_shape (Optional[Tuple, None]) – 模型的输入形状，用于在不调用 fit_from_dataset 的情况下实例化模型，并对用于重新训练或预测的新训练/推理数据集进行完整性检查。

• loss_fn (_Loss) – 用于训练的 PyTorch 损失函数。如果指定了 likelihood 参数，则此参数将被概率模型忽略。默认值：torch.nn.MSELoss()。

• torch_metrics (Union[Metric, MetricCollection, None]) – 用于评估的 torch 指标或 MetricCollection。可用指标的完整列表可以在 https://torchmetrics.readthedocs.io/en/latest/ 找到。默认值：None。

• likelihood (Optional[Likelihood, None]) – Darts 的 似然模型 之一，用于概率预测。默认值：None。

• optimizer_cls (Optimizer) – 要使用的 PyTorch 优化器类。默认值：torch.optim.Adam。

• optimizer_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 优化器的键值参数（例如，{'lr': 1e-3} 用于指定学习率）。否则将使用所选 optimizer_cls 的默认值。默认值：None。

• lr_scheduler_cls (Optional[_LRScheduler, None]) – 可选地，指定要使用的 PyTorch 学习率调度器类。指定 None 对应于使用恒定的学习率。默认值：None。

• lr_scheduler_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 学习率调度器的键值参数。默认值：None。

• use_reversible_instance_norm (bool) – 是否使用可逆实例归一化 RINorm 来对抗分布偏移，如 [1] 所示。它仅应用于目标序列的特征，而不应用于协变量。

引用

1

T. Kim et al. “Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift”, https://openreview.net/forum?id=cGDAkQo1C0p

属性

automatic_optimization	如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。
current_epoch	Trainer 中的当前时期，如果未附加则为 0。
device_mesh	像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。
example_input_array	示例输入数组是模块在 forward() 方法中可以消费的规范。
global_rank	当前进程在所有节点和设备中的索引。
global_step	所有epoch中看到的总训练批次。
hparams	使用 save_hyperparameters() 保存的超参数集合。
hparams_initial	使用 save_hyperparameters() 保存的超参数集合。
local_rank	当前进程在单个节点中的索引。
logger	Trainer 中对 logger 对象的引用。
loggers	参考 Trainer 中的日志记录器列表。
on_gpu	如果该模型当前位于GPU上，则返回 True。
output_chunk_length	模型一次预测的时间步数。
strict_loading	确定 Lightning 如何使用 .load_state_dict(..., strict=model.strict_loading) 加载此模型。

设备
dtype
已训练的轮数
fabric
supports_probabilistic_prediction
训练者

方法

add_module(name, module)	将一个子模块添加到当前模块。
all_gather(data[, group, sync_grads])	从多个进程中收集张量或张量集合。
apply(fn)	递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。
backward(loss, *args, **kwargs)	在 training_step() 中返回的损失上调用以执行反向传播。
bfloat16()	将所有浮点参数和缓冲区转换为 bfloat16 数据类型。
buffers([recurse])	返回一个遍历模块缓冲区的迭代器。
children()	返回一个遍历直接子模块的迭代器。
clip_gradients(optimizer[, ...])	内部处理梯度裁剪。
compile(*args, **kwargs)	使用 torch.compile() 编译此模块的 forward。
configure_callbacks()	配置模型特定的回调。
configure_gradient_clipping(optimizer[, ...])	对优化器参数执行梯度裁剪。
configure_model()	在策略和精度感知的环境中创建模块的钩子。
configure_optimizers()	配置优化器和学习率调度器以优化模型。
configure_sharded_model()	已弃用。
configure_torch_metrics(torch_metrics)	处理 torch_metrics 参数。
cpu()	参见 torch.nn.Module.cpu()。
cuda([device])	将所有模型参数和缓冲区移动到GPU。
double()	参见 torch.nn.Module.double()。
eval()	将模块设置为评估模式。
extra_repr()	设置模块的额外表示。
float()	参见 torch.nn.Module.float()。
forward(*args, **kwargs)	与 torch.nn.Module.forward() 相同。
freeze()	冻结所有参数以进行推理。
get_buffer(target)	如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。
get_extra_state()	返回要包含在模块的 state_dict 中的任何额外状态。
get_parameter(target)	如果存在，则返回由 target 给出的参数，否则抛出错误。
get_submodule(target)	如果存在，返回由 target 指定的子模块，否则抛出错误。
half()	参见 torch.nn.Module.half()。
ipu([device])	将所有模型参数和缓冲区移动到 IPU。
load_from_checkpoint(checkpoint_path[, ...])	从检查点加载模型的主要方式。
load_state_dict(state_dict[, strict, assign])	将参数和缓冲区从 state_dict 复制到此模块及其后代中。
log(name, value[, prog_bar, logger, ...])	记录一个键值对。
log_dict(dictionary[, prog_bar, logger, ...])	一次性记录一个值的字典。
lr_scheduler_step(scheduler, metric)	重写此方法以调整 Trainer 调用每个调度器的默认方式。
lr_schedulers()	返回在训练期间使用的学习率调度器。
manual_backward(loss, *args, **kwargs)	在进行手动优化时，直接从 training_step() 调用此方法。
modules()	返回网络中所有模块的迭代器。
named_buffers([prefix, recurse, ...])	返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。
named_children()	返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。
named_modules([memo, prefix, remove_duplicate])	返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。
named_parameters([prefix, recurse, ...])	返回一个模块参数的迭代器，同时生成参数的名称和参数本身。
on_after_backward()	在 loss.backward() 之后和优化器步进之前调用。
on_after_batch_transfer(batch, dataloader_idx)	覆盖此方法以在批次转移到设备后对其进行批量增强或修改。
on_before_backward(loss)	在 loss.backward() 之前调用。
on_before_batch_transfer(batch, dataloader_idx)	在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。
on_before_optimizer_step(optimizer)	在 optimizer.step() 之前调用。
on_before_zero_grad(optimizer)	在 training_step() 之后和 optimizer.zero_grad() 之前调用。
on_fit_end()	在拟合的最后调用。
on_fit_start()	在拟合开始时调用。
on_load_checkpoint(checkpoint)	由 Lightning 调用以恢复您的模型。
on_predict_batch_end(outputs, batch, batch_idx)	在批次之后调用预测循环。
on_predict_batch_start(batch, batch_idx[, ...])	在预测循环中，在该批次发生任何事情之前调用。
on_predict_end()	在预测结束时调用。
on_predict_epoch_end()	在预测结束时调用。
on_predict_epoch_start()	在预测开始时调用。
on_predict_model_eval()	在预测循环开始时调用。
on_predict_start()	在预测开始时调用。
on_save_checkpoint(checkpoint)	当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。
on_test_batch_end(outputs, batch, batch_idx)	在批次之后在测试循环中调用。
on_test_batch_start(batch, batch_idx[, ...])	在测试循环中，在该批次发生任何事情之前调用。
on_test_end()	在测试结束时调用。
on_test_epoch_end()	在每个epoch的测试循环结束时调用。
on_test_epoch_start()	在epoch开始时在测试循环中调用。
on_test_model_eval()	在测试循环开始时调用。
on_test_model_train()	在测试循环结束时调用。
on_test_start()	在测试开始时调用。
on_train_batch_end(outputs, batch, batch_idx)	在训练循环中批次处理后调用。
on_train_batch_start(batch, batch_idx)	在训练循环中，在该批次发生任何事情之前调用。
on_train_end()	在训练结束前调用，在记录器实验关闭之前。
on_train_epoch_end()	在训练循环的每个epoch结束时调用。
on_train_epoch_start()	在训练循环中，在epoch开始时调用。
on_train_start()	在完整性检查后，训练开始时调用。
on_validation_batch_end(outputs, batch, ...)	在批次验证循环后调用。
on_validation_batch_start(batch, batch_idx)	在验证循环中，对该批次进行任何操作之前调用。
on_validation_end()	在验证结束时调用。
on_validation_epoch_end()	在epoch结束时的验证循环中调用。
on_validation_epoch_start()	在epoch开始时验证循环中调用。
on_validation_model_eval()	在验证循环开始时调用。
on_validation_model_train()	当验证循环结束时调用。
on_validation_model_zero_grad()	在进入验证循环之前，由训练循环调用来释放梯度。
on_validation_start()	在验证开始时调用。
optimizer_step(epoch, batch_idx, optimizer)	重写此方法以调整 Trainer 调用优化器的默认方式。
optimizer_zero_grad(epoch, batch_idx, optimizer)	重写此方法以改变 optimizer.zero_grad() 的默认行为。
optimizers([use_pl_optimizer])	返回训练期间使用的优化器。
parameters([recurse])	返回一个遍历模块参数的迭代器。
predict_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。
predict_step(batch, batch_idx[, dataloader_idx])	执行预测步骤
prepare_data()	使用此功能下载并准备数据。
print(*args, **kwargs)	仅从进程 0 打印。
register_backward_hook(hook)	在模块上注册一个反向钩子。
register_buffer(name, tensor[, persistent])	向模块添加一个缓冲区。
register_forward_hook(hook, *[, prepend, ...])	在模块上注册一个前向钩子。
register_forward_pre_hook(hook, *[, ...])	在模块上注册一个前向预钩子。
register_full_backward_hook(hook[, prepend])	在模块上注册一个反向钩子。
register_full_backward_pre_hook(hook[, prepend])	在模块上注册一个反向预钩子。
register_load_state_dict_post_hook(hook)	注册一个在模块的 load_state_dict 调用后运行的后置钩子。
register_module(name, module)	别名 add_module()。
register_parameter(name, param)	向模块添加一个参数。
register_state_dict_pre_hook(hook)	为 state_dict() 方法注册一个预钩子。
requires_grad_([requires_grad])	更改是否应在此模块的参数上记录自动梯度操作。
save_hyperparameters(*args[, ignore, frame, ...])	将参数保存到 hparams 属性中。
set_extra_state(state)	设置包含在加载的 state_dict 中的额外状态。
set_predict_parameters(n, num_samples, ...)	在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。
setup(stage)	在 fit（训练 + 验证）、验证、测试或预测开始时调用。
share_memory()	参见 torch.Tensor.share_memory_()。
state_dict(*args[, destination, prefix, ...])	返回一个包含模块整体状态引用的字典。
teardown(stage)	在 fit（训练 + 验证）、验证、测试或预测结束时调用。
test_dataloader()	一个可迭代对象或多个可迭代对象的集合，指定测试样本。
test_step(*args, **kwargs)	对测试集中的单个批次数据进行操作。
to(*args, **kwargs)	参见 torch.nn.Module.to()。
to_dtype(dtype)	将模块精度（默认为 float32）转换为另一种精度。
to_empty(*, device[, recurse])	将参数和缓冲区移动到指定设备，而不复制存储。
to_onnx(file_path[, input_sample])	将模型保存为 ONNX 格式。
to_torchscript([file_path, method, ...])	默认情况下，将整个模型编译为 ScriptModule。
toggle_optimizer(optimizer)	确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。
train([mode])	设置模块为训练模式。
train_dataloader()	指定训练样本的可迭代对象或可迭代对象的集合。
training_step(train_batch, batch_idx)	执行训练步骤
transfer_batch_to_device(batch, device, ...)	如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。
type(dst_type)	参见 torch.nn.Module.type()。
unfreeze()	解冻所有参数以进行训练。
untoggle_optimizer(optimizer)	重置通过 toggle_optimizer() 切换的所需梯度的状态。
val_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。
validation_step(val_batch, batch_idx)	执行验证步骤
xpu([device])	将所有模型参数和缓冲区移动到XPU。
zero_grad([set_to_none])	重置所有模型参数的梯度。

__call__
set_mc_dropout

CHECKPOINT_HYPER_PARAMS_KEY = 'hyper_parameters'

CHECKPOINT_HYPER_PARAMS_NAME = 'hparams_name'

CHECKPOINT_HYPER_PARAMS_TYPE = 'hparams_type'

T_destination

TypeVar(‘T_destination’, bound=Dict[str, Any]) 的别名

add_module(name, module)

将一个子模块添加到当前模块。

可以通过给定的名称将模块作为属性访问。

参数

• name (str) – 子模块的名称。可以使用给定的名称从此模块访问子模块。

• module (Module) – 要添加到模块中的子模块。

返回类型

None

all_gather(data, group=None, sync_grads=False)

从多个进程中收集张量或张量集合。

此方法需要在所有进程上调用，并且所有进程上的张量需要具有相同的形状，否则程序将永远停滞。

参数

• data (Union[Tensor, Dict, List, Tuple]) – int, float, 形状为 (batch, …) 的张量，或可能是嵌套的集合。

• group (Optional[Any, None]) – 要从中收集结果的进程组。默认为所有进程（world）

• sync_grads (bool) – 允许用户为 all_gather 操作同步梯度的标志

返回类型

Union[Tensor, Dict, List, Tuple]

返回

一个形状为 (world_size, batch, …) 的张量，或者如果输入是一个集合，输出也将是一个具有这种形状张量的集合。对于 world_size 为 1 的特殊情况，张量不会增加额外的维度。

allow_zero_length_dataloader_with_multiple_devices: bool

apply(fn)

递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。

典型用途包括初始化模型的参数（另见 nn-init-doc）。

参数

fn (Module -> None) – 要应用于每个子模块的函数

返回

自己

返回类型

Module

示例:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

property automatic_optimization: bool

如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。

返回类型

bool

backward(loss, *args, **kwargs)

在 training_step() 中返回的损失上调用以执行反向传播。如果需要，请使用您自己的实现覆盖此钩子。

参数

loss (Tensor) – 由 training_step() 返回的损失张量。如果使用梯度累积，这里的损失值是归一化的（按 1 / 累积步数缩放）。

示例:

def backward(self, loss):
    loss.backward()

返回类型

None

bfloat16()

将所有浮点参数和缓冲区转换为 bfloat16 数据类型。

备注

此方法就地修改模块。

返回

自己

返回类型

Module

buffers(recurse=True)

返回一个遍历模块缓冲区的迭代器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。

生成器

torch.Tensor – 模块缓冲区

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Tensor]

call_super_init: bool = False

children()

返回一个遍历直接子模块的迭代器。

生成器

模块 – 子模块

返回类型

Iterator[Module]

clip_gradients(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

内部处理梯度裁剪。

备注

• 不要重写此方法。如果你想自定义梯度裁剪，考虑使用 configure_gradient_clipping() 方法。

• 对于手动优化（self.automatic_optimization = False），如果你想使用梯度裁剪，可以在训练步骤中手动调用 self.clip_gradients(opt, gradient_clip_val=0.5, gradient_clip_algorithm="norm")。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。

• gradient_clip_algorithm (Optional[str, None]) – 使用的梯度裁剪算法。传递 gradient_clip_algorithm="value" 按值裁剪，传递 gradient_clip_algorithm="norm" 按范数裁剪。

返回类型

None

compile(*args, **kwargs)

使用 torch.compile() 编译此模块的 forward。

此模块的 __call__ 方法被编译，所有参数都按原样传递给 torch.compile()。

有关此函数的参数详情，请参阅 torch.compile()。

configure_callbacks()

配置模型特定的回调。当模型被附加时，例如当调用 .fit() 或 .test() 时，这里返回的回调列表或回调将与传递给 Trainer 的 callbacks 参数的回调列表合并。如果这里返回的回调与 Trainer 的回调列表中已有的一个或多个回调类型相同，它将优先并替换它们。此外，Lightning 将确保 ModelCheckpoint 回调最后运行。

返回类型

Union[Sequence[Callback], Callback]

返回

一个回调函数或回调函数列表，这些回调函数将扩展 Trainer 中的回调函数列表。

示例:

def configure_callbacks(self):
    early_stop = EarlyStopping(monitor="val_acc", mode="max")
    checkpoint = ModelCheckpoint(monitor="val_loss")
    return [early_stop, checkpoint]

configure_gradient_clipping(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

对优化器参数执行梯度裁剪。在调用 optimizer_step() 之前执行。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。默认情况下，传递给 Trainer 的值将在此处可用。

• gradient_clip_algorithm (Optional[str, None]) – 要使用的梯度裁剪算法。默认情况下，传递给 Trainer 的值将在此处可用。

示例:

def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm):
    # Implement your own custom logic to clip gradients
    # You can call `self.clip_gradients` with your settings:
    self.clip_gradients(
        optimizer,
        gradient_clip_val=gradient_clip_val,
        gradient_clip_algorithm=gradient_clip_algorithm
    )

返回类型

None

configure_model()

在策略和精度感知的环境中创建模块的钩子。

这在使用分片策略（FSDP 和 DeepSpeed）时特别有用，我们希望立即分片模型以节省内存和初始化时间。对于非分片策略，您可以选择覆盖此钩子或在 init_module() 上下文管理器下初始化您的模型。

此钩子在每个 fit/val/test/predict 阶段都会被调用，因此请确保此钩子的实现是 幂等的，即，在第一次调用钩子后，后续的调用应该是无操作的。

返回类型

None

configure_optimizers()

配置优化器和学习率调度器以优化模型。

configure_sharded_model()

已弃用。

使用 configure_model() 代替。

返回类型

None

static configure_torch_metrics(torch_metrics)

处理 torch_metrics 参数。

返回类型

MetricCollection

cpu()

参见 torch.nn.Module.cpu()。

返回类型

Self

cuda(device=None)

将所有模型参数和缓冲区移动到GPU。这也会使相关的参数和缓冲区成为不同的对象。因此，如果模块在优化期间将驻留在GPU上，则应在构造优化器之前调用此方法。

参数

device (Union[device, int, None]) – 如果指定，所有参数将被复制到该设备。如果为 None，将使用当前的 CUDA 设备索引。

返回

自己

返回类型

Module

property current_epoch: int

Trainer 中的当前时期，如果未附加则为 0。

返回类型

int

property device: device

返回类型

device

property device_mesh: Optional[DeviceMesh]

像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。

返回类型

Optional[ForwardRef, None]

double()

参见 torch.nn.Module.double()。

返回类型

Self

property dtype: Union[str, dtype]

返回类型

Union[str, dtype]

dump_patches: bool = False

property epochs_trained

eval()

将模块设置为评估模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

这等同于 self.train(False)。

请参阅 locally-disable-grad-doc 以比较 .eval() 和几个可能与之混淆的类似机制。

返回

自己

返回类型

Module

property example_input_array: Optional[Union[Tensor, Tuple, Dict]]

示例输入数组是模块在 forward() 方法中可以使用的规范。返回类型解释如下：

• 单个张量：假设模型接受一个参数，即 model.forward(model.example_input_array)

• 元组：输入数组应被解释为位置参数的序列，即 model.forward(*model.example_input_array)

• 字典：输入数组表示命名关键字参数，即 model.forward(**model.example_input_array)

返回类型

Union[Tensor, Tuple, Dict, None]

extra_repr()

设置模块的额外表示。

要打印自定义的额外信息，您应该在您自己的模块中重新实现此方法。单行和多行字符串都是可接受的。

返回类型

str

property fabric: Optional[Fabric]

返回类型

Optional[Fabric, None]

float()

参见 torch.nn.Module.float()。

返回类型

Self

abstract forward(*args, **kwargs)

与 torch.nn.Module.forward() 相同。

参数

• *args – 无论你决定传入 forward 方法的内容。

• **kwargs – 关键字参数也是可能的。

返回类型

Any

返回

你的模型的输出

freeze()

冻结所有参数以进行推理。

示例:

model = MyLightningModule(...)
model.freeze()

返回类型

None

get_buffer(target)

如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的缓冲区的完全限定字符串名称。（有关如何指定完全限定字符串，请参见 get_submodule。）

返回

target 引用的缓冲区

返回类型

torch.Tensor

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是缓冲区的内容

get_extra_state()

返回要包含在模块的 state_dict 中的任何额外状态。

如果需要存储额外状态，请在你的模块中实现此功能以及相应的 set_extra_state()。此函数在构建模块的 state_dict() 时被调用。

请注意，额外的状态应该是可pickle的，以确保state_dict的序列化工作。我们仅对序列化张量提供向后兼容性保证；如果其他对象的序列化pickle形式发生变化，它们可能会破坏向后兼容性。

返回

模块的 state_dict 中存储的任何额外状态

返回类型

object

get_parameter(target)

如果存在，则返回由 target 给出的参数，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的参数的全限定字符串名称。（关于如何指定全限定字符串，请参见 get_submodule。）

返回

target 引用的参数

返回类型

torch.nn.Parameter

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Parameter 的内容

get_submodule(target)

如果存在，返回由 target 指定的子模块，否则抛出错误。

例如，假设你有一个 nn.Module A ，它看起来像这样：

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(该图展示了一个 nn.Module A。A 有一个嵌套的子模块 net_b，它本身有两个子模块 net_c 和 linear。net_c 然后有一个子模块 conv。)

要检查我们是否拥有 linear 子模块，我们会调用 get_submodule("net_b.linear")。要检查我们是否拥有 conv 子模块，我们会调用 get_submodule("net_b.net_c.conv")。

get_submodule 的运行时间受限于 target 中的模块嵌套程度。对 named_modules 的查询可以实现相同的结果，但它是 O(N) 的，其中 N 是传递模块的数量。因此，对于检查某个子模块是否存在的简单检查，应始终使用 get_submodule。

参数

target (str) – 要查找的子模块的完全限定字符串名称。（请参见上面的示例，了解如何指定完全限定的字符串。）

返回

由 target 引用的子模块

返回类型

torch.nn.Module

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Module 的内容

property global_rank: int

当前进程在所有节点和设备中的索引。

返回类型

int

property global_step: int

所有epoch中看到的总训练批次。

如果没有连接训练器，此属性为 0。

返回类型

int

half()

参见 torch.nn.Module.half()。

返回类型

Self

property hparams: Union[AttributeDict, MutableMapping]

通过 save_hyperparameters() 保存的超参数集合。用户可以对其进行修改。对于初始超参数的冻结集，请使用 hparams_initial。

返回类型

Union[AttributeDict, MutableMapping]

返回

可变超参数字典

property hparams_initial: AttributeDict

通过 save_hyperparameters() 保存的超参数集合。这些内容是只读的。可以通过 hparams 手动更新保存的超参数。

返回

不可变的初始超参数

返回类型

AttributeDict

ipu(device=None)

将所有模型参数和缓冲区移动到 IPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在IPU上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs)

从检查点加载模型的主要方式。当Lightning保存检查点时，它会将传递给 __init__ 的参数存储在检查点中的 "hyper_parameters" 下。

通过 **kwargs 指定的任何参数将覆盖存储在 "hyper_parameters" 中的参数。

参数

• checkpoint_path (Union[str, Path, IO]) – 检查点的路径。这也可以是一个URL，或类似文件的对象。

• map_location (Union[device, str, int, Callable[[UntypedStorage, str], Optional[UntypedStorage, None]], Dict[Union[device, str, int], Union[device, str, int]], None]) – 如果你的检查点保存了一个GPU模型，而你现在在CPU上或不同数量的GPU上加载，使用这个来映射到新的设置。其行为与 torch.load() 相同。

• hparams_file (Union[str, Path, None]) –

  可选路径到一个 .yaml 或 .csv 文件，其层次结构如本示例所示:

  drop_prob: 0.2
  dataloader:
      batch_size: 32
  

  您很可能不需要这个，因为 Lightning 总是会将超参数保存到检查点。然而，如果您的检查点权重没有保存超参数，请使用此方法传入一个包含您希望使用的超参数的 .yaml 文件。这些将被转换为 dict 并传递到您的 LightningModule 中使用。

  如果你的模型的 hparams 参数是 Namespace 并且 .yaml 文件具有层次结构，你需要重构你的模型以将 hparams 视为 dict。

• strict (Optional[bool, None]) – 是否严格要求 checkpoint_path 中的键与该模块状态字典返回的键匹配。默认为 True ，除非设置了 LightningModule.strict_loading ，在这种情况下，它默认为 LightningModule.strict_loading 的值。

• **kwargs – 初始化模型所需的任何额外关键字参数。也可以用来覆盖已保存的超参数值。

返回类型

Self

返回

LightningModule 实例，加载了权重和超参数（如果可用）。

备注

load_from_checkpoint 是一个 类 方法。你应该使用你的 LightningModule 类 来调用它，而不是 LightningModule 实例，否则会引发 TypeError。

备注

为了确保所有层都能从检查点加载，如果您的 LightningModule 中重写了此钩子，此函数将在实例化模型后直接调用 configure_model()。但是，请注意 load_from_checkpoint 不支持加载分片检查点，如果模型太大，您可能会耗尽内存。在这种情况下，请考虑通过 .fit(ckpt_path=...) 通过 Trainer 加载。

示例:

# load weights without mapping ...
model = MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')

# or load weights mapping all weights from GPU 1 to GPU 0 ...
map_location = {'cuda:1':'cuda:0'}
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    map_location=map_location
)

# or load weights and hyperparameters from separate files.
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    hparams_file='/path/to/hparams_file.yaml'
)

# override some of the params with new values
model = MyLightningModule.load_from_checkpoint(
    PATH,
    num_layers=128,
    pretrained_ckpt_path=NEW_PATH,
)

# predict
pretrained_model.eval()
pretrained_model.freeze()
y_hat = pretrained_model(x)

load_state_dict(state_dict, strict=True, assign=False)

将参数和缓冲区从 state_dict 复制到此模块及其后代中。

如果 strict 是 True，那么 state_dict 的键必须与该模块的 state_dict() 函数返回的键完全匹配。

警告

如果 assign 是 True，优化器必须在调用 load_state_dict 之后创建，除非 get_swap_module_params_on_conversion() 是 True。

参数

• state_dict (dict) – 包含参数和持久缓冲区的字典。

• strict (bool, optional) – 是否严格要求 state_dict 中的键与该模块的 state_dict() 函数返回的键匹配。默认值：True

• assign (bool, optional) – 当 False 时，当前模块中张量的属性会被保留，而当 True 时，状态字典中张量的属性会被保留。唯一的例外是 Parameter 的 requires_grad 字段，该字段的值会从模块中保留。默认值：False

返回

• missing_keys 是一个包含任何预期键的 str 列表

  此模块中存在但未在提供的 state_dict 中出现的。

• unexpected_keys 是一个包含未预料到的键的 str 列表。

  此模块预期但未在提供的 state_dict 中出现。

返回类型

带有 missing_keys 和 unexpected_keys 字段的 NamedTuple

备注

如果一个参数或缓冲区被注册为 None ，并且其对应的键存在于 state_dict 中，load_state_dict() 将引发一个 RuntimeError 。

property local_rank: int

当前进程在单个节点中的索引。

返回类型

int

log(name, value, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, metric_attribute=None, rank_zero_only=False)

记录一个键值对。

示例:

self.log('train_loss', loss)

每个钩子的默认行为在此处有文档说明：extensions/logging:自动日志记录。

参数

• name (str) – 日志的键。如果使用DDP或其他分布式策略，必须在所有进程中保持一致。

• value (Union[Metric, Tensor, int, float]) – 要记录的值。可以是 float、Tensor 或 Metric。

• prog_bar (bool) – 如果 True 则记录到进度条。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• on_epoch (Optional[bool, None]) – 如果 True ，记录每个epoch累积的指标。默认值由钩子决定。详情请参见 extensions/logging:自动日志记录 。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图表。

• sync_dist (bool) – 如果 True，将在设备间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的DDP组。

• add_dataloader_idx (bool) – 如果 True，将当前数据加载器的索引附加到名称（当使用多个数据加载器时）。如果为 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的 batch_size。这将直接从加载的批次中推断出来，但对于某些数据结构，您可能需要显式提供它。

• metric_attribute (Optional[str, None]) – 要恢复指标状态，Lightning 需要模型中 torchmetrics.Metric 的引用。如果它是模型属性，则会自动找到。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

log_dict(dictionary, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, rank_zero_only=False)

一次性记录一个值的字典。

示例:

values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n}
self.log_dict(values)

参数

• dictionary (Union[Mapping[str, Union[Metric, Tensor, int, float]], MetricCollection]) – 键值对。如果使用DDP或其他分布式策略，键必须在所有进程中保持一致。值可以是 float 、 Tensor 、 Metric 或 MetricCollection 。

• prog_bar (bool) – 如果 True 则记录到进度基础。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。None 自动记录训练步骤，但不记录验证/测试步骤。默认值由钩子决定。详情请参阅 extensions/logging:自动记录。

• on_epoch (Optional[bool, None]) – 如果 True 记录累积的 epoch 指标。None 自动记录 val/test 步骤的日志，但不包括 training_step。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图形

• sync_dist (bool) – 如果 True，将在GPU/TPU之间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的 ddp 组。

• add_dataloader_idx (bool) – 如果 True，则将当前数据加载器的索引附加到名称（在使用多个数据加载器时）。如果 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的批次大小。这将直接从加载的批次中推断出来，但某些数据结构可能需要显式提供它。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

property logger: Optional[Union[Logger, Logger]]

Trainer 中对 logger 对象的引用。

返回类型

Union[Logger, Logger, None]

property loggers: Union[List[Logger], List[Logger]]

参考 Trainer 中的日志记录器列表。

返回类型

Union[List[Logger], List[Logger]]

lr_scheduler_step(scheduler, metric)

重写此方法以调整 Trainer 调用每个调度器的默认方式。默认情况下，Lightning 会根据每个调度器的 interval 调用 step() 方法，如示例所示。

参数

• scheduler (Union[LRScheduler, ReduceLROnPlateau]) – 学习率调度器。

• metric (Optional[Any, None]) – 用于 ReduceLROnPlateau 等调度器的监视器值。

示例:

# DEFAULT
def lr_scheduler_step(self, scheduler, metric):
    if metric is None:
        scheduler.step()
    else:
        scheduler.step(metric)

# Alternative way to update schedulers if it requires an epoch value
def lr_scheduler_step(self, scheduler, metric):
    scheduler.step(epoch=self.current_epoch)

返回类型

None

lr_schedulers()

返回在训练期间使用的学习率调度器。对于手动优化很有用。

返回类型

Union[None, List[Union[LRScheduler, ReduceLROnPlateau]], LRScheduler, ReduceLROnPlateau]

返回

一个调度器，或者在存在多个调度器的情况下是一个调度器列表，如果在 configure_optimizers() 中没有返回调度器，则为 None。

manual_backward(loss, *args, **kwargs)

当你手动进行优化时，请直接从你的 training_step() 调用此方法。通过使用此方法，Lightning 可以确保在使用混合精度时应用所有适当的缩放。

更多示例请参见 手动优化。

示例:

def training_step(...):
    opt = self.optimizers()
    loss = ...
    opt.zero_grad()
    # automatically applies scaling, etc...
    self.manual_backward(loss)
    opt.step()

参数

• loss (Tensor) – 要计算梯度的张量。必须附加一个图。

• *args – 要转发给 backward() 的附加位置参数

• **kwargs – 要转发给 backward() 的额外关键字参数

返回类型

None

modules()

返回网络中所有模块的迭代器。

生成器

模块 – 网络中的一个模块

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

返回类型

Iterator[Module]

named_buffers(prefix='', recurse=True, remove_duplicate=True)

返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。

参数

• prefix (str) – 要添加到所有缓冲区名称的前缀。

• recurse (bool, optional) – 如果为 True，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。默认为 True。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的缓冲区。默认为 True。

生成器

(str, torch.Tensor) – 包含名称和缓冲区的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

返回类型

Iterator[Tuple[str, Tensor]]

named_children()

返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。

生成器

(str, Module) – 包含名称和子模块的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

返回类型

Iterator[Tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)

返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。

参数

• memo (Optional[Set[Module], None]) – 一个备忘录，用于存储已添加到结果中的模块集合

• prefix (str) – 一个将添加到模块名称前的前缀

• remove_duplicate (bool) – 是否在结果中移除重复的模块实例

生成器

(str, Module) – 名称和模块的元组

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)

返回一个模块参数的迭代器，同时生成参数的名称和参数本身。

参数

• prefix (str) – 要添加到所有参数名称前的前缀。

• recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的参数。默认为 True。

生成器

(str, Parameter) – 包含名称和参数的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

返回类型

Iterator[Tuple[str, Parameter]]

on_after_backward()

在 loss.backward() 之后和优化器步进之前调用。

备注

如果使用原生 AMP，梯度在此刻将不会被缩放。如果你需要未缩放的梯度，请使用 on_before_optimizer_step。

返回类型

None

on_after_batch_transfer(batch, dataloader_idx)

覆盖此方法以在批次转移到设备后对其进行批量增强或修改。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_after_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = gpu_transforms(batch['x'])
    return batch

参见

• on_before_batch_transfer()

• transfer_batch_to_device()

on_before_backward(loss)

在 loss.backward() 之前调用。

参数

loss (Tensor) – 损失除以梯度累积的批次数量，如果使用AMP则进行缩放。

返回类型

None

on_before_batch_transfer(batch, dataloader_idx)

在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_before_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = transforms(batch['x'])
    return batch

参见

• on_after_batch_transfer()

• transfer_batch_to_device()

on_before_optimizer_step(optimizer)

在 optimizer.step() 之前调用。

如果使用梯度累积，钩子会在梯度累积完成后被调用。参见：:paramref:`~pytorch_lightning.trainer.trainer.Trainer.accumulate_grad_batches`。

如果使用AMP，损失将在调用此钩子之前未缩放。有关梯度缩放的更多信息，请参阅这些 文档。

如果裁剪梯度，梯度尚未被裁剪。

参数

optimizer (Optimizer) – 当前使用的优化器。

示例:

def on_before_optimizer_step(self, optimizer):
    # example to inspect gradient information in tensorboard
    if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
        for k, v in self.named_parameters():
            self.logger.experiment.add_histogram(
                tag=k, values=v.grad, global_step=self.trainer.global_step
            )

返回类型

None

on_before_zero_grad(optimizer)

在 training_step() 之后和 optimizer.zero_grad() 之前调用。

在训练循环中，调用优化器步骤之后和清零梯度之前。这是一个检查更新后权重信息的好地方。

这是调用它的地方:

for optimizer in optimizers:
    out = training_step(...)

    model.on_before_zero_grad(optimizer) # < ---- called here
    optimizer.zero_grad()

    backward()

参数

optimizer (Optimizer) – 应该被置零的梯度优化器。

返回类型

None

on_fit_end()

在拟合的最后调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

on_fit_start()

在拟合开始时调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

property on_gpu: bool

如果该模型当前位于GPU上，则返回 True。

在LightningModule周围设置标志以区分CPU和GPU行为非常有用。

返回类型

bool

on_load_checkpoint(checkpoint)

由 Lightning 调用以恢复您的模型。如果您使用 on_save_checkpoint() 保存了某些内容，这是您恢复这些内容的机会。

参数

checkpoint (Dict[str, Any]) – 加载的检查点

示例:

def on_load_checkpoint(self, checkpoint):
    # 99% of the time you don't need to implement this method
    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']

备注

Lightning 自动恢复全局步数、轮次和训练状态，包括 amp 缩放。您无需恢复任何与训练相关的内容。

返回类型

None

on_predict_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后调用预测循环。

参数

• outputs (Optional[Any, None]) – predict_step(x) 的输出

• batch (Any) – 由预测 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_batch_start(batch, batch_idx, dataloader_idx=0)

在预测循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_start()

在预测开始时调用。

返回类型

None

on_predict_model_eval()

在预测循环开始时调用。

默认情况下，预测循环在开始之前会在 LightningModule 上调用 .eval()。重写此钩子以更改行为。

返回类型

None

on_predict_start()

在预测开始时调用。

返回类型

None

on_save_checkpoint(checkpoint)

当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。

参数

checkpoint (Dict[str, Any]) – 在转储到文件之前的完整检查点字典。此钩子的实现可以将额外数据插入此字典中。

示例:

def on_save_checkpoint(self, checkpoint):
    # 99% of use cases you don't need to implement this method
    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object

备注

Lightning 保存了训练的所有方面（如 epoch、全局步数等），包括 amp 缩放。您无需存储任何关于训练的信息。

返回类型

None

on_test_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后在测试循环中调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – test_step(x) 的输出

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_batch_start(batch, batch_idx, dataloader_idx=0)

在测试循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_end()

在测试结束时调用。

返回类型

None

on_test_epoch_end()

在每个epoch的测试循环结束时调用。

返回类型

None

on_test_epoch_start()

在epoch开始时在测试循环中调用。

返回类型

None

on_test_model_eval()

在测试循环开始时调用。

默认情况下，测试循环在开始前会调用 .eval() 方法在 LightningModule 上。重写此钩子以改变行为。另请参见 on_test_model_train()。

返回类型

None

on_test_model_train()

在测试循环结束时调用。

测试循环默认会将 LightningModule 的 training 模式恢复到开始测试之前的状态。覆盖此钩子以更改行为。另请参见 on_test_model_eval()。

返回类型

None

on_test_start()

在测试开始时调用。

返回类型

None

on_train_batch_end(outputs, batch, batch_idx)

在训练循环中批次处理后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – 训练步骤(x)的输出

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

备注

这里的 outputs["loss"] 值将是相对于 accumulate_grad_batches 的损失归一化值，该损失从 training_step 返回。

返回类型

None

on_train_batch_start(batch, batch_idx)

在训练循环中，在该批次发生任何事情之前调用。

如果你在这里返回 -1，你将跳过当前轮次的剩余训练。

参数

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

返回类型

Optional[int, None]

on_train_end()

在训练结束前调用，在记录器实验关闭之前。

返回类型

None

on_train_epoch_end()

在训练循环的每个epoch结束时调用。

要在每个epoch结束时访问所有批处理输出，您可以将步骤输出缓存为 LightningModule 的一个属性，并在此钩子中访问它们：

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # free up the memory
        self.training_step_outputs.clear()

on_train_epoch_start()

在训练循环中，在epoch开始时调用。

返回类型

None

on_train_start()

在完整性检查后，训练开始时调用。

返回类型

None

on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次验证循环后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – validation_step(x) 的输出

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_batch_start(batch, batch_idx, dataloader_idx=0)

在验证循环中，对该批次进行任何操作之前调用。

参数

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_end()

在验证结束时调用。

返回类型

None

on_validation_epoch_end()

在epoch结束时的验证循环中调用。

on_validation_epoch_start()

在epoch开始时验证循环中调用。

返回类型

None

on_validation_model_eval()

在验证循环开始时调用。

默认情况下，验证循环在开始前会在 LightningModule 上调用 .eval()。覆盖此钩子以更改行为。另请参见 on_validation_model_train()。

返回类型

None

on_validation_model_train()

当验证循环结束时调用。

默认情况下，验证循环会将 LightningModule 的 training 模式恢复到开始验证之前的状态。重写此钩子以更改行为。另请参见 on_validation_model_eval()。

返回类型

None

on_validation_model_zero_grad()

在进入验证循环之前，由训练循环调用来释放梯度。

返回类型

None

on_validation_start()

在验证开始时调用。

返回类型

None

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)

重写此方法以调整 Trainer 调用优化器的默认方式。

默认情况下，Lightning 会调用 step() 和 zero_grad()，如示例所示。当 Trainer(accumulate_grad_batches != 1) 时，在累积阶段不会调用此方法（和 zero_grad()）。手动优化时，重写此钩子没有好处。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Union[Optimizer, LightningOptimizer]) – 一个 PyTorch 优化器

• optimizer_closure (Optional[Callable[[], Any], None]) – 优化器闭包。此闭包必须执行，因为它包含了 training_step()、optimizer.zero_grad() 和 backward() 的调用。

示例:

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    # Add your custom logic to run directly before `optimizer.step()`

    optimizer.step(closure=optimizer_closure)

    # Add your custom logic to run directly after `optimizer.step()`

返回类型

None

optimizer_zero_grad(epoch, batch_idx, optimizer)

重写此方法以改变 optimizer.zero_grad() 的默认行为。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Optimizer) – 一个 PyTorch 优化器

示例:

# DEFAULT
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad()

# Set gradients to `None` instead of zero to improve performance (not required on `torch>=2.0.0`).
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad(set_to_none=True)

参见 torch.optim.Optimizer.zero_grad() 以了解上述示例的解释。

返回类型

None

optimizers(use_pl_optimizer=True)

返回训练期间使用的优化器。对于手动优化很有用。

参数

use_pl_optimizer (bool) – 如果 True，将会把优化器包装在 LightningOptimizer 中，以自动处理精度、分析和步数调用的计数，以进行适当的日志记录和检查点保存。它特别包装了 step 方法，不支持没有此方法的自定义优化器。

返回类型

Union[Optimizer, LightningOptimizer, _FabricOptimizer, List[Optimizer], List[LightningOptimizer], List[_FabricOptimizer]]

返回

一个单一的优化器，或者在存在多个优化器的情况下，一个优化器列表。

property output_chunk_length: Optional[int]

模型一次预测的时间步数。

返回类型

Optional[int, None]

parameters(recurse=True)

返回一个遍历模块参数的迭代器。

这通常传递给一个优化器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

生成器

参数 – 模块参数

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Parameter]

pred_batch_size: Optional[int]

pred_mc_dropout: Optional[bool]

pred_n: Optional[int]

pred_n_jobs: Optional[int]

pred_num_samples: Optional[int]

pred_roll_size: Optional[int]

predict_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。

有关多个数据加载器的更多信息，请参阅此 部分。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• predict()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

返回

一个 torch.utils.data.DataLoader 或它们的序列，指定预测样本。

predict_likelihood_parameters: Optional[bool]

predict_step(batch, batch_idx, dataloader_idx=None)

执行预测步骤

批处理

Darts 的 InferenceDataset 的输出 - 包含 (过去的目标, 过去的协变量, 历史未来协变量, 未来协变量, 未来过去协变量, 输入时间序列, 预测开始时间步) 的元组

batch_idx

当前批次的批次索引

dataloader_idx

数据加载器索引

返回类型

Sequence[TimeSeries]

prepare_data()

使用此方法下载和准备数据。使用多个进程（分布式设置）下载和保存数据将导致数据损坏。Lightning 确保此方法仅在一个进程内调用，因此您可以安全地在其中添加下载逻辑。

警告

不要将状态设置到模型中（使用 setup 代替），因为这不会在每个设备上被调用

示例:

def prepare_data(self):
    # good
    download_data()
    tokenize()
    etc()

    # bad
    self.split = data_split
    self.some_state = some_other_state()

在分布式环境中，prepare_data 可以通过两种方式调用（使用 prepare_data_per_node）

1. 每个节点一次。这是默认设置，仅在 LOCAL_RANK=0 时调用。

2. 总共一次。仅在 GLOBAL_RANK=0 时调用。

示例:

# DEFAULT
# called once per node on LOCAL_RANK=0 of that node
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = True

# call on GLOBAL_RANK=0 (great for shared file systems)
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = False

在请求数据加载器之前调用此方法：

model.prepare_data()
initialize_distributed()
model.setup(stage)
model.train_dataloader()
model.val_dataloader()
model.test_dataloader()
model.predict_dataloader()

返回类型

None

prepare_data_per_node: bool

print(*args, **kwargs)

仅从进程 0 打印。在任何分布式模式下使用此功能以仅记录一次。

参数

• *args – 要打印的内容。与Python内置的print函数相同。

• **kwargs – 与Python内置的print函数相同。

示例:

def forward(self, x):
    self.print(x, 'in forward')

返回类型

None

register_backward_hook(hook)

在模块上注册一个反向钩子。

此函数已被 register_full_backward_hook() 取代，并且此函数的行为将在未来版本中发生变化。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_buffer(name, tensor, persistent=True)

向模块添加一个缓冲区。

这通常用于注册一个不应被视为模型参数的缓冲区。例如，BatchNorm 的 running_mean 不是参数，而是模块状态的一部分。缓冲区默认是持久的，并将与参数一起保存。可以通过将 persistent 设置为 False 来改变此行为。持久缓冲区和非持久缓冲区之间的唯一区别是后者不会成为此模块的 state_dict 的一部分。

缓冲区可以通过给定的名称作为属性访问。

参数

• name (str) – 缓冲区的名称。可以通过给定的名称从此模块访问缓冲区。

• tensor (Tensor or None) – 要注册的缓冲区。如果为 None ，则运行在缓冲区上的操作，例如 cuda ，将被忽略。如果为 None ，该缓冲区将 不 包含在模块的 state_dict 中。

• persistent (bool) – 缓冲区是否是此模块的 state_dict 的一部分。

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

返回类型

None

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

在模块上注册一个前向钩子。

钩子将在每次 forward() 计算输出后被调用。

如果 with_kwargs 是 False 或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward 。钩子可以修改输出。它可以就地修改输入，但这不会影响前向传播，因为这是在调用 forward() 之后调用的。钩子应具有以下签名:

hook(module, args, output) -> None or modified output

如果 with_kwargs 是 True，前向钩子将被传递给前向函数的 kwargs，并期望返回可能被修改的输出。钩子应具有以下签名:

hook(module, args, kwargs, output) -> None or modified output

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果 True，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之后触发。请注意，使用 register_module_forward_hook() 注册的全局 forward 钩子将在通过此方法注册的所有钩子之前触发。默认值：False

• with_kwargs (bool) – 如果 True，hook 将会传递给 forward 函数的 kwargs。默认值：False

• always_call (bool) – 如果 True ，无论在调用模块时是否引发异常，都将运行 hook 。默认值： False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

在模块上注册一个前向预钩子。

钩子将在每次调用 forward() 之前被调用。

如果 with_kwargs 为假或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward。钩子可以修改输入。用户可以在钩子中返回一个元组或单个修改后的值。如果返回单个值（除非该值已经是元组），我们将把该值包装成一个元组。钩子应具有以下签名:

hook(module, args) -> None or modified input

如果 with_kwargs 为真，前向预钩子将传递给前向函数的 kwargs。如果钩子修改了输入，则应返回 args 和 kwargs。钩子应具有以下签名:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之后触发。请注意，通过 register_module_forward_pre_hook() 注册的全局 forward_pre 钩子将在所有通过此方法注册的钩子之前触发。默认值：False

• with_kwargs (bool) – 如果为真，hook 将被传递给 forward 函数的 kwargs。默认值：False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)

在模块上注册一个反向钩子。

每次计算与模块相关的梯度时，钩子都会被调用，即当且仅当计算与模块输出相关的梯度时，钩子才会执行。钩子应具有以下签名:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input 和 grad_output 是包含相对于输入和输出的梯度的元组。钩子不应修改其参数，但可以选择返回一个新的相对于输入的梯度，该梯度将用于替代 grad_input 在后续计算中。grad_input 将仅对应于作为位置参数给出的输入，所有 kwarg 参数都将被忽略。grad_input 和 grad_output 中的条目将为 None 用于所有非 Tensor 参数。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入或输出，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之后触发。请注意，通过 register_module_full_backward_hook() 注册的全局 backward 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

在模块上注册一个反向预钩子。

每次计算模块的梯度时，都会调用该钩子。钩子应具有以下签名:

hook(module, grad_output) -> tuple[Tensor] or None

The grad_output 是一个元组。钩子不应修改其参数，但它可以选择返回一个新的相对于输出的梯度，该梯度将用于替代 grad_output 在后续计算中。grad_output 中的条目对于所有非 Tensor 参数将为 None。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之后触发。请注意，通过 register_module_full_backward_pre_hook() 注册的全局 backward_pre 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)

注册一个在模块的 load_state_dict 调用后运行的后置钩子。

它应该有以下签名:

hook(module, incompatible_keys) -> None

module 参数是当前注册此钩子的模块，而 incompatible_keys 参数是一个由属性 missing_keys 和 unexpected_keys 组成的 NamedTuple。missing_keys 是一个包含缺失键的 str 的 list，而 unexpected_keys 是一个包含意外键的 str 的 list。

如果需要，可以就地修改给定的 incompatible_keys。

请注意，当使用 strict=True 调用 load_state_dict() 时执行的检查会受到钩子对 missing_keys 或 unexpected_keys 所做修改的影响，这是预期的。对任一组键的添加将导致在 strict=True 时抛出错误，而清除掉缺失和意外的键将避免错误。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_module(name, module)

别名 add_module()。

返回类型

None

register_parameter(name, param)

向模块添加一个参数。

该参数可以通过给定的名称作为属性访问。

参数

• name (str) – 参数的名称。可以使用给定的名称从此模块访问该参数。

• param (Parameter or None) – 要添加到模块的参数。如果为 None ，则运行在参数上的操作，例如 cuda ，将被忽略。如果为 None ，该参数将 不 包含在模块的 state_dict 中。

返回类型

None

register_state_dict_pre_hook(hook)

为 state_dict() 方法注册一个预钩子。

这些钩子将在调用 self 上的 state_dict 之前，使用参数 self、prefix 和 keep_vars 进行调用。注册的钩子可以用于在调用 state_dict 之前执行预处理。

requires_grad_(requires_grad=True)

更改是否应在此模块的参数上记录自动梯度操作。

此方法就地设置参数的 requires_grad 属性。

这种方法有助于冻结模块的一部分以进行微调或单独训练模型的一部分（例如，GAN训练）。

请参阅 locally-disable-grad-doc 以比较 .requires_grad_() 和几个可能与之混淆的类似机制。

参数

requires_grad (bool) – 是否应在此模块的参数上记录自动求导操作。默认值：True。

返回

自己

返回类型

Module

save_hyperparameters(*args, ignore=None, frame=None, logger=True)

将参数保存到 hparams 属性中。

参数

• args (Any) – dict、NameSpace 或 OmegaConf 的单个对象，或来自类 __init__ 的字符串名称或参数

• ignore (Union[Sequence[str], str, None]) – 一个参数名称或从类 __init__ 中忽略的参数名称列表

• frame (Optional[frame, None]) – 一个框架对象。默认是 None

• logger (bool) – 是否将超参数发送到记录器。默认值：True

示例:

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # manually assign arguments
...         self.save_hyperparameters('arg1', 'arg3')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class AutomaticArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # equivalent automatic
...         self.save_hyperparameters()
...     def forward(self, *args, **kwargs):
...         ...
>>> model = AutomaticArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg2": abc
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class SingleArgModel(HyperparametersMixin):
...     def __init__(self, params):
...         super().__init__()
...         # manually assign single argument
...         self.save_hyperparameters(params)
...     def forward(self, *args, **kwargs):
...         ...
>>> model = SingleArgModel(Namespace(p1=1, p2='abc', p3=3.14))
>>> model.hparams
"p1": 1
"p2": abc
"p3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # pass argument(s) to ignore as a string or in a list
...         self.save_hyperparameters(ignore='arg2')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

返回类型

None

set_extra_state(state)

设置包含在加载的 state_dict 中的额外状态。

此函数从 load_state_dict() 调用，以处理在 state_dict 中找到的任何额外状态。如果你的模块需要在 state_dict 中存储额外状态，请实现此函数以及相应的 get_extra_state()。

参数

state (dict) – 来自 state_dict 的额外状态

返回类型

None

set_mc_dropout(active)

set_predict_parameters(n, num_samples, roll_size, batch_size, n_jobs, predict_likelihood_parameters, mc_dropout)

在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。

返回类型

None

setup(stage)

在 fit（训练 + 验证）、验证、测试或预测开始时调用。当你需要动态构建模型或调整模型时，这是一个很好的钩子。在使用 DDP 时，此钩子会在每个进程上调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

示例:

class LitModel(...):
    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

        # don't do this
        self.something = else

    def setup(self, stage):
        data = load_data(...)
        self.l1 = nn.Linear(28, data.num_classes)

返回类型

None

share_memory()

参见 torch.Tensor.share_memory_()。

返回类型

~T

state_dict(*args, destination=None, prefix='', keep_vars=False)

返回一个包含模块整体状态引用的字典。

参数和持久缓冲区（例如运行平均值）都包含在内。键是对应的参数和缓冲区名称。设置为 None 的参数和缓冲区不包含在内。

备注

返回的对象是一个浅拷贝。它包含对模块参数和缓冲区的引用。

警告

目前，state_dict() 也按顺序接受 destination、prefix 和 keep_vars 的位置参数。然而，这正在被弃用，未来版本将强制使用关键字参数。

警告

请避免使用参数 destination，因为它不是为最终用户设计的。

参数

• destination (dict, optional) – 如果提供了，模块的状态将被更新到字典中，并返回相同的对象。否则，将创建并返回一个 OrderedDict。默认值：None。

• prefix (str, optional) – 在 state_dict 中组合键时添加到参数和缓冲区名称的前缀。默认值：''。

• keep_vars (bool, optional) – 默认情况下，状态字典中返回的 Tensor 是与自动求导分离的。如果设置为 True，则不会执行分离操作。默认值：False。

返回

包含模块整个状态的字典

返回类型

dict

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

property strict_loading: bool

确定 Lightning 如何使用 .load_state_dict(…, strict=model.strict_loading) 加载此模型。

返回类型

bool

property supports_probabilistic_prediction: bool

返回类型

bool

teardown(stage)

在 fit（训练 + 验证）、验证、测试或预测结束时调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

返回类型

None

test_dataloader()

一个可迭代对象或多个可迭代对象的集合，指定测试样本。

有关多个数据加载器的更多信息，请参阅此 部分。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• test()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要测试数据集和 test_step() ，你不需要实现这个方法。

返回类型

Any

test_step(*args, **kwargs)

对测试集中的单个批次数据进行操作。在这一步中，您通常会生成示例或计算任何感兴趣的内容，例如准确性。

参数

• batch – 你的数据可迭代对象的输出，通常是一个 DataLoader。

• batch_idx – 这一批的索引。

• dataloader_idx – 生成此批次的 dataloader 的索引。（仅在使用多个 dataloader 时）

返回类型

Union[Tensor, Mapping[str, Any], None]

返回

• Tensor - 损失张量

• dict - 一个字典。可以包含任何键，但必须包含键 'loss'。

• None - 跳到下一个批次。

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...

# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

示例:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

如果你传入多个测试数据加载器，test_step() 将会有一个额外的参数。我们建议将默认值设置为0，这样你可以快速在单个和多个数据加载器之间切换。

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

备注

如果你不需要测试，你就不需要实现这个方法。

备注

当调用 test_step() 时，模型已处于评估模式，并且已禁用 PyTorch 梯度。在测试轮次结束时，模型返回到训练模式，并且启用梯度。

to(*args, **kwargs)

参见 torch.nn.Module.to()。

返回类型

Self

to_dtype(dtype)

将模块精度（默认为 float32）转换为另一种精度。

to_empty(*, device, recurse=True)

将参数和缓冲区移动到指定设备，而不复制存储。

参数

• device (torch.device) – 此模块中参数和缓冲区的所需设备。

• recurse (bool) – 子模块的参数和缓冲区是否应递归地移动到指定设备。

返回

自己

返回类型

Module

to_onnx(file_path, input_sample=None, **kwargs)

将模型保存为 ONNX 格式。

参数

• file_path (Union[str, Path]) – 应保存 onnx 模型的文件路径。

• input_sample (Optional[Any, None]) – 用于跟踪的输入。默认值：None（使用 self.example_input_array）

• **kwargs – 将被传递给 torch.onnx.export 函数。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)

model = SimpleModel()
input_sample = torch.randn(1, 64)
model.to_onnx("export.onnx", input_sample, export_params=True)

返回类型

None

to_torchscript(file_path=None, method='script', example_inputs=None, **kwargs)

默认情况下，将整个模型编译为 ScriptModule。如果你想使用追踪，请提供参数 method='trace' 并确保提供了 example_inputs 参数，或者模型设置了 example_input_array。如果你想自定义被脚本的模块，你应该重写这个方法。如果你想要返回多个模块，我们建议使用字典。

参数

• file_path (Union[str, Path, None]) – 保存 torchscript 的路径。默认：None（不保存文件）。

• method (Optional[str, None]) – 是否使用 TorchScript 的 script 或 trace 方法。默认值：’script’

• example_inputs (Optional[Any, None]) – 当方法设置为’trace’时，用于进行追踪的输入。默认值：None（使用 example_input_array）

• **kwargs – 将传递给 torch.jit.script() 或 torch.jit.trace() 函数的额外参数。

备注

• 需要实现 forward() 方法。

• 导出的脚本将被设置为评估模式。

• 建议您安装最新支持版本的 PyTorch 以无限制地使用此功能。另请参阅 torch.jit 文档以了解支持的功能。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

model = SimpleModel()
model.to_torchscript(file_path="model.pt")

torch.jit.save(model.to_torchscript(
    file_path="model_trace.pt", method='trace', example_inputs=torch.randn(1, 64))
)

返回类型

Union[ScriptModule, Dict[str, ScriptModule]]

返回

无论 file_path 是否定义，此 LightningModule 都可以作为 torchscript 使用。

toggle_optimizer(optimizer)

确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。

它与 untoggle_optimizer() 一起工作，以确保 param_requires_grad_state 被正确重置。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要切换的优化器。

返回类型

None

train(mode=True)

设置模块为训练模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

参数

mode (bool) – 是否设置训练模式（True）或评估模式（False）。默认值：True。

返回

自己

返回类型

Module

train_criterion_reduction: Optional[str]

train_dataloader()

指定训练样本的可迭代对象或可迭代对象的集合。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• fit()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

property trainer: Trainer

返回类型

Trainer

training: bool

training_step(train_batch, batch_idx)

执行训练步骤

返回类型

Tensor

transfer_batch_to_device(batch, device, dataloader_idx)

如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。

以下列出的数据类型（以及它们的任意嵌套）都是开箱即用的：

• torch.Tensor 或任何实现了 .to(…) 的对象

• 列表

• 字典

• 元组

对于其他情况，你需要定义数据如何移动到目标设备（CPU、GPU、TPU等）。

备注

此钩子应仅传输数据，不应修改数据，也不应将数据移动到除作为参数传入的设备之外的任何其他设备（除非你知道自己在做什么）。要检查此钩子的当前执行状态，可以使用 self.trainer.training/testing/validating/predicting，以便可以根据你的需求添加不同的逻辑。

参数

• batch (Any) – 一批需要转移到新设备的数据。

• device (device) – 在 PyTorch 中定义的目标设备。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

对新设备上数据的引用。

示例:

def transfer_batch_to_device(self, batch, device, dataloader_idx):
    if isinstance(batch, CustomBatch):
        # move all tensors in your custom data structure to the device
        batch.samples = batch.samples.to(device)
        batch.targets = batch.targets.to(device)
    elif dataloader_idx == 0:
        # skip device transfer for the first dataloader or anything you wish
        pass
    else:
        batch = super().transfer_batch_to_device(batch, device, dataloader_idx)
    return batch

参见

• move_data_to_device()

• apply_to_collection()

type(dst_type)

参见 torch.nn.Module.type()。

返回类型

Self

unfreeze()

解冻所有参数以进行训练。

model = MyLightningModule(...)
model.unfreeze()

返回类型

None

untoggle_optimizer(optimizer)

重置通过 toggle_optimizer() 切换的所需梯度的状态。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要取消选择的优化器。

返回类型

None

val_criterion_reduction: Optional[str]

val_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• fit()

• validate()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要验证数据集和 validation_step()，你不需要实现这个方法。

返回类型

Any

validation_step(val_batch, batch_idx)

执行验证步骤

返回类型

Tensor

xpu(device=None)

将所有模型参数和缓冲区移动到XPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在 XPU 上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

zero_grad(set_to_none=True)

重置所有模型参数的梯度。

更多上下文请参见 torch.optim.Optimizer 下的类似函数。

参数

set_to_none (bool) – 而不是设置为零，将梯度设置为 None。详情请参阅 torch.optim.Optimizer.zero_grad()。

返回类型

None

class darts.models.forecasting.pl_forecasting_module.PLMixedCovariatesModule(input_chunk_length, output_chunk_length, output_chunk_shift=0, train_sample_shape=None, loss_fn=MSELoss(), torch_metrics=None, likelihood=None, optimizer_cls=<class 'torch.optim.adam.Adam'>, optimizer_kwargs=None, lr_scheduler_cls=None, lr_scheduler_kwargs=None, use_reversible_instance_norm=False)

基类：PLForecastingModule, ABC

基于 PyTorch Lightning 的预测模块。

这个类旨在通过继承来创建一个新的基于 PyTorch Lightning 的预测模块。当子类化这个类时，请确保添加以下具有给定签名的方法：

• PLForecastingModule.__init__()

• PLForecastingModule.forward()

• PLForecastingModule._produce_train_output()

• PLForecastingModule._get_batch_prediction()

在子类 MyModel 的 __init__() 函数中调用 super(MyModel, self).__init__(**kwargs)，其中 kwargs 是 PLForecastingModule 的参数。

参数

• input_chunk_length (int) – 过去作为模型输入的时间步数（每个块）。适用于目标序列，以及过去和/或未来协变量（如果模型支持）。

• output_chunk_length (int) – 内部模型一次预测的时间步数（每个块）。此外，模型输入中使用的未来协变量的未来值数量（如果模型支持未来协变量）。这与 predict() 中使用的预测范围 n 不同，后者是使用一次性或自回归预测生成的所需预测点数。设置 n <= output_chunk_length 可以防止自回归。当协变量不足以延伸到未来，或禁止模型使用过去和/或未来协变量的未来值进行预测时（取决于模型的协变量支持），这很有用。

• train_sample_shape (Optional[Tuple, None]) – 模型的输入形状，用于在不调用 fit_from_dataset 的情况下实例化模型，并对用于重新训练或预测的新训练/推理数据集进行完整性检查。

• loss_fn (_Loss) – 用于训练的 PyTorch 损失函数。如果指定了 likelihood 参数，则此参数将被概率模型忽略。默认值：torch.nn.MSELoss()。

• torch_metrics (Union[Metric, MetricCollection, None]) – 用于评估的 torch 指标或 MetricCollection。可用指标的完整列表可以在 https://torchmetrics.readthedocs.io/en/latest/ 找到。默认值：None。

• likelihood (Optional[Likelihood, None]) – Darts 的 似然模型 之一，用于概率预测。默认值：None。

• optimizer_cls (Optimizer) – 要使用的 PyTorch 优化器类。默认值：torch.optim.Adam。

• optimizer_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 优化器的键值参数（例如，{'lr': 1e-3} 用于指定学习率）。否则将使用所选 optimizer_cls 的默认值。默认值：None。

• lr_scheduler_cls (Optional[_LRScheduler, None]) – 可选地，指定要使用的 PyTorch 学习率调度器类。指定 None 对应于使用恒定的学习率。默认值：None。

• lr_scheduler_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 学习率调度器的键值参数。默认值：None。

• use_reversible_instance_norm (bool) – 是否使用可逆实例归一化 RINorm 来对抗分布偏移，如 [1] 所示。它仅应用于目标序列的特征，而不应用于协变量。

引用

1

T. Kim et al. “Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift”, https://openreview.net/forum?id=cGDAkQo1C0p

属性

automatic_optimization	如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。
current_epoch	Trainer 中的当前时期，如果未附加则为 0。
device_mesh	像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。
example_input_array	示例输入数组是模块在 forward() 方法中可以消费的规范。
global_rank	当前进程在所有节点和设备中的索引。
global_step	所有epoch中看到的总训练批次。
hparams	使用 save_hyperparameters() 保存的超参数集合。
hparams_initial	使用 save_hyperparameters() 保存的超参数集合。
local_rank	当前进程在单个节点中的索引。
logger	Trainer 中对 logger 对象的引用。
loggers	参考 Trainer 中的日志记录器列表。
on_gpu	如果该模型当前位于GPU上，则返回 True。
output_chunk_length	模型一次预测的时间步数。
strict_loading	确定 Lightning 如何使用 .load_state_dict(..., strict=model.strict_loading) 加载此模型。

设备
dtype
已训练的轮数
fabric
supports_probabilistic_prediction
训练者

方法

add_module(name, module)	将一个子模块添加到当前模块。
all_gather(data[, group, sync_grads])	从多个进程中收集张量或张量集合。
apply(fn)	递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。
backward(loss, *args, **kwargs)	在 training_step() 中返回的损失上调用以执行反向传播。
bfloat16()	将所有浮点参数和缓冲区转换为 bfloat16 数据类型。
buffers([recurse])	返回一个遍历模块缓冲区的迭代器。
children()	返回一个遍历直接子模块的迭代器。
clip_gradients(optimizer[, ...])	内部处理梯度裁剪。
compile(*args, **kwargs)	使用 torch.compile() 编译此模块的 forward。
configure_callbacks()	配置模型特定的回调。
configure_gradient_clipping(optimizer[, ...])	对优化器参数执行梯度裁剪。
configure_model()	在策略和精度感知的环境中创建模块的钩子。
configure_optimizers()	配置优化器和学习率调度器以优化模型。
configure_sharded_model()	已弃用。
configure_torch_metrics(torch_metrics)	处理 torch_metrics 参数。
cpu()	参见 torch.nn.Module.cpu()。
cuda([device])	将所有模型参数和缓冲区移动到GPU。
double()	参见 torch.nn.Module.double()。
eval()	将模块设置为评估模式。
extra_repr()	设置模块的额外表示。
float()	参见 torch.nn.Module.float()。
forward(*args, **kwargs)	与 torch.nn.Module.forward() 相同。
freeze()	冻结所有参数以进行推理。
get_buffer(target)	如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。
get_extra_state()	返回要包含在模块的 state_dict 中的任何额外状态。
get_parameter(target)	如果存在，则返回由 target 给出的参数，否则抛出错误。
get_submodule(target)	如果存在，返回由 target 指定的子模块，否则抛出错误。
half()	参见 torch.nn.Module.half()。
ipu([device])	将所有模型参数和缓冲区移动到 IPU。
load_from_checkpoint(checkpoint_path[, ...])	从检查点加载模型的主要方式。
load_state_dict(state_dict[, strict, assign])	将参数和缓冲区从 state_dict 复制到此模块及其后代中。
log(name, value[, prog_bar, logger, ...])	记录一个键值对。
log_dict(dictionary[, prog_bar, logger, ...])	一次性记录一个值的字典。
lr_scheduler_step(scheduler, metric)	重写此方法以调整 Trainer 调用每个调度器的默认方式。
lr_schedulers()	返回在训练期间使用的学习率调度器。
manual_backward(loss, *args, **kwargs)	在进行手动优化时，直接从 training_step() 调用此方法。
modules()	返回网络中所有模块的迭代器。
named_buffers([prefix, recurse, ...])	返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。
named_children()	返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。
named_modules([memo, prefix, remove_duplicate])	返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。
named_parameters([prefix, recurse, ...])	返回一个模块参数的迭代器，同时生成参数的名称和参数本身。
on_after_backward()	在 loss.backward() 之后和优化器步进之前调用。
on_after_batch_transfer(batch, dataloader_idx)	覆盖此方法以在批次转移到设备后对其进行批量增强或修改。
on_before_backward(loss)	在 loss.backward() 之前调用。
on_before_batch_transfer(batch, dataloader_idx)	在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。
on_before_optimizer_step(optimizer)	在 optimizer.step() 之前调用。
on_before_zero_grad(optimizer)	在 training_step() 之后和 optimizer.zero_grad() 之前调用。
on_fit_end()	在拟合的最后调用。
on_fit_start()	在拟合开始时调用。
on_load_checkpoint(checkpoint)	由 Lightning 调用以恢复您的模型。
on_predict_batch_end(outputs, batch, batch_idx)	在批次之后调用预测循环。
on_predict_batch_start(batch, batch_idx[, ...])	在预测循环中，在该批次发生任何事情之前调用。
on_predict_end()	在预测结束时调用。
on_predict_epoch_end()	在预测结束时调用。
on_predict_epoch_start()	在预测开始时调用。
on_predict_model_eval()	在预测循环开始时调用。
on_predict_start()	在预测开始时调用。
on_save_checkpoint(checkpoint)	当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。
on_test_batch_end(outputs, batch, batch_idx)	在批次之后在测试循环中调用。
on_test_batch_start(batch, batch_idx[, ...])	在测试循环中，在该批次发生任何事情之前调用。
on_test_end()	在测试结束时调用。
on_test_epoch_end()	在每个epoch的测试循环结束时调用。
on_test_epoch_start()	在epoch开始时在测试循环中调用。
on_test_model_eval()	在测试循环开始时调用。
on_test_model_train()	在测试循环结束时调用。
on_test_start()	在测试开始时调用。
on_train_batch_end(outputs, batch, batch_idx)	在训练循环中批次处理后调用。
on_train_batch_start(batch, batch_idx)	在训练循环中，在该批次发生任何事情之前调用。
on_train_end()	在训练结束前调用，在记录器实验关闭之前。
on_train_epoch_end()	在训练循环的每个epoch结束时调用。
on_train_epoch_start()	在训练循环中，在epoch开始时调用。
on_train_start()	在完整性检查后，训练开始时调用。
on_validation_batch_end(outputs, batch, ...)	在批次验证循环后调用。
on_validation_batch_start(batch, batch_idx)	在验证循环中，对该批次进行任何操作之前调用。
on_validation_end()	在验证结束时调用。
on_validation_epoch_end()	在epoch结束时的验证循环中调用。
on_validation_epoch_start()	在epoch开始时验证循环中调用。
on_validation_model_eval()	在验证循环开始时调用。
on_validation_model_train()	当验证循环结束时调用。
on_validation_model_zero_grad()	在进入验证循环之前，由训练循环调用来释放梯度。
on_validation_start()	在验证开始时调用。
optimizer_step(epoch, batch_idx, optimizer)	重写此方法以调整 Trainer 调用优化器的默认方式。
optimizer_zero_grad(epoch, batch_idx, optimizer)	重写此方法以改变 optimizer.zero_grad() 的默认行为。
optimizers([use_pl_optimizer])	返回训练期间使用的优化器。
parameters([recurse])	返回一个遍历模块参数的迭代器。
predict_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。
predict_step(batch, batch_idx[, dataloader_idx])	执行预测步骤
prepare_data()	使用此功能下载并准备数据。
print(*args, **kwargs)	仅从进程 0 打印。
register_backward_hook(hook)	在模块上注册一个反向钩子。
register_buffer(name, tensor[, persistent])	向模块添加一个缓冲区。
register_forward_hook(hook, *[, prepend, ...])	在模块上注册一个前向钩子。
register_forward_pre_hook(hook, *[, ...])	在模块上注册一个前向预钩子。
register_full_backward_hook(hook[, prepend])	在模块上注册一个反向钩子。
register_full_backward_pre_hook(hook[, prepend])	在模块上注册一个反向预钩子。
register_load_state_dict_post_hook(hook)	注册一个在模块的 load_state_dict 调用后运行的后置钩子。
register_module(name, module)	别名 add_module()。
register_parameter(name, param)	向模块添加一个参数。
register_state_dict_pre_hook(hook)	为 state_dict() 方法注册一个预钩子。
requires_grad_([requires_grad])	更改是否应在此模块的参数上记录自动梯度操作。
save_hyperparameters(*args[, ignore, frame, ...])	将参数保存到 hparams 属性中。
set_extra_state(state)	设置包含在加载的 state_dict 中的额外状态。
set_predict_parameters(n, num_samples, ...)	在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。
setup(stage)	在 fit（训练 + 验证）、验证、测试或预测开始时调用。
share_memory()	参见 torch.Tensor.share_memory_()。
state_dict(*args[, destination, prefix, ...])	返回一个包含模块整体状态引用的字典。
teardown(stage)	在 fit（训练 + 验证）、验证、测试或预测结束时调用。
test_dataloader()	一个可迭代对象或多个可迭代对象的集合，指定测试样本。
test_step(*args, **kwargs)	对测试集中的单个批次数据进行操作。
to(*args, **kwargs)	参见 torch.nn.Module.to()。
to_dtype(dtype)	将模块精度（默认为 float32）转换为另一种精度。
to_empty(*, device[, recurse])	将参数和缓冲区移动到指定设备，而不复制存储。
to_onnx(file_path[, input_sample])	将模型保存为 ONNX 格式。
to_torchscript([file_path, method, ...])	默认情况下，将整个模型编译为 ScriptModule。
toggle_optimizer(optimizer)	确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。
train([mode])	设置模块为训练模式。
train_dataloader()	指定训练样本的可迭代对象或可迭代对象的集合。
training_step(train_batch, batch_idx)	执行训练步骤
transfer_batch_to_device(batch, device, ...)	如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。
type(dst_type)	参见 torch.nn.Module.type()。
unfreeze()	解冻所有参数以进行训练。
untoggle_optimizer(optimizer)	重置通过 toggle_optimizer() 切换的所需梯度的状态。
val_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。
validation_step(val_batch, batch_idx)	执行验证步骤
xpu([device])	将所有模型参数和缓冲区移动到XPU。
zero_grad([set_to_none])	重置所有模型参数的梯度。

__call__
set_mc_dropout

CHECKPOINT_HYPER_PARAMS_KEY = 'hyper_parameters'

CHECKPOINT_HYPER_PARAMS_NAME = 'hparams_name'

CHECKPOINT_HYPER_PARAMS_TYPE = 'hparams_type'

T_destination

TypeVar(‘T_destination’, bound=Dict[str, Any]) 的别名

add_module(name, module)

将一个子模块添加到当前模块。

可以通过给定的名称将模块作为属性访问。

参数

• name (str) – 子模块的名称。可以使用给定的名称从此模块访问子模块。

• module (Module) – 要添加到模块中的子模块。

返回类型

None

all_gather(data, group=None, sync_grads=False)

从多个进程中收集张量或张量集合。

此方法需要在所有进程上调用，并且所有进程上的张量需要具有相同的形状，否则程序将永远停滞。

参数

• data (Union[Tensor, Dict, List, Tuple]) – int, float, 形状为 (batch, …) 的张量，或可能是嵌套的集合。

• group (Optional[Any, None]) – 要从中收集结果的进程组。默认为所有进程（world）

• sync_grads (bool) – 允许用户为 all_gather 操作同步梯度的标志

返回类型

Union[Tensor, Dict, List, Tuple]

返回

一个形状为 (world_size, batch, …) 的张量，或者如果输入是一个集合，输出也将是一个具有这种形状张量的集合。对于 world_size 为 1 的特殊情况，张量不会增加额外的维度。

allow_zero_length_dataloader_with_multiple_devices: bool

apply(fn)

递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。

典型用途包括初始化模型的参数（另见 nn-init-doc）。

参数

fn (Module -> None) – 要应用于每个子模块的函数

返回

自己

返回类型

Module

示例:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

property automatic_optimization: bool

如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。

返回类型

bool

backward(loss, *args, **kwargs)

在 training_step() 中返回的损失上调用以执行反向传播。如果需要，请使用您自己的实现覆盖此钩子。

参数

loss (Tensor) – 由 training_step() 返回的损失张量。如果使用梯度累积，这里的损失值是归一化的（按 1 / 累积步数缩放）。

示例:

def backward(self, loss):
    loss.backward()

返回类型

None

bfloat16()

将所有浮点参数和缓冲区转换为 bfloat16 数据类型。

备注

此方法就地修改模块。

返回

自己

返回类型

Module

buffers(recurse=True)

返回一个遍历模块缓冲区的迭代器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。

生成器

torch.Tensor – 模块缓冲区

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Tensor]

call_super_init: bool = False

children()

返回一个遍历直接子模块的迭代器。

生成器

模块 – 子模块

返回类型

Iterator[Module]

clip_gradients(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

内部处理梯度裁剪。

备注

• 不要重写此方法。如果你想自定义梯度裁剪，考虑使用 configure_gradient_clipping() 方法。

• 对于手动优化（self.automatic_optimization = False），如果你想使用梯度裁剪，可以在训练步骤中手动调用 self.clip_gradients(opt, gradient_clip_val=0.5, gradient_clip_algorithm="norm")。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。

• gradient_clip_algorithm (Optional[str, None]) – 使用的梯度裁剪算法。传递 gradient_clip_algorithm="value" 按值裁剪，传递 gradient_clip_algorithm="norm" 按范数裁剪。

返回类型

None

compile(*args, **kwargs)

使用 torch.compile() 编译此模块的 forward。

此模块的 __call__ 方法被编译，所有参数都按原样传递给 torch.compile()。

有关此函数的参数详情，请参阅 torch.compile()。

configure_callbacks()

配置模型特定的回调。当模型被附加时，例如当调用 .fit() 或 .test() 时，这里返回的回调列表或回调将与传递给 Trainer 的 callbacks 参数的回调列表合并。如果这里返回的回调与 Trainer 的回调列表中已有的一个或多个回调类型相同，它将优先并替换它们。此外，Lightning 将确保 ModelCheckpoint 回调最后运行。

返回类型

Union[Sequence[Callback], Callback]

返回

一个回调函数或回调函数列表，这些回调函数将扩展 Trainer 中的回调函数列表。

示例:

def configure_callbacks(self):
    early_stop = EarlyStopping(monitor="val_acc", mode="max")
    checkpoint = ModelCheckpoint(monitor="val_loss")
    return [early_stop, checkpoint]

configure_gradient_clipping(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

对优化器参数执行梯度裁剪。在调用 optimizer_step() 之前执行。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。默认情况下，传递给 Trainer 的值将在此处可用。

• gradient_clip_algorithm (Optional[str, None]) – 要使用的梯度裁剪算法。默认情况下，传递给 Trainer 的值将在此处可用。

示例:

def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm):
    # Implement your own custom logic to clip gradients
    # You can call `self.clip_gradients` with your settings:
    self.clip_gradients(
        optimizer,
        gradient_clip_val=gradient_clip_val,
        gradient_clip_algorithm=gradient_clip_algorithm
    )

返回类型

None

configure_model()

在策略和精度感知的环境中创建模块的钩子。

这在使用分片策略（FSDP 和 DeepSpeed）时特别有用，我们希望立即分片模型以节省内存和初始化时间。对于非分片策略，您可以选择覆盖此钩子或在 init_module() 上下文管理器下初始化您的模型。

此钩子在每个 fit/val/test/predict 阶段都会被调用，因此请确保此钩子的实现是 幂等的，即，在第一次调用钩子后，后续的调用应该是无操作的。

返回类型

None

configure_optimizers()

配置优化器和学习率调度器以优化模型。

configure_sharded_model()

已弃用。

使用 configure_model() 代替。

返回类型

None

static configure_torch_metrics(torch_metrics)

处理 torch_metrics 参数。

返回类型

MetricCollection

cpu()

参见 torch.nn.Module.cpu()。

返回类型

Self

cuda(device=None)

将所有模型参数和缓冲区移动到GPU。这也会使相关的参数和缓冲区成为不同的对象。因此，如果模块在优化期间将驻留在GPU上，则应在构造优化器之前调用此方法。

参数

device (Union[device, int, None]) – 如果指定，所有参数将被复制到该设备。如果为 None，将使用当前的 CUDA 设备索引。

返回

自己

返回类型

Module

property current_epoch: int

Trainer 中的当前时期，如果未附加则为 0。

返回类型

int

property device: device

返回类型

device

property device_mesh: Optional[DeviceMesh]

像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。

返回类型

Optional[ForwardRef, None]

double()

参见 torch.nn.Module.double()。

返回类型

Self

property dtype: Union[str, dtype]

返回类型

Union[str, dtype]

dump_patches: bool = False

property epochs_trained

eval()

将模块设置为评估模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

这等同于 self.train(False)。

请参阅 locally-disable-grad-doc 以比较 .eval() 和几个可能与之混淆的类似机制。

返回

自己

返回类型

Module

property example_input_array: Optional[Union[Tensor, Tuple, Dict]]

示例输入数组是模块在 forward() 方法中可以使用的规范。返回类型解释如下：

• 单个张量：假设模型接受一个参数，即 model.forward(model.example_input_array)

• 元组：输入数组应被解释为位置参数的序列，即 model.forward(*model.example_input_array)

• 字典：输入数组表示命名关键字参数，即 model.forward(**model.example_input_array)

返回类型

Union[Tensor, Tuple, Dict, None]

extra_repr()

设置模块的额外表示。

要打印自定义的额外信息，您应该在您自己的模块中重新实现此方法。单行和多行字符串都是可接受的。

返回类型

str

property fabric: Optional[Fabric]

返回类型

Optional[Fabric, None]

float()

参见 torch.nn.Module.float()。

返回类型

Self

abstract forward(*args, **kwargs)

与 torch.nn.Module.forward() 相同。

参数

• *args – 无论你决定传入 forward 方法的内容。

• **kwargs – 关键字参数也是可能的。

返回类型

Any

返回

你的模型的输出

freeze()

冻结所有参数以进行推理。

示例:

model = MyLightningModule(...)
model.freeze()

返回类型

None

get_buffer(target)

如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的缓冲区的完全限定字符串名称。（有关如何指定完全限定字符串，请参见 get_submodule。）

返回

target 引用的缓冲区

返回类型

torch.Tensor

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是缓冲区的内容

get_extra_state()

返回要包含在模块的 state_dict 中的任何额外状态。

如果需要存储额外状态，请在你的模块中实现此功能以及相应的 set_extra_state()。此函数在构建模块的 state_dict() 时被调用。

请注意，额外的状态应该是可pickle的，以确保state_dict的序列化工作。我们仅对序列化张量提供向后兼容性保证；如果其他对象的序列化pickle形式发生变化，它们可能会破坏向后兼容性。

返回

模块的 state_dict 中存储的任何额外状态

返回类型

object

get_parameter(target)

如果存在，则返回由 target 给出的参数，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的参数的全限定字符串名称。（关于如何指定全限定字符串，请参见 get_submodule。）

返回

target 引用的参数

返回类型

torch.nn.Parameter

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Parameter 的内容

get_submodule(target)

如果存在，返回由 target 指定的子模块，否则抛出错误。

例如，假设你有一个 nn.Module A ，它看起来像这样：

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(该图展示了一个 nn.Module A。A 有一个嵌套的子模块 net_b，它本身有两个子模块 net_c 和 linear。net_c 然后有一个子模块 conv。)

要检查我们是否拥有 linear 子模块，我们会调用 get_submodule("net_b.linear")。要检查我们是否拥有 conv 子模块，我们会调用 get_submodule("net_b.net_c.conv")。

get_submodule 的运行时间受限于 target 中的模块嵌套程度。对 named_modules 的查询可以实现相同的结果，但它是 O(N) 的，其中 N 是传递模块的数量。因此，对于检查某个子模块是否存在的简单检查，应始终使用 get_submodule。

参数

target (str) – 要查找的子模块的完全限定字符串名称。（请参见上面的示例，了解如何指定完全限定的字符串。）

返回

由 target 引用的子模块

返回类型

torch.nn.Module

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Module 的内容

property global_rank: int

当前进程在所有节点和设备中的索引。

返回类型

int

property global_step: int

所有epoch中看到的总训练批次。

如果没有连接训练器，此属性为 0。

返回类型

int

half()

参见 torch.nn.Module.half()。

返回类型

Self

property hparams: Union[AttributeDict, MutableMapping]

通过 save_hyperparameters() 保存的超参数集合。用户可以对其进行修改。对于初始超参数的冻结集，请使用 hparams_initial。

返回类型

Union[AttributeDict, MutableMapping]

返回

可变超参数字典

property hparams_initial: AttributeDict

通过 save_hyperparameters() 保存的超参数集合。这些内容是只读的。可以通过 hparams 手动更新保存的超参数。

返回

不可变的初始超参数

返回类型

AttributeDict

ipu(device=None)

将所有模型参数和缓冲区移动到 IPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在IPU上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs)

从检查点加载模型的主要方式。当Lightning保存检查点时，它会将传递给 __init__ 的参数存储在检查点中的 "hyper_parameters" 下。

通过 **kwargs 指定的任何参数将覆盖存储在 "hyper_parameters" 中的参数。

参数

• checkpoint_path (Union[str, Path, IO]) – 检查点的路径。这也可以是一个URL，或类似文件的对象。

• map_location (Union[device, str, int, Callable[[UntypedStorage, str], Optional[UntypedStorage, None]], Dict[Union[device, str, int], Union[device, str, int]], None]) – 如果你的检查点保存了一个GPU模型，而你现在在CPU上或不同数量的GPU上加载，使用这个来映射到新的设置。其行为与 torch.load() 相同。

• hparams_file (Union[str, Path, None]) –

  可选路径到一个 .yaml 或 .csv 文件，其层次结构如本示例所示:

  drop_prob: 0.2
  dataloader:
      batch_size: 32
  

  您很可能不需要这个，因为 Lightning 总是会将超参数保存到检查点。然而，如果您的检查点权重没有保存超参数，请使用此方法传入一个包含您希望使用的超参数的 .yaml 文件。这些将被转换为 dict 并传递到您的 LightningModule 中使用。

  如果你的模型的 hparams 参数是 Namespace 并且 .yaml 文件具有层次结构，你需要重构你的模型以将 hparams 视为 dict。

• strict (Optional[bool, None]) – 是否严格要求 checkpoint_path 中的键与该模块状态字典返回的键匹配。默认为 True ，除非设置了 LightningModule.strict_loading ，在这种情况下，它默认为 LightningModule.strict_loading 的值。

• **kwargs – 初始化模型所需的任何额外关键字参数。也可以用来覆盖已保存的超参数值。

返回类型

Self

返回

LightningModule 实例，加载了权重和超参数（如果可用）。

备注

load_from_checkpoint 是一个 类 方法。你应该使用你的 LightningModule 类 来调用它，而不是 LightningModule 实例，否则会引发 TypeError。

备注

为了确保所有层都能从检查点加载，如果您的 LightningModule 中重写了此钩子，此函数将在实例化模型后直接调用 configure_model()。但是，请注意 load_from_checkpoint 不支持加载分片检查点，如果模型太大，您可能会耗尽内存。在这种情况下，请考虑通过 .fit(ckpt_path=...) 通过 Trainer 加载。

示例:

# load weights without mapping ...
model = MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')

# or load weights mapping all weights from GPU 1 to GPU 0 ...
map_location = {'cuda:1':'cuda:0'}
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    map_location=map_location
)

# or load weights and hyperparameters from separate files.
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    hparams_file='/path/to/hparams_file.yaml'
)

# override some of the params with new values
model = MyLightningModule.load_from_checkpoint(
    PATH,
    num_layers=128,
    pretrained_ckpt_path=NEW_PATH,
)

# predict
pretrained_model.eval()
pretrained_model.freeze()
y_hat = pretrained_model(x)

load_state_dict(state_dict, strict=True, assign=False)

将参数和缓冲区从 state_dict 复制到此模块及其后代中。

如果 strict 是 True，那么 state_dict 的键必须与该模块的 state_dict() 函数返回的键完全匹配。

警告

如果 assign 是 True，优化器必须在调用 load_state_dict 之后创建，除非 get_swap_module_params_on_conversion() 是 True。

参数

• state_dict (dict) – 包含参数和持久缓冲区的字典。

• strict (bool, optional) – 是否严格要求 state_dict 中的键与该模块的 state_dict() 函数返回的键匹配。默认值：True

• assign (bool, optional) – 当 False 时，当前模块中张量的属性会被保留，而当 True 时，状态字典中张量的属性会被保留。唯一的例外是 Parameter 的 requires_grad 字段，该字段的值会从模块中保留。默认值：False

返回

• missing_keys 是一个包含任何预期键的 str 列表

  此模块中存在但未在提供的 state_dict 中出现的。

• unexpected_keys 是一个包含未预料到的键的 str 列表。

  此模块预期但未在提供的 state_dict 中出现。

返回类型

带有 missing_keys 和 unexpected_keys 字段的 NamedTuple

备注

如果一个参数或缓冲区被注册为 None ，并且其对应的键存在于 state_dict 中，load_state_dict() 将引发一个 RuntimeError 。

property local_rank: int

当前进程在单个节点中的索引。

返回类型

int

log(name, value, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, metric_attribute=None, rank_zero_only=False)

记录一个键值对。

示例:

self.log('train_loss', loss)

每个钩子的默认行为在此处有文档说明：extensions/logging:自动日志记录。

参数

• name (str) – 日志的键。如果使用DDP或其他分布式策略，必须在所有进程中保持一致。

• value (Union[Metric, Tensor, int, float]) – 要记录的值。可以是 float、Tensor 或 Metric。

• prog_bar (bool) – 如果 True 则记录到进度条。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• on_epoch (Optional[bool, None]) – 如果 True ，记录每个epoch累积的指标。默认值由钩子决定。详情请参见 extensions/logging:自动日志记录 。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图表。

• sync_dist (bool) – 如果 True，将在设备间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的DDP组。

• add_dataloader_idx (bool) – 如果 True，将当前数据加载器的索引附加到名称（当使用多个数据加载器时）。如果为 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的 batch_size。这将直接从加载的批次中推断出来，但对于某些数据结构，您可能需要显式提供它。

• metric_attribute (Optional[str, None]) – 要恢复指标状态，Lightning 需要模型中 torchmetrics.Metric 的引用。如果它是模型属性，则会自动找到。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

log_dict(dictionary, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, rank_zero_only=False)

一次性记录一个值的字典。

示例:

values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n}
self.log_dict(values)

参数

• dictionary (Union[Mapping[str, Union[Metric, Tensor, int, float]], MetricCollection]) – 键值对。如果使用DDP或其他分布式策略，键必须在所有进程中保持一致。值可以是 float 、 Tensor 、 Metric 或 MetricCollection 。

• prog_bar (bool) – 如果 True 则记录到进度基础。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。None 自动记录训练步骤，但不记录验证/测试步骤。默认值由钩子决定。详情请参阅 extensions/logging:自动记录。

• on_epoch (Optional[bool, None]) – 如果 True 记录累积的 epoch 指标。None 自动记录 val/test 步骤的日志，但不包括 training_step。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图形

• sync_dist (bool) – 如果 True，将在GPU/TPU之间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的 ddp 组。

• add_dataloader_idx (bool) – 如果 True，则将当前数据加载器的索引附加到名称（在使用多个数据加载器时）。如果 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的批次大小。这将直接从加载的批次中推断出来，但某些数据结构可能需要显式提供它。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

property logger: Optional[Union[Logger, Logger]]

Trainer 中对 logger 对象的引用。

返回类型

Union[Logger, Logger, None]

property loggers: Union[List[Logger], List[Logger]]

参考 Trainer 中的日志记录器列表。

返回类型

Union[List[Logger], List[Logger]]

lr_scheduler_step(scheduler, metric)

重写此方法以调整 Trainer 调用每个调度器的默认方式。默认情况下，Lightning 会根据每个调度器的 interval 调用 step() 方法，如示例所示。

参数

• scheduler (Union[LRScheduler, ReduceLROnPlateau]) – 学习率调度器。

• metric (Optional[Any, None]) – 用于 ReduceLROnPlateau 等调度器的监视器值。

示例:

# DEFAULT
def lr_scheduler_step(self, scheduler, metric):
    if metric is None:
        scheduler.step()
    else:
        scheduler.step(metric)

# Alternative way to update schedulers if it requires an epoch value
def lr_scheduler_step(self, scheduler, metric):
    scheduler.step(epoch=self.current_epoch)

返回类型

None

lr_schedulers()

返回在训练期间使用的学习率调度器。对于手动优化很有用。

返回类型

Union[None, List[Union[LRScheduler, ReduceLROnPlateau]], LRScheduler, ReduceLROnPlateau]

返回

一个调度器，或者在存在多个调度器的情况下是一个调度器列表，如果在 configure_optimizers() 中没有返回调度器，则为 None。

manual_backward(loss, *args, **kwargs)

当你手动进行优化时，请直接从你的 training_step() 调用此方法。通过使用此方法，Lightning 可以确保在使用混合精度时应用所有适当的缩放。

更多示例请参见 手动优化。

示例:

def training_step(...):
    opt = self.optimizers()
    loss = ...
    opt.zero_grad()
    # automatically applies scaling, etc...
    self.manual_backward(loss)
    opt.step()

参数

• loss (Tensor) – 要计算梯度的张量。必须附加一个图。

• *args – 要转发给 backward() 的附加位置参数

• **kwargs – 要转发给 backward() 的额外关键字参数

返回类型

None

modules()

返回网络中所有模块的迭代器。

生成器

模块 – 网络中的一个模块

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

返回类型

Iterator[Module]

named_buffers(prefix='', recurse=True, remove_duplicate=True)

返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。

参数

• prefix (str) – 要添加到所有缓冲区名称的前缀。

• recurse (bool, optional) – 如果为 True，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。默认为 True。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的缓冲区。默认为 True。

生成器

(str, torch.Tensor) – 包含名称和缓冲区的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

返回类型

Iterator[Tuple[str, Tensor]]

named_children()

返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。

生成器

(str, Module) – 包含名称和子模块的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

返回类型

Iterator[Tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)

返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。

参数

• memo (Optional[Set[Module], None]) – 一个备忘录，用于存储已添加到结果中的模块集合

• prefix (str) – 一个将添加到模块名称前的前缀

• remove_duplicate (bool) – 是否在结果中移除重复的模块实例

生成器

(str, Module) – 名称和模块的元组

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)

返回一个模块参数的迭代器，同时生成参数的名称和参数本身。

参数

• prefix (str) – 要添加到所有参数名称前的前缀。

• recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的参数。默认为 True。

生成器

(str, Parameter) – 包含名称和参数的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

返回类型

Iterator[Tuple[str, Parameter]]

on_after_backward()

在 loss.backward() 之后和优化器步进之前调用。

备注

如果使用原生 AMP，梯度在此刻将不会被缩放。如果你需要未缩放的梯度，请使用 on_before_optimizer_step。

返回类型

None

on_after_batch_transfer(batch, dataloader_idx)

覆盖此方法以在批次转移到设备后对其进行批量增强或修改。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_after_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = gpu_transforms(batch['x'])
    return batch

参见

• on_before_batch_transfer()

• transfer_batch_to_device()

on_before_backward(loss)

在 loss.backward() 之前调用。

参数

loss (Tensor) – 损失除以梯度累积的批次数量，如果使用AMP则进行缩放。

返回类型

None

on_before_batch_transfer(batch, dataloader_idx)

在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_before_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = transforms(batch['x'])
    return batch

参见

• on_after_batch_transfer()

• transfer_batch_to_device()

on_before_optimizer_step(optimizer)

在 optimizer.step() 之前调用。

如果使用梯度累积，钩子会在梯度累积完成后被调用。参见：:paramref:`~pytorch_lightning.trainer.trainer.Trainer.accumulate_grad_batches`。

如果使用AMP，损失将在调用此钩子之前未缩放。有关梯度缩放的更多信息，请参阅这些 文档。

如果裁剪梯度，梯度尚未被裁剪。

参数

optimizer (Optimizer) – 当前使用的优化器。

示例:

def on_before_optimizer_step(self, optimizer):
    # example to inspect gradient information in tensorboard
    if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
        for k, v in self.named_parameters():
            self.logger.experiment.add_histogram(
                tag=k, values=v.grad, global_step=self.trainer.global_step
            )

返回类型

None

on_before_zero_grad(optimizer)

在 training_step() 之后和 optimizer.zero_grad() 之前调用。

在训练循环中，调用优化器步骤之后和清零梯度之前。这是一个检查更新后权重信息的好地方。

这是调用它的地方:

for optimizer in optimizers:
    out = training_step(...)

    model.on_before_zero_grad(optimizer) # < ---- called here
    optimizer.zero_grad()

    backward()

参数

optimizer (Optimizer) – 应该被置零的梯度优化器。

返回类型

None

on_fit_end()

在拟合的最后调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

on_fit_start()

在拟合开始时调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

property on_gpu: bool

如果该模型当前位于GPU上，则返回 True。

在LightningModule周围设置标志以区分CPU和GPU行为非常有用。

返回类型

bool

on_load_checkpoint(checkpoint)

由 Lightning 调用以恢复您的模型。如果您使用 on_save_checkpoint() 保存了某些内容，这是您恢复这些内容的机会。

参数

checkpoint (Dict[str, Any]) – 加载的检查点

示例:

def on_load_checkpoint(self, checkpoint):
    # 99% of the time you don't need to implement this method
    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']

备注

Lightning 自动恢复全局步数、轮次和训练状态，包括 amp 缩放。您无需恢复任何与训练相关的内容。

返回类型

None

on_predict_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后调用预测循环。

参数

• outputs (Optional[Any, None]) – predict_step(x) 的输出

• batch (Any) – 由预测 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_batch_start(batch, batch_idx, dataloader_idx=0)

在预测循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_start()

在预测开始时调用。

返回类型

None

on_predict_model_eval()

在预测循环开始时调用。

默认情况下，预测循环在开始之前会在 LightningModule 上调用 .eval()。重写此钩子以更改行为。

返回类型

None

on_predict_start()

在预测开始时调用。

返回类型

None

on_save_checkpoint(checkpoint)

当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。

参数

checkpoint (Dict[str, Any]) – 在转储到文件之前的完整检查点字典。此钩子的实现可以将额外数据插入此字典中。

示例:

def on_save_checkpoint(self, checkpoint):
    # 99% of use cases you don't need to implement this method
    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object

备注

Lightning 保存了训练的所有方面（如 epoch、全局步数等），包括 amp 缩放。您无需存储任何关于训练的信息。

返回类型

None

on_test_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后在测试循环中调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – test_step(x) 的输出

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_batch_start(batch, batch_idx, dataloader_idx=0)

在测试循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_end()

在测试结束时调用。

返回类型

None

on_test_epoch_end()

在每个epoch的测试循环结束时调用。

返回类型

None

on_test_epoch_start()

在epoch开始时在测试循环中调用。

返回类型

None

on_test_model_eval()

在测试循环开始时调用。

默认情况下，测试循环在开始前会调用 .eval() 方法在 LightningModule 上。重写此钩子以改变行为。另请参见 on_test_model_train()。

返回类型

None

on_test_model_train()

在测试循环结束时调用。

测试循环默认会将 LightningModule 的 training 模式恢复到开始测试之前的状态。覆盖此钩子以更改行为。另请参见 on_test_model_eval()。

返回类型

None

on_test_start()

在测试开始时调用。

返回类型

None

on_train_batch_end(outputs, batch, batch_idx)

在训练循环中批次处理后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – 训练步骤(x)的输出

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

备注

这里的 outputs["loss"] 值将是相对于 accumulate_grad_batches 的损失归一化值，该损失从 training_step 返回。

返回类型

None

on_train_batch_start(batch, batch_idx)

在训练循环中，在该批次发生任何事情之前调用。

如果你在这里返回 -1，你将跳过当前轮次的剩余训练。

参数

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

返回类型

Optional[int, None]

on_train_end()

在训练结束前调用，在记录器实验关闭之前。

返回类型

None

on_train_epoch_end()

在训练循环的每个epoch结束时调用。

要在每个epoch结束时访问所有批处理输出，您可以将步骤输出缓存为 LightningModule 的一个属性，并在此钩子中访问它们：

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # free up the memory
        self.training_step_outputs.clear()

on_train_epoch_start()

在训练循环中，在epoch开始时调用。

返回类型

None

on_train_start()

在完整性检查后，训练开始时调用。

返回类型

None

on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次验证循环后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – validation_step(x) 的输出

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_batch_start(batch, batch_idx, dataloader_idx=0)

在验证循环中，对该批次进行任何操作之前调用。

参数

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_end()

在验证结束时调用。

返回类型

None

on_validation_epoch_end()

在epoch结束时的验证循环中调用。

on_validation_epoch_start()

在epoch开始时验证循环中调用。

返回类型

None

on_validation_model_eval()

在验证循环开始时调用。

默认情况下，验证循环在开始前会在 LightningModule 上调用 .eval()。覆盖此钩子以更改行为。另请参见 on_validation_model_train()。

返回类型

None

on_validation_model_train()

当验证循环结束时调用。

默认情况下，验证循环会将 LightningModule 的 training 模式恢复到开始验证之前的状态。重写此钩子以更改行为。另请参见 on_validation_model_eval()。

返回类型

None

on_validation_model_zero_grad()

在进入验证循环之前，由训练循环调用来释放梯度。

返回类型

None

on_validation_start()

在验证开始时调用。

返回类型

None

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)

重写此方法以调整 Trainer 调用优化器的默认方式。

默认情况下，Lightning 会调用 step() 和 zero_grad()，如示例所示。当 Trainer(accumulate_grad_batches != 1) 时，在累积阶段不会调用此方法（和 zero_grad()）。手动优化时，重写此钩子没有好处。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Union[Optimizer, LightningOptimizer]) – 一个 PyTorch 优化器

• optimizer_closure (Optional[Callable[[], Any], None]) – 优化器闭包。此闭包必须执行，因为它包含了 training_step()、optimizer.zero_grad() 和 backward() 的调用。

示例:

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    # Add your custom logic to run directly before `optimizer.step()`

    optimizer.step(closure=optimizer_closure)

    # Add your custom logic to run directly after `optimizer.step()`

返回类型

None

optimizer_zero_grad(epoch, batch_idx, optimizer)

重写此方法以改变 optimizer.zero_grad() 的默认行为。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Optimizer) – 一个 PyTorch 优化器

示例:

# DEFAULT
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad()

# Set gradients to `None` instead of zero to improve performance (not required on `torch>=2.0.0`).
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad(set_to_none=True)

参见 torch.optim.Optimizer.zero_grad() 以了解上述示例的解释。

返回类型

None

optimizers(use_pl_optimizer=True)

返回训练期间使用的优化器。对于手动优化很有用。

参数

use_pl_optimizer (bool) – 如果 True，将会把优化器包装在 LightningOptimizer 中，以自动处理精度、分析和步数调用的计数，以进行适当的日志记录和检查点保存。它特别包装了 step 方法，不支持没有此方法的自定义优化器。

返回类型

Union[Optimizer, LightningOptimizer, _FabricOptimizer, List[Optimizer], List[LightningOptimizer], List[_FabricOptimizer]]

返回

一个单一的优化器，或者在存在多个优化器的情况下，一个优化器列表。

property output_chunk_length: Optional[int]

模型一次预测的时间步数。

返回类型

Optional[int, None]

parameters(recurse=True)

返回一个遍历模块参数的迭代器。

这通常传递给一个优化器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

生成器

参数 – 模块参数

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Parameter]

pred_batch_size: Optional[int]

pred_mc_dropout: Optional[bool]

pred_n: Optional[int]

pred_n_jobs: Optional[int]

pred_num_samples: Optional[int]

pred_roll_size: Optional[int]

predict_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。

有关多个数据加载器的更多信息，请参阅此 部分。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• predict()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

返回

一个 torch.utils.data.DataLoader 或它们的序列，指定预测样本。

predict_likelihood_parameters: Optional[bool]

predict_step(batch, batch_idx, dataloader_idx=None)

执行预测步骤

批处理

Darts 的 InferenceDataset 的输出 - 包含 (过去的目标, 过去的协变量, 历史未来协变量, 未来协变量, 未来过去协变量, 输入时间序列, 预测开始时间步) 的元组

batch_idx

当前批次的批次索引

dataloader_idx

数据加载器索引

返回类型

Sequence[TimeSeries]

prepare_data()

使用此方法下载和准备数据。使用多个进程（分布式设置）下载和保存数据将导致数据损坏。Lightning 确保此方法仅在一个进程内调用，因此您可以安全地在其中添加下载逻辑。

警告

不要将状态设置到模型中（使用 setup 代替），因为这不会在每个设备上被调用

示例:

def prepare_data(self):
    # good
    download_data()
    tokenize()
    etc()

    # bad
    self.split = data_split
    self.some_state = some_other_state()

在分布式环境中，prepare_data 可以通过两种方式调用（使用 prepare_data_per_node）

1. 每个节点一次。这是默认设置，仅在 LOCAL_RANK=0 时调用。

2. 总共一次。仅在 GLOBAL_RANK=0 时调用。

示例:

# DEFAULT
# called once per node on LOCAL_RANK=0 of that node
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = True

# call on GLOBAL_RANK=0 (great for shared file systems)
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = False

在请求数据加载器之前调用此方法：

model.prepare_data()
initialize_distributed()
model.setup(stage)
model.train_dataloader()
model.val_dataloader()
model.test_dataloader()
model.predict_dataloader()

返回类型

None

prepare_data_per_node: bool

print(*args, **kwargs)

仅从进程 0 打印。在任何分布式模式下使用此功能以仅记录一次。

参数

• *args – 要打印的内容。与Python内置的print函数相同。

• **kwargs – 与Python内置的print函数相同。

示例:

def forward(self, x):
    self.print(x, 'in forward')

返回类型

None

register_backward_hook(hook)

在模块上注册一个反向钩子。

此函数已被 register_full_backward_hook() 取代，并且此函数的行为将在未来版本中发生变化。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_buffer(name, tensor, persistent=True)

向模块添加一个缓冲区。

这通常用于注册一个不应被视为模型参数的缓冲区。例如，BatchNorm 的 running_mean 不是参数，而是模块状态的一部分。缓冲区默认是持久的，并将与参数一起保存。可以通过将 persistent 设置为 False 来改变此行为。持久缓冲区和非持久缓冲区之间的唯一区别是后者不会成为此模块的 state_dict 的一部分。

缓冲区可以通过给定的名称作为属性访问。

参数

• name (str) – 缓冲区的名称。可以通过给定的名称从此模块访问缓冲区。

• tensor (Tensor or None) – 要注册的缓冲区。如果为 None ，则运行在缓冲区上的操作，例如 cuda ，将被忽略。如果为 None ，该缓冲区将 不 包含在模块的 state_dict 中。

• persistent (bool) – 缓冲区是否是此模块的 state_dict 的一部分。

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

返回类型

None

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

在模块上注册一个前向钩子。

钩子将在每次 forward() 计算输出后被调用。

如果 with_kwargs 是 False 或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward 。钩子可以修改输出。它可以就地修改输入，但这不会影响前向传播，因为这是在调用 forward() 之后调用的。钩子应具有以下签名:

hook(module, args, output) -> None or modified output

如果 with_kwargs 是 True，前向钩子将被传递给前向函数的 kwargs，并期望返回可能被修改的输出。钩子应具有以下签名:

hook(module, args, kwargs, output) -> None or modified output

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果 True，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之后触发。请注意，使用 register_module_forward_hook() 注册的全局 forward 钩子将在通过此方法注册的所有钩子之前触发。默认值：False

• with_kwargs (bool) – 如果 True，hook 将会传递给 forward 函数的 kwargs。默认值：False

• always_call (bool) – 如果 True ，无论在调用模块时是否引发异常，都将运行 hook 。默认值： False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

在模块上注册一个前向预钩子。

钩子将在每次调用 forward() 之前被调用。

如果 with_kwargs 为假或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward。钩子可以修改输入。用户可以在钩子中返回一个元组或单个修改后的值。如果返回单个值（除非该值已经是元组），我们将把该值包装成一个元组。钩子应具有以下签名:

hook(module, args) -> None or modified input

如果 with_kwargs 为真，前向预钩子将传递给前向函数的 kwargs。如果钩子修改了输入，则应返回 args 和 kwargs。钩子应具有以下签名:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之后触发。请注意，通过 register_module_forward_pre_hook() 注册的全局 forward_pre 钩子将在所有通过此方法注册的钩子之前触发。默认值：False

• with_kwargs (bool) – 如果为真，hook 将被传递给 forward 函数的 kwargs。默认值：False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)

在模块上注册一个反向钩子。

每次计算与模块相关的梯度时，钩子都会被调用，即当且仅当计算与模块输出相关的梯度时，钩子才会执行。钩子应具有以下签名:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input 和 grad_output 是包含相对于输入和输出的梯度的元组。钩子不应修改其参数，但可以选择返回一个新的相对于输入的梯度，该梯度将用于替代 grad_input 在后续计算中。grad_input 将仅对应于作为位置参数给出的输入，所有 kwarg 参数都将被忽略。grad_input 和 grad_output 中的条目将为 None 用于所有非 Tensor 参数。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入或输出，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之后触发。请注意，通过 register_module_full_backward_hook() 注册的全局 backward 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

在模块上注册一个反向预钩子。

每次计算模块的梯度时，都会调用该钩子。钩子应具有以下签名:

hook(module, grad_output) -> tuple[Tensor] or None

The grad_output 是一个元组。钩子不应修改其参数，但它可以选择返回一个新的相对于输出的梯度，该梯度将用于替代 grad_output 在后续计算中。grad_output 中的条目对于所有非 Tensor 参数将为 None。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之后触发。请注意，通过 register_module_full_backward_pre_hook() 注册的全局 backward_pre 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)

注册一个在模块的 load_state_dict 调用后运行的后置钩子。

它应该有以下签名:

hook(module, incompatible_keys) -> None

module 参数是当前注册此钩子的模块，而 incompatible_keys 参数是一个由属性 missing_keys 和 unexpected_keys 组成的 NamedTuple。missing_keys 是一个包含缺失键的 str 的 list，而 unexpected_keys 是一个包含意外键的 str 的 list。

如果需要，可以就地修改给定的 incompatible_keys。

请注意，当使用 strict=True 调用 load_state_dict() 时执行的检查会受到钩子对 missing_keys 或 unexpected_keys 所做修改的影响，这是预期的。对任一组键的添加将导致在 strict=True 时抛出错误，而清除掉缺失和意外的键将避免错误。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_module(name, module)

别名 add_module()。

返回类型

None

register_parameter(name, param)

向模块添加一个参数。

该参数可以通过给定的名称作为属性访问。

参数

• name (str) – 参数的名称。可以使用给定的名称从此模块访问该参数。

• param (Parameter or None) – 要添加到模块的参数。如果为 None ，则运行在参数上的操作，例如 cuda ，将被忽略。如果为 None ，该参数将 不 包含在模块的 state_dict 中。

返回类型

None

register_state_dict_pre_hook(hook)

为 state_dict() 方法注册一个预钩子。

这些钩子将在调用 self 上的 state_dict 之前，使用参数 self、prefix 和 keep_vars 进行调用。注册的钩子可以用于在调用 state_dict 之前执行预处理。

requires_grad_(requires_grad=True)

更改是否应在此模块的参数上记录自动梯度操作。

此方法就地设置参数的 requires_grad 属性。

这种方法有助于冻结模块的一部分以进行微调或单独训练模型的一部分（例如，GAN训练）。

请参阅 locally-disable-grad-doc 以比较 .requires_grad_() 和几个可能与之混淆的类似机制。

参数

requires_grad (bool) – 是否应在此模块的参数上记录自动求导操作。默认值：True。

返回

自己

返回类型

Module

save_hyperparameters(*args, ignore=None, frame=None, logger=True)

将参数保存到 hparams 属性中。

参数

• args (Any) – dict、NameSpace 或 OmegaConf 的单个对象，或来自类 __init__ 的字符串名称或参数

• ignore (Union[Sequence[str], str, None]) – 一个参数名称或从类 __init__ 中忽略的参数名称列表

• frame (Optional[frame, None]) – 一个框架对象。默认是 None

• logger (bool) – 是否将超参数发送到记录器。默认值：True

示例:

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # manually assign arguments
...         self.save_hyperparameters('arg1', 'arg3')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class AutomaticArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # equivalent automatic
...         self.save_hyperparameters()
...     def forward(self, *args, **kwargs):
...         ...
>>> model = AutomaticArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg2": abc
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class SingleArgModel(HyperparametersMixin):
...     def __init__(self, params):
...         super().__init__()
...         # manually assign single argument
...         self.save_hyperparameters(params)
...     def forward(self, *args, **kwargs):
...         ...
>>> model = SingleArgModel(Namespace(p1=1, p2='abc', p3=3.14))
>>> model.hparams
"p1": 1
"p2": abc
"p3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # pass argument(s) to ignore as a string or in a list
...         self.save_hyperparameters(ignore='arg2')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

返回类型

None

set_extra_state(state)

设置包含在加载的 state_dict 中的额外状态。

此函数从 load_state_dict() 调用，以处理在 state_dict 中找到的任何额外状态。如果你的模块需要在 state_dict 中存储额外状态，请实现此函数以及相应的 get_extra_state()。

参数

state (dict) – 来自 state_dict 的额外状态

返回类型

None

set_mc_dropout(active)

set_predict_parameters(n, num_samples, roll_size, batch_size, n_jobs, predict_likelihood_parameters, mc_dropout)

在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。

返回类型

None

setup(stage)

在 fit（训练 + 验证）、验证、测试或预测开始时调用。当你需要动态构建模型或调整模型时，这是一个很好的钩子。在使用 DDP 时，此钩子会在每个进程上调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

示例:

class LitModel(...):
    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

        # don't do this
        self.something = else

    def setup(self, stage):
        data = load_data(...)
        self.l1 = nn.Linear(28, data.num_classes)

返回类型

None

share_memory()

参见 torch.Tensor.share_memory_()。

返回类型

~T

state_dict(*args, destination=None, prefix='', keep_vars=False)

返回一个包含模块整体状态引用的字典。

参数和持久缓冲区（例如运行平均值）都包含在内。键是对应的参数和缓冲区名称。设置为 None 的参数和缓冲区不包含在内。

备注

返回的对象是一个浅拷贝。它包含对模块参数和缓冲区的引用。

警告

目前，state_dict() 也按顺序接受 destination、prefix 和 keep_vars 的位置参数。然而，这正在被弃用，未来版本将强制使用关键字参数。

警告

请避免使用参数 destination，因为它不是为最终用户设计的。

参数

• destination (dict, optional) – 如果提供了，模块的状态将被更新到字典中，并返回相同的对象。否则，将创建并返回一个 OrderedDict。默认值：None。

• prefix (str, optional) – 在 state_dict 中组合键时添加到参数和缓冲区名称的前缀。默认值：''。

• keep_vars (bool, optional) – 默认情况下，状态字典中返回的 Tensor 是与自动求导分离的。如果设置为 True，则不会执行分离操作。默认值：False。

返回

包含模块整个状态的字典

返回类型

dict

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

property strict_loading: bool

确定 Lightning 如何使用 .load_state_dict(…, strict=model.strict_loading) 加载此模型。

返回类型

bool

property supports_probabilistic_prediction: bool

返回类型

bool

teardown(stage)

在 fit（训练 + 验证）、验证、测试或预测结束时调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

返回类型

None

test_dataloader()

一个可迭代对象或多个可迭代对象的集合，指定测试样本。

有关多个数据加载器的更多信息，请参阅此 部分。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• test()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要测试数据集和 test_step() ，你不需要实现这个方法。

返回类型

Any

test_step(*args, **kwargs)

对测试集中的单个批次数据进行操作。在这一步中，您通常会生成示例或计算任何感兴趣的内容，例如准确性。

参数

• batch – 你的数据可迭代对象的输出，通常是一个 DataLoader。

• batch_idx – 这一批的索引。

• dataloader_idx – 生成此批次的 dataloader 的索引。（仅在使用多个 dataloader 时）

返回类型

Union[Tensor, Mapping[str, Any], None]

返回

• Tensor - 损失张量

• dict - 一个字典。可以包含任何键，但必须包含键 'loss'。

• None - 跳到下一个批次。

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...

# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

示例:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

如果你传入多个测试数据加载器，test_step() 将会有一个额外的参数。我们建议将默认值设置为0，这样你可以快速在单个和多个数据加载器之间切换。

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

备注

如果你不需要测试，你就不需要实现这个方法。

备注

当调用 test_step() 时，模型已处于评估模式，并且已禁用 PyTorch 梯度。在测试轮次结束时，模型返回到训练模式，并且启用梯度。

to(*args, **kwargs)

参见 torch.nn.Module.to()。

返回类型

Self

to_dtype(dtype)

将模块精度（默认为 float32）转换为另一种精度。

to_empty(*, device, recurse=True)

将参数和缓冲区移动到指定设备，而不复制存储。

参数

• device (torch.device) – 此模块中参数和缓冲区的所需设备。

• recurse (bool) – 子模块的参数和缓冲区是否应递归地移动到指定设备。

返回

自己

返回类型

Module

to_onnx(file_path, input_sample=None, **kwargs)

将模型保存为 ONNX 格式。

参数

• file_path (Union[str, Path]) – 应保存 onnx 模型的文件路径。

• input_sample (Optional[Any, None]) – 用于跟踪的输入。默认值：None（使用 self.example_input_array）

• **kwargs – 将被传递给 torch.onnx.export 函数。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)

model = SimpleModel()
input_sample = torch.randn(1, 64)
model.to_onnx("export.onnx", input_sample, export_params=True)

返回类型

None

to_torchscript(file_path=None, method='script', example_inputs=None, **kwargs)

默认情况下，将整个模型编译为 ScriptModule。如果你想使用追踪，请提供参数 method='trace' 并确保提供了 example_inputs 参数，或者模型设置了 example_input_array。如果你想自定义被脚本的模块，你应该重写这个方法。如果你想要返回多个模块，我们建议使用字典。

参数

• file_path (Union[str, Path, None]) – 保存 torchscript 的路径。默认：None（不保存文件）。

• method (Optional[str, None]) – 是否使用 TorchScript 的 script 或 trace 方法。默认值：’script’

• example_inputs (Optional[Any, None]) – 当方法设置为’trace’时，用于进行追踪的输入。默认值：None（使用 example_input_array）

• **kwargs – 将传递给 torch.jit.script() 或 torch.jit.trace() 函数的额外参数。

备注

• 需要实现 forward() 方法。

• 导出的脚本将被设置为评估模式。

• 建议您安装最新支持版本的 PyTorch 以无限制地使用此功能。另请参阅 torch.jit 文档以了解支持的功能。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

model = SimpleModel()
model.to_torchscript(file_path="model.pt")

torch.jit.save(model.to_torchscript(
    file_path="model_trace.pt", method='trace', example_inputs=torch.randn(1, 64))
)

返回类型

Union[ScriptModule, Dict[str, ScriptModule]]

返回

无论 file_path 是否定义，此 LightningModule 都可以作为 torchscript 使用。

toggle_optimizer(optimizer)

确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。

它与 untoggle_optimizer() 一起工作，以确保 param_requires_grad_state 被正确重置。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要切换的优化器。

返回类型

None

train(mode=True)

设置模块为训练模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

参数

mode (bool) – 是否设置训练模式（True）或评估模式（False）。默认值：True。

返回

自己

返回类型

Module

train_criterion_reduction: Optional[str]

train_dataloader()

指定训练样本的可迭代对象或可迭代对象的集合。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• fit()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

property trainer: Trainer

返回类型

Trainer

training: bool

training_step(train_batch, batch_idx)

执行训练步骤

返回类型

Tensor

transfer_batch_to_device(batch, device, dataloader_idx)

如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。

以下列出的数据类型（以及它们的任意嵌套）都是开箱即用的：

• torch.Tensor 或任何实现了 .to(…) 的对象

• 列表

• 字典

• 元组

对于其他情况，你需要定义数据如何移动到目标设备（CPU、GPU、TPU等）。

备注

此钩子应仅传输数据，不应修改数据，也不应将数据移动到除作为参数传入的设备之外的任何其他设备（除非你知道自己在做什么）。要检查此钩子的当前执行状态，可以使用 self.trainer.training/testing/validating/predicting，以便可以根据你的需求添加不同的逻辑。

参数

• batch (Any) – 一批需要转移到新设备的数据。

• device (device) – 在 PyTorch 中定义的目标设备。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

对新设备上数据的引用。

示例:

def transfer_batch_to_device(self, batch, device, dataloader_idx):
    if isinstance(batch, CustomBatch):
        # move all tensors in your custom data structure to the device
        batch.samples = batch.samples.to(device)
        batch.targets = batch.targets.to(device)
    elif dataloader_idx == 0:
        # skip device transfer for the first dataloader or anything you wish
        pass
    else:
        batch = super().transfer_batch_to_device(batch, device, dataloader_idx)
    return batch

参见

• move_data_to_device()

• apply_to_collection()

type(dst_type)

参见 torch.nn.Module.type()。

返回类型

Self

unfreeze()

解冻所有参数以进行训练。

model = MyLightningModule(...)
model.unfreeze()

返回类型

None

untoggle_optimizer(optimizer)

重置通过 toggle_optimizer() 切换的所需梯度的状态。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要取消选择的优化器。

返回类型

None

val_criterion_reduction: Optional[str]

val_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• fit()

• validate()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要验证数据集和 validation_step()，你不需要实现这个方法。

返回类型

Any

validation_step(val_batch, batch_idx)

执行验证步骤

返回类型

Tensor

xpu(device=None)

将所有模型参数和缓冲区移动到XPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在 XPU 上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

zero_grad(set_to_none=True)

重置所有模型参数的梯度。

更多上下文请参见 torch.optim.Optimizer 下的类似函数。

参数

set_to_none (bool) – 而不是设置为零，将梯度设置为 None。详情请参阅 torch.optim.Optimizer.zero_grad()。

返回类型

None

class darts.models.forecasting.pl_forecasting_module.PLPastCovariatesModule(input_chunk_length, output_chunk_length, output_chunk_shift=0, train_sample_shape=None, loss_fn=MSELoss(), torch_metrics=None, likelihood=None, optimizer_cls=<class 'torch.optim.adam.Adam'>, optimizer_kwargs=None, lr_scheduler_cls=None, lr_scheduler_kwargs=None, use_reversible_instance_norm=False)

基类：PLForecastingModule, ABC

基于 PyTorch Lightning 的预测模块。

这个类旨在通过继承来创建一个新的基于 PyTorch Lightning 的预测模块。当子类化这个类时，请确保添加以下具有给定签名的方法：

• PLForecastingModule.__init__()

• PLForecastingModule.forward()

• PLForecastingModule._produce_train_output()

• PLForecastingModule._get_batch_prediction()

在子类 MyModel 的 __init__() 函数中调用 super(MyModel, self).__init__(**kwargs)，其中 kwargs 是 PLForecastingModule 的参数。

参数

• input_chunk_length (int) – 过去作为模型输入的时间步数（每个块）。适用于目标序列，以及过去和/或未来协变量（如果模型支持）。

• output_chunk_length (int) – 内部模型一次预测的时间步数（每个块）。此外，模型输入中使用的未来协变量的未来值数量（如果模型支持未来协变量）。这与 predict() 中使用的预测范围 n 不同，后者是使用一次性或自回归预测生成的所需预测点数。设置 n <= output_chunk_length 可以防止自回归。当协变量不足以延伸到未来，或禁止模型使用过去和/或未来协变量的未来值进行预测时（取决于模型的协变量支持），这很有用。

• train_sample_shape (Optional[Tuple, None]) – 模型的输入形状，用于在不调用 fit_from_dataset 的情况下实例化模型，并对用于重新训练或预测的新训练/推理数据集进行完整性检查。

• loss_fn (_Loss) – 用于训练的 PyTorch 损失函数。如果指定了 likelihood 参数，则此参数将被概率模型忽略。默认值：torch.nn.MSELoss()。

• torch_metrics (Union[Metric, MetricCollection, None]) – 用于评估的 torch 指标或 MetricCollection。可用指标的完整列表可以在 https://torchmetrics.readthedocs.io/en/latest/ 找到。默认值：None。

• likelihood (Optional[Likelihood, None]) – Darts 的 似然模型 之一，用于概率预测。默认值：None。

• optimizer_cls (Optimizer) – 要使用的 PyTorch 优化器类。默认值：torch.optim.Adam。

• optimizer_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 优化器的键值参数（例如，{'lr': 1e-3} 用于指定学习率）。否则将使用所选 optimizer_cls 的默认值。默认值：None。

• lr_scheduler_cls (Optional[_LRScheduler, None]) – 可选地，指定要使用的 PyTorch 学习率调度器类。指定 None 对应于使用恒定的学习率。默认值：None。

• lr_scheduler_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 学习率调度器的键值参数。默认值：None。

• use_reversible_instance_norm (bool) – 是否使用可逆实例归一化 RINorm 来对抗分布偏移，如 [1] 所示。它仅应用于目标序列的特征，而不应用于协变量。

引用

1

T. Kim et al. “Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift”, https://openreview.net/forum?id=cGDAkQo1C0p

属性

automatic_optimization	如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。
current_epoch	Trainer 中的当前时期，如果未附加则为 0。
device_mesh	像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。
example_input_array	示例输入数组是模块在 forward() 方法中可以消费的规范。
global_rank	当前进程在所有节点和设备中的索引。
global_step	所有epoch中看到的总训练批次。
hparams	使用 save_hyperparameters() 保存的超参数集合。
hparams_initial	使用 save_hyperparameters() 保存的超参数集合。
local_rank	当前进程在单个节点中的索引。
logger	Trainer 中对 logger 对象的引用。
loggers	参考 Trainer 中的日志记录器列表。
on_gpu	如果该模型当前位于GPU上，则返回 True。
output_chunk_length	模型一次预测的时间步数。
strict_loading	确定 Lightning 如何使用 .load_state_dict(..., strict=model.strict_loading) 加载此模型。

设备
dtype
已训练的轮数
fabric
supports_probabilistic_prediction
训练者

方法

add_module(name, module)	将一个子模块添加到当前模块。
all_gather(data[, group, sync_grads])	从多个进程中收集张量或张量集合。
apply(fn)	递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。
backward(loss, *args, **kwargs)	在 training_step() 中返回的损失上调用以执行反向传播。
bfloat16()	将所有浮点参数和缓冲区转换为 bfloat16 数据类型。
buffers([recurse])	返回一个遍历模块缓冲区的迭代器。
children()	返回一个遍历直接子模块的迭代器。
clip_gradients(optimizer[, ...])	内部处理梯度裁剪。
compile(*args, **kwargs)	使用 torch.compile() 编译此模块的 forward。
configure_callbacks()	配置模型特定的回调。
configure_gradient_clipping(optimizer[, ...])	对优化器参数执行梯度裁剪。
configure_model()	在策略和精度感知的环境中创建模块的钩子。
configure_optimizers()	配置优化器和学习率调度器以优化模型。
configure_sharded_model()	已弃用。
configure_torch_metrics(torch_metrics)	处理 torch_metrics 参数。
cpu()	参见 torch.nn.Module.cpu()。
cuda([device])	将所有模型参数和缓冲区移动到GPU。
double()	参见 torch.nn.Module.double()。
eval()	将模块设置为评估模式。
extra_repr()	设置模块的额外表示。
float()	参见 torch.nn.Module.float()。
forward(*args, **kwargs)	与 torch.nn.Module.forward() 相同。
freeze()	冻结所有参数以进行推理。
get_buffer(target)	如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。
get_extra_state()	返回要包含在模块的 state_dict 中的任何额外状态。
get_parameter(target)	如果存在，则返回由 target 给出的参数，否则抛出错误。
get_submodule(target)	如果存在，返回由 target 指定的子模块，否则抛出错误。
half()	参见 torch.nn.Module.half()。
ipu([device])	将所有模型参数和缓冲区移动到 IPU。
load_from_checkpoint(checkpoint_path[, ...])	从检查点加载模型的主要方式。
load_state_dict(state_dict[, strict, assign])	将参数和缓冲区从 state_dict 复制到此模块及其后代中。
log(name, value[, prog_bar, logger, ...])	记录一个键值对。
log_dict(dictionary[, prog_bar, logger, ...])	一次性记录一个值的字典。
lr_scheduler_step(scheduler, metric)	重写此方法以调整 Trainer 调用每个调度器的默认方式。
lr_schedulers()	返回在训练期间使用的学习率调度器。
manual_backward(loss, *args, **kwargs)	在进行手动优化时，直接从 training_step() 调用此方法。
modules()	返回网络中所有模块的迭代器。
named_buffers([prefix, recurse, ...])	返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。
named_children()	返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。
named_modules([memo, prefix, remove_duplicate])	返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。
named_parameters([prefix, recurse, ...])	返回一个模块参数的迭代器，同时生成参数的名称和参数本身。
on_after_backward()	在 loss.backward() 之后和优化器步进之前调用。
on_after_batch_transfer(batch, dataloader_idx)	覆盖此方法以在批次转移到设备后对其进行批量增强或修改。
on_before_backward(loss)	在 loss.backward() 之前调用。
on_before_batch_transfer(batch, dataloader_idx)	在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。
on_before_optimizer_step(optimizer)	在 optimizer.step() 之前调用。
on_before_zero_grad(optimizer)	在 training_step() 之后和 optimizer.zero_grad() 之前调用。
on_fit_end()	在拟合的最后调用。
on_fit_start()	在拟合开始时调用。
on_load_checkpoint(checkpoint)	由 Lightning 调用以恢复您的模型。
on_predict_batch_end(outputs, batch, batch_idx)	在批次之后调用预测循环。
on_predict_batch_start(batch, batch_idx[, ...])	在预测循环中，在该批次发生任何事情之前调用。
on_predict_end()	在预测结束时调用。
on_predict_epoch_end()	在预测结束时调用。
on_predict_epoch_start()	在预测开始时调用。
on_predict_model_eval()	在预测循环开始时调用。
on_predict_start()	在预测开始时调用。
on_save_checkpoint(checkpoint)	当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。
on_test_batch_end(outputs, batch, batch_idx)	在批次之后在测试循环中调用。
on_test_batch_start(batch, batch_idx[, ...])	在测试循环中，在该批次发生任何事情之前调用。
on_test_end()	在测试结束时调用。
on_test_epoch_end()	在每个epoch的测试循环结束时调用。
on_test_epoch_start()	在epoch开始时在测试循环中调用。
on_test_model_eval()	在测试循环开始时调用。
on_test_model_train()	在测试循环结束时调用。
on_test_start()	在测试开始时调用。
on_train_batch_end(outputs, batch, batch_idx)	在训练循环中批次处理后调用。
on_train_batch_start(batch, batch_idx)	在训练循环中，在该批次发生任何事情之前调用。
on_train_end()	在训练结束前调用，在记录器实验关闭之前。
on_train_epoch_end()	在训练循环的每个epoch结束时调用。
on_train_epoch_start()	在训练循环中，在epoch开始时调用。
on_train_start()	在完整性检查后，训练开始时调用。
on_validation_batch_end(outputs, batch, ...)	在批次验证循环后调用。
on_validation_batch_start(batch, batch_idx)	在验证循环中，对该批次进行任何操作之前调用。
on_validation_end()	在验证结束时调用。
on_validation_epoch_end()	在epoch结束时的验证循环中调用。
on_validation_epoch_start()	在epoch开始时验证循环中调用。
on_validation_model_eval()	在验证循环开始时调用。
on_validation_model_train()	当验证循环结束时调用。
on_validation_model_zero_grad()	在进入验证循环之前，由训练循环调用来释放梯度。
on_validation_start()	在验证开始时调用。
optimizer_step(epoch, batch_idx, optimizer)	重写此方法以调整 Trainer 调用优化器的默认方式。
optimizer_zero_grad(epoch, batch_idx, optimizer)	重写此方法以改变 optimizer.zero_grad() 的默认行为。
optimizers([use_pl_optimizer])	返回训练期间使用的优化器。
parameters([recurse])	返回一个遍历模块参数的迭代器。
predict_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。
predict_step(batch, batch_idx[, dataloader_idx])	执行预测步骤
prepare_data()	使用此功能下载并准备数据。
print(*args, **kwargs)	仅从进程 0 打印。
register_backward_hook(hook)	在模块上注册一个反向钩子。
register_buffer(name, tensor[, persistent])	向模块添加一个缓冲区。
register_forward_hook(hook, *[, prepend, ...])	在模块上注册一个前向钩子。
register_forward_pre_hook(hook, *[, ...])	在模块上注册一个前向预钩子。
register_full_backward_hook(hook[, prepend])	在模块上注册一个反向钩子。
register_full_backward_pre_hook(hook[, prepend])	在模块上注册一个反向预钩子。
register_load_state_dict_post_hook(hook)	注册一个在模块的 load_state_dict 调用后运行的后置钩子。
register_module(name, module)	别名 add_module()。
register_parameter(name, param)	向模块添加一个参数。
register_state_dict_pre_hook(hook)	为 state_dict() 方法注册一个预钩子。
requires_grad_([requires_grad])	更改是否应在此模块的参数上记录自动梯度操作。
save_hyperparameters(*args[, ignore, frame, ...])	将参数保存到 hparams 属性中。
set_extra_state(state)	设置包含在加载的 state_dict 中的额外状态。
set_predict_parameters(n, num_samples, ...)	在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。
setup(stage)	在 fit（训练 + 验证）、验证、测试或预测开始时调用。
share_memory()	参见 torch.Tensor.share_memory_()。
state_dict(*args[, destination, prefix, ...])	返回一个包含模块整体状态引用的字典。
teardown(stage)	在 fit（训练 + 验证）、验证、测试或预测结束时调用。
test_dataloader()	一个可迭代对象或多个可迭代对象的集合，指定测试样本。
test_step(*args, **kwargs)	对测试集中的单个批次数据进行操作。
to(*args, **kwargs)	参见 torch.nn.Module.to()。
to_dtype(dtype)	将模块精度（默认为 float32）转换为另一种精度。
to_empty(*, device[, recurse])	将参数和缓冲区移动到指定设备，而不复制存储。
to_onnx(file_path[, input_sample])	将模型保存为 ONNX 格式。
to_torchscript([file_path, method, ...])	默认情况下，将整个模型编译为 ScriptModule。
toggle_optimizer(optimizer)	确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。
train([mode])	设置模块为训练模式。
train_dataloader()	指定训练样本的可迭代对象或可迭代对象的集合。
training_step(train_batch, batch_idx)	执行训练步骤
transfer_batch_to_device(batch, device, ...)	如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。
type(dst_type)	参见 torch.nn.Module.type()。
unfreeze()	解冻所有参数以进行训练。
untoggle_optimizer(optimizer)	重置通过 toggle_optimizer() 切换的所需梯度的状态。
val_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。
validation_step(val_batch, batch_idx)	执行验证步骤
xpu([device])	将所有模型参数和缓冲区移动到XPU。
zero_grad([set_to_none])	重置所有模型参数的梯度。

__call__
set_mc_dropout

CHECKPOINT_HYPER_PARAMS_KEY = 'hyper_parameters'

CHECKPOINT_HYPER_PARAMS_NAME = 'hparams_name'

CHECKPOINT_HYPER_PARAMS_TYPE = 'hparams_type'

T_destination

TypeVar(‘T_destination’, bound=Dict[str, Any]) 的别名

add_module(name, module)

将一个子模块添加到当前模块。

可以通过给定的名称将模块作为属性访问。

参数

• name (str) – 子模块的名称。可以使用给定的名称从此模块访问子模块。

• module (Module) – 要添加到模块中的子模块。

返回类型

None

all_gather(data, group=None, sync_grads=False)

从多个进程中收集张量或张量集合。

此方法需要在所有进程上调用，并且所有进程上的张量需要具有相同的形状，否则程序将永远停滞。

参数

• data (Union[Tensor, Dict, List, Tuple]) – int, float, 形状为 (batch, …) 的张量，或可能是嵌套的集合。

• group (Optional[Any, None]) – 要从中收集结果的进程组。默认为所有进程（world）

• sync_grads (bool) – 允许用户为 all_gather 操作同步梯度的标志

返回类型

Union[Tensor, Dict, List, Tuple]

返回

一个形状为 (world_size, batch, …) 的张量，或者如果输入是一个集合，输出也将是一个具有这种形状张量的集合。对于 world_size 为 1 的特殊情况，张量不会增加额外的维度。

allow_zero_length_dataloader_with_multiple_devices: bool

apply(fn)

递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。

典型用途包括初始化模型的参数（另见 nn-init-doc）。

参数

fn (Module -> None) – 要应用于每个子模块的函数

返回

自己

返回类型

Module

示例:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

property automatic_optimization: bool

如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。

返回类型

bool

backward(loss, *args, **kwargs)

在 training_step() 中返回的损失上调用以执行反向传播。如果需要，请使用您自己的实现覆盖此钩子。

参数

loss (Tensor) – 由 training_step() 返回的损失张量。如果使用梯度累积，这里的损失值是归一化的（按 1 / 累积步数缩放）。

示例:

def backward(self, loss):
    loss.backward()

返回类型

None

bfloat16()

将所有浮点参数和缓冲区转换为 bfloat16 数据类型。

备注

此方法就地修改模块。

返回

自己

返回类型

Module

buffers(recurse=True)

返回一个遍历模块缓冲区的迭代器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。

生成器

torch.Tensor – 模块缓冲区

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Tensor]

call_super_init: bool = False

children()

返回一个遍历直接子模块的迭代器。

生成器

模块 – 子模块

返回类型

Iterator[Module]

clip_gradients(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

内部处理梯度裁剪。

备注

• 不要重写此方法。如果你想自定义梯度裁剪，考虑使用 configure_gradient_clipping() 方法。

• 对于手动优化（self.automatic_optimization = False），如果你想使用梯度裁剪，可以在训练步骤中手动调用 self.clip_gradients(opt, gradient_clip_val=0.5, gradient_clip_algorithm="norm")。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。

• gradient_clip_algorithm (Optional[str, None]) – 使用的梯度裁剪算法。传递 gradient_clip_algorithm="value" 按值裁剪，传递 gradient_clip_algorithm="norm" 按范数裁剪。

返回类型

None

compile(*args, **kwargs)

使用 torch.compile() 编译此模块的 forward。

此模块的 __call__ 方法被编译，所有参数都按原样传递给 torch.compile()。

有关此函数的参数详情，请参阅 torch.compile()。

configure_callbacks()

配置模型特定的回调。当模型被附加时，例如当调用 .fit() 或 .test() 时，这里返回的回调列表或回调将与传递给 Trainer 的 callbacks 参数的回调列表合并。如果这里返回的回调与 Trainer 的回调列表中已有的一个或多个回调类型相同，它将优先并替换它们。此外，Lightning 将确保 ModelCheckpoint 回调最后运行。

返回类型

Union[Sequence[Callback], Callback]

返回

一个回调函数或回调函数列表，这些回调函数将扩展 Trainer 中的回调函数列表。

示例:

def configure_callbacks(self):
    early_stop = EarlyStopping(monitor="val_acc", mode="max")
    checkpoint = ModelCheckpoint(monitor="val_loss")
    return [early_stop, checkpoint]

configure_gradient_clipping(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

对优化器参数执行梯度裁剪。在调用 optimizer_step() 之前执行。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。默认情况下，传递给 Trainer 的值将在此处可用。

• gradient_clip_algorithm (Optional[str, None]) – 要使用的梯度裁剪算法。默认情况下，传递给 Trainer 的值将在此处可用。

示例:

def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm):
    # Implement your own custom logic to clip gradients
    # You can call `self.clip_gradients` with your settings:
    self.clip_gradients(
        optimizer,
        gradient_clip_val=gradient_clip_val,
        gradient_clip_algorithm=gradient_clip_algorithm
    )

返回类型

None

configure_model()

在策略和精度感知的环境中创建模块的钩子。

这在使用分片策略（FSDP 和 DeepSpeed）时特别有用，我们希望立即分片模型以节省内存和初始化时间。对于非分片策略，您可以选择覆盖此钩子或在 init_module() 上下文管理器下初始化您的模型。

此钩子在每个 fit/val/test/predict 阶段都会被调用，因此请确保此钩子的实现是 幂等的，即，在第一次调用钩子后，后续的调用应该是无操作的。

返回类型

None

configure_optimizers()

配置优化器和学习率调度器以优化模型。

configure_sharded_model()

已弃用。

使用 configure_model() 代替。

返回类型

None

static configure_torch_metrics(torch_metrics)

处理 torch_metrics 参数。

返回类型

MetricCollection

cpu()

参见 torch.nn.Module.cpu()。

返回类型

Self

cuda(device=None)

将所有模型参数和缓冲区移动到GPU。这也会使相关的参数和缓冲区成为不同的对象。因此，如果模块在优化期间将驻留在GPU上，则应在构造优化器之前调用此方法。

参数

device (Union[device, int, None]) – 如果指定，所有参数将被复制到该设备。如果为 None，将使用当前的 CUDA 设备索引。

返回

自己

返回类型

Module

property current_epoch: int

Trainer 中的当前时期，如果未附加则为 0。

返回类型

int

property device: device

返回类型

device

property device_mesh: Optional[DeviceMesh]

像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。

返回类型

Optional[ForwardRef, None]

double()

参见 torch.nn.Module.double()。

返回类型

Self

property dtype: Union[str, dtype]

返回类型

Union[str, dtype]

dump_patches: bool = False

property epochs_trained

eval()

将模块设置为评估模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

这等同于 self.train(False)。

请参阅 locally-disable-grad-doc 以比较 .eval() 和几个可能与之混淆的类似机制。

返回

自己

返回类型

Module

property example_input_array: Optional[Union[Tensor, Tuple, Dict]]

示例输入数组是模块在 forward() 方法中可以使用的规范。返回类型解释如下：

• 单个张量：假设模型接受一个参数，即 model.forward(model.example_input_array)

• 元组：输入数组应被解释为位置参数的序列，即 model.forward(*model.example_input_array)

• 字典：输入数组表示命名关键字参数，即 model.forward(**model.example_input_array)

返回类型

Union[Tensor, Tuple, Dict, None]

extra_repr()

设置模块的额外表示。

要打印自定义的额外信息，您应该在您自己的模块中重新实现此方法。单行和多行字符串都是可接受的。

返回类型

str

property fabric: Optional[Fabric]

返回类型

Optional[Fabric, None]

float()

参见 torch.nn.Module.float()。

返回类型

Self

abstract forward(*args, **kwargs)

与 torch.nn.Module.forward() 相同。

参数

• *args – 无论你决定传入 forward 方法的内容。

• **kwargs – 关键字参数也是可能的。

返回类型

Any

返回

你的模型的输出

freeze()

冻结所有参数以进行推理。

示例:

model = MyLightningModule(...)
model.freeze()

返回类型

None

get_buffer(target)

如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的缓冲区的完全限定字符串名称。（有关如何指定完全限定字符串，请参见 get_submodule。）

返回

target 引用的缓冲区

返回类型

torch.Tensor

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是缓冲区的内容

get_extra_state()

返回要包含在模块的 state_dict 中的任何额外状态。

如果需要存储额外状态，请在你的模块中实现此功能以及相应的 set_extra_state()。此函数在构建模块的 state_dict() 时被调用。

请注意，额外的状态应该是可pickle的，以确保state_dict的序列化工作。我们仅对序列化张量提供向后兼容性保证；如果其他对象的序列化pickle形式发生变化，它们可能会破坏向后兼容性。

返回

模块的 state_dict 中存储的任何额外状态

返回类型

object

get_parameter(target)

如果存在，则返回由 target 给出的参数，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的参数的全限定字符串名称。（关于如何指定全限定字符串，请参见 get_submodule。）

返回

target 引用的参数

返回类型

torch.nn.Parameter

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Parameter 的内容

get_submodule(target)

如果存在，返回由 target 指定的子模块，否则抛出错误。

例如，假设你有一个 nn.Module A ，它看起来像这样：

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(该图展示了一个 nn.Module A。A 有一个嵌套的子模块 net_b，它本身有两个子模块 net_c 和 linear。net_c 然后有一个子模块 conv。)

要检查我们是否拥有 linear 子模块，我们会调用 get_submodule("net_b.linear")。要检查我们是否拥有 conv 子模块，我们会调用 get_submodule("net_b.net_c.conv")。

get_submodule 的运行时间受限于 target 中的模块嵌套程度。对 named_modules 的查询可以实现相同的结果，但它是 O(N) 的，其中 N 是传递模块的数量。因此，对于检查某个子模块是否存在的简单检查，应始终使用 get_submodule。

参数

target (str) – 要查找的子模块的完全限定字符串名称。（请参见上面的示例，了解如何指定完全限定的字符串。）

返回

由 target 引用的子模块

返回类型

torch.nn.Module

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Module 的内容

property global_rank: int

当前进程在所有节点和设备中的索引。

返回类型

int

property global_step: int

所有epoch中看到的总训练批次。

如果没有连接训练器，此属性为 0。

返回类型

int

half()

参见 torch.nn.Module.half()。

返回类型

Self

property hparams: Union[AttributeDict, MutableMapping]

通过 save_hyperparameters() 保存的超参数集合。用户可以对其进行修改。对于初始超参数的冻结集，请使用 hparams_initial。

返回类型

Union[AttributeDict, MutableMapping]

返回

可变超参数字典

property hparams_initial: AttributeDict

通过 save_hyperparameters() 保存的超参数集合。这些内容是只读的。可以通过 hparams 手动更新保存的超参数。

返回

不可变的初始超参数

返回类型

AttributeDict

ipu(device=None)

将所有模型参数和缓冲区移动到 IPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在IPU上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs)

从检查点加载模型的主要方式。当Lightning保存检查点时，它会将传递给 __init__ 的参数存储在检查点中的 "hyper_parameters" 下。

通过 **kwargs 指定的任何参数将覆盖存储在 "hyper_parameters" 中的参数。

参数

• checkpoint_path (Union[str, Path, IO]) – 检查点的路径。这也可以是一个URL，或类似文件的对象。

• map_location (Union[device, str, int, Callable[[UntypedStorage, str], Optional[UntypedStorage, None]], Dict[Union[device, str, int], Union[device, str, int]], None]) – 如果你的检查点保存了一个GPU模型，而你现在在CPU上或不同数量的GPU上加载，使用这个来映射到新的设置。其行为与 torch.load() 相同。

• hparams_file (Union[str, Path, None]) –

  可选路径到一个 .yaml 或 .csv 文件，其层次结构如本示例所示:

  drop_prob: 0.2
  dataloader:
      batch_size: 32
  

  您很可能不需要这个，因为 Lightning 总是会将超参数保存到检查点。然而，如果您的检查点权重没有保存超参数，请使用此方法传入一个包含您希望使用的超参数的 .yaml 文件。这些将被转换为 dict 并传递到您的 LightningModule 中使用。

  如果你的模型的 hparams 参数是 Namespace 并且 .yaml 文件具有层次结构，你需要重构你的模型以将 hparams 视为 dict。

• strict (Optional[bool, None]) – 是否严格要求 checkpoint_path 中的键与该模块状态字典返回的键匹配。默认为 True ，除非设置了 LightningModule.strict_loading ，在这种情况下，它默认为 LightningModule.strict_loading 的值。

• **kwargs – 初始化模型所需的任何额外关键字参数。也可以用来覆盖已保存的超参数值。

返回类型

Self

返回

LightningModule 实例，加载了权重和超参数（如果可用）。

备注

load_from_checkpoint 是一个 类 方法。你应该使用你的 LightningModule 类 来调用它，而不是 LightningModule 实例，否则会引发 TypeError。

备注

为了确保所有层都能从检查点加载，如果您的 LightningModule 中重写了此钩子，此函数将在实例化模型后直接调用 configure_model()。但是，请注意 load_from_checkpoint 不支持加载分片检查点，如果模型太大，您可能会耗尽内存。在这种情况下，请考虑通过 .fit(ckpt_path=...) 通过 Trainer 加载。

示例:

# load weights without mapping ...
model = MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')

# or load weights mapping all weights from GPU 1 to GPU 0 ...
map_location = {'cuda:1':'cuda:0'}
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    map_location=map_location
)

# or load weights and hyperparameters from separate files.
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    hparams_file='/path/to/hparams_file.yaml'
)

# override some of the params with new values
model = MyLightningModule.load_from_checkpoint(
    PATH,
    num_layers=128,
    pretrained_ckpt_path=NEW_PATH,
)

# predict
pretrained_model.eval()
pretrained_model.freeze()
y_hat = pretrained_model(x)

load_state_dict(state_dict, strict=True, assign=False)

将参数和缓冲区从 state_dict 复制到此模块及其后代中。

如果 strict 是 True，那么 state_dict 的键必须与该模块的 state_dict() 函数返回的键完全匹配。

警告

如果 assign 是 True，优化器必须在调用 load_state_dict 之后创建，除非 get_swap_module_params_on_conversion() 是 True。

参数

• state_dict (dict) – 包含参数和持久缓冲区的字典。

• strict (bool, optional) – 是否严格要求 state_dict 中的键与该模块的 state_dict() 函数返回的键匹配。默认值：True

• assign (bool, optional) – 当 False 时，当前模块中张量的属性会被保留，而当 True 时，状态字典中张量的属性会被保留。唯一的例外是 Parameter 的 requires_grad 字段，该字段的值会从模块中保留。默认值：False

返回

• missing_keys 是一个包含任何预期键的 str 列表

  此模块中存在但未在提供的 state_dict 中出现的。

• unexpected_keys 是一个包含未预料到的键的 str 列表。

  此模块预期但未在提供的 state_dict 中出现。

返回类型

带有 missing_keys 和 unexpected_keys 字段的 NamedTuple

备注

如果一个参数或缓冲区被注册为 None ，并且其对应的键存在于 state_dict 中，load_state_dict() 将引发一个 RuntimeError 。

property local_rank: int

当前进程在单个节点中的索引。

返回类型

int

log(name, value, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, metric_attribute=None, rank_zero_only=False)

记录一个键值对。

示例:

self.log('train_loss', loss)

每个钩子的默认行为在此处有文档说明：extensions/logging:自动日志记录。

参数

• name (str) – 日志的键。如果使用DDP或其他分布式策略，必须在所有进程中保持一致。

• value (Union[Metric, Tensor, int, float]) – 要记录的值。可以是 float、Tensor 或 Metric。

• prog_bar (bool) – 如果 True 则记录到进度条。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• on_epoch (Optional[bool, None]) – 如果 True ，记录每个epoch累积的指标。默认值由钩子决定。详情请参见 extensions/logging:自动日志记录 。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图表。

• sync_dist (bool) – 如果 True，将在设备间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的DDP组。

• add_dataloader_idx (bool) – 如果 True，将当前数据加载器的索引附加到名称（当使用多个数据加载器时）。如果为 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的 batch_size。这将直接从加载的批次中推断出来，但对于某些数据结构，您可能需要显式提供它。

• metric_attribute (Optional[str, None]) – 要恢复指标状态，Lightning 需要模型中 torchmetrics.Metric 的引用。如果它是模型属性，则会自动找到。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

log_dict(dictionary, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, rank_zero_only=False)

一次性记录一个值的字典。

示例:

values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n}
self.log_dict(values)

参数

• dictionary (Union[Mapping[str, Union[Metric, Tensor, int, float]], MetricCollection]) – 键值对。如果使用DDP或其他分布式策略，键必须在所有进程中保持一致。值可以是 float 、 Tensor 、 Metric 或 MetricCollection 。

• prog_bar (bool) – 如果 True 则记录到进度基础。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。None 自动记录训练步骤，但不记录验证/测试步骤。默认值由钩子决定。详情请参阅 extensions/logging:自动记录。

• on_epoch (Optional[bool, None]) – 如果 True 记录累积的 epoch 指标。None 自动记录 val/test 步骤的日志，但不包括 training_step。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图形

• sync_dist (bool) – 如果 True，将在GPU/TPU之间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的 ddp 组。

• add_dataloader_idx (bool) – 如果 True，则将当前数据加载器的索引附加到名称（在使用多个数据加载器时）。如果 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的批次大小。这将直接从加载的批次中推断出来，但某些数据结构可能需要显式提供它。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

property logger: Optional[Union[Logger, Logger]]

Trainer 中对 logger 对象的引用。

返回类型

Union[Logger, Logger, None]

property loggers: Union[List[Logger], List[Logger]]

参考 Trainer 中的日志记录器列表。

返回类型

Union[List[Logger], List[Logger]]

lr_scheduler_step(scheduler, metric)

重写此方法以调整 Trainer 调用每个调度器的默认方式。默认情况下，Lightning 会根据每个调度器的 interval 调用 step() 方法，如示例所示。

参数

• scheduler (Union[LRScheduler, ReduceLROnPlateau]) – 学习率调度器。

• metric (Optional[Any, None]) – 用于 ReduceLROnPlateau 等调度器的监视器值。

示例:

# DEFAULT
def lr_scheduler_step(self, scheduler, metric):
    if metric is None:
        scheduler.step()
    else:
        scheduler.step(metric)

# Alternative way to update schedulers if it requires an epoch value
def lr_scheduler_step(self, scheduler, metric):
    scheduler.step(epoch=self.current_epoch)

返回类型

None

lr_schedulers()

返回在训练期间使用的学习率调度器。对于手动优化很有用。

返回类型

Union[None, List[Union[LRScheduler, ReduceLROnPlateau]], LRScheduler, ReduceLROnPlateau]

返回

一个调度器，或者在存在多个调度器的情况下是一个调度器列表，如果在 configure_optimizers() 中没有返回调度器，则为 None。

manual_backward(loss, *args, **kwargs)

当你手动进行优化时，请直接从你的 training_step() 调用此方法。通过使用此方法，Lightning 可以确保在使用混合精度时应用所有适当的缩放。

更多示例请参见 手动优化。

示例:

def training_step(...):
    opt = self.optimizers()
    loss = ...
    opt.zero_grad()
    # automatically applies scaling, etc...
    self.manual_backward(loss)
    opt.step()

参数

• loss (Tensor) – 要计算梯度的张量。必须附加一个图。

• *args – 要转发给 backward() 的附加位置参数

• **kwargs – 要转发给 backward() 的额外关键字参数

返回类型

None

modules()

返回网络中所有模块的迭代器。

生成器

模块 – 网络中的一个模块

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

返回类型

Iterator[Module]

named_buffers(prefix='', recurse=True, remove_duplicate=True)

返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。

参数

• prefix (str) – 要添加到所有缓冲区名称的前缀。

• recurse (bool, optional) – 如果为 True，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。默认为 True。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的缓冲区。默认为 True。

生成器

(str, torch.Tensor) – 包含名称和缓冲区的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

返回类型

Iterator[Tuple[str, Tensor]]

named_children()

返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。

生成器

(str, Module) – 包含名称和子模块的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

返回类型

Iterator[Tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)

返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。

参数

• memo (Optional[Set[Module], None]) – 一个备忘录，用于存储已添加到结果中的模块集合

• prefix (str) – 一个将添加到模块名称前的前缀

• remove_duplicate (bool) – 是否在结果中移除重复的模块实例

生成器

(str, Module) – 名称和模块的元组

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)

返回一个模块参数的迭代器，同时生成参数的名称和参数本身。

参数

• prefix (str) – 要添加到所有参数名称前的前缀。

• recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的参数。默认为 True。

生成器

(str, Parameter) – 包含名称和参数的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

返回类型

Iterator[Tuple[str, Parameter]]

on_after_backward()

在 loss.backward() 之后和优化器步进之前调用。

备注

如果使用原生 AMP，梯度在此刻将不会被缩放。如果你需要未缩放的梯度，请使用 on_before_optimizer_step。

返回类型

None

on_after_batch_transfer(batch, dataloader_idx)

覆盖此方法以在批次转移到设备后对其进行批量增强或修改。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_after_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = gpu_transforms(batch['x'])
    return batch

参见

• on_before_batch_transfer()

• transfer_batch_to_device()

on_before_backward(loss)

在 loss.backward() 之前调用。

参数

loss (Tensor) – 损失除以梯度累积的批次数量，如果使用AMP则进行缩放。

返回类型

None

on_before_batch_transfer(batch, dataloader_idx)

在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_before_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = transforms(batch['x'])
    return batch

参见

• on_after_batch_transfer()

• transfer_batch_to_device()

on_before_optimizer_step(optimizer)

在 optimizer.step() 之前调用。

如果使用梯度累积，钩子会在梯度累积完成后被调用。参见：:paramref:`~pytorch_lightning.trainer.trainer.Trainer.accumulate_grad_batches`。

如果使用AMP，损失将在调用此钩子之前未缩放。有关梯度缩放的更多信息，请参阅这些 文档。

如果裁剪梯度，梯度尚未被裁剪。

参数

optimizer (Optimizer) – 当前使用的优化器。

示例:

def on_before_optimizer_step(self, optimizer):
    # example to inspect gradient information in tensorboard
    if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
        for k, v in self.named_parameters():
            self.logger.experiment.add_histogram(
                tag=k, values=v.grad, global_step=self.trainer.global_step
            )

返回类型

None

on_before_zero_grad(optimizer)

在 training_step() 之后和 optimizer.zero_grad() 之前调用。

在训练循环中，调用优化器步骤之后和清零梯度之前。这是一个检查更新后权重信息的好地方。

这是调用它的地方:

for optimizer in optimizers:
    out = training_step(...)

    model.on_before_zero_grad(optimizer) # < ---- called here
    optimizer.zero_grad()

    backward()

参数

optimizer (Optimizer) – 应该被置零的梯度优化器。

返回类型

None

on_fit_end()

在拟合的最后调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

on_fit_start()

在拟合开始时调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

property on_gpu: bool

如果该模型当前位于GPU上，则返回 True。

在LightningModule周围设置标志以区分CPU和GPU行为非常有用。

返回类型

bool

on_load_checkpoint(checkpoint)

由 Lightning 调用以恢复您的模型。如果您使用 on_save_checkpoint() 保存了某些内容，这是您恢复这些内容的机会。

参数

checkpoint (Dict[str, Any]) – 加载的检查点

示例:

def on_load_checkpoint(self, checkpoint):
    # 99% of the time you don't need to implement this method
    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']

备注

Lightning 自动恢复全局步数、轮次和训练状态，包括 amp 缩放。您无需恢复任何与训练相关的内容。

返回类型

None

on_predict_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后调用预测循环。

参数

• outputs (Optional[Any, None]) – predict_step(x) 的输出

• batch (Any) – 由预测 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_batch_start(batch, batch_idx, dataloader_idx=0)

在预测循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_start()

在预测开始时调用。

返回类型

None

on_predict_model_eval()

在预测循环开始时调用。

默认情况下，预测循环在开始之前会在 LightningModule 上调用 .eval()。重写此钩子以更改行为。

返回类型

None

on_predict_start()

在预测开始时调用。

返回类型

None

on_save_checkpoint(checkpoint)

当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。

参数

checkpoint (Dict[str, Any]) – 在转储到文件之前的完整检查点字典。此钩子的实现可以将额外数据插入此字典中。

示例:

def on_save_checkpoint(self, checkpoint):
    # 99% of use cases you don't need to implement this method
    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object

备注

Lightning 保存了训练的所有方面（如 epoch、全局步数等），包括 amp 缩放。您无需存储任何关于训练的信息。

返回类型

None

on_test_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后在测试循环中调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – test_step(x) 的输出

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_batch_start(batch, batch_idx, dataloader_idx=0)

在测试循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_end()

在测试结束时调用。

返回类型

None

on_test_epoch_end()

在每个epoch的测试循环结束时调用。

返回类型

None

on_test_epoch_start()

在epoch开始时在测试循环中调用。

返回类型

None

on_test_model_eval()

在测试循环开始时调用。

默认情况下，测试循环在开始前会调用 .eval() 方法在 LightningModule 上。重写此钩子以改变行为。另请参见 on_test_model_train()。

返回类型

None

on_test_model_train()

在测试循环结束时调用。

测试循环默认会将 LightningModule 的 training 模式恢复到开始测试之前的状态。覆盖此钩子以更改行为。另请参见 on_test_model_eval()。

返回类型

None

on_test_start()

在测试开始时调用。

返回类型

None

on_train_batch_end(outputs, batch, batch_idx)

在训练循环中批次处理后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – 训练步骤(x)的输出

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

备注

这里的 outputs["loss"] 值将是相对于 accumulate_grad_batches 的损失归一化值，该损失从 training_step 返回。

返回类型

None

on_train_batch_start(batch, batch_idx)

在训练循环中，在该批次发生任何事情之前调用。

如果你在这里返回 -1，你将跳过当前轮次的剩余训练。

参数

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

返回类型

Optional[int, None]

on_train_end()

在训练结束前调用，在记录器实验关闭之前。

返回类型

None

on_train_epoch_end()

在训练循环的每个epoch结束时调用。

要在每个epoch结束时访问所有批处理输出，您可以将步骤输出缓存为 LightningModule 的一个属性，并在此钩子中访问它们：

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # free up the memory
        self.training_step_outputs.clear()

on_train_epoch_start()

在训练循环中，在epoch开始时调用。

返回类型

None

on_train_start()

在完整性检查后，训练开始时调用。

返回类型

None

on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次验证循环后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – validation_step(x) 的输出

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_batch_start(batch, batch_idx, dataloader_idx=0)

在验证循环中，对该批次进行任何操作之前调用。

参数

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_end()

在验证结束时调用。

返回类型

None

on_validation_epoch_end()

在epoch结束时的验证循环中调用。

on_validation_epoch_start()

在epoch开始时验证循环中调用。

返回类型

None

on_validation_model_eval()

在验证循环开始时调用。

默认情况下，验证循环在开始前会在 LightningModule 上调用 .eval()。覆盖此钩子以更改行为。另请参见 on_validation_model_train()。

返回类型

None

on_validation_model_train()

当验证循环结束时调用。

默认情况下，验证循环会将 LightningModule 的 training 模式恢复到开始验证之前的状态。重写此钩子以更改行为。另请参见 on_validation_model_eval()。

返回类型

None

on_validation_model_zero_grad()

在进入验证循环之前，由训练循环调用来释放梯度。

返回类型

None

on_validation_start()

在验证开始时调用。

返回类型

None

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)

重写此方法以调整 Trainer 调用优化器的默认方式。

默认情况下，Lightning 会调用 step() 和 zero_grad()，如示例所示。当 Trainer(accumulate_grad_batches != 1) 时，在累积阶段不会调用此方法（和 zero_grad()）。手动优化时，重写此钩子没有好处。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Union[Optimizer, LightningOptimizer]) – 一个 PyTorch 优化器

• optimizer_closure (Optional[Callable[[], Any], None]) – 优化器闭包。此闭包必须执行，因为它包含了 training_step()、optimizer.zero_grad() 和 backward() 的调用。

示例:

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    # Add your custom logic to run directly before `optimizer.step()`

    optimizer.step(closure=optimizer_closure)

    # Add your custom logic to run directly after `optimizer.step()`

返回类型

None

optimizer_zero_grad(epoch, batch_idx, optimizer)

重写此方法以改变 optimizer.zero_grad() 的默认行为。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Optimizer) – 一个 PyTorch 优化器

示例:

# DEFAULT
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad()

# Set gradients to `None` instead of zero to improve performance (not required on `torch>=2.0.0`).
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad(set_to_none=True)

参见 torch.optim.Optimizer.zero_grad() 以了解上述示例的解释。

返回类型

None

optimizers(use_pl_optimizer=True)

返回训练期间使用的优化器。对于手动优化很有用。

参数

use_pl_optimizer (bool) – 如果 True，将会把优化器包装在 LightningOptimizer 中，以自动处理精度、分析和步数调用的计数，以进行适当的日志记录和检查点保存。它特别包装了 step 方法，不支持没有此方法的自定义优化器。

返回类型

Union[Optimizer, LightningOptimizer, _FabricOptimizer, List[Optimizer], List[LightningOptimizer], List[_FabricOptimizer]]

返回

一个单一的优化器，或者在存在多个优化器的情况下，一个优化器列表。

property output_chunk_length: Optional[int]

模型一次预测的时间步数。

返回类型

Optional[int, None]

parameters(recurse=True)

返回一个遍历模块参数的迭代器。

这通常传递给一个优化器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

生成器

参数 – 模块参数

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Parameter]

pred_batch_size: Optional[int]

pred_mc_dropout: Optional[bool]

pred_n: Optional[int]

pred_n_jobs: Optional[int]

pred_num_samples: Optional[int]

pred_roll_size: Optional[int]

predict_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。

有关多个数据加载器的更多信息，请参阅此 部分。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• predict()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

返回

一个 torch.utils.data.DataLoader 或它们的序列，指定预测样本。

predict_likelihood_parameters: Optional[bool]

predict_step(batch, batch_idx, dataloader_idx=None)

执行预测步骤

批处理

Darts 的 InferenceDataset 的输出 - 包含 (过去的目标, 过去的协变量, 历史未来协变量, 未来协变量, 未来过去协变量, 输入时间序列, 预测开始时间步) 的元组

batch_idx

当前批次的批次索引

dataloader_idx

数据加载器索引

返回类型

Sequence[TimeSeries]

prepare_data()

使用此方法下载和准备数据。使用多个进程（分布式设置）下载和保存数据将导致数据损坏。Lightning 确保此方法仅在一个进程内调用，因此您可以安全地在其中添加下载逻辑。

警告

不要将状态设置到模型中（使用 setup 代替），因为这不会在每个设备上被调用

示例:

def prepare_data(self):
    # good
    download_data()
    tokenize()
    etc()

    # bad
    self.split = data_split
    self.some_state = some_other_state()

在分布式环境中，prepare_data 可以通过两种方式调用（使用 prepare_data_per_node）

1. 每个节点一次。这是默认设置，仅在 LOCAL_RANK=0 时调用。

2. 总共一次。仅在 GLOBAL_RANK=0 时调用。

示例:

# DEFAULT
# called once per node on LOCAL_RANK=0 of that node
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = True

# call on GLOBAL_RANK=0 (great for shared file systems)
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = False

在请求数据加载器之前调用此方法：

model.prepare_data()
initialize_distributed()
model.setup(stage)
model.train_dataloader()
model.val_dataloader()
model.test_dataloader()
model.predict_dataloader()

返回类型

None

prepare_data_per_node: bool

print(*args, **kwargs)

仅从进程 0 打印。在任何分布式模式下使用此功能以仅记录一次。

参数

• *args – 要打印的内容。与Python内置的print函数相同。

• **kwargs – 与Python内置的print函数相同。

示例:

def forward(self, x):
    self.print(x, 'in forward')

返回类型

None

register_backward_hook(hook)

在模块上注册一个反向钩子。

此函数已被 register_full_backward_hook() 取代，并且此函数的行为将在未来版本中发生变化。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_buffer(name, tensor, persistent=True)

向模块添加一个缓冲区。

这通常用于注册一个不应被视为模型参数的缓冲区。例如，BatchNorm 的 running_mean 不是参数，而是模块状态的一部分。缓冲区默认是持久的，并将与参数一起保存。可以通过将 persistent 设置为 False 来改变此行为。持久缓冲区和非持久缓冲区之间的唯一区别是后者不会成为此模块的 state_dict 的一部分。

缓冲区可以通过给定的名称作为属性访问。

参数

• name (str) – 缓冲区的名称。可以通过给定的名称从此模块访问缓冲区。

• tensor (Tensor or None) – 要注册的缓冲区。如果为 None ，则运行在缓冲区上的操作，例如 cuda ，将被忽略。如果为 None ，该缓冲区将 不 包含在模块的 state_dict 中。

• persistent (bool) – 缓冲区是否是此模块的 state_dict 的一部分。

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

返回类型

None

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

在模块上注册一个前向钩子。

钩子将在每次 forward() 计算输出后被调用。

如果 with_kwargs 是 False 或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward 。钩子可以修改输出。它可以就地修改输入，但这不会影响前向传播，因为这是在调用 forward() 之后调用的。钩子应具有以下签名:

hook(module, args, output) -> None or modified output

如果 with_kwargs 是 True，前向钩子将被传递给前向函数的 kwargs，并期望返回可能被修改的输出。钩子应具有以下签名:

hook(module, args, kwargs, output) -> None or modified output

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果 True，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之后触发。请注意，使用 register_module_forward_hook() 注册的全局 forward 钩子将在通过此方法注册的所有钩子之前触发。默认值：False

• with_kwargs (bool) – 如果 True，hook 将会传递给 forward 函数的 kwargs。默认值：False

• always_call (bool) – 如果 True ，无论在调用模块时是否引发异常，都将运行 hook 。默认值： False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

在模块上注册一个前向预钩子。

钩子将在每次调用 forward() 之前被调用。

如果 with_kwargs 为假或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward。钩子可以修改输入。用户可以在钩子中返回一个元组或单个修改后的值。如果返回单个值（除非该值已经是元组），我们将把该值包装成一个元组。钩子应具有以下签名:

hook(module, args) -> None or modified input

如果 with_kwargs 为真，前向预钩子将传递给前向函数的 kwargs。如果钩子修改了输入，则应返回 args 和 kwargs。钩子应具有以下签名:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之后触发。请注意，通过 register_module_forward_pre_hook() 注册的全局 forward_pre 钩子将在所有通过此方法注册的钩子之前触发。默认值：False

• with_kwargs (bool) – 如果为真，hook 将被传递给 forward 函数的 kwargs。默认值：False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)

在模块上注册一个反向钩子。

每次计算与模块相关的梯度时，钩子都会被调用，即当且仅当计算与模块输出相关的梯度时，钩子才会执行。钩子应具有以下签名:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input 和 grad_output 是包含相对于输入和输出的梯度的元组。钩子不应修改其参数，但可以选择返回一个新的相对于输入的梯度，该梯度将用于替代 grad_input 在后续计算中。grad_input 将仅对应于作为位置参数给出的输入，所有 kwarg 参数都将被忽略。grad_input 和 grad_output 中的条目将为 None 用于所有非 Tensor 参数。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入或输出，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之后触发。请注意，通过 register_module_full_backward_hook() 注册的全局 backward 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

在模块上注册一个反向预钩子。

每次计算模块的梯度时，都会调用该钩子。钩子应具有以下签名:

hook(module, grad_output) -> tuple[Tensor] or None

The grad_output 是一个元组。钩子不应修改其参数，但它可以选择返回一个新的相对于输出的梯度，该梯度将用于替代 grad_output 在后续计算中。grad_output 中的条目对于所有非 Tensor 参数将为 None。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之后触发。请注意，通过 register_module_full_backward_pre_hook() 注册的全局 backward_pre 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)

注册一个在模块的 load_state_dict 调用后运行的后置钩子。

它应该有以下签名:

hook(module, incompatible_keys) -> None

module 参数是当前注册此钩子的模块，而 incompatible_keys 参数是一个由属性 missing_keys 和 unexpected_keys 组成的 NamedTuple。missing_keys 是一个包含缺失键的 str 的 list，而 unexpected_keys 是一个包含意外键的 str 的 list。

如果需要，可以就地修改给定的 incompatible_keys。

请注意，当使用 strict=True 调用 load_state_dict() 时执行的检查会受到钩子对 missing_keys 或 unexpected_keys 所做修改的影响，这是预期的。对任一组键的添加将导致在 strict=True 时抛出错误，而清除掉缺失和意外的键将避免错误。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_module(name, module)

别名 add_module()。

返回类型

None

register_parameter(name, param)

向模块添加一个参数。

该参数可以通过给定的名称作为属性访问。

参数

• name (str) – 参数的名称。可以使用给定的名称从此模块访问该参数。

• param (Parameter or None) – 要添加到模块的参数。如果为 None ，则运行在参数上的操作，例如 cuda ，将被忽略。如果为 None ，该参数将 不 包含在模块的 state_dict 中。

返回类型

None

register_state_dict_pre_hook(hook)

为 state_dict() 方法注册一个预钩子。

这些钩子将在调用 self 上的 state_dict 之前，使用参数 self、prefix 和 keep_vars 进行调用。注册的钩子可以用于在调用 state_dict 之前执行预处理。

requires_grad_(requires_grad=True)

更改是否应在此模块的参数上记录自动梯度操作。

此方法就地设置参数的 requires_grad 属性。

这种方法有助于冻结模块的一部分以进行微调或单独训练模型的一部分（例如，GAN训练）。

请参阅 locally-disable-grad-doc 以比较 .requires_grad_() 和几个可能与之混淆的类似机制。

参数

requires_grad (bool) – 是否应在此模块的参数上记录自动求导操作。默认值：True。

返回

自己

返回类型

Module

save_hyperparameters(*args, ignore=None, frame=None, logger=True)

将参数保存到 hparams 属性中。

参数

• args (Any) – dict、NameSpace 或 OmegaConf 的单个对象，或来自类 __init__ 的字符串名称或参数

• ignore (Union[Sequence[str], str, None]) – 一个参数名称或从类 __init__ 中忽略的参数名称列表

• frame (Optional[frame, None]) – 一个框架对象。默认是 None

• logger (bool) – 是否将超参数发送到记录器。默认值：True

示例:

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # manually assign arguments
...         self.save_hyperparameters('arg1', 'arg3')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class AutomaticArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # equivalent automatic
...         self.save_hyperparameters()
...     def forward(self, *args, **kwargs):
...         ...
>>> model = AutomaticArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg2": abc
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class SingleArgModel(HyperparametersMixin):
...     def __init__(self, params):
...         super().__init__()
...         # manually assign single argument
...         self.save_hyperparameters(params)
...     def forward(self, *args, **kwargs):
...         ...
>>> model = SingleArgModel(Namespace(p1=1, p2='abc', p3=3.14))
>>> model.hparams
"p1": 1
"p2": abc
"p3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # pass argument(s) to ignore as a string or in a list
...         self.save_hyperparameters(ignore='arg2')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

返回类型

None

set_extra_state(state)

设置包含在加载的 state_dict 中的额外状态。

此函数从 load_state_dict() 调用，以处理在 state_dict 中找到的任何额外状态。如果你的模块需要在 state_dict 中存储额外状态，请实现此函数以及相应的 get_extra_state()。

参数

state (dict) – 来自 state_dict 的额外状态

返回类型

None

set_mc_dropout(active)

set_predict_parameters(n, num_samples, roll_size, batch_size, n_jobs, predict_likelihood_parameters, mc_dropout)

在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。

返回类型

None

setup(stage)

在 fit（训练 + 验证）、验证、测试或预测开始时调用。当你需要动态构建模型或调整模型时，这是一个很好的钩子。在使用 DDP 时，此钩子会在每个进程上调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

示例:

class LitModel(...):
    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

        # don't do this
        self.something = else

    def setup(self, stage):
        data = load_data(...)
        self.l1 = nn.Linear(28, data.num_classes)

返回类型

None

share_memory()

参见 torch.Tensor.share_memory_()。

返回类型

~T

state_dict(*args, destination=None, prefix='', keep_vars=False)

返回一个包含模块整体状态引用的字典。

参数和持久缓冲区（例如运行平均值）都包含在内。键是对应的参数和缓冲区名称。设置为 None 的参数和缓冲区不包含在内。

备注

返回的对象是一个浅拷贝。它包含对模块参数和缓冲区的引用。

警告

目前，state_dict() 也按顺序接受 destination、prefix 和 keep_vars 的位置参数。然而，这正在被弃用，未来版本将强制使用关键字参数。

警告

请避免使用参数 destination，因为它不是为最终用户设计的。

参数

• destination (dict, optional) – 如果提供了，模块的状态将被更新到字典中，并返回相同的对象。否则，将创建并返回一个 OrderedDict。默认值：None。

• prefix (str, optional) – 在 state_dict 中组合键时添加到参数和缓冲区名称的前缀。默认值：''。

• keep_vars (bool, optional) – 默认情况下，状态字典中返回的 Tensor 是与自动求导分离的。如果设置为 True，则不会执行分离操作。默认值：False。

返回

包含模块整个状态的字典

返回类型

dict

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

property strict_loading: bool

确定 Lightning 如何使用 .load_state_dict(…, strict=model.strict_loading) 加载此模型。

返回类型

bool

property supports_probabilistic_prediction: bool

返回类型

bool

teardown(stage)

在 fit（训练 + 验证）、验证、测试或预测结束时调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

返回类型

None

test_dataloader()

一个可迭代对象或多个可迭代对象的集合，指定测试样本。

有关多个数据加载器的更多信息，请参阅此 部分。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• test()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要测试数据集和 test_step() ，你不需要实现这个方法。

返回类型

Any

test_step(*args, **kwargs)

对测试集中的单个批次数据进行操作。在这一步中，您通常会生成示例或计算任何感兴趣的内容，例如准确性。

参数

• batch – 你的数据可迭代对象的输出，通常是一个 DataLoader。

• batch_idx – 这一批的索引。

• dataloader_idx – 生成此批次的 dataloader 的索引。（仅在使用多个 dataloader 时）

返回类型

Union[Tensor, Mapping[str, Any], None]

返回

• Tensor - 损失张量

• dict - 一个字典。可以包含任何键，但必须包含键 'loss'。

• None - 跳到下一个批次。

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...

# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

示例:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

如果你传入多个测试数据加载器，test_step() 将会有一个额外的参数。我们建议将默认值设置为0，这样你可以快速在单个和多个数据加载器之间切换。

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

备注

如果你不需要测试，你就不需要实现这个方法。

备注

当调用 test_step() 时，模型已处于评估模式，并且已禁用 PyTorch 梯度。在测试轮次结束时，模型返回到训练模式，并且启用梯度。

to(*args, **kwargs)

参见 torch.nn.Module.to()。

返回类型

Self

to_dtype(dtype)

将模块精度（默认为 float32）转换为另一种精度。

to_empty(*, device, recurse=True)

将参数和缓冲区移动到指定设备，而不复制存储。

参数

• device (torch.device) – 此模块中参数和缓冲区的所需设备。

• recurse (bool) – 子模块的参数和缓冲区是否应递归地移动到指定设备。

返回

自己

返回类型

Module

to_onnx(file_path, input_sample=None, **kwargs)

将模型保存为 ONNX 格式。

参数

• file_path (Union[str, Path]) – 应保存 onnx 模型的文件路径。

• input_sample (Optional[Any, None]) – 用于跟踪的输入。默认值：None（使用 self.example_input_array）

• **kwargs – 将被传递给 torch.onnx.export 函数。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)

model = SimpleModel()
input_sample = torch.randn(1, 64)
model.to_onnx("export.onnx", input_sample, export_params=True)

返回类型

None

to_torchscript(file_path=None, method='script', example_inputs=None, **kwargs)

默认情况下，将整个模型编译为 ScriptModule。如果你想使用追踪，请提供参数 method='trace' 并确保提供了 example_inputs 参数，或者模型设置了 example_input_array。如果你想自定义被脚本的模块，你应该重写这个方法。如果你想要返回多个模块，我们建议使用字典。

参数

• file_path (Union[str, Path, None]) – 保存 torchscript 的路径。默认：None（不保存文件）。

• method (Optional[str, None]) – 是否使用 TorchScript 的 script 或 trace 方法。默认值：’script’

• example_inputs (Optional[Any, None]) – 当方法设置为’trace’时，用于进行追踪的输入。默认值：None（使用 example_input_array）

• **kwargs – 将传递给 torch.jit.script() 或 torch.jit.trace() 函数的额外参数。

备注

• 需要实现 forward() 方法。

• 导出的脚本将被设置为评估模式。

• 建议您安装最新支持版本的 PyTorch 以无限制地使用此功能。另请参阅 torch.jit 文档以了解支持的功能。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

model = SimpleModel()
model.to_torchscript(file_path="model.pt")

torch.jit.save(model.to_torchscript(
    file_path="model_trace.pt", method='trace', example_inputs=torch.randn(1, 64))
)

返回类型

Union[ScriptModule, Dict[str, ScriptModule]]

返回

无论 file_path 是否定义，此 LightningModule 都可以作为 torchscript 使用。

toggle_optimizer(optimizer)

确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。

它与 untoggle_optimizer() 一起工作，以确保 param_requires_grad_state 被正确重置。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要切换的优化器。

返回类型

None

train(mode=True)

设置模块为训练模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

参数

mode (bool) – 是否设置训练模式（True）或评估模式（False）。默认值：True。

返回

自己

返回类型

Module

train_criterion_reduction: Optional[str]

train_dataloader()

指定训练样本的可迭代对象或可迭代对象的集合。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• fit()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

property trainer: Trainer

返回类型

Trainer

training: bool

training_step(train_batch, batch_idx)

执行训练步骤

返回类型

Tensor

transfer_batch_to_device(batch, device, dataloader_idx)

如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。

以下列出的数据类型（以及它们的任意嵌套）都是开箱即用的：

• torch.Tensor 或任何实现了 .to(…) 的对象

• 列表

• 字典

• 元组

对于其他情况，你需要定义数据如何移动到目标设备（CPU、GPU、TPU等）。

备注

此钩子应仅传输数据，不应修改数据，也不应将数据移动到除作为参数传入的设备之外的任何其他设备（除非你知道自己在做什么）。要检查此钩子的当前执行状态，可以使用 self.trainer.training/testing/validating/predicting，以便可以根据你的需求添加不同的逻辑。

参数

• batch (Any) – 一批需要转移到新设备的数据。

• device (device) – 在 PyTorch 中定义的目标设备。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

对新设备上数据的引用。

示例:

def transfer_batch_to_device(self, batch, device, dataloader_idx):
    if isinstance(batch, CustomBatch):
        # move all tensors in your custom data structure to the device
        batch.samples = batch.samples.to(device)
        batch.targets = batch.targets.to(device)
    elif dataloader_idx == 0:
        # skip device transfer for the first dataloader or anything you wish
        pass
    else:
        batch = super().transfer_batch_to_device(batch, device, dataloader_idx)
    return batch

参见

• move_data_to_device()

• apply_to_collection()

type(dst_type)

参见 torch.nn.Module.type()。

返回类型

Self

unfreeze()

解冻所有参数以进行训练。

model = MyLightningModule(...)
model.unfreeze()

返回类型

None

untoggle_optimizer(optimizer)

重置通过 toggle_optimizer() 切换的所需梯度的状态。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要取消选择的优化器。

返回类型

None

val_criterion_reduction: Optional[str]

val_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• fit()

• validate()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要验证数据集和 validation_step()，你不需要实现这个方法。

返回类型

Any

validation_step(val_batch, batch_idx)

执行验证步骤

返回类型

Tensor

xpu(device=None)

将所有模型参数和缓冲区移动到XPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在 XPU 上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

zero_grad(set_to_none=True)

重置所有模型参数的梯度。

更多上下文请参见 torch.optim.Optimizer 下的类似函数。

参数

set_to_none (bool) – 而不是设置为零，将梯度设置为 None。详情请参阅 torch.optim.Optimizer.zero_grad()。

返回类型

None

class darts.models.forecasting.pl_forecasting_module.PLSplitCovariatesModule(input_chunk_length, output_chunk_length, output_chunk_shift=0, train_sample_shape=None, loss_fn=MSELoss(), torch_metrics=None, likelihood=None, optimizer_cls=<class 'torch.optim.adam.Adam'>, optimizer_kwargs=None, lr_scheduler_cls=None, lr_scheduler_kwargs=None, use_reversible_instance_norm=False)

基类：PLForecastingModule, ABC

基于 PyTorch Lightning 的预测模块。

这个类旨在通过继承来创建一个新的基于 PyTorch Lightning 的预测模块。当子类化这个类时，请确保添加以下具有给定签名的方法：

• PLForecastingModule.__init__()

• PLForecastingModule.forward()

• PLForecastingModule._produce_train_output()

• PLForecastingModule._get_batch_prediction()

在子类 MyModel 的 __init__() 函数中调用 super(MyModel, self).__init__(**kwargs)，其中 kwargs 是 PLForecastingModule 的参数。

参数

• input_chunk_length (int) – 过去作为模型输入的时间步数（每个块）。适用于目标序列，以及过去和/或未来协变量（如果模型支持）。

• output_chunk_length (int) – 内部模型一次预测的时间步数（每个块）。此外，模型输入中使用的未来协变量的未来值数量（如果模型支持未来协变量）。这与 predict() 中使用的预测范围 n 不同，后者是使用一次性或自回归预测生成的所需预测点数。设置 n <= output_chunk_length 可以防止自回归。当协变量不足以延伸到未来，或禁止模型使用过去和/或未来协变量的未来值进行预测时（取决于模型的协变量支持），这很有用。

• train_sample_shape (Optional[Tuple, None]) – 模型的输入形状，用于在不调用 fit_from_dataset 的情况下实例化模型，并对用于重新训练或预测的新训练/推理数据集进行完整性检查。

• loss_fn (_Loss) – 用于训练的 PyTorch 损失函数。如果指定了 likelihood 参数，则此参数将被概率模型忽略。默认值：torch.nn.MSELoss()。

• torch_metrics (Union[Metric, MetricCollection, None]) – 用于评估的 torch 指标或 MetricCollection。可用指标的完整列表可以在 https://torchmetrics.readthedocs.io/en/latest/ 找到。默认值：None。

• likelihood (Optional[Likelihood, None]) – Darts 的 似然模型 之一，用于概率预测。默认值：None。

• optimizer_cls (Optimizer) – 要使用的 PyTorch 优化器类。默认值：torch.optim.Adam。

• optimizer_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 优化器的键值参数（例如，{'lr': 1e-3} 用于指定学习率）。否则将使用所选 optimizer_cls 的默认值。默认值：None。

• lr_scheduler_cls (Optional[_LRScheduler, None]) – 可选地，指定要使用的 PyTorch 学习率调度器类。指定 None 对应于使用恒定的学习率。默认值：None。

• lr_scheduler_kwargs (Optional[Dict, None]) – 可选地，一些用于 PyTorch 学习率调度器的键值参数。默认值：None。

• use_reversible_instance_norm (bool) – 是否使用可逆实例归一化 RINorm 来对抗分布偏移，如 [1] 所示。它仅应用于目标序列的特征，而不应用于协变量。

引用

1

T. Kim et al. “Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift”, https://openreview.net/forum?id=cGDAkQo1C0p

属性

automatic_optimization	如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。
current_epoch	Trainer 中的当前时期，如果未附加则为 0。
device_mesh	像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。
example_input_array	示例输入数组是模块在 forward() 方法中可以消费的规范。
global_rank	当前进程在所有节点和设备中的索引。
global_step	所有epoch中看到的总训练批次。
hparams	使用 save_hyperparameters() 保存的超参数集合。
hparams_initial	使用 save_hyperparameters() 保存的超参数集合。
local_rank	当前进程在单个节点中的索引。
logger	Trainer 中对 logger 对象的引用。
loggers	参考 Trainer 中的日志记录器列表。
on_gpu	如果该模型当前位于GPU上，则返回 True。
output_chunk_length	模型一次预测的时间步数。
strict_loading	确定 Lightning 如何使用 .load_state_dict(..., strict=model.strict_loading) 加载此模型。

设备
dtype
已训练的轮数
fabric
supports_probabilistic_prediction
训练者

方法

add_module(name, module)	将一个子模块添加到当前模块。
all_gather(data[, group, sync_grads])	从多个进程中收集张量或张量集合。
apply(fn)	递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。
backward(loss, *args, **kwargs)	在 training_step() 中返回的损失上调用以执行反向传播。
bfloat16()	将所有浮点参数和缓冲区转换为 bfloat16 数据类型。
buffers([recurse])	返回一个遍历模块缓冲区的迭代器。
children()	返回一个遍历直接子模块的迭代器。
clip_gradients(optimizer[, ...])	内部处理梯度裁剪。
compile(*args, **kwargs)	使用 torch.compile() 编译此模块的 forward。
configure_callbacks()	配置模型特定的回调。
configure_gradient_clipping(optimizer[, ...])	对优化器参数执行梯度裁剪。
configure_model()	在策略和精度感知的环境中创建模块的钩子。
configure_optimizers()	配置优化器和学习率调度器以优化模型。
configure_sharded_model()	已弃用。
configure_torch_metrics(torch_metrics)	处理 torch_metrics 参数。
cpu()	参见 torch.nn.Module.cpu()。
cuda([device])	将所有模型参数和缓冲区移动到GPU。
double()	参见 torch.nn.Module.double()。
eval()	将模块设置为评估模式。
extra_repr()	设置模块的额外表示。
float()	参见 torch.nn.Module.float()。
forward(*args, **kwargs)	与 torch.nn.Module.forward() 相同。
freeze()	冻结所有参数以进行推理。
get_buffer(target)	如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。
get_extra_state()	返回要包含在模块的 state_dict 中的任何额外状态。
get_parameter(target)	如果存在，则返回由 target 给出的参数，否则抛出错误。
get_submodule(target)	如果存在，返回由 target 指定的子模块，否则抛出错误。
half()	参见 torch.nn.Module.half()。
ipu([device])	将所有模型参数和缓冲区移动到 IPU。
load_from_checkpoint(checkpoint_path[, ...])	从检查点加载模型的主要方式。
load_state_dict(state_dict[, strict, assign])	将参数和缓冲区从 state_dict 复制到此模块及其后代中。
log(name, value[, prog_bar, logger, ...])	记录一个键值对。
log_dict(dictionary[, prog_bar, logger, ...])	一次性记录一个值的字典。
lr_scheduler_step(scheduler, metric)	重写此方法以调整 Trainer 调用每个调度器的默认方式。
lr_schedulers()	返回在训练期间使用的学习率调度器。
manual_backward(loss, *args, **kwargs)	在进行手动优化时，直接从 training_step() 调用此方法。
modules()	返回网络中所有模块的迭代器。
named_buffers([prefix, recurse, ...])	返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。
named_children()	返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。
named_modules([memo, prefix, remove_duplicate])	返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。
named_parameters([prefix, recurse, ...])	返回一个模块参数的迭代器，同时生成参数的名称和参数本身。
on_after_backward()	在 loss.backward() 之后和优化器步进之前调用。
on_after_batch_transfer(batch, dataloader_idx)	覆盖此方法以在批次转移到设备后对其进行批量增强或修改。
on_before_backward(loss)	在 loss.backward() 之前调用。
on_before_batch_transfer(batch, dataloader_idx)	在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。
on_before_optimizer_step(optimizer)	在 optimizer.step() 之前调用。
on_before_zero_grad(optimizer)	在 training_step() 之后和 optimizer.zero_grad() 之前调用。
on_fit_end()	在拟合的最后调用。
on_fit_start()	在拟合开始时调用。
on_load_checkpoint(checkpoint)	由 Lightning 调用以恢复您的模型。
on_predict_batch_end(outputs, batch, batch_idx)	在批次之后调用预测循环。
on_predict_batch_start(batch, batch_idx[, ...])	在预测循环中，在该批次发生任何事情之前调用。
on_predict_end()	在预测结束时调用。
on_predict_epoch_end()	在预测结束时调用。
on_predict_epoch_start()	在预测开始时调用。
on_predict_model_eval()	在预测循环开始时调用。
on_predict_start()	在预测开始时调用。
on_save_checkpoint(checkpoint)	当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。
on_test_batch_end(outputs, batch, batch_idx)	在批次之后在测试循环中调用。
on_test_batch_start(batch, batch_idx[, ...])	在测试循环中，在该批次发生任何事情之前调用。
on_test_end()	在测试结束时调用。
on_test_epoch_end()	在每个epoch的测试循环结束时调用。
on_test_epoch_start()	在epoch开始时在测试循环中调用。
on_test_model_eval()	在测试循环开始时调用。
on_test_model_train()	在测试循环结束时调用。
on_test_start()	在测试开始时调用。
on_train_batch_end(outputs, batch, batch_idx)	在训练循环中批次处理后调用。
on_train_batch_start(batch, batch_idx)	在训练循环中，在该批次发生任何事情之前调用。
on_train_end()	在训练结束前调用，在记录器实验关闭之前。
on_train_epoch_end()	在训练循环的每个epoch结束时调用。
on_train_epoch_start()	在训练循环中，在epoch开始时调用。
on_train_start()	在完整性检查后，训练开始时调用。
on_validation_batch_end(outputs, batch, ...)	在批次验证循环后调用。
on_validation_batch_start(batch, batch_idx)	在验证循环中，对该批次进行任何操作之前调用。
on_validation_end()	在验证结束时调用。
on_validation_epoch_end()	在epoch结束时的验证循环中调用。
on_validation_epoch_start()	在epoch开始时验证循环中调用。
on_validation_model_eval()	在验证循环开始时调用。
on_validation_model_train()	当验证循环结束时调用。
on_validation_model_zero_grad()	在进入验证循环之前，由训练循环调用来释放梯度。
on_validation_start()	在验证开始时调用。
optimizer_step(epoch, batch_idx, optimizer)	重写此方法以调整 Trainer 调用优化器的默认方式。
optimizer_zero_grad(epoch, batch_idx, optimizer)	重写此方法以改变 optimizer.zero_grad() 的默认行为。
optimizers([use_pl_optimizer])	返回训练期间使用的优化器。
parameters([recurse])	返回一个遍历模块参数的迭代器。
predict_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。
predict_step(batch, batch_idx[, dataloader_idx])	执行预测步骤
prepare_data()	使用此功能下载并准备数据。
print(*args, **kwargs)	仅从进程 0 打印。
register_backward_hook(hook)	在模块上注册一个反向钩子。
register_buffer(name, tensor[, persistent])	向模块添加一个缓冲区。
register_forward_hook(hook, *[, prepend, ...])	在模块上注册一个前向钩子。
register_forward_pre_hook(hook, *[, ...])	在模块上注册一个前向预钩子。
register_full_backward_hook(hook[, prepend])	在模块上注册一个反向钩子。
register_full_backward_pre_hook(hook[, prepend])	在模块上注册一个反向预钩子。
register_load_state_dict_post_hook(hook)	注册一个在模块的 load_state_dict 调用后运行的后置钩子。
register_module(name, module)	别名 add_module()。
register_parameter(name, param)	向模块添加一个参数。
register_state_dict_pre_hook(hook)	为 state_dict() 方法注册一个预钩子。
requires_grad_([requires_grad])	更改是否应在此模块的参数上记录自动梯度操作。
save_hyperparameters(*args[, ignore, frame, ...])	将参数保存到 hparams 属性中。
set_extra_state(state)	设置包含在加载的 state_dict 中的额外状态。
set_predict_parameters(n, num_samples, ...)	在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。
setup(stage)	在 fit（训练 + 验证）、验证、测试或预测开始时调用。
share_memory()	参见 torch.Tensor.share_memory_()。
state_dict(*args[, destination, prefix, ...])	返回一个包含模块整体状态引用的字典。
teardown(stage)	在 fit（训练 + 验证）、验证、测试或预测结束时调用。
test_dataloader()	一个可迭代对象或多个可迭代对象的集合，指定测试样本。
test_step(*args, **kwargs)	对测试集中的单个批次数据进行操作。
to(*args, **kwargs)	参见 torch.nn.Module.to()。
to_dtype(dtype)	将模块精度（默认为 float32）转换为另一种精度。
to_empty(*, device[, recurse])	将参数和缓冲区移动到指定设备，而不复制存储。
to_onnx(file_path[, input_sample])	将模型保存为 ONNX 格式。
to_torchscript([file_path, method, ...])	默认情况下，将整个模型编译为 ScriptModule。
toggle_optimizer(optimizer)	确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。
train([mode])	设置模块为训练模式。
train_dataloader()	指定训练样本的可迭代对象或可迭代对象的集合。
training_step(train_batch, batch_idx)	执行训练步骤
transfer_batch_to_device(batch, device, ...)	如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。
type(dst_type)	参见 torch.nn.Module.type()。
unfreeze()	解冻所有参数以进行训练。
untoggle_optimizer(optimizer)	重置通过 toggle_optimizer() 切换的所需梯度的状态。
val_dataloader()	一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。
validation_step(val_batch, batch_idx)	执行验证步骤
xpu([device])	将所有模型参数和缓冲区移动到XPU。
zero_grad([set_to_none])	重置所有模型参数的梯度。

__call__
set_mc_dropout

CHECKPOINT_HYPER_PARAMS_KEY = 'hyper_parameters'

CHECKPOINT_HYPER_PARAMS_NAME = 'hparams_name'

CHECKPOINT_HYPER_PARAMS_TYPE = 'hparams_type'

T_destination

TypeVar(‘T_destination’, bound=Dict[str, Any]) 的别名

add_module(name, module)

将一个子模块添加到当前模块。

可以通过给定的名称将模块作为属性访问。

参数

• name (str) – 子模块的名称。可以使用给定的名称从此模块访问子模块。

• module (Module) – 要添加到模块中的子模块。

返回类型

None

all_gather(data, group=None, sync_grads=False)

从多个进程中收集张量或张量集合。

此方法需要在所有进程上调用，并且所有进程上的张量需要具有相同的形状，否则程序将永远停滞。

参数

• data (Union[Tensor, Dict, List, Tuple]) – int, float, 形状为 (batch, …) 的张量，或可能是嵌套的集合。

• group (Optional[Any, None]) – 要从中收集结果的进程组。默认为所有进程（world）

• sync_grads (bool) – 允许用户为 all_gather 操作同步梯度的标志

返回类型

Union[Tensor, Dict, List, Tuple]

返回

一个形状为 (world_size, batch, …) 的张量，或者如果输入是一个集合，输出也将是一个具有这种形状张量的集合。对于 world_size 为 1 的特殊情况，张量不会增加额外的维度。

allow_zero_length_dataloader_with_multiple_devices: bool

apply(fn)

递归地将 fn 应用于每个子模块（由 .children() 返回）以及自身。

典型用途包括初始化模型的参数（另见 nn-init-doc）。

参数

fn (Module -> None) – 要应用于每个子模块的函数

返回

自己

返回类型

Module

示例:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

property automatic_optimization: bool

如果设置为 False ，你需要负责调用 .backward() 、 .step() 、 .zero_grad() 。

返回类型

bool

backward(loss, *args, **kwargs)

在 training_step() 中返回的损失上调用以执行反向传播。如果需要，请使用您自己的实现覆盖此钩子。

参数

loss (Tensor) – 由 training_step() 返回的损失张量。如果使用梯度累积，这里的损失值是归一化的（按 1 / 累积步数缩放）。

示例:

def backward(self, loss):
    loss.backward()

返回类型

None

bfloat16()

将所有浮点参数和缓冲区转换为 bfloat16 数据类型。

备注

此方法就地修改模块。

返回

自己

返回类型

Module

buffers(recurse=True)

返回一个遍历模块缓冲区的迭代器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。

生成器

torch.Tensor – 模块缓冲区

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Tensor]

call_super_init: bool = False

children()

返回一个遍历直接子模块的迭代器。

生成器

模块 – 子模块

返回类型

Iterator[Module]

clip_gradients(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

内部处理梯度裁剪。

备注

• 不要重写此方法。如果你想自定义梯度裁剪，考虑使用 configure_gradient_clipping() 方法。

• 对于手动优化（self.automatic_optimization = False），如果你想使用梯度裁剪，可以在训练步骤中手动调用 self.clip_gradients(opt, gradient_clip_val=0.5, gradient_clip_algorithm="norm")。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。

• gradient_clip_algorithm (Optional[str, None]) – 使用的梯度裁剪算法。传递 gradient_clip_algorithm="value" 按值裁剪，传递 gradient_clip_algorithm="norm" 按范数裁剪。

返回类型

None

compile(*args, **kwargs)

使用 torch.compile() 编译此模块的 forward。

此模块的 __call__ 方法被编译，所有参数都按原样传递给 torch.compile()。

有关此函数的参数详情，请参阅 torch.compile()。

configure_callbacks()

配置模型特定的回调。当模型被附加时，例如当调用 .fit() 或 .test() 时，这里返回的回调列表或回调将与传递给 Trainer 的 callbacks 参数的回调列表合并。如果这里返回的回调与 Trainer 的回调列表中已有的一个或多个回调类型相同，它将优先并替换它们。此外，Lightning 将确保 ModelCheckpoint 回调最后运行。

返回类型

Union[Sequence[Callback], Callback]

返回

一个回调函数或回调函数列表，这些回调函数将扩展 Trainer 中的回调函数列表。

示例:

def configure_callbacks(self):
    early_stop = EarlyStopping(monitor="val_acc", mode="max")
    checkpoint = ModelCheckpoint(monitor="val_loss")
    return [early_stop, checkpoint]

configure_gradient_clipping(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

对优化器参数执行梯度裁剪。在调用 optimizer_step() 之前执行。

参数

• optimizer (Optimizer) – 当前使用的优化器。

• gradient_clip_val (Union[int, float, None]) – 要裁剪梯度的值。默认情况下，传递给 Trainer 的值将在此处可用。

• gradient_clip_algorithm (Optional[str, None]) – 要使用的梯度裁剪算法。默认情况下，传递给 Trainer 的值将在此处可用。

示例:

def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm):
    # Implement your own custom logic to clip gradients
    # You can call `self.clip_gradients` with your settings:
    self.clip_gradients(
        optimizer,
        gradient_clip_val=gradient_clip_val,
        gradient_clip_algorithm=gradient_clip_algorithm
    )

返回类型

None

configure_model()

在策略和精度感知的环境中创建模块的钩子。

这在使用分片策略（FSDP 和 DeepSpeed）时特别有用，我们希望立即分片模型以节省内存和初始化时间。对于非分片策略，您可以选择覆盖此钩子或在 init_module() 上下文管理器下初始化您的模型。

此钩子在每个 fit/val/test/predict 阶段都会被调用，因此请确保此钩子的实现是 幂等的，即，在第一次调用钩子后，后续的调用应该是无操作的。

返回类型

None

configure_optimizers()

配置优化器和学习率调度器以优化模型。

configure_sharded_model()

已弃用。

使用 configure_model() 代替。

返回类型

None

static configure_torch_metrics(torch_metrics)

处理 torch_metrics 参数。

返回类型

MetricCollection

cpu()

参见 torch.nn.Module.cpu()。

返回类型

Self

cuda(device=None)

将所有模型参数和缓冲区移动到GPU。这也会使相关的参数和缓冲区成为不同的对象。因此，如果模块在优化期间将驻留在GPU上，则应在构造优化器之前调用此方法。

参数

device (Union[device, int, None]) – 如果指定，所有参数将被复制到该设备。如果为 None，将使用当前的 CUDA 设备索引。

返回

自己

返回类型

Module

property current_epoch: int

Trainer 中的当前时期，如果未附加则为 0。

返回类型

int

property device: device

返回类型

device

property device_mesh: Optional[DeviceMesh]

像 ModelParallelStrategy 这样的策略将创建一个设备网格，可以在 configure_model() 钩子中访问该网格以并行化 LightningModule。

返回类型

Optional[ForwardRef, None]

double()

参见 torch.nn.Module.double()。

返回类型

Self

property dtype: Union[str, dtype]

返回类型

Union[str, dtype]

dump_patches: bool = False

property epochs_trained

eval()

将模块设置为评估模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

这等同于 self.train(False)。

请参阅 locally-disable-grad-doc 以比较 .eval() 和几个可能与之混淆的类似机制。

返回

自己

返回类型

Module

property example_input_array: Optional[Union[Tensor, Tuple, Dict]]

示例输入数组是模块在 forward() 方法中可以使用的规范。返回类型解释如下：

• 单个张量：假设模型接受一个参数，即 model.forward(model.example_input_array)

• 元组：输入数组应被解释为位置参数的序列，即 model.forward(*model.example_input_array)

• 字典：输入数组表示命名关键字参数，即 model.forward(**model.example_input_array)

返回类型

Union[Tensor, Tuple, Dict, None]

extra_repr()

设置模块的额外表示。

要打印自定义的额外信息，您应该在您自己的模块中重新实现此方法。单行和多行字符串都是可接受的。

返回类型

str

property fabric: Optional[Fabric]

返回类型

Optional[Fabric, None]

float()

参见 torch.nn.Module.float()。

返回类型

Self

abstract forward(*args, **kwargs)

与 torch.nn.Module.forward() 相同。

参数

• *args – 无论你决定传入 forward 方法的内容。

• **kwargs – 关键字参数也是可能的。

返回类型

Any

返回

你的模型的输出

freeze()

冻结所有参数以进行推理。

示例:

model = MyLightningModule(...)
model.freeze()

返回类型

None

get_buffer(target)

如果 target 指定的缓冲区存在，则返回该缓冲区，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的缓冲区的完全限定字符串名称。（有关如何指定完全限定字符串，请参见 get_submodule。）

返回

target 引用的缓冲区

返回类型

torch.Tensor

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是缓冲区的内容

get_extra_state()

返回要包含在模块的 state_dict 中的任何额外状态。

如果需要存储额外状态，请在你的模块中实现此功能以及相应的 set_extra_state()。此函数在构建模块的 state_dict() 时被调用。

请注意，额外的状态应该是可pickle的，以确保state_dict的序列化工作。我们仅对序列化张量提供向后兼容性保证；如果其他对象的序列化pickle形式发生变化，它们可能会破坏向后兼容性。

返回

模块的 state_dict 中存储的任何额外状态

返回类型

object

get_parameter(target)

如果存在，则返回由 target 给出的参数，否则抛出错误。

有关此方法功能的更详细说明以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数

target (str) – 要查找的参数的全限定字符串名称。（关于如何指定全限定字符串，请参见 get_submodule。）

返回

target 引用的参数

返回类型

torch.nn.Parameter

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Parameter 的内容

get_submodule(target)

如果存在，返回由 target 指定的子模块，否则抛出错误。

例如，假设你有一个 nn.Module A ，它看起来像这样：

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(该图展示了一个 nn.Module A。A 有一个嵌套的子模块 net_b，它本身有两个子模块 net_c 和 linear。net_c 然后有一个子模块 conv。)

要检查我们是否拥有 linear 子模块，我们会调用 get_submodule("net_b.linear")。要检查我们是否拥有 conv 子模块，我们会调用 get_submodule("net_b.net_c.conv")。

get_submodule 的运行时间受限于 target 中的模块嵌套程度。对 named_modules 的查询可以实现相同的结果，但它是 O(N) 的，其中 N 是传递模块的数量。因此，对于检查某个子模块是否存在的简单检查，应始终使用 get_submodule。

参数

target (str) – 要查找的子模块的完全限定字符串名称。（请参见上面的示例，了解如何指定完全限定的字符串。）

返回

由 target 引用的子模块

返回类型

torch.nn.Module

引发

AttributeError – 如果目标字符串引用了一个无效的路径或解析为不是 nn.Module 的内容

property global_rank: int

当前进程在所有节点和设备中的索引。

返回类型

int

property global_step: int

所有epoch中看到的总训练批次。

如果没有连接训练器，此属性为 0。

返回类型

int

half()

参见 torch.nn.Module.half()。

返回类型

Self

property hparams: Union[AttributeDict, MutableMapping]

通过 save_hyperparameters() 保存的超参数集合。用户可以对其进行修改。对于初始超参数的冻结集，请使用 hparams_initial。

返回类型

Union[AttributeDict, MutableMapping]

返回

可变超参数字典

property hparams_initial: AttributeDict

通过 save_hyperparameters() 保存的超参数集合。这些内容是只读的。可以通过 hparams 手动更新保存的超参数。

返回

不可变的初始超参数

返回类型

AttributeDict

ipu(device=None)

将所有模型参数和缓冲区移动到 IPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在IPU上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs)

从检查点加载模型的主要方式。当Lightning保存检查点时，它会将传递给 __init__ 的参数存储在检查点中的 "hyper_parameters" 下。

通过 **kwargs 指定的任何参数将覆盖存储在 "hyper_parameters" 中的参数。

参数

• checkpoint_path (Union[str, Path, IO]) – 检查点的路径。这也可以是一个URL，或类似文件的对象。

• map_location (Union[device, str, int, Callable[[UntypedStorage, str], Optional[UntypedStorage, None]], Dict[Union[device, str, int], Union[device, str, int]], None]) – 如果你的检查点保存了一个GPU模型，而你现在在CPU上或不同数量的GPU上加载，使用这个来映射到新的设置。其行为与 torch.load() 相同。

• hparams_file (Union[str, Path, None]) –

  可选路径到一个 .yaml 或 .csv 文件，其层次结构如本示例所示:

  drop_prob: 0.2
  dataloader:
      batch_size: 32
  

  您很可能不需要这个，因为 Lightning 总是会将超参数保存到检查点。然而，如果您的检查点权重没有保存超参数，请使用此方法传入一个包含您希望使用的超参数的 .yaml 文件。这些将被转换为 dict 并传递到您的 LightningModule 中使用。

  如果你的模型的 hparams 参数是 Namespace 并且 .yaml 文件具有层次结构，你需要重构你的模型以将 hparams 视为 dict。

• strict (Optional[bool, None]) – 是否严格要求 checkpoint_path 中的键与该模块状态字典返回的键匹配。默认为 True ，除非设置了 LightningModule.strict_loading ，在这种情况下，它默认为 LightningModule.strict_loading 的值。

• **kwargs – 初始化模型所需的任何额外关键字参数。也可以用来覆盖已保存的超参数值。

返回类型

Self

返回

LightningModule 实例，加载了权重和超参数（如果可用）。

备注

load_from_checkpoint 是一个 类 方法。你应该使用你的 LightningModule 类 来调用它，而不是 LightningModule 实例，否则会引发 TypeError。

备注

为了确保所有层都能从检查点加载，如果您的 LightningModule 中重写了此钩子，此函数将在实例化模型后直接调用 configure_model()。但是，请注意 load_from_checkpoint 不支持加载分片检查点，如果模型太大，您可能会耗尽内存。在这种情况下，请考虑通过 .fit(ckpt_path=...) 通过 Trainer 加载。

示例:

# load weights without mapping ...
model = MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')

# or load weights mapping all weights from GPU 1 to GPU 0 ...
map_location = {'cuda:1':'cuda:0'}
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    map_location=map_location
)

# or load weights and hyperparameters from separate files.
model = MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    hparams_file='/path/to/hparams_file.yaml'
)

# override some of the params with new values
model = MyLightningModule.load_from_checkpoint(
    PATH,
    num_layers=128,
    pretrained_ckpt_path=NEW_PATH,
)

# predict
pretrained_model.eval()
pretrained_model.freeze()
y_hat = pretrained_model(x)

load_state_dict(state_dict, strict=True, assign=False)

将参数和缓冲区从 state_dict 复制到此模块及其后代中。

如果 strict 是 True，那么 state_dict 的键必须与该模块的 state_dict() 函数返回的键完全匹配。

警告

如果 assign 是 True，优化器必须在调用 load_state_dict 之后创建，除非 get_swap_module_params_on_conversion() 是 True。

参数

• state_dict (dict) – 包含参数和持久缓冲区的字典。

• strict (bool, optional) – 是否严格要求 state_dict 中的键与该模块的 state_dict() 函数返回的键匹配。默认值：True

• assign (bool, optional) – 当 False 时，当前模块中张量的属性会被保留，而当 True 时，状态字典中张量的属性会被保留。唯一的例外是 Parameter 的 requires_grad 字段，该字段的值会从模块中保留。默认值：False

返回

• missing_keys 是一个包含任何预期键的 str 列表

  此模块中存在但未在提供的 state_dict 中出现的。

• unexpected_keys 是一个包含未预料到的键的 str 列表。

  此模块预期但未在提供的 state_dict 中出现。

返回类型

带有 missing_keys 和 unexpected_keys 字段的 NamedTuple

备注

如果一个参数或缓冲区被注册为 None ，并且其对应的键存在于 state_dict 中，load_state_dict() 将引发一个 RuntimeError 。

property local_rank: int

当前进程在单个节点中的索引。

返回类型

int

log(name, value, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, metric_attribute=None, rank_zero_only=False)

记录一个键值对。

示例:

self.log('train_loss', loss)

每个钩子的默认行为在此处有文档说明：extensions/logging:自动日志记录。

参数

• name (str) – 日志的键。如果使用DDP或其他分布式策略，必须在所有进程中保持一致。

• value (Union[Metric, Tensor, int, float]) – 要记录的值。可以是 float、Tensor 或 Metric。

• prog_bar (bool) – 如果 True 则记录到进度条。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• on_epoch (Optional[bool, None]) – 如果 True ，记录每个epoch累积的指标。默认值由钩子决定。详情请参见 extensions/logging:自动日志记录 。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图表。

• sync_dist (bool) – 如果 True，将在设备间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的DDP组。

• add_dataloader_idx (bool) – 如果 True，将当前数据加载器的索引附加到名称（当使用多个数据加载器时）。如果为 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的 batch_size。这将直接从加载的批次中推断出来，但对于某些数据结构，您可能需要显式提供它。

• metric_attribute (Optional[str, None]) – 要恢复指标状态，Lightning 需要模型中 torchmetrics.Metric 的引用。如果它是模型属性，则会自动找到。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

log_dict(dictionary, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, rank_zero_only=False)

一次性记录一个值的字典。

示例:

values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n}
self.log_dict(values)

参数

• dictionary (Union[Mapping[str, Union[Metric, Tensor, int, float]], MetricCollection]) – 键值对。如果使用DDP或其他分布式策略，键必须在所有进程中保持一致。值可以是 float 、 Tensor 、 Metric 或 MetricCollection 。

• prog_bar (bool) – 如果 True 则记录到进度基础。

• logger (Optional[bool, None]) – 如果 True 则记录到日志记录器。

• on_step (Optional[bool, None]) – 如果 True 在此步骤记录日志。None 自动记录训练步骤，但不记录验证/测试步骤。默认值由钩子决定。详情请参阅 extensions/logging:自动记录。

• on_epoch (Optional[bool, None]) – 如果 True 记录累积的 epoch 指标。None 自动记录 val/test 步骤的日志，但不包括 training_step。默认值由钩子决定。详情请参阅 extensions/logging:自动日志记录。

• reduce_fx (Union[str, Callable]) – 在每个epoch结束时对步骤值进行缩减的函数。默认使用 torch.mean()。

• enable_graph (bool) – 如果 True，将不会自动分离图形

• sync_dist (bool) – 如果 True，将在GPU/TPU之间减少指标。请谨慎使用，因为这可能会导致显著的通信开销。

• sync_dist_group (Optional[Any, None]) – 要同步的 ddp 组。

• add_dataloader_idx (bool) – 如果 True，则将当前数据加载器的索引附加到名称（在使用多个数据加载器时）。如果 False，用户需要为每个数据加载器提供唯一的名称，以避免混合值。

• batch_size (Optional[int, None]) – 当前的批次大小。这将直接从加载的批次中推断出来，但某些数据结构可能需要显式提供它。

• rank_zero_only (bool) – 告诉 Lightning 你是否从每个进程（默认）或仅从 rank 0 调用 self.log。如果为 True，你将无法在回调中使用此指标作为监视器（例如，早期停止）。警告：不当使用可能导致死锁！更多详情请参见 高级日志记录。

返回类型

None

property logger: Optional[Union[Logger, Logger]]

Trainer 中对 logger 对象的引用。

返回类型

Union[Logger, Logger, None]

property loggers: Union[List[Logger], List[Logger]]

参考 Trainer 中的日志记录器列表。

返回类型

Union[List[Logger], List[Logger]]

lr_scheduler_step(scheduler, metric)

重写此方法以调整 Trainer 调用每个调度器的默认方式。默认情况下，Lightning 会根据每个调度器的 interval 调用 step() 方法，如示例所示。

参数

• scheduler (Union[LRScheduler, ReduceLROnPlateau]) – 学习率调度器。

• metric (Optional[Any, None]) – 用于 ReduceLROnPlateau 等调度器的监视器值。

示例:

# DEFAULT
def lr_scheduler_step(self, scheduler, metric):
    if metric is None:
        scheduler.step()
    else:
        scheduler.step(metric)

# Alternative way to update schedulers if it requires an epoch value
def lr_scheduler_step(self, scheduler, metric):
    scheduler.step(epoch=self.current_epoch)

返回类型

None

lr_schedulers()

返回在训练期间使用的学习率调度器。对于手动优化很有用。

返回类型

Union[None, List[Union[LRScheduler, ReduceLROnPlateau]], LRScheduler, ReduceLROnPlateau]

返回

一个调度器，或者在存在多个调度器的情况下是一个调度器列表，如果在 configure_optimizers() 中没有返回调度器，则为 None。

manual_backward(loss, *args, **kwargs)

当你手动进行优化时，请直接从你的 training_step() 调用此方法。通过使用此方法，Lightning 可以确保在使用混合精度时应用所有适当的缩放。

更多示例请参见 手动优化。

示例:

def training_step(...):
    opt = self.optimizers()
    loss = ...
    opt.zero_grad()
    # automatically applies scaling, etc...
    self.manual_backward(loss)
    opt.step()

参数

• loss (Tensor) – 要计算梯度的张量。必须附加一个图。

• *args – 要转发给 backward() 的附加位置参数

• **kwargs – 要转发给 backward() 的额外关键字参数

返回类型

None

modules()

返回网络中所有模块的迭代器。

生成器

模块 – 网络中的一个模块

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

返回类型

Iterator[Module]

named_buffers(prefix='', recurse=True, remove_duplicate=True)

返回一个模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。

参数

• prefix (str) – 要添加到所有缓冲区名称的前缀。

• recurse (bool, optional) – 如果为 True，则生成此模块及其所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。默认为 True。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的缓冲区。默认为 True。

生成器

(str, torch.Tensor) – 包含名称和缓冲区的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

返回类型

Iterator[Tuple[str, Tensor]]

named_children()

返回一个遍历直接子模块的迭代器，生成模块的名称以及模块本身。

生成器

(str, Module) – 包含名称和子模块的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

返回类型

Iterator[Tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)

返回一个遍历网络中所有模块的迭代器，同时生成模块的名称和模块本身。

参数

• memo (Optional[Set[Module], None]) – 一个备忘录，用于存储已添加到结果中的模块集合

• prefix (str) – 一个将添加到模块名称前的前缀

• remove_duplicate (bool) – 是否在结果中移除重复的模块实例

生成器

(str, Module) – 名称和模块的元组

备注

重复的模块只会返回一次。在下面的例子中，l 只会返回一次。

示例:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='', recurse=True, remove_duplicate=True)

返回一个模块参数的迭代器，同时生成参数的名称和参数本身。

参数

• prefix (str) – 要添加到所有参数名称前的前缀。

• recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

• remove_duplicate (bool, optional) – 是否在结果中移除重复的参数。默认为 True。

生成器

(str, Parameter) – 包含名称和参数的元组

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

返回类型

Iterator[Tuple[str, Parameter]]

on_after_backward()

在 loss.backward() 之后和优化器步进之前调用。

备注

如果使用原生 AMP，梯度在此刻将不会被缩放。如果你需要未缩放的梯度，请使用 on_before_optimizer_step。

返回类型

None

on_after_batch_transfer(batch, dataloader_idx)

覆盖此方法以在批次转移到设备后对其进行批量增强或修改。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_after_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = gpu_transforms(batch['x'])
    return batch

参见

• on_before_batch_transfer()

• transfer_batch_to_device()

on_before_backward(loss)

在 loss.backward() 之前调用。

参数

loss (Tensor) – 损失除以梯度累积的批次数量，如果使用AMP则进行缩放。

返回类型

None

on_before_batch_transfer(batch, dataloader_idx)

在将批次数据传输到设备之前，重写此方法以更改或应用批次增强。

备注

要检查此钩子的当前执行状态，您可以使用 self.trainer.training/testing/validating/predicting，以便根据您的需求添加不同的逻辑。

参数

• batch (Any) – 需要被修改或增强的一批数据。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

一批数据

示例:

def on_before_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = transforms(batch['x'])
    return batch

参见

• on_after_batch_transfer()

• transfer_batch_to_device()

on_before_optimizer_step(optimizer)

在 optimizer.step() 之前调用。

如果使用梯度累积，钩子会在梯度累积完成后被调用。参见：:paramref:`~pytorch_lightning.trainer.trainer.Trainer.accumulate_grad_batches`。

如果使用AMP，损失将在调用此钩子之前未缩放。有关梯度缩放的更多信息，请参阅这些 文档。

如果裁剪梯度，梯度尚未被裁剪。

参数

optimizer (Optimizer) – 当前使用的优化器。

示例:

def on_before_optimizer_step(self, optimizer):
    # example to inspect gradient information in tensorboard
    if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
        for k, v in self.named_parameters():
            self.logger.experiment.add_histogram(
                tag=k, values=v.grad, global_step=self.trainer.global_step
            )

返回类型

None

on_before_zero_grad(optimizer)

在 training_step() 之后和 optimizer.zero_grad() 之前调用。

在训练循环中，调用优化器步骤之后和清零梯度之前。这是一个检查更新后权重信息的好地方。

这是调用它的地方:

for optimizer in optimizers:
    out = training_step(...)

    model.on_before_zero_grad(optimizer) # < ---- called here
    optimizer.zero_grad()

    backward()

参数

optimizer (Optimizer) – 应该被置零的梯度优化器。

返回类型

None

on_fit_end()

在拟合的最后调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

on_fit_start()

在拟合开始时调用。

如果在DDP上，它会在每个进程上调用

返回类型

None

property on_gpu: bool

如果该模型当前位于GPU上，则返回 True。

在LightningModule周围设置标志以区分CPU和GPU行为非常有用。

返回类型

bool

on_load_checkpoint(checkpoint)

由 Lightning 调用以恢复您的模型。如果您使用 on_save_checkpoint() 保存了某些内容，这是您恢复这些内容的机会。

参数

checkpoint (Dict[str, Any]) – 加载的检查点

示例:

def on_load_checkpoint(self, checkpoint):
    # 99% of the time you don't need to implement this method
    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']

备注

Lightning 自动恢复全局步数、轮次和训练状态，包括 amp 缩放。您无需恢复任何与训练相关的内容。

返回类型

None

on_predict_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后调用预测循环。

参数

• outputs (Optional[Any, None]) – predict_step(x) 的输出

• batch (Any) – 由预测 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_batch_start(batch, batch_idx, dataloader_idx=0)

在预测循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_predict_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_end()

在预测结束时调用。

返回类型

None

on_predict_epoch_start()

在预测开始时调用。

返回类型

None

on_predict_model_eval()

在预测循环开始时调用。

默认情况下，预测循环在开始之前会在 LightningModule 上调用 .eval()。重写此钩子以更改行为。

返回类型

None

on_predict_start()

在预测开始时调用。

返回类型

None

on_save_checkpoint(checkpoint)

当Lightning保存检查点时调用，给你一个机会来存储你可能想要保存的任何其他内容。

参数

checkpoint (Dict[str, Any]) – 在转储到文件之前的完整检查点字典。此钩子的实现可以将额外数据插入此字典中。

示例:

def on_save_checkpoint(self, checkpoint):
    # 99% of use cases you don't need to implement this method
    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object

备注

Lightning 保存了训练的所有方面（如 epoch、全局步数等），包括 amp 缩放。您无需存储任何关于训练的信息。

返回类型

None

on_test_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次之后在测试循环中调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – test_step(x) 的输出

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_batch_start(batch, batch_idx, dataloader_idx=0)

在测试循环中，在该批次发生任何事情之前调用。

参数

• batch (Any) – 测试 DataLoader 返回的批量数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_test_end()

在测试结束时调用。

返回类型

None

on_test_epoch_end()

在每个epoch的测试循环结束时调用。

返回类型

None

on_test_epoch_start()

在epoch开始时在测试循环中调用。

返回类型

None

on_test_model_eval()

在测试循环开始时调用。

默认情况下，测试循环在开始前会调用 .eval() 方法在 LightningModule 上。重写此钩子以改变行为。另请参见 on_test_model_train()。

返回类型

None

on_test_model_train()

在测试循环结束时调用。

测试循环默认会将 LightningModule 的 training 模式恢复到开始测试之前的状态。覆盖此钩子以更改行为。另请参见 on_test_model_eval()。

返回类型

None

on_test_start()

在测试开始时调用。

返回类型

None

on_train_batch_end(outputs, batch, batch_idx)

在训练循环中批次处理后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – 训练步骤(x)的输出

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

备注

这里的 outputs["loss"] 值将是相对于 accumulate_grad_batches 的损失归一化值，该损失从 training_step 返回。

返回类型

None

on_train_batch_start(batch, batch_idx)

在训练循环中，在该批次发生任何事情之前调用。

如果你在这里返回 -1，你将跳过当前轮次的剩余训练。

参数

• batch (Any) – 训练 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

返回类型

Optional[int, None]

on_train_end()

在训练结束前调用，在记录器实验关闭之前。

返回类型

None

on_train_epoch_end()

在训练循环的每个epoch结束时调用。

要在每个epoch结束时访问所有批处理输出，您可以将步骤输出缓存为 LightningModule 的一个属性，并在此钩子中访问它们：

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # free up the memory
        self.training_step_outputs.clear()

on_train_epoch_start()

在训练循环中，在epoch开始时调用。

返回类型

None

on_train_start()

在完整性检查后，训练开始时调用。

返回类型

None

on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx=0)

在批次验证循环后调用。

参数

• outputs (Union[Tensor, Mapping[str, Any], None]) – validation_step(x) 的输出

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_batch_start(batch, batch_idx, dataloader_idx=0)

在验证循环中，对该批次进行任何操作之前调用。

参数

• batch (Any) – 验证 DataLoader 返回的批处理数据。

• batch_idx (int) – 批次索引

• dataloader_idx (int) – 数据加载器的索引

返回类型

None

on_validation_end()

在验证结束时调用。

返回类型

None

on_validation_epoch_end()

在epoch结束时的验证循环中调用。

on_validation_epoch_start()

在epoch开始时验证循环中调用。

返回类型

None

on_validation_model_eval()

在验证循环开始时调用。

默认情况下，验证循环在开始前会在 LightningModule 上调用 .eval()。覆盖此钩子以更改行为。另请参见 on_validation_model_train()。

返回类型

None

on_validation_model_train()

当验证循环结束时调用。

默认情况下，验证循环会将 LightningModule 的 training 模式恢复到开始验证之前的状态。重写此钩子以更改行为。另请参见 on_validation_model_eval()。

返回类型

None

on_validation_model_zero_grad()

在进入验证循环之前，由训练循环调用来释放梯度。

返回类型

None

on_validation_start()

在验证开始时调用。

返回类型

None

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)

重写此方法以调整 Trainer 调用优化器的默认方式。

默认情况下，Lightning 会调用 step() 和 zero_grad()，如示例所示。当 Trainer(accumulate_grad_batches != 1) 时，在累积阶段不会调用此方法（和 zero_grad()）。手动优化时，重写此钩子没有好处。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Union[Optimizer, LightningOptimizer]) – 一个 PyTorch 优化器

• optimizer_closure (Optional[Callable[[], Any], None]) – 优化器闭包。此闭包必须执行，因为它包含了 training_step()、optimizer.zero_grad() 和 backward() 的调用。

示例:

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    # Add your custom logic to run directly before `optimizer.step()`

    optimizer.step(closure=optimizer_closure)

    # Add your custom logic to run directly after `optimizer.step()`

返回类型

None

optimizer_zero_grad(epoch, batch_idx, optimizer)

重写此方法以改变 optimizer.zero_grad() 的默认行为。

参数

• epoch (int) – 当前时期

• batch_idx (int) – 当前批次的索引

• optimizer (Optimizer) – 一个 PyTorch 优化器

示例:

# DEFAULT
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad()

# Set gradients to `None` instead of zero to improve performance (not required on `torch>=2.0.0`).
def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
    optimizer.zero_grad(set_to_none=True)

参见 torch.optim.Optimizer.zero_grad() 以了解上述示例的解释。

返回类型

None

optimizers(use_pl_optimizer=True)

返回训练期间使用的优化器。对于手动优化很有用。

参数

use_pl_optimizer (bool) – 如果 True，将会把优化器包装在 LightningOptimizer 中，以自动处理精度、分析和步数调用的计数，以进行适当的日志记录和检查点保存。它特别包装了 step 方法，不支持没有此方法的自定义优化器。

返回类型

Union[Optimizer, LightningOptimizer, _FabricOptimizer, List[Optimizer], List[LightningOptimizer], List[_FabricOptimizer]]

返回

一个单一的优化器，或者在存在多个优化器的情况下，一个优化器列表。

property output_chunk_length: Optional[int]

模型一次预测的时间步数。

返回类型

Optional[int, None]

parameters(recurse=True)

返回一个遍历模块参数的迭代器。

这通常传递给一个优化器。

参数

recurse (bool) – 如果为真，则生成此模块及其所有子模块的参数。否则，仅生成此模块的直接成员参数。

生成器

参数 – 模块参数

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

返回类型

Iterator[Parameter]

pred_batch_size: Optional[int]

pred_mc_dropout: Optional[bool]

pred_n: Optional[int]

pred_n_jobs: Optional[int]

pred_num_samples: Optional[int]

pred_roll_size: Optional[int]

predict_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定预测样本。

有关多个数据加载器的更多信息，请参阅此 部分。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• predict()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

返回

一个 torch.utils.data.DataLoader 或它们的序列，指定预测样本。

predict_likelihood_parameters: Optional[bool]

predict_step(batch, batch_idx, dataloader_idx=None)

执行预测步骤

批处理

Darts 的 InferenceDataset 的输出 - 包含 (过去的目标, 过去的协变量, 历史未来协变量, 未来协变量, 未来过去协变量, 输入时间序列, 预测开始时间步) 的元组

batch_idx

当前批次的批次索引

dataloader_idx

数据加载器索引

返回类型

Sequence[TimeSeries]

prepare_data()

使用此方法下载和准备数据。使用多个进程（分布式设置）下载和保存数据将导致数据损坏。Lightning 确保此方法仅在一个进程内调用，因此您可以安全地在其中添加下载逻辑。

警告

不要将状态设置到模型中（使用 setup 代替），因为这不会在每个设备上被调用

示例:

def prepare_data(self):
    # good
    download_data()
    tokenize()
    etc()

    # bad
    self.split = data_split
    self.some_state = some_other_state()

在分布式环境中，prepare_data 可以通过两种方式调用（使用 prepare_data_per_node）

1. 每个节点一次。这是默认设置，仅在 LOCAL_RANK=0 时调用。

2. 总共一次。仅在 GLOBAL_RANK=0 时调用。

示例:

# DEFAULT
# called once per node on LOCAL_RANK=0 of that node
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = True

# call on GLOBAL_RANK=0 (great for shared file systems)
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = False

在请求数据加载器之前调用此方法：

model.prepare_data()
initialize_distributed()
model.setup(stage)
model.train_dataloader()
model.val_dataloader()
model.test_dataloader()
model.predict_dataloader()

返回类型

None

prepare_data_per_node: bool

print(*args, **kwargs)

仅从进程 0 打印。在任何分布式模式下使用此功能以仅记录一次。

参数

• *args – 要打印的内容。与Python内置的print函数相同。

• **kwargs – 与Python内置的print函数相同。

示例:

def forward(self, x):
    self.print(x, 'in forward')

返回类型

None

register_backward_hook(hook)

在模块上注册一个反向钩子。

此函数已被 register_full_backward_hook() 取代，并且此函数的行为将在未来版本中发生变化。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_buffer(name, tensor, persistent=True)

向模块添加一个缓冲区。

这通常用于注册一个不应被视为模型参数的缓冲区。例如，BatchNorm 的 running_mean 不是参数，而是模块状态的一部分。缓冲区默认是持久的，并将与参数一起保存。可以通过将 persistent 设置为 False 来改变此行为。持久缓冲区和非持久缓冲区之间的唯一区别是后者不会成为此模块的 state_dict 的一部分。

缓冲区可以通过给定的名称作为属性访问。

参数

• name (str) – 缓冲区的名称。可以通过给定的名称从此模块访问缓冲区。

• tensor (Tensor or None) – 要注册的缓冲区。如果为 None ，则运行在缓冲区上的操作，例如 cuda ，将被忽略。如果为 None ，该缓冲区将 不 包含在模块的 state_dict 中。

• persistent (bool) – 缓冲区是否是此模块的 state_dict 的一部分。

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

返回类型

None

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

在模块上注册一个前向钩子。

钩子将在每次 forward() 计算输出后被调用。

如果 with_kwargs 是 False 或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward 。钩子可以修改输出。它可以就地修改输入，但这不会影响前向传播，因为这是在调用 forward() 之后调用的。钩子应具有以下签名:

hook(module, args, output) -> None or modified output

如果 with_kwargs 是 True，前向钩子将被传递给前向函数的 kwargs，并期望返回可能被修改的输出。钩子应具有以下签名:

hook(module, args, kwargs, output) -> None or modified output

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果 True，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之后触发。请注意，使用 register_module_forward_hook() 注册的全局 forward 钩子将在通过此方法注册的所有钩子之前触发。默认值：False

• with_kwargs (bool) – 如果 True，hook 将会传递给 forward 函数的 kwargs。默认值：False

• always_call (bool) – 如果 True ，无论在调用模块时是否引发异常，都将运行 hook 。默认值： False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

在模块上注册一个前向预钩子。

钩子将在每次调用 forward() 之前被调用。

如果 with_kwargs 为假或未指定，输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward。钩子可以修改输入。用户可以在钩子中返回一个元组或单个修改后的值。如果返回单个值（除非该值已经是元组），我们将把该值包装成一个元组。钩子应具有以下签名:

hook(module, args) -> None or modified input

如果 with_kwargs 为真，前向预钩子将传递给前向函数的 kwargs。如果钩子修改了输入，则应返回 args 和 kwargs。钩子应具有以下签名:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 forward_pre 钩子之后触发。请注意，通过 register_module_forward_pre_hook() 注册的全局 forward_pre 钩子将在所有通过此方法注册的钩子之前触发。默认值：False

• with_kwargs (bool) – 如果为真，hook 将被传递给 forward 函数的 kwargs。默认值：False

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)

在模块上注册一个反向钩子。

每次计算与模块相关的梯度时，钩子都会被调用，即当且仅当计算与模块输出相关的梯度时，钩子才会执行。钩子应具有以下签名:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input 和 grad_output 是包含相对于输入和输出的梯度的元组。钩子不应修改其参数，但可以选择返回一个新的相对于输入的梯度，该梯度将用于替代 grad_input 在后续计算中。grad_input 将仅对应于作为位置参数给出的输入，所有 kwarg 参数都将被忽略。grad_input 和 grad_output 中的条目将为 None 用于所有非 Tensor 参数。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入或输出，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上所有现有的 backward 钩子之后触发。请注意，通过 register_module_full_backward_hook() 注册的全局 backward 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

在模块上注册一个反向预钩子。

每次计算模块的梯度时，都会调用该钩子。钩子应具有以下签名:

hook(module, grad_output) -> tuple[Tensor] or None

The grad_output 是一个元组。钩子不应修改其参数，但它可以选择返回一个新的相对于输出的梯度，该梯度将用于替代 grad_output 在后续计算中。grad_output 中的条目对于所有非 Tensor 参数将为 None。

出于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。同样，调用者将接收模块前向函数返回的每个张量的视图。

警告

在使用反向钩子时，不允许就地修改输入，否则将引发错误。

参数

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – 如果为真，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之后触发。请注意，通过 register_module_full_backward_pre_hook() 注册的全局 backward_pre 钩子将在通过此方法注册的所有钩子之前触发。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)

注册一个在模块的 load_state_dict 调用后运行的后置钩子。

它应该有以下签名:

hook(module, incompatible_keys) -> None

module 参数是当前注册此钩子的模块，而 incompatible_keys 参数是一个由属性 missing_keys 和 unexpected_keys 组成的 NamedTuple。missing_keys 是一个包含缺失键的 str 的 list，而 unexpected_keys 是一个包含意外键的 str 的 list。

如果需要，可以就地修改给定的 incompatible_keys。

请注意，当使用 strict=True 调用 load_state_dict() 时执行的检查会受到钩子对 missing_keys 或 unexpected_keys 所做修改的影响，这是预期的。对任一组键的添加将导致在 strict=True 时抛出错误，而清除掉缺失和意外的键将避免错误。

返回

一个可以通过调用 handle.remove() 来移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_module(name, module)

别名 add_module()。

返回类型

None

register_parameter(name, param)

向模块添加一个参数。

该参数可以通过给定的名称作为属性访问。

参数

• name (str) – 参数的名称。可以使用给定的名称从此模块访问该参数。

• param (Parameter or None) – 要添加到模块的参数。如果为 None ，则运行在参数上的操作，例如 cuda ，将被忽略。如果为 None ，该参数将 不 包含在模块的 state_dict 中。

返回类型

None

register_state_dict_pre_hook(hook)

为 state_dict() 方法注册一个预钩子。

这些钩子将在调用 self 上的 state_dict 之前，使用参数 self、prefix 和 keep_vars 进行调用。注册的钩子可以用于在调用 state_dict 之前执行预处理。

requires_grad_(requires_grad=True)

更改是否应在此模块的参数上记录自动梯度操作。

此方法就地设置参数的 requires_grad 属性。

这种方法有助于冻结模块的一部分以进行微调或单独训练模型的一部分（例如，GAN训练）。

请参阅 locally-disable-grad-doc 以比较 .requires_grad_() 和几个可能与之混淆的类似机制。

参数

requires_grad (bool) – 是否应在此模块的参数上记录自动求导操作。默认值：True。

返回

自己

返回类型

Module

save_hyperparameters(*args, ignore=None, frame=None, logger=True)

将参数保存到 hparams 属性中。

参数

• args (Any) – dict、NameSpace 或 OmegaConf 的单个对象，或来自类 __init__ 的字符串名称或参数

• ignore (Union[Sequence[str], str, None]) – 一个参数名称或从类 __init__ 中忽略的参数名称列表

• frame (Optional[frame, None]) – 一个框架对象。默认是 None

• logger (bool) – 是否将超参数发送到记录器。默认值：True

示例:

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # manually assign arguments
...         self.save_hyperparameters('arg1', 'arg3')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class AutomaticArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # equivalent automatic
...         self.save_hyperparameters()
...     def forward(self, *args, **kwargs):
...         ...
>>> model = AutomaticArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg2": abc
"arg3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class SingleArgModel(HyperparametersMixin):
...     def __init__(self, params):
...         super().__init__()
...         # manually assign single argument
...         self.save_hyperparameters(params)
...     def forward(self, *args, **kwargs):
...         ...
>>> model = SingleArgModel(Namespace(p1=1, p2='abc', p3=3.14))
>>> model.hparams
"p1": 1
"p2": abc
"p3": 3.14

>>> from pytorch_lightning.core.mixins import HyperparametersMixin
>>> class ManuallyArgsModel(HyperparametersMixin):
...     def __init__(self, arg1, arg2, arg3):
...         super().__init__()
...         # pass argument(s) to ignore as a string or in a list
...         self.save_hyperparameters(ignore='arg2')
...     def forward(self, *args, **kwargs):
...         ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg3": 3.14

返回类型

None

set_extra_state(state)

设置包含在加载的 state_dict 中的额外状态。

此函数从 load_state_dict() 调用，以处理在 state_dict 中找到的任何额外状态。如果你的模块需要在 state_dict 中存储额外状态，请实现此函数以及相应的 get_extra_state()。

参数

state (dict) – 来自 state_dict 的额外状态

返回类型

None

set_mc_dropout(active)

set_predict_parameters(n, num_samples, roll_size, batch_size, n_jobs, predict_likelihood_parameters, mc_dropout)

在调用 trainer.predict() 之前从 TorchForecastingModel 设置，并在 self.on_predict_end() 时重置。

返回类型

None

setup(stage)

在 fit（训练 + 验证）、验证、测试或预测开始时调用。当你需要动态构建模型或调整模型时，这是一个很好的钩子。在使用 DDP 时，此钩子会在每个进程上调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

示例:

class LitModel(...):
    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

        # don't do this
        self.something = else

    def setup(self, stage):
        data = load_data(...)
        self.l1 = nn.Linear(28, data.num_classes)

返回类型

None

share_memory()

参见 torch.Tensor.share_memory_()。

返回类型

~T

state_dict(*args, destination=None, prefix='', keep_vars=False)

返回一个包含模块整体状态引用的字典。

参数和持久缓冲区（例如运行平均值）都包含在内。键是对应的参数和缓冲区名称。设置为 None 的参数和缓冲区不包含在内。

备注

返回的对象是一个浅拷贝。它包含对模块参数和缓冲区的引用。

警告

目前，state_dict() 也按顺序接受 destination、prefix 和 keep_vars 的位置参数。然而，这正在被弃用，未来版本将强制使用关键字参数。

警告

请避免使用参数 destination，因为它不是为最终用户设计的。

参数

• destination (dict, optional) – 如果提供了，模块的状态将被更新到字典中，并返回相同的对象。否则，将创建并返回一个 OrderedDict。默认值：None。

• prefix (str, optional) – 在 state_dict 中组合键时添加到参数和缓冲区名称的前缀。默认值：''。

• keep_vars (bool, optional) – 默认情况下，状态字典中返回的 Tensor 是与自动求导分离的。如果设置为 True，则不会执行分离操作。默认值：False。

返回

包含模块整个状态的字典

返回类型

dict

示例:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

property strict_loading: bool

确定 Lightning 如何使用 .load_state_dict(…, strict=model.strict_loading) 加载此模型。

返回类型

bool

property supports_probabilistic_prediction: bool

返回类型

bool

teardown(stage)

在 fit（训练 + 验证）、验证、测试或预测结束时调用。

参数

stage (str) – 要么 'fit'，要么 'validate'，要么 'test'，要么 'predict'

返回类型

None

test_dataloader()

一个可迭代对象或多个可迭代对象的集合，指定测试样本。

有关多个数据加载器的更多信息，请参阅此 部分。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• test()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要测试数据集和 test_step() ，你不需要实现这个方法。

返回类型

Any

test_step(*args, **kwargs)

对测试集中的单个批次数据进行操作。在这一步中，您通常会生成示例或计算任何感兴趣的内容，例如准确性。

参数

• batch – 你的数据可迭代对象的输出，通常是一个 DataLoader。

• batch_idx – 这一批的索引。

• dataloader_idx – 生成此批次的 dataloader 的索引。（仅在使用多个 dataloader 时）

返回类型

Union[Tensor, Mapping[str, Any], None]

返回

• Tensor - 损失张量

• dict - 一个字典。可以包含任何键，但必须包含键 'loss'。

• None - 跳到下一个批次。

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...

# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

示例:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

如果你传入多个测试数据加载器，test_step() 将会有一个额外的参数。我们建议将默认值设置为0，这样你可以快速在单个和多个数据加载器之间切换。

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

备注

如果你不需要测试，你就不需要实现这个方法。

备注

当调用 test_step() 时，模型已处于评估模式，并且已禁用 PyTorch 梯度。在测试轮次结束时，模型返回到训练模式，并且启用梯度。

to(*args, **kwargs)

参见 torch.nn.Module.to()。

返回类型

Self

to_dtype(dtype)

将模块精度（默认为 float32）转换为另一种精度。

to_empty(*, device, recurse=True)

将参数和缓冲区移动到指定设备，而不复制存储。

参数

• device (torch.device) – 此模块中参数和缓冲区的所需设备。

• recurse (bool) – 子模块的参数和缓冲区是否应递归地移动到指定设备。

返回

自己

返回类型

Module

to_onnx(file_path, input_sample=None, **kwargs)

将模型保存为 ONNX 格式。

参数

• file_path (Union[str, Path]) – 应保存 onnx 模型的文件路径。

• input_sample (Optional[Any, None]) – 用于跟踪的输入。默认值：None（使用 self.example_input_array）

• **kwargs – 将被传递给 torch.onnx.export 函数。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)

model = SimpleModel()
input_sample = torch.randn(1, 64)
model.to_onnx("export.onnx", input_sample, export_params=True)

返回类型

None

to_torchscript(file_path=None, method='script', example_inputs=None, **kwargs)

默认情况下，将整个模型编译为 ScriptModule。如果你想使用追踪，请提供参数 method='trace' 并确保提供了 example_inputs 参数，或者模型设置了 example_input_array。如果你想自定义被脚本的模块，你应该重写这个方法。如果你想要返回多个模块，我们建议使用字典。

参数

• file_path (Union[str, Path, None]) – 保存 torchscript 的路径。默认：None（不保存文件）。

• method (Optional[str, None]) – 是否使用 TorchScript 的 script 或 trace 方法。默认值：’script’

• example_inputs (Optional[Any, None]) – 当方法设置为’trace’时，用于进行追踪的输入。默认值：None（使用 example_input_array）

• **kwargs – 将传递给 torch.jit.script() 或 torch.jit.trace() 函数的额外参数。

备注

• 需要实现 forward() 方法。

• 导出的脚本将被设置为评估模式。

• 建议您安装最新支持版本的 PyTorch 以无限制地使用此功能。另请参阅 torch.jit 文档以了解支持的功能。

示例:

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(in_features=64, out_features=4)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

model = SimpleModel()
model.to_torchscript(file_path="model.pt")

torch.jit.save(model.to_torchscript(
    file_path="model_trace.pt", method='trace', example_inputs=torch.randn(1, 64))
)

返回类型

Union[ScriptModule, Dict[str, ScriptModule]]

返回

无论 file_path 是否定义，此 LightningModule 都可以作为 torchscript 使用。

toggle_optimizer(optimizer)

确保在训练步骤中仅计算当前优化器的参数的梯度，以防止在多优化器设置中出现悬空梯度。

它与 untoggle_optimizer() 一起工作，以确保 param_requires_grad_state 被正确重置。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要切换的优化器。

返回类型

None

train(mode=True)

设置模块为训练模式。

这仅对某些模块有影响。如果受到影响，请参阅特定模块的文档以了解其在训练/评估模式下的行为详情，例如 Dropout, BatchNorm 等。

参数

mode (bool) – 是否设置训练模式（True）或评估模式（False）。默认值：True。

返回

自己

返回类型

Module

train_criterion_reduction: Optional[str]

train_dataloader()

指定训练样本的可迭代对象或可迭代对象的集合。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

对于数据处理，请使用以下模式：

• 在 prepare_data() 中下载

• 在 setup() 中处理和分割

然而，上述内容仅对分布式处理是必要的。

警告

不要在 prepare_data 中分配状态

• fit()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

返回类型

Any

property trainer: Trainer

返回类型

Trainer

training: bool

training_step(train_batch, batch_idx)

执行训练步骤

返回类型

Tensor

transfer_batch_to_device(batch, device, dataloader_idx)

如果你的 DataLoader 返回的是封装在自定义数据结构中的张量，请重写此钩子。

以下列出的数据类型（以及它们的任意嵌套）都是开箱即用的：

• torch.Tensor 或任何实现了 .to(…) 的对象

• 列表

• 字典

• 元组

对于其他情况，你需要定义数据如何移动到目标设备（CPU、GPU、TPU等）。

备注

此钩子应仅传输数据，不应修改数据，也不应将数据移动到除作为参数传入的设备之外的任何其他设备（除非你知道自己在做什么）。要检查此钩子的当前执行状态，可以使用 self.trainer.training/testing/validating/predicting，以便可以根据你的需求添加不同的逻辑。

参数

• batch (Any) – 一批需要转移到新设备的数据。

• device (device) – 在 PyTorch 中定义的目标设备。

• dataloader_idx (int) – 批次所属的数据加载器的索引。

返回类型

Any

返回

对新设备上数据的引用。

示例:

def transfer_batch_to_device(self, batch, device, dataloader_idx):
    if isinstance(batch, CustomBatch):
        # move all tensors in your custom data structure to the device
        batch.samples = batch.samples.to(device)
        batch.targets = batch.targets.to(device)
    elif dataloader_idx == 0:
        # skip device transfer for the first dataloader or anything you wish
        pass
    else:
        batch = super().transfer_batch_to_device(batch, device, dataloader_idx)
    return batch

参见

• move_data_to_device()

• apply_to_collection()

type(dst_type)

参见 torch.nn.Module.type()。

返回类型

Self

unfreeze()

解冻所有参数以进行训练。

model = MyLightningModule(...)
model.unfreeze()

返回类型

None

untoggle_optimizer(optimizer)

重置通过 toggle_optimizer() 切换的所需梯度的状态。

参数

optimizer (Union[Optimizer, LightningOptimizer]) – 要取消选择的优化器。

返回类型

None

val_criterion_reduction: Optional[str]

val_dataloader()

一个可迭代对象或多个可迭代对象的集合，用于指定验证样本。

有关多个数据加载器的更多信息，请参阅此 部分。

你返回的数据加载器不会被重新加载，除非你将 :paramref:`~pytorch_lightning.trainer.trainer.Trainer.reload_dataloaders_every_n_epochs` 设置为一个正整数。

建议所有数据下载和准备工作在 prepare_data() 中进行。

• fit()

• validate()

• prepare_data()

• setup()

备注

Lightning 尝试为分布式和任意硬件添加正确的采样器。无需自行设置。

备注

如果你不需要验证数据集和 validation_step()，你不需要实现这个方法。

返回类型

Any

validation_step(val_batch, batch_idx)

执行验证步骤

返回类型

Tensor

xpu(device=None)

将所有模型参数和缓冲区移动到XPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块将在 XPU 上进行优化，则应在构建优化器之前调用此方法。

备注

此方法就地修改模块。

参数

device (int, optional) – 如果指定，所有参数将被复制到该设备

返回

自己

返回类型

Module

zero_grad(set_to_none=True)

重置所有模型参数的梯度。

更多上下文请参见 torch.optim.Optimizer 下的类似函数。

参数

set_to_none (bool) – 而不是设置为零，将梯度设置为 None。详情请参阅 torch.optim.Optimizer.zero_grad()。

返回类型

None

darts.models.forecasting.pl_forecasting_module.io_processor(forward)

对 PLForecastingModule.forward 应用一些输入/输出处理。请注意，此包装器必须添加到 PLForecastinModule 的每个子类的 forward 方法中。以下是如何添加装饰器的示例：

```python

@io_processor def forward(self, *args, **kwargs)

pass

toctree 是一个 reStructuredText 指令 ，这是一个非常多功能的标记。指令可以有参数、选项和内容。

上一页

N-线性

下一页

Facebook Prophet