wtjiang98 / PSGAN

PyTorch code for "PSGAN: Pose and Expression Robust Spatial-Aware GAN for Customizable Makeup Transfer" (CVPR 2020 Oral)

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

ImportError: cannot import name 'load_state_dict_from_url' from 'torchvision.models.vgg'

china-mashina opened this issue · comments

When I tried to run demo.py I got the following error.
The problem is that load_state_dict_from_url, model_urls, cfgs were deprecated in more recent version of torchvision.
You can resolve this by changing demo.py in the following way:
replace
from torchvision.models.vgg import load_state_dict_from_url, model_urls, cfgs
with
import torchvision.models as models

and replace

def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs):
    if pretrained:
        kwargs['init_weights'] = False
    model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
    if pretrained:
        state_dict = load_state_dict_from_url(model_urls[arch],
                                              progress=progress)
        model.load_state_dict(state_dict)
    return model
def vgg16(pretrained=False, progress=True, **kwargs):
    r"""VGG 16-layer model (configuration "D")
    `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _vgg('vgg16', 'D', False, pretrained, progress, **kwargs)

with

def vgg16(pretrained=False, progress=True, **kwargs):
    """
    VGG 16-layer model (configuration "D")
    "Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    # The `models.vgg16` function now directly supports loading a pretrained model
    # with the `pretrained` and `progress` arguments.
    model = models.vgg16(pretrained=pretrained, progress=progress, **kwargs)
    return model

I spent a lot of time to figure this out, hopefully someone else will read this and won't loose as much time.