自动混合精度(amp)与梯度缩放

torch.cuda.amp包提供便利的混合精度操作方法,即对应不同的层采取不一样的浮点数精度(torch.float32、torch.float16、torch.bfloat16),以减小模型计算时间与内存资源为目的。一般生产情况下, torch.autocast 与 torch.cuda.amp.GradScaler配合完成自动混合精度的功能。

在训练loop中官方推荐方式如下:

for epoch in range(xxx): 
	for input, target in zip(data, targets):
		with torch.autocast(device_type=device, dtype=torch.float16):
			#1 step
			forward pass

			#2 step
			compute loss

		#3 step
		#Scales loss. Calls ``backward()`` on scaled loss to create scaled gradients.
		#Gradient scaling helps prevent gradients with small magnitudes from flushing to zero (“underflowing”) when training with mixed precision
		scaler.scale(loss).backward()
		
		#4 step
		# Unscales the gradients of optimizer's assigned parameters in-place
		scaler.unscale_(opt)
			
		#5 step
		# Since the gradients of optimizer's assigned parameters are now unscaled, clips as usual.
		torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=0.1)
		  
		#6 step
		# ``scaler.step()`` first unscales the gradients of the optimizer's assigned parameters.
		# If these gradients do not contain ``inf``s or ``NaN``s, optimizer.step() is then called,otherwise, optimizer.step() is skipped.
		scaler.step(opt)

		#7 step
		# Updates the scale for next iteration.
		scaler.update()
		
		#8 step
		opt.zero_grad()

### YOLO 中的自动混合精度训练(AMP)实现方法 #### 背景说明 自动混合精度训练(Automatic Mixed Precision, AMP)是一种通过动态调整模型权重和梯度计算的数值精度来加速深度学习模型训练的技术。它能够在保持模型收敛性能的同时减少内存占用并提升训练速度[^1]。 在 PyTorch 的支持下,AMP 可以轻松集成到基于 PyTorch 构建的目标检测框架中,例如 YOLO 系列模型。以下是具体实现方式: --- #### 方法一:命令行配置 如果使用的是脚本化工具或预定义的训练流程,可以通过命令行参数控制是否启用 AMP 训练。默认情况下,许多现代版本的 YOLO 已经启用了 AMP 功能。要禁用该功能,可以在启动训练时添加 `--unamp` 参数[^3]: ```bash python train.py --data dataset.yaml --epochs 100 --batch-size 16 --unamp ``` 此操作会告诉程序不使用 AMP 进行训练。反之,如果不传递 `--unamp` 参数,则默认开启 AMP 支持。 --- #### 方法二:代码级修改 对于手动编写训练逻辑的情况,可以直接在代码层面引入 PyTorch 提供的 `torch.cuda.amp` 模块[^2]。以下是一个典型的实现示例: ```python import torch from torch.cuda.amp import GradScaler, autocast # 初始化缩放器用于梯度放大 scaler = GradScaler() def train_step(model, optimizer, data_loader, device): model.train() for images, targets in data_loader: images = list(image.to(device) for image in images) targets = [{k: v.to(device) for k, v in t.items()} for t in targets] # 使用 autocast 上下文管理器进行前向传播 with autocast(): outputs = model(images, targets) loss_dict = outputs['losses'] losses = sum(loss for loss in loss_dict.values()) # 清除之前的梯度 optimizer.zero_grad() # 缩放损失以便反向传播更稳定 scaler.scale(losses).backward() # 更新权重 scaler.step(optimizer) # 更新缩放器状态 scaler.update() ``` 上述代码片段展示了如何利用 `autocast` 和 `GradScaler` 来完成 AMP 的核心部分。其中: - **`autocast()`**:负责自动切换张量的操作精度,在适当的地方采用 FP16 或 FP32。 - **`GradScaler`**:解决由于低精度运算可能导致的梯度过小问题,从而提高稳定性[^4]。 --- #### 方法三:YOLO 配置文件设置 某些高级版的 YOLO 实现允许通过 YAML 配置文件指定是否启用 AMP。通常会在配置文件中增加如下字段: ```yaml train: amp: true # 是否启用自动混合精度训练 ``` 当设为 `true` 时,默认打开 AMP;而将其改为 `false` 则可关闭这一特性。 --- #### 总结 无论是通过命令行选项、代码内部调优还是借助外部配置文件的方式,都可以灵活地控制 YOLO 模型中的 AMP 行为。这些技术手段不仅简化了开发者的调试过程,还显著提升了 GPU 设备上的推理效率资源利用率. ---
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值