Pytorch Hub is a pre-trained model repository designed to facilitate research reproducibility.
Pytorch Hub supports publishing pre-trained models(model definitions and pre-trained weights) to a github repository by adding a simple hubconf.py
file;
hubconf.py
can have multiple entrypoints. Each entrypoint is defined as a python function (example: a pre-trained model you want to publish).
def entrypoint_name(*args, **kwargs):
# args & kwargs are optional, for models which take positional/keyword arguments.
...
Here is a code snippet specifies an entrypoint for resnet18
model if we expand the implementation in pytorch/vision/hubconf.py
. In most case importing the right function in hubconf.py
is sufficient. Here we just want to use the expanded version as an example to show how it works. You can see the full script in pytorch/vision repo
dependencies = ['torch']
from torchvision.models.resnet import resnet18 as _resnet18
# resnet18 is the name of entrypoint
def resnet18(pretrained=False, **kwargs):
""" # This docstring shows up in hub.help()
Resnet18 model
pretrained (bool): kwargs, load pretrained weights into the model
"""
# Call the model, load pretrained weights
model = _resnet18(pretrained=pretrained, **kwargs)
return model
dependencies
variable is a list of package names required to load the model. Note this might be slightly different from dependencies required for training a model.args
and kwargs
are passed along to the real callable function.torch.hub.list()
.torch.hub.load_state_dict_from_url()
. If less than 2GB, it’s recommended to attach it to a project release and use the url from the release. In the example above torchvision.models.resnet.resnet18
handles pretrained
, alternatively you can put the following logic in the entrypoint definition.if pretrained:
# For checkpoint saved in local github repo, e.g. <RELATIVE_PATH_TO_CHECKPOINT>=weights/save.pth
dirname = os.path.dirname(__file__)
checkpoint = os.path.join(dirname, <RELATIVE_PATH_TO_CHECKPOINT>)
state_dict = torch.load(checkpoint)
model.load_state_dict(state_dict)
# For checkpoint saved elsewhere
checkpoint = 'https://download.pytorch.org/models/resnet18-5c106cde.pth'
model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=False))
Pytorch Hub provides convenient APIs to explore all available models in hub through torch.hub.list()
, show docstring and examples through torch.hub.help()
and load the pre-trained models using torch.hub.load()
torch.hub.list
(github, force_reload=False)[source]
List all entrypoints available in github hubconf.
Parameters:
Returns:
Return type:
Example:
>>> entrypoints = torch.hub.list('pytorch/vision', force_reload=True)
torch.hub.help
(github, model, force_reload=False)[source]
Show the docstring of entrypoint model.
Parameters:
Example:
>>> print(torch.hub.help('pytorch/vision', 'resnet18', force_reload=True))
torch.hub.load
(github, model, *args, **kwargs)[source]
Load a model from a github repo, with pretrained weights.
Parameters:
Returns:
Example:
>>> model = torch.hub.load('pytorch/vision', 'resnet50', pretrained=True)
torch.hub.download_url_to_file
(url, dst, hash_prefix=None, progress=True)[source]
Download object at the given URL to a local path.
Parameters:
Example:
>>> torch.hub.download_url_to_file('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth', '/tmp/temporary_file')
torch.hub.load_state_dict_from_url
(url, model_dir=None, map_location=None, progress=True, check_hash=False)[source]
Loads the Torch serialized object at the given URL.
If downloaded file is a zip file, it will be automatically decompressed.
If the object is already present in model_dir, it’s deserialized and returned. The default value of model_dir is $TORCH_HOME/checkpoints
where environment variable $TORCH_HOME
defaults to $XDG_CACHE_HOME/torch
. $XDG_CACHE_HOME
follows the X Design Group specification of the Linux filesytem layout, with a default value ~/.cache
if not set.
Parameters:
filename-<sha256>.ext
where <sha256>
is the first eight or more digits of the SHA256 hash of the contents of the file. The hash is used to ensure unique names and to verify the contents of the file. Default: FalseExample:
>>> state_dict = torch.hub.load_state_dict_from_url('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth')
Note that *args, **kwargs
in torch.load()
are used to instantiate a model. After you loaded a model, how can you find out what you can do with the model? A suggested workflow is
dir(model)
to see all available methods of the model.help(model.foo)
to check what arguments model.foo
takes to run.To help users explore without referring to documentation back and forth, we strongly recommend repo owners make function help messages clear and succinct. It’s also helpful to include a minimal working example.
The locations are used in the order of
hub.set_dir(<PATH_TO_HUB_DIR>)
$TORCH_HOME/hub
, if environment variable TORCH_HOME
is set.$XDG_CACHE_HOME/torch/hub
, if environment variable XDG_CACHE_HOME
is set.~/.cache/torch/hub
torch.hub.set_dir
(d)[source]
Optionally set hub_dir to a local dir to save downloaded models & weights.
If set_dir
is not called, default path is $TORCH_HOME/hub
where environment variable $TORCH_HOME
defaults to $XDG_CACHE_HOME/torch
. $XDG_CACHE_HOME
follows the X Design Group specification of the Linux filesytem layout, with a default value ~/.cache
if the environment variable is not set.
Parameters:
By default, we don’t clean up files after loading it. Hub uses the cache by default if it already exists in hub_dir
.
Users can force a reload by calling hub.load(..., force_reload=True)
. This will delete the existing github folder and downloaded weights, reinitialize a fresh download. This is useful when updates are published to the same branch, users can keep up with the latest release.
Torch hub works by importing the package as if it was installed. There’re some side effects introduced by importing in Python. For example, you can see new items in Python caches sys.modules
and sys.path_importer_cache
which is normal Python behavior.
A known limitation that worth mentioning here is user CANNOT load two different branches of the same repo in the same python process. It’s just like installing two packages with the same name in Python, which is not good. Cache might join the party and give you surprises if you actually try that. Of course it’s totally fine to load them in separate processes.