简单的使用PyTorch训练模型的项目示例代码

PyTorch是一个开源的深度学习框架，它提供了一系列先进的技术和功能，使得机器学习和深度学习模型的开发、训练、调试和部署变得更加高效和灵活。

本博客持续关注和应用云原生、AI、数字孪生、智慧城市等技术，请关注其他文章或持续关注本博客的技术更新。

这里从项目中抽象出一个如何在 PyTorch 中进行一般的图像分类任务训练示例：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import transforms

# 定义一个简单的卷积神经网络模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc = nn.Linear(16 * 16 * 16, num_classes)  # 假设输出类别数为num_classes

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(-1, 16 * 16 * 16)
        x = self.fc(x)
        return x

# 初始化模型、损失函数和优化器
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 数据预处理及加载数据集
data_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_dataset = ImageFolder('path/to/train_data', transform=data_transforms)
test_dataset = ImageFolder('path/to/test_data', transform=data_transforms)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

# 开始训练过程
num_epochs = 10
for epoch in range(num_epochs):
    model.train()  # 进入训练模式
    for inputs, labels in train_loader:
        optimizer.zero_grad()  # 清零梯度
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()  # 反向传播计算梯度
        optimizer.step()  # 更新权重参数

    print(f'Epoch {epoch + 1}/{num_epochs} - Loss: {loss.item():.4f}')

# 在测试集上评估模型
model.eval()  # 进入评估模式
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'Test Accuracy: {accuracy}%')