devbox
笔触狂放9
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

深度学习训练模型时，GPU显存不够怎么办？

作者：笔触狂放9 | 2024-08-20 20:03:04

赞

踩

gpu显存不够怎么办

点击上方“小白学视觉”，选择加"星标"或“置顶”

重磅干货，第一时间送达

作者丨游客26024@知乎（已授权）

来源丨https://www.zhihu.com/question/461811359/answer/2492822726

编辑丨极市平台

极市导读

 

此篇博文以AlexNet为网络架构(其需要输入的图像大小为227x227x3)，CIFAR10为数据集，Adamw为梯度下降函数，学习率机制为ReduceLROnPlateau举例。旨为如何让网络模型加速训练，而非去了解其原理。

题外话，我为什么要写这篇博客，就是因为我穷！没钱！租的服务器使用多GPU时一会钱就烧没了（gpu内存不用），急需要一种trick，来降低内存加速。

回到正题，如果我们使用的数据集较大，且网络较深，则会造成训练较慢，此时我们要想加速训练可以使用Pytorch的AMP（autocast与Gradscaler）；本文便是依据此写出的博文，对Pytorch的AMP(autocast与Gradscaler进行对比)自动混合精度对模型训练加速。

注意Pytorch1.6+，已经内置torch.cuda.amp，因此便不需要加载NVIDIA的apex库（半精度加速），为方便我们便不使用NVIDIA的apex库(安装麻烦)，转而使用torch.cuda.amp。

AMP (Automatic mixed precision): 自动混合精度，那什么是自动混合精度？

先来梳理一下历史：先有NVIDIA的apex，之后NVIDIA的开发人员将其贡献到Pytorch 1.6+产生了torch.cuda.amp[这是笔者梳理，可能有误，请留言]

详细讲：默认情况下，大多数深度学习框架都采用32位浮点算法进行训练。2017年，NVIDIA研究了一种用于混合精度训练的方法（apex），该方法在训练网络时将单精度（FP32）与半精度(FP16)结合在一起，并使用相同的超参数实现了与FP32几乎相同的精度，且速度比之前快了不少

之后，来到了AMP时代（特指torch.cuda.amp），此有两个关键词：自动与混合精度（Pytorch 1.6+中的torch.cuda.amp）其中，自动表现在Tensor的dtype类型会自动变化，框架按需自动调整tensor的dtype，可能有些地方需要手动干预；混合精度表现在采用不止一种精度的Tensor, torch.FloatTensor与torch.HalfTensor。并且从名字可以看出torch.cuda.amp，这个功能只能在cuda上使用！

为什么我们要使用AMP自动混合精度？

1.减少显存占用（FP16优势）

2.加快训练和推断的计算（FP16优势）

3.张量核心的普及（NVIDIA Tensor Core），低精度（FP16优势）

4. 混合精度训练缓解舍入误差问题，（FP16有此劣势，但是FP32可以避免此）

5.损失放大，可能使用混合精度还会出现无法收敛的问题[其原因时激活梯度值较小]，造成了溢出，则可以通过使用torch.cuda.amp.GradScaler放大损失来防止梯度的下溢

申明此篇博文主旨为如何让网络模型加速训练，而非去了解其原理，且其以AlexNet为网络架构(其需要输入的图像大小为227x227x3)，CIFAR10为数据集，Adamw为梯度下降函数，学习率机制为ReduceLROnPlateau举例。使用的电脑是2060的拯救者，虽然渣，但是还是可以搞搞这些测试。

本文从1.没使用DDP与DP训练与评估代码（之后加入amp），2.分布式DP训练与评估代码（之后加入amp），3.单进程占用多卡DDP训练与评估代码（之后加入amp） 角度讲解。

运行此程序时，文件的结构：

  1. D:/PycharmProject/Simple-CV-Pytorch-master
  2. |
  3. |
  4. |
  5. |----AMP(train_without.py、train_DP.py、train_autocast.py、train_GradScaler.py、eval_XXX.py
  6. |等,之后加入的alexnet也在这里，alexnet.py)
  7. |
  8. |
  9. |
  10. |----tensorboard(保存tensorboard的文件夹)
  11. |
  12. |
  13. |
  14. |----checkpoint(保存模型的文件夹)
  15. |
  16. |
  17. |
  18. |----data(数据集所在文件夹)

1.没使用DDP与DP训练与评估代码

没使用DDP与DP的训练与评估实验，作为我们实验的参照组

(1)原本模型的训练与评估源码：

训练源码：

注意：此段代码无比简陋，仅为代码的雏形，大致能理解尚可！

train_without.py

  1. import time
  2. import torch
  3. import torchvision
  4. from torch import nn
  5. from torch.utils.data import DataLoader
  6. from torchvision.models import alexnet
  7. from torchvision import transforms
  8. from torch.utils.tensorboard import SummaryWriter
  9. import numpy as np
  10. import argparse
  11.  
  12.  
  13. def parse_args():
  14.     parser = argparse.ArgumentParser(description='CV Train')
  15.     parser.add_mutually_exclusive_group()
  16.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  17.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  18.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  19.     parser.add_argument('--tensorboard', type=str, default=True, help='Use tensorboard for loss visualization')
  20.     parser.add_argument('--tensorboard_log', type=str, default='../tensorboard', help='tensorboard folder')
  21.     parser.add_argument('--cuda', type=str, default=True, help='if is cuda available')
  22.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  23.     parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
  24.     parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')
  25.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  26.     return parser.parse_args()
  27.  
  28.  
  29. args = parse_args()
  30.  
  31. # 1.Create SummaryWriter
  32. if args.tensorboard:
  33.     writer = SummaryWriter(args.tensorboard_log)
  34.  
  35. # 2.Ready dataset
  36. if args.dataset == 'CIFAR10':
  37.     train_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=True, transform=transforms.Compose(
  38.         [transforms.Resize(args.img_size), transforms.ToTensor()]), download=True)
  39. else:
  40.     raise ValueError("Dataset is not CIFAR10")
  41. cuda = torch.cuda.is_available()
  42. print('CUDA available: {}'.format(cuda))
  43.  
  44. # 3.Length
  45. train_dataset_size = len(train_dataset)
  46. print("the train dataset size is {}".format(train_dataset_size))
  47.  
  48. # 4.DataLoader
  49. train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size)
  50.  
  51. # 5.Create model
  52. model = alexnet()
  53.  
  54. if args.cuda == cuda:
  55.     model = model.cuda()
  56.  
  57. # 6.Create loss
  58. cross_entropy_loss = nn.CrossEntropyLoss()
  59.  
  60. # 7.Optimizer
  61. optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
  62. scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=3, verbose=True)
  63. # 8. Set some parameters to control loop
  64. # epoch
  65. iter = 0
  66. t0 = time.time()
  67. for epoch in range(args.epochs):
  68.     t1 = time.time()
  69.     print(" -----------------the {} number of training epoch --------------".format(epoch))
  70.     model.train()
  71.     for data in train_dataloader:
  72.         loss = 0
  73.         imgs, targets = data
  74.         if args.cuda == cuda:
  75.             cross_entropy_loss = cross_entropy_loss.cuda()
  76.             imgs, targets = imgs.cuda(), targets.cuda()
  77.         outputs = model(imgs)
  78.         loss_train = cross_entropy_loss(outputs, targets)
  79.         loss = loss_train.item() + loss
  80.         if args.tensorboard:
  81.             writer.add_scalar("train_loss", loss_train.item(), iter)
  82.  
  83.         optim.zero_grad()
  84.         loss_train.backward()
  85.         optim.step()
  86.         iter = iter + 1
  87.         if iter % 100 == 0:
  88.             print(
  89.                 "Epoch: {} | Iteration: {} | lr: {} | loss: {} | np.mean(loss): {} "
  90.                     .format(epoch, iter, optim.param_groups[0]['lr'], loss_train.item(),
  91.                             np.mean(loss)))
  92.     if args.tensorboard:
  93.         writer.add_scalar("lr", optim.param_groups[0]['lr'], epoch)
  94.     scheduler.step(np.mean(loss))
  95.     t2 = time.time()
  96.     h = (t2 - t1) // 3600
  97.     m = ((t2 - t1) % 3600) // 60
  98.     s = ((t2 - t1) % 3600) % 60
  99.     print("epoch {} is finished, and time is {}h{}m{}s".format(epoch, int(h), int(m), int(s)))
  100.  
  101.     if epoch % 1 == 0:
  102.         print("Save state, iter: {} ".format(epoch))
  103.         torch.save(model.state_dict(), "{}/AlexNet_{}.pth".format(args.checkpoint, epoch))
  104.  
  105. torch.save(model.state_dict(), "{}/AlexNet.pth".format(args.checkpoint))
  106. t3 = time.time()
  107. h_t = (t3 - t0) // 3600
  108. m_t = ((t3 - t0) % 3600) // 60
  109. s_t = ((t3 - t0) % 3600) // 60
  110. print("The finished time is {}h{}m{}s".format(int(h_t), int(m_t), int(s_t)))
  111. if args.tensorboard:
  112.     writer.close()

运行结果：

30aa626509d35f8f96f6bb1d9dd848e6.jpeg 376de9524dede231b30839125beb62b2.jpeg

Tensorboard观察：

