Repository: Lightning-Universe/lightning-flash

Issue:
Flash Requires Internet to Load a Local Checkpoint 
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

Steps to reproduce the behavior:

1. Turn of internet on device
2. Load checkpoint from local file
```python
clf = ImageClassifier.load_from_checkpoint('/kaggle/input/baserazcrmodels/razcr_resnet50_base_model.pt')
```
3. Error caused by ImageClassifer Init which pulls from torchvision on init even though it's loading from a local checkpoint
```
---------------------------------------------------------------------------
gaierror                                  Traceback (most recent call last)
/opt/conda/lib/python3.7/urllib/request.py in do_open(self, http_class, req, **http_conn_args)
   1349                 h.request(req.get_method(), req.selector, req.data, headers,
-> 1350                           encode_chunked=req.has_header('Transfer-encoding'))
   1351             except OSError as err: # timeout error
/opt/conda/lib/python3.7/http/client.py in request(self, method, url, body, headers, encode_chunked)
   1276         """Send a complete request to the server."""
-> 1277         self._send_request(method, url, body, headers, encode_chunked)
   1278 
/opt/conda/lib/python3.7/http/client.py in _send_request(self, method, url, body, headers, encode_chunked)
   1322             body = _encode(body, 'body')
-> 1323         self.endheaders(body, encode_chunked=encode_chunked)
   1324 
/opt/conda/lib/python3.7/http/client.py in endheaders(self, message_body, encode_chunked)
   1271             raise CannotSendHeader()
-> 1272         self._send_output(message_body, encode_chunked=encode_chunked)
   1273 
/opt/conda/lib/python3.7/http/client.py in _send_output(self, message_body, encode_chunked)
   1031         del self._buffer[:]
-> 1032         self.send(msg)
   1033 
/opt/conda/lib/python3.7/http/client.py in send(self, data)
    971             if self.auto_open:
--> 972                 self.connect()
    973             else:
/opt/conda/lib/python3.7/http/client.py in connect(self)
   1438 
-> 1439             super().connect()
   1440 
/opt/conda/lib/python3.7/http/client.py in connect(self)
    943         self.sock = self._create_connection(
--> 944             (self.host,self.port), self.timeout, self.source_address)
    945         self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
/opt/conda/lib/python3.7/socket.py in create_connection(address, timeout, source_address)
    706     err = None
--> 707     for res in getaddrinfo(host, port, 0, SOCK_STREAM):
    708         af, socktype, proto, canonname, sa = res
/opt/conda/lib/python3.7/socket.py in getaddrinfo(host, port, family, type, proto, flags)
    751     addrlist = []
--> 752     for res in _socket.getaddrinfo(host, port, family, type, proto, flags):
    753         af, socktype, proto, canonname, sa = res
gaierror: [Errno -3] Temporary failure in name resolution
During handling of the above exception, another exception occurred:
URLError                                  Traceback (most recent call last)
<ipython-input-8-d3a1d7810d85> in <module>
      5 #                       num_classes=len(columns))
      6 
----> 7 clf = ImageClassifier.load_from_checkpoint("../input/baserazcrmodels/razcr_resnet50_base_model.pt")
/opt/conda/lib/python3.7/site-packages/pytorch_lightning/core/saving.py in load_from_checkpoint(cls, checkpoint_path, map_location, hparams_file, strict, **kwargs)
    155         checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
    156 
--> 157         model = cls._load_model_state(checkpoint, strict=strict, **kwargs)
    158         return model
    159 
/opt/conda/lib/python3.7/site-packages/pytorch_lightning/core/saving.py in _load_model_state(cls, checkpoint, strict, **cls_kwargs_new)
    196             _cls_kwargs = {k: v for k, v in _cls_kwargs.items() if k in cls_init_args_name}
    197 
--> 198         model = cls(**_cls_kwargs)
    199 
    200         # give model a chance to load something
/opt/conda/lib/python3.7/site-packages/flash/vision/classification/model.py in __init__(self, num_classes, backbone, pretrained, loss_fn, optimizer, metrics, learning_rate, multilabel)
     61         self.save_hyperparameters()
     62 
---> 63         self.backbone, num_features = backbone_and_num_features(backbone, pretrained=pretrained)
     64 
     65         self.head = nn.Sequential(
/opt/conda/lib/python3.7/site-packages/flash/vision/backbones.py in backbone_and_num_features(model_name, fpn, pretrained, trainable_backbone_layers, **kwargs)
     69 
     70     if model_name in TORCHVISION_MODELS:
---> 71         return torchvision_backbone_and_num_features(model_name, pretrained)
     72 
     73     raise ValueError(f"{model_name} is not supported yet.")
/opt/conda/lib/python3.7/site-packages/flash/vision/backbones.py in torchvision_backbone_and_num_features(model_name, pretrained)
    128 
    129     elif model_name in RESNET_MODELS:
--> 130         model = model(pretrained=pretrained)
    131         # remove the last two layers & turn it into a Sequential model
    132         backbone = nn.Sequential(*list(model.children())[:-2])
/opt/conda/lib/python3.7/site-packages/torchvision/models/resnet.py in resnet50(pretrained, progress, **kwargs)
    263     """
    264     return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
--> 265                    **kwargs)
    266 
    267 
/opt/conda/lib/python3.7/site-packages/torchvision/models/resnet.py in _resnet(arch, block, layers, pretrained, progress, **kwargs)
    225     if pretrained:
    226         state_dict = load_state_dict_from_url(model_urls[arch],
--> 227                                               progress=progress)
    228         model.load_state_dict(state_dict)
    229     return model
/opt/conda/lib/python3.7/site-packages/torch/hub.py in load_state_dict_from_url(url, model_dir, map_location, progress, check_hash, file_name)
    553             r = HASH_REGEX.search(filename)  # r is Optional[Match[str]]
    554             hash_prefix = r.group(1) if r else None
--> 555         download_url_to_file(url, cached_file, hash_prefix, progress=progress)
    556 
    557     if _is_legacy_zip_format(cached_file):
/opt/conda/lib/python3.7/site-packages/torch/hub.py in download_url_to_file(url, dst, hash_prefix, progress)
    423     # certificates in older Python
    424     req = Request(url, headers={"User-Agent": "torch.hub"})
--> 425     u = urlopen(req)
    426     meta = u.info()
    427     if hasattr(meta, 'getheaders'):
/opt/conda/lib/python3.7/urllib/request.py in urlopen(url, data, timeout, cafile, capath, cadefault, context)
    220     else:
    221         opener = _opener
--> 222     return opener.open(url, data, timeout)
    223 
    224 def install_opener(opener):
/opt/conda/lib/python3.7/urllib/request.py in open(self, fullurl, data, timeout)
    523             req = meth(req)
    524 
--> 525         response = self._open(req, data)
    526 
    527         # post-process response
/opt/conda/lib/python3.7/urllib/request.py in _open(self, req, data)
    541         protocol = req.type
    542         result = self._call_chain(self.handle_open, protocol, protocol +
--> 543                                   '_open', req)
    544         if result:
    545             return result
/opt/conda/lib/python3.7/urllib/request.py in _call_chain(self, chain, kind, meth_name, *args)
    501         for handler in handlers:
    502             func = getattr(handler, meth_name)
--> 503             result = func(*args)
    504             if result is not None:
    505                 return result
/opt/conda/lib/python3.7/urllib/request.py in https_open(self, req)
   1391         def https_open(self, req):
   1392             return self.do_open(http.client.HTTPSConnection, req,
-> 1393                 context=self._context, check_hostname=self._check_hostname)
   1394 
   1395         https_request = AbstractHTTPHandler.do_request_
/opt/conda/lib/python3.7/urllib/request.py in do_open(self, http_class, req, **http_conn_args)
   1350                           encode_chunked=req.has_header('Transfer-encoding'))
   1351             except OSError as err: # timeout error
-> 1352                 raise URLError(err)
   1353             r = h.getresponse()
   1354         except:
URLError: <urlopen error [Errno -3] Temporary failure in name resolution>
```


