Do not load pretrained VGG (open-mmlab#466)

* Don not load pretrained VGG

* Add load_pretrained_vgg

Co-authored-by: liyinshuo <liyinshuo@sensetime.com>

Author: Yshuo-Li

mmedit/models/extractors/lte.py

@@ -18,9 +18,18 @@ class LTE(nn.Module):
         requires_grad (bool): Require grad or not. Default: True.
         pixel_range (float): Pixel range of geature. Default: 1.
         pretrained (str): Path for pretrained model. Default: None.
+        load_pretrained_vgg (bool): Load pretrained VGG from torchvision.
+            Default: True.
+            Train: must load pretrained VGG
+            Eval: needn't load pretrained VGG, because we will load pretrained
+                LTE.
     """
 
-    def __init__(self, requires_grad=True, pixel_range=1., pretrained=None):
+    def __init__(self,
+                 requires_grad=True,
+                 pixel_range=1.,
+                 pretrained=None,
+                 load_pretrained_vgg=True):
         super().__init__()
 
         vgg_mean = (0.485, 0.456, 0.406)
@@ -30,7 +39,8 @@ def __init__(self, requires_grad=True, pixel_range=1., pretrained=None):
             pixel_range=pixel_range, img_mean=vgg_mean, img_std=vgg_std)
 
         # use vgg19 weights to initialize
-        vgg_pretrained_features = models.vgg19(pretrained=True).features
+        vgg_pretrained_features = models.vgg19(
+            pretrained=load_pretrained_vgg).features
 
         self.slice1 = torch.nn.Sequential()
         self.slice2 = torch.nn.Sequential()

tests/test_models/test_extractors/test_lte.py

@@ -6,7 +6,11 @@
 
 def test_lte():
     model_cfg = dict(
-        type='LTE', requires_grad=False, pixel_range=1., pretrained=None)
+        type='LTE',
+        requires_grad=False,
+        pixel_range=1.,
+        pretrained=None,
+        load_pretrained_vgg=False)
 
     lte = build_component(model_cfg)
     assert lte.__class__.__name__ == 'LTE'

tests/test_models/test_restorers/test_ttsr.py

@@ -70,7 +70,7 @@ def test_ttsr():
             out_channels=3,
             mid_channels=64,
             num_blocks=(16, 16, 8, 4)),
-        extractor=dict(type='LTE'),
+        extractor=dict(type='LTE', load_pretrained_vgg=False),
         transformer=dict(type='SearchTransformer'),
         discriminator=dict(type='TTSRDiscriminator', in_size=64),
         pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'),