7ebfc1cd8617d8e1a75266eba3417902.jpeg 543ea92802e1a1bd8bef8c02932ca603.jpeg

评估源码：

代码特别粗犷，尤其是device与精度计算，仅供参考，切勿模仿！

eval_without.py

  1. import torch
  2. import torchvision
  3. from torch.utils.data import DataLoader
  4. from torchvision.transforms import transforms
  5. from alexnet import alexnet
  6. import argparse
  7.  
  8.  
  9. # eval
  10. def parse_args():
  11.     parser = argparse.ArgumentParser(description='CV Evaluation')
  12.     parser.add_mutually_exclusive_group()
  13.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  14.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  15.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  16.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  17.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  18.     return parser.parse_args()
  19.  
  20.  
  21. args = parse_args()
  22. # 1.Create model
  23. model = alexnet()
  24.  
  25.  
  26. # 2.Ready Dataset
  27. if args.dataset == 'CIFAR10':
  28.     test_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=False,
  29.                                                 transform=transforms.Compose(
  30.                                                     [transforms.Resize(args.img_size),
  31.                                                      transforms.ToTensor()]),
  32.                                                 download=True)
  33. else:
  34.     raise ValueError("Dataset is not CIFAR10")
  35. # 3.Length
  36. test_dataset_size = len(test_dataset)
  37. print("the test dataset size is {}".format(test_dataset_size))
  38.  
  39. # 4.DataLoader
  40. test_dataloader = DataLoader(dataset=test_dataset, batch_size=args.batch_size)
  41.  
  42. # 5. Set some parameters for testing the network
  43. total_accuracy = 0
  44.  
  45. # test
  46. model.eval()
  47. with torch.no_grad():
  48.     for data in test_dataloader:
  49.         imgs, targets = data
  50.         device = torch.device('cpu')
  51.         imgs, targets = imgs.to(device), targets.to(device)
  52.         model_load = torch.load("{}/AlexNet.pth".format(args.checkpoint), map_location=device)
  53.         model.load_state_dict(model_load)
  54.         outputs = model(imgs)
  55.         outputs = outputs.to(device)
  56.         accuracy = (outputs.argmax(1) == targets).sum()
  57.         total_accuracy = total_accuracy + accuracy
  58.         accuracy = total_accuracy / test_dataset_size
  59.     print("the total accuracy is {}".format(accuracy))

运行结果：

23bb06c9bc8cc2812b7a7aaa0fc581c5.jpeg

分析：

原本模型训练完20个epochs花费了22分22秒，得到的准确率为0.8191

(2)原本模型加入autocast的训练与评估源码：

训练源码：

训练大致代码流程：

  1. from torch.cuda.amp import autocast as autocast
  2.  
  3. ...
  4.  
  5. # Create model, default torch.FloatTensor
  6. model = Net().cuda()
  7.  
  8. # SGD,Adm, Admw,...
  9. optim = optim.XXX(model.parameters(),..)
  10.  
  11. ...
  12.  
  13. for imgs,targets in dataloader:
  14.     imgs,targets = imgs.cuda(),targets.cuda()
  15.  
  16.     ....
  17.     with autocast():
  18.         outputs = model(imgs)
  19.         loss = loss_fn(outputs,targets)
  20.    ...
  21.     optim.zero_grad()
  22.     loss.backward()
  23.     optim.step()
  24.  
  25. ...

train_autocast_without.py

  1. import time
  2. import torch
  3. import torchvision
  4. from torch import nn
  5. from torch.cuda.amp import autocast
  6. from torchvision import transforms
  7. from torchvision.models import alexnet
  8. from torch.utils.data import DataLoader
  9. from torch.utils.tensorboard import SummaryWriter
  10. import numpy as np
  11. import argparse
  12.  
  13.  
  14. def parse_args():
  15.     parser = argparse.ArgumentParser(description='CV Train')
  16.     parser.add_mutually_exclusive_group()
  17.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  18.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  19.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  20.     parser.add_argument('--tensorboard', type=str, default=True, help='Use tensorboard for loss visualization')
  21.     parser.add_argument('--tensorboard_log', type=str, default='../tensorboard', help='tensorboard folder')
  22.     parser.add_argument('--cuda', type=str, default=True, help='if is cuda available')
  23.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  24.     parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
  25.     parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')
  26.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  27.     return parser.parse_args()
  28.  
  29.  
  30. args = parse_args()
  31.  
  32. # 1.Create SummaryWriter
  33. if args.tensorboard:
  34.     writer = SummaryWriter(args.tensorboard_log)
  35.  
  36. # 2.Ready dataset
  37. if args.dataset == 'CIFAR10':
  38.     train_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=True, transform=transforms.Compose(
  39.         [transforms.Resize(args.img_size), transforms.ToTensor()]), download=True)
  40. else:
  41.     raise ValueError("Dataset is not CIFAR10")
  42. cuda = torch.cuda.is_available()
  43. print('CUDA available: {}'.format(cuda))
  44.  
  45. # 3.Length
  46. train_dataset_size = len(train_dataset)
  47. print("the train dataset size is {}".format(train_dataset_size))
  48.  
  49. # 4.DataLoader
  50. train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size)
  51.  
  52. # 5.Create model
  53. model = alexnet()
  54.  
  55. if args.cuda == cuda:
  56.     model = model.cuda()
  57.  
  58. # 6.Create loss
  59. cross_entropy_loss = nn.CrossEntropyLoss()
  60.  
  61. # 7.Optimizer
  62. optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
  63. scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=3, verbose=True)
  64. # 8. Set some parameters to control loop
  65. # epoch
  66. iter = 0
  67. t0 = time.time()
  68. for epoch in range(args.epochs):
  69.     t1 = time.time()
  70.     print(" -----------------the {} number of training epoch --------------".format(epoch))
  71.     model.train()
  72.     for data in train_dataloader:
  73.         loss = 0
  74.         imgs, targets = data
  75.         if args.cuda == cuda:
  76.             cross_entropy_loss = cross_entropy_loss.cuda()
  77.             imgs, targets = imgs.cuda(), targets.cuda()
  78.         with autocast():
  79.             outputs = model(imgs)
  80.             loss_train = cross_entropy_loss(outputs, targets)
  81.         loss = loss_train.item() + loss
  82.         if args.tensorboard:
  83.             writer.add_scalar("train_loss", loss_train.item(), iter)
  84.  
  85.         optim.zero_grad()
  86.         loss_train.backward()
  87.         optim.step()
  88.         iter = iter + 1
  89.         if iter % 100 == 0:
  90.             print(
  91.                 "Epoch: {} | Iteration: {} | lr: {} | loss: {} | np.mean(loss): {} "
  92.                     .format(epoch, iter, optim.param_groups[0]['lr'], loss_train.item(),
  93.                             np.mean(loss)))
  94.     if args.tensorboard:
  95.         writer.add_scalar("lr", optim.param_groups[0]['lr'], epoch)
  96.     scheduler.step(np.mean(loss))
  97.     t2 = time.time()
  98.     h = (t2 - t1) // 3600
  99.     m = ((t2 - t1) % 3600) // 60
  100.     s = ((t2 - t1) % 3600) % 60
  101.     print("epoch {} is finished, and time is {}h{}m{}s".format(epoch, int(h), int(m), int(s)))
  102.  
  103.     if epoch % 1 == 0:
  104.         print("Save state, iter: {} ".format(epoch))
  105.         torch.save(model.state_dict(), "{}/AlexNet_{}.pth".format(args.checkpoint, epoch))
  106.  
  107. torch.save(model.state_dict(), "{}/AlexNet.pth".format(args.checkpoint))
  108. t3 = time.time()
  109. h_t = (t3 - t0) // 3600
  110. m_t = ((t3 - t0) % 3600) // 60
  111. s_t = ((t3 - t0) % 3600) // 60
  112. print("The finished time is {}h{}m{}s".format(int(h_t), int(m_t), int(s_t)))
  113. if args.tensorboard:
  114.     writer.close()

运行结果：

11e219be0d405c9dcc9a2ca9dc425aba.jpeg 9f97c180ee5efda20afc054d24f53c5d.jpeg

Tensorboard观察：

c5e5386aeeb4f040cd8d90fb78ad7596.jpeg db180d836b93c89d4ef137228157fddb.jpeg

评估源码：

eval_without.py 和 1.（1）一样

运行结果：

a76a95c19e2a0343730b046f26f15716.jpeg

分析：

原本模型训练完20个epochs花费了22分22秒，加入autocast之后模型花费的时间为21分21秒，说明模型速度增加了，并且准确率从之前的0.8191提升到0.8403

(3)原本模型加入autocast与GradScaler的训练与评估源码：

使用torch.cuda.amp.GradScaler是放大损失值来防止梯度的下溢

训练源码：

训练大致代码流程：

  1. from torch.cuda.amp import autocast as autocast
  2. from torch.cuda.amp import GradScaler as GradScaler
  3. ...
  4.  
  5. # Create model, default torch.FloatTensor
  6. model = Net().cuda()
  7.  
  8. # SGD,Adm, Admw,...
  9. optim = optim.XXX(model.parameters(),..)
  10. scaler = GradScaler()
  11.  
  12. ...
  13.  
  14. for imgs,targets in dataloader:
  15.     imgs,targets = imgs.cuda(),targets.cuda()
  16.     ...
  17.     optim.zero_grad()
  18.     ....
  19.     with autocast():
  20.         outputs = model(imgs)
  21.         loss = loss_fn(outputs,targets)
  22.  
  23.     scaler.scale(loss).backward()
  24.     scaler.step(optim)
  25.     scaler.update()
  26. ...

