前言
曾在游戏世界挥洒创意,也曾在前端和后端的浪潮间穿梭,如今,而立的我仰望AI的璀璨星空,心潮澎湃,步履不停!愿你我皆乘风破浪,逐梦星辰!
混合精度:
FL32是目前模型存储数据一个比较普遍的格式。有时候过于浪费,根本用不着那么多。所以如果有一种方法能动态调整存储数据的大小,就能节省不少显存占用,从来提高批次大小,加速学习。
混合精度简单简单的来说,就是用FL16+FL32来代替原来的清一色FL32数据。具体实现是,在初始化好scaler以后,在遍历批次中的数据的时候,使用autocast自动对前向传播和损失处理混合精度。然后再使用梯度缩放器来缩放损失,并反向传播。
from torch.cuda.amp import GradScaler, autocastscaler = GradScaler()
# 使用autocast自动处理混合精度
with autocast():out = model(**data) # 前向传播loss = out['loss'] # 获取损失
# 使用梯度缩放器来缩放损失,并反向传播
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()