PyTorch-Lightning:trining_step的自动优化

2024-04-12 09:44

本文主要是介绍PyTorch-Lightning:trining_step的自动优化,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

文章目录

  • PyTorch-Lightning:trining_step的自动优化
      • 总结:
    • class _ AutomaticOptimization()
      • def run
      • def _make_closure
      • def _training_step
        • class ClosureResult():
          • def from_training_step_output
      • class Closure

PyTorch-Lightning:trining_step的自动优化

使用PyTorch-Lightning时,在trining_step定义损失,在没有定义损失,没有任何返回的情况下没有报错,在定义一个包含loss的多个元素字典返回时,也可以正常训练,那么到底lightning是怎么完成训练过程的。

总结:

在自动优化中,training_step必须返回一个tensor或者dict或者None(跳过),对于简单的使用,在training_step可以return一个tensor会作为Loss回传,也可以return一个字典,其中必须包括key"loss",字典中的"loss"会提取出来作为Loss回传,具体过程主要包含在lightning\pytorch\loop\sautomatic.py中的_ AutomaticOptimization()类。

在这里插入图片描述

class _ AutomaticOptimization()

实现自动优化(前向,梯度清零,后向,optimizer step)

在training_epoch_loop中会调用这个类的run函数。

def run

首先通过 _make_closure得到一个closure,详见def _make_closure,最后返回一个字典,如果我们在training_step只return了一个loss tensor则字典只有一个’loss’键值对,如果return了一个字典,则包含其他键值对。

可以看到调用了_ optimizer_step,_ optimizer_step经过层层调用,最后会调用torch默认的optimizer.zero_grad,而我们通过 make_closure得到的closure作为参数传入,具体而言是调用了closure类的_ call __()方法。

