当前位置：首页 > news >正文

神经网络课设

news 来源：原创 2025/6/17 10:38:23

1. 使用 `ImageFolder` 加载按文件夹分类的数据集

适用于数据集按 “类别 / 图像” 的结构组织（如 CIFAR-2、Caltech101 等）。

目录结构要求：

plaintext

dataset/
├── train/
│   ├── class1/
│   │   ├── img1.jpg
│   │   └── ...
│   └── class2/
│       ├── img2.jpg
│       └── ...
└── test/├── class1/└── class2/

代码示例：

from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, Resize, ToTensor, Normalize
from torch.utils.data import DataLoader# 定义数据预处理
transform = Compose([Resize((224, 224)),  # 调整图像大小为224×224ToTensor(),          # 转换为Tensor并归一化到[0,1]Normalize(           # 标准化（使用ImageNet的均值和标准差）mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
])# 加载数据集
train_dataset = ImageFolder(root='./dataset/train',transform=transform
)test_dataset = ImageFolder(root='./dataset/test',transform=transform
)# 创建数据加载器
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
)# 查看数据集信息
print(f"训练集大小: {len(train_dataset)}")
print(f"类别映射: {train_dataset.class_to_idx}")# 示例：获取一个批次的数据
images, labels = next(iter(train_loader))
print(f"图像形状: {images.shape}")  # 输出: [32, 3, 224, 224]
print(f"标签: {labels}")

完整代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
import time
import os
from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, Resize, ToTensor, Normalize# 解决Windows多进程问题
if __name__ == '__main__':import torch.multiprocessing as mpmp.set_start_method('spawn', force=True)# 定义数据预处理
transform = Compose([Resize((32, 32)),  # 调整为原始CIFAR尺寸，避免计算量过大ToTensor(),Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
])# 加载本地2分类数据集
train_dataset = ImageFolder(root='D:\\BaiduNetdiskDownload\\cifar2_datasets\\train',transform=transform
)
test_dataset = ImageFolder(root='D:\\BaiduNetdiskDownload\\cifar2_datasets\\test',transform=transform
)# 创建数据加载器
train_loader = DataLoader(train_dataset,batch_size=100,shuffle=True,num_workers=4,pin_memory=True
)
test_loader = DataLoader(test_dataset,batch_size=100,shuffle=False,num_workers=4,pin_memory=True
)# 查看数据集信息
def look_data():print(f"训练集大小: {len(train_dataset)}")print(f"类别映射: {train_dataset.class_to_idx}")print(f"测试集大小: {len(test_dataset)}")print(f"类别映射: {test_dataset.class_to_idx}")# 验证数据加载images, labels = next(iter(train_loader))print(f"图像形状: {images.shape}")print(f"标签: {labels}")# 定义2分类模型
class ImageClassification(nn.Module):def __init__(self):super(ImageClassification, self).__init__()# 卷积池化层self.conv1 = nn.Conv2d(3, 6, kernel_size=3, stride=1, padding=1)self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)self.conv2 = nn.Conv2d(6, 16, kernel_size=3, stride=1, padding=1)self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)# 计算全连接层输入维度（根据图像尺寸32x32计算）# 32x32 -> 卷积1(3x3,p=1) -> 32x32 -> 池化1 -> 16x16# 卷积2(3x3,p=1) -> 16x16 -> 池化2 -> 8x8# 8x8x16 = 1024self.linear1 = nn.Linear(1024, 120)self.linear2 = nn.Linear(120, 84)self.out = nn.Linear(84, 2)  # 输出层为2分类def forward(self, x):x = self.pool1(F.relu(self.conv1(x)))x = self.pool2(F.relu(self.conv2(x)))x = x.view(x.size(0), -1)x = F.relu(self.linear1(x))x = F.relu(self.linear2(x))return self.out(x)# 训练函数（修改为使用本地数据集）
def train():print("开始训练2分类模型...")model = ImageClassification()criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=1e-3)epochs = 10# 创建保存模型的目录if not os.path.exists('model'):os.makedirs('model')for epoch_idx in range(epochs):model.train()  # 设置为训练模式start = time.time()total_loss = 0.0correct = 0total_samples = 0for images, labels in train_loader:# 前向传播outputs = model(images)loss = criterion(outputs, labels)# 反向传播和优化optimizer.zero_grad()loss.backward()optimizer.step()# 统计total_loss += loss.item() * len(labels)correct += (torch.argmax(outputs, dim=1) == labels).sum().item()total_samples += len(labels)# 打印训练进度epoch_loss = total_loss / total_samplesepoch_acc = correct / total_samplesprint(f'epoch: {epoch_idx + 1}/{epochs}  loss: {epoch_loss:.4f}  acc: {epoch_acc:.4f}  time: {time.time() - start:.2f}s')# 保存模型torch.save(model.state_dict(), 'model/cifar2_classification.pth')print("模型保存至: model/cifar2_classification.pth")# 测试函数
def test():print("开始测试模型...")model = ImageClassification()# 加载模型权重if os.path.exists('model/cifar2_classification.pth'):model.load_state_dict(torch.load('model/cifar2_classification.pth'))else:print("错误：未找到模型文件，请先训练模型")returnmodel.eval()  # 设置为评估模式correct = 0total_samples = 0with torch.no_grad():  # 不计算梯度，节省内存for images, labels in test_loader:outputs = model(images)correct += (torch.argmax(outputs, dim=1) == labels).sum().item()total_samples += len(labels)# 计算准确率test_acc = correct / total_samplesprint(f"测试集准确率: {test_acc:.4f}")# # 计算混淆矩阵（可选）# if total_samples > 0:#     from sklearn.metrics import confusion_matrix#     import numpy as np##     all_preds = []#     all_labels = []#     for images, labels in test_loader:#         outputs = model(images)#         all_preds.extend(torch.argmax(outputs, dim=1).cpu().numpy())#         all_labels.extend(labels.cpu().numpy())##     cm = confusion_matrix(all_labels, all_preds)#     print("混淆矩阵:")#     print(cm)#     class_names = list(train_dataset.class_to_idx.keys())#     if len(class_names) == 2:#         print(f"类别 {class_names[0]} 预测为: {class_names[0]} {cm[0, 0]} 次, {class_names[1]} {cm[0, 1]} 次")#         print(f"类别 {class_names[1]} 预测为: {class_names[0]} {cm[1, 0]} 次, {class_names[1]} {cm[1, 1]} 次")if __name__ == '__main__':# 查看数据look_data()# 训练模型train()# 测试模型test()