PyTorch 浮点数精度全景:从 float16/bfloat16 到 float64 及混合精度实战
PyTorch 在深度学习中提供了多种 IEEE 754 二进制浮点格式的支持,包括半精度(float16
)、Brain‑float(bfloat16
)、单精度(float32
)和双精度(float64
),并通过统一的 torch.dtype
接口进行管理citeturn0search0turn0search3。用户可利用 torch.finfo
查询各类型的数值极限(如最大值、最小值、机器 ε 等),通过 torch.set_default_dtype
/torch.get_default_dtype
设置或获取全局默认浮点精度,并使用 torch.promote_types
控制运算中的类型提升规则citeturn0search2turn0search4。在现代 GPU 上,PyTorch 提供了 torch.amp.autocast
和 torch.amp.GradScaler
等自动混合精度(AMP)工具,能够在保证数值稳定性的前提下,大幅提升训练速度和降低显存占用citeturn0search6turn0search11。
PyTorch 浮点类型对比
类型 (torch.dtype ) | 别名 | 位宽 | 符号位 | 指数位 | 尾数位 (显式) | 有效精度 (含隐含位) | 典型用途 |
---|---|---|---|---|---|---|---|
torch.float16 | torch.half | 16 | 1 | 5 | 10 | 11 位 (~3.3 十进制位) | 推理加速,对精度要求不高的场景 |
torch.bfloat16 | — | 16 | 1 | 8 | 7 | 8 位 (~2.4 十进制位) | 大规模训练(TPU、支持 BF16 的 GPU) |
torch.float32 | torch.float | 32 | 1 | 8 | 23 | 24 位 (~7.2 十进制位) | 深度学习训练/推理的标准精度 |
torch.float64 | torch.double | 64 | 1 | 11 | 52 | 53 位 (~15.9 十进制位) | 科学计算、高精度数值分析 |
上表位宽、指数位、尾数位数据遵循 IEEE 754 标准:二进制16(binary16)格式指数 5 位、尾数 10 位citeturn1search0;二进制32(binary32)格式指数 8 位、尾数 23 位citeturn1search8;二进制64(binary64)格式指数 11 位、尾数 52 位citeturn1search8。
数值属性查询
torch.finfo(dtype)
:返回指定浮点类型的数值极限信息,包括:bits
:总位宽eps
:机器 ε,即最小增量min
/max
:可表示的最小/最大值tiny
/smallest_normal
:最小非规范/规范化值 citeturn0search2。
import torch
print(torch.finfo(torch.float32))
# finfo(resolution=1e-06, min=-3.40282e+38, max=3.40282e+38, eps=1.19209e-07, smallest_normal=1.17549e-38, tiny=1.17549e-38, dtype=float32)
默认精度与类型提升
-
全局默认浮点精度
torch.get_default_dtype()
:获取当前默认浮点类型,初始值为torch.float32
citeturn0search9。torch.set_default_dtype(d)
:设置默认浮点类型,仅支持浮点类型输入;后续通过 Pythonfloat
构造的张量将采用该类型citeturn0search4。
-
类型提升 (Type Promotion)
torch.promote_types(type1, type2)
:返回在保证不降低精度与范围的前提下,最小的可兼容浮点类型,用于混合类型运算时的结果类型推断citeturn0search5。
自动混合精度(AMP)
PyTorch 的 AMP 机制在 前向/反向传播 中自动选择低精度(float16
或 bfloat16
)计算,而在 权重更新 等关键环节保留 float32
,以兼顾性能与数值稳定性。
torch.amp.autocast
:上下文管理器,针对支持的设备(如 CUDA GPU 或 CPU)自动切换运算精度;在 CUDA 上默认使用float16
,在 CPU 上可指定dtype=torch.bfloat16
citeturn0search6。torch.amp.GradScaler
:动态缩放梯度,避免低精度下的梯度下溢,实现稳定训练;与autocast
搭配使用可获显著加速(1.5–2×)和显存节省citeturn0search11。
示例(CUDA 上的混合精度训练):
from torch.amp import autocast, GradScaler
scaler = GradScaler()
for data, target in loader:optimizer.zero_grad()with autocast():output = model(data)loss = loss_fn(output, target)scaler.scale(loss).backward()scaler.step(optimizer)scaler.update()
实践建议
- 开发与调试阶段:优先使用
float32
,确保数值稳定。 - 大规模训练:若硬件支持 BF16,可尝试
bfloat16
训练;否则在 GPU 上结合 AMP 使用float16
。 3. 部署推理:在对精度容忍度高的场景下采用float16
,监控精度变化。 - 默认设置优化:根据项目需求使用
torch.set_default_dtype
控制全局默认精度,并结合torch.promote_types
处理跨类型运算。