def run(self, optimizer: Optimizer, batch_idx: int, kwargs: OrderedDict) -> _OUTPUTS_TYPE:closure = self._make_closure(kwargs, optimizer, batch_idx)if (# when the strategy handles accumulation, we want to always call the optimizer stepnot self.trainer.strategy.handles_gradient_accumulation and self.trainer.fit_loop._should_accumulate()):# For gradient accumulation# -------------------# calculate loss (train step + train step end)# -------------------# automatic_optimization=True: perform ddp sync only when performing optimizer_stepwith _block_parallel_sync_behavior(self.trainer.strategy, block=True):closure()# ------------------------------# BACKWARD PASS# ------------------------------# gradient update with accumulated gradientselse:self._optimizer_step(batch_idx, closure)result = closure.consume_result()if result.loss is None:return {}return result.asdict()

def _make_closure

创建一个closure对象,来捕捉给定的参数并且运行’training_step’和可选的其他如backword和zero_grad函数

比较重要的是step_fn,在这里调用了_training_step,得到的是一个存储我们在定义模型时重写的training step的输出所构成ClosureResult数据类。具体见def _training_step

def _make_closure(self, kwargs: OrderedDict, optimizer: Optimizer, batch_idx: int) -> Closure:step_fn = self._make_step_fn(kwargs)backward_fn = self._make_backward_fn(optimizer)zero_grad_fn = self._make_zero_grad_fn(batch_idx, optimizer)return Closure(step_fn=step_fn, backward_fn=backward_fn, zero_grad_fn=zero_grad_fn)

def _training_step

通过hook函数实现真正的训练step,返回一个存储training step输出的ClosureResult数据类。

将我们在定义模型时定义的lightning.pytorch.core.LightningModule.training_step的输出作为参数传入存储容器class ClosureResult的from_training_step_output方法,见class Closure

class ClosureResult():

一个数据类,包含closure_loss,loss,extra

    closure_loss: Optional[Tensor]loss: Optional[Tensor] = field(init=False, default=None)extra: Dict[str, Any] = field(default_factory=dict)
def from_training_step_output

一个类方法,如果我们在training_step定义的返回是一个字典,则我们会将key值为"loss"的value赋值给closure_loss,而其余的键值对赋值给extra字典,如果返回的既不是包含"loss"的字典也不是tensor,则会报错。当我们在training_step不设定返回,则自然为None,但是不会报错。

class Closure

闭包是将外部作用域中的变量绑定到对这些变量进行计算的函数变量,而不将它们明确地作为输入。这样做的好处是可以将闭包传递给对象,之后可以像函数一样调用它,但不需要传入任何参数。

在lightning,用于自动优化的Closure类将training_step和backward, zero_grad三个基本的闭包结合在一起。

这个Closure得到training循环中的结果之后传入torch.optim.Optimizer.step。

参数:

  • step_fn: 这里是一个存储了training step输出的ClosureResult数据类,见def _training_step
  • backward_fn: 梯度回传函数
  • zero_grad_fn: 梯度清零函数

按照顺序,会先检查得到loss,之后调用梯度清零函数,最后调用梯度回传函数

class Closure(AbstractClosure[ClosureResult]):warning_cache = WarningCache()def __init__(self,step_fn: Callable[[], ClosureResult],backward_fn: Optional[Callable[[Tensor], None]] = None,zero_grad_fn: Optional[Callable[[], None]] = None,):super().__init__()self._step_fn = step_fnself._backward_fn = backward_fnself._zero_grad_fn = zero_grad_fn@override@torch.enable_grad()def closure(self, *args: Any, **kwargs: Any) -> ClosureResult:step_output = self._step_fn()if step_output.closure_loss is None:self.warning_cache.warn("`training_step` returned `None`. If this was on purpose, ignore this warning...")if self._zero_grad_fn is not None:self._zero_grad_fn()if self._backward_fn is not None and step_output.closure_loss is not None:self._backward_fn(step_output.closure_loss)return step_output@overridedef __call__(self, *args: Any, **kwargs: Any) -> Optional[Tensor]:self._result = self.closure(*args, **kwargs)return self._result.loss

这篇关于PyTorch-Lightning:trining_step的自动优化的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!


原文地址:
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.chinasem.cn/article/896785

相关文章

Flask 验证码自动生成的实现示例

《Flask验证码自动生成的实现示例》本文主要介绍了Flask验证码自动生成的实现示例,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习... 目录生成图片以及结果处理验证码蓝图html页面展示想必验证码大家都有所了解,但是可以自己定义图片验证码

pytorch+torchvision+python版本对应及环境安装

《pytorch+torchvision+python版本对应及环境安装》本文主要介绍了pytorch+torchvision+python版本对应及环境安装,安装过程中需要注意Numpy版本的降级,... 目录一、版本对应二、安装命令(pip)1. 版本2. 安装全过程3. 命令相关解释参考文章一、版本对

Python Excel实现自动添加编号

《PythonExcel实现自动添加编号》这篇文章主要为大家详细介绍了如何使用Python在Excel中实现自动添加编号效果,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下... 目录1、背景介绍2、库的安装3、核心代码4、完整代码1、背景介绍简单的说,就是在Excel中有一列h=会有重复

Java嵌套for循环优化方案分享

《Java嵌套for循环优化方案分享》介绍了Java中嵌套for循环的优化方法,包括减少循环次数、合并循环、使用更高效的数据结构、并行处理、预处理和缓存、算法优化、尽量减少对象创建以及本地变量优化,通... 目录Java 嵌套 for 循环优化方案1. 减少循环次数2. 合并循环3. 使用更高效的数据结构4

Springboot的自动配置是什么及注意事项

《Springboot的自动配置是什么及注意事项》SpringBoot的自动配置(Auto-configuration)是指框架根据项目的依赖和应用程序的环境自动配置Spring应用上下文中的Bean... 目录核心概念:自动配置的关键特点:自动配置工作原理:示例:需要注意的点1.默认配置可能不适合所有场景

从零教你安装pytorch并在pycharm中使用

《从零教你安装pytorch并在pycharm中使用》本文详细介绍了如何使用Anaconda包管理工具创建虚拟环境,并安装CUDA加速平台和PyTorch库,同时在PyCharm中配置和使用PyTor... 目录背景介绍安装Anaconda安装CUDA安装pytorch报错解决——fbgemm.dll连接p

pycharm远程连接服务器运行pytorch的过程详解

《pycharm远程连接服务器运行pytorch的过程详解》:本文主要介绍在Linux环境下使用Anaconda管理不同版本的Python环境,并通过PyCharm远程连接服务器来运行PyTorc... 目录linux部署pytorch背景介绍Anaconda安装Linux安装pytorch虚拟环境安装cu

Java中实现订单超时自动取消功能(最新推荐)

《Java中实现订单超时自动取消功能(最新推荐)》本文介绍了Java中实现订单超时自动取消功能的几种方法,包括定时任务、JDK延迟队列、Redis过期监听、Redisson分布式延迟队列、Rocket... 目录1、定时任务2、JDK延迟队列 DelayQueue(1)定义实现Delayed接口的实体类 (

shell脚本自动删除30天以前的文件(最新推荐)

《shell脚本自动删除30天以前的文件(最新推荐)》该文章介绍了如何使用Shell脚本自动删除指定目录下30天以前的文件,并通过crontab设置定时任务,此外,还提供了如何使用Shell脚本删除E... 目录shell脚本自动删除30天以前的文件linux按照日期定时删除elasticsearch索引s

Deepseek使用指南与提问优化策略方式

《Deepseek使用指南与提问优化策略方式》本文介绍了DeepSeek语义搜索引擎的核心功能、集成方法及优化提问策略,通过自然语言处理和机器学习提供精准搜索结果,适用于智能客服、知识库检索等领域... 目录序言1. DeepSeek 概述2. DeepSeek 的集成与使用2.1 DeepSeek API