【pytorch实用小技巧】单gpu与多gpu训练与评估-Toy模板网

这篇具有很好参考价值的文章主要介绍了【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)