如何将预训练AlexNet从1000分类适配为3分类任务？

Hey there! Adapting a pre-trained AlexNet for your 3-class classification task is a classic transfer learning problem, and it's totally manageable once you know the key steps. Let's break this down clearly, starting with PyTorch (the most common framework for this work), then a quick note for TensorFlow/Keras users.

The pre-trained AlexNet ends with a fully connected layer that outputs 1000 classes (matching the ImageNet dataset). Our goal is to replace this final layer with one that outputs 3 classes, then set up training to either keep the pre-trained weights frozen (for small datasets) or fine-tune them (for larger datasets).

1. Load the Pre-trained AlexNet #

First, grab the model with its pre-trained ImageNet weights:

import torch
import torchvision.models as models

# Load AlexNet with pre-trained ImageNet weights
alexnet = models.alexnet(pretrained=True)

This pulls in the full model, including the 1000-class final layer.

2. Replace the Final Fully Connected Layer #

AlexNet's classifier is a sequential module, and the 6th element (index 6) is the final output layer. We need to swap this out for a layer that outputs 3 classes.

First, get the number of input features the final layer expects (this doesn't change, since it's fed by the previous layers):

# Get the input feature count for the final layer
num_input_features = alexnet.classifier[6].in_features

Then replace the layer:

# Replace the 1000-class layer with a 3-class layer
alexnet.classifier[6] = torch.nn.Linear(num_input_features, 3)

3. Initialize the New Layer's Weights #

The pre-trained weights don't apply to our new layer, so we should initialize it properly to help training converge faster. A common approach is Xavier uniform initialization for weights and zero initialization for biases:

# Initialize weights for the new layer
torch.nn.init.xavier_uniform_(alexnet.classifier[6].weight)
# Set biases to 0
torch.nn.init.zeros_(alexnet.classifier[6].bias)

4. Choose Your Training Strategy (Freeze vs. Fine-Tune) #

This depends entirely on the size of your 3-class dataset:

Small Dataset (Few Hundred Samples) #

Freeze all pre-trained layers so we only train the new final layer. This avoids overfitting and leverages the pre-trained feature extraction power:

# Freeze all layers in the feature extractor (the first part of AlexNet)
for param in alexnet.features.parameters():
    param.requires_grad = False

# Freeze all classifier layers except the new final one
for param in alexnet.classifier[:6].parameters():
    param.requires_grad = False

# Ensure only the new final layer gets updated during training
for param in alexnet.classifier[6].parameters():
    param.requires_grad = True

Larger Dataset (Thousands of Samples) #

You can fine-tune some or all of the pre-trained layers after first training the final layer alone. Once the final layer is stable, unfreeze a few top feature layers and train with a small learning rate (to avoid destroying the good pre-trained weights):

# After training the final layer alone, unfreeze some top feature layers
for param in alexnet.features[-4:].parameters():
    param.requires_grad = True

# Use a small learning rate for fine-tuning
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, alexnet.parameters()), lr=1e-4)

5. Finish Training Setup #

Now you're ready to train like any other classification model:

• Use CrossEntropyLoss as your loss function (perfect for multi-class tasks)
• Pick an optimizer (SGD or Adam work well; use a higher lr for the final layer alone, lower for fine-tuning)
• Train on your dataset, validate on a holdout set, and adjust as needed

Quick TensorFlow/Keras Version #

If you prefer Keras, the process is similar:

import tensorflow as tf
from tensorflow.keras.applications import AlexNet
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model

# Load pre-trained AlexNet without the top classification layer
alexnet = AlexNet(weights='imagenet', include_top=False, input_shape=(224, 224, 3))

# Add a global pooling layer and our new 3-class output layer
x = alexnet.output
x = GlobalAveragePooling2D()(x)
predictions = Dense(3, activation='softmax')(x)

# Build the full model
model = Model(inputs=alexnet.input, outputs=predictions)

# Freeze pre-trained layers for small datasets
for layer in alexnet.layers:
    layer.trainable = False

# Compile and train
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
# model.fit(...)

• You don't need to modify the existing pre-trained weights—just swap out the final layer and initialize its weights.
• Freeze pre-trained layers if your dataset is small to avoid overfitting.
• Fine-tune with a small learning rate if you have enough data to improve feature extraction for your specific task.

内容的提问来源于stack exchange，提问作者Pooja