train_GradScaler_without.py

  1. import time
  2. import torch
  3. import torchvision
  4. from torch import nn
  5. from torch.cuda.amp import autocast, GradScaler
  6. from torchvision import transforms
  7. from torchvision.models import alexnet
  8. from torch.utils.data import DataLoader
  9. from torch.utils.tensorboard import SummaryWriter
  10. import numpy as np
  11. import argparse
  12.  
  13.  
  14. def parse_args():
  15.     parser = argparse.ArgumentParser(description='CV Train')
  16.     parser.add_mutually_exclusive_group()
  17.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  18.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  19.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  20.     parser.add_argument('--tensorboard', type=str, default=True, help='Use tensorboard for loss visualization')
  21.     parser.add_argument('--tensorboard_log', type=str, default='../tensorboard', help='tensorboard folder')
  22.     parser.add_argument('--cuda', type=str, default=True, help='if is cuda available')
  23.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  24.     parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
  25.     parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')
  26.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  27.     return parser.parse_args()
  28.  
  29.  
  30. args = parse_args()
  31.  
  32. # 1.Create SummaryWriter
  33. if args.tensorboard:
  34.     writer = SummaryWriter(args.tensorboard_log)
  35.  
  36. # 2.Ready dataset
  37. if args.dataset == 'CIFAR10':
  38.     train_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=True, transform=transforms.Compose(
  39.         [transforms.Resize(args.img_size), transforms.ToTensor()]), download=True)
  40. else:
  41.     raise ValueError("Dataset is not CIFAR10")
  42. cuda = torch.cuda.is_available()
  43. print('CUDA available: {}'.format(cuda))
  44.  
  45. # 3.Length
  46. train_dataset_size = len(train_dataset)
  47. print("the train dataset size is {}".format(train_dataset_size))
  48.  
  49. # 4.DataLoader
  50. train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size)
  51.  
  52. # 5.Create model
  53. model = alexnet()
  54.  
  55. if args.cuda == cuda:
  56.     model = model.cuda()
  57.  
  58. # 6.Create loss
  59. cross_entropy_loss = nn.CrossEntropyLoss()
  60.  
  61. # 7.Optimizer
  62. optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
  63. scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=3, verbose=True)
  64. scaler = GradScaler()
  65. # 8. Set some parameters to control loop
  66. # epoch
  67. iter = 0
  68. t0 = time.time()
  69. for epoch in range(args.epochs):
  70.     t1 = time.time()
  71.     print(" -----------------the {} number of training epoch --------------".format(epoch))
  72.     model.train()
  73.     for data in train_dataloader:
  74.         loss = 0
  75.         imgs, targets = data
  76.         optim.zero_grad()
  77.         if args.cuda == cuda:
  78.             cross_entropy_loss = cross_entropy_loss.cuda()
  79.             imgs, targets = imgs.cuda(), targets.cuda()
  80.         with autocast():
  81.             outputs = model(imgs)
  82.             loss_train = cross_entropy_loss(outputs, targets)
  83.             loss = loss_train.item() + loss
  84.         if args.tensorboard:
  85.             writer.add_scalar("train_loss", loss_train.item(), iter)
  86.  
  87.         scaler.scale(loss_train).backward()
  88.         scaler.step(optim)
  89.         scaler.update()
  90.         iter = iter + 1
  91.         if iter % 100 == 0:
  92.             print(
  93.                 "Epoch: {} | Iteration: {} | lr: {} | loss: {} | np.mean(loss): {} "
  94.                     .format(epoch, iter, optim.param_groups[0]['lr'], loss_train.item(),
  95.                             np.mean(loss)))
  96.     if args.tensorboard:
  97.         writer.add_scalar("lr", optim.param_groups[0]['lr'], epoch)
  98.     scheduler.step(np.mean(loss))
  99.     t2 = time.time()
  100.     h = (t2 - t1) // 3600
  101.     m = ((t2 - t1) % 3600) // 60
  102.     s = ((t2 - t1) % 3600) % 60
  103.     print("epoch {} is finished, and time is {}h{}m{}s".format(epoch, int(h), int(m), int(s)))
  104.  
  105.     if epoch % 1 == 0:
  106.         print("Save state, iter: {} ".format(epoch))
  107.         torch.save(model.state_dict(), "{}/AlexNet_{}.pth".format(args.checkpoint, epoch))
  108.  
  109. torch.save(model.state_dict(), "{}/AlexNet.pth".format(args.checkpoint))
  110. t3 = time.time()
  111. h_t = (t3 - t0) // 3600
  112. m_t = ((t3 - t0) % 3600) // 60
  113. s_t = ((t3 - t0) % 3600) // 60
  114. print("The finished time is {}h{}m{}s".format(int(h_t), int(m_t), int(s_t)))
  115. if args.tensorboard:
  116.     writer.close()

运行结果：

e48c8c92a03208d989fe5c2c70284270.jpeg bea7e94e7d0cffaed31ee4acd7106506.jpeg

Tensorboard观察：

fd696dda0cd658b6c25a613287af965c.jpeg 3107c999fe3be0371ab3682b5b7d1e01.jpeg

评估源码：

eval_without.py 和 1.（1）一样

运行结果：

bfa9634dc25c482c2fe7caad62c03442.jpeg

分析：

为什么，我们训练完20个epochs花费了27分27秒，比之前原模型未使用任何amp的时间（22分22秒）都多了？

这是因为我们使用了GradScaler放大了损失降低了模型训练的速度，还有个原因可能是笔者自身的显卡太小，没有起到加速的作用

2.分布式DP训练与评估代码

(1)DP原本模型的训练与评估源码：

训练源码：

train_DP.py

  1. import time
  2. import torch
  3. import torchvision
  4. from torch import nn
  5. from torch.utils.data import DataLoader
  6. from torchvision.models import alexnet
  7. from torchvision import transforms
  8. from torch.utils.tensorboard import SummaryWriter
  9. import numpy as np
  10. import argparse
  11.  
  12.  
  13. def parse_args():
  14.     parser = argparse.ArgumentParser(description='CV Train')
  15.     parser.add_mutually_exclusive_group()
  16.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  17.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  18.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  19.     parser.add_argument('--tensorboard', type=str, default=True, help='Use tensorboard for loss visualization')
  20.     parser.add_argument('--tensorboard_log', type=str, default='../tensorboard', help='tensorboard folder')
  21.     parser.add_argument('--cuda', type=str, default=True, help='if is cuda available')
  22.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  23.     parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
  24.     parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')
  25.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  26.     return parser.parse_args()
  27.  
  28.  
  29. args = parse_args()
  30.  
  31. # 1.Create SummaryWriter
  32. if args.tensorboard:
  33.     writer = SummaryWriter(args.tensorboard_log)
  34.  
  35. # 2.Ready dataset
  36. if args.dataset == 'CIFAR10':
  37.     train_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=True, transform=transforms.Compose(
  38.         [transforms.Resize(args.img_size), transforms.ToTensor()]), download=True)
  39. else:
  40.     raise ValueError("Dataset is not CIFAR10")
  41. cuda = torch.cuda.is_available()
  42. print('CUDA available: {}'.format(cuda))
  43.  
  44. # 3.Length
  45. train_dataset_size = len(train_dataset)
  46. print("the train dataset size is {}".format(train_dataset_size))
  47.  
  48. # 4.DataLoader
  49. train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size)
  50.  
  51. # 5.Create model
  52. model = alexnet()
  53.  
  54. if args.cuda == cuda:
  55.     model = model.cuda()
  56.     model = torch.nn.DataParallel(model).cuda()
  57. else:
  58.     model = torch.nn.DataParallel(model)
  59.  
  60. # 6.Create loss
  61. cross_entropy_loss = nn.CrossEntropyLoss()
  62.  
  63. # 7.Optimizer
  64. optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
  65. scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=3, verbose=True)
  66. # 8. Set some parameters to control loop
  67. # epoch
  68. iter = 0
  69. t0 = time.time()
  70. for epoch in range(args.epochs):
  71.     t1 = time.time()
  72.     print(" -----------------the {} number of training epoch --------------".format(epoch))
  73.     model.train()
  74.     for data in train_dataloader:
  75.         loss = 0
  76.         imgs, targets = data
  77.         if args.cuda == cuda:
  78.             cross_entropy_loss = cross_entropy_loss.cuda()
  79.             imgs, targets = imgs.cuda(), targets.cuda()
  80.         outputs = model(imgs)
  81.         loss_train = cross_entropy_loss(outputs, targets)
  82.         loss = loss_train.item() + loss
  83.         if args.tensorboard:
  84.             writer.add_scalar("train_loss", loss_train.item(), iter)
  85.  
  86.         optim.zero_grad()
  87.         loss_train.backward()
  88.         optim.step()
  89.         iter = iter + 1
  90.         if iter % 100 == 0:
  91.             print(
  92.                 "Epoch: {} | Iteration: {} | lr: {} | loss: {} | np.mean(loss): {} "
  93.                     .format(epoch, iter, optim.param_groups[0]['lr'], loss_train.item(),
  94.                             np.mean(loss)))
  95.     if args.tensorboard:
  96.         writer.add_scalar("lr", optim.param_groups[0]['lr'], epoch)
  97.     scheduler.step(np.mean(loss))
  98.     t2 = time.time()
  99.     h = (t2 - t1) // 3600
  100.     m = ((t2 - t1) % 3600) // 60
  101.     s = ((t2 - t1) % 3600) % 60
  102.     print("epoch {} is finished, and time is {}h{}m{}s".format(epoch, int(h), int(m), int(s)))
  103.  
  104.     if epoch % 1 == 0:
  105.         print("Save state, iter: {} ".format(epoch))
  106.         torch.save(model.state_dict(), "{}/AlexNet_{}.pth".format(args.checkpoint, epoch))
  107.  
  108. torch.save(model.state_dict(), "{}/AlexNet.pth".format(args.checkpoint))
  109. t3 = time.time()
  110. h_t = (t3 - t0) // 3600
  111. m_t = ((t3 - t0) % 3600) // 60
  112. s_t = ((t3 - t0) % 3600) // 60
  113. print("The finished time is {}h{}m{}s".format(int(h_t), int(m_t), int(s_t)))
  114. if args.tensorboard:
  115.     writer.close()

