from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Dense, Flatten, Dropout, ReLU, AveragePooling2D


# 使用 Keras 重新定义 AlexNet 模型
def create_model(input_shape=(224, 224, 3), num_classes=10):
    model = Sequential([
        # 特征提取部分
        Conv2D(64, kernel_size=(11, 11), strides=(4, 4), padding="same", input_shape=input_shape),
        ReLU(),
        MaxPooling2D(pool_size=(3, 3), strides=(2, 2)),

        Conv2D(192, kernel_size=(5, 5), padding="same"),
        ReLU(),
        MaxPooling2D(pool_size=(3, 3), strides=(2, 2)),

        Conv2D(384, kernel_size=(3, 3), padding="same"),
        ReLU(),

        Conv2D(256, kernel_size=(3, 3), padding="same"),
        ReLU(),

        Conv2D(256, kernel_size=(3, 3), padding="same"),
        ReLU(),
        MaxPooling2D(pool_size=(3, 3), strides=(2, 2)),

        # 全局平均池化部分
        AveragePooling2D(pool_size=(6, 6)),

        # 分类部分
        Flatten(),
        Dropout(0.5),
        Dense(4096),
        ReLU(),
        Dropout(0.5),
        Dense(4096),
        ReLU(),
        Dense(num_classes, activation='softmax')
    ])
    return model