I have trained and saved my model locally, but how do I reload it and make predictions?

Beginners
1

Hi,

I am working with an image classification model where I have used Trainer, TrainingArguments, and load_dataset("image_folder") to train my model.
The set up for training from scratch and predicting on the trained model looks like this:

model = MobileViTForImageClassification.from_pretrained(
                "apple/mobilevit-xx-small",
                ...)

training_args = TrainingArguments(**parameters)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=collate_fn,
    compute_metrics=self.compute_metrics,
    train_dataset=prepared_train_ds["train"],
    eval_dataset=prepared_train_ds["validation"],
    tokenizer=processor
)

train_results = trainer.train()
trainer.save_model("model_20231102") # Saving the model and I want to use this later
prediction_results = trainer.predict(prepared_test_ds) # works well

So this method works nice and I achieve really good results. However I do not want to re-train the model everytime I am classifying images, so I want to reload the saved model.

# Reload the model
model = MobileViTForImageClassification.from_pretrained("model_20231102")
model(prepared_test_ds)

The test set looks like this:

In [2]: prepared_test_ds
Out[2]: 
Dataset({
    features: ['image', 'image_path', 'label'],
    num_rows: 4997
})

In [3]: prepared_test_ds[0]
Out[3]: 
{'pixel_values': tensor([[[0.7804, 0.7804, 0.7843,  ..., 0.7961, 0.8000, 0.7961],
          [0.7843, 0.7843, 0.7843,  ..., 0.7961, 0.7961, 0.7961],
          [0.7843, 0.7843, 0.7843,  ..., 0.7961, 0.7961, 0.7922],
          ...,
          [0.1765, 0.1922, 0.2588,  ..., 0.4392, 0.4078, 0.5922],
          [0.0549, 0.1059, 0.1490,  ..., 0.4863, 0.4235, 0.5569],
          [0.2078, 0.2745, 0.2980,  ..., 0.4902, 0.4392, 0.5490]],
 
         [[0.5333, 0.5333, 0.5373,  ..., 0.5373, 0.5412, 0.5373],
          [0.5373, 0.5373, 0.5373,  ..., 0.5412, 0.5451, 0.5412],
          [0.5373, 0.5373, 0.5373,  ..., 0.5451, 0.5451, 0.5412],
          ...,
          [0.2196, 0.2353, 0.3059,  ..., 0.4941, 0.4863, 0.6902],
          [0.0784, 0.1333, 0.1804,  ..., 0.5333, 0.4863, 0.6353],
          [0.2000, 0.2706, 0.2980,  ..., 0.5412, 0.4941, 0.6157]],
 
         [[0.3373, 0.3373, 0.3412,  ..., 0.3294, 0.3333, 0.3294],
          [0.3373, 0.3373, 0.3373,  ..., 0.3333, 0.3333, 0.3294],
          [0.3373, 0.3373, 0.3373,  ..., 0.3333, 0.3333, 0.3294],
          ...,
          [0.2235, 0.2471, 0.3176,  ..., 0.5255, 0.5255, 0.7451],
          [0.0706, 0.1294, 0.1843,  ..., 0.5647, 0.5255, 0.6863],
          [0.1569, 0.2392, 0.2745,  ..., 0.5725, 0.5333, 0.6627]]]),
 'labels': 1,
 'image_path': '114806403_image36.jpg'}

But this does not seem to work… I get the following error:

TypeError: conv2d() received an invalid combination of arguments - got (Dataset, Parameter, NoneType, tuple, tuple, tuple, int), but expected one of:
 * (Tensor input, Tensor weight, Tensor bias, tuple of ints stride, tuple of ints padding, tuple of ints dilation, int groups)
      didn't match because some of the arguments have invalid types: (Dataset, Parameter, NoneType, tuple of (int, int), tuple of (int, int), tuple of (int, int), int)
 * (Tensor input, Tensor weight, Tensor bias, tuple of ints stride, str padding, tuple of ints dilation, int groups)
      didn't match because some of the arguments have invalid types: (Dataset, Parameter, NoneType, tuple of (int, int), tuple of (int, int), tuple of (int, int), int)

How to fix this? is there a way to reload the trainer such that I can reuse it for predictions?