torch 指定gpu进行训练

Handsome_zhou · 发表于 2023-7-26 16:57:27

马上注册，结交更多好友，享用更多功能^_^

您需要登录才可以下载或查看，没有账号？立即注册

x

本帖最后由 Handsome_zhou 于 2023-9-21 10:52 编辑

总结：如果设置了环境变量，模型和张量就用.cuda()方法进行GPU迁移。否则使用.to(device)

1、sh文件中使用CUDA_VISIBLE_DEVICE环境变量:
在训练脚本运行之前，可以设置'CUDA_VISIBLE_DEVICES'来指定使用哪个GPU。这是最简单的方式，适用于单个或多个GPU的情况。
例如，当有单机多GPU时，可以在sh文件中将以下命令添加到训练脚本之前，以选择第一块GPU：

export CUDA_VISIBLE_DEVICES=0
python your_training_script.py

复制代码

2、py文件中使用torch.device对象:
可以使用torch.device对象将模型和张量移动到特定的GPU上。这样可以在代码中灵活地控制GPU的选择。

import torch
device = torch.device("cuda:0") # 选择第一块 GPU
#或者使用 device = 0 if torch.cuda.is_available() else -1
model = YourModel().to(device)
input_tensor = input_tensor.to(device)

复制代码

3、在命令行中指定GPU设备:
可以在命令行中通过传递'--device'或者类似的参数来指定要使用的GPU设备。然后在训练脚本中解析这些参数并将模型和数据移到相应的设备上。

python your_training_script.py --device 0

复制代码

4、使用torch.nn.DataParallel或者torch.nn.parallel.DistributedDataParallel进行多GPU训练。
这两种方法可以在多个GPU上进行模型的正向和反向传播，使训练更快。

import torch
from torch.nn import DataParallel
device_ids = [0, 1] # 指定要使用的 GPU 设备
model = YourModel()
model = DataParallel(model, device_ids=device_ids)
model.to(device_ids[0]) # 将模型移动到第一个 GPU

复制代码

4.1、使用torch.nn.DataParallel进行多GPU训练的示例：

# 测试可用
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
import os
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
# 设置CUDA可见设备
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
modules = [
nn.Linear(10, 3),
nn.Linear(3, 4),
nn.Linear(4, 5),
]
self.net = nn.ModuleList(modules)
def forward(self, inputs):
for i, n in enumerate(self.net):
inputs = n(inputs)
return inputs
def main():
X = np.random.uniform(-1, 1, (15, 10)).astype(np.float32)
y = np.random.randint(0, 5, (15,))
print(X.shape)
print(y.shape)
model = Net()
loss = nn.CrossEntropyLoss()
print('Model:', type(model))
print('Loss:', type(loss))
X = torch.from_numpy(X)
y = torch.from_numpy(y)
print('X', X.size(), 'y', y.size())
if torch.cuda.is_available():
# 将模型移到GPU上，并使用DataParallel
model = nn.DataParallel(model, device_ids=[0, 1])
model = model.cuda()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
loss = loss.cuda()
X = X.cuda()
y = y.cuda()
print('Using', torch.cuda.device_count(), 'GPUs')
else:
print('No GPUs available')
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
epochs = 10000
for epoch in range(epochs):
running_loss = 0.0
for i, data in enumerate(dataloader, 0):
inputs, labels = data
optimizer.zero_grad()
# 前向传播
outputs = model(inputs)
loss_value = loss(outputs, labels)
# 反向传播和优化
loss_value.backward()
optimizer.step()
running_loss += loss_value.item()
# print(f"第 {i+1} 批，损失: {loss_value.item()}")
print(f"第 {epoch+1} 轮，平均损失: {running_loss / len(dataloader)}")
print("训练结束")
if __name__ == '__main__':
main()

复制代码

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DataParallel
# 设置 CUDA_VISIBLE_DEVICES 环境变量来指定使用 GPU 0 和 1
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
# 定义一个简单的模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
# 创建模型并将其包装在 DataParallel 中
model = SimpleModel()
model = DataParallel(model)
# 创建数据加载器和优化器
# 注意：您需要适应您的数据加载逻辑和损失函数
dataloader = YourDataLoader()
optimizer = optim.SGD(model.parameters(), lr=0.001)
# 模拟训练循环
for epoch in range(num_epochs):
for batch_data, batch_labels in dataloader:
# batch_data, batch_labels 已经在 DataParallel 内部自动分发到多个 GPU
optimizer.zero_grad()
outputs = model(batch_data)
loss = YourLossFunction(outputs, batch_labels)
loss.backward()
optimizer.step()
# 保存模型
torch.save(model.module.state_dict(), 'model.pth') # 使用 model.module 获得原始模型的状态字典

复制代码

4.2、使用torch.nn.parallel.DistributedDataParallel进行多GPU训练的示例：
使用 torch.nn.parallel.DistributedDataParallel（DDP）进行多 GPU 训练通常需要在多个进程中运行，其中每个进程可以使用一个或多个 GPU。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel
# 定义一个简单的模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
# 定义一个训练函数，该函数在每个进程中运行
def train(rank, world_size):
# 设置每个进程的 GPU 设备，这里假设有4个GPU
os.environ['CUDA_VISIBLE_DEVICES'] = str(rank)
# 初始化分布式环境
dist.init_process_group(backend='nccl', init_method='env://')
# 创建模型并将其包装在 DistributedDataParallel 中
model = SimpleModel().cuda()
model = DistributedDataParallel(model)
# 创建数据加载器和优化器
# 注意：您需要适应您的数据加载逻辑和损失函数
dataloader = YourDataLoader()
optimizer = optim.SGD(model.parameters(), lr=0.001)
# 模拟训练循环
for epoch in range(num_epochs):
for batch_data, batch_labels in dataloader:
optimizer.zero_grad()
outputs = model(batch_data)
loss = YourLossFunction(outputs, batch_labels)
loss.backward()
optimizer.step()
if __name__ == '__main__':
# 设置分布式训练的参数
num_processes = 4 # 四个进程，每个进程使用一个 GPU
mp.spawn(train, args=(num_processes, ), nprocs=num_processes)

复制代码

5、py文件中通过os.environ['CUDA_VISIBLE_DEVICES']指定GPU进行训练
使用'CUDA_VISIBLE_DEVICES'环境变量来选择要使用的GPU，然后使用.cuda()方法将模型和张量移到指定的GPU上。

import os
import torch
# 设置 CUDA_VISIBLE_DEVICES 环境变量来指定 GPU
os.environ['CUDA_VISIBLE_DEVICES'] = '0' # 选择第一块 GPU
# 创建一个简单的模型并将其移到 GPU 上
model = torch.nn.Linear(10, 1).cuda() # 这里将模型移到指定的 GPU 上
# 创建一个示例输入并将其移到 GPU 上
input_tensor = torch.randn(1, 10).cuda()
# 运行模型
output = model(input_tensor)
# 确保输出是在指定的 GPU 上
print(output.device) # 输出应该是 "cuda:0"

复制代码

歌者文明清理员 · 发表于 2023-7-26 17:06:21

你这个头像绝了

账号		自动登录	找回密码
密码			立即注册

[学习笔记] torch 指定gpu进行训练

马上注册，结交更多好友，享用更多功能^_^

浏览过的版块