91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
""" OpenAI pretrained model functions
|
|
|
|
Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
|
|
"""
|
|
|
|
import os
|
|
import warnings
|
|
from typing import List, Optional, Union
|
|
|
|
import torch
|
|
|
|
from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
|
|
from .model import build_model_from_openai_state_dict, convert_weights_to_lp, get_cast_dtype
|
|
from .pretrained import get_pretrained_url, list_pretrained_models_by_tag, download_pretrained_from_url
|
|
|
|
__all__ = ["list_openai_models", "load_openai_model"]
|
|
|
|
|
|
def list_openai_models() -> List[str]:
|
|
"""Returns the names of available CLIP models"""
|
|
return list_pretrained_models_by_tag('openai')
|
|
|
|
|
|
def load_openai_model(
|
|
name: str,
|
|
precision: Optional[str] = None,
|
|
device: Optional[Union[str, torch.device]] = None,
|
|
cache_dir: Optional[str] = None,
|
|
):
|
|
"""Load a CLIP model
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
|
|
precision: str
|
|
Model precision, if None defaults to 'fp32' if device == 'cpu' else 'fp16'.
|
|
device : Union[str, torch.device]
|
|
The device to put the loaded model
|
|
cache_dir : Optional[str]
|
|
The directory to cache the downloaded model weights
|
|
|
|
Returns
|
|
-------
|
|
model : torch.nn.Module
|
|
The CLIP model
|
|
preprocess : Callable[[PIL.Image], torch.Tensor]
|
|
A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
|
|
"""
|
|
if device is None:
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
if precision is None:
|
|
precision = 'fp32' if device == 'cpu' else 'fp16'
|
|
|
|
if get_pretrained_url(name, 'openai'):
|
|
model_path = download_pretrained_from_url(get_pretrained_url(name, 'openai'), cache_dir=cache_dir)
|
|
elif os.path.isfile(name):
|
|
model_path = name
|
|
else:
|
|
raise RuntimeError(f"Model {name} not found; available models = {list_openai_models()}")
|
|
|
|
try:
|
|
# loading JIT archive
|
|
model = torch.jit.load(model_path, map_location="cpu").eval()
|
|
state_dict = None
|
|
except RuntimeError:
|
|
# loading saved state dict
|
|
state_dict = torch.load(model_path, map_location="cpu")
|
|
|
|
# Build a non-jit model from the OpenAI jitted model state dict
|
|
cast_dtype = get_cast_dtype(precision)
|
|
try:
|
|
model = build_model_from_openai_state_dict(state_dict or model.state_dict(), cast_dtype=cast_dtype)
|
|
except KeyError:
|
|
sd = {k[7:]: v for k, v in state_dict["state_dict"].items()}
|
|
model = build_model_from_openai_state_dict(sd, cast_dtype=cast_dtype)
|
|
|
|
# model from OpenAI state dict is in manually cast fp16 mode, must be converted for AMP/fp32/bf16 use
|
|
model = model.to(device)
|
|
# FIXME support pure fp16/bf16 precision modes
|
|
if precision != 'fp16':
|
|
model.float()
|
|
if precision == 'bf16':
|
|
# for bf16, convert back to low-precision
|
|
convert_weights_to_lp(model, dtype=torch.bfloat16)
|
|
|
|
# add mean / std attributes for consistency with OpenCLIP models
|
|
model.visual.image_mean = OPENAI_DATASET_MEAN
|
|
model.visual.image_std = OPENAI_DATASET_STD
|
|
return model
|