基于PyTorch将预训练ResNet50替换MNIST自定义网络适配(200,200,3)输入

1. 核心模型修改 #

ResNet50预训练模型默认适配ImageNet 1000分类任务，而MNIST是10分类任务，因此只需替换模型的最后全连接层即可。你的输入尺寸(200,200,3)完全符合ResNet的输入要求（ResNet支持任意大于32×32的RGB图像，内部自适应池化会处理尺寸差异），无需修改前面的卷积层结构。

代码示例：

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

# 加载预训练ResNet50
model = models.resnet50(pretrained=True)

# 获取原全连接层的输入特征维度
num_input_features = model.fc.in_features

# 替换为10分类的全连接层
model.fc = nn.Linear(num_input_features, 10)

2. 数据输入适配（针对MNIST数据集） #

MNIST原始数据是单通道灰度图，而ResNet要求3通道RGB输入，因此需要在数据预处理阶段添加通道转换：

from torchvision import transforms, datasets

# 定义数据变换流程
transform = transforms.Compose([
    transforms.Resize((200, 200)),  # 将图像调整为目标输入尺寸
    transforms.Grayscale(num_output_channels=3),  # 单通道转3通道
    transforms.ToTensor(),
    # 使用ResNet预训练对应的归一化参数
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载处理后的MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

如果你的输入本身就是(200,200,3)的RGB图像，只需保留Resize（若原始尺寸不符）和归一化步骤，移除Grayscale转换即可。

3. 训练策略调整 #

• 冻结预训练层（仅训练新层）：若想快速收敛，可先冻结ResNet的预训练参数，只训练替换后的全连接层：

  # 冻结所有预训练层参数
  for param in model.parameters():
      param.requires_grad = False
  
  # 解冻最后全连接层的参数，允许训练
  for param in model.fc.parameters():
      param.requires_grad = True
  

• 微调整个网络：若追求更好性能，可解冻所有参数，使用较小的学习率（如1e-4）进行全网络微调，注意此时会增加训练时长和显存占用。

训练循环部分可直接复用原自定义网络的代码，只需将模型实例替换为上述修改后的ResNet50即可。

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