如何为Fastai vision_learner训练的ResNet50构建TorchServe适配模型类

核心思路 #

Fastai的vision_learner（基于ResNet50）本质是ResNet backbone + Fastai自定义分类头的组合，我们需要用纯PyTorch类复现这个结构，让TorchServe能正确加载权重并执行推理。

自定义模型类实现 #

import torch
import torch.nn as nn
from fastai.vision.models import resnet50
from fastai.layers import AdaptiveConcatPool2d, Flatten

class FastaiResNet50Classifier(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        # 初始化ResNet50 backbone（移除原生分类层）
        self.backbone = resnet50(pretrained=False)
        self.backbone.fc = nn.Identity()
        
        # 复现Fastai默认分类头结构
        self.head = nn.Sequential(
            AdaptiveConcatPool2d(),  # Fastai特有的自适应拼接池化
            Flatten(),
            nn.Linear(2048 * 2, 512),  # ResNet50输出维度2048，拼接池化后翻倍
            nn.ReLU(),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        # 输入需符合PyTorch格式：(batch_size, 3, height, width)
        features = self.backbone(x)
        logits = self.head(features)
        return logits

权重加载与验证 #

1. 从Fastai导出权重
   假设你已经有训练好的learner，执行以下代码导出权重文件：

   # 导出模型权重（注意是learner.model的state_dict）
   torch.save(learner.model.state_dict(), "resnet50_fastai.pth")
   

2. 加载权重到自定义模型

   # 替换num_classes为你的任务类别数
   model = FastaiResNet50Classifier(num_classes=10)
   model.load_state_dict(torch.load("resnet50_fastai.pth"))
   model.eval()
   

关于DynamicUnet的说明 #

DynamicUnet是Fastai为分割任务设计的模型类（位于fastai.vision.models.unet），仅用于unet_learner。你用的是vision_learner（分类任务），因此不需要涉及这个类。

TorchServe适配注意事项 #

• 确保你的处理器（预处理逻辑）与Fastai训练时的预处理完全一致（比如图像归一化的均值、标准差、尺寸等）。
• 打包TorchServe模型时，需将自定义模型类代码包含在model.py中，同时提供权重文件和handler.py（或使用默认图像分类handler）。

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