Pytorch自动混合精度(AMP)介绍与使用 - autocast和Gradscaler

>>import torch
>>tensor1=torch.zeros(30,20)
>>tensor1.type()
'torch.FloatTensor'
>>tensor2=torch.Tensor([1,2])
>>tensor2.type()
'torch.FlatTensor'

torch.FloatTensor(32bit floating point)
torch.DoubleTensor(64bit floating point)
torch.HalfTensor(16bit floating piont1)
torch.BFloat16Tensor(16bit floating piont2)
torch.ByteTensor(8bit integer(unsigned)
torch.CharTensor(8bit integer(signed))
torch.ShortTensor(16bit integer(signed))
torch.IntTensor(32bit integer(signed))
torch.LongTensor(64bit integer(signed))
torch.BoolTensor(Boolean)
# 默认Tensor是32bit floating point，这就是32位浮点型精度的tensor。

__matmul__
addbmm
addmm
addmv
addr
baddbmm
bmm
chain_matmul
conv1d
conv2d
conv3d
conv_transpose1d
conv_transpose2d
conv_transpose3d
linear
matmul
mm
mv
prelu

from apex import amp
model,optimizer = amp.initial(model,optimizer,opt_level="O1")   #注意是O,不是０
with amp.scale_loss(loss,optimizer) as scaled_loss:
    scaled_loss.backward()取代loss.backward()

import argparse
import apex import amp
import apex.parallel import convert_syncbn_model
import apex.parallel import DistributedDataParallel as DDP
定义超参数：
def parse():
    parser=argparse.ArgumentParser()
    parser.add_argument('--local_rank',type=int, default=0)  #local_rank指定了输出设备，默认为GPU可用列表中的第一个GPU，必须加上。
    ...
    args = parser.parser.parse_args()
    return args
主函数写：
def main():
    args = parse()
    torch.cuda.set_device(args.local_rank)  #必须写在下一句的前面
   torch.distributed.init_process_group(
       'nccl',
       init_method='env://')
导入数据接口，需要用DistributedSampler
    dataset = ...
    num_workers = 4 if cuda else 0
    train_sampler=torch.utils.data.distributed.DistributedSampler(dataset)
    loader = DataLoader(dataset, batch_size=batchsize, shuflle=False, num_worker=num_workers,pin_memory=cuda, drop_last=True, sampler=train_sampler)
定义模型：
net = XXXNet(using_amp=True)
net.train()
net= convert_syncbn_model(net)
device=torch.device('cuda:{}'.format(args.local_rank))
net=net.to(device)
定义优化器，损失函数，定义优化器一定要把模型搬运到GPU之上
apt = Adam([{'params':params_low_lr,'lr':4e-5},
    {'params':params_high_lr,'lr':1e-4}],weight_decay=settings.WEIGHT_DECAY)
crit = nn.BCELoss().to(device)
多GPU设置import torch.nn.parallel.DistributedDataParallel as DDP
net,opt = amp.initialize(net,opt,opt_level='o1')
net=DDP(net,delay_allreduce=True)loss使用方法：opt.zero_grad()with amp.scale_loss(loss, opt) as scaled_loss:    scaled_loss.backward()opt.step()加入主入口：if __name__ == '__main__':    main()无论是apex支持的DDP还是pytorch自身支持的DDP,都需使用torch.distributed.launch来使用，方法如下：CUDA_VISIBLE_DIVECES=1,2,4 python -m torch.distributed.launch --nproc_per_node=3 train.py1,2,4是GPU编号，nproc_per_node是指定用了哪些GPU,记得开头说的local_rank，是因为torch.distributed.launch会调用这个local_ran

虽然是多个进程，但每个进程上模型的参数值都是一样的，而默认代号为０的进程是主进程if arg.local_rank == 0:
    torch.save(xxx)

from torch.cuda.amp import autocast as autocast
model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)
for input,target in data:
  optimizer.zero_grad()
  with autocast():
    output=model(input)
    loss = loss_fn(output,target)
  loss.backward()
  optimizer.step()

Traceback (most recent call last):
......
 File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
  result = self.forward(*input, ** kwargs)
......
RuntimeError: expected scalar type float but found c10::Half

from torch.cuda.amp import autocast as autocast
model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)
scaler = GradScaler() #训练前实例化一个GradScaler对象
for epoch in epochs:
  for input,target in data:
    optimizer.zero_grad()
    with autocast():　＃前后开启autocast
      output=model(input)
      loss = loss_fn(output,targt)
    scaler.scale(loss).backward()  #为了梯度放大
    #scaler.step()　首先把梯度值unscale回来，如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重，否则，忽略step调用，从而保证权重不更新。　　 scaler.step(optimizer)
    scaler.update()  #准备着，看是否要增大scaler

model = MyModel()
dp_model = nn.DataParallel(model)
with autocast():
    output=dp_model(input)
loss=loss_fn(output)

MyModel(nn.Module):
    @autocast()
    def forward(self, input):
        ...
#alternatively
MyModel(nn.Module):
    def forward(self, input):
        with autocast():
            ...
model = MyModel()
dp_model=nn.DataParallel(model)
with autocast():
    output=dp_model(input)
    loss = loss_fn(output)

gtx 1080ti:
半精度浮点数：0.17TFLOPS
单精度浮点数：11.34TFLOPS
双精度浮点数：0.33TFLOPS
rtx 2080ti:
半精度浮点数：20.14TFLOPS
单精度浮点数：10.07TFLOPS
双精度浮点数：0.31TFLOPS