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

摸鱼仙人~2024-08-15 15:24

文章目录

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

1. 数据准备

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

示例：加载 CIFAR-10 数据集
python 复制代码 
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 方法中定义前向传播过程。

示例：定义一个简单的卷积神经网络
python 复制代码 
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()

3. 损失函数和优化器

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

示例：定义损失函数和优化器
python 复制代码 
import torch.optim as optim

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 定义优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)

4. 训练循环

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

示例：训练循环
python 复制代码 
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. 评估和测试

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

示例：评估模型
python 复制代码 
def evaluate(model, test_loader, criterion):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0
    with 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. 保存和加载模型

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

示例：保存和加载模型
python 复制代码 
# 保存模型
torch.save(model.state_dict(), 'model.pth')

# 加载模型
model = SimpleCNN()
model.load_state_dict(torch.load('model.pth'))

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

python 复制代码 
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 x

model = 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.0
    correct = 0
    total = 0
    with 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. 保存和加载模型：训练完成后，保存模型参数以便后续使用。
相关推荐
说私域2 分钟前
基于开源2 + 1链动模式AI智能名片S2B2C商城小程序的内容创作与传播效能探究
人工智能·小程序·开源
想拿高薪的韭菜1 小时前
人工智能第2章-知识点与学习笔记
人工智能·笔记·学习
雾岛心情1 小时前
【AIGC专栏】AI在自然语言中的应用场景
人工智能·chatgpt·aigc
Jet45052 小时前
玩转ChatGPT：DeepSeek测评（科研思路梳理）
人工智能·chatgpt·kimi·deepseek-r1
雾散睛明2 小时前
尝试ai生成figma设计
人工智能·figma
车载诊断技术3 小时前
基于新一代电子电器架构的SOA服务设计方法
人工智能·架构·汽车·计算机外设·ecu故障诊断指南
Luzem03193 小时前
使用朴素贝叶斯对自定义数据集进行分类
人工智能·机器学习
小菜鸟博士3 小时前
手撕Vision Transformer -- Day1 -- 基础原理
人工智能·深度学习·学习·算法·面试
找方案3 小时前
智慧城市（城市大脑）建设方案
人工智能·智慧城市·城市大脑
老艾的AI世界3 小时前
AI定制祝福视频，广州塔、动态彩灯、LED表白，直播互动新玩法（附下载链接）
图像处理·人工智能·深度学习·神经网络·目标检测·机器学习·ai·ai视频·ai视频生成·ai视频制作
热门推荐
01DeepSeek本地部署详细指南02【deepseek】deepseek-r1本地部署-第一步：下载LM Studio03DeepSeek r1本地安装全指南04将DeepSeek接入Word，打造AI办公助手05使用Ollama 在Ubuntu运行deepseek大模型：以DeepSeek-coder为例06DeepSeek学术写作测评第一弹：论文润色，中译英效果如何？07DeepSeek 云端部署，释放无限 AI 潜力！08新鲜速递：DeepSeek-R1开源大模型本地部署实战—Ollama + MaxKB 搭建RAG检索增强生成应用09Dell服务器升级ubuntu 22.04失败解决10如何本地部署DeepSeek