赞
踩
在 PyTorch 中,使用 GPU 进行模型训练与仅使用 CPU 相比,需要对代码进行一些特定的修改。这些修改主要涉及以下几个方面:
首先,需要检查系统中是否有可用的 GPU,并选择合适的设备进行训练。
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
将定义好的模型移动到 GPU 上。
model.to(device)
在训练和评估过程中,需要将输入数据和标签移动到 GPU 上。
images, labels = images.to(device), labels.to(device)
在训练循环中,确保每次迭代都将数据移动到 GPU 上。
def train(model, train_loader, criterion, optimizer, num_epochs, device):
model.train()
for epoch in range(num_epochs):
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
在评估循环中,同样需要将数据移动到 GPU 上。
def evaluate(model, test_loader, criterion, device):
model.eval()
total_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
total_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Test Loss: {total_loss/len(test_loader):.4f}, Accuracy: {100 * correct / total:.2f}%')
在保存和加载模型时,确保模型参数在 GPU 上。
# 保存模型
torch.save(model.state_dict(), 'model.pth')
# 加载模型
model = SimpleCNN()
model.load_state_dict(torch.load('model.pth'))
model.to(device)
使用 GPU 进行训练时,需要修改的具体部分包括:
在 PyTorch 中,使用 GPU 进行模型训练可以显著提高训练速度,尤其是在处理大规模数据集和复杂模型时。以下是基于 GPU 训练的一般套路,包括数据准备、模型定义、损失函数和优化器的选择、训练循环、评估和测试,以及模型保存和加载。
首先,需要准备好训练和测试数据。通常使用 torchvision.datasets
加载内置数据集,或者使用自定义数据集。数据加载后,使用 torch.utils.data.DataLoader
进行批量加载。
from torchvision import datasets, transforms from torch.utils.data import DataLoader # 定义图像转换 transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 加载数据集 train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) # 使用 DataLoader 加载数据 train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
定义一个神经网络模型,通常继承自 torch.nn.Module
,并在 __init__
方法中定义网络层,在 forward
方法中定义前向传播过程。
import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) self.fc1 = nn.Linear(64 * 56 * 56, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = F.relu(self.conv1(x)) x = F.max_pool2d(x, 2) x = F.relu(self.conv2(x)) x = F.max_pool2d(x, 2) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = self.fc2(x) return x model = SimpleCNN()
在训练之前,需要检查是否有可用的 GPU,并将模型和数据移动到 GPU 上。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
选择合适的损失函数和优化器。常见的损失函数包括 nn.CrossEntropyLoss
用于分类任务,nn.MSELoss
用于回归任务。优化器通常使用 torch.optim
模块中的优化器,如 optim.SGD
或 optim.Adam
。
import torch.optim as optim
# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)
编写训练循环,包括前向传播、计算损失、反向传播和参数更新。在训练过程中,需要将数据移动到 GPU 上。
def train(model, train_loader, criterion, optimizer, num_epochs, device):
model.train()
for epoch in range(num_epochs):
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
# 训练模型
train(model, train_loader, criterion, optimizer, num_epochs=10, device=device)
在训练完成后,使用测试数据集评估模型的性能。在评估过程中,同样需要将数据移动到 GPU 上。
def evaluate(model, test_loader, criterion, device): model.eval() total_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for images, labels in test_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) total_loss += loss.item() _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f'Test Loss: {total_loss/len(test_loader):.4f}, Accuracy: {100 * correct / total:.2f}%') # 评估模型 evaluate(model, test_loader, criterion, device)
训练完成后,可以保存模型参数以便后续使用。在加载模型时,需要将模型移动到 GPU 上。
# 保存模型
torch.save(model.state_dict(), 'model.pth')
# 加载模型
model = SimpleCNN()
model.load_state_dict(torch.load('model.pth'))
model.to(device)
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。