From 6e3c3b62b220eb6c0de73671a65cabf070b7e1f2 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 11 Jul 2020 15:03:12 -0700 Subject: [PATCH] expired pytorch hub cache bug fix #365 --- hubconf.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/hubconf.py b/hubconf.py index 7ca9d93..8455ae2 100644 --- a/hubconf.py +++ b/hubconf.py @@ -27,15 +27,19 @@ def create(name, pretrained, channels, classes): Returns: pytorch model """ - config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) # model.yaml path - model = Model(config, channels, classes) - if pretrained: - ckpt = '%s.pt' % name # checkpoint filename - google_utils.attempt_download(ckpt) # download if not found locally - state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32 - state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter - model.load_state_dict(state_dict, strict=False) # load - return model + try: + config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) # model.yaml path + model = Model(config, channels, classes) + if pretrained: + ckpt = '%s.pt' % name # checkpoint filename + google_utils.attempt_download(ckpt) # download if not found locally + state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32 + state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter + model.load_state_dict(state_dict, strict=False) # load + return model + except Exception as e: + help_url = 'https://github.com/ultralytics/yolov5/issues/36' + print('%s\nCache is out of date. Delete cache and retry. See %s for help.' % (e, help_url)) def yolov5s(pretrained=False, channels=3, classes=80):