Spaces:
Sleeping
Sleeping
import torch | |
import os | |
import logging | |
from collections import OrderedDict | |
from timm.models import load_checkpoint | |
try: | |
from torch.hub import load_state_dict_from_url | |
except ImportError: | |
from torch.utils.model_zoo import load_url as load_state_dict_from_url | |
def load_pretrained(model, url, filter_fn=None, strict=True): | |
if not url: | |
logging.warning("Pretrained model URL is empty, using random initialization. " | |
"Did you intend to use a `tf_` variant of the model?") | |
return | |
state_dict = load_state_dict_from_url(url, progress=False, map_location='cpu') | |
if filter_fn is not None: | |
state_dict = filter_fn(state_dict) | |
model.load_state_dict(state_dict, strict=strict) | |