Python脚本数据图像分类CNN如何搭建

wen 实用脚本 2026-07-18 1

本文目录导读：

基础CNN模型搭建
更复杂的CNN架构
数据增强和预处理
迁移学习（使用预训练模型）
完整的训练脚本
使用建议

我来详细介绍如何使用Python搭建CNN进行图像分类。

基础CNN模型搭建

使用TensorFlow/Keras

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, models
import numpy as np
import matplotlib.pyplot as plt
# 1. 构建CNN模型
def create_cnn_model(input_shape=(32, 32, 3), num_classes=10):
    model = models.Sequential([
        # 卷积层1
        layers.Conv2D(32, (3, 3), activation='relu', input_shape=input_shape),
        layers.MaxPooling2D((2, 2)),
        # 卷积层2
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.MaxPooling2D((2, 2)),
        # 卷积层3
        layers.Conv2D(64, (3, 3), activation='relu'),
        # 全连接层
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation='softmax')
    ])
    return model
# 2. 编译模型
def compile_model(model):
    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )
    return model
# 3. 训练模型
def train_model(model, train_images, train_labels, test_images, test_labels):
    history = model.fit(
        train_images, 
        train_labels,
        epochs=10,
        validation_data=(test_images, test_labels),
        batch_size=32,
        verbose=1
    )
    return history
# 4. 评估模型
def evaluate_model(model, test_images, test_labels):
    test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)
    print(f'\nTest accuracy: {test_acc:.4f}')
    return test_loss, test_acc
# 示例：使用CIFAR-10数据集
def main():
    # 加载CIFAR-10数据集
    (train_images, train_labels), (test_images, test_labels) = keras.datasets.cifar10.load_data()
    # 数据预处理
    train_images = train_images.astype('float32') / 255.0
    test_images = test_images.astype('float32') / 255.0
    # 创建模型
    model = create_cnn_model()
    model = compile_model(model)
    model.summary()
    # 训练模型
    history = train_model(model, train_images, train_labels, test_images, test_labels)
    # 评估模型
    evaluate_model(model, test_images, test_labels)
    # 绘制训练曲线
    plot_training_history(history)
def plot_training_history(history):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    # 准确率曲线
    ax1.plot(history.history['accuracy'], label='Training Accuracy')
    ax1.plot(history.history['val_accuracy'], label='Validation Accuracy')
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Accuracy')
    ax1.legend()
    ax1.set_title('Model Accuracy')
    # 损失曲线
    ax2.plot(history.history['loss'], label='Training Loss')
    ax2.plot(history.history['val_loss'], label='Validation Loss')
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Loss')
    ax2.legend()
    ax2.set_title('Model Loss')
    plt.tight_layout()
    plt.show()
if __name__ == "__main__":
    main()

更复杂的CNN架构

使用PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
# 定义CNN模型
class CNNClassifier(nn.Module):
    def __init__(self, num_classes=10):
        super(CNNClassifier, self).__init__()
        # 卷积层
        self.conv_layers = nn.Sequential(
            # 第一层
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # 第二层
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # 第三层
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(128),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        # 全连接层
        self.fc_layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes)
        )
    def forward(self, x):
        x = self.conv_layers(x)
        x = self.fc_layers(x)
        return x
# 训练函数
def train_pytorch_model():
    # 数据预处理
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(10),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    # 加载数据集
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)
    trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
    testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                           download=True, transform=transform)
    testloader = DataLoader(testset, batch_size=64, shuffle=False)
    # 初始化模型
    model = CNNClassifier(num_classes=10)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    # 训练循环
    num_epochs = 10
    for epoch in range(num_epochs):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            if i % 200 == 199:
                print(f'[Epoch {epoch + 1}, Batch {i + 1}] loss: {running_loss / 200:.3f}')
                running_loss = 0.0
    print('Finished Training')
    return model
# 评估模型
def evaluate_pytorch_model(model, testloader):
    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f'Accuracy: {100 * correct / total:.2f}%')

