How to Share and Load PyTorch Models Through Torch Hub?

Making Models Easily Accessible

6 min readApr 1, 2023

To bridge research and production, we want a way to easily publish and access models. PyTorch provides these functionalities through Torch Hub.

It also follows the “Don’t Repeat Yourself” (DRY) principle of programming. I sometimes find myself spending much time implementing someone else’s model given their codebase and trying to replicate their results. So when I learned about Torch Hub, I was somehow overjoyed.

With Torch Hub, authors can also specifically publish an easy way of loading their models (which in turn increases the user base of their models and *ehem*, citations), and furthermore, leaves the user only one line of code to do that. Sweet deal, right?

In this post, we will only go through the basics, enough for us to play with “more flexible” situations in the future.

We will cover the following:

Publishing a custom model to a github repository
Loading an existing model from a github repository

The materials for this tutorial can be found in this repository: github.com/lbleal1/torch-hub-test.

Big Picture

This big picture only shows the basic components to publish, and access and load models. For publishing, you at least need three things: model parameters, model class , hubconf.py before you push it to your github repository. For accessing and loading, you need to have the repository name and entrypoint name, then it will be cached into your computer locally. Finally, you can use the one-liner torch.hub.load to load the model.

Publishing a custom model

The minimum components in publishing a custom model again are:

model.py — contains the class defining the model structure and function for loading the model.

2. model_urls where can download the model parameters — I find this better than accessing the model locally because the github repository will only contain the code and you do not need to change anything. Of course, there are cases where accessing locally can be beneficial too.

Make sure to try the url such that if you enter it in your search bar, it should prompt a download instead of going into some specific site.

3. hubconf.py — among other files, this file should be named exactly as it is and should be placed in the outermost part of the repository as seen below in our example:

If you have these, then you can push it now to github.

Load an existing model

There is a general way of doing this but I divided this into two parts to show how to load from the PyTorch repository and from someone’s repository for custom models.

Loading from the PyTorch repository

For the example below, we will be loading resnet18 from the pytorch/vision repository. In most tutorials, you will see that this line # for the HTTP Error 403 torch.hub._validate_not_a_forked_repo=lambda a,b,c: True is not usually added. Upon looking at it, it seems it is a recurring error as many people have been reporting it. A useful thread can be found here for more details.

The next line, model = torch.hub.load(‘pytorch/vision:v0.9.0’, ‘resnet18’, pretrained=True) is quite self-explanatory. Usually the first two consists of the repository branch‘pytorch/vision:v0.9.0’ and the entrypoint name ‘resnet18’ . You usually find the entrypoint names by looking at hubconf.py. Other options such as pretrained=True can also be added.

import torch

# for the HTTP Error 403
torch.hub._validate_not_a_forked_repo=lambda a,b,c: True

model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet18', pretrained=True)
print(model)

Output:

Downloading: "https://github.com/pytorch/vision/zipball/v0.9.0" to C:\Users\ASUS/.cache\torch\hub\v0.9.0.zip
ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (1): BasicBlock(
      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer2): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer3): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer4): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=512, out_features=1000, bias=True)
)

Running this code will cache the repository as seen in this part: Downloading: “https://github.com/pytorch/vision/zipball/v0.9.0" to C:\Users\ASUS/.cache\torch\hub\v0.9.0.zip. You can try looking at this folder <userpath>\.cache\torch\hub and you’ll see the repository folders you’ve used so far.

Loading from someone’s repository

This is no different from the previous one but you can remove the line that deals with the error.

import torch

model = torch.hub.load('lbleal1/torch-hub-test', 
                       'model', pretrained=True)
print(model)

Output:

Downloading: "https://github.com/lbleal1/torch-hub-test/zipball/main" to C:\Users\ASUS/.cache\torch\hub\main.zip
Downloading: "https://github.com/lbleal1/torch-hub-test/raw/main/model_resources/iris_classifier_params.pt" to C:\Users\ASUS/.cache\torch\hub\checkpoints\iris_classifier_params.pt
100%|██████████████████████████████████████████████████████████████████| 2.24k/2.24k [00:00<00:00, 1.71MB/s] 
Model(
  (layer1): Linear(in_features=4, out_features=16, bias=True)
  (layer2): Linear(in_features=16, out_features=3, bias=True)
)

You should note that if the repository is updated and you want to use the changes, you should remove cached repository folder in <userpath>\.cache\torch\hub ;otherwise, it will keep using the old repository.

With these, you successfully loaded the model and decide what to do next.

Other Notes

If you need more examples, I find that pytorch/vision repository is a good guide. You can also look at torch.hub documentation for more details.

I am keeping this for my own notes and hopefully, it will save your time too.