yolo ultralytics之amp（Automatic Mixed Precision，自动混合精度）

原创已于 2025-10-17 14:46:14 修改 · 640 阅读

20 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #深度学习 #计算机视觉

于 2025-10-17 14:45:08 首次发布

在Ultralytics YOLO中，setup_train方法中的amp指的是自动混合精度（Automatic Mixed Precision）训练。详细解释：

1. AMP基本概念

# setup_train中的AMP设置
def setup_train(self, weights, data, device, ...):
    self.amp = self.amp and check_version('0.0.0', 
                  torch.__version__, '1.10.0')  # 检查版本兼容性

AMP = Automatic Mixed Precision，是PyTorch提供的自动混合精度训练功能。

2. AMP的作用和原理

为什么要使用AMP？

# 传统训练（FP32）
with torch.no_grad():
    predictions = model(images)  # 全部使用FP32精度

# AMP训练（混合精度）
with torch.cuda.amp.autocast():
    predictions = model(images)  # 自动选择FP16/FP32

AMP的工作方式：

前向传播: 使用FP16（半精度）计算，加快速度
反向传播: 使用FP16计算梯度
优化器更新: 使用FP32（单精度）更新参数，保持数值稳定性

3. AMP在Ultralytics中的配置

自动检测和设置

class BaseTrainer:
    def setup_train(self, ...):
        # AMP自动配置
        self.amp = (self.amp and 
                   device.type != 'cpu' and 
                   check_version(torch.__version__, '1.10.0'))
        
        if self.amp:
            self.scaler = torch.cuda.amp.GradScaler(enabled=True)
            LOGGER.info('AMP enabled ✅')
        else:
            self.scaler = None

训练参数影响

# 在训练命令或配置文件中
amp: True  # 启用自动混合精度
# 或者
yolo train model=yolov8n.pt data=coco128.yaml amp=True

4. AMP带来的好处

性能提升对比

# 启用AMP的效果
训练速度: +30% ~ 50% 提升
GPU内存使用: -40% ~ 60% 减少
精度损失: 通常 < 1% mAP

实际训练中的体现

def train_one_epoch(self):
    for batch in self.train_loader:
        # AMP训练块
        with torch.cuda.amp.autocast(enabled=self.amp):
            predictions = self.model(batch['img'])
            loss = self.criterion(predictions, batch)
        
        # 梯度缩放和更新
        self.scaler.scale(loss).backward()
        self.scaler.step(self.optimizer)
        self.scaler.update()

5. AMP的启用条件

自动检查逻辑

def check_amp(model):
    """检查模型是否支持AMP"""
    enabled = True
    # 检查设备支持
    if device.type == 'cpu':
        enabled = False
        LOGGER.warning('AMP not supported on CPU')
    
    # 检查PyTorch版本
    if not check_version(torch.__version__, '1.10.0'):
        enabled = False
        LOGGER.warning('AMP requires PyTorch >= 1.10.0')
    
    # 检查模型兼容性
    if hasattr(model, 'amp') and not model.amp:
        enabled = False
    
    return enabled

6. AMP训练流程详解

class Trainer:
    def __init__(self):
        self.amp = True  # 默认启用
        self.setup_amp()
    
    def setup_amp(self):
        """设置AMP相关组件"""
        if self.amp:
            # 1. 创建梯度缩放器
            self.scaler = torch.cuda.amp.GradScaler(
                enabled=True,
                init_scale=65536.0,  # 初始缩放因子
                growth_interval=2000  # 缩放因子更新间隔
            )
        else:
            self.scaler = None
    
    def train_step(self, batch):
        """AMP训练步骤"""
        # 清空梯度
        self.optimizer.zero_grad()
        
        # AMP前向传播
        with torch.cuda.amp.autocast(enabled=self.amp):
            predictions = self.model(batch['img'])
            loss, loss_items = self.criterion(predictions, batch)
        
        # AMP反向传播和优化
        if self.amp:
            self.scaler.scale(loss).backward()
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0)
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            # 普通训练
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0)
            self.optimizer.step()

7. AMP的注意事项

兼容性问题

# 某些操作在AMP中需要特别注意
if self.amp:
    # 可能需要手动处理某些层
    with torch.cuda.amp.autocast():
        # 大部分层自动处理
        x = self.conv_layers(x)
        # 某些特殊操作可能需要FP32
        x = x.float() if x.dtype == torch.float16 else x

调试和监控

# 监控AMP训练状态
if self.amp:
    scale = self.scaler.get_scale()
    if scale < 1.0:
        LOGGER.warning(f'Gradient scale dropped to {scale}, possible numerical issues')