运行结果：

145244898d3bc86a29471b6936d8fbef.jpeg 1e74d6a3048a5de96d314cbc27f93801.jpeg

Tensorboard观察：

dd013bdee4d3bcce10c41d5fa9f371cd.jpeg cf279ec91ed3f252e3b8bb0df1f54df1.jpeg

评估源码：

eval_DP.py

  1. import torch
  2. import torchvision
  3. from torch.utils.data import DataLoader
  4. from torchvision.transforms import transforms
  5. from alexnet import alexnet
  6. import argparse
  7.  
  8.  
  9. # eval
  10. def parse_args():
  11.     parser = argparse.ArgumentParser(description='CV Evaluation')
  12.     parser.add_mutually_exclusive_group()
  13.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  14.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  15.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  16.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  17.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  18.     return parser.parse_args()
  19.  
  20.  
  21. args = parse_args()
  22. # 1.Create model
  23. model = alexnet()
  24. model = torch.nn.DataParallel(model)
  25.  
  26. # 2.Ready Dataset
  27. if args.dataset == 'CIFAR10':
  28.     test_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=False,
  29.                                                 transform=transforms.Compose(
  30.                                                     [transforms.Resize(args.img_size),
  31.                                                      transforms.ToTensor()]),
  32.                                                 download=True)
  33. else:
  34.     raise ValueError("Dataset is not CIFAR10")
  35. # 3.Length
  36. test_dataset_size = len(test_dataset)
  37. print("the test dataset size is {}".format(test_dataset_size))
  38.  
  39. # 4.DataLoader
  40. test_dataloader = DataLoader(dataset=test_dataset, batch_size=args.batch_size)
  41.  
  42. # 5. Set some parameters for testing the network
  43. total_accuracy = 0
  44.  
  45. # test
  46. model.eval()
  47. with torch.no_grad():
  48.     for data in test_dataloader:
  49.         imgs, targets = data
  50.         device = torch.device('cpu')
  51.         imgs, targets = imgs.to(device), targets.to(device)
  52.         model_load = torch.load("{}/AlexNet.pth".format(args.checkpoint), map_location=device)
  53.         model.load_state_dict(model_load)
  54.         outputs = model(imgs)
  55.         outputs = outputs.to(device)
  56.         accuracy = (outputs.argmax(1) == targets).sum()
  57.         total_accuracy = total_accuracy + accuracy
  58.         accuracy = total_accuracy / test_dataset_size
  59.     print("the total accuracy is {}".format(accuracy))

运行结果：

b440808411626261e52c05bbbdba7d44.jpeg

(2)DP使用autocast的训练与评估源码：

训练源码:

如果你这样写代码，那么你的代码无效！！！

  1. ...
  2.     model = Model()
  3.     model = torch.nn.DataParallel(model)
  4.     ...
  5.     with autocast():
  6.         output = model(imgs)
  7.         loss = loss_fn(output)

正确写法，训练大致流程代码：

  1. 1.Model(nn.Module):
  2.       @autocast()
  3.       def forward(self, input):
  4.       ...
  5.  
  6. 2.Model(nn.Module):
  7.       def foward(self, input):
  8.           with autocast():
  9.               ...

1与2皆可，之后：

  1. ...
  2. model = Model()
  3. model = torch.nn.DataParallel(model)
  4. with autocast():
  5.     output = model(imgs)
  6.     loss = loss_fn(output)
  7. ...

模型：

须在forward函数上加入@autocast()或者在forward里面最上面加入with autocast():

alexnet.py

  1. import torch
  2. import torch.nn as nn
  3. from torchvision.models.utils import load_state_dict_from_url
  4. from torch.cuda.amp import autocast
  5. from typing import Any
  6.  
  7. __all__ = ['AlexNet', 'alexnet']
  8.  
  9. model_urls = {
  10.     'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth',
  11. }
  12.  
  13.  
  14. class AlexNet(nn.Module):
  15.  
  16.     def __init__(self, num_classes: int = 1000) -> None:
  17.         super(AlexNet, self).__init__()
  18.         self.features = nn.Sequential(
  19.             nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
  20.             nn.ReLU(inplace=True),
  21.             nn.MaxPool2d(kernel_size=3, stride=2),
  22.             nn.Conv2d(64, 192, kernel_size=5, padding=2),
  23.             nn.ReLU(inplace=True),
  24.             nn.MaxPool2d(kernel_size=3, stride=2),
  25.             nn.Conv2d(192, 384, kernel_size=3, padding=1),
  26.             nn.ReLU(inplace=True),
  27.             nn.Conv2d(384, 256, kernel_size=3, padding=1),
  28.             nn.ReLU(inplace=True),
  29.             nn.Conv2d(256, 256, kernel_size=3, padding=1),
  30.             nn.ReLU(inplace=True),
  31.             nn.MaxPool2d(kernel_size=3, stride=2),
  32.         )
  33.         self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
  34.         self.classifier = nn.Sequential(
  35.             nn.Dropout(),
  36.             nn.Linear(256 * 6 * 6, 4096),
  37.             nn.ReLU(inplace=True),
  38.             nn.Dropout(),
  39.             nn.Linear(4096, 4096),
  40.             nn.ReLU(inplace=True),
  41.             nn.Linear(4096, num_classes),
  42.         )
  43.  
  44.     @autocast()
  45.     def forward(self, x: torch.Tensor) -> torch.Tensor:
  46.         x = self.features(x)
  47.         x = self.avgpool(x)
  48.         x = torch.flatten(x, 1)
  49.         x = self.classifier(x)
  50.         return x
  51.  
  52.  
  53. def alexnet(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> AlexNet:
  54.     r"""AlexNet model architecture from the
  55.     `"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
  56.     Args:
  57.         pretrained (bool): If True, returns a model pre-trained on ImageNet
  58.         progress (bool): If True, displays a progress bar of the download to stderr
  59.     """
  60.     model = AlexNet(**kwargs)
  61.     if pretrained:
  62.         state_dict = load_state_dict_from_url(model_urls["alexnet"],
  63.                                               progress=progress)
  64.         model.load_state_dict(state_dict)
  65.     return model

train_DP_autocast.py 导入自己的alexnet.py

  1. import time
  2. import torch
  3. from alexnet import alexnet
  4. import torchvision
  5. from torch import nn
  6. from torch.utils.data import DataLoader
  7. from torchvision import transforms
  8. from torch.cuda.amp import autocast as autocast
  9. from torch.utils.tensorboard import SummaryWriter
  10. import numpy as np
  11. import argparse
  12.  
  13.  
  14. def parse_args():
  15.     parser = argparse.ArgumentParser(description='CV Train')
  16.     parser.add_mutually_exclusive_group()
  17.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  18.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  19.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  20.     parser.add_argument('--tensorboard', type=str, default=True, help='Use tensorboard for loss visualization')
  21.     parser.add_argument('--tensorboard_log', type=str, default='../tensorboard', help='tensorboard folder')
  22.     parser.add_argument('--cuda', type=str, default=True, help='if is cuda available')
  23.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  24.     parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
  25.     parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')
  26.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  27.     return parser.parse_args()
  28.  
  29.  
  30. args = parse_args()
  31.  
  32. # 1.Create SummaryWriter
  33. if args.tensorboard:
  34.     writer = SummaryWriter(args.tensorboard_log)
  35.  
  36. # 2.Ready dataset
  37. if args.dataset == 'CIFAR10':
  38.     train_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=True, transform=transforms.Compose(
  39.         [transforms.Resize(args.img_size), transforms.ToTensor()]), download=True)
  40. else:
  41.     raise ValueError("Dataset is not CIFAR10")
  42. cuda = torch.cuda.is_available()
  43. print('CUDA available: {}'.format(cuda))
  44.  
  45. # 3.Length
  46. train_dataset_size = len(train_dataset)
  47. print("the train dataset size is {}".format(train_dataset_size))
  48.  
  49. # 4.DataLoader
  50. train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size)
  51.  
  52. # 5.Create model
  53. model = alexnet()
  54.  
  55. if args.cuda == cuda:
  56.     model = model.cuda()
  57.     model = torch.nn.DataParallel(model).cuda()
  58. else:
  59.     model = torch.nn.DataParallel(model)
  60.  
  61. # 6.Create loss
  62. cross_entropy_loss = nn.CrossEntropyLoss()
  63.  
  64. # 7.Optimizer
  65. optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
  66. scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=3, verbose=True)
  67. # 8. Set some parameters to control loop
  68. # epoch
  69. iter = 0
  70. t0 = time.time()
  71. for epoch in range(args.epochs):
  72.     t1 = time.time()
  73.     print(" -----------------the {} number of training epoch --------------".format(epoch))
  74.     model.train()
  75.     for data in train_dataloader:
  76.         loss = 0
  77.         imgs, targets = data
  78.         if args.cuda == cuda:
  79.             cross_entropy_loss = cross_entropy_loss.cuda()
  80.             imgs, targets = imgs.cuda(), targets.cuda()
  81.         with autocast():
  82.             outputs = model(imgs)
  83.             loss_train = cross_entropy_loss(outputs, targets)
  84.         loss = loss_train.item() + loss
  85.         if args.tensorboard:
  86.             writer.add_scalar("train_loss", loss_train.item(), iter)
  87.  
  88.         optim.zero_grad()
  89.         loss_train.backward()
  90.         optim.step()
  91.         iter = iter + 1
  92.         if iter % 100 == 0:
  93.             print(
  94.                 "Epoch: {} | Iteration: {} | lr: {} | loss: {} | np.mean(loss): {} "
  95.                     .format(epoch, iter, optim.param_groups[0]['lr'], loss_train.item(),
  96.                             np.mean(loss)))
  97.     if args.tensorboard:
  98.         writer.add_scalar("lr", optim.param_groups[0]['lr'], epoch)
  99.     scheduler.step(np.mean(loss))
  100.     t2 = time.time()
  101.     h = (t2 - t1) // 3600
  102.     m = ((t2 - t1) % 3600) // 60
  103.     s = ((t2 - t1) % 3600) % 60
  104.     print("epoch {} is finished, and time is {}h{}m{}s".format(epoch, int(h), int(m), int(s)))
  105.  
  106.     if epoch % 1 == 0:
  107.         print("Save state, iter: {} ".format(epoch))
  108.         torch.save(model.state_dict(), "{}/AlexNet_{}.pth".format(args.checkpoint, epoch))
  109.  
  110. torch.save(model.state_dict(), "{}/AlexNet.pth".format(args.checkpoint))
  111. t3 = time.time()
  112. h_t = (t3 - t0) // 3600
  113. m_t = ((t3 - t0) % 3600) // 60
  114. s_t = ((t3 - t0) % 3600) // 60
  115. print("The finished time is {}h{}m{}s".format(int(h_t), int(m_t), int(s_t)))
  116. if args.tensorboard:
  117.     writer.close()

