Pytorch - Streamlit - custom model

Hello everyone … Online I found this guide that explains how to create a simple webapp for image recognition. In this case it uses the ImageNet model already pretrained …
Here the github site:

I would like to use a model that I have already made the train of which I have the .pth file
Where can I edit the file? For convenience I paste it below:

from torchvision import models, transforms
import torch
from PIL import Image
def predict(image_path):
resnet = models.resnet101(pretrained=True)

transform = transforms.Compose([
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]

img =
batch_t = torch.unsqueeze(transform(img), 0)

out = resnet(batch_t)

with open('imagenet_classes.txt') as f:
    classes = [line.strip() for line in f.readlines()]

prob = torch.nn.functional.softmax(out, dim=1)[0] * 100
_, indices = torch.sort(out, descending=True)
return [(classes[idx], prob[idx].item()) for idx in indices[0][:5]]

The predict method seems to initialize the complete model, transformation, loads the data, processes the forward pass, and returns the class probabilities for a single input image.

I’m not a webapp expert, but this seems very inefficient. I would recommend to take a look at chapter 15 in Deep Learning with PyTorch (which is still free) by @elistevens, @lantiga, and @tom.

Thanks a lot! you have been very kind. I’m going to read now …

To echo @ptrblck’s comment, there is nothing wrong with streamlit if you want the UI, but hopefully the examples will give you some idea about the parts involved, some thoughts on efficiency, when to do what etc.

Best regards


Thank you for your answer
I would like to create a webapp-beta equal to that of the example I have linked above, with the difference I would like to use my .pth model… Nothing different for now…
bye bye

Yes, probably the first step would be to just replace the resnet with your own but also to move everything above and the reading of the classes out of the predict function.
If you’re looking at multithreading or so, using the JIT is important (it was mainly my theory when I wrote that chapter, but we’ve seen it in the wild in these forums since).