# Load and re-use a PyTorch model[¶](https://developer.dataiku.com/latest/tutorials/machine-learning/code-env-resources/pytorch-resources/index.html#load-and-re-use-a-pytorch-model "Permalink to this heading")

Pre-requisites

* Dataiku >= 10.0.0.

* A Code Environment with the following packages:

§ torch==1.9.1

§ torchvision==0.10.0

Machine learning use cases can involve a lot of input data and compute-heavy thus expensive model training. You might not want to retrain a model from scratch for common tasks like processing images/ text or during your initial experiments. Instead, you can load *pre-trained models* retrieved from remote repositories and use them for generating predictions.

In this tutorial, you will use Dataiku’s Code Environment resources feature to download and save a pre-trained image classification model from PyTorch Hub. You will then re-use that model to predict the class of a downloaded image.

## Loading the pre-trained model[¶](https://developer.dataiku.com/latest/tutorials/machine-learning/code-env-resources/pytorch-resources/index.html#loading-the-pre-trained-model "Permalink to this heading")

The first step is to download the required assets for your pre-trained model. To do so, in the *Resources* screen of your Code Environment, input the following **initialization script** then click on *Update*:

§ ## Base imports

§ from dataiku.code\_env\_resources import clear\_all\_env\_vars

§ from dataiku.code\_env\_resources import set\_env\_path

§ from dataiku.code\_env\_resources import set\_env\_var

§ from dataiku.code\_env\_resources import grant\_permissions

§ # Import torchvision models

§ import torchvision.models as models

§ # Clears all environment variables defined by previously run script

§ clear\_all\_env\_vars()

§ ## PyTorch

§ # Set PyTorch cache directory

§ set\_env\_path("TORCH\_HOME", "pytorch")

§ # Download pretrained model: automatically managed by PyTorch,

§ # does not download anything if model is already in TORCH\_HOME

§ resnet18 = models.resnet18(pretrained=True)

§ # Grant everyone read access to pretrained models in pytorch/ folder

§ # (by default, PyTorch makes them only readable by the owner)

§ grant\_permissions("pytorch")

This script will retrieve a ResNet18 model from PyTorch Hub and store it on the Dataiku Instance.

Note that it will only need to run once, after that all users allowed to use the Code Environment will be able to leverage the pre-trained model with re-downloading it again.

## Using the pre-trained model for inference[¶](https://developer.dataiku.com/latest/tutorials/machine-learning/code-env-resources/pytorch-resources/index.html#using-the-pre-trained-model-for-inference "Permalink to this heading")

You can now re-use this pre-trained model in your Dataiku Project’s Python Recipe or notebook. Here is an example adapted from a tutorial on the PyTorch website that downloads a sample image and predicts its class from the ImageNet labels.

§ import torch

§ from torchvision import models, transforms

§ from PIL import Image

§ # Import pre-trained model from cache & set to evaluation mode

§ model = models.resnet18(pretrained=True)

§ model.eval()

§ # Download example image from pytorch (it's a doggie, but what kind?)

§ img\_url = "https://github.com/pytorch/hub/raw/master/images/dog.jpg"

§ img\_file = "dog.jpg"

§ torch.hub.download\_url\_to\_file(img\_url, img\_file)

§ # Pre-process image & create a mini-batch as expected by the model

§ input\_image = Image.open(img\_file)

§ preprocess = transforms.Compose([

§ transforms.Resize(256),

§ transforms.CenterCrop(224),

§ transforms.ToTensor(),

§ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),

§ ])

§ input\_tensor = preprocess(input\_image)

§ input\_batch = input\_tensor.unsqueeze(0)

§ # Run softmax to get probabilities since the output has unnormalized scores

§ with torch.no\_grad():

§ output = model(input\_batch)

§ probabilities = torch.nn.functional.softmax(output[0], dim=0)

§ # Download ImageNet class labels

§ classes\_url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet\_classes.txt"

§ classes\_file = "imagenet\_classes.txt"

§ torch.hub.download\_url\_to\_file(classes\_url, classes\_file)

§ # Map prediction to class labels and print top 5 predicted classes

§ with open("imagenet\_classes.txt", "r") as f:

§ categories = [s.strip() for s in f.readlines()]

§ top5\_prob, top5\_catid = torch.topk(probabilities, 5)

§ for i in range(top5\_prob.size(0)):

§ print(categories[top5\_catid[i]], top5\_prob[i].item())

Running this code should give you an output similar to this:

§ Samoyed 0.8846230506896973

§ Arctic fox 0.0458049401640892

§ white wolf 0.044276054948568344

§ Pomeranian 0.00562133826315403

§ Great Pyrenees 0.004651993978768587

## Wrapping up[¶](https://developer.dataiku.com/latest/tutorials/machine-learning/code-env-resources/pytorch-resources/index.html#wrapping-up "Permalink to this heading")

Your pre-trained model is now operational! From there you can easily reuse it, e.g. to directly classify other images stored in a Managed Folder or to fine-tune it for a more specific task.