运行结果：

065e1d8b984b2f351b873e7aa543a917.jpeg bb059d779d0ca53995f794746ffa3494.jpeg

Tensorboard观察：

41bd1764511b5e3aa28843e810344605.jpeg 6bd74693f8dc6c32aa72d79d5ba44f42.jpeg

评估源码：

eval_DP.py 相比与2. （1）导入自己的alexnet.py

运行结果：

7798f3f0cac02ae1c9af0d04d83a7098.jpeg

分析：

可以看出DP使用autocast训练完20个epochs时需要花费的时间是21分21秒，相比与之前DP没有使用的时间(22分22秒)快了1分1秒

之前DP未使用amp能达到准确率0.8216，而现在准确率降低到0.8188，说明还是使用自动混合精度加速还是对模型的准确率有所影响，后期可通过增大batch_sizel让运行时间和之前一样，但是准确率上升，来降低此影响

(3)DP使用autocast与GradScaler的训练与评估源码：

训练源码:

train_DP_GradScaler.py 导入自己的alexnet.py

  1. import time
  2. import torch
  3. from alexnet import alexnet
  4. import torchvision
  5. from torch import nn
  6. from torch.utils.data import DataLoader
  7. from torchvision import transforms
  8. from torch.cuda.amp import autocast as autocast
  9. from torch.cuda.amp import GradScaler as GradScaler
  10. from torch.utils.tensorboard import SummaryWriter
  11. import numpy as np
  12. import argparse
  13.  
  14.  
  15. def parse_args():
  16.     parser = argparse.ArgumentParser(description='CV Train')
  17.     parser.add_mutually_exclusive_group()
  18.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  19.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  20.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  21.     parser.add_argument('--tensorboard', type=str, default=True, help='Use tensorboard for loss visualization')
  22.     parser.add_argument('--tensorboard_log', type=str, default='../tensorboard', help='tensorboard folder')
  23.     parser.add_argument('--cuda', type=str, default=True, help='if is cuda available')
  24.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  25.     parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
  26.     parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')
  27.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  28.     return parser.parse_args()
  29.  
  30.  
  31. args = parse_args()
  32.  
  33. # 1.Create SummaryWriter
  34. if args.tensorboard:
  35.     writer = SummaryWriter(args.tensorboard_log)
  36.  
  37. # 2.Ready dataset
  38. if args.dataset == 'CIFAR10':
  39.     train_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=True, transform=transforms.Compose(
  40.         [transforms.Resize(args.img_size), transforms.ToTensor()]), download=True)
  41. else:
  42.     raise ValueError("Dataset is not CIFAR10")
  43. cuda = torch.cuda.is_available()
  44. print('CUDA available: {}'.format(cuda))
  45.  
  46. # 3.Length
  47. train_dataset_size = len(train_dataset)
  48. print("the train dataset size is {}".format(train_dataset_size))
  49.  
  50. # 4.DataLoader
  51. train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size)
  52.  
  53. # 5.Create model
  54. model = alexnet()
  55.  
  56. if args.cuda == cuda:
  57.     model = model.cuda()
  58.     model = torch.nn.DataParallel(model).cuda()
  59. else:
  60.     model = torch.nn.DataParallel(model)
  61.  
  62. # 6.Create loss
  63. cross_entropy_loss = nn.CrossEntropyLoss()
  64.  
  65. # 7.Optimizer
  66. optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
  67. scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=3, verbose=True)
  68. scaler = GradScaler()
  69. # 8. Set some parameters to control loop
  70. # epoch
  71. iter = 0
  72. t0 = time.time()
  73. for epoch in range(args.epochs):
  74.     t1 = time.time()
  75.     print(" -----------------the {} number of training epoch --------------".format(epoch))
  76.     model.train()
  77.     for data in train_dataloader:
  78.         loss = 0
  79.         imgs, targets = data
  80.         optim.zero_grad()
  81.         if args.cuda == cuda:
  82.             cross_entropy_loss = cross_entropy_loss.cuda()
  83.             imgs, targets = imgs.cuda(), targets.cuda()
  84.         with autocast():
  85.             outputs = model(imgs)
  86.             loss_train = cross_entropy_loss(outputs, targets)
  87.             loss = loss_train.item() + loss
  88.         if args.tensorboard:
  89.             writer.add_scalar("train_loss", loss_train.item(), iter)
  90.  
  91.         scaler.scale(loss_train).backward()
  92.         scaler.step(optim)
  93.         scaler.update()
  94.  
  95.         iter = iter + 1
  96.         if iter % 100 == 0:
  97.             print(
  98.                 "Epoch: {} | Iteration: {} | lr: {} | loss: {} | np.mean(loss): {} "
  99.                     .format(epoch, iter, optim.param_groups[0]['lr'], loss_train.item(),
  100.                             np.mean(loss)))
  101.     if args.tensorboard:
  102.         writer.add_scalar("lr", optim.param_groups[0]['lr'], epoch)
  103.     scheduler.step(np.mean(loss))
  104.     t2 = time.time()
  105.     h = (t2 - t1) // 3600
  106.     m = ((t2 - t1) % 3600) // 60
  107.     s = ((t2 - t1) % 3600) % 60
  108.     print("epoch {} is finished, and time is {}h{}m{}s".format(epoch, int(h), int(m), int(s)))
  109.  
  110.     if epoch % 1 == 0:
  111.         print("Save state, iter: {} ".format(epoch))
  112.         torch.save(model.state_dict(), "{}/AlexNet_{}.pth".format(args.checkpoint, epoch))
  113.  
  114. torch.save(model.state_dict(), "{}/AlexNet.pth".format(args.checkpoint))
  115. t3 = time.time()
  116. h_t = (t3 - t0) // 3600
  117. m_t = ((t3 - t0) % 3600) // 60
  118. s_t = ((t3 - t0) % 3600) // 60
  119. print("The finished time is {}h{}m{}s".format(int(h_t), int(m_t), int(s_t)))
  120. if args.tensorboard:
  121.     writer.close()

运行结果：

6b8d533c6551096f7a895b5b266661e9.jpeg 74de466fd614d54a051d0925d3a1a255.jpeg

Tensorboard观察：

