Function load_model

examples/common/model_loader.py:339–501

(version: str, adaptor_names: str = None, use_huggingface: bool = False, use_local_lib: bool = True,
               device: torch.device = None, return_spatial_features: bool = True, force_reload: bool = False, neck_name: str = None,
               torchhub_repo="NVlabs/RADIO", **kwargs)

# and once it completes, it allows all other ranks to execute, using the now cached weights.
@rank_gate
def load_model(version: str, adaptor_names: str = None, use_huggingface: bool = False, use_local_lib: bool = True,
               device: torch.device = None, return_spatial_features: bool = True, force_reload: bool = False, neck_name: str = None,
               torchhub_repo="NVlabs/RADIO", **kwargs):
    kwargs = {k: v for k, v in kwargs.items() if v is not None}

    if os.path.isfile(version) or 'radio' in version:
        if use_huggingface:
            from transformers import AutoModel, AutoConfig
            hf_repo = 'E-RADIO' if 'eradio' in version else 'RADIO'
            hf_repo = f"nvidia/{hf_repo}"
            config = AutoConfig.from_pretrained(
                hf_repo,
                trust_remote_code=True,
                version=version,
                adaptor_names=adaptor_names,
                **kwargs,
            )
            model: nn.Module = AutoModel.from_pretrained(hf_repo, config=config, trust_remote_code=True, **kwargs)
        elif use_local_lib:
            from hubconf import radio_model
            model, chk = radio_model(version=version, progress=True, adaptor_names=adaptor_names, return_checkpoint=True, neck_name=neck_name, **kwargs)
        else:
            model, chk = torch.hub.load(torchhub_repo, 'radio_model', version=version, progress=True,
                                              adaptor_names=adaptor_names, return_spatial_features=return_spatial_features,
                                              force_reload=force_reload,
                                              return_checkpoint=True, neck_name=neck_name, **kwargs,
            )

        preprocessor = model.make_preprocessor_external()
        info = ModelInfo(model_class='RADIO', model_subtype=version.replace('/', '_'), checkpoint=chk)
    elif version.startswith('dinov2'):
        model = torch.hub.load('facebookresearch/dinov2', version, pretrained=True, force_reload=force_reload, **kwargs)
        model = DinoWrapper(model)

        preprocessor = InputConditioner(1.0, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
        info = ModelInfo(model_class='DINOv2', model_subtype=version.replace('dinov2_', ''))
    elif version.startswith('dinov3'):
        _, chk_path = version.split(',')
        fname = os.path.basename(chk_path)
        model_version = fname[:fname.index('_pretrain')]
        model: nn.Module = torch.hub.load('facebookresearch/dinov3', model_version, pretrained=False)

        chk = torch.load(chk_path, map_location='cpu')
        model.load_state_dict(chk, strict=True)
        model = DinoWrapper(model, patch_attn=False)

        preprocessor = InputConditioner(1.0, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
        info = ModelInfo(model_class='DINOv3', model_subtype=model_version.replace('dinov3_', ''))
    elif version.startswith('open_clip'):
        import open_clip
        _, model_arch, pretrained = version.split(',')
        model = open_clip.create_model(model_arch, pretrained, device=device)
        viz_model = model.visual

        preprocessor = InputConditioner(1.0,
            getattr(viz_model, 'image_mean', open_clip.OPENAI_DATASET_MEAN),
            getattr(viz_model, 'image_std', open_clip.OPENAI_DATASET_STD),
        )