指定GPU运行和训练 Python 程序,深度学习单卡、多卡训练GPU设置
在进行深度学习任务时,GPU的使用是提高训练速度和效果的重要手段之一。在Python中,我们可以通过一些方法来指定GPU的运行和训练。
指定GPU运行
当我们使用多个GPU进行训练时,有时需要手动指定程序运行在哪个GPU上。这可以通过设置CUDA_VISIBLE_DEVICES环境变量来实现。下面是一个示例代码:
import os
# 设置CUDA_VISIBLE_DEVICES环境变量,指定程序运行在第一块GPU上
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# 后续的代码将会在第一块GPU上运行
在这个示例中,我们使用了os.environ
来设置环境变量CUDA_VISIBLE_DEVICES
,并将其值设为0
,表示程序将会在第一块GPU上运行。
深度学习单卡训练GPU设置
在进行深度学习训练时,我们通常会使用GPU来加速模型的训练过程。下面是一个示例代码,展示了如何将深度学习模型加载到GPU上进行训练:
import torch
# 创建一个深度学习模型
model = MyModel()
# 将模型加载到GPU上
model = model.cuda()
# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 加载训练数据
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
# 进行训练
for epoch in range(10):
for images, labels in train_loader:
# 将数据加载到GPU上
images = images.cuda()
labels = labels.cuda()
# 前向传播
outputs = model(images)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
在这个示例中,我们首先创建了一个深度学习模型MyModel
,然后使用model.cuda()
将模型加载到GPU上。接下来,我们定义了损失函数和优化器,并使用torch.utils.data.DataLoader
加载训练数据。在训练过程中,我们将数据也加载到GPU上,通过images.cuda()
和labels.cuda()
实现。这样,模型的前向传播、计算损失和反向传播的过程都会在GPU上进行,加速了训练过程。
深度学习多卡训练GPU设置
当我们有多个GPU可用时,可以使用数据并行的方式在多个GPU上进行训练。下面是一个示例代码,展示了如何在多个GPU上进行深度学习模型的训练:
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
# 定义一个多卡模型
class MultiGPUModel(nn.Module):
def __init__(self):
super(MultiGPUModel, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
self.fc1 = nn.Linear(64 * 16 * 16, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = x.view(x.size(0), -1)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
# 创建一个多卡模型实例
model = MultiGPUModel()
# 使用数据并行的方式在多个GPU上进行训练
model = nn.DataParallel(model)
# 将模型加载到GPU上
model = model.cuda()
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 加载训练数据
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)