您的位置:首页 > 科技 > IT业 > Pytorch常用训练套路框架(CPU)

Pytorch常用训练套路框架(CPU)

2024/12/23 10:47:48 来源:https://blog.csdn.net/qq_52964132/article/details/141181457  浏览:    关键词:Pytorch常用训练套路框架(CPU)

文章目录

      • 1. 数据准备
        • 示例:加载 CIFAR-10 数据集
      • 2. 模型定义
        • 示例:定义一个简单的卷积神经网络
      • 3. 损失函数和优化器
        • 示例:定义损失函数和优化器
      • 4. 训练循环
        • 示例:训练循环
      • 5. 评估和测试
        • 示例:评估模型
      • 6. 保存和加载模型
        • 示例:保存和加载模型
      • 7. 完整案例:训练 CIFAR-10 分类模型
      • 解释


在 PyTorch 中,模型训练通常遵循一个标准的流程,包括数据准备、模型定义、损失函数和优化器的选择、训练循环以及评估和测试。以下是一个详细的步骤介绍:

1. 数据准备

首先,需要准备好训练和测试数据。通常使用 torchvision.datasets 加载内置数据集,或者使用自定义数据集。数据加载后,使用 torch.utils.data.DataLoader 进行批量加载。

示例:加载 CIFAR-10 数据集
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)

2. 模型定义

定义一个神经网络模型,通常继承自 torch.nn.Module,并在 __init__ 方法中定义网络层,在 forward 方法中定义前向传播过程。

示例:定义一个简单的卷积神经网络
import torch.nn as nn
import torch.nn.functional as Fclass 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 xmodel = SimpleCNN()

3. 损失函数和优化器

选择合适的损失函数和优化器。常见的损失函数包括 nn.CrossEntropyLoss 用于分类任务,nn.MSELoss 用于回归任务。优化器通常使用 torch.optim 模块中的优化器,如 optim.SGDoptim.Adam

示例:定义损失函数和优化器
import torch.optim as optim# 定义损失函数
criterion = nn.CrossEntropyLoss()# 定义优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)

4. 训练循环

编写训练循环,包括前向传播、计算损失、反向传播和参数更新。通常还会包括模型保存和日志记录。

示例:训练循环
def train(model, train_loader, criterion, optimizer, num_epochs):model.train()for epoch in range(num_epochs):for images, labels in train_loader:# 前向传播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)

5. 评估和测试

在训练完成后,使用测试数据集评估模型的性能。通常包括计算准确率、损失等指标。

示例:评估模型
def evaluate(model, test_loader, criterion):model.eval()total_loss = 0.0correct = 0total = 0with torch.no_grad():for images, labels in test_loader: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)

6. 保存和加载模型

训练完成后,可以保存模型参数以便后续使用。

示例:保存和加载模型
# 保存模型
torch.save(model.state_dict(), 'model.pth')# 加载模型
model = SimpleCNN()
model.load_state_dict(torch.load('model.pth'))

7. 完整案例:训练 CIFAR-10 分类模型

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader# 1. 数据准备
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)train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)# 2. 模型定义
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 xmodel = SimpleCNN()# 3. 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 4. 训练循环
def train(model, train_loader, criterion, optimizer, num_epochs):model.train()for epoch in range(num_epochs):for images, labels in train_loader: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)# 5. 评估和测试
def evaluate(model, test_loader, criterion):model.eval()total_loss = 0.0correct = 0total = 0with torch.no_grad():for images, labels in test_loader: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)# 6. 保存和加载模型
torch.save(model.state_dict(), 'model.pth')model = SimpleCNN()
model.load_state_dict(torch.load('model.pth'))

解释

  1. 数据准备:加载 CIFAR-10 数据集,并应用一系列图像转换操作。
  2. 模型定义:定义一个简单的卷积神经网络 SimpleCNN
  3. 损失函数和优化器:选择交叉熵损失函数和 Adam 优化器。
  4. 训练循环:编写训练循环,包括前向传播、计算损失、反向传播和参数更新。
  5. 评估和测试:使用测试数据集评估模型的性能,并计算准确率和损失。
  6. 保存和加载模型:训练完成后,保存模型参数以便后续使用。

版权声明:

本网仅为发布的内容提供存储空间,不对发表、转载的内容提供任何形式的保证。凡本网注明“来源:XXX网络”的作品,均转载自其它媒体,著作权归作者所有,商业转载请联系作者获得授权,非商业转载请注明出处。

我们尊重并感谢每一位作者,均已注明文章来源和作者。如因作品内容、版权或其它问题,请及时与我们联系,联系邮箱:809451989@qq.com,投稿邮箱:809451989@qq.com