为自己的 inicls 框架加上 fp16 训练

# 为自己的 inicls 框架加上 fp16 训练

# 01、如何实现

不得不说现在的库封装的都特别好，核心代码就加了几行

from torch.cuda.amp import GradScaler, autocast


for iteration in range(max_iteration):
    optimizer.zero_grad()

    if cfg.fp16 is True:
        with autocast():
            logits = model(images)
            loss = criterion(logits, labels).mean()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
    else:
        logits = model(images)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

if cfg.fp16:
	with autocast():
		logits = model(images)
else:
	logits = model(images)

1
2
3
4
5

# 02、效果比较

# 2.1 脚本介绍

实验的脚本非常简单，不得不说搭了这个框架之后做一些普通实验是真的很方便

export CUDA_VISIBLE_DEVICES=0
python train.py config/resnet/resnet18_b16x8_kaggle_leaves.py --tag resnet_fp16 --options "fp16=True" "data.train.ann_file=train_fold0.csv" "data.val.ann_file=valid_fold0.csv"

python train.py config/resnet/resnet18_b16x8_kaggle_leaves.py --tag resnet_not_fp16 --options "data.train.ann_file=train_fold0.csv" "data.val.ann_file=valid_fold0.csv"

1
2
3
4

# 2.2 速度、显存、性能对比

实验名称	学习率	验证集上最优精度	训练时间	显存占用	备注
resnet_fp16_lr_0.1	0.1	0.9553	13m27s	2300MiB	在 6k iteration 之后 loss 变为 nan
resnet_fp16_lr_0.01	0.01	0.9455	12m59s	2300MiB	在 5k iteration 之后 loss 变为 nan
resnet_fp16_lr_0.001	0.001	0.7126	13m31s	2300MiB	在 4k iteration 之后 loss 变为 nan
resnet_nofp16_lr_0.1	0.1	0.9651	19m19s		未使用 fp16
resnet_fp16_lr_0.1_v3	0.1	0.9635	12m10s	4360MiB	batch_size 改为 256，并且 loss 并未变为 nan
resnet_fp16_lr_0.1_v4	0.1	0.9624	50m29s	4360MiB	batch_size 改为 256，并且将epoch 改为 200，在 8k iteration 之后 loss 变为 nan
resnet_fp16_lr_0.1_v7_apex	0.1	0.933	16m2s	2300MiB	使用apex.amp，模式为01
resnet_fp16_lr_0.1_v8	0.1	0.9654	13m28s	2300MiB	将 `optimizer.step()` 写到 `else` 语句里面

# 03、loss 变为nan

在我用默认配置跑的时候，发现在 7k 左右的 iteration ，loss 会突变为 nan，目前尚不知是什么原因，先尝试了把学习率给调低，并没有什么影响

然后把 batch_size调高了发现不会出现nan了，

下面尝试讲 batch_size 调高的同时加大epoch数

破案了，发现是 optimizer 重复step了

if cfg.fp16 is True:
    with autocast():
        logits = model(images)
        loss = criterion(logits, labels)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
else:
    logits = model(images)
    loss = criterion(logits, labels)
    loss.backward()
optimizer.step()

1
2
3
4
5
6
7
8
9
10
11
12

将 optimizer.step() 写到里面即可

# 04、mmcls 是如何实现 fp16 训练的呢

参考资料：

上次更新: 2021/08/02, 21:04:52

← mmcls 是如何能够通过config 就搭建好一个模型的？为自己的 inicls 框架集成 Horovod→