Discussion:
Temp fix sent to Ari to unblock for the Task.
> Temp fix sent to Ari to unblock for the Task.

Can you please link the PR? 
@Borda The temp fix is this, to check if it's called from `_load_model_state`.

```python
calling_function = inspect.getframeinfo(sys._getframe(1))[2]
if calling_function == "_load_model_state":
   pretrained = False
```
@aribornstein can you confirm that with `0.2.2rc2` if you still get this error ?
@aribornstein I have been able to reproduce this. For anyone interested you have to delete the cached torchvision models from `~/.cache/torch/hub/checkpoints` then the error appears. Working on a fix now :smiley:

Resolving patch:
diff --git a/flash/vision/backbones.py b/flash/vision/backbones.py
--- a/flash/vision/backbones.py
+++ b/flash/vision/backbones.py
@@ -11,14 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import functools
 import os
+import urllib.error
 import warnings
 from functools import partial
 from typing import Tuple
 
 from pytorch_lightning import LightningModule
-from pytorch_lightning.utilities import _BOLTS_AVAILABLE
+from pytorch_lightning.utilities import _BOLTS_AVAILABLE, rank_zero_warn
 from torch import nn as nn
 from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
 
@@ -51,6 +52,24 @@
 OBJ_DETECTION_BACKBONES = FlashRegistry("backbones")
 
 
