CNN卷积神经网络学习笔记

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.io import read_image, ImageReadMode
import os
import argparse
import json

# 标签编码器：将类别名称转换为数字编码
class SimpleLabelEncoder:
    def __init__(self):
        self.label_to_idx = {}
        self.idx_to_label = {}
    
    def fit(self, labels):
        unique_labels = sorted(list(set(labels)))
        self.label_to_idx = {label: idx for idx, label in enumerate(unique_labels)}
        self.idx_to_label = {idx: label for label, idx in self.label_to_idx.items()}
    
    def transform(self, label):
        return self.label_to_idx[label]
    
    def inverse_transform(self, idx):
        return self.idx_to_label[idx]

# 自定义数据集类：加载花卉图像与标签
class CustomImageDataset(Dataset):
    def __init__(self, img_dir, transform=None, target_transform=None, encoder=None):
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform
        self.encoder = encoder
        self.img_paths = []
        self.labels = []
        
        # 遍历数据集目录，获取图像路径和标签
        for label_dir in os.listdir(img_dir):
            label_path = os.path.join(img_dir, label_dir)
            if os.path.isdir(label_path):
                for img_name in os.listdir(label_path):
                    img_path = os.path.join(label_path, img_name)
                    self.img_paths.append(img_path)
                    self.labels.append(label_dir)
    
    def __len__(self):
        return len(self.img_paths)
    
    def __getitem__(self, idx):
        # 读取图像并转换为张量
        img_tensor = read_image(self.img_paths[idx], mode=ImageReadMode.RGB)
        label = self.labels[idx]
        # 标签编码
        if self.encoder:
            label = self.encoder.transform(label)
        # 图像变换
        if self.transform:
            img_tensor = self.transform(img_tensor)
        if self.target_transform:
            label = self.target_transform(label)
        return img_tensor, label

# 训练一个epoch的函数
def train_epoch(model, dataloader, criterion, optimizer, device):
    model.train()
    total_loss = 0.0
    total_correct = 0
    total_samples = 0
    
    for inputs, labels in dataloader:
        inputs, labels = inputs.to(device), labels.to(device)
        
        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
        # 反向传播与参数更新
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # 计算损失和准确率
        total_loss += loss.item() * inputs.size(0)
        _, predicted = torch.max(outputs, 1)
        total_correct += (predicted == labels).sum().item()
        total_samples += inputs.size(0)
    
    avg_loss = total_loss / total_samples
    accuracy = total_correct / total_samples
    return avg_loss, accuracy

# 评估模型的函数
def evaluate(model, dataloader, criterion, device):
    model.eval()
    total_loss = 0.0
    total_correct = 0
    total_samples = 0
    
    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            total_loss += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs, 1)
            total_correct += (predicted == labels).sum().item()
            total_samples += inputs.size(0)
    
    avg_loss = total_loss / total_samples
    accuracy = total_correct / total_samples
    return avg_loss, accuracy

# 保存模型为Hugging Face格式（便于部署）
def save_huggingface_format(model, save_dir, encoder):
    os.makedirs(save_dir, exist_ok=True)
    # 保存模型参数
    torch.save(model.state_dict(), os.path.join(save_dir, "pytorch_model.bin"))
    # 保存标签编码器
    with open(os.path.join(save_dir, "label_encoder.json"), "w") as f:
        json.dump({"label_to_idx": encoder.label_to_idx, "idx_to_label": encoder.idx_to_label}, f)

# 图像预测函数
def predict_image(image_path, model, encoder, device, img_size=150):
    # 图像预处理：与训练集保持一致
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ConvertImageDtype(torch.float),
        transforms.Normalize(mean=[0.0, 0.0, 0.0], std=[1.0, 1.0, 1.0])  # 与训练时的归一化一致
    ])
    
    # 读取并预处理图像
    img_tensor = read_image(image_path, mode=ImageReadMode.RGB)
    img_tensor = transform(img_tensor).unsqueeze(0).to(device)  # 增加批次维度
    
    # 模型预测
    model.eval()
    with torch.no_grad():
        outputs = model(img_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        _, predicted_idx = torch.max(outputs, 1)
        predicted_label = encoder.inverse_transform(predicted_idx.item())
        confidence = probabilities[0][predicted_idx.item()].item()
    
    return predicted_label, confidence

# 优化后的CNN模型类
class OptimizedCNN(nn.Module):
    def __init__(self, num_classes=4, img_size=150):
        super(OptimizedCNN, self).__init__()
        self.img_size = img_size
        
        # 第1个卷积块：conv → bn → relu → pool
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)
        
        # 第2个卷积块：conv → bn → relu → pool
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)
        
        # 第3个卷积块：conv → bn → relu → pool
        self.conv3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(64)
        self.relu3 = nn.ReLU()
        self.pool3 = nn.MaxPool2d(2)
        
        # 第4个卷积块：conv → bn → relu → pool
        self.conv4 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
        self.bn4 = nn.BatchNorm2d(128)
        self.relu4 = nn.ReLU()
        self.pool4 = nn.MaxPool2d(2)
        
        # 计算全连接层输入维度（经过4次池化，尺寸缩小16倍）
        feature_h = img_size // 16
        feature_w = img_size // 16
        self.fc_input_dim = 128 * feature_h * feature_w
        
        # 全连接层与Dropout
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(self.fc_input_dim, 512)
        self.relu_fc1 = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(512, num_classes)
    
    def forward(self, x):
        # 4个卷积块前向传播
        x = self.pool1(self.relu1(self.bn1(self.conv1(x))))
        x = self.pool2(self.relu2(self.bn2(self.conv2(x))))
        x = self.pool3(self.relu3(self.bn3(self.conv3(x))))
        x = self.pool4(self.relu4(self.bn4(self.conv4(x))))
        
        # 拉平特征图，输入全连接层
        x = self.flatten(x)
        x = self.dropout(self.relu_fc1(self.fc1(x)))
        x = self.fc2(x)
        return x

