美文网首页
PyTorch-混合精度训练(Mixed-Precisioin

PyTorch-混合精度训练(Mixed-Precisioin

作者: 侠之大者_7d3f | 来源:发表于2023-06-30 22:14 被阅读0次

混合精度训练介绍

Mixed-Precision Training是指在深度学习AI模型训练过程中不同的层Layer采用不同的数据精度进行训练, 最终使得训练过程中的资源消耗(GPU显存,GPU 算力)降低, 同时保证训练可收敛,模型精度与高精度FP32的结果接近。


CNN ResNet 混合精度训练

  1. 导入torch.cuda.amp package
    由于CNN训练要求大量算力, 因此一般混合精度需要使用 NVIDIA Automatic Mixed Precision (AMP)包, NVIDIA的AMP以及集成到了Pyorch, 因此直接调用 torch.cuda.amp APIs.

混合精度主要用到 Loss-Scaling (损失缩放) + Auto-cast (自动精度选择/转换)

# Mixed-Precision Training
from torch.cuda.amp.grad_scaler import GradScaler
from torch.cuda.amp.autocast_mode import autocast

# 实例化一个GradeScaler对象
scaler = GradScaler()
  1. 对Training Loop进行修改, 修改2个地方
  • 添加autocast(): autocast是一个Python context Manager, autocast 作用区域的代码在运行的时候会跟据OP的类型,自动转换为预定义好的低精度类型 (比如FP16)
    *注意: autocast一般作用的代码区域为 Forward, Backward 阶段需要指定autocast, 因此在Forward阶段不同的layer (op)以及被设置了各自的精度模式, 在Backward阶段,采用和Forward相同的精度进行计算。

  • 添加GradeScalar: GradeScalar的目的是对权重的梯度矩阵值G_W进行缩放(扩大化), 因为一般情况下G_W的值非常小,如果采用低精度类型 (FP16),则导致G_W下溢underflow. 解决方法之一就是希望将 G_W进行缩放变大; 由于Loss函数导数具有线性性质, 因此也可以对 Loss进行缩放,实际等价于对梯度值进行了放大。

 for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        # move data to the same device as model
        images = images.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)

        # compute output
        with autocast(enabled=args.mixed_precision, dtype=torch.float16):
                output = model(images)
                loss = criterion(output, target)

        # measure accuracy and record loss
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        # losses.update(loss.item(), images.size(0))
        top1.update(acc1[0], images.size(0))
        top5.update(acc5[0], images.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        # loss.backward()
        # optimizer.step()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
         ...

为了便于观察训练过程, 在代码中添加了Pytorch Profiler进行可视化:
为了说明情况, 可以只跑少数的几个batch即可

image.png image.png image.png

对应的CUDA kernel函数, FP16类型的


image.png

采用BFLOAT16进行混合精度

方法很简单,在autocast的dtype设置为torch.bfloat16。 除此之外,需要采用支持BFLOAT16类型的计算设备(TPU, >=NVIDIA Ampere/Volta 架构的GPU, 比如NVIDIA V100, A100, RTX 30/40系列)

with autocast(enabled=args.mixed_precision, dtype=torch.bfloat16):
                output = model(images)
                loss = criterion(output, target)
image.png
image.png

相关文章

  • 混合精度训练

    原文来自于机器学习算法与自然语言处理公众号 混合精度训练 《MIXED PRECISION TRAINING》这篇...

  • 混合精度训练

    减少所需显存。减少了IO. 深度学习训练使用16bit表示/运算正逐渐成为主流。 低精度带来了性能、功耗优势,但需...

  • 【学习tensorflow2】有用的API汇总

    动态显存分配 随机数种子: 为 混合精度 中断训练与继续训练 日志可视化

  • 全网最全-混合精度训练原理

    如图所示在网络模型训练阶段, 激活函数的梯度分布式中,有67%的精度小于2^24,直接使用FP16进行表示会截断下...

  • Pytroch使用amp进行混合精度训练

    简介 AMP:Automatic mixed precision,自动混合精度,可以在神经网络推理过程中,针对不同...

  • Pytorch-训练模型

    上一节已经介绍了怎么定义网络,计算损失函数以及更新权重,这一节将介绍如何开始训练模型。 What about da...

  • PyTorch自动混合精度

    https://zhuanlan.zhihu.com/p/165152789

  • 分布式训练、低比特训练、混合精度训练

    暂时不更新了,工作面试先不准备,赶紧先安心把MM的论文写好,把毕设论文写好再搞别的。剩余几天只有两件事,两篇论文!...

  • 过拟合及其解决方式

    过拟合(overfit)是指模型过度拟合训练数据,即在训练数据上,精度表现的特别高;在测试数据上,精度远不如在训练...

  • 验证方法

    留出验证集 训练模型 绘制训练损失和验证损失 绘制训练精度和验证精度 可以看出在第9轮之后开始过拟合。 重新开始训...

网友评论

      本文标题:PyTorch-混合精度训练(Mixed-Precisioin

      本文链接:https://www.haomeiwen.com/subject/zystudtx.html