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:

2

You need to pass here input tensor(s) for your model in torch format

torch.jit.trace(model,(torch.tensor,))

If your model has multiple inputs they must passed in order that is defined by your model, for example:

torch.jit.trace(model,(torch.tensor1, torch.tensor2, torch.tensorN))

Please refer for mans for more details TorchScript — PyTorch 2.0 documentation