152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
import warnings
|
|
from dataclasses import dataclass
|
|
from functools import cached_property
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .util import get_model_path, save_config
|
|
|
|
|
|
@dataclass
|
|
class OpenCLIPModelConfig:
|
|
name: str
|
|
pretrained: str
|
|
|
|
@cached_property
|
|
def model_config(self) -> dict[str, Any]:
|
|
import open_clip
|
|
|
|
config: dict[str, Any] | None = open_clip.get_model_config(self.name)
|
|
if config is None:
|
|
raise ValueError(f"Unknown model {self.name}")
|
|
return config
|
|
|
|
@property
|
|
def image_size(self) -> int:
|
|
image_size: int = self.model_config["vision_cfg"]["image_size"]
|
|
return image_size
|
|
|
|
@property
|
|
def sequence_length(self) -> int:
|
|
context_length: int = self.model_config["text_cfg"].get("context_length", 77)
|
|
return context_length
|
|
|
|
|
|
def to_onnx(
|
|
model_cfg: OpenCLIPModelConfig,
|
|
opset_version: int,
|
|
output_dir_visual: Path | str | None = None,
|
|
output_dir_textual: Path | str | None = None,
|
|
cache: bool = True,
|
|
) -> tuple[Path | None, Path | None]:
|
|
visual_path = None
|
|
textual_path = None
|
|
if output_dir_visual is not None:
|
|
output_dir_visual = Path(output_dir_visual)
|
|
visual_path = get_model_path(output_dir_visual)
|
|
|
|
if output_dir_textual is not None:
|
|
output_dir_textual = Path(output_dir_textual)
|
|
textual_path = get_model_path(output_dir_textual)
|
|
|
|
if cache and ((textual_path is None or textual_path.exists()) and (visual_path is None or visual_path.exists())):
|
|
print(f"Models {textual_path} and {visual_path} already exist, skipping")
|
|
return visual_path, textual_path
|
|
|
|
import open_clip
|
|
import torch
|
|
from transformers import AutoTokenizer
|
|
|
|
torch.backends.mha.set_fastpath_enabled(False)
|
|
|
|
model = open_clip.create_model(
|
|
model_cfg.name,
|
|
pretrained=model_cfg.pretrained,
|
|
jit=False,
|
|
require_pretrained=True,
|
|
)
|
|
|
|
text_vision_cfg = open_clip.get_model_config(model_cfg.name)
|
|
|
|
model.eval()
|
|
for param in model.parameters():
|
|
param.requires_grad_(False)
|
|
|
|
if visual_path is not None and output_dir_visual is not None:
|
|
if not cache or not visual_path.exists():
|
|
save_config(
|
|
open_clip.get_model_preprocess_cfg(model),
|
|
output_dir_visual / "preprocess_cfg.json",
|
|
)
|
|
save_config(text_vision_cfg, output_dir_visual.parent / "config.json")
|
|
_export_image_encoder(model, model_cfg, visual_path, opset_version)
|
|
else:
|
|
print(f"Model {visual_path} already exists, skipping")
|
|
|
|
if textual_path is not None and output_dir_textual is not None:
|
|
if not cache or not textual_path.exists():
|
|
tokenizer_name = text_vision_cfg["text_cfg"].get("hf_tokenizer_name", "openai/clip-vit-base-patch32")
|
|
AutoTokenizer.from_pretrained(tokenizer_name).save_pretrained(output_dir_textual)
|
|
_export_text_encoder(model, model_cfg, textual_path, opset_version)
|
|
else:
|
|
print(f"Model {textual_path} already exists, skipping")
|
|
return visual_path, textual_path
|
|
|
|
|
|
def _export_image_encoder(
|
|
model: Any, model_cfg: OpenCLIPModelConfig, output_path: Path | str, opset_version: int
|
|
) -> None:
|
|
import torch
|
|
|
|
output_path = Path(output_path)
|
|
|
|
def encode_image(image: torch.Tensor) -> torch.Tensor:
|
|
output = model.encode_image(image, normalize=True)
|
|
assert isinstance(output, torch.Tensor)
|
|
return output
|
|
|
|
model.forward = encode_image
|
|
|
|
args = (torch.randn(1, 3, model_cfg.image_size, model_cfg.image_size),)
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore", UserWarning)
|
|
torch.onnx.export(
|
|
model,
|
|
args,
|
|
output_path.as_posix(),
|
|
input_names=["image"],
|
|
output_names=["embedding"],
|
|
opset_version=opset_version,
|
|
# dynamic_axes={"image": {0: "batch_size"}},
|
|
)
|
|
|
|
|
|
def _export_text_encoder(
|
|
model: Any, model_cfg: OpenCLIPModelConfig, output_path: Path | str, opset_version: int
|
|
) -> None:
|
|
import torch
|
|
|
|
output_path = Path(output_path)
|
|
|
|
def encode_text(text: torch.Tensor) -> torch.Tensor:
|
|
output = model.encode_text(text, normalize=True)
|
|
assert isinstance(output, torch.Tensor)
|
|
return output
|
|
|
|
model.forward = encode_text
|
|
|
|
args = (torch.ones(1, model_cfg.sequence_length, dtype=torch.int32),)
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore", UserWarning)
|
|
torch.onnx.export(
|
|
model,
|
|
args,
|
|
output_path.as_posix(),
|
|
input_names=["text"],
|
|
output_names=["embedding"],
|
|
opset_version=opset_version,
|
|
# dynamic_axes={"text": {0: "batch_size"}},
|
|
)
|