AMP Tensor Cores节省内存PyTorch模型详解

目录
  • 导读
  • 什么是Tensor Cores?
  • 那么,我们如何使用Tensor Cores?
  • 使用PyTorch进行混合精度训练:
  • 基准测试

导读

只需要添加几行代码,就可以得到更快速,更省显存的PyTorch模型。

你知道吗,在1986年Geoffrey Hinton就在Nature论文中给出了反向传播算法?

此外,卷积网络最早是由Yann le cun在1998年提出的,用于数字分类,他使用了一个卷积层。但是直到2012年晚些时候,Alexnet才通过使用多个卷积层来实现最先进的imagenet。

那么,是什么让他们现在如此出名,而不是之前呢?

只有在我们拥有大量计算资源的情况下,我们才能够在最近的过去试验和充分利用深度学习的潜力。

但是,我们是否已经足够好地使用了我们的计算资源呢?我们能做得更好吗?

这篇文章的主要内容是关于如何利用Tensor Cores和自动混合精度更快地训练深度学习网络。

什么是Tensor Cores?

根据NVIDIA的网站:

NVIDIA Turing和Volta GPUs都是由Tensor Cores驱动的,这是一项突破性的技术,提供了突破性的AI性能。Tensor Cores可以加速AI核心的大矩阵运算,在一次运算中就可以完成混合精度的矩阵乘法和累加运算。在一个NVIDIA GPU上有数百个Tensor Cores并行运行,这大大提高了吞吐量和效率。

简单地说,它们是专门的cores,非常适合特定类型的矩阵操作。

我们可以将两个FP16矩阵相乘,并将其添加到一个FP16/FP32矩阵中,从而得到一个FP16/FP32矩阵。Tensor cores支持混合精度数学,即以半精度(FP16)进行输入,以全精度(FP32)进行输出。上述类型的操作对许多深度学习任务具有内在价值,而Tensor cores为这种操作提供了专门的硬件。

现在,使用FP16和FP32主要有两个好处。

  • FP16需要更少的内存,因此更容易训练和部署大型神经网络。它还只需要较少的数据移动。
  • 数学运算在降低精度的Tensor cores运行得更快。NVIDIA给出的Volta GPU的确切数字是:FP16的125 TFlops vs FP32的15.7 TFlops(8倍加速)。

但也有缺点。当我们从FP32转到FP16时,我们需要降低精度。

FP32 vs FP16: FP32 有8个指数位和23个分数位,而FP16有5个指数位和10个分数位。

但是FP32真的有必要吗?

实际上,FP16可以很好地表示大多数权重和梯度。所以存储和使用FP32是很浪费的。

那么,我们如何使用Tensor Cores?

我检查了一下我的Titan RTX GPU有576个tensor cores和4608个NVIDIA CUDA核心。但是我如何使用这些tensor cores呢?

坦白地说,NVIDIA用几行代码就能提供自动混合精度,因此使用tensor cores很简单。我们需要在代码中做两件事:

  • 需要用到FP32的运算比如Softmax之类的就分配用FP32,而Conv之类的操作可以用FP16的则被自动分配用FP16。

  • 使用损失缩放 为了保留小的梯度值。梯度值可能落在FP16的范围之外。在这种情况下,梯度值被缩放,使它们落在FP16范围内。

如果你还不了解背景细节也没关系,代码实现相对简单。

使用PyTorch进行混合精度训练:

让我们从PyTorch中的一个基本网络开始。

N, D_in, D_out = 64, 1024, 512
x = torch.randn(N, D_in, device="cuda")
y = torch.randn(N, D_out, device="cuda")
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
for to in range(500):
   y_pred = model(x)
   loss = torch.nn.functional.mse_loss(y_pred, y)
   optimizer.zero_grad()
   loss.backward()
   optimizer.step()

为了充分利用自动混合精度训练的优势,我们首先需要安装apex库。只需在终端中运行以下命令。

$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

然后,我们只需向神经网络代码中添加几行代码,就可以利用自动混合精度(AMP)。

from apex import amp
N, D_in, D_out = 64, 1024, 512
x = torch.randn(N, D_in, device="cuda")
y = torch.randn(N, D_out, device="cuda")
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
for to in range(500):
   y_pred = model(x)
   loss = torch.nn.functional.mse_loss(y_pred, y)
   optimizer.zero_grad()
   with amp.scale_loss(loss, optimizer) as scaled_loss:
      scaled_loss.backward()
   optimizer.step()

在这里你可以看到我们用amp.initialize初始化了我们的模型。我们还使用amp.scale_loss来指定损失缩放。

基准测试

