如何将Hugging Face AutoTrain的PyTorch模型转为TorchScript(.pt)用于部署？

背景 #

你通过Hugging Face AutoTrain训练了图像分类模型，得到三个文件：

• config.json
• preprocessor_config.json
• pytorch_model.bin

配置文件内容 #

preprocessor_config.json:

{
  "do_normalize": true,
  "do_rescale": true,
  "do_resize": true,
  "feature_extractor_type": "ViTFeatureExtractor",
  "image_mean": [0.485, 0.456, 0.406],
  "image_processor_type": "ViTImageProcessor",
  "image_std": [0.229, 0.224, 0.225],
  "resample": 3,
  "rescale_factor": 0.00392156862745098,
  "size": {
    "height": 224,
    "width": 224
  }
}

config.json:

{
  "_name_or_path": "AutoTrain",
  "architectures": ["SwinForImageClassification"],
  "attention_probs_dropout_prob": 0.0,
  "depths": [2, 2, 18, 2],
  "drop_path_rate": 0.1,
  "embed_dim": 128,
  "encoder_stride": 32,
  "hidden_act": "gelu",
  "hidden_dropout_prob": 0.0,
  "hidden_size": 1024,
  "id2label": {
    "hello": "0",
    "world": "1",
    // approx 60 more labels here...
  },
  "image_size": 224,
  "initializer_range": 0.02,
  "label2id": {
    "hello": "0",
    "world": "1",
    // approx 60 more labels here...
  },
  "layer_norm_eps": 1e-5,
  "max_length": 128,
  "mlp_ratio": 4.0,
  "model_type": "swin",
  "num_channels": 3,
  "num_heads": [4, 8, 16, 32],
  "num_layers": 4,
  "padding": "max_length",
  "patch_size": 4,
  "path_norm": true,
  "problem_type": "single_label_classification",
  "qkv_bias": true,
  "torch_dtype": "float32",
  "transformers_version": "4.25.1",
  "use_absolute_embeddings": false,
  "window_size": 7
}

你已能通过以下代码成功推理：

from PIL import Image
from transformers import pipeline

classifier = pipeline(
    "image-classification",
    "path/to/model",
)

image = Image.open("./test.jpg").convert("RGB")

print(classifier(image))

但在导出TorchScript时卡在创建trace的步骤，不清楚需要传入什么格式的虚拟数据。

解决方案 #

要完成TorchScript转换，需生成符合模型输入要求的虚拟张量。根据配置文件，模型期望输入为形状(1, 3, 224, 224)的张量（1代表批量大小，3是RGB通道，224×224是图像尺寸）。以下是两种可行方法：

方法1：直接创建虚拟张量 #

import torch
from transformers import AutoModelForImageClassification

# 加载模型
model = AutoModelForImageClassification.from_pretrained(
    "path/to/model",
    torchscript=True,
)
model.eval()

# 创建符合要求的虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 生成trace并保存
traced_model = torch.jit.trace(model, dummy_input)
torch.jit.save(traced_model, "my_converted_model.pt")

方法2：用图像处理器生成标准输入（更严谨） #

该方法模拟真实推理的预处理流程，确保输入格式完全匹配：

import torch
from PIL import Image
from transformers import AutoModelForImageClassification, AutoImageProcessor

# 加载模型和图像处理器
model = AutoModelForImageClassification.from_pretrained(
    "path/to/model",
    torchscript=True,
)
processor = AutoImageProcessor.from_pretrained("path/to/model")
model.eval()

# 创建虚拟图像（或使用真实图像）
dummy_image = Image.new("RGB", (224, 224))
# 预处理得到模型所需输入张量
inputs = processor(dummy_image, return_tensors="pt")

# 生成trace并保存
traced_model = torch.jit.trace(model, inputs["pixel_values"])
torch.jit.save(traced_model, "my_converted_model.pt")

验证转换结果 #

可加载保存的TorchScript模型测试是否正常工作：

traced_model = torch.jit.load("my_converted_model.pt")
traced_model.eval()

# 预处理测试图像
test_image = Image.open("./test.jpg").convert("RGB")
inputs = processor(test_image, return_tensors="pt")

# 推理并解析结果
outputs = traced_model(inputs["pixel_values"])
predicted_class_idx = outputs.logits.argmax(-1).item()
print(model.config.id2label[str(predicted_class_idx)])

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