Need some help converting this Pytorch model to TorchScript (.pt) for deployment! 🙏

Beginners
1

Hi folks,

I trained a model here using AutoTrain which left me with the following files:

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

I can make inferernces on the model successfully using the following code:

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) for deployment via torchserve. I’m trying to follow the guide here:

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

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??...>])
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 so I kind of have a little black box on my hands.

Here’s what the 2 json files I received with the model looks 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",
    // more labels...
  },
  "image_size": 224,
  "initializer_range": 0.02,
  "label2id": {
    "hello": "0",
    "world": "1",
    // more labels...
  },
  "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
}

If someone would be so kind enough to tell me what to plugin here:

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

so I can export the model as a .pt file I wold be forever greatfull! :pray: