| """Run checkpoint-backed RemoteCLIP retrieval inference.""" |
|
|
| import argparse |
| import importlib.util |
| from pathlib import Path |
| import numpy as np |
| import torch |
| import yaml |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml") |
| parser.add_argument("--data", type=Path); parser.add_argument("--checkpoint", type=Path) |
| parser.add_argument("--output-dir", type=Path); parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto") |
| args = parser.parse_args(); config = yaml.safe_load(args.config.read_text()) |
| checkpoint_path = args.checkpoint or ROOT / config["paths"]["checkpoint"] |
| if not checkpoint_path.is_file(): raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}") |
| spec = importlib.util.spec_from_file_location("remoteclip", ROOT / "model/remoteclip.py") |
| module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module) |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| model = module.RemoteCLIP(vocabulary_size=config["data"]["vocabulary_size"], |
| context_length=config["data"]["context_length"], |
| eot_token_id=config["data"]["eot_token_id"], **config["model"]) |
| model.load_state_dict(checkpoint["model"]) |
| use_cuda = torch.cuda.is_available() and args.device != "cpu" |
| if args.device == "cuda" and not use_cuda: raise RuntimeError("CUDA requested but unavailable") |
| device = torch.device("cuda" if use_cuda else "cpu"); model.to(device).eval() |
| data_path = args.data or ROOT / config["data"]["root"] / "test.npz" |
| train_spec = importlib.util.spec_from_file_location("remoteclip_train", ROOT / "scripts/train.py") |
| train_module = importlib.util.module_from_spec(train_spec); train_spec.loader.exec_module(train_module) |
| dataset = train_module.PairDataset(data_path, config); archive = np.load(data_path) |
| with torch.inference_mode(): |
| image_features = model.encode_image(torch.from_numpy(archive["images"]).to(device)) |
| text_features = model.encode_text(torch.from_numpy(archive["tokens"]).to(device)) |
| output_dir = args.output_dir or ROOT / config["paths"]["inference_dir"]; output_dir.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output_dir / "retrieval.npz", similarities=(image_features @ text_features.T).cpu().numpy(), |
| image_features=image_features.cpu().numpy(), text_features=text_features.cpu().numpy(), |
| pair_ids=archive["pair_ids"], images=archive["images"], checkpoint=np.asarray(str(checkpoint_path)), |
| data_source=archive["data_source"] if "data_source" in archive else np.asarray("provided"), |
| protocol=archive["protocol"] if "protocol" in archive else np.asarray("provided_npz")) |
| print(f"inference={output_dir / 'retrieval.npz'} checkpoint={checkpoint_path}") |
|
|
|
|
| if __name__ == "__main__": main() |
|
|