25b6ba00a00304cab3183fb1623c3376.jpeg dcdaab86b7a4bafaabdf1e3fa5b1ba13.jpeg

评估源码：

eval_DP.py 相比与2. （1）导入自己的alexnet.py

运行结果：

310a8f2c174b0fce70f9e96742f4a89b.jpeg

分析：

跟之前一样，DP使用了GradScaler放大了损失降低了模型训练的速度

现在DP使用了autocast与GradScaler的准确率为0.8409，相比与DP只使用autocast准确率0.8188还是有所上升，并且之前DP未使用amp是准确率（0.8216）也提高了不少

3.单进程占用多卡DDP训练与评估代码

(1)DDP原模型训练与评估源码：

训练源码：

train_DDP.py

  1. import time
  2. import torch
  3. from torchvision.models.alexnet import alexnet
  4. import torchvision
  5. from torch import nn
  6. import torch.distributed as dist
  7. from torchvision import transforms
  8. from torch.utils.data import DataLoader
  9. from torch.utils.tensorboard import SummaryWriter
  10. import numpy as np
  11. import argparse
  12.  
  13.  
  14. def parse_args():
  15.     parser = argparse.ArgumentParser(description='CV Train')
  16.     parser.add_mutually_exclusive_group()
  17.     parser.add_argument("--rank", type=int, default=0)
  18.     parser.add_argument("--world_size", type=int, default=1)
  19.     parser.add_argument("--master_addr", type=str, default="127.0.0.1")
  20.     parser.add_argument("--master_port", type=str, default="12355")
  21.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  22.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  23.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  24.     parser.add_argument('--tensorboard', type=str, default=True, help='Use tensorboard for loss visualization')
  25.     parser.add_argument('--tensorboard_log', type=str, default='../tensorboard', help='tensorboard folder')
  26.     parser.add_argument('--cuda', type=str, default=True, help='if is cuda available')
  27.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  28.     parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
  29.     parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')
  30.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  31.     return parser.parse_args()
  32.  
  33.  
  34. args = parse_args()
  35.  
  36.  
  37. def train():
  38.     dist.init_process_group("gloo", init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
  39.                             rank=args.rank,
  40.                             world_size=args.world_size)
  41.     # 1.Create SummaryWriter
  42.     if args.tensorboard:
  43.         writer = SummaryWriter(args.tensorboard_log)
  44.  
  45.     # 2.Ready dataset
  46.     if args.dataset == 'CIFAR10':
  47.         train_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=True, transform=transforms.Compose(
  48.             [transforms.Resize(args.img_size), transforms.ToTensor()]), download=True)
  49.  
  50.     else:
  51.         raise ValueError("Dataset is not CIFAR10")
  52.  
  53.     cuda = torch.cuda.is_available()
  54.     print('CUDA available: {}'.format(cuda))
  55.  
  56.     # 3.Length
  57.     train_dataset_size = len(train_dataset)
  58.     print("the train dataset size is {}".format(train_dataset_size))
  59.  
  60.     train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
  61.     # 4.DataLoader
  62.     train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, sampler=train_sampler,
  63.                                   num_workers=2,
  64.                                   pin_memory=True)
  65.  
  66.     # 5.Create model
  67.     model = alexnet()
  68.  
  69.     if args.cuda == cuda:
  70.         model = model.cuda()
  71.         model = torch.nn.parallel.DistributedDataParallel(model).cuda()
  72.     else:
  73.         model = torch.nn.parallel.DistributedDataParallel(model)
  74.  
  75.     # 6.Create loss
  76.     cross_entropy_loss = nn.CrossEntropyLoss()
  77.  
  78.     # 7.Optimizer
  79.     optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
  80.     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=3, verbose=True)
  81.  
  82.     # 8. Set some parameters to control loop
  83.     # epoch
  84.     iter = 0
  85.     t0 = time.time()
  86.     for epoch in range(args.epochs):
  87.         t1 = time.time()
  88.         print(" -----------------the {} number of training epoch --------------".format(epoch))
  89.         model.train()
  90.         for data in train_dataloader:
  91.             loss = 0
  92.             imgs, targets = data
  93.             if args.cuda == cuda:
  94.                 cross_entropy_loss = cross_entropy_loss.cuda()
  95.                 imgs, targets = imgs.cuda(), targets.cuda()
  96.             outputs = model(imgs)
  97.             loss_train = cross_entropy_loss(outputs, targets)
  98.             loss = loss_train.item() + loss
  99.             if args.tensorboard:
  100.                 writer.add_scalar("train_loss", loss_train.item(), iter)
  101.  
  102.             optim.zero_grad()
  103.             loss_train.backward()
  104.             optim.step()
  105.             iter = iter + 1
  106.             if iter % 100 == 0:
  107.                 print(
  108.                     "Epoch: {} | Iteration: {} | lr: {} | loss: {} | np.mean(loss): {} "
  109.                         .format(epoch, iter, optim.param_groups[0]['lr'], loss_train.item(),
  110.                                 np.mean(loss)))
  111.         if args.tensorboard:
  112.             writer.add_scalar("lr", optim.param_groups[0]['lr'], epoch)
  113.         scheduler.step(np.mean(loss))
  114.         t2 = time.time()
  115.         h = (t2 - t1) // 3600
  116.         m = ((t2 - t1) % 3600) // 60
  117.         s = ((t2 - t1) % 3600) % 60
  118.         print("epoch {} is finished, and time is {}h{}m{}s".format(epoch, int(h), int(m), int(s)))
  119.  
  120.         if epoch % 1 == 0:
  121.             print("Save state, iter: {} ".format(epoch))
  122.             torch.save(model.state_dict(), "{}/AlexNet_{}.pth".format(args.checkpoint, epoch))
  123.  
  124.     torch.save(model.state_dict(), "{}/AlexNet.pth".format(args.checkpoint))
  125.     t3 = time.time()
  126.     h_t = (t3 - t0) // 3600
  127.     m_t = ((t3 - t0) % 3600) // 60
  128.     s_t = ((t3 - t0) % 3600) // 60
  129.     print("The finished time is {}h{}m{}s".format(int(h_t), int(m_t), int(s_t)))
  130.     if args.tensorboard:
  131.         writer.close()
  132.  
  133.  
  134. if __name__ == "__main__":
  135.     local_size = torch.cuda.device_count()
  136.     print("local_size: ".format(local_size))
  137.     train()

运行结果：

ff6576acd54d11f0c3c402a2caf675d5.jpeg e03dfbb12ba42433b60140639f68785a.jpeg

Tensorboard观察：

50977e5715772f1c70a94eda0ea37c8b.jpeg d4cc6a32f9ca9fa78f45f4ffdfdc74a8.jpeg

评估源码：

eval_DDP.py

  1. import torch
  2. import torchvision
  3. import torch.distributed as dist
  4. from torch.utils.data import DataLoader
  5. from torchvision.transforms import transforms
  6. # from alexnet import alexnet
  7. from torchvision.models.alexnet import alexnet
  8. import argparse
  9.  
  10.  
  11. # eval
  12. def parse_args():
  13.     parser = argparse.ArgumentParser(description='CV Evaluation')
  14.     parser.add_mutually_exclusive_group()
  15.     parser.add_argument("--rank", type=int, default=0)
  16.     parser.add_argument("--world_size", type=int, default=1)
  17.     parser.add_argument("--master_addr", type=str, default="127.0.0.1")
  18.     parser.add_argument("--master_port", type=str, default="12355")
  19.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  20.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  21.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  22.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  23.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  24.     return parser.parse_args()
  25.  
  26.  
  27. args = parse_args()
  28.  
  29.  
  30. def eval():
  31.     dist.init_process_group("gloo", init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
  32.                             rank=args.rank,
  33.                             world_size=args.world_size)
  34.     # 1.Create model
  35.     model = alexnet()
  36.     model = torch.nn.parallel.DistributedDataParallel(model)
  37.  
  38.     # 2.Ready Dataset
  39.     if args.dataset == 'CIFAR10':
  40.         test_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=False,
  41.                                                     transform=transforms.Compose(
  42.                                                         [transforms.Resize(args.img_size),
  43.                                                          transforms.ToTensor()]),
  44.                                                     download=True)
  45.  
  46.     else:
  47.         raise ValueError("Dataset is not CIFAR10")
  48.  
  49.     # 3.Length
  50.     test_dataset_size = len(test_dataset)
  51.     print("the test dataset size is {}".format(test_dataset_size))
  52.     test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)
  53.  
  54.     # 4.DataLoader
  55.     test_dataloader = DataLoader(dataset=test_dataset, sampler=test_sampler, batch_size=args.batch_size,
  56.                                  num_workers=2,
  57.                                  pin_memory=True)
  58.  
  59.     # 5. Set some parameters for testing the network
  60.     total_accuracy = 0
  61.  
  62.     # test
  63.     model.eval()
  64.     with torch.no_grad():
  65.         for data in test_dataloader:
  66.             imgs, targets = data
  67.             device = torch.device('cpu')
  68.             imgs, targets = imgs.to(device), targets.to(device)
  69.             model_load = torch.load("{}/AlexNet.pth".format(args.checkpoint), map_location=device)
  70.             model.load_state_dict(model_load)
  71.             outputs = model(imgs)
  72.             outputs = outputs.to(device)
  73.             accuracy = (outputs.argmax(1) == targets).sum()
  74.             total_accuracy = total_accuracy + accuracy
  75.             accuracy = total_accuracy / test_dataset_size
  76.         print("the total accuracy is {}".format(accuracy))
  77.  
  78.  
  79. if __name__ == "__main__":
  80.     local_size = torch.cuda.device_count()
  81.     print("local_size: ".format(local_size))
  82.     eval()

