Other models from torchvision

[madiltalay]

Hi, I wanted to use light-weight base models from torchvision like mobilenet, squeezenet and densenet. Any tips on how to edit the code?

[potterhsu]

To add a new backbone, you should do the following steps:

1. Create an inherited class from backbone.base.Base, then override the method features

   class DenseNet161(backbone.base.Base):
       def __init__(self, pretrained: bool):
           super().__init__(pretrained)
   
       def features(self) -> Tuple[nn.Module, nn.Module, int, int]:
           densenet161 = torchvision.models.densenet161(pretrained=self._pretrained)
           ...

2. In backbone/base.py, add a new branch for your new backbone, and don't forget to extend the options

3. Specific new backbone to script

   $ python train.py -s=xxx -b=densenet161
   

Hope this helps.

[madiltalay]

Thanks for the prompt reply, I'm now working to get the right dimensions for

1. input
2. features
3. hidden