git clone https://github.com/MLWhiz/data_science_blogs
cd data_science_blogs/amp/pytorch-apex-experiment/
python run_benchmark.py
python make_plot.py --GPU 'RTX' --method 'FP32' 'FP16' 'amp' --batch 128 256 512 1024 2048

这会在home目录中生成下面的图:

在这里,我使用不同的精度和批大小设置训练了同一个模型的多个实例。我们可以看到,从FP32到amp,内存需求减少,而精度保持大致相同。时间也会减少,但不会减少那么多。这可能是由于数据集或模型太简单。

根据NVIDIA给出的基准测试,AMP比标准的FP32快3倍左右,如下图所示。

在单精度和自动混合精度两种精度下,加速比为固定周期训练的时间比。

以上就是AMP Tensor Cores节省内存PyTorch模型详解的详细内容,更多关于AMP Tensor Cores 内存模型PyTorch的资料请关注我们其它相关文章!

(0)

相关推荐

  • pytorch tensor内所有元素相乘实例

    目录 tensor内所有元素相乘 tensor乘法运算汇总与解析 元素一一相乘 向量点乘 矩阵乘法 vector 与 matrix 相乘 matrix 与 vector 相乘 带有batch_size 的 broad cast乘法 tensor内所有元素相乘 a = torch.Tensor([1,2,3]) print(torch.prod(a)) 输出 tensor(6.) tensor乘法运算汇总与解析 元素一一相乘 该操作又称作 “哈达玛积”, 简单来说就是 tensor 元素逐个相乘.

  • pytorch tensor计算三通道均值方式

    目录 tensor计算三通道均值 第一种思路 Pytorch tensor的运算 tensor操作 tensor计算三通道均值 今天用pytorch处理图像时,涉及到了计算均值的问题,整理一下解决思路. 第一种思路 tensor转换为numpy再进行处理 import torch import cv2 img = cv2.imread("image path") tensor_img = torch.from_numpy((img[:, :, ::-1] / 255.0)[None,

  • pytorch 如何使用amp进行混合精度训练

    简介 AMP:Automatic mixed precision,自动混合精度,可以在神经网络推理过程中,针对不同的层,采用不同的数据精度进行计算,从而实现节省显存和加快速度的目的. 在Pytorch 1.5版本及以前,通过NVIDIA提供的apex库可以实现amp功能.但是在使用过程中会伴随着一些版本兼容和奇怪的报错问题. 从1.6版本开始,Pytorch原生支持自动混合精度训练,并已进入稳定阶段,AMP 训练能在 Tensor Core GPU 上实现更高的性能并节省多达 50% 的内存.

  • pytorch中Tensor.to(device)和model.to(device)的区别及说明

    目录 Tensor.to(device)和model.to(device)的区别 区别所在 举例 pytorch学习笔记--to(device)用法 这段代码到底有什么用呢? 为什么要在GPU上做运算呢? .cuda()和.to(device)的效果一样吗?为什么后者更好? 如果你有多个GPU Tensor.to(device)和model.to(device)的区别 区别所在 使用GPU训练的时候,需要将Module对象和Tensor类型的数据送入到device.通常会使用 to.(devic

  • pytorch随机采样操作SubsetRandomSampler()

    这篇文章记录一个采样器都随机地从原始的数据集中抽样数据.抽样数据采用permutation. 生成任意一个下标重排,从而利用下标来提取dataset中的数据的方法 需要的库 import torch 使用方法 这里以MNIST举例 train_dataset = dsets.MNIST(root='./data', #文件存放路径 train=True, #提取训练集 transform=transforms.ToTensor(), #将图像转化为Tensor download=True) sa

  • pytorch中dataloader 的sampler 参数详解

    目录 1. dataloader() 初始化函数 2. shuffle 与sample 之间的关系 3. sample 的定义方法 3.1 sampler 参数的使用 4. batch 生成过程 1. dataloader() 初始化函数 def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, collate_fn=None, pin_mem

  • AMP Tensor Cores节省内存PyTorch模型详解

    目录 导读 什么是Tensor Cores? 那么,我们如何使用Tensor Cores? 使用PyTorch进行混合精度训练: 基准测试 导读 只需要添加几行代码,就可以得到更快速,更省显存的PyTorch模型. 你知道吗,在1986年Geoffrey Hinton就在Nature论文中给出了反向传播算法? 此外,卷积网络最早是由Yann le cun在1998年提出的,用于数字分类,他使用了一个卷积层.但是直到2012年晚些时候,Alexnet才通过使用多个卷积层来实现最先进的imagene

  • Java 中的io模型详解

    1. BIO 我们先看一个 Java 例子: package cn.bridgeli.demo;   import java.io.BufferedReader; import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; import java.net.ServerSocket; import java.net.Socket;   /**  * @author bridgel

  • JVM入门之JVM内存结构内容详解

    一.java代码编译执行过程 源码编译:通过Java源码编译器将Java代码编译成JVM字节码(.class文件) 类加载:通过ClassLoader及其子类来完成JVM的类加载 类执行:字节码被装入内存,进入JVM虚拟机,被解释器解释执行   注:Java平台由Java虚拟机和Java应用程序接口搭建,Java语言则是进入这个平台的通道,   用Java语言编写并编译的程序可以运行在这个平台上 二.JVM简介 1.java程序经过一次编译之后,将java代码编译为字节码也就是class文件,然

  • 内存溢出和内存泄漏的详解及区别

    内存溢出和内存泄漏的详解及区别 内存溢出 out of memory,是指程序在申请内存时,没有足够的内存空间供其使用,出现out of memory:比如申请了一个integer,但给它存了long才能存下的数,那就是内存溢出. 内存泄露 memory leak,是指程序在申请内存后,无法释放已申请的内存空间,一次内存泄露危害可以忽略,但内存泄露堆积后果很严重,无论多少内存,迟早会被占光. memory leak会最终会导致out of memory! 内存溢出就是你要求分配的内存超出了系统能

  • C语言动态内存分配的详解

    C语言动态内存分配的详解 1.为什么使用动态内存分配 数组在使用的时候可能造成内存浪费,使用动态内存分配可以解决这个问题. 2. malloc和free C函数库提供了两个函数,malloc和free,分别用于执行动态内存分配和释放. (1)void *malloc(size_t size); malloc的参数就是需要分配的内存字节数.malloc分配一块连续的内存.如果操作系统无法向malloc提供更多的内存,malloc就返回一个NULL指针. (2)void free(void *poi

  • C++中的内存对齐实例详解

    C++中的内存对齐实例详解 内存对齐 在我们的程序中,数据结构还有变量等等都需要占有内存,在很多系统中,它都要求内存分配的时候要对齐,这样做的好处就是可以提高访问内存的速度. 我们还是先来看一段简单的程序: 程序一 #include <iostream> using namespace std; struct X1 { int i;//4个字节 char c1;//1个字节 char c2;//1个字节 }; struct X2 { char c1;//1个字节 int i;//4个字节 ch

  • Android 单线程模型详解及实例

    Android 单线程模型详解及实例 我们今天将会在这篇文章中为大家详细介绍有关Android单线程模型的相关内容.希望初学者们可以通过本文介绍的内容对这一概念有一个充分的认识,并从中对这一系统有一个深刻的认识. 当第一次启动一个Android程序时,Android会自动创建一个称为"main"主线程的线程.这个主线程(也称为UI线程)很重要,因为它负责把事件分派到相应的控件,其中就包括屏幕绘图事件,它同样是用户与Andriod控件交互的线程.比如,当你在屏幕上按下一个按钮后,UI线程

  • Java GC 机制与内存分配策略详解

    Java GC 机制与内存分配策略详解 收集算法是内存回收的方法论,垃圾收集器是内存回收的具体实现 自动内存管理解决的是:给对象分配内存 以及 回收分配给对象的内存 为什么我们要了解学习 GC 与内存分配呢? 在 JVM 自动内存管理机制的帮助下,不再需要为每一个new操作写配对的delete/free代码.但出现内存泄漏和溢出的问题时,如果不了解虚拟机是怎样使用内存的,那么排查错误将是一项非常艰难的工作. GC(垃圾收集器)在对堆进行回收前,会先确定哪些对象"存活",哪些已经&quo

  • Android 消息队列模型详解及实例

    Android 消息队列模型详解及实例 Android系统的消息队列和消息循环都是针对具体线程的,一个线程可以存在(当然也可以不存在)一个消息队列(Message Queue)和一个消息循环(Looper).Android中除了UI线程(主线程),创建的工作线程默认是没有消息循环和消息队列的.如果想让该线程具有消息队列和消息循环,并具有消息处理机制,就需要在线程中首先调用Looper.prepare()来创建消息队列,然后调用Looper.loop()进入消息循环.如以下代码所示: class

  • C++ 中继承与动态内存分配的详解

    C++ 中继承与动态内存分配的详解 继承是怎样与动态内存分配进行互动的呢?例如,如果基类使用动态内存分配,并重新定义赋值和复制构造函数,这将怎样影响派生类的实现呢?这个问题的答案取决于派生类的属性.如果派生类也使用动态内存分配,那么就需要学习几个新的小技巧.下面来看看这两种情况: 一.派生类不使用new 派生类是否需要为显示定义析构函数,复制构造函数和赋值操作符呢? 不需要! 首先,来看是否需要析构函数,如果没有定义析构函数,编译器将定义一个不执行任何操作的默认构造函数.实际上,派生类的默认构造

随机推荐