如何使用ResNet或VGGNet进行图像分类并手动设置输出类别

核心原理 #

预训练的ResNet/VGG默认最后一层是适配ImageNet 1000分类的，你需要替换这一层为2个输出单元的全连接层（对应道路、草地二分类），同时可选择冻结前面的特征提取层或后续微调。

PyTorch实现步骤 #

1. 加载预训练模型并替换输出层 #

以ResNet50和VGG16为例：

import torch
import torch.nn as nn
from torchvision import models

# ResNet50修改
model = models.resnet50(pretrained=True)
# 冻结所有特征提取层（数据集较小时推荐先冻结，后续再微调）
for param in model.parameters():
    param.requires_grad = False
# 获取原最后一层输入特征数，替换为二分类层
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)

# VGG16修改
vgg_model = models.vgg16(pretrained=True)
for param in vgg_model.parameters():
    param.requires_grad = False
# VGG的分类层在classifier第6位
num_ftrs_vgg = vgg_model.classifier[6].in_features
vgg_model.classifier[6] = nn.Linear(num_ftrs_vgg, 2)

2. 损失函数与优化器设置 #

二分类场景用CrossEntropyLoss即可，若冻结了特征层，优化器仅更新新替换的层：

criterion = nn.CrossEntropyLoss()
# 只优化最后一层参数
optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.001)

3. 训练注意事项 #

确保数据集标签为0（道路）、1（草地）的整数类型，按常规流程完成前向传播、损失计算与反向传播即可。

TensorFlow/Keras实现步骤 #

1. 加载预训练模型并替换输出层 #

以ResNet50和VGG16为例：

import tensorflow as tf
from tensorflow.keras.applications import ResNet50, VGG16
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Flatten
from tensorflow.keras.models import Model

# ResNet50修改
base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(224,224,3))
base_model.trainable = False  # 冻结特征层
# 添加自定义顶层
x = base_model.output
x = GlobalAveragePooling2D()(x)
predictions = Dense(2, activation='softmax')(x)
model = Model(inputs=base_model.input, outputs=predictions)

# VGG16修改
vgg_base = VGG16(weights='imagenet', include_top=False, input_shape=(224,224,3))
vgg_base.trainable = False
x_vgg = vgg_base.output
x_vgg = Flatten()(x_vgg)
predictions_vgg = Dense(2, activation='softmax')(x_vgg)
vgg_model = Model(inputs=vgg_base.input, outputs=predictions_vgg)

2. 编译模型 #

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',  # 标签为整数时用该损失，one-hot标签用categorical_crossentropy
              metrics=['accuracy'])

额外关键提示 #

• 数据预处理：必须匹配预训练模型要求，比如PyTorch中需将图像归一化到mean=[0.485, 0.456, 0.406]、std=[0.229, 0.224, 0.225]；Keras中直接调用preprocess_input函数处理。
• 微调策略：若冻结特征层训练后精度不足，可解冻部分顶层特征层，用更小的学习率（如1e-5）继续训练，让模型适配你的自定义数据集。
• 数据集划分：确保训练集、验证集、测试集划分合理，避免过拟合。

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