【pytorch实用小技巧】单gpu与多gpu训练与评估

1、单gpu

• 首先检查GPU是否可用，并将模型、输入数据和目标标签移动到GPU上。
• 然后，定义损失函数和优化器。在训练循环中，将模型设置为训练模式，进行前向传播、计算损失、反向传播和参数更新。
• 在测试阶段，将模型设置为评估模式，并在测试数据上进行推断。

import torch
import torch.nn as nn
import torch.optim as optim

# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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

# 创建模型实例并将其移动到GPU上
model = Net().to(device)

# 创建输入数据和目标标签，并将其移动到GPU上
input_data = torch.randn(100, 10).to(device)
target = torch.randn(100, 1).to(device)

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

# 进行训练
for epoch in range(100):
    # 将模型设置为训练模式
    model.train()

    # 清零梯度
    optimizer.zero_grad()
    
    # 前向传播
    output = model(input_data)
    
    # 计算损失
    loss = criterion(output, target)
    
    # 反向传播
    loss.backward()
    
    # 更新模型参数
    optimizer.step()

    # 输出当前训练的损失值
    print(f"Epoch [{
      
      epoch+1}/100], Loss: {
      
      loss.item():.4f}")

# 将模型设置为评估模式
model.eval()

# 在测试数据上进行推断
test_data = torch.randn(10, 10).to(device)
with torch.no_grad():
    output = model(test_data)
    print("Inference result:", output)

2、多gpu

• 首先检查是否有多个GPU可用，如果没有至少两个GPU，则会引发错误。
• 然后创建模型实例，并使用nn.DataParallel进行多GPU并行训练。将模型移动到GPU上。输入数据和目标标签也要移动到GPU上。
• 然后定义损失函数和优化器。在训练循环中，将模型设置为训练模式，进行前向传播、计算损失、反向传播和参数更新。
• 在测试阶段，将模型设置为评估模式，并在测试数据上进行推断。请确保您的计算机上至少有两个可用的GPU才能运行此示例。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.parallel

# 检查是否有多个GPU可用
if torch.cuda.device_count() < 2:
    raise RuntimeError("需要至少两个GPU来运行此示例。")

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

# 创建模型实例并将其移到多个GPU上
model = Net()
model = nn.DataParallel(model)  # 使用DataParallel进行多GPU并行训练
model = model.cuda()  # 将模型移动到GPU上

# 创建输入数据和目标标签，并将其移到GPU上
input_data = torch.randn(100, 10).cuda()
target = torch.randn(100, 1).cuda()

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

# 进行训练
for epoch in range(100):
    # 将模型设置为训练模式
    model.train()

    # 清零梯度
    optimizer.zero_grad()
    
    # 前向传播
    output = model(input_data)
    
    # 计算损失
    loss = criterion(output, target)
    
    # 反向传播
    loss.backward()
    
    # 更新模型参数
    optimizer.step()

    # 输出当前训练的损失值
    print(f"Epoch [{
      
      epoch+1}/100], Loss: {
      
      loss.item():.4f}")

# 将模型设置为评估模式
model.eval()

# 在测试数据上进行推断
test_data = torch.randn(10, 10).cuda()
with torch.no_grad():
    output = model(test_data)
    print("Inference result:", output)