import os os.environ['CUDA_VISIBLE_DEVICES'] = '0' import functools from typing import List import jittor as jt import numpy as np from tqdm import tqdm from PIL import Image import jclip as clip from tabulate import tabulate from template import templatesv1, templatesv2, templatesv3 from path_config import DATASET_ROOT_DIR from utils.logger import setup_logger from utils.color_print import yellow_print jt.flags.use_cuda = 1 @functools.lru_cache() def get_valid_dataset(): dataset_path = os.path.join(DATASET_ROOT_DIR, 'valid.lst') img_label_pairs = open(dataset_path, 'r').read().splitlines() image_list = [l.split(' ')[0] for l in img_label_pairs] id_list = [int(l.split(' ')[1]) for l in img_label_pairs] return image_list, id_list def split_valid_dataset(image_list, id_list): valid_dataset = { 'animal': { 'image_list': [], 'id_list': [] }, 'caltech': { 'image_list': [], 'id_list': [] }, 'dogs': { 'image_list': [], 'id_list': [] }, 'food': { 'image_list': [], 'id_list': [] } } for image, label in zip(image_list, id_list): if label < 52: valid_dataset['animal']['image_list'].append(image) valid_dataset['animal']['id_list'].append(label) elif label < 143: valid_dataset['caltech']['image_list'].append(image) valid_dataset['caltech']['id_list'].append(label) elif label < 244: valid_dataset['food']['image_list'].append(image) valid_dataset['food']['id_list'].append(label) else: valid_dataset['dogs']['image_list'].append(image) valid_dataset['dogs']['id_list'].append(label) return valid_dataset class TestSet(jt.dataset.Dataset): def __init__(self, image_list, id_list, transform, batch_size=256, num_workers=8): super(TestSet, self).__init__() self.image_list = image_list self.id_list = id_list self.batch_size = batch_size self.num_workers = num_workers self.transform = transform self.set_attrs( batch_size=self.batch_size, total_len=len(self.image_list), num_workers=self.num_workers, buffer_size=1024 * 1024 * 1024 ) def __getitem__(self, idx): image_path = self.image_list[idx] label = self.id_list[idx] image = Image.open(f'{DATASET_ROOT_DIR}/{image_path}').convert('RGB') if self.transform is not None: image = self.transform(image) img = np.asarray(image) return img, label @functools.lru_cache() def get_classnames() -> List[str]: """映射：从类别（数字）到文本""" with open(os.path.join(DATASET_ROOT_DIR, 'classes.txt'), 'r') as file: classes = file.read().splitlines() map_dict = {} id_list = [int(line.split(' ')[1]) for line in classes] classnames = [line.split(' ')[0] for line in classes] for idx, name in zip(id_list, classnames): if idx < 52: map_dict[idx] = ' '.join(name[7:].lower().split('_')) elif idx < 143: map_dict[idx] = ' '.join(name[12:].lower().split('_')) elif idx < 244: map_dict[idx] = ' '.join(name[9:].lower().split('_')) else: map_dict[idx] = ' '.join(name[8:].lower().split('_')) return list(map_dict.values()) def zeroshot_classifier(model: clip.model.CLIP, classnames: List[str] = get_classnames(), weights_version: int = 1) -> jt.Var: """ 使用 CLIP 模型进行零样本分类器的构建。 Args: - model (clip.model.CLIP): 加载的 CLIP 模型实例。 - classnames (list): 包含所有类别名称的列表。 - templates (list): 包含模板字符串的列表，用于生成每个类别的文本输入。 Returns: - torch.Tensor: 形状为 (embedding_size, num_classes) 的零样本权重张量。 注： - 此函数假设模型已经在 GPU 上，并且模型可以处理字符串的 tokenization 和文本嵌入。 - 输出的张量将包含每个类别的平均嵌入向量，并进行了归一化处理。 """ if weights_version == 1: templates = templatesv1 elif weights_version == 2: templates = templatesv2 elif weights_version == 3: templates = templatesv3 else: raise ValueError("weights_version must be 1, 2, or 3") model.eval() with jt.no_grad(): zeroshot_weights = [] for classname in tqdm(classnames, desc='Extracting class embeddings'): texts = [template.format(classname) for template in templates] #format with class texts = clip.tokenize(texts) #tokenize class_embeddings = model.encode_text(texts) #embed with text encoder class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True) class_embedding = class_embeddings.mean(dim=0) class_embedding /= class_embedding.norm() zeroshot_weights.append(class_embedding) zeroshot_weights = jt.stack(zeroshot_weights, dim=1) return zeroshot_weights def evaluate(model, dataloader, zeroshot_weights, name): model.eval() corrct = 0 total_count = 0 with jt.no_grad(): print(f"\nTesting on {name}") bar = tqdm(dataloader) for i, batch in enumerate(bar): images, targets = batch total_count += len(images) image_features = model.encode_image(images) image_features = image_features / image_features.norm(dim=1, keepdim=True) logits = (100 * image_features @ zeroshot_weights).softmax(dim=-1) preds = jt.argmax(logits, dim=1)[0] corrct += jt.equal(preds, targets).sum().item() bar.set_description(f'{name} : {corrct}/{total_count} Acc: {corrct / total_count:.4f}') bar.update(1) bar.close() return corrct def main(best_model_path, weights_version): model, transform = clip.load("/jittor-competiiton/pretrained/ViT-B-32.pkl") if best_model_path is not None: model.load_state_dict(jt.load(best_model_path)) yellow_print('Loaded weights from {}'.format(best_model_path)) else: yellow_print('weights not loaded! using pretrained weights') zeroshot_weights = zeroshot_classifier(model, weights_version=weights_version) image_list, id_list = get_valid_dataset() valid_dataset = split_valid_dataset(image_list, id_list) batch_sizes = [256, 256, 512, 512] num_works = [4, 4, 8, 8] animal = TestSet(valid_dataset['animal']['image_list'], valid_dataset['animal']['id_list'], transform=transform, batch_size=batch_sizes[0], num_workers=num_works[0]) caltech = TestSet(valid_dataset['caltech']['image_list'], valid_dataset['caltech']['id_list'], transform=transform, batch_size=batch_sizes[1], num_workers=num_works[1]) food = TestSet(valid_dataset['food']['image_list'], valid_dataset['food']['id_list'], transform=transform, batch_size=batch_sizes[2], num_workers=num_works[2]) dogs = TestSet(valid_dataset['dogs']['image_list'], valid_dataset['dogs']['id_list'], transform=transform, batch_size=batch_sizes[3], num_workers=num_works[3] ) animal_total = len(valid_dataset['animal']['image_list']) caltech_total = len(valid_dataset['caltech']['image_list']) food_total = len(valid_dataset['food']['image_list']) dogs_total = len(valid_dataset['dogs']['image_list']) total = animal_total + caltech_total + food_total + dogs_total animal_correct = evaluate(model, animal, zeroshot_weights, 'animal') caltech_correct = evaluate(model, caltech, zeroshot_weights, 'caltech') food_correct = evaluate(model, food, zeroshot_weights, 'food') dogs_correct = evaluate(model, dogs, zeroshot_weights, 'dogs') correct_total = animal_correct + caltech_correct + food_correct + dogs_correct metrics = [animal_correct/ animal_total, caltech_correct/ caltech_total, food_correct/ food_total, dogs_correct/ dogs_total, correct_total / total] print('Average Acc: ', metrics[-1]) return [round(acc, 4) for acc in metrics] if __name__ == '__main__': # 待测试的模型路径，是一个文件夹 model_dir = '/jittor-competiiton/ckptFE/07-30/version_1' logger = setup_logger(model_dir, type_='test') # 需要测试的模型文件名， 如min_loss, 20, 30 model_name = ['min_loss', 20, 50, 70, 90, 100] # 需要测试的提示词版本 # 1. basic: a photo of # 2. custom: # 3. from imagenet test_weights_version = [1, 2, 3] table_header = ['Epoch', '提示词', 'Animal', 'Caltech', 'Food', 'Dogs', 'Total'] table_data = [] promot = { 1: 'basic', 2: 'custom', 3: 'imagenet' } for epoch in model_name: if isinstance(epoch, str): model_path = os.path.join(model_dir, 'min_loss.pth') elif isinstance(epoch, int): model_path = os.path.join(model_dir, f'epoch_{epoch}.pth') print(f'Testing with {model_path}') logger.info(f'Testing with {model_path}') for weights_version in test_weights_version: metrics = main(model_path, weights_version) metrics.insert(0, promot[weights_version]) metrics.insert(0, epoch) table_data.append(metrics) print(tabulate(table_data, headers=table_header, tablefmt='fancy_grid')) logger.info(tabulate(table_data, headers=table_header, tablefmt='fancy_grid'))