How to convert a Hugging Face Pytorch model (AutoTrain) to TorchScript (.pt) for deployment?

Question

I trained an image classification model using Hugging Face's AutoTrain service which left me with the following three files:

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

Here's what the 2 json files look like:

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
}

Without changing anything I can make inferences on the model successfully using the following code:

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))

But now I need to export the model to torchscript (.pt file format) for deployment via torchserve. I’m trying to follow the guide here:

https://huggingface.co/docs/transformers/torchscript

But am stuck at the “Creating the trace” part below:

import torch
from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained(
    "path/to/model",
    torchscript=True,
)

# The model needs to be in evaluation mode
model.eval()

# Creating the trace
traced_model = torch.jit.trace(model, [<some_tensors_here??...>]) # <---- HERE!
torch.jit.save(traced_model, "my_converted_model.pt")

From what I understand, I need to pass some dummy data to the model in a shape that it expects. How can I find out what this data structure should look like?

I’m pretty clueless when it comes to ML (just trying to get a little model working for my app). I also didn’t train the model myself but used Autotrain (which is a automl GUI tool) so I'm kind of left in the dark about how the model was made and what input data it uses/expects.

If someone would be so kind enough to tell me what to plug in to get this to work I'd really appreciate it! The goal is to convert the format I have now to torchscript so that I can either deploy it on a server using torchserve or run it on device by converting it to Playtorch to use in my React Native app. Both of these things expect a .pt file for the model as a starting point.