+def catch_url_error(fn):
+
+    @functools.wraps(fn)
+    def wrapper(pretrained=False, **kwargs):
+        try:
+            return fn(pretrained=pretrained, **kwargs)
+        except urllib.error.URLError:
+            result = fn(pretrained=False, **kwargs)
+            rank_zero_warn(
+                "Failed to download pretrained weights for the selected backbone. The backbone has been created with"
+                " `pretrained=False` instead. If you are loading from a local checkpoint, this warning can be safely"
+                " ignored.", UserWarning
+            )
+            return result
+
+    return wrapper
+
+
 @IMAGE_CLASSIFIER_BACKBONES(name="simclr-imagenet", namespace="vision", package="bolts")
 def load_simclr_imagenet(path_or_url: str = f"{ROOT_S3_BUCKET}/simclr/bolts_simclr_imagenet/simclr_imagenet.ckpt", **_):
     simclr: LightningModule = SimCLR.load_from_checkpoint(path_or_url, strict=False)
@@ -83,7 +102,7 @@ def _fn_mobilenet_vgg(model_name: str, pretrained: bool = True) -> Tuple[nn.Modu
         _type = "mobilenet" if model_name in MOBILENET_MODELS else "vgg"
 
         IMAGE_CLASSIFIER_BACKBONES(
-            fn=partial(_fn_mobilenet_vgg, model_name),
+            fn=catch_url_error(partial(_fn_mobilenet_vgg, model_name)),
             name=model_name,
             namespace="vision",
             package="torchvision",
@@ -99,7 +118,7 @@ def _fn_resnet(model_name: str, pretrained: bool = True) -> Tuple[nn.Module, int
             return backbone, num_features
 
         IMAGE_CLASSIFIER_BACKBONES(
-            fn=partial(_fn_resnet, model_name),
+            fn=catch_url_error(partial(_fn_resnet, model_name)),
             name=model_name,
             namespace="vision",
             package="torchvision",
@@ -118,7 +137,10 @@ def _fn_resnet_fpn(
             return backbone, 256
 
         OBJ_DETECTION_BACKBONES(
-            fn=partial(_fn_resnet_fpn, model_name), name=model_name, package="torchvision", type="resnet-fpn"
+            fn=catch_url_error(partial(_fn_resnet_fpn, model_name)),
+            name=model_name,
+            package="torchvision",
+            type="resnet-fpn"
         )
 
     for model_name in DENSENET_MODELS:
@@ -130,7 +152,7 @@ def _fn_densenet(model_name: str, pretrained: bool = True) -> Tuple[nn.Module, i
             return backbone, num_features
 
         IMAGE_CLASSIFIER_BACKBONES(
-            fn=partial(_fn_densenet, model_name),
+            fn=catch_url_error(partial(_fn_densenet, model_name)),
             name=model_name,
             namespace="vision",
             package="torchvision",
@@ -156,5 +178,5 @@ def _fn_timm(
             return backbone, num_features
 
         IMAGE_CLASSIFIER_BACKBONES(
-            fn=partial(_fn_timm, model_name), name=model_name, namespace="vision", package="timm"
+            fn=catch_url_error(partial(_fn_timm, model_name)), name=model_name, namespace="vision", package="timm"
         )
diff --git a/flash/vision/detection/model.py b/flash/vision/detection/model.py
--- a/flash/vision/detection/model.py
+++ b/flash/vision/detection/model.py
@@ -137,8 +137,8 @@ def get_model(
                 )
         else:
             backbone_model, num_features = ObjectDetector.backbones.get(backbone)(
-                pretrained_backbone,
-                trainable_backbone_layers,
+                pretrained=pretrained_backbone,
+                trainable_layers=trainable_backbone_layers,
                 **kwargs,
             )
             backbone_model.out_channels = num_features


Test patch:
diff --git a/tests/vision/test_backbones.py b/tests/vision/test_backbones.py
--- a/tests/vision/test_backbones.py
+++ b/tests/vision/test_backbones.py
@@ -1,8 +1,10 @@
+import urllib.error
+
 import pytest
 from pytorch_lightning.utilities import _BOLTS_AVAILABLE, _TORCHVISION_AVAILABLE
 
 from flash.utils.imports import _TIMM_AVAILABLE
-from flash.vision.backbones import IMAGE_CLASSIFIER_BACKBONES
+from flash.vision.backbones import catch_url_error, IMAGE_CLASSIFIER_BACKBONES
 
 
 @pytest.mark.parametrize(["backbone", "expected_num_features"], [
@@ -17,3 +19,13 @@ def test_image_classifier_backbones_registry(backbone, expected_num_features):
     backbone_model, num_features = backbone_fn(pretrained=False)
     assert backbone_model
     assert num_features == expected_num_features
+
+
+def test_pretrained_backbones_catch_url_error():
+
+    def raise_error_if_pretrained(pretrained=False):
+        if pretrained:
+            raise urllib.error.URLError('Test error')
+
+    with pytest.warns(UserWarning, match="Failed to download pretrained weights"):
+        catch_url_error(raise_error_if_pretrained)(pretrained=True)