基于PyTorch的彩色图片分类卷积神经网络

import os
import torch
from torch import nn, optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 设置训练集和测试集的文件夹路径
train_dir = 'path/to/train/directory'
test_dir = 'path/to/test/directory'

# 获取训练集图片文件名列表
train_files = os.listdir(train_dir)

# 推断测试集图片的命名规律
test_files = []
for file_name in train_files:
    test_file_name = file_name.replace('train', 'test')
    test_files.append(test_file_name)

# 数据预处理
transform = transforms.Compose([
    transforms.Resize(28),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载训练集和测试集
train_data = datasets.ImageFolder(train_dir, transform=transform)
test_data = datasets.ImageFolder(test_dir, transform=transform)

trainloader = DataLoader(train_data, batch_size=64, shuffle=True)
testloader = DataLoader(test_data, batch_size=64, shuffle=True)

# 定义卷积神经网络模型
class ColorCNN(nn.Module):
    def __init__(self):
        super(ColorCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        self.dropout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(128, 2)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 32 * 7 * 7)
        x = self.dropout(x)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 创建模型实例并定义优化器和损失函数
model = ColorCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(10):
    running_loss = 0.0
    for images, labels in trainloader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch+1} - Training loss: {running_loss / len(trainloader)}')

# 模型评估
correct = 0
total = 0
with torch.no_grad():
    for images, labels in testloader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += torch.eq(predicted, labels).sum().item()
accuracy = correct / total
print(f'Test accuracy: {accuracy}')

# 打印测试集图片文件名列表
print('测试集图片文件名列表:', test_files)

代码说明:

导入必要库: 导入 os, torch, torchvision 和 torch.utils.data 库，用于文件操作、深度学习、图像处理和数据加载。
设置路径: 设置训练集和测试集的文件夹路径。
获取文件名: 获取训练集和测试集的图片文件名列表。
数据预处理: 使用 transforms.Compose 定义数据预处理步骤，包括图像大小调整、转换为张量和归一化。
加载数据: 使用 datasets.ImageFolder 加载训练集和测试集，并使用 DataLoader 创建数据加载器，用于批量加载数据。
定义模型: 定义一个名为 ColorCNN 的卷积神经网络模型，该模型包含两个卷积层、两个池化层、两个全连接层、ReLU 激活函数和 Dropout 层。
创建实例: 创建模型实例，定义优化器和损失函数。
训练模型: 循环遍历训练集，进行模型训练，并打印每个 epoch 的训练损失。
模型评估: 在测试集上评估模型性能，计算准确率。
打印文件名: 打印测试集图片文件名列表，用于验证数据加载。