单卡训练自动迁移

本自动迁移样例使用APEX混合精度模块开启混合精度。

1. 在训练脚本中添加如下库代码：

   import torch
   ......
   import torch_npu
   from apex import amp
   from torch_npu.contrib import transfer_to_npu

2. 在模型和优化器定义之后添加初始化APEX模块的代码：

   ......
   criterion = nn.CrossEntropyLoss().to(device)
   
   optimizer = torch.optim.SGD(model.parameters(), args.lr,
                               momentum=args.momentum,
                               weight_decay=args.weight_decay)
   model, optimizer = amp.initialize(model, optimizer, combine_grad=True)

3. 改写梯度反向传播的代码。

   原代码如下：

   ......
   # compute gradient and do SGD step
   optimizer.zero_grad()
   loss.backward()
   optimizer.step()

   修改后代码如下：

   ......
   # compute gradient and do SGD step
   optimizer.zero_grad()
   with amp.scale_loss(loss, optimizer) as scaled_loss:
       scaled_loss.backward() 
   optimizer.step()

4. 参考模型训练，拉起训练进程。