FP16也称为半精度浮点数是一种计算机使用的二进制浮点数数据类型,使用 2 字节(16 位)存储
FP16的表示范围( 6 ∗ 1 0 − 8 → 65504 6*10^{-8} \to 65504 6∗10−8→65504),FP32表示范围( 1.4 ∗ 1 0 − 45 → 1.7 ∗ 1 0 38 1.4*10^{-45} \to 1.7*10^{38} 1.4∗10−45→1.7∗1038)
FP16的作用
FP16的问题
FP16 带来的问题主要有两个:溢出错误和舍入误差
为了解决FP16的问题,提出了混合精度训练+动态损失放大的方法
# 1. 导入包
from apex import amp
# 2. 模型和优化器的初始化
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”,不是“零一”
# 3. 损失函数的反向传播
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
# 4. 梯度裁剪(如果有)
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
模型和优化器初始化中 opt_level
参数说明:总共可选[O0, O1,O2, O3]
O0
:纯 FP32 训练,可以作为 accuracy 的 baseline
O1
:混合精度训练(推荐使用),根据黑白名单自动决定使用 FP16(GEMM, 卷积)还是 FP32(Softmax)进行计算
O2
:“几乎 FP16”混合精度训练,不存在黑白名单,除了 Batch Norm,几乎都是用 FP16 计算
O3
:纯 FP16 训练,很不稳定,但是可以作为 speed 的 baseline
使用例子
def train(trainset, evalset, model, tokenizer, model_dir, lr, epochs, device):
optimizer = AdamW(model.parameters(), lr=lr)
model, optimizer = amp.initialize(model, optimizer, opt_level="O3")
for epoch in tqdm(range(epochs), desc="epoch"):
train_loss, steps = 0, 0
for batch in tqdm(trainset, desc="train"):
batch = tuple(input_tensor.to(device) for input_tensor in batch if isinstance(input_tensor, torch.Tensor))
input_ids, label, mc_ids = batch
steps += 1
model.train()
loss, logits = model(input_ids=input_ids, mc_token_ids=mc_ids, labels=label)
# loss.backward()
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
train_loss += loss.item()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), 5)
# torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
optimizer.step()
optimizer.zero_grad()
if steps % 500 == 0:
print("step:%d avg_loss:%.3f"%(steps, train_loss/steps))
eval_res = evaluate(evalset, model, device)
os.makedirs(model_dir, exist_ok=True)
model_path = os.path.join(model_dir, "gpt2clsnews.model%d.ckpt"%epoch)
model.save_pretrained(model_path)
tokenizer.save_pretrained(os.path.join(model_dir,"gpt2clsnews.tokinizer"))
logging.info("checkpoint saved in %s"%model_dir)
结果对比
不使用apm | O0 | O1 | O3 | |
---|---|---|---|---|
训练耗时 | 3m48s | 3m45s | 3m18s | 1m56s |
预测耗时 | 1m12s | 1m12s | 50s | 34s |
accuracy | 0.859 | 0.849 | 0.875 | 0.004 |
显存占用 | 13.76G | 13.76G | 12.05G | 7.79G |
上面O3情况下loss 为nan
, 应该是出错了,从上表中可以看出使用混合精度O1可以节省部分时间和显存,其中accuracy并不相关,最开始使用了1000个样本测时不使用apm的accuracy比O1高,后来表中换成3000个样本时,O1比不使用apm高
PyTorch必备神器 | 唯快不破:基于Apex的混合精度加速