数据增强和预处理

import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# 数据增强配置
def create_data_augmentation():
    datagen = ImageDataGenerator(
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        horizontal_flip=True,
        zoom_range=0.2,
        shear_range=0.2,
        fill_mode='nearest'
    )
    return datagen
# 使用数据增强训练
def train_with_augmentation(model, train_images, train_labels, test_images, test_labels):
    datagen = create_data_augmentation()
    datagen.fit(train_images)
    history = model.fit(
        datagen.flow(train_images, train_labels, batch_size=32),
        epochs=20,
        validation_data=(test_images, test_labels),
        steps_per_epoch=len(train_images) // 32,
        verbose=1
    )
    return history

迁移学习（使用预训练模型）

from tensorflow.keras.applications import VGG16, ResNet50
from tensorflow.keras import layers, models
# 使用预训练模型进行迁移学习
def create_transfer_learning_model(input_shape=(224, 224, 3), num_classes=10):
    # 加载预训练模型（去掉顶层）
    base_model = VGG16(
        weights='imagenet',
        include_top=False,
        input_shape=input_shape
    )
    # 冻结预训练层
    base_model.trainable = False
    # 添加自定义分类层
    model = models.Sequential([
        base_model,
        layers.GlobalAveragePooling2D(),
        layers.Dense(256, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(128, activation='relu'),
        layers.Dropout(0.3),
        layers.Dense(num_classes, activation='softmax')
    ])
    return model

完整的训练脚本

import os
import argparse
import numpy as np
import tensorflow as tf
from tensorflow import keras
def create_custom_cnn(input_shape, num_classes):
    """创建自定义CNN模型"""
    model = keras.Sequential([
        keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=input_shape),
        keras.layers.MaxPooling2D((2, 2)),
        keras.layers.Conv2D(64, (3, 3), activation='relu'),
        keras.layers.MaxPooling2D((2, 2)),
        keras.layers.Conv2D(128, (3, 3), activation='relu'),
        keras.layers.Flatten(),
        keras.layers.Dense(128, activation='relu'),
        keras.layers.Dropout(0.5),
        keras.layers.Dense(num_classes, activation='softmax')
    ])
    return model
def load_and_preprocess_data():
    """加载并预处理数据"""
    (train_images, train_labels), (test_images, test_labels) = keras.datasets.cifar10.load_data()
    # 归一化
    train_images = train_images.astype('float32') / 255.0
    test_images = test_images.astype('float32') / 255.0
    return train_images, train_labels, test_images, test_labels
def main():
    parser = argparse.ArgumentParser(description='CNN Image Classifier')
    parser.add_argument('--epochs', type=int, default=10, help='Number of epochs')
    parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
    parser.add_argument('--save_model', action='store_true', help='Save trained model')
    args = parser.parse_args()
    # 加载数据
    print("Loading data...")
    train_images, train_labels, test_images, test_labels = load_and_preprocess_data()
    # 创建模型
    print("Creating model...")
    model = create_custom_cnn((32, 32, 3), 10)
    # 编译模型
    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )
    # 显示模型结构
    model.summary()
    # 训练模型
    print("Training model...")
    history = model.fit(
        train_images,
        train_labels,
        batch_size=args.batch_size,
        epochs=args.epochs,
        validation_data=(test_images, test_labels)
    )
    # 评估模型
    test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=0)
    print(f'\nTest accuracy: {test_acc:.4f}')
    # 保存模型
    if args.save_model:
        model.save('cnn_image_classifier.h5')
        print("Model saved as 'cnn_image_classifier.h5'")
if __name__ == "__main__":
    main()

使用建议

开始简单：从基础CNN开始，逐步增加复杂度
数据预处理：归一化、数据增强都很重要
监控过拟合：使用验证集监控，添加Dropout和正则化
调整超参数：学习率、批量大小、层数等
使用GPU：CNN训练计算量大，建议使用GPU

这个框架可以根据你的具体需求进行调整和扩展。