使用PyTorch 构建一个简单的图片分类模型。一个使用卷积神经网络（CNN）进行图像分类的基本示例。包括模型的定义、数据加载和训练过程。不多说直接上代码
首先，确保您已经安装了 PyTorch。如果还没有安装，可以通过以下命令安装：
pip install torch torchvision
import torch
2import torch.nn as nn
3import torch.optim as optim
4from torchvision import datasets, transforms
5from torch.utils.data import DataLoader
6
7# 定义一个简单的CNN模型
8class SimpleCNN(nn.Module):
9 def __init__(self, num_classes=10):
10 super(SimpleCNN, self).__init__()
11 self.layer1 = nn.Sequential(
12 nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),
13 nn.BatchNorm2d(16),
14 nn.ReLU(),
15 nn.MaxPool2d(kernel_size=2, stride=2)
16 )
17 self.layer2 = nn.Sequential(
18 nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
19 nn.BatchNorm2d(32),
20 nn.ReLU(),
21 nn.MaxPool2d(kernel_size=2, stride=2)
22 )
23 self.fc = nn.Linear(32 * 16 * 16, num_classes)
24
25 def forward(self, x):
26 out = self.layer1(x)
27 out = self.layer2(out)
28 out = out.reshape(out.size(0), -1)
29 out = self.fc(out)
30 return out
31
32# 数据预处理和加载
33transform = transforms.Compose([
34 transforms.Resize((32, 32)), # 根据您的数据调整大小
35 transforms.ToTensor(),
36 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
37])
38
39# 假设您使用的是CIFAR-10数据集
40train_dataset = datasets.CIFAR10(root='./data', train=True,
41 download=True, transform=transform)
42test_dataset = datasets.CIFAR10(root='./data', train=False,
43 download=True, transform=transform)
44
45train_loader = DataLoader(train_dataset, batch_size=100, shuffle=True, num_workers=2)
46test_loader = DataLoader(test_dataset, batch_size=100, shuffle=False, num_workers=2)
47
48# 初始化模型、损失函数和优化器
49device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
50model = SimpleCNN(num_classes=10).to(device)
51criterion = nn.CrossEntropyLoss()
52optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
53
54# 训练模型
55num_epochs = 10
56for epoch in range(num_epochs):
57 running_loss = 0.0
58 for i, data in enumerate(train_loader, 0):
59 inputs, labels = data
60 inputs, labels = inputs.to(device), labels.to(device)
61
62 optimizer.zero_grad()
63
64 outputs = model(inputs)
65 loss = criterion(outputs, labels)
66 loss.backward()
67 optimizer.step()
68
69 running_loss += loss.item()
70 print(f'Epoch {epoch + 1}, Loss: {running_loss / (i + 1)}')
71
72print('Finished Training.')
73
74# 这里可以添加测试代码来评估模型性能
