File size: 2,499 Bytes
af84b0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """Run Rockchip's exporter using its shipped CLIP embeddings for the dry run."""
import os
from pathlib import Path
import runpy
import sys
from tempfile import NamedTemporaryFile
import numpy as np
import torch
import mmdet.apis
from mmengine.config import Config
root = Path(__file__).resolve().parent
upstream = root / "upstream"
zoo = root.parent.parent.parent / "rknn_model_zoo"
assets = zoo / "examples/yolo_world/model"
embeddings = torch.from_numpy(np.load(assets / "coco_text_outp.npy")).reshape(1, 80, 512)
config_name = "yolo_world_v2_m_vlpan_bn_noeinsum_2e-3_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.py"
output = root / "onnx/yolo_world"
output.mkdir(parents=True, exist_ok=True)
os.chdir(upstream)
sys.path[:0] = [str(upstream / "deploy"), str(upstream)]
# The detector takes text embeddings as a separate input. The official dry run
# needs their shape, so reuse Rockchip's COCO embeddings and skip fetching CLIP.
init_detector = mmdet.apis.init_detector
def init_with_local_text(config_path, checkpoint, device):
source = Path(config_path)
fallback = root / "mmyolo-source"
base = "configs/yolov8/yolov8_m_syncbn_fast_8xb16-500e_coco.py"
if (upstream / "third_party/mmyolo" / base).is_file():
config = Config.fromfile(source)
elif (fallback / base).is_file():
with NamedTemporaryFile(mode="w", suffix=".py", dir=source.parent) as copy:
copy.write(source.read_text().replace("../../third_party/mmyolo/", f"{fallback}/"))
copy.flush()
config = Config.fromfile(copy.name)
else:
raise FileNotFoundError("Missing the pinned MMYOLO configuration")
config.model.backbone.with_text_model = False
# init_detector reads dataset metadata even though export uses no images
# or labels from the evaluation dataset.
config.test_dataloader.dataset = dict(type="mmdet.CocoDataset", lazy_init=True)
model = init_detector(config, checkpoint, device=device)
model.reparameterize = lambda texts: setattr(model, "text_feats", embeddings.to(device))
return model
mmdet.apis.init_detector = init_with_local_text
script = str(upstream / "deploy/export_onnx.py")
sys.argv = [script, str(upstream / "configs/pretrain" / config_name),
str(root / "weights/yolo_world_v2m.pth"), "--custom-text",
str(upstream / "data/texts/coco_class_texts.json"), "--opset", "11",
"--model-only", "--work-dir", str(output)]
runpy.run_path(script, run_name="__main__")
|