A simple workflow and code for training a custom model and running inference.

[binominerigike]

Hi, so I've gotten timm to run well in the past, but had to primarily rely on the train.py and inference.py scripts in command line.

As they're a bit complex and have a lot of options I likely won't change, I'm curious if anyone has any example code of loading them in like libraries and training and more importantly, running inference.

For the inference portion, I tried the following, which is how I'd do this in other pytorch programs:

import torch
import timm
import numpy as np
from PIL import Image


def classify_image(model_checkpoint_path, model_checkpoint_architecture, class_names_file, image_file):
    # Load the class names
    with open(class_names_file, 'r') as f:
        class_names = [line.strip() for line in f.readlines()]

    # Load the model architecture and weights
    model = timm.create_model(model_checkpoint_architecture, num_classes=len(class_names), checkpoint_path=model_checkpoint_path, pretrained=False)
    state_dict = torch.load(model_checkpoint_path)
    model.load_state_dict(state_dict)

    # Load the image and preprocess it
    img = Image.open(image_file)
    img = img.resize((224, 224))
    img = np.array(img).astype(np.float32) / 255.0
    img = np.transpose(img, (2, 0, 1))
    img = np.expand_dims(img, axis=0)
    img = torch.from_numpy(img)

    # Make a prediction and print the result
    with torch.no_grad():
        model.eval()
        logits = model(img)
        probs = torch.softmax(logits, dim=1)
        topk = probs.topk(1)

    class_index = topk.indices.squeeze().item()
    class_name = class_names[class_index]
    confidence = topk.values.squeeze().item()

    print(f'Prediction: {class_name} ({confidence:.2f})')


model = "./output/train/20230224-154752-seresnet34-224/model_best.pth.tar"
arch = "seresnet34"
classes_file = "classes.txt"
image_file = "an_image.jpg"

classify_image(model, arch, classes_file, image_file)

However, I am getting the following:

Traceback (most recent call last):
  File "/nas/nns/img_class/timm/pytorch-image-models/timm_predict_cgpt2.py", line 44, in <module>
    classify_image(model, arch, classes_file, image_file)
  File "/nas/nns/img_class/timm/pytorch-image-models/timm_predict_cgpt2.py", line 13, in classify_image
    model = timm.create_model(model_checkpoint_architecture, num_classes=2, checkpoint_path=model_checkpoint_path, pretrained=False)
  File "/nas/nns/img_class/timm/pytorch-image-models/timm/models/_factory.py", line 101, in create_model
    load_checkpoint(model, checkpoint_path)
  File "/nas/nns/img_class/timm/pytorch-image-models/timm/models/_helpers.py", line 70, in load_checkpoint
    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
  File "/home/binominerigike/.pyenv/versions/3.10.4/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1671, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for ResNet:
	size mismatch for fc.weight: copying a param with shape torch.Size([1000, 512]) from checkpoint, the shape in current model is torch.Size([2, 512]).
	size mismatch for fc.bias: copying a param with shape torch.Size([1000]) from checkpoint, the shape in current model is torch.Size([2]).

Which is confusing since I made sure I used --num-classes 2 and --class-map classes.txt while training. I've read through the docs on huggingface and fastai but they don't seem to have an example similar to what I'm attempting which is very strange to me.

If anyone has any suggestions, I'm most appreciative. Again it's also helpful if this were very clean and simple code like yolov8's quickstart example. I'm preparing this for a lesson for teaching students so it's helpful to keep it very clean.