Can't load densenet model properly

[humzaiqbal]

Hi, I'm trying to load some of the pretrained DenseNet models from this repository using the make_and_restore_model utility, but when I do I notice the following error

/usr/local/lib/python3.6/dist-packages/robustness/datasets.py in get_model(self, arch, pretrained)
    208         """
    209         return imagenet_models.__dict__[arch](num_classes=self.num_classes, 
--> 210                                         pretrained=pretrained)
    211 
    212 class Places365(DataSet):

TypeError: __init__() got an unexpected keyword argument 'pretrained'

after looking at the densenet.py file I notice something, for the DenseNet option that would actually point to the base DenseNet class which isn't configured as the other classes like densenet121 as an example. Looking further at the code it seems there is a simple solution which is to use _densenet instead as that seems to load a base DenseNet model using the proper arguments. What are your thoughts on this? Thanks much am a big fan of the library!

[andrewilyas]

Hi @humzaiqbal thanks for raising this! I will look into it, will probably require a quick update to the library. For now your fix sounds like the appropriate way to deal with it.

[andrewilyas]

Actually, can you provide the code that you are running and what error you get?

[humzaiqbal]

Sure here is a snippet. The densenet_l2_eps0.ckpt that I reference is this from this repository

OUT_DIR = '/tmp/'
NUM_WORKERS = 16
BATCH_SIZE = 512

from robustness import model_utils, datasets, train, defaults
from robustness.datasets import CIFAR, ImageNet
import torch
from cox.utils import Parameters
import cox.store
from torchvision import transforms
imagenet_ds = ImageNet('/tmp/')
densenet , _ = model_utils.make_and_restore_model(arch='DenseNet', dataset=imagenet_ds, 
                                                          resume_path='densenet_l2_eps0.ckpt', parallel=False)

And the error I get is

/usr/local/lib/python3.6/dist-packages/robustness/datasets.py in get_model(self, arch, pretrained)
    208         """
    209         return imagenet_models.__dict__[arch](num_classes=self.num_classes, 
--> 210                                         pretrained=pretrained)
    211 
    212 class Places365(DataSet):

TypeError: __init__() got an unexpected keyword argument 'pretrained'

[Hadisalman]

Hey @humzaiqbal, the models from the robust-models-transfer repo are official PyTorch models and can be loaded using the following code:

OUT_DIR = '/tmp/'
NUM_WORKERS = 16
BATCH_SIZE = 512

from robustness import model_utils, datasets, train, defaults
from robustness.datasets import CIFAR, ImageNet
import torch
from cox.utils import Parameters
import cox.store
from torchvision import transforms
from torchvision import models

imagenet_ds = ImageNet('/tmp/')
densenet , _ = model_utils.make_and_restore_model(arch=models.densenet161(), dataset=imagenet_ds, 
                                                          resume_path='densenet_l2_eps0.ckpt', parallel=False)

Hope this helps!

[humzaiqbal]

Thanks much @Hadisalman :) I think I was confused because in the example notebook I saw the model being loaded by passing a string name for the architecture and assumed thats how it always worked