# 主函数：整合训练、评估与预测流程
def main(args):
    # 设置设备（GPU优先）
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {DEVICE}")
    
    # 图像尺寸与批次大小
    IMG_SIZE = 150
    BATCH_SIZE = 32
    
    # 1. 数据加载与预处理
    # 标签编码器初始化与拟合
    encoder = SimpleLabelEncoder()
    # 遍历数据集获取所有标签并拟合编码器
    all_labels = []
    for label_dir in os.listdir(args.data_dir):
        all_labels.append(label_dir)
    encoder.fit(all_labels)
    num_classes = len(encoder.label_to_idx)
    
    # 数据变换（训练集增强，测试集仅归一化）
    train_transforms = transforms.Compose([
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.RandomRotation(15),      # 随机旋转
        transforms.ConvertImageDtype(torch.float),
        transforms.Normalize(mean=[0.0, 0.0, 0.0], std=[1.0, 1.0, 1.0])  # 归一化到0-1
    ])
    
    test_transforms = transforms.Compose([
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.ConvertImageDtype(torch.float),
        transforms.Normalize(mean=[0.0, 0.0, 0.0], std=[1.0, 1.0, 1.0])
    ])
    
    # 加载数据集并划分训练集、测试集（8:2）
    full_dataset = CustomImageDataset(
        img_dir=args.data_dir,
        encoder=encoder
    )
    train_size = int(0.8 * len(full_dataset))
    test_size = len(full_dataset) - train_size
    train_dataset, test_dataset = torch.utils.data.random_split(full_dataset, [train_size, test_size])
    
    # 重新设置数据集的transform（因为split后transform会丢失）
    train_dataset.dataset.transform = train_transforms
    test_dataset.dataset.transform = test_transforms
    
    # 创建数据加载器
    train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
    
    # 2. 模型初始化
    model = OptimizedCNN(num_classes=num_classes, img_size=IMG_SIZE).to(DEVICE)
    criterion = nn.CrossEntropyLoss()  # 交叉熵损失（多分类）
    optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器
    
    # 3. 训练模型（若为训练模式）
    if args.mode == "train":
        epochs = 100
        best_test_acc = 0.0
        print("Starting training...")
        for epoch in range(1, epochs + 1):
            train_loss, train_acc = train_epoch(model, train_dataloader, criterion, optimizer, DEVICE)
            test_loss, test_acc = evaluate(model, test_dataloader, criterion, DEVICE)
            
            # 打印训练日志
            print(f"Epoch [{epoch}/{epochs}], Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, "
                  f"Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}")
            
            # 保存最优模型
            if test_acc > best_test_acc:
                best_test_acc = test_acc
                torch.save(model.state_dict(), args.model_path)
                print(f"Best model saved (Test Acc: {best_test_acc:.4f})")
        
        # 保存Hugging Face格式模型
        save_huggingface_format(model, args.save_dir, encoder)
        print(f"Training completed. Best Test Accuracy: {best_test_acc:.4f}")
    
    # 4. 预测图像（若为推理模式）
    elif args.mode == "infer":
        # 加载训练好的模型
        model.load_state_dict(torch.load(args.model_path, map_location=DEVICE))
        # 加载标签编码器
        with open(os.path.join(args.save_dir, "label_encoder.json"), "r") as f:
            encoder_dict = json.load(f)
            encoder.label_to_idx = encoder_dict["label_to_idx"]
            encoder.idx_to_label = encoder_dict["idx_to_label"]
        
        # 预测图像
        predicted_label, confidence = predict_image(args.image_path, model, encoder, DEVICE, IMG_SIZE)
        print(f"Prediction Result: {predicted_label}, Confidence: {confidence:.4f}")

if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="CNN Flower Classification")
    parser.add_argument("--mode", type=str, choices=["train", "infer"], required=True, help="Mode: train or infer")
    parser.add_argument("--data_dir", type=str, default="../flowers", help="Directory of flower dataset")
    parser.add_argument("--model_path", type=str, default="cnn_flower_model.pth", help="Path to save/load model")
    parser.add_argument("--save_dir", type=str, default="cnn_flower_model_hf", help="Directory to save Hugging Face format model")
    parser.add_argument("--image_path", type=str, help="Path of image to infer (required in infer mode)")
    
    args = parser.parse_args()
    main(args)

self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)

self.fc_input_dim = 128 * feature_h * feature_w  # 128×9×9=10368
self.fc1 = nn.Linear(self.fc_input_dim, 512)  # 从10368维压缩到512维
self.relu_fc1 = nn.ReLU()
self.dropout = nn.Dropout(0.5)
self.fc2 = nn.Linear(512, num_classes)  # 从512维映射到4个类别

import kagglehub

# Download latest version
path = kagglehub.dataset_download("alxmamaev/flowers-recognition")

print("Path to dataset files:", path)

def forward(self, x):
    # 4个卷积块：conv → bn → relu → pool
    x = self.pool1(self.relu1(self.bn1(self.conv1(x))))
    x = self.pool2(self.relu2(self.bn2(self.conv2(x))))
    x = self.pool3(self.relu3(self.bn3(self.conv3(x))))
    x = self.pool4(self.relu4(self.bn4(self.conv4(x))))
    
    # 拉平特征图，输入全连接层
    x = self.flatten(x)
    x = self.dropout(self.relu_fc1(self.fc1(x)))
    x = self.fc2(x)
    return x

model = OptimizedCNN(num_classes=num_classes).to(DEVICE)  # 加载模型到指定设备（GPU/CPU）
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数，适用于多分类任务
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器，用于参数更新