如何在Torch中实现多GPU训练-开发者知识库平台

如何在Torch中实现多GPU训练

Torch

775

2024/3/27 19:40:00

栏目: 深度学习

在Torch中实现多GPU训练可以通过使用nn.DataParallel模块来实现。nn.DataParallel模块可以将模型复制到多个GPU上，并自动将数据分发到每个GPU上进行计算，并最终将结果进行合并。以下是一个简单示例：

import torch
import torch.nn as nn

# 定义模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

model = Model()

# 检查是否有多个GPU可用
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

# 将模型加载到GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 定义优化器和损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# 训练模型
for epoch in range(num_epochs):
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

在上面的示例中，首先检查是否有多个GPU可用，如果有，则将模型放入nn.DataParallel中。然后将模型加载到GPU上进行训练。在训练过程中，数据会自动分发到每个GPU上，并在多个GPU上并行计算，从而加快训练速度。

辰迅云「云服务器」，即开即用、新一代英特尔至强铂金CPU、三副本存储NVMe SSD云盘，价格低至29元/月。点击查看>>

如何在Torch中实现多GPU训练

最新知识库

相关标签