运行结果：

9a8eb57e6a05572a381456442fdff286.jpeg

(2)DDP使用autocast的训练与评估源码：

训练源码：

train_DDP_autocast.py 导入自己的alexnet.py

  1. import time
  2. import torch
  3. from alexnet import alexnet
  4. import torchvision
  5. from torch import nn
  6. import torch.distributed as dist
  7. from torchvision import transforms
  8. from torch.utils.data import DataLoader
  9. from torch.cuda.amp import autocast as autocast
  10. from torch.utils.tensorboard import SummaryWriter
  11. import numpy as np
  12. import argparse
  13.  
  14.  
  15. def parse_args():
  16.     parser = argparse.ArgumentParser(description='CV Train')
  17.     parser.add_mutually_exclusive_group()
  18.     parser.add_argument("--rank", type=int, default=0)
  19.     parser.add_argument("--world_size", type=int, default=1)
  20.     parser.add_argument("--master_addr", type=str, default="127.0.0.1")
  21.     parser.add_argument("--master_port", type=str, default="12355")
  22.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  23.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  24.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  25.     parser.add_argument('--tensorboard', type=str, default=True, help='Use tensorboard for loss visualization')
  26.     parser.add_argument('--tensorboard_log', type=str, default='../tensorboard', help='tensorboard folder')
  27.     parser.add_argument('--cuda', type=str, default=True, help='if is cuda available')
  28.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  29.     parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
  30.     parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')
  31.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  32.     return parser.parse_args()
  33.  
  34.  
  35. args = parse_args()
  36.  
  37.  
  38. def train():
  39.     dist.init_process_group("gloo", init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
  40.                             rank=args.rank,
  41.                             world_size=args.world_size)
  42.     # 1.Create SummaryWriter
  43.     if args.tensorboard:
  44.         writer = SummaryWriter(args.tensorboard_log)
  45.  
  46.     # 2.Ready dataset
  47.     if args.dataset == 'CIFAR10':
  48.         train_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=True, transform=transforms.Compose(
  49.             [transforms.Resize(args.img_size), transforms.ToTensor()]), download=True)
  50.  
  51.     else:
  52.         raise ValueError("Dataset is not CIFAR10")
  53.  
  54.     cuda = torch.cuda.is_available()
  55.     print('CUDA available: {}'.format(cuda))
  56.  
  57.     # 3.Length
  58.     train_dataset_size = len(train_dataset)
  59.     print("the train dataset size is {}".format(train_dataset_size))
  60.  
  61.     train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
  62.     # 4.DataLoader
  63.     train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, sampler=train_sampler,
  64.                                   num_workers=2,
  65.                                   pin_memory=True)
  66.  
  67.     # 5.Create model
  68.     model = alexnet()
  69.  
  70.     if args.cuda == cuda:
  71.         model = model.cuda()
  72.         model = torch.nn.parallel.DistributedDataParallel(model).cuda()
  73.     else:
  74.         model = torch.nn.parallel.DistributedDataParallel(model)
  75.  
  76.     # 6.Create loss
  77.     cross_entropy_loss = nn.CrossEntropyLoss()
  78.  
  79.     # 7.Optimizer
  80.     optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
  81.     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=3, verbose=True)
  82.  
  83.     # 8. Set some parameters to control loop
  84.     # epoch
  85.     iter = 0
  86.     t0 = time.time()
  87.     for epoch in range(args.epochs):
  88.         t1 = time.time()
  89.         print(" -----------------the {} number of training epoch --------------".format(epoch))
  90.         model.train()
  91.         for data in train_dataloader:
  92.             loss = 0
  93.             imgs, targets = data
  94.             if args.cuda == cuda:
  95.                 cross_entropy_loss = cross_entropy_loss.cuda()
  96.                 imgs, targets = imgs.cuda(), targets.cuda()
  97.             with autocast():
  98.                 outputs = model(imgs)
  99.                 loss_train = cross_entropy_loss(outputs, targets)
  100.             loss = loss_train.item() + loss
  101.             if args.tensorboard:
  102.                 writer.add_scalar("train_loss", loss_train.item(), iter)
  103.  
  104.             optim.zero_grad()
  105.             loss_train.backward()
  106.             optim.step()
  107.             iter = iter + 1
  108.             if iter % 100 == 0:
  109.                 print(
  110.                     "Epoch: {} | Iteration: {} | lr: {} | loss: {} | np.mean(loss): {} "
  111.                         .format(epoch, iter, optim.param_groups[0]['lr'], loss_train.item(),
  112.                                 np.mean(loss)))
  113.         if args.tensorboard:
  114.             writer.add_scalar("lr", optim.param_groups[0]['lr'], epoch)
  115.         scheduler.step(np.mean(loss))
  116.         t2 = time.time()
  117.         h = (t2 - t1) // 3600
  118.         m = ((t2 - t1) % 3600) // 60
  119.         s = ((t2 - t1) % 3600) % 60
  120.         print("epoch {} is finished, and time is {}h{}m{}s".format(epoch, int(h), int(m), int(s)))
  121.  
  122.         if epoch % 1 == 0:
  123.             print("Save state, iter: {} ".format(epoch))
  124.             torch.save(model.state_dict(), "{}/AlexNet_{}.pth".format(args.checkpoint, epoch))
  125.  
  126.     torch.save(model.state_dict(), "{}/AlexNet.pth".format(args.checkpoint))
  127.     t3 = time.time()
  128.     h_t = (t3 - t0) // 3600
  129.     m_t = ((t3 - t0) % 3600) // 60
  130.     s_t = ((t3 - t0) % 3600) // 60
  131.     print("The finished time is {}h{}m{}s".format(int(h_t), int(m_t), int(s_t)))
  132.     if args.tensorboard:
  133.         writer.close()
  134.  
  135.  
  136. if __name__ == "__main__":
  137.     local_size = torch.cuda.device_count()
  138.     print("local_size: ".format(local_size))
  139.     train()

运行结果：

f9a9ee93723966af31d820a142f04bea.jpeg 83c31645c6057eb82541c18258943e11.jpeg

Tensorboard观察：

eedaf9af23a013a8af9f11995fb335ac.jpeg c1a867bd471b8158b96385f105d83097.jpeg

评估源码：

eval_DDP.py 导入自己的alexnet.py

  1. import torch
  2. import torchvision
  3. import torch.distributed as dist
  4. from torch.utils.data import DataLoader
  5. from torchvision.transforms import transforms
  6. from alexnet import alexnet
  7. # from torchvision.models.alexnet import alexnet
  8. import argparse
  9.  
  10.  
  11. # eval
  12. def parse_args():
  13.     parser = argparse.ArgumentParser(description='CV Evaluation')
  14.     parser.add_mutually_exclusive_group()
  15.     parser.add_argument("--rank", type=int, default=0)
  16.     parser.add_argument("--world_size", type=int, default=1)
  17.     parser.add_argument("--master_addr", type=str, default="127.0.0.1")
  18.     parser.add_argument("--master_port", type=str, default="12355")
  19.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  20.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  21.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  22.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  23.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  24.     return parser.parse_args()
  25.  
  26.  
  27. args = parse_args()
  28.  
  29.  
  30. def eval():
  31.     dist.init_process_group("gloo", init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
  32.                             rank=args.rank,
  33.                             world_size=args.world_size)
  34.     # 1.Create model
  35.     model = alexnet()
  36.     model = torch.nn.parallel.DistributedDataParallel(model)
  37.  
  38.     # 2.Ready Dataset
  39.     if args.dataset == 'CIFAR10':
  40.         test_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=False,
  41.                                                     transform=transforms.Compose(
  42.                                                         [transforms.Resize(args.img_size),
  43.                                                          transforms.ToTensor()]),
  44.                                                     download=True)
  45.  
  46.     else:
  47.         raise ValueError("Dataset is not CIFAR10")
  48.  
  49.     # 3.Length
  50.     test_dataset_size = len(test_dataset)
  51.     print("the test dataset size is {}".format(test_dataset_size))
  52.     test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)
  53.  
  54.     # 4.DataLoader
  55.     test_dataloader = DataLoader(dataset=test_dataset, sampler=test_sampler, batch_size=args.batch_size,
  56.                                  num_workers=2,
  57.                                  pin_memory=True)
  58.  
  59.     # 5. Set some parameters for testing the network
  60.     total_accuracy = 0
  61.  
  62.     # test
  63.     model.eval()
  64.     with torch.no_grad():
  65.         for data in test_dataloader:
  66.             imgs, targets = data
  67.             device = torch.device('cpu')
  68.             imgs, targets = imgs.to(device), targets.to(device)
  69.             model_load = torch.load("{}/AlexNet.pth".format(args.checkpoint), map_location=device)
  70.             model.load_state_dict(model_load)
  71.             outputs = model(imgs)
  72.             outputs = outputs.to(device)
  73.             accuracy = (outputs.argmax(1) == targets).sum()
  74.             total_accuracy = total_accuracy + accuracy
  75.             accuracy = total_accuracy / test_dataset_size
  76.         print("the total accuracy is {}".format(accuracy))
  77.  
  78.  
  79. if __name__ == "__main__":
  80.     local_size = torch.cuda.device_count()
  81.     print("local_size: ".format(local_size))
  82.     eval()

