文章目录
什么是显存?
显存(GPU Memory)是图像处理器(GPU)上的专用内存,用于存储和处理图形数据及其他需要高并行计算的数据。在深度学习中,显存主要用于存储神经网络的模型参数、输入数据、中间计算结果(激活值)和梯度信息等。
显存的重要性
- 高并行计算:GPU具有大量的计算核心,能够并行处理大量数据,显存的高速访问能力支持这一特性。
- 存储大量数据:神经网络尤其是深层网络,包含大量参数和中间计算结果,显存需要足够大以容纳这些数据。
- 加速训练和推理:足够的显存可以让数据和模型尽可能多地驻留在GPU上,减少数据传输的瓶颈,从而加快训练和推理速度。
类比理解
想象显存是你的工作台:
- 模型参数:桌子上的工具和设备。
- 输入数据:正在处理的材料。
- 激活值:加工过程中的中间产物。
- 梯度信息:对工具进行调整的记录。
工作台空间越大,可以同时处理的材料和工具就越多,效率也越高。
显存与CPU内存的区别
- 速度:显存的访问速度远高于CPU内存,适合高并行计算任务。
- 容量:一般情况下,显存容量较小(例如8GB、16GB和24GB),而CPU内存可以更大(例如64GB、128GB)。
- 用途:显存专门用于GPU上的计算任务,而CPU内存用于存储和处理更广泛的数据。
什么是神经网络?
神经网络是一种模仿人脑神经元连接方式的计算模型,广泛应用于图像识别、自然语言处理等领域。它由多个层(Layer)组成,每一层包含多个神经元(Neuron),神经元之间通过权重(Weights)连接。
神经网络的基本结构
- 输入层(Input Layer):
- 接收原始数据,例如图像的像素值、文本的词向量等。
- 不进行任何计算,进作为数据的入口。
- 隐藏层(Hidden Layers):
- 进行特征提取和转换。
- 可以有多个隐藏层,每一层通过激活函数引入非线性。
- 隐藏层的数量和每层的神经元数量决定了网络的复杂度。
- 输出层(Output Layer):
- 输出最终结果,例如分类的概率分布、回归的数值等。
- 输出的形式取决于具体任务(分类、回归等)。
神经网络的工作原理
- 前向传播(Forward Propagation):
- 输入数据通过网络层层传递,经过加权和激活函数,生成输出结果。
- 每一层的输出作为下一层的输入。
- 损失计算(Loss Calculation):
- 比较网络输出与真实标签,计算损失值(误差)。
- 反向传播(Backward Propagation):
- 通过损失值反向计算梯度,更新模型参数(权重和偏置),以减小损失。
- 优化(Optimization):
- 使用优化算法(如SGD、Adam)根据梯度更新参数,使模型逐步学习数据的模式。
神经网络的类型
- 全连接网络(Fully Connected Networks):每一层的神经元与下一层的每个神经元完全连接。
- 卷积神经网络(Convolutional Neural Networks,CNNs):主要用于处理图像数据,利用卷积层提取空间特征。
- 循环神经网络(Recurrent Neural Networks,RNNs):适用于处理序列数据,如文本、时间序列。
- 生成对抗网络(Generative Adversarial Networks,GANs):由生成器和判别器组成,用于生成逼真的数据。
神经网络的参数
- 权重(Weights):
- 连接两个神经元之间的参数,决定输入信号的影响力。
- 网络学习的核心,通过训练不断调整。
- 偏置(Biases):
- 每个神经元的附加参数,用于调整激活函数的输入,增强模型的表达能力。
- 激活函数(Activation Functions):
- 引入非线性,使神经网络能够学习和表示复杂的模式。
- 常见的激活函数有ReLU、Sigmoid、Tanh等。
神经网络如何使用显存?
神经网络在运行时,显存主要用于以下几个方面:
- 模型参数:存储神经网络的权重和偏置。
- 输入数据:存储每批次(Batch)的输入数据。
- 激活值:存储每一层的输出,用于后向传播计算。
- 梯度信息:存储每个参数的梯度,用于优化模型。
- 优化器状态:如动量、学习率等额外信息。
为了更好的理解显存的占用,下面我们以一个具体的例子来计算显存使用情况。假设我们的神经网络包含两个卷积层和两个激活层,具体参数如下:
- 输入数据:
- 批量大小(Batch Size):64
- 通道数(Channels):3
- 高度(Height):32
- 宽度(Width):32
- 第一卷积层(Conv1):
- 卷积核大小:3×3
- 输入通道数(C_in):3
- 输出通道数(C_out):64
- 是否包含偏置:是
- 激活函数1(ReLU1)
- 第二卷积层(Conv2):
- 卷积核大小:3×3
- 输入通道数(C_in):64
- 输出通道数(C_out):128
- 是否包含偏置:是
- 激活函数2(ReLU2)
- 全连接层(FC):
- 输入特征数:128×28×28
- 输出特征数:10
模型参数的显存占用
模型参数包括每一层的权重和偏置。这些参数在训练过程中需要存储在显存中,以便在前向和反向传播中使用。
参数占用计算
显存占用量可以通过以下公式计算: 显存占用(字节) = 参数数量 × 每个参数的字节数 显存占用(字节)=参数数量×每个参数的字节数 显存占用(字节)=参数数量×每个参数的字节数
常用的数据类型及其字节数:
- Float32(FP32):4字节
- Float16(FP16):2字节
示例:两个3×3卷积层
第一卷积层(Conv1)
- 参数数量:
- 权重:3(输入通道)×3(高度)×3(宽度)×64(输出通道)=1728
- 偏置:64(输出通道)=64
- 总参数数量:1728+64=1792
- 显存占用:
- 使用FP32: 1792 ∗ 4 = 7168 字节 ≈ 7 K B 1792*4=7168字节≈7KB 1792∗4=7168字节≈7KB
- 使用FP16: 1792 ∗ 2 = 3584 字节 ≈ 3.5 K B 1792*2=3584字节≈3.5KB 1792∗2=3584字节≈3.5KB
第二卷积层(Conv2)
- 参数数量:
- 权重:64(输入通道)×3(高度)×3(宽度)×128(输出通道)=73728
- 偏置:128(输出通道)=128
- 总参数数量:73728+128=73856
- 显存占用:
- 使用FP32: 73856 ∗ 4 = 295424 字节 ≈ 288 K B 73856*4=295424字节≈288KB 73856∗4=295424字节≈288KB
- 使用FP16: 73856 ∗ 2 = 147712 字节 ≈ 144 K B 73856*2=147712字节≈144KB 73856∗2=147712字节≈144KB
全连接层(FC)
- 参数数量:
- 权重:128×28×28×10=1003520
- 偏置:10
- 总参数数量:1003520+10=1003530
- 显存占用:
- 使用FP32: 1003530 ∗ 4 = 4014120 字节 ≈ 3920 K B 1003530*4=4014120字节≈3920KB 1003530∗4=4014120字节≈3920KB
- 使用FP16: 1003530 ∗ 2 = 2007060 字节 ≈ 1960 K B 1003530*2=2007060字节≈1960KB 1003530∗2=2007060字节≈1960KB
总参数显存占用(假设使用FP32): 7 K B + 288 K B + 3920 K B = 4215 K B ≈ 4.12 M B 7KB+288KB+3920KB=4215KB≈4.12MB 7KB+288KB+3920KB=4215KB≈4.12MB
激活值的显存占用
激活值是每一层在前向传播过程中生成的输出数据。这些激活值需要暂时存储,以供后续层使用,并在反向传播时计算梯度。
激活占用计算: 显存占用(字节) = 激活数量 × 每个激活的字节数 显存占用(字节)=激活数量×每个激活的字节数 显存占用(字节)=激活数量×每个激活的字节数
激活数量取决于:
- 批量大小(Batch Size)
- 输出特征图的尺寸(Height×Width)
- 输出通道数(C_out)
示例:两个卷积层的激活
第一卷积层(Conv1)
- 输出特征图尺寸:
- 输入尺寸:32×32
- 卷积核大小:3×3
- 无填充(Padding=0),步幅为1
- 输出尺寸:32-3+1=30
- 即:30×30
- 激活数量: 64 ( 输出通道 ) × 30 (高度) × 30 (宽度) × 64 (批量大小) = 3686400 64(输出通道)×30(高度)×30(宽度)×64(批量大小)=3686400 64(输出通道)×30(高度)×30(宽度)×64(批量大小)=3686400
- 显存占用
- 使用FP32: 3686400 × 4 = 14745600 字节 ≈ 14 M B 3686400×4=14745600字节≈14MB 3686400×4=14745600字节≈14MB
- 使用FP16: 3686400 × 2 = 7372800 字节 ≈ 7 M B 3686400×2=7372800字节≈7MB 3686400×2=7372800字节≈7MB
第二卷积层(Conv2)
- 输出特征图尺寸:
- 输入尺寸:30×30
- 卷积核大小:3×3
- 无填充(Padding=0),步幅为1
- 输出尺寸:30-3+1=28
- 即:28×28
- 激活数量: 128 ( 输出通道 ) × 28 (高度) × 28 (宽度) × 64 (批量大小) = 6422528 128(输出通道)×28(高度)×28(宽度)×64(批量大小)=6422528 128(输出通道)×28(高度)×28(宽度)×64(批量大小)=6422528
- 显存占用
- 使用FP32: 6422528 × 4 = 25690112 字节 ≈ 24.5 M B 6422528×4=25690112字节≈24.5MB 6422528×4=25690112字节≈24.5MB
- 使用FP16: 6422528 × 2 = 12845056 字节 ≈ 12.25 M B 6422528×2=12845056字节≈12.25MB 6422528×2=12845056字节≈12.25MB
总激活显存占用(假设使用FP32): 14 M B + 24.5 M B = 38.5 M B 14MB+24.5MB=38.5MB 14MB+24.5MB=38.5MB
梯度信息的显存占用
在训练过程中,除了存储参数和激活值,还需要存储每个参数的梯度,以及优化器的状态(如动量)。
梯度占用计算
梯度的显存占用与参数的显存占用相同,因为每个参数对应一个梯度。
优化器状态占用计算
以Adam优化器为例,每个参数除了梯度外,还需要存储动量 m m m 和二阶动量 v v v,因此每个参数需要额外存储两倍的显存。
示例
继续使用之前的三个层:
- 总参数显存(FP32):4.12MB
- 梯度显存占用:4.12MB
- Adam优化器状态显存占用: 4.12 M B × 2 = 8.24 M B 4.12MB×2=8.24MB 4.12MB×2=8.24MB
总梯度和优化器显存占用(FP32): 4.12 K B + 8.24 M B = 12.36 M B 4.12KB+8.24MB=12.36MB 4.12KB+8.24MB=12.36MB
实际计算示例
让我们综合计算一个简单神经网络的显存占用。假设我们有以下网络结构:
- 输入层:
- 输入尺寸:64(批量大小) × 3(通道) × 32(高度) × 32(宽度)
- 第一卷积层(Conv1):
- 卷积核:3×3
- 输入通道:3
- 输出通道:64
- 是否包含偏置:是
- 激活函数1(ReLU1)
- 第二卷积层(Conv2):
- 卷积核:3×3
- 输入通道:64
- 输出通道:128
- 是否包含偏置:是
- 激活函数2(ReLU2)
- 全连接层(FC):
- 输入特征数:128 × 28 × 28
- 输出特征数:10
计算步骤
- 模型参数显存占用
- Conv1:
- 参数数量: 3 × 3 × 3 × 64 + 64 = 1728 + 64 = 1792 3 × 3 × 3 × 64 + 64 = 1728 + 64 = 1792 3×3×3×64+64=1728+64=1792
- 显存占用(FP32): 1792 × 4 = 7168 字节 ≈ 7 K B 1792 × 4 = 7168 字节 ≈ 7 KB 1792×4=7168字节≈7KB
- Conv2:
- 参数数量: 64 × 3 × 3 × 128 + 128 = 73728 + 128 = 73856 64 × 3 × 3 × 128 + 128 = 73728 + 128 = 73856 64×3×3×128+128=73728+128=73856
- 显存占用(FP32): 73856 × 4 = 295424 字节 ≈ 288 K B 73856 × 4 = 295424 字节 ≈ 288 KB 73856×4=295424字节≈288KB
- FC:
- 参数数量: 128 × 28 × 28 × 10 + 10 = 1003520 + 10 = 1003530 128 × 28 × 28 × 10 + 10 = 1003520 + 10 = 1003530 128×28×28×10+10=1003520+10=1003530
- 显存占用(FP32): 1003530 × 4 = 4014120 字节 ≈ 3920 K B 1003530 × 4 = 4014120 字节 ≈ 3920KB 1003530×4=4014120字节≈3920KB
- 总参数显存: 7 K B + 288 K B + 3920 K B = 4215 K B ≈ 4.12 M B 7KB+288KB+3920KB=4215KB≈4.12MB 7KB+288KB+3920KB=4215KB≈4.12MB
- 激活值显存占用
- Conv1 输出:
- 输出尺寸:64 × 64 × 30 × 30
- 激活数量:64 × 30 × 30 × 64 = 3686400
- 显存占用(FP32): 3686400 × 4 = 14745600 字节 ≈ 14 M B 3686400 × 4 = 14745600 字节 ≈ 14 MB 3686400×4=14745600字节≈14MB
- Conv2 输出:
- 输出尺寸:64 × 128 × 28 × 28
- 激活数量:64 × 128 × 28 × 28 = 6422528
- 显存占用(FP32): 6422528 × 4 = 25690112 字节 ≈ 24.5 M B 6422528 × 4 = 25690112 字节 ≈ 24.5 MB 6422528×4=25690112字节≈24.5MB
- 总激活显存: 14 M B + 24.5 M B = 38.5 M B B 14MB+24.5MB=38.5MBB 14MB+24.5MB=38.5MBB
- 梯度和优化器状态显存占用
- 梯度显存: 4.12 M B 4.12MB 4.12MB
- Adam优化器状态显存: 4.12 M B × 2 = 8.24 M B 4.12MB×2=8.24MB 4.12MB×2=8.24MB
- 总梯度和优化器显存: 4.12 K B + 8.24 M B = 12.36 M B 4.12KB+8.24MB=12.36MB 4.12KB+8.24MB=12.36MB
- 输入数据显存占用
- 输入数据:
- 尺寸:64 × 3 × 32 × 32
- 激活数量: 64 × 3 × 32 × 32 = 196608 64 × 3 × 32 × 32 = 196608 64×3×32×32=196608
- 显存占用(FP32): 196608 × 4 = 786 , 432 字节 ≈ 0.75 M B 196608 × 4 = 786,432 字节 ≈ 0.75 MB 196608×4=786,432字节≈0.75MB
- 总显存占用
总显存 = 模型参数 + 激活值 + 梯度和优化器 + 输入数据 总显存=模型参数+激活值+梯度和优化器+输入数据 总显存=模型参数+激活值+梯度和优化器+输入数据
总显存 = 4.12 M B + 38.5 M B + 12.36 M B + 0.75 M B ≈ 55.73 M B 总显存=4.12MB+38.5MB+12.36MB+0.75MB≈55.73MB 总显存=4.12MB+38.5MB+12.36MB+0.75MB≈55.73MB
注意:这是一个简化的计算,实际情况中,显存还会被框架本身、缓存和其他中间变量占用。
代码示例
环境准备
为了运行代码示例,需要我们本地有Python(>3.8)、PyTorch。接下来将以一个简单的神经网络,并观察它是如何使用显存的。
导入必要的库:
import torch
import torch.nn as nn
import torch.optim as optim
检查GPU是否可用:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用的设备: {device}")
如果输出“使用的设备:cuda“,说明GPU可用,否则使用CPU
定义一个简单的神经网络:
class SimpleConvNet(nn.Module):def __init__(self, input_channels, num_classes):super(SimpleConvNet, self).__init__()self.conv1 = nn.Conv2d(input_channels, 64, kernel_size=3, bias=True) # 第一卷积层self.relu1 = nn.ReLU() # 激活函数1self.conv2 = nn.Conv2d(64, 128, kernel_size=3, bias=True) # 第二卷积层self.relu2 = nn.ReLU() # 激活函数2self.fc = nn.Linear(128 * 28 * 28, num_classes) # 全连接层def forward(self, x):out = self.conv1(x)out = self.relu1(out)out = self.conv2(out)out = self.relu2(out)out = out.view(out.size(0), -1) # 展平out = self.fc(out)return out# 实例化模型
input_channels = 3
num_classes = 10
model = SimpleConvNet(input_channels, num_classes).to(device)
print(model)
以上的输出结果为:
SimpleConvNet((conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))(relu1): ReLU()(conv2): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))(relu2): ReLU()(fc): Linear(in_features=25088, out_features=10, bias=True)
)
创建输入数据和标签:
# 批量大小
batch_size = 64# 输入数据尺寸:批量大小 × 通道数 × 高度 × 宽度
inputs = torch.randn(batch_size, input_channels, 32, 32).to(device)# 随机生成标签(假设有10个类别)
labels = torch.randint(0, num_classes, (batch_size,)).to(device)
定义损失函数和优化器:
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
前向传播和反向传播:
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
print(f"损失: {loss.item()}")# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
完整代码如下:
将以上代码放在一个Python脚本中运行,你将会看到模型结构和损失值输出。
import torch
import torch.nn as nn
import torch.optim as optim# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用的设备: {device}")# 定义神经网络
class SimpleConvNet(nn.Module):def __init__(self, input_channels, num_classes):super(SimpleConvNet, self).__init__()self.conv1 = nn.Conv2d(input_channels, 64, kernel_size=3, bias=True)self.relu1 = nn.ReLU()self.conv2 = nn.Conv2d(64, 128, kernel_size=3, bias=True)self.relu2 = nn.ReLU()self.fc = nn.Linear(128 * 28 * 28, num_classes)def forward(self, x):out = self.conv1(x)out = self.relu1(out)out = self.conv2(out)out = self.relu2(out)out = out.view(out.size(0), -1)out = self.fc(out)return out# 实例化模型
input_channels = 3
num_classes = 10
model = SimpleConvNet(input_channels, num_classes).to(device)
print(model)# 创建输入数据和标签
batch_size = 64
inputs = torch.randn(batch_size, input_channels, 32, 32).to(device)
labels = torch.randint(0, num_classes, (batch_size,)).to(device)# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
print(f"损失: {loss.item()}")# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
以上代码的运行结果为:
使用的设备: cuda
SimpleConvNet((conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))(relu1): ReLU()(conv2): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))(relu2): ReLU()(fc): Linear(in_features=25088, out_features=10, bias=True)
)
损失: 2.3979134559631348
注:由于设备和代码运行环境不同,损失值不相同,这个无关紧要,因为在正常网络训练过程中可以看到损失值是逐渐收敛的。
如何查看显存使用情况
在训练神经网络时,我们可以使用PyTorch提供的工具来查看显存使用情况。
使用torch.cuda.memory_allocated
torch.cuda.memory_allocated(device)
返回指定设备上已分配的显存(以字节为单位)。
if torch.cuda.is_available():allocated = torch.cuda.memory_allocated(device) / (1024 ** 2) # 转换为MBprint(f"已分配的显存: {allocated:.2f} MB")
使用torch.cuda.memory_reserved
torch.cuda.memory_reserved(device)
返回PyTorch在指定设备上保留的显存总量(以字节为单位)。这包括被分配和未被分配的显存。
if torch.cuda.is_available():reserved = torch.cuda.memory_reserved(device) / (1024 ** 2) # 转换为MBprint(f"保留的显存: {reserved:.2f} MB")
示例:在训练过程中查看显存
在以上代码的基础上,我们假设在每个epoch结束时查看显存使用情况。
import torch
import torch.nn as nn
import torch.optim as optim# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用的设备: {device}")# 定义神经网络
class SimpleConvNet(nn.Module):def __init__(self, input_channels, num_classes):super(SimpleConvNet, self).__init__()self.conv1 = nn.Conv2d(input_channels, 64, kernel_size=3, bias=True)self.relu1 = nn.ReLU()self.conv2 = nn.Conv2d(64, 128, kernel_size=3, bias=True)self.relu2 = nn.ReLU()self.fc = nn.Linear(128 * 28 * 28, num_classes)def forward(self, x):out = self.conv1(x)out = self.relu1(out)out = self.conv2(out)out = self.relu2(out)out = out.view(out.size(0), -1)out = self.fc(out)return out# 实例化模型
input_channels = 3
num_classes = 10
model = SimpleConvNet(input_channels, num_classes).to(device)
print(model)# 创建输入数据和标签
batch_size = 64
inputs = torch.randn(batch_size, input_channels, 32, 32).to(device)
labels = torch.randint(0, num_classes, (batch_size,)).to(device)# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练多个epoch,并查看显存使用情况
num_epochs = 2
for epoch in range(num_epochs):# 前向传播outputs = model(inputs)loss = criterion(outputs, labels)print(f"Epoch {epoch+1}, 损失: {loss.item()}")# 反向传播optimizer.zero_grad()loss.backward()optimizer.step()# 查看显存使用情况if torch.cuda.is_available():allocated = torch.cuda.memory_allocated(device) / (1024 ** 2)reserved = torch.cuda.memory_reserved(device) / (1024 ** 2)print(f"Epoch {epoch+1}: 已分配显存: {allocated:.2f} MB, 保留显存: {reserved:.2f} MB")
运行结果:
使用的设备: cuda
SimpleConvNet((conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))(relu1): ReLU()(conv2): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))(relu2): ReLU()(fc): Linear(in_features=25088, out_features=10, bias=True)
)
Epoch 1, 损失: 2.3979134559631348
Epoch 1: 已分配显存: 54.12 MB, 保留显存: 54.12 MB
Epoch 2, 损失: 2.364225387573242
Epoch 2: 已分配显存: 54.12 MB, 保留显存: 54.12 MB
说明:
- 已分配显存(memory_allocated)表示当前被PyTorch实际使用的显存。
- 保留显存(memory_reserved)表示PyTorch为未来的分配预留的显存。这部分显存可能尚未被实际使用。
根据运行结果,有小伙伴可能注意到控制台打印的已分配显存和我们上面计算的显存并不一致,这是因为在网络训练过程中,由于代码运行环境(Pytorch、系统、Python版本等)以及其他程序运行等原因,PyTorch可能会有不同的显存优化策略从而与理论计算结果不一致。
优化显存的方法
当我们手里拥有的显存有限,尤其是在使用大模型或大批量数据时,以下是一些常见的优化方案:
减少批量大小(Batch Size)
批量大小越大,占用的显存越多。通过减少批量大小,可以降低显存使用。
batch_size = 32 # 从64减半
inputs = torch.randn(batch_size, input_channels, 32, 32).to(device)
labels = torch.randint(0, num_classes, (batch_size,)).to(device)
使用更小的模型
减少神经网络的层数或每层的神经元数量。
# 修改隐藏层通道数
class SmallerConvNet(nn.Module):def __init__(self, input_channels, num_classes):super(SmallerConvNet, self).__init__()self.conv1 = nn.Conv2d(input_channels, 32, kernel_size=3, bias=True) # 从64减半self.relu1 = nn.ReLU()self.conv2 = nn.Conv2d(32, 64, kernel_size=3, bias=True) # 从128减半self.relu2 = nn.ReLU()self.fc = nn.Linear(64 * 28 * 28, num_classes) # 全连接层尺寸相应调整def forward(self, x):out = self.conv1(x)out = self.relu1(out)out = self.conv2(out)out = self.relu2(out)out = out.view(out.size(0), -1)out = self.fc(out)return outmodel = SmallerConvNet(input_channels, num_classes).to(device)
print(model)
效果:减少每层的通道数,可以显著降低模型参数和激活值的显存占用。
混合精度训练
使用16位浮点数(FP16)代替32位浮点数(FP32),可以减少显存占用,并加快计算速度。
安装torch.cuda.amp
PyTorch 1.6及以上版本已经内置了自动混合精度(AMP)。
import torch
import torch.nn as nn
import torch.optim as optim# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用的设备: {device}")# 定义神经网络
class SimpleConvNet(nn.Module):def __init__(self, input_channels, num_classes):super(SimpleConvNet, self).__init__()self.conv1 = nn.Conv2d(input_channels, 64, kernel_size=3, bias=True)self.relu1 = nn.ReLU()self.conv2 = nn.Conv2d(64, 128, kernel_size=3, bias=True)self.relu2 = nn.ReLU()self.fc = nn.Linear(128 * 28 * 28, num_classes)def forward(self, x):out = self.conv1(x)out = self.relu1(out)out = self.conv2(out)out = self.relu2(out)out = out.view(out.size(0), -1)out = self.fc(out)return out# 实例化模型
input_channels = 3
num_classes = 10
model = SimpleConvNet(input_channels, num_classes).to(device)
print(model)# 创建输入数据和标签
batch_size = 64
inputs = torch.randn(batch_size, input_channels, 32, 32).to(device)
labels = torch.randint(0, num_classes, (batch_size,)).to(device)# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 初始化GradScaler
scaler = torch.cuda.amp.GradScaler()# 训练多个epoch,并查看显存使用情况
num_epochs = 2
for epoch in range(num_epochs):optimizer.zero_grad()with torch.cuda.amp.autocast():outputs = model(inputs)loss = criterion(outputs, labels)scaler.scale(loss).backward()scaler.step(optimizer)scaler.update()print(f"Epoch {epoch+1}, 损失: {loss.item()}")# 查看显存使用情况if torch.cuda.is_available():allocated = torch.cuda.memory_allocated(device) / (1024 ** 2)reserved = torch.cuda.memory_reserved(device) / (1024 ** 2)print(f"Epoch {epoch+1}: 已分配显存: {allocated:.2f} MB, 保留显存: {reserved:.2f} MB")
效果:通过使用混合精度训练,可以显著减少显存占用,大约减少一半,同时加快训练速度。
梯度累计
在显存受限的情况下,通过多次小批量计算梯度来模拟更大的批量大小。
import torch
import torch.nn as nn
import torch.optim as optim# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用的设备: {device}")# 定义神经网络
class SimpleConvNet(nn.Module):def __init__(self, input_channels, num_classes):super(SimpleConvNet, self).__init__()self.conv1 = nn.Conv2d(input_channels, 64, kernel_size=3, bias=True)self.relu1 = nn.ReLU()self.conv2 = nn.Conv2d(64, 128, kernel_size=3, bias=True)self.relu2 = nn.ReLU()self.fc = nn.Linear(128 * 28 * 28, num_classes)def forward(self, x):out = self.conv1(x)out = self.relu1(out)out = self.conv2(out)out = self.relu2(out)out = out.view(out.size(0), -1)out = self.fc(out)return out# 实例化模型
input_channels = 3
num_classes = 10
model = SimpleConvNet(input_channels, num_classes).to(device)
print(model)# 创建输入数据和标签
batch_size = 16 # 减小批量大小
inputs = torch.randn(batch_size, input_channels, 32, 32).to(device)
labels = torch.randint(0, num_classes, (batch_size,)).to(device)# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 梯度累积步数
accumulation_steps = 4 # 累积4次梯度,相当于批量大小64# 训练多个epoch,并查看显存使用情况
num_epochs = 2
for epoch in range(num_epochs):optimizer.zero_grad()for step in range(accumulation_steps):# 前向传播outputs = model(inputs)loss = criterion(outputs, labels)loss = loss / accumulation_steps # 平均损失loss.backward()optimizer.step()print(f"Epoch {epoch+1}, 损失: {loss.item() * accumulation_steps}")# 查看显存使用情况if torch.cuda.is_available():allocated = torch.cuda.memory_allocated(device) / (1024 ** 2)reserved = torch.cuda.memory_reserved(device) / (1024 ** 2)print(f"Epoch {epoch+1}: 已分配显存: {allocated:.2f} MB, 保留显存: {reserved:.2f} MB")
效果:通过梯度累积,可以在保持较小批量大小的同时,实现更大的有效批量大小,从而在显存受限的情况下提高模型性能。
使用torch.no_grad()
进行评估
在评估模型时,不需要计算梯度,可以节省显存。
with torch.no_grad():outputs = model(inputs)loss = criterion(outputs, labels)
效果:在评估阶段,显存使用会有所减少,因为不需要存储激活值和梯度信息。
使用梯度检查点(Gradient Checkpointing)
梯度检查点通过在反向传播时重新计算部分激活来节省显存,但会增加计算时间。PyTorch提供了 torch.utils.checkpoint
模块来实现这一功能。
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.checkpoint as checkpoint# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用的设备: {device}")# 定义使用梯度检查点的神经网络
class CheckpointConvNet(nn.Module):def __init__(self, input_channels, num_classes):super(CheckpointConvNet, self).__init__()self.conv1 = nn.Conv2d(input_channels, 64, kernel_size=3, bias=True)self.relu1 = nn.ReLU()self.conv2 = nn.Conv2d(64, 128, kernel_size=3, bias=True)self.relu2 = nn.ReLU()self.fc = nn.Linear(128 * 28 * 28, num_classes)def forward(self, x):out = checkpoint.checkpoint(self.conv1, x)out = self.relu1(out)out = checkpoint.checkpoint(self.conv2, out)out = self.relu2(out)out = out.view(out.size(0), -1)out = self.fc(out)return out# 实例化模型
input_channels = 3
num_classes = 10
model = CheckpointConvNet(input_channels, num_classes).to(device)
print(model)# 创建输入数据和标签
batch_size = 64
inputs = torch.randn(batch_size, input_channels, 32, 32).to(device)
labels = torch.randint(0, num_classes, (batch_size,)).to(device)# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练多个epoch,并查看显存使用情况
num_epochs = 2
for epoch in range(num_epochs):optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()print(f"Epoch {epoch+1}, 损失: {loss.item()}")# 查看显存使用情况if torch.cuda.is_available():allocated = torch.cuda.memory_allocated(device) / (1024 ** 2)reserved = torch.cuda.memory_reserved(device) / (1024 ** 2)print(f"Epoch {epoch+1}: 已分配显存: {allocated:.2f} MB, 保留显存: {reserved:.2f} MB")
效果:使用梯度检查点可以显著减少显存占用,特别是在深层网络中。但需要注意,计算时间会增加,因为部分激活在反向传播时需要重新计算。
总结
以上就是关于神经网络中如何占用显存的计算以及代码示例,需要说明的是理论和现实仍有一定的不同,因为在实际代码运行过程中,可能会由于不同的运行环境、不同的PyTorch版本以及不同的系统会有不同的优化策略。大家在实际网络训练过程中,要结合自己的任务来调整配置来更好的训练自己的网络。