Python PyTorch checkpoint_sequential用法及代码示例

本文简要介绍python语言中 torch.utils.checkpoint.checkpoint_sequential 的用法。

用法:

torch.utils.checkpoint.checkpoint_sequential(functions, segments, input, **kwargs)

参数：

• functions-A torch.nn.Sequential 或模块或函数列表(包括模型)以顺序运行。

• segments-在模型中创建的块数

• input-输入到functions 的张量

• preserve_rng_state(bool,可选的,默认=真) -在每个检查点期间省略存储和恢复 RNG 状态。

返回：

在 *inputs 上按顺序运行 functions 的输出

用于检查点顺序模型的辅助函数。

顺序模型按顺序(按顺序)执行模块/函数列表。因此，我们可以将这样的模型划分为各个段，并对每个段进行检查点。除最后一个段外，所有段都将以 torch.no_grad() 方式运行，即不存储中间激活。每个检查点段的输入将被保存，以便在反向传递中重新运行该段。

请参阅 checkpoint() 了解检查点的工作原理。

警告

检查点当前仅支持 torch.autograd.backward() ，并且仅在其inputs 参数未通过时才支持。不支持 torch.autograd.grad() 。

示例

>>> model = nn.Sequential(...)
>>> input_var = checkpoint_sequential(model, chunks, input_var)