在PyTorch中实现一个卷积神经网络(CNN)来处理图像分类问题,并展示如何训练和评估该模型。
时间: 2024-11-11 22:32:55 浏览: 7
要实现一个卷积神经网络(CNN)来处理图像分类问题,首先需要了解PyTorch中CNN的基本构建块。卷积层、池化层和全连接层是CNN的主要组成部分。下面将展示如何定义一个简单的CNN模型,并进行训练和评估的步骤。
参考资源链接:[深度学习入门:PyTorch实战教程](https://wenku.csdn.net/doc/3eev71xbzt?spm=1055.2569.3001.10343)
首先,确保已经安装了PyTorch,并且导入了必要的模块。
```python
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 定义CNN模型
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
self.fc1 = nn.Linear(64 * 8 * 8, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = torch.relu(self.conv1(x))
x = torch.max_pool2d(x, kernel_size=2)
x = torch.relu(self.conv2(x))
x = torch.max_pool2d(x, kernel_size=2)
x = x.view(x.size(0), -1) # Flatten the tensor
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 创建模型实例
model = SimpleCNN()
```
接下来,准备数据集和数据加载器:
```python
transform = ***pose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
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(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)
```
定义损失函数和优化器:
```python
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
```
训练模型:
```python
def train(model, train_loader, optimizer):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
# 训练模型几个周期
for epoch in range(1, epochs + 1):
train(model, train_loader, optimizer)
```
最后,评估模型:
```python
def evaluate(model, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
output = model(data)
test_loss += criterion(output, target).item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print(f
参考资源链接:[深度学习入门:PyTorch实战教程](https://wenku.csdn.net/doc/3eev71xbzt?spm=1055.2569.3001.10343)
阅读全文