Search code examples
pythondeep-learningpytorchpre-trained-modeltorchvision

Is there any way I can download the pre-trained models available in PyTorch to a specific path?


I am referring to the models that can be found here: https://pytorch.org/docs/stable/torchvision/models.html#torchvision-models


Solution

  • As, @dennlinger mentioned in his answer : torch.utils.model_zoo, is being internally called when you load a pre-trained model.

    More specifically, the method: torch.utils.model_zoo.load_url() is being called every time a pre-trained model is loaded. The documentation for the same, mentions:

    The default value of model_dir is $TORCH_HOME/models where $TORCH_HOME defaults to ~/.torch.

    The default directory can be overridden with the $TORCH_HOME environment variable.

    This can be done as follows:

    import torch 
    import torchvision
    import os
    
    # Suppose you are trying to load pre-trained resnet model in directory- models\resnet
    
    os.environ['TORCH_HOME'] = 'models\\resnet' #setting the environment variable
    resnet = torchvision.models.resnet18(pretrained=True)
    

    I came across the above solution by raising an issue in the PyTorch's GitHub repository: https://github.com/pytorch/vision/issues/616

    This led to an improvement in the documentation i.e. the solution mentioned above.