使用 PyTorch 实现 CNN 图像分类
5星
- 浏览量: 0
- 大小:None
- 文件类型：PDF
简介：
本项目利用PyTorch框架实现卷积神经网络(CNN)进行图像分类任务。通过训练CNN模型，能够有效识别和分类不同类别的图像数据。
在4*4的图片中比较外围黑色像素点与内圈黑色像素点的数量，并根据数量差异将图片分类为两类：如果外围黑色像素点多于内圈，则归类为0，反之则为1。
为了实现这一任务，可以使用numpy和PIL库生成随机的二值图像数据集。首先通过`buildDataset`函数创建一个包含4*4大小、黑白两种颜色（分别用数值0和1表示）的数据集，并将这些图片保存为.jpg格式文件；同时计算每个图中外围与内圈黑色像素的数量差，以此作为标签信息。
接下来需要设计自定义数据集类`MyDataset`继承于`torch.utils.data.Dataset`。该类从CSV文件加载图像路径及对应的分类标签，并提供必要的方法支持批量读取和处理功能，例如使用预处理器调整图片尺寸或归一化等操作。
在构建CNN模型时有两种方案：一种是在4*4的输入上直接应用1x1卷积层来提取特征；另一种是通过给原始图像添加padding使其变为6*6大小后采用2x2的卷积核进行处理，最终输出一个3*3的结果图。这两种方法都可以连接全连接层完成分类任务。
具体的PyTorch模型定义如下：
```python
import torch.nn as nn
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
# 1x1卷积方案：
# self.conv = nn.Conv2d(1, 8, kernel_size=1)
# 或者
self.conv1 = nn.Conv2d(1, 8, kernel_size=2, padding=1)
self.relu = nn.ReLU()
self.fc1 = nn.Linear(8 * 3 * 3, 50)
self.fc2 = nn.Linear(50, 2)
def forward(self, x):
# 使用卷积层
x = self.conv1(x)
x = self.relu(x)
# 展平特征图并进行全连接操作
x = torch.flatten(x, start_dim=1)
x = self.fc1(x)
output = self.fc2(x)
return output
```
在训练阶段，使用交叉熵损失函数和随机梯度下降（SGD）优化器。通过`DataLoader`加载数据集，并进行多个epoch的迭代以更新模型参数。
```python
device = torch.device(cuda if torch.cuda.is_available() else cpu)
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for epoch in range(epochs):
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 验证阶段
with torch.no_grad():
total_correct = 0
for val_images, val_labels in val_loader:
output_val = model(val_images.to(device))
_, predicted_label = torch.max(output_val.data, dim=1)
total_correct += (predicted_label.cpu() == val_labels).sum().item()
accuracy = total_correct / len(valData)
print(fEpoch {epoch+1}/{epochs}, Validation Accuracy: {accuracy})
```
此示例展示了如何从头开始构建一个图像分类任务的全部流程，包括数据集生成、自定义数据类的设计以及CNN模型架构的选择和实现。
全部评论 (0)