运行结果：

0520145065e8ec9ef8f1f43bc1699ff2.jpeg

分析：

从DDP未使用amp花费21分21秒，DDP使用autocast花费20分20秒，说明速度提升了

DDP未使用amp的准确率0.8224，之后DDP使用了autocast准确率下降到0.8162

(3)DDP使用autocast与GradScaler的训练与评估源码

训练源码：

train_DDP_GradScaler.py 导入自己的alexnet.py

  1. import time
  2. import torch
  3. from alexnet import alexnet
  4. import torchvision
  5. from torch import nn
  6. import torch.distributed as dist
  7. from torchvision import transforms
  8. from torch.utils.data import DataLoader
  9. from torch.cuda.amp import autocast as autocast
  10. from torch.cuda.amp import GradScaler as GradScaler
  11. from torch.utils.tensorboard import SummaryWriter
  12. import numpy as np
  13. import argparse
  14.  
  15.  
  16. def parse_args():
  17.     parser = argparse.ArgumentParser(description='CV Train')
  18.     parser.add_mutually_exclusive_group()
  19.     parser.add_argument("--rank", type=int, default=0)
  20.     parser.add_argument("--world_size", type=int, default=1)
  21.     parser.add_argument("--master_addr", type=str, default="127.0.0.1")
  22.     parser.add_argument("--master_port", type=str, default="12355")
  23.     parser.add_argument('--dataset', type=str, default='CIFAR10', help='CIFAR10')
  24.     parser.add_argument('--dataset_root', type=str, default='../data', help='Dataset root directory path')
  25.     parser.add_argument('--img_size', type=int, default=227, help='image size')
  26.     parser.add_argument('--tensorboard', type=str, default=True, help='Use tensorboard for loss visualization')
  27.     parser.add_argument('--tensorboard_log', type=str, default='../tensorboard', help='tensorboard folder')
  28.     parser.add_argument('--cuda', type=str, default=True, help='if is cuda available')
  29.     parser.add_argument('--batch_size', type=int, default=64, help='batch size')
  30.     parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
  31.     parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')
  32.     parser.add_argument('--checkpoint', type=str, default='../checkpoint', help='Save .pth fold')
  33.     return parser.parse_args()
  34.  
  35.  
  36. args = parse_args()
  37.  
  38.  
  39. def train():
  40.     dist.init_process_group("gloo", init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
  41.                             rank=args.rank,
  42.                             world_size=args.world_size)
  43.     # 1.Create SummaryWriter
  44.     if args.tensorboard:
  45.         writer = SummaryWriter(args.tensorboard_log)
  46.  
  47.     # 2.Ready dataset
  48.     if args.dataset == 'CIFAR10':
  49.         train_dataset = torchvision.datasets.CIFAR10(root=args.dataset_root, train=True, transform=transforms.Compose(
  50.             [transforms.Resize(args.img_size), transforms.ToTensor()]), download=True)
  51.     else:
  52.         raise ValueError("Dataset is not CIFAR10")
  53.  
  54.     cuda = torch.cuda.is_available()
  55.     print('CUDA available: {}'.format(cuda))
  56.  
  57.     # 3.Length
  58.     train_dataset_size = len(train_dataset)
  59.     print("the train dataset size is {}".format(train_dataset_size))
  60.  
  61.     train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
  62.     # 4.DataLoader
  63.     train_dataloader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, sampler=train_sampler,
  64.                                   num_workers=2,
  65.                                   pin_memory=True)
  66.  
  67.     # 5.Create model
  68.     model = alexnet()
  69.  
  70.     if args.cuda == cuda:
  71.         model = model.cuda()
  72.         model = torch.nn.parallel.DistributedDataParallel(model).cuda()
  73.     else:
  74.         model = torch.nn.parallel.DistributedDataParallel(model)
  75.  
  76.     # 6.Create loss
  77.     cross_entropy_loss = nn.CrossEntropyLoss()
  78.  
  79.     # 7.Optimizer
  80.     optim = torch.optim.AdamW(model.parameters(), lr=args.lr)
  81.     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=3, verbose=True)
  82.     scaler = GradScaler()
  83.     # 8. Set some parameters to control loop
  84.     # epoch
  85.     iter = 0
  86.     t0 = time.time()
  87.     for epoch in range(args.epochs):
  88.         t1 = time.time()
  89.         print(" -----------------the {} number of training epoch --------------".format(epoch))
  90.         model.train()
  91.         for data in train_dataloader:
  92.             loss = 0
  93.             imgs, targets = data
  94.             optim.zero_grad()
  95.             if args.cuda == cuda:
  96.                 cross_entropy_loss = cross_entropy_loss.cuda()
  97.                 imgs, targets = imgs.cuda(), targets.cuda()
  98.             with autocast():
  99.                 outputs = model(imgs)
  100.                 loss_train = cross_entropy_loss(outputs, targets)
  101.                 loss = loss_train.item() + loss
  102.             if args.tensorboard:
  103.                 writer.add_scalar("train_loss", loss_train.item(), iter)
  104.  
  105.             scaler.scale(loss_train).backward()
  106.             scaler.step(optim)
  107.             scaler.update()
  108.  
  109.             iter = iter + 1
  110.             if iter % 100 == 0:
  111.                 print(
  112.                     "Epoch: {} | Iteration: {} | lr: {} | loss: {} | np.mean(loss): {} "
  113.                         .format(epoch, iter, optim.param_groups[0]['lr'], loss_train.item(),
  114.                                 np.mean(loss)))
  115.         if args.tensorboard:
  116.             writer.add_scalar("lr", optim.param_groups[0]['lr'], epoch)
  117.         scheduler.step(np.mean(loss))
  118.         t2 = time.time()
  119.         h = (t2 - t1) // 3600
  120.         m = ((t2 - t1) % 3600) // 60
  121.         s = ((t2 - t1) % 3600) % 60
  122.         print("epoch {} is finished, and time is {}h{}m{}s".format(epoch, int(h), int(m), int(s)))
  123.  
  124.         if epoch % 1 == 0:
  125.             print("Save state, iter: {} ".format(epoch))
  126.             torch.save(model.state_dict(), "{}/AlexNet_{}.pth".format(args.checkpoint, epoch))
  127.  
  128.     torch.save(model.state_dict(), "{}/AlexNet.pth".format(args.checkpoint))
  129.     t3 = time.time()
  130.     h_t = (t3 - t0) // 3600
  131.     m_t = ((t3 - t0) % 3600) // 60
  132.     s_t = ((t3 - t0) % 3600) // 60
  133.     print("The finished time is {}h{}m{}s".format(int(h_t), int(m_t), int(s_t)))
  134.     if args.tensorboard:
  135.         writer.close()
  136.  
  137.  
  138. if __name__ == "__main__":
  139.     local_size = torch.cuda.device_count()
  140.     print("local_size: ".format(local_size))
  141.     train()

运行结果：

9db69b3648a430895a9445ae03cf2579.jpeg 418e59c25f72b4583f207f1f914c606f.jpeg

Tensorboard观察：

e422b625c024b294db77d3033ab374e4.jpeg 11151bed790f79206d6e3910b9d3b7c8.jpeg

评估源码：

eval_DDP.py 与3. (2) 一样，导入自己的alexnet.py

运行结果：

f0fd086fb0693f855a7136055e9972e0.jpeg

分析：

运行起来了，速度也比DDP未使用amp（用时21分21秒）快了不少（用时20分20秒），之前DDP未使用amp准确率到达0.8224，现在DDP使用了autocast与GradScaler的准确率达到0.8252，提升了

参考：

1.Pytorch自动混合精度(AMP)训练：https://blog.csdn.net/ytusdc/article/details/122152244

2.PyTorch分布式训练基础--DDP使用：https://zhuanlan.zhihu.com/p/358974461

  1. 下载1：OpenCV-Contrib扩展模块中文版教程
  2.  
  3. 在「小白学视觉」公众号后台回复：扩展模块中文教程，即可下载全网第一份OpenCV扩展模块教程中文版，涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。
  4.  
  5.  
  6. 下载2：Python视觉实战项目52讲
  7. 在「小白学视觉」公众号后台回复：Python视觉实战项目，即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目，助力快速学校计算机视觉。
  8.  
  9.  
  10. 下载3：OpenCV实战项目20讲
  11. 在「小白学视觉」公众号后台回复：OpenCV实战项目20讲，即可下载含有20个基于OpenCV实现20个实战项目，实现OpenCV学习进阶。
  12.  
  13.  
  14. 交流群
  15.  
  16. 欢迎加入公众号读者群一起和同行交流，目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群（以后会逐渐细分），请扫描下面微信号加群，备注：”昵称+学校/公司+研究方向“，例如：”张三 + 上海交大 + 视觉SLAM“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告，否则会请出群，谢谢理解~
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/笔触狂放9/article/detail/1008553
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号