generate_class_embedding.py

import argparse
import json
import os
import random
from collections import OrderedDict
from typing import List, Sequence, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    from transformers import AutoTokenizer, AutoConfig, XLMRobertaModel
except ImportError:
    AutoTokenizer = None
    HFBertModel = None


class XLMRobertaLanguageBackbone(nn.Module):

    def __init__(
        self,
        ckpt_path,
        frozen_modules: Sequence[str] = (),
        dropout: float = 0.0,
        init_cfg= None,
    ) -> None:

        super().__init__()
        if 'base' in ckpt_path:
            self.head = nn.Linear(768, 768, bias=True) # XLarge
            model_name = "./xlm-roberta-base/"
        elif 'large' in ckpt_path:
            self.head = nn.Linear(1024, 768, bias=True) # XLarge
            model_name = "./xlm-roberta-large/"

        self.frozen_modules = frozen_modules
        cfg = AutoConfig.from_pretrained(model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = XLMRobertaModel(cfg)
        self.language_dim = cfg.hidden_size
        
        
        # 加载 model 权重
        new_state_dict = OrderedDict()
        state_dict = torch.load(
            ckpt_path,
            map_location="cpu",
            weights_only=False,
        )['state_dict']
        for k, v in state_dict.items():
            if k.startswith('backbone.text_model.'):
                name = k.split("backbone.text_model.")[-1]
                new_state_dict[name] = v
        msg = self.load_state_dict(new_state_dict, strict=True)
        print(msg)

        print("TEXT-ENCODER xlm-roberta-base LOADING WEIGHTS !!!!")



    def forward(self, text: List[str], max_seq_len: int = 32):
        text = self.tokenizer(text=text, return_tensors="pt",
                              padding="max_length", max_length=max_seq_len)
        text = text.to(device=self.model.device)

        txt_feats = self.model(**text)["last_hidden_state"][:, 0]
        txt_feats = self.head(txt_feats)

        return txt_feats


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--wedetect_checkpoint', type=str, default='checkpoints/wedetect_base.pth')
    parser.add_argument('--classname_file', type=str, default='data/texts/coco_zh_class_texts.json')
    parser.add_argument('--max-seq-len', type=int, default=32,
                        help='Fixed token length (must match ONNX export).')
    parser.add_argument('--num-classes-per-group', type=int, default=4,
                        help='Number of classes per group npy.')
    parser.add_argument('--num-groups', type=int, default=64,
                        help='Number of random groups to generate.')
    parser.add_argument('--calib-dir', type=str, default='calib_data',
                        help='Directory for text-encoder quantisation calibration data.')
    args = parser.parse_args()

    with open(args.classname_file) as f:
        name_chinese = json.load(f)
    name_chinese = [name[0] for name in name_chinese]

    language_encoder = XLMRobertaLanguageBackbone(args.wedetect_checkpoint).cuda()

    # Generate random groups: each group picks 4 random classes → shape (1, 4, 768)
    total_classes = len(name_chinese)
    print(f"Total classes: {total_classes}  → Generating {args.num_groups} random groups")
    # Generate calibration data for text-encoder quantisation
    # Directories: calib_dir/input_ids/  calib_dir/attention_mask/
    calib_input_ids = os.path.join(args.calib_dir, "input_ids")
    calib_attn_mask = os.path.join(args.calib_dir, "attention_mask")
    for d in (calib_input_ids, calib_attn_mask):
        os.makedirs(d, exist_ok=True)

    tokenizer = language_encoder.tokenizer

    for g in range(args.num_groups):
        idx = random.sample(range(total_classes), args.num_classes_per_group)
        group_texts = [name_chinese[i] for i in idx]
        tokens = tokenizer(group_texts, padding="max_length",
                           max_length=args.max_seq_len, return_tensors="np")

        np.save(os.path.join(calib_input_ids, f"{g:03d}.npy"),
                tokens["input_ids"].astype(np.int64))
        np.save(os.path.join(calib_attn_mask, f"{g:03d}.npy"),
                tokens["attention_mask"].astype(np.int64))
        print(f"calib [{g:03d}] input_ids: {tokens['input_ids'].shape}  "
              f"classes: {group_texts}")

    # Compress each subdirectory to .tar.gz
    import tarfile
    for sub_name in ("input_ids", "attention_mask"):
        sub_dir = os.path.join(args.calib_dir, sub_name)
        tar_path = os.path.join(args.calib_dir, f"{sub_name}.tar.gz")
        with tarfile.open(tar_path, "w:gz") as tar:
            for fname in sorted(os.listdir(sub_dir)):
                tar.add(os.path.join(sub_dir, fname), arcname=fname)
        print(f"Compressed: {tar_path}")

    print(f"Saved calibration data to {args.calib_dir}/")

    # -------------------------------------------------------------------
    # Generate 64 random 4-class text embeddings  →  class_embedding_4cls/
    # Each file: (1, 4, 768) float32, L2-normalised (same as the ONNX
    # image encoder expects via text_features input).
    # -------------------------------------------------------------------
    embed_dir = os.path.join(args.calib_dir, "class_embedding_4cls")
    os.makedirs(embed_dir, exist_ok=True)

    print(f"\nGenerating {args.num_groups} random {args.num_classes_per_group}-class "
          f"text embeddings → {embed_dir}/")
    for g in range(args.num_groups):
        idx = random.sample(range(total_classes), args.num_classes_per_group)
        group_texts = [name_chinese[i] for i in idx]
        with torch.no_grad():
            feats = language_encoder(group_texts, max_seq_len=args.max_seq_len)
            feats = F.normalize(feats, dim=-1).unsqueeze(0)  # (1, 4, 768)
        fpath = os.path.join(embed_dir, f"{g:03d}.npy")
        np.save(fpath, feats.cpu().numpy().astype(np.float32))
        if (g + 1) % 16 == 0 or g == args.num_groups - 1:
            print(f"  [{g + 1:3d}/{args.num_groups}]  shape={feats.shape}  "
                  f"classes: {group_texts}")

    # Compress
    tar_path = os.path.join(args.calib_dir, "class_embedding_4cls.tar.gz")
    with tarfile.open(tar_path, "w:gz") as tar:
        for fname in sorted(os.listdir(embed_dir)):
            tar.add(os.path.join(embed_dir, fname), arcname=fname)
    print(f"Compressed: {tar_path}")

    print(f"Saved calibration data to {args.calib_dir}/")