diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 00000000..3f0f9b4f --- /dev/null +++ b/.gitattributes @@ -0,0 +1,16 @@ +* text=auto + +*.py text eol=lf +*.yaml text eol=lf +*.yml text eol=lf +*.md text eol=lf +*.txt text eol=lf + +*.png binary +*.jpg binary +*.jpeg binary +*.gif binary +*.pdf binary +*.pt binary +*.pth binary +*.tar binary diff --git a/.gitignore b/.gitignore index 2db7330a..648b78b3 100644 --- a/.gitignore +++ b/.gitignore @@ -106,6 +106,8 @@ celerybeat.pid # Environments .env +.env.* +!.env.example .venv env/ venv/ @@ -135,3 +137,14 @@ dmypy.json .idea/ .vscode/ results/ +results_smoke/ + +data/ +pretrain/ +tools/ + +# Model checkpoints and local experiment artifacts +*.ckpt +*.pt +*.pth +*.tar diff --git a/README.md b/README.md index 72460dbe..3d1ad4b7 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,6 @@ # [LibFewShot](https://arxiv.org/abs/2109.04898) Make few-shot learning easy. - [LibFewShot: A Comprehensive Library for Few-shot Learning](https://www.computer.org/csdl/journal/tp/5555/01/10239698/1Qck3o85t60). Wenbin Li, Ziyi Wang, Xuesong Yang, Chuanqi Dong, Pinzhuo Tian, Tiexin Qin, Jing Huo, Yinghuan Shi, Lei Wang, Yang Gao, Jiebo Luo. In TPAMI 2023.
@@ -49,6 +48,32 @@ Please refer to [install.md](https://libfewshot-en.readthedocs.io/en/latest/inst Complete tutorials can be found at [document](https://libfewshot-en.readthedocs.io/en/latest/)([中文文档](https://libfewshot-en.readthedocs.io/zh_CN/latest/index.html)). +## Command-line Usage + +All standard entry points are config-driven; no local experiment path is +embedded in the scripts. + +Train a model: + +```bash +python run_trainer.py --config ./config/proto.yaml +``` + +Evaluate a result directory (the saved `config.yaml` is loaded automatically): + +```bash +python run_test.py --result-path ./results/your-experiment +``` + +Resume training: + +```bash +python run_trainer_resume.py --result-path ./results/your-experiment +``` + +Dataset paths and other common settings can still be overridden on the command +line. For example: `--data_root ./data/miniImageNet--ravi --n_gpu 1`. + ## Reproduction We provide some validated configs in [reproduce](./reproduce/), please refer to `./reproduce//README.md` for further infomations. The meanings of the symbols are as follows: diff --git a/config/backbones/resnet10.yaml b/config/backbones/resnet10.yaml new file mode 100644 index 00000000..755c7995 --- /dev/null +++ b/config/backbones/resnet10.yaml @@ -0,0 +1,6 @@ +backbone: + name: resnet10 + kwargs: + avg_pool: True + is_flatten: True + is_feature: False diff --git a/config/boil.yaml b/config/boil.yaml index 13aad64e..5bfdf0fc 100644 --- a/config/boil.yaml +++ b/config/boil.yaml @@ -10,7 +10,7 @@ includes: way_num: 5 shot_num: 5 query_num: 15 -data_root: /data/wzy/miniImageNet--ravi +data_root: ./data/miniImageNet--ravi image_size: 84 use_memory: False augment: True diff --git a/config/classifiers/LDC.yaml b/config/classifiers/LDC.yaml new file mode 100644 index 00000000..ced9edf5 --- /dev/null +++ b/config/classifiers/LDC.yaml @@ -0,0 +1,4 @@ +classifier: + name: LDC + kwargs: + paper_mode: true diff --git a/config/classifiers/LDPNet.yaml b/config/classifiers/LDPNet.yaml new file mode 100644 index 00000000..9d92f6f8 --- /dev/null +++ b/config/classifiers/LDPNet.yaml @@ -0,0 +1,14 @@ +classifier: + name: LDPNet + kwargs: + lamba1: 1.0 + lamba2: 0.15 + momentum: 0.998 + beta: 0.5 + query_aug_times: 1 + use_lr_classifier: True + transductive: False + transductive_iter: 7 + transductive_k: 10 + lr_max_iter: 1000 + local_query_crops: 0 diff --git a/config/frn.yaml b/config/frn.yaml index e373c639..a9466fcc 100644 --- a/config/frn.yaml +++ b/config/frn.yaml @@ -23,7 +23,7 @@ augment_method: DeepBdcAugment # pre_training: True # data_root: /data/fewshot/miniImageNet--ravi/ -data_root: /data/fewshot/tiered_imagenet/ +data_root: ./data/tiered_imagenet backbone: name: resnet12 kwargs: diff --git a/config/headers/data.yaml b/config/headers/data.yaml index b59b6b05..e495d72e 100644 --- a/config/headers/data.yaml +++ b/config/headers/data.yaml @@ -1,8 +1,9 @@ -data_root: /data/fewshot/miniImageNet--ravi +data_root: ./data/miniImageNet--ravi image_size: 84 use_memory: False augment: True augment_times: 1 augment_times_query: 1 workers: 8 # number of workers for dataloader in all threads +pin_memory: True dataloader_num: 1 diff --git a/config/ldc.yaml b/config/ldc.yaml new file mode 100644 index 00000000..a9004e42 --- /dev/null +++ b/config/ldc.yaml @@ -0,0 +1,62 @@ +includes: + - headers/data.yaml + - headers/device.yaml + - headers/misc.yaml + - headers/model.yaml + - headers/optimizer.yaml + - classifiers/LDC.yaml + - backbones/Conv64F.yaml + +# LDC paper-mode example. Change these fields when switching datasets/shots. +# The author repository can also be selected with the LDC_ROOT environment variable. +ldc_root: ../LDC +data_root: ./data/EuroSAT +dataset_name: eurosat +path_prefix: "2750" + +clip_model: RN50 +device: cuda +device_ids: 0 +n_gpu: 1 +seed: 2024 +deterministic: True + +shot_num: 1 +way_num: 10 +query_num: 1 +test_way: 10 +test_shot: 1 +test_query: 1 +episode_size: 1 +train_episode: 1 +test_episode: 1 + +epoch: 50 +batch_size: 64 +eval_batch_size: 64 +workers: 0 +num_workers: 0 +pin_memory: False +image_size: 224 + +lr: 0.001 +weight_decay: 0.0001 +eps: 0.0001 +fuse_type: 2 +loss_lambda: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + +optimizer: + name: AdamW + kwargs: + lr: 0.0001 + weight_decay: 0.0001 + +lr_scheduler: + name: CosineAnnealingLR + kwargs: + T_max: 50 diff --git a/config/ldp_net.yaml b/config/ldp_net.yaml new file mode 100644 index 00000000..8331d36f --- /dev/null +++ b/config/ldp_net.yaml @@ -0,0 +1,70 @@ +includes: + - headers/data.yaml + - headers/device.yaml + - headers/misc.yaml + - headers/model.yaml + - headers/optimizer.yaml + - classifiers/LDPNet.yaml + - backbones/resnet10.yaml + +classifier: + name: LDPNet + kwargs: + lamba1: 1.0 + lamba2: 0.15 + momentum: 0.998 + beta: 0.5 + query_aug_times: 1 + use_lr_classifier: True + transductive: False + transductive_iter: 7 + transductive_k: 10 + lr_max_iter: 1000 + local_query_crops: 6 + +data_root: ./data/miniImageNet--ravi +image_size: 224 +augment: True +augment_times: 1 +augment_times_query: 1 +augment_method: S2M2Augment +ldp_local_crops: 6 +ldp_local_size: 96 +use_memory: False +workers: 4 +pin_memory: False + +device_ids: 0 +n_gpu: 1 +seed: 1111 +deterministic: True +epoch: 100 +val_per_epoch: 101 +way_num: 5 +shot_num: 5 +query_num: 15 +test_way: 5 +test_shot: 5 +test_query: 15 +episode_size: 1 +train_episode: 100 +test_episode: 600 + +pretrain_path: ./pretrain/399.tar +result_root: ./results +save_interval: 10 +save_part: + - emb_func +test_checkpoint: model_last.pth + +optimizer: + name: Adam + kwargs: + lr: 0.001 + other: ~ + +lr_scheduler: + name: StepLR + kwargs: + gamma: 1.0 + step_size: 100 diff --git a/config/matchingnet.yaml b/config/matchingnet.yaml index 08aafb09..f6b305ec 100644 --- a/config/matchingnet.yaml +++ b/config/matchingnet.yaml @@ -29,7 +29,7 @@ backbone: avg_pool: True is_flatten: True -pretrain_path: /data1/student/stx/ML/FSL/pretrained/IfslPretrain-miniImageNet-resnet12-5-1-Jun-02-2024-19-27-07/checkpoints/emb_func_best.pth +pretrain_path: ./pretrain/IfslPretrain-miniImageNet-resnet12/emb_func_best.pth classifier: name: DMatchingNet kwargs: @@ -58,4 +58,4 @@ classifier: use_x_only: False normalize_before_center: False normalize_d: False - normalize_ed: False \ No newline at end of file + normalize_ed: False diff --git a/config/metal.yaml b/config/metal.yaml index 78acf266..56e5e03c 100644 --- a/config/metal.yaml +++ b/config/metal.yaml @@ -5,7 +5,7 @@ includes: - headers/model.yaml - headers/optimizer.yaml -data_root: /data/Libfewshot/tiered_imagenet +data_root: ./data/tiered_imagenet way_num: 5 shot_num: 1 diff --git a/config/negative_margin.yaml b/config/negative_margin.yaml index b8ea35bd..83782f5b 100644 --- a/config/negative_margin.yaml +++ b/config/negative_margin.yaml @@ -1,53 +1,53 @@ -includes: - - headers/data.yaml - - headers/device.yaml - #- headers/losses.yaml - - headers/misc.yaml - - headers/model.yaml - - headers/optimizer.yaml - - classifiers/Negative_Margin.yaml - - backbones/resnet12.yaml - -backbone: -# kwargs: -# avg_pool: true -# is_feature: false -# is_flatten: true -# name: resnet18 - - name: resnet12 - kwargs: ~ - -# kwargs: -# is_feature: false -# is_flatten: true -# last_pool: true -# leaky_relu: false -# negative_slope: 0.2 -# name: Conv64F - - - - - -epoch: 200 -batch_size: 128 -image_size: 84 - - -optimizer: - name: Adam - kwargs: - lr: 3e-3 - weight_decay: 1e-4 - other: - - -warmup_params: - multiplier: 16 - epoch: 0 - -data_root: /data/yxs/tiered_imagenet -test_episode: 600 -shot_num: 1 +includes: + - headers/data.yaml + - headers/device.yaml + #- headers/losses.yaml + - headers/misc.yaml + - headers/model.yaml + - headers/optimizer.yaml + - classifiers/Negative_Margin.yaml + - backbones/resnet12.yaml + +backbone: +# kwargs: +# avg_pool: true +# is_feature: false +# is_flatten: true +# name: resnet18 + + name: resnet12 + kwargs: ~ + +# kwargs: +# is_feature: false +# is_flatten: true +# last_pool: true +# leaky_relu: false +# negative_slope: 0.2 +# name: Conv64F + + + + + +epoch: 200 +batch_size: 128 +image_size: 84 + + +optimizer: + name: Adam + kwargs: + lr: 3e-3 + weight_decay: 1e-4 + other: + + +warmup_params: + multiplier: 16 + epoch: 0 + +data_root: ./data/tiered_imagenet +test_episode: 600 +shot_num: 1 device_ids: 1 diff --git a/config/renet.yaml b/config/renet.yaml index 74bebbb2..0b6cdff6 100644 --- a/config/renet.yaml +++ b/config/renet.yaml @@ -20,7 +20,7 @@ classifier: temperature: 0.2 temperature_attn: 5.0 name: RENet -data_root: /data/fewshot/miniImageNet--ravi +data_root: ./data/miniImageNet--ravi deterministic: true device_ids: 3 episode_size: 1 diff --git a/core/data/collates/__init__.py b/core/data/collates/__init__.py index b3829c47..d026d506 100644 --- a/core/data/collates/__init__.py +++ b/core/data/collates/__init__.py @@ -1,5 +1,9 @@ # -*- coding: utf-8 -*- -from .collate_functions import GeneralCollateFunction, FewShotAugCollateFunction +from .collate_functions import ( + GeneralCollateFunction, + FewShotAugCollateFunction, + LDPNetTrainCollateFunction, +) from .contrib import get_augment_method,get_mean_std from ...utils import ModelType @@ -23,7 +27,22 @@ def get_collate_function(config, trfms, mode, model_type): model_type != ModelType.ABSTRACT ), "model_type should not be ModelType.ABSTRACT" - if mode == "train" and model_type == ModelType.FINETUNING: + if ( + mode == "train" + and config.get("classifier", {}).get("name") == "LDPNet" + and config.get("ldp_local_crops", 0) > 0 + ): + collate_function = LDPNetTrainCollateFunction( + trfms, + config["augment_times"], + config["way_num"], + config["shot_num"], + config["query_num"], + config["ldp_local_crops"], + config.get("ldp_local_size", 96), + config.get("image_size", 224), + ) + elif mode == "train" and model_type == ModelType.FINETUNING: collate_function = GeneralCollateFunction(trfms, config["augment_times"]) else: collate_function = FewShotAugCollateFunction( diff --git a/core/data/collates/collate_functions.py b/core/data/collates/collate_functions.py index fff4005d..e9ce3c3a 100644 --- a/core/data/collates/collate_functions.py +++ b/core/data/collates/collate_functions.py @@ -1,8 +1,61 @@ # -*- coding: utf-8 -*- import itertools +import random from collections import Iterable +import numpy as np import torch +from PIL import ImageEnhance, ImageFile, ImageFilter +from torchvision import transforms + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +TRANSFORM_TYPE_DICT = dict( + Brightness=ImageEnhance.Brightness, + Contrast=ImageEnhance.Contrast, + Sharpness=ImageEnhance.Sharpness, + Color=ImageEnhance.Color, +) + + +class ImageJitter(object): + def __init__(self, transform_dict): + self.transforms = [ + (TRANSFORM_TYPE_DICT[k], transform_dict[k]) for k in transform_dict + ] + + def __call__(self, img): + out = img + rand_tensor = torch.rand(len(self.transforms)) + for i, (transformer, alpha) in enumerate(self.transforms): + r = alpha * (rand_tensor[i] * 2.0 - 1.0) + 1 + out = transformer(out).enhance(r).convert("RGB") + return out + + +class PILRandomGaussianBlur(object): + def __init__(self, p=0.5, radius_min=0.1, radius_max=2.0): + self.prob = p + self.radius_min = radius_min + self.radius_max = radius_max + + def __call__(self, img): + if np.random.rand() > self.prob: + return img + + return img.filter( + ImageFilter.GaussianBlur( + radius=random.uniform(self.radius_min, self.radius_max) + ) + ) + + +def get_color_distortion(s=0.5): + color_jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s) + rnd_color_jitter = transforms.RandomApply([color_jitter], p=0.8) + rnd_gray = transforms.RandomGrayscale(p=0.2) + return transforms.Compose([rnd_color_jitter, rnd_gray]) class GeneralCollateFunction(object): @@ -179,3 +232,102 @@ def method(self, batch): def __call__(self, batch): return self.method(batch) + + +class LDPNetTrainCollateFunction(FewShotAugCollateFunction): + """Collate function for LDP-Net training. + + It returns the standard 224x224 support/query branch plus several 96x96 + augmented query crops used by the EMA branch. + """ + + def __init__( + self, + trfms, + times, + way_num, + shot_num, + query_num, + local_crops=6, + local_size=96, + image_size=224, + ): + super(LDPNetTrainCollateFunction, self).__init__( + trfms, times, 1, way_num, shot_num, query_num + ) + self.local_crops = local_crops + self.image_size = image_size + self.mean = [0.485, 0.456, 0.406] + self.std = [0.229, 0.224, 0.225] + + color_transform = [get_color_distortion(), PILRandomGaussianBlur()] + self.multicrop_trfms = [] + for size, nmb, min_scale, max_scale in zip( + [image_size, local_size], + [2, local_crops], + [0.14, 0.05], + [1.0, 0.14], + ): + random_resized_crop = transforms.RandomResizedCrop( + size, + scale=(min_scale, max_scale), + ) + crop_transform = transforms.Compose( + [ + random_resized_crop, + transforms.RandomHorizontalFlip(p=0.5), + transforms.Compose(color_transform), + transforms.ToTensor(), + transforms.Normalize(mean=self.mean, std=self.std), + ] + ) + self.multicrop_trfms.extend([crop_transform] * nmb) + + jitter_param = dict(Brightness=0.4, Contrast=0.4, Color=0.4) + self.global_trfms = transforms.Compose( + [ + transforms.Resize([image_size, image_size]), + ImageJitter(jitter_param), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize(mean=self.mean, std=self.std), + ] + ) + + def _apply_train_dataset_transforms(self, image): + multi_crops = [trfm(image) for trfm in self.multicrop_trfms] + raw_image = self.global_trfms(image) + return multi_crops[2:], raw_image + + def method(self, batch): + images, labels = zip(*batch) + images_split_by_label = [ + images[index : index + self.shot_num + self.query_num] + for index in range(0, len(images), self.shot_num + self.query_num) + ] + + anchor_images = [] + query_aug_images = [[] for _ in range(self.local_crops)] + for class_images in images_split_by_label: + for image_idx, image in enumerate(class_images): + local_crops, global_image = self._apply_train_dataset_transforms(image) + anchor_images.append(global_image) + if image_idx >= self.shot_num: + for crop_idx, crop in enumerate(local_crops): + query_aug_images[crop_idx].append(crop) + + anchor_images = torch.stack(anchor_images) + query_aug_images = torch.stack( + [torch.stack(crops) for crops in query_aug_images] + ) + + global_labels = torch.tensor(labels, dtype=torch.int64).reshape( + -1, self.way_num, self.shot_num + self.query_num + ) + global_labels = ( + global_labels[..., 0] + .unsqueeze(-1) + .repeat(1, 1, self.shot_num * self.times + self.query_num) + ) + + return anchor_images, global_labels, query_aug_images diff --git a/core/data/collates/contrib/__init__.py b/core/data/collates/contrib/__init__.py index c3916dbc..c5167279 100644 --- a/core/data/collates/contrib/__init__.py +++ b/core/data/collates/contrib/__init__.py @@ -75,7 +75,16 @@ def get_augment_method( ] else: - if config["image_size"] == 224: + if ( + config.get("classifier", {}).get("name") == "LDPNet" + and config["image_size"] == 224 + ): + resize_size = int(config["image_size"] * 1.15) + trfms_list = [ + transforms.Resize((resize_size, resize_size)), + transforms.CenterCrop((config["image_size"], config["image_size"])), + ] + elif config["image_size"] == 224: trfms_list = [ transforms.Resize((256, 256)), transforms.CenterCrop((224, 224)), @@ -147,4 +156,4 @@ def get_mean_std( elif config["augment_method"]== "S2M2Augment": MEAN= [0.485, 0.456, 0.406] STD=[0.229, 0.224, 0.225] - return MEAN,STD \ No newline at end of file + return MEAN,STD diff --git a/core/data/dataloader.py b/core/data/dataloader.py index 9e198e9f..6f6728ef 100644 --- a/core/data/dataloader.py +++ b/core/data/dataloader.py @@ -47,8 +47,6 @@ def get_dataloader(config, mode, model_type, distribute): # TODO classifier_name = config.get("classifier", {}).get("name", "") use_fgfl = "GAIN" in classifier_name.upper() - # Check if FGFL mode is enabled - print("use_fgfl:", use_fgfl) if use_fgfl: # For FGFL, we need frequency domain processing in dataset @@ -89,6 +87,7 @@ def get_dataloader(config, mode, model_type, distribute): data_scale = 1 if config["n_gpu"] == 0 else config["n_gpu"] workers = config["workers"] // data_scale + pin_memory = config.get("pin_memory", True) if workers == 0: print( "with zero workers, the training phase will be very slow", @@ -105,7 +104,7 @@ def get_dataloader(config, mode, model_type, distribute): shuffle=False if few_shot or distribute else True, num_workers=workers, # num_workers for each gpu drop_last=False if few_shot else True, - pin_memory=True, + pin_memory=pin_memory, collate_fn=collate_function, ) @@ -120,11 +119,12 @@ def get_dataloader(config, mode, model_type, distribute): mode=mode, config=config, ) + pin_memory = config.get("pin_memory", True) dataloader = DataLoader( dataset, batch_sampler=sampler, num_workers=config["n_gpu"] * 4, - pin_memory=True, + pin_memory=pin_memory, collate_fn=collate_function, ) collate_function = get_collate_function( @@ -143,7 +143,7 @@ def get_dataloader(config, mode, model_type, distribute): shuffle=True, num_workers=config["n_gpu"] * 4, drop_last=True, - pin_memory=True, + pin_memory=pin_memory, collate_fn=collate_function, ) diff --git a/core/data/dataset.py b/core/data/dataset.py index fe8e2331..9113ea4d 100644 --- a/core/data/dataset.py +++ b/core/data/dataset.py @@ -115,7 +115,7 @@ def _generate_data_list(self): data_list = [] label_list = [] class_label_dict = dict() - with open(meta_csv) as f_csv: + with open(meta_csv, encoding="utf-8-sig") as f_csv: f_train = csv.reader(f_csv, delimiter=",") for row in f_train: if f_train.line_num == 1: @@ -271,7 +271,7 @@ def _generate_csv_data_list(self, dataset_config): label_list = [] class_label_dict = dict() - with open(meta_csv) as f_csv: + with open(meta_csv, encoding="utf-8-sig") as f_csv: f_train = csv.reader(f_csv, delimiter=",") for row in f_train: if f_train.line_num == 1: diff --git a/core/data/samplers.py b/core/data/samplers.py index 7706a8cf..f8a1d536 100644 --- a/core/data/samplers.py +++ b/core/data/samplers.py @@ -10,8 +10,6 @@ def get_sampler(dataset, few_shot, distribute, mode, config): # Check if FGFL mode is enabled by classifier name classifier_name = config.get("classifier", {}).get("name", "") use_fgfl = "GAIN" in classifier_name.upper() - print(f"Sampler debug: classifier_name={classifier_name}, use_fgfl={use_fgfl}") - if distribute: # Choose distributed sampler based on FGFL mode sampler_class = FGFLDistributedCategoriesSampler if use_fgfl else DistributedCategoriesSampler diff --git a/core/model/backbone/__init__.py b/core/model/backbone/__init__.py index 68ebbae7..eb8af7ea 100644 --- a/core/model/backbone/__init__.py +++ b/core/model/backbone/__init__.py @@ -2,6 +2,7 @@ from .conv_four import Conv32F, Conv64F, R2D2Embedding from .conv_four_mcl import Conv64F_MCL from .resnet_12 import resnet12, resnet12woLSC +from .resnet_10 import resnet10 from .resnet_12_mcl import resnet12_mcl, resnet12_r2d2 from .resnet_18 import resnet18 from .wrn import WRN diff --git a/core/model/backbone/resnet_10.py b/core/model/backbone/resnet_10.py new file mode 100644 index 00000000..6d0992b0 --- /dev/null +++ b/core/model/backbone/resnet_10.py @@ -0,0 +1,140 @@ +# -*- coding: utf-8 -*- +import torch +from torch import nn + + +def conv3x3(in_planes, out_planes, stride=1, dilation=1): + return nn.Conv2d( + in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + dilation=dilation, + bias=False, + ) + + +def conv1x1(in_planes, out_planes, stride=1): + return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) + + +class SimpleBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, dilation=1): + super(SimpleBlock, self).__init__() + self.conv1 = conv3x3(inplanes, planes, stride=stride, dilation=dilation) + self.bn1 = nn.BatchNorm2d(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes) + self.bn2 = nn.BatchNorm2d(planes) + + if stride != 1 or inplanes != planes: + self.downsample = nn.Sequential( + conv1x1(inplanes, planes, stride), + nn.BatchNorm2d(planes), + ) + else: + self.downsample = None + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + return out + + +class ResNet10(nn.Module): + def __init__( + self, + out_dims=(64, 128, 256, 512), + strides=(1, 2, 2, 2), + dilations=(1, 1, 1, 1), + avg_pool=True, + is_flatten=True, + is_feature=False, + ): + super(ResNet10, self).__init__() + self.avg_pool = avg_pool + self.is_flatten = is_flatten + self.is_feature = is_feature + + self.stem = nn.Sequential( + nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + nn.MaxPool2d(kernel_size=3, stride=2, padding=1), + ) + + blocks = [] + inplanes = 64 + for planes, stride, dilation in zip(out_dims, strides, dilations): + blocks.append(SimpleBlock(inplanes, planes, stride, dilation)) + inplanes = planes + self.layers = nn.ModuleList(blocks) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + + for module in self.modules(): + if isinstance(module, nn.Conv2d): + nn.init.kaiming_normal_( + module.weight, mode="fan_out", nonlinearity="relu" + ) + elif isinstance(module, nn.BatchNorm2d): + nn.init.constant_(module.weight, 1) + nn.init.constant_(module.bias, 0) + + def forward(self, x): + features = [] + out = self.stem(x) + for layer in self.layers: + out = layer(out) + features.append(out) + + if self.avg_pool: + out = self.avgpool(out) + if self.is_flatten: + out = torch.flatten(out, 1) + if self.is_feature: + features[-1] = out + return tuple(features) + return out + + def convert_state_dict(self, state_dict): + """Convert the original LDP-Net ResNet10 checkpoint keys.""" + converted = {} + for key, value in state_dict.items(): + new_key = key + if new_key.startswith("trunk.0."): + new_key = new_key.replace("trunk.0.", "stem.0.", 1) + elif new_key.startswith("trunk.1."): + new_key = new_key.replace("trunk.1.", "stem.1.", 1) + elif new_key.startswith("trunk."): + parts = new_key.split(".", 2) + if len(parts) == 3 and parts[1].isdigit(): + block_idx = int(parts[1]) - 4 + if 0 <= block_idx < len(self.layers): + new_key = "layers.{}.{}".format(block_idx, parts[2]) + + new_key = new_key.replace(".C1.", ".conv1.") + new_key = new_key.replace(".BN1.", ".bn1.") + new_key = new_key.replace(".C2.", ".conv2.") + new_key = new_key.replace(".BN2.", ".bn2.") + new_key = new_key.replace(".shortcut.", ".downsample.0.") + new_key = new_key.replace(".BNshortcut.", ".downsample.1.") + converted[new_key] = value + return converted + + +def resnet10(**kwargs): + return ResNet10(**kwargs) diff --git a/core/model/metric/__init__.py b/core/model/metric/__init__.py index 55e0136c..8e8965e4 100644 --- a/core/model/metric/__init__.py +++ b/core/model/metric/__init__.py @@ -15,3 +15,5 @@ from .meta_baseline import MetaBaseline from .mcl import MCL from .fgfl_net import GAINModel +from .ldp_net import LDPNet +from .ldc import LDC diff --git a/core/model/metric/ldc.py b/core/model/metric/ldc.py new file mode 100644 index 00000000..22cf402b --- /dev/null +++ b/core/model/metric/ldc.py @@ -0,0 +1,723 @@ +# -*- coding: utf-8 -*- +"""Logits DeConfusion with CLIP for few-shot learning. + +This module contains the paper-style LDC implementation used by LibFewShot. +The CLIP architecture is loaded from the author's reference repository rather +than copied into LibFewShot. Set ``ldc_root`` or the ``LDC_ROOT`` environment +variable to point to that repository. +""" +import csv +import importlib +import json +import os +import random +import sys +from collections import defaultdict +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn.functional as F +import torchvision.transforms as T +import yaml +from PIL import Image +from torch.utils.data import DataLoader, Dataset +from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor + +from core.utils import accuracy +from .metric_model import MetricModel + +try: + from torchvision.transforms import InterpolationMode + + BICUBIC = InterpolationMode.BICUBIC +except ImportError: + BICUBIC = Image.BICUBIC + + +EUROSAT_CLASSNAMES = { + "AnnualCrop": "Annual Crop Land", + "Forest": "Forest", + "HerbaceousVegetation": "Herbaceous Vegetation Land", + "Highway": "Highway or Road", + "Industrial": "Industrial Buildings", + "Pasture": "Pasture Land", + "PermanentCrop": "Permanent Crop Land", + "Residential": "Residential Buildings", + "River": "River", + "SeaLake": "Sea or Lake", +} + +DATASET_TEMPLATES = { + "caltech101": ["a photo of a {}."], + "dtd": ["{} texture."], + "eurosat": ["a centered satellite photo of {}."], + "fgvc": ["a photo of a {}, a type of aircraft."], + "food101": ["a photo of {}, a type of food."], + "oxford_flowers": ["a photo of a {}, a type of flower."], + "oxford_pets": ["a photo of a {}, a type of pet."], + "stanford_cars": ["a photo of a {}."], + "stanfordcar": ["a photo of a {}."], + "sun397": ["a photo of a {}."], + "ucf101": ["a photo of a person doing {}."], +} + +DATASET_SPLIT_JSON = { + "caltech101": "split_zhou_Caltech101.json", + "dtd": "split_zhou_DescribableTextures.json", + "eurosat": "split_zhou_EuroSAT.json", + "food101": "split_zhou_Food101.json", + "oxford_flowers": "split_zhou_OxfordFlowers.json", + "oxford_pets": "split_zhou_OxfordPets.json", + "stanford_cars": "split_zhou_StanfordCars.json", + "stanfordcar": "split_zhou_StanfordCars.json", + "sun397": "split_zhou_SUN397.json", + "ucf101": "split_zhou_UCF101.json", +} + +DATASET_PATH_PREFIX = { + "caltech101": "101_ObjectCategories", + "dtd": "images", + "eurosat": "2750", + "food101": "images", + "oxford_flowers": "jpg", + "oxford_pets": "images", + "stanford_cars": "", + "stanfordcar": "", + "sun397": "SUN397", + "ucf101": "UCF-101-midframes", +} + + +@dataclass +class Datum: + impath: str + label: int + classname: str + + +class LDCImageDataset(Dataset): + def __init__(self, items, transform): + self.items = items + self.transform = transform + + def __len__(self): + return len(self.items) + + def __getitem__(self, index): + item = self.items[index] + image = Image.open(item.impath).convert("RGB") + return self.transform(image), item.label + + +class LDCConfig(object): + def __init__(self, values): + for key, value in values.items(): + setattr(self, key, value) + + def as_dict(self): + return dict(self.__dict__) + + +class LDC(MetricModel): + """LibFewShot classifier entry for LDC. + + The paper reproduction path is launched through :func:`run_ldc_from_config`. + The inherited episodic forward methods are intentionally unavailable because + the CVPR 2025 method trains CLIP adapters on fixed dataset splits rather + than LibFewShot N-way episodic batches. + """ + + def __init__(self, paper_mode=True, **kwargs): + super(LDC, self).__init__(**kwargs) + self.paper_mode = paper_mode + + def set_forward(self, batch): + raise RuntimeError( + "LDC paper reproduction uses CLIP fixed-split training. " + "Run it with `python run_ldc.py --config ./config/ldc.yaml`." + ) + + def set_forward_loss(self, batch): + raise RuntimeError( + "LDC paper reproduction uses CLIP fixed-split training. " + "Run it with `python run_ldc.py --config ./config/ldc.yaml`." + ) + + +def is_ldc_paper_config(config): + classifier = config.get("classifier", {}) + kwargs = classifier.get("kwargs") or {} + return classifier.get("name") == "LDC" and kwargs.get("paper_mode", False) + + +def load_clip_backend(ldc_root): + """Import ``clip_ldc`` from an explicit author-repository path.""" + ldc_root = os.path.abspath(os.path.expanduser(os.path.expandvars(ldc_root))) + package_dir = os.path.join(ldc_root, "clip_ldc") + if not os.path.isdir(package_dir): + raise ImportError( + "Cannot find clip_ldc under {}. Set ldc_root in the config or set " + "the LDC_ROOT environment variable.".format(ldc_root) + ) + if ldc_root not in sys.path: + sys.path.insert(0, ldc_root) + return importlib.import_module("clip_ldc") + + +def run_ldc_from_config( + config, + eval_only=False, + checkpoint=None, + check_data=False, + smoke_forward=False, +): + ldc_config = build_ldc_config(config) + setup_seed(ldc_config.seed) + os.makedirs(ldc_config.output_dir, exist_ok=True) + + with open(os.path.join(ldc_config.output_dir, "config.yaml"), "w", encoding="utf-8") as fout: + yaml.safe_dump(config, fout, allow_unicode=True) + + train_items, val_items, test_items, classnames = build_datasets(ldc_config) + ldc_config.num_classes = len(classnames) + missing = count_missing_paths(train_items + val_items + test_items) + print( + "dataset={} classes={} shots={} train/val/test={}/{}/{}".format( + ldc_config.dataset_name, + ldc_config.num_classes, + ldc_config.shots, + len(train_items), + len(val_items), + len(test_items), + ) + ) + print("missing_paths={}".format(missing)) + print("classnames={}".format(classnames[:10])) + if check_data: + if missing: + raise FileNotFoundError( + "{} paths from the selected split are missing under {}".format( + missing, ldc_config.data_root + ) + ) + return None + if missing: + raise FileNotFoundError( + "{} paths from the selected split are missing under {}".format( + missing, ldc_config.data_root + ) + ) + clip_backend = load_clip_backend(ldc_config.ldc_root) + + device = torch.device( + ldc_config.device if torch.cuda.is_available() and ldc_config.device != "cpu" else "cpu" + ) + model, _ = clip_backend.load( + ldc_config.clip_model, + download_root=ldc_config.clip_cache_dir, + num_classes=ldc_config.num_classes, + config=ldc_config, + device=device, + ) + model.eval() + + template = getattr( + ldc_config, + "template", + DATASET_TEMPLATES.get(ldc_config.dataset_name.lower(), ["a photo of a {}."]), + ) + text_cache = os.path.join( + ldc_config.output_dir, + "{}_{}_textfeats.pt".format( + ldc_config.dataset_name, ldc_config.clip_model.replace("/", "-") + ), + ) + text_feats = clip_classifier(text_cache, classnames, template, model, device, clip_backend).to( + device + ) + + train_loader = build_loader( + train_items, + transform_train(ldc_config.image_size), + ldc_config.batch_size, + ldc_config.num_workers, + shuffle=True, + ) + val_loader = build_loader( + val_items, + transform_test(ldc_config.image_size), + ldc_config.eval_batch_size, + ldc_config.num_workers, + shuffle=False, + ) + test_loader = build_loader( + test_items, + transform_test(ldc_config.image_size), + ldc_config.eval_batch_size, + ldc_config.num_workers, + shuffle=False, + ) + + if smoke_forward: + smoke_ldc_forward(model, train_loader, text_feats, device) + return None + if eval_only: + checkpoint = checkpoint or os.path.join(ldc_config.output_dir, "model_best.pth") + payload = torch.load(checkpoint, map_location=device) + model.load_state_dict(payload["model"]) + beta = payload.get("beta") + _, test_result = evaluate(model, test_loader, text_feats, device, beta=beta) + print("Test result: {}".format(test_result)) + return test_result + + return train_ldc( + ldc_config, + model, + train_loader, + val_loader, + test_loader, + text_feats, + device, + ) + + +def build_ldc_config(config): + classifier_kwargs = config.get("classifier", {}).get("kwargs") or {} + ldc_values = dict(classifier_kwargs.get("ldc", {})) + + aliases = { + "ldc_root": "ldc_root", + "clip_model": "clip_model", + "clip_cache_dir": "clip_cache_dir", + "dataset_name": "dataset_name", + "data_root": "data_root", + "split_json": "split_json", + "path_prefix": "path_prefix", + "output_dir": "output_dir", + "shots": "shots", + "seed": "seed", + "device": "device", + "image_size": "image_size", + "batch_size": "batch_size", + "eval_batch_size": "eval_batch_size", + "num_workers": "num_workers", + "lr": "lr", + "weight_decay": "weight_decay", + "eps": "eps", + "epochs": "epochs", + "loss_lambda": "loss_lambda", + "fuse_type": "fuse_type", + "template": "template", + } + for source, target in aliases.items(): + if source in config: + ldc_values[target] = config[source] + + ldc_values.setdefault("dataset_name", "eurosat") + ldc_values.setdefault("data_root", "./data/EuroSAT") + ldc_values.setdefault("clip_model", "RN50") + environment_root = os.environ.get("LDC_ROOT") + if environment_root: + ldc_values["ldc_root"] = environment_root + else: + ldc_values.setdefault("ldc_root", "../LDC") + ldc_values.setdefault("seed", config.get("seed", 2024)) + ldc_values.setdefault("shots", config.get("shot_num", 1)) + ldc_values.setdefault("epochs", config.get("epoch", 50)) + ldc_values.setdefault("batch_size", config.get("batch_size", 64)) + ldc_values.setdefault("eval_batch_size", 64) + ldc_values.setdefault("num_workers", config.get("workers", 0)) + ldc_values.setdefault("lr", 0.001) + ldc_values.setdefault("weight_decay", 0.0001) + ldc_values.setdefault("eps", 0.0001) + ldc_values.setdefault("device", "cuda") + ldc_values.setdefault("image_size", 224) + ldc_values.setdefault("fuse_type", 2) + ldc_values.setdefault("loss_lambda", [1.0, 1.0, 1.0, 1.0, 1.0]) + + dataset_key = ldc_values["dataset_name"].lower() + ldc_values["ldc_root"] = os.path.normpath( + os.path.expanduser(os.path.expandvars(ldc_values["ldc_root"])) + ) + ldc_values.setdefault("clip_cache_dir", os.path.join(ldc_values["ldc_root"], "model", "clip")) + split_filename = DATASET_SPLIT_JSON.get(dataset_key) + if split_filename: + ldc_values.setdefault( + "split_json", + os.path.join(ldc_values["ldc_root"], "datasets", "jsons", split_filename), + ) + ldc_values.setdefault( + "output_dir", + "./results/LDC-CLIP-{}-{}shot".format(ldc_values["dataset_name"], ldc_values["shots"]), + ) + ldc_values.setdefault("path_prefix", DATASET_PATH_PREFIX.get(dataset_key, "images")) + if isinstance(ldc_values.get("template"), str): + ldc_values["template"] = [ldc_values["template"]] + for key in ["data_root", "clip_cache_dir", "output_dir"]: + ldc_values[key] = os.path.normpath(ldc_values[key]) + if ldc_values.get("split_json"): + ldc_values["split_json"] = os.path.normpath(ldc_values["split_json"]) + return LDCConfig(ldc_values) + + +def setup_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + + +def convert_image_to_rgb(image): + return image.convert("RGB") + + +def transform_train(size, scale=(0.8, 1.0)): + return Compose( + [ + T.RandomResizedCrop(size=size, scale=scale, interpolation=BICUBIC), + T.RandomHorizontalFlip(p=0.5), + convert_image_to_rgb, + ToTensor(), + Normalize( + (0.48145466, 0.4578275, 0.40821073), + (0.26862954, 0.26130258, 0.27577711), + ), + ] + ) + + +def transform_test(size): + return Compose( + [ + Resize(size, interpolation=BICUBIC), + CenterCrop(size), + convert_image_to_rgb, + ToTensor(), + Normalize( + (0.48145466, 0.4578275, 0.40821073), + (0.26862954, 0.26130258, 0.27577711), + ), + ] + ) + + +def clean_csv_text(value): + return str(value).strip().lstrip("\ufeff").strip('"').strip("'") + + +def normalize_classname(label): + label = str(label) + return EUROSAT_CLASSNAMES.get(label, label.replace("_", " ")) + + +def read_csv_rows(data_root, split): + csv_path = os.path.join(data_root, "{}.csv".format(split)) + rows = [] + with open(csv_path, "r", encoding="utf-8") as fin: + reader = csv.DictReader(fin) + for row in reader: + clean_row = {clean_csv_text(key): clean_csv_text(value) for key, value in row.items()} + rows.append((clean_row["filename"], clean_row["label"])) + return rows + + +def build_label_space(data_root): + labels = set() + for split in ["train", "val", "test"]: + labels.update(label for _, label in read_csv_rows(data_root, split)) + labels = sorted(labels) + label_to_idx = {label: idx for idx, label in enumerate(labels)} + idx_to_classname = {idx: normalize_classname(label) for label, idx in label_to_idx.items()} + return label_to_idx, idx_to_classname + + +def read_csv_split(data_root, split, label_to_idx, idx_to_classname): + image_dir = os.path.join(data_root, "images") + rows = read_csv_rows(data_root, split) + return [ + Datum( + impath=os.path.join(image_dir, filename), + label=label_to_idx[label], + classname=idx_to_classname[label_to_idx[label]], + ) + for filename, label in rows + ] + + +def build_datasets(config): + if getattr(config, "split_json", None): + if not os.path.isfile(config.split_json): + raise FileNotFoundError("Dataset split file not found: {}".format(config.split_json)) + train_all, val, test, classnames = build_datasets_from_split_json( + config.data_root, + config.split_json, + path_prefix=getattr(config, "path_prefix", "images"), + ) + train = sample_fewshot(train_all, config.shots, config.seed) + return train, val, test, classnames + + label_to_idx, idx_to_classname = build_label_space(config.data_root) + train_all = read_csv_split(config.data_root, "train", label_to_idx, idx_to_classname) + val = read_csv_split(config.data_root, "val", label_to_idx, idx_to_classname) + test = read_csv_split(config.data_root, "test", label_to_idx, idx_to_classname) + train = sample_fewshot(train_all, config.shots, config.seed) + classnames = [idx_to_classname[idx] for idx in range(len(idx_to_classname))] + return train, val, test, classnames + + +def build_datasets_from_split_json(data_root, split_json, path_prefix="images"): + with open(split_json, "r", encoding="utf-8") as fin: + split = json.load(fin) + + image_dir = os.path.join(data_root, str(path_prefix)) + + def convert(items): + output = [] + for rel_path, label, classname in items: + output.append( + Datum( + impath=resolve_image_path(image_dir, rel_path, data_root), + label=int(label), + classname=classname, + ) + ) + return output + + train = convert(split["train"]) + val = convert(split["val"]) + test = convert(split["test"]) + return train, val, test, classnames_from_items(train + val + test) + + +def resolve_image_path(image_dir, rel_path, data_root=None): + candidates = [ + os.path.join(image_dir, rel_path), + os.path.join(image_dir, rel_path.replace("/", "__").replace("\\", "__")), + os.path.join(image_dir, os.path.basename(rel_path)), + ] + if data_root is not None: + candidates.extend( + [ + os.path.join(data_root, rel_path), + os.path.join(data_root, rel_path.replace("/", "__").replace("\\", "__")), + os.path.join(data_root, "images", rel_path), + os.path.join( + data_root, + "images", + rel_path.replace("/", "__").replace("\\", "__"), + ), + ] + ) + for path in candidates: + if os.path.exists(path): + return path + return candidates[0] + + +def classnames_from_items(items): + mapping = {} + for item in items: + mapping[item.label] = item.classname + return [mapping[label] for label in sorted(mapping)] + + +def sample_fewshot(items, shots, seed): + if shots < 1: + return items + rng = random.Random(seed) + by_label = defaultdict(list) + for item in items: + by_label[item.label].append(item) + sampled = [] + for label in sorted(by_label): + candidates = by_label[label] + if len(candidates) >= shots: + sampled.extend(rng.sample(candidates, shots)) + else: + sampled.extend(rng.choices(candidates, k=shots)) + return sampled + + +def build_loader(items, transform, batch_size, workers, shuffle): + return DataLoader( + LDCImageDataset(items, transform), + batch_size=batch_size, + shuffle=shuffle, + drop_last=False, + num_workers=workers, + pin_memory=torch.cuda.is_available(), + ) + + +def clip_classifier(cache_path, classnames, template, clip_model, device, clip_backend): + if os.path.exists(cache_path): + return torch.load(cache_path, map_location=device) + + os.makedirs(os.path.dirname(cache_path), exist_ok=True) + with torch.no_grad(): + weights = [] + for classname in classnames: + texts = [prompt.format(classname.replace("_", " ")) for prompt in template] + texts = clip_backend.tokenize(texts).to(device) + class_embeddings = clip_model.encode_text(texts) + class_embeddings = class_embeddings / class_embeddings.norm(dim=-1, keepdim=True) + class_embedding = class_embeddings.mean(dim=0) + class_embedding = class_embedding / class_embedding.norm() + weights.append(class_embedding) + weights = torch.stack(weights, dim=1) + torch.save(weights.cpu(), cache_path) + return weights + + +def compute_loss(labels, clip_logits, mlp_logits, ada_logits, total_logits, lambdas): + ce_mlp = F.cross_entropy(mlp_logits, labels) * lambdas[0] + ce_ada = F.cross_entropy(ada_logits, labels) * lambdas[1] + ce_total = F.cross_entropy(total_logits, labels) * lambdas[2] + l1_mlp = F.l1_loss(mlp_logits, clip_logits) * lambdas[3] + l1_ada = F.l1_loss(ada_logits, clip_logits) * lambdas[4] + return l1_mlp + l1_ada + ce_mlp + ce_ada + ce_total + + +def fuse_logits(mlp_logits, ada_logits, beta): + return beta * mlp_logits + (1.0 - beta) * ada_logits + + +def search_beta(mlp_logits, ada_logits, labels, steps=50): + best_beta, best_acc = 0.0, -1.0 + for idx in range(steps + 1): + beta = idx / steps + acc = accuracy(fuse_logits(mlp_logits, ada_logits, beta), labels) + if acc > best_acc: + best_beta, best_acc = beta, acc + return best_beta, best_acc + + +@torch.no_grad() +def evaluate(model, loader, text_feats, device, beta=None): + model.eval() + all_labels, all_mlp, all_ada, all_total, all_clip = [], [], [], [], [] + for images, labels in loader: + images = images.to(device) + labels = labels.to(device) + clip_logits, mlp_logits, ada_logits, total_logits = model.my_forward(images, text_feats) + all_labels.append(labels) + all_clip.append(clip_logits) + all_mlp.append(mlp_logits) + all_ada.append(ada_logits) + all_total.append(total_logits) + + labels = torch.cat(all_labels) + clip_logits = torch.cat(all_clip) + mlp_logits = torch.cat(all_mlp) + ada_logits = torch.cat(all_ada) + total_logits = torch.cat(all_total) + + result = { + "clip_logits": accuracy(clip_logits, labels), + "mlp_logits": accuracy(mlp_logits, labels), + "ada_logits": accuracy(ada_logits, labels), + "tot_logits": accuracy(total_logits, labels), + } + if beta is None: + beta, result["acc"] = search_beta(mlp_logits, ada_logits, labels) + else: + result["acc"] = accuracy(fuse_logits(mlp_logits, ada_logits, beta), labels) + return beta, result + + +def train_ldc(config, model, train_loader, val_loader, test_loader, text_feats, device): + for param in model.parameters(): + param.requires_grad = False + for name, param in model.named_parameters(): + if "adapter" in name: + param.requires_grad = True + + optimizer = torch.optim.AdamW( + model.parameters(), + lr=config.lr / 10, + weight_decay=config.weight_decay, + eps=config.eps, + ) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, + max(1, config.epochs * len(train_loader)), + ) + + best_val_acc = -1.0 + best_payload = None + for epoch in range(config.epochs): + model.visual.adapter.train() + model.adapter.train() + loss_sum, acc_sum, count = 0.0, 0.0, 0 + for images, labels in train_loader: + images = images.to(device) + labels = labels.to(device) + clip_logits, mlp_logits, ada_logits, total_logits = model.my_forward(images, text_feats) + loss = compute_loss( + labels, + clip_logits, + mlp_logits, + ada_logits, + total_logits, + config.loss_lambda, + ) + optimizer.zero_grad() + loss.backward() + optimizer.step() + scheduler.step() + + loss_sum += loss.item() + acc_sum += accuracy(mlp_logits, labels) + count += 1 + + beta, val_result = evaluate(model, val_loader, text_feats, device) + _, test_result = evaluate(model, test_loader, text_feats, device, beta=beta) + print( + "Epoch {:03d}: train_loss={:.4f} train_acc={:.2f} " + "val_acc={:.2f} test_acc={:.2f} beta={:.2f}".format( + epoch, + loss_sum / max(count, 1), + acc_sum / max(count, 1), + val_result["acc"], + test_result["acc"], + beta, + ) + ) + if val_result["acc"] > best_val_acc: + best_val_acc = val_result["acc"] + best_payload = { + "model": model.state_dict(), + "epoch": epoch, + "beta": beta, + "val": val_result, + "test": test_result, + } + torch.save(best_payload, os.path.join(config.output_dir, "model_best.pth")) + + torch.save({"model": model.state_dict()}, os.path.join(config.output_dir, "model_last.pth")) + print("Best payload: {}".format(best_payload)) + return best_payload + + +@torch.no_grad() +def smoke_ldc_forward(model, loader, text_feats, device): + model.eval() + images, labels = next(iter(loader)) + images = images.to(device) + clip_logits, mlp_logits, ada_logits, total_logits = model.my_forward(images, text_feats) + print( + "smoke_forward batch={} labels={} logits={}/{}/{}/{}".format( + images.size(0), + tuple(labels.shape), + tuple(clip_logits.shape), + tuple(mlp_logits.shape), + tuple(ada_logits.shape), + tuple(total_logits.shape), + ) + ) + + +def count_missing_paths(items): + return sum(not os.path.exists(item.impath) for item in items) diff --git a/core/model/metric/ldp_net.py b/core/model/metric/ldp_net.py new file mode 100644 index 00000000..e6372957 --- /dev/null +++ b/core/model/metric/ldp_net.py @@ -0,0 +1,330 @@ +# -*- coding: utf-8 -*- +""" +LDP-Net: Revisiting Prototypical Network for Cross Domain Few-Shot Learning. + +This implementation follows LibFewShot's episodic metric-model interface while +adapting the original LDP-Net training and testing ideas. +""" +import copy + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from core.utils import accuracy +from .metric_model import MetricModel + + +class LDPNet(MetricModel): + def __init__( + self, + lamba1=1.0, + lamba2=0.15, + momentum=0.998, + beta=0.5, + query_aug_times=1, + use_lr_classifier=True, + transductive=False, + transductive_iter=7, + transductive_k=10, + lr_max_iter=1000, + local_query_crops=0, + eps=1e-6, + **kwargs + ): + super(LDPNet, self).__init__(**kwargs) + self.lamba1 = lamba1 + self.lamba2 = lamba2 + self.momentum = momentum + self.beta = beta + self.query_aug_times = max(int(query_aug_times), 1) + self.use_lr_classifier = use_lr_classifier + self.transductive = transductive + self.transductive_iter = transductive_iter + self.transductive_k = transductive_k + self.lr_max_iter = lr_max_iter + self.local_query_crops = local_query_crops + self.eps = eps + + self.loss_func = nn.CrossEntropyLoss() + self.ema_emb_func = copy.deepcopy(self.emb_func) + for param in self.ema_emb_func.parameters(): + param.requires_grad = False + + if self.query_aug_times > 1 and self.query_num % self.query_aug_times != 0: + raise ValueError( + "query_num must be divisible by query_aug_times in LDPNet. " + "Set classifier.kwargs.query_aug_times to match augment_times_query." + ) + + def train(self, mode=True): + super(LDPNet, self).train(mode) + self.ema_emb_func.train(mode) + return self + + @torch.no_grad() + def _momentum_update_ema(self): + for param, ema_param in zip( + self.emb_func.parameters(), self.ema_emb_func.parameters() + ): + ema_param.data.mul_(self.momentum).add_( + param.data, alpha=1.0 - self.momentum + ) + + def _flatten_feat(self, feat): + if feat.dim() > 2: + feat = F.adaptive_avg_pool2d(feat, 1) + feat = torch.flatten(feat, 1) + return feat + + def _power_transform(self, feat): + if self.beta is None: + return feat + return torch.pow(feat.clamp_min(0.0), self.beta) + + def _proto_logits(self, support_feat, query_feat): + episode_size, _, feat_dim = support_feat.size() + support_feat = support_feat.view( + episode_size, self.way_num, self.shot_num, feat_dim + ) + proto_feat = torch.mean(support_feat, dim=2) + return -torch.sum( + torch.pow(query_feat.unsqueeze(2) - proto_feat.unsqueeze(1), 2), + dim=3, + ) + + def _anchor_query(self, query_feat): + if self.query_aug_times <= 1: + return query_feat + episode_size, _, feat_dim = query_feat.size() + query_per_aug = self.query_num // self.query_aug_times + query_feat = query_feat.view( + episode_size, self.way_num, self.query_aug_times, query_per_aug, feat_dim + ) + query_feat = query_feat.permute(0, 2, 1, 3, 4).contiguous() + return query_feat[:, 0].view(episode_size, self.way_num * query_per_aug, feat_dim) + + def _aug_query(self, query_feat): + if self.query_aug_times <= 1: + return None + episode_size, _, feat_dim = query_feat.size() + query_per_aug = self.query_num // self.query_aug_times + query_feat = query_feat.view( + episode_size, self.way_num, self.query_aug_times, query_per_aug, feat_dim + ) + query_feat = query_feat.permute(0, 2, 1, 3, 4).contiguous() + return query_feat[:, 1:].view( + episode_size, + (self.query_aug_times - 1) * self.way_num * query_per_aug, + feat_dim, + ) + + def _anchor_target(self, query_target): + if self.query_aug_times <= 1: + return query_target + episode_size = query_target.size(0) + query_per_aug = self.query_num // self.query_aug_times + query_target = query_target.view( + episode_size, self.way_num, self.query_aug_times, query_per_aug + ) + query_target = query_target.permute(0, 2, 1, 3).contiguous() + return query_target[:, 0].view(episode_size, self.way_num * query_per_aug) + + def _consistency_loss(self, anchor_output, aug_output, aug_num=None): + if aug_output is None: + return anchor_output.new_tensor(0.0), anchor_output.new_tensor(0.0) + + episode_size, anchor_num, way_num = anchor_output.size() + aug_num = self.query_aug_times - 1 if aug_num is None else aug_num + query_per_aug = anchor_num // self.way_num + + anchor_prob = F.softmax(anchor_output, dim=-1) + aug_prob = F.softmax(aug_output, dim=-1) + + anchor_repeat = anchor_prob.unsqueeze(1).repeat(1, aug_num, 1, 1) + anchor_repeat = anchor_repeat.view(episode_size, aug_num * anchor_num, way_num) + self_image_loss = -torch.mean( + torch.sum(anchor_repeat * torch.log(aug_prob.clamp_min(self.eps)), dim=-1) + ) + + anchor_by_class = anchor_prob.view( + episode_size, self.way_num, query_per_aug, way_num + ) + aug_by_class = aug_prob.view( + episode_size, aug_num, self.way_num, query_per_aug, way_num + ) + # Match the author's implementation: the two query indices are sampled + # from NumPy's RNG. Using the CUDA RNG here changes the fixed-seed + # training trajectory even when every other paper setting is identical. + global_idx = [ + np.random.permutation(query_per_aug)[0] for _ in range(episode_size) + ] + local_idx = [ + np.random.permutation(query_per_aug)[0] for _ in range(episode_size) + ] + + global_prob = [] + local_prob = [] + for epi in range(episode_size): + global_prob.append(anchor_by_class[epi, :, global_idx[epi], :]) + local_prob.append(aug_by_class[epi, :, :, local_idx[epi], :]) + global_prob = torch.stack(global_prob, dim=0) + local_prob = torch.stack(local_prob, dim=0) + global_prob = global_prob.unsqueeze(1).expand(-1, aug_num, -1, -1) + cross_image_loss = -torch.mean( + torch.sum(global_prob * torch.log(local_prob.clamp_min(self.eps)), dim=-1) + ) + return self_image_loss, cross_image_loss + + def set_forward_loss(self, batch): + images, global_targets = batch[:2] + explicit_query_aug_images = batch[2] if len(batch) > 2 else None + images = images.to(self.device) + + feat = self._flatten_feat(self.emb_func(images)) + support_feat, query_feat, support_target, query_target = self.split_by_episode( + feat, mode=1 + ) + anchor_query_feat = self._anchor_query(query_feat) + anchor_query_target = self._anchor_target(query_target) + + anchor_output = self._proto_logits(support_feat, anchor_query_feat) + loss = self.loss_func( + anchor_output.reshape(-1, self.way_num), anchor_query_target.reshape(-1) + ) + + if explicit_query_aug_images is not None: + explicit_query_aug_images = explicit_query_aug_images.to(self.device) + local_crops, query_count, c, h, w = explicit_query_aug_images.size() + with torch.no_grad(): + aug_feat = self._flatten_feat( + self.ema_emb_func( + explicit_query_aug_images.view(local_crops * query_count, c, h, w) + ) + ) + aug_feat = aug_feat.view( + local_crops, + -1, + self.way_num * (self.query_num // self.query_aug_times), + aug_feat.size(-1), + ) + aug_feat = aug_feat.permute(1, 0, 2, 3).contiguous() + aug_query_feat = aug_feat.view( + aug_feat.size(0), local_crops * aug_feat.size(2), aug_feat.size(3) + ) + aug_output = self._proto_logits(support_feat, aug_query_feat) + self_image_loss, cross_image_loss = self._consistency_loss( + anchor_output, aug_output, aug_num=local_crops + ) + loss = loss + self.lamba1 * self_image_loss + self.lamba2 * cross_image_loss + elif self.query_aug_times > 1: + with torch.no_grad(): + ema_feat = self._flatten_feat(self.ema_emb_func(images)) + _, ema_query_feat, _, _ = self.split_by_episode(ema_feat, mode=1) + aug_query_feat = self._aug_query(ema_query_feat) + aug_output = self._proto_logits(support_feat, aug_query_feat) + self_image_loss, cross_image_loss = self._consistency_loss( + anchor_output, aug_output + ) + loss = loss + self.lamba1 * self_image_loss + self.lamba2 * cross_image_loss + + output = anchor_output.reshape(-1, self.way_num) + acc = accuracy(output, anchor_query_target.reshape(-1)) + return output, acc, loss + + def set_forward(self, batch): + images, global_targets = batch + images = images.to(self.device) + feat = self._power_transform(self._flatten_feat(self.emb_func(images))) + support_feat, query_feat, support_target, query_target = self.split_by_episode( + feat, mode=1 + ) + + query_feat = self._anchor_query(query_feat) + query_target = self._anchor_target(query_target) + + if self.use_lr_classifier: + output = self._logistic_regression_forward( + support_feat, support_target, query_feat + ) + else: + output = self._proto_logits(support_feat, query_feat).reshape( + -1, self.way_num + ) + acc = accuracy(output, query_target.reshape(-1)) + return output, acc + + def _logistic_regression_forward(self, support_feat, support_target, query_feat): + from sklearn.linear_model import LogisticRegression + + output_list = [] + episode_size = support_feat.size(0) + for epi in range(episode_size): + support_np = support_feat[epi].detach().cpu().numpy() + query_np = query_feat[epi].detach().cpu().numpy() + target_np = support_target[epi].detach().cpu().numpy() + + classifier = LogisticRegression(max_iter=self.lr_max_iter) + classifier.fit(support_np, target_np) + prob = classifier.predict_proba(query_np) + + if self.transductive: + prob = self._transductive_refine( + classifier, + support_feat[epi], + support_target[epi], + query_feat[epi], + prob, + ) + + output_list.append(torch.from_numpy(prob).to(self.device).float()) + return torch.cat(output_list, dim=0) + + def _transductive_refine( + self, classifier, support_feat, support_target, query_feat, prob + ): + from sklearn.linear_model import LogisticRegression + + feat_dim = support_feat.size(-1) + support_by_class = support_feat.view(self.way_num, self.shot_num, feat_dim) + y_proto = np.arange(self.way_num) + + for _ in range(self.transductive_iter): + # The reference Tr() applies softmax to sklearn's predict_proba + # output and keeps its float64 precision for confidence ranking. + prob_tensor = torch.from_numpy(prob).to(self.device) + prob_tensor = F.softmax(prob_tensor, dim=1) + score, pred = prob_tensor.max(dim=1) + refined_feat = [] + refined_target = [] + for cls_idx in range(self.way_num): + selected = torch.nonzero(pred == cls_idx, as_tuple=False).view(-1) + if selected.numel() > self.transductive_k: + selected_score = score[selected] + _, top_idx = torch.topk(selected_score, self.transductive_k) + selected = selected[top_idx] + + if self.shot_num == 1: + class_feat = support_by_class[cls_idx] + if selected.numel() > 0: + class_feat = torch.cat([class_feat, query_feat[selected]], dim=0) + refined_feat.append(class_feat.mean(dim=0, keepdim=True)) + else: + class_feat = support_by_class[cls_idx] + if selected.numel() > 0: + class_feat = torch.cat([class_feat, query_feat[selected]], dim=0) + refined_feat.append(class_feat) + refined_target.extend([cls_idx] * class_feat.size(0)) + + if self.shot_num == 1: + train_x = torch.cat(refined_feat, dim=0).detach().cpu().numpy() + train_y = y_proto + else: + train_x = torch.cat(refined_feat, dim=0).detach().cpu().numpy() + train_y = np.asarray(refined_target) + + classifier = LogisticRegression(max_iter=self.lr_max_iter) + classifier.fit(train_x, train_y) + prob = classifier.predict_proba(query_feat.detach().cpu().numpy()) + return prob diff --git a/core/test.py b/core/test.py index 0da3167a..37cde941 100644 --- a/core/test.py +++ b/core/test.py @@ -193,7 +193,8 @@ def _init_files(self, config): rank=self.rank, ) - state_dict_path = os.path.join(result_path, "checkpoints", "model_best.pth") + checkpoint_name = config.get("test_checkpoint", "model_best.pth") + state_dict_path = os.path.join(result_path, "checkpoints", checkpoint_name) if self.rank == 0: create_dirs([result_path, log_path, viz_path]) @@ -204,15 +205,14 @@ def _init_logger(self): self.logger = getLogger(__name__) # Hack print - def use_logger(msg, level="info"): + def use_logger(*msg, level="info"): if self.rank != 0: return - if level == "info": - self.logger.info(msg) - elif level == "warning": - self.logger.warning(msg) - else: - raise ("Not implemente {} level log".format(level)) + try: + text = " ".join(map(str, msg)) + getattr(self.logger, level)(text) + except AttributeError: + raise ValueError("Not implemente {} level log".format(level)) builtins.print = use_logger @@ -280,10 +280,10 @@ def _init_model(self, config): model_kwargs = { "way_num": config["way_num"], "shot_num": config["shot_num"] * config["augment_times"], - "query_num": config["query_num"], + "query_num": config["query_num"] * config["augment_times_query"], "test_way": config["test_way"], "test_shot": config["test_shot"] * config["augment_times"], - "test_query": config["test_query"], + "test_query": config["test_query"] * config["augment_times_query"], "emb_func": emb_func, "device": self.device, } @@ -292,7 +292,11 @@ def _init_model(self, config): print(model) print("Trainable params in the model: {}.".format(count_parameters(model))) print("Loading the state dict from {}.".format(self.state_dict_path)) - state_dict = torch.load(self.state_dict_path, map_location="cpu") + state_dict = torch.load( + self.state_dict_path, map_location="cpu", weights_only=False + ) + if isinstance(state_dict, dict) and "model" in state_dict: + state_dict = state_dict["model"] model.load_state_dict(state_dict) if self.distribute: diff --git a/core/trainer.py b/core/trainer.py index 612894f8..96fcacd8 100644 --- a/core/trainer.py +++ b/core/trainer.py @@ -194,6 +194,9 @@ def _train(self, epoch_idx): # if (param.grad != param.grad).float().sum() != 0: # nan detected # param.grad.zero_() self.optimizer.step() + model_for_ema = self.model.module if self.distribute else self.model + if hasattr(model_for_ema, "_momentum_update_ema"): + model_for_ema._momentum_update_ema() meter.update("calc_time", time() - calc_begin) # measure accuracy and record loss @@ -442,10 +445,10 @@ def _init_model(self, config): model_kwargs = { "way_num": config["way_num"], "shot_num": config["shot_num"] * config["augment_times"], - "query_num": config["query_num"], + "query_num": config["query_num"] * config["augment_times_query"], "test_way": config["test_way"], "test_shot": config["test_shot"] * config["augment_times"], - "test_query": config["test_query"], + "test_query": config["test_query"] * config["augment_times_query"], "emb_func": emb_func, "device": self.device, } @@ -460,7 +463,13 @@ def _init_model(self, config): "load pretraining emb_func from {}".format(self.config["pretrain_path"]) ) state_dict = torch.load(self.config["pretrain_path"], map_location="cpu") + if isinstance(state_dict, dict) and "state" in state_dict: + state_dict = state_dict["state"] + if hasattr(model.emb_func, "convert_state_dict"): + state_dict = model.emb_func.convert_state_dict(state_dict) msg = model.emb_func.load_state_dict(state_dict, strict=False) + if hasattr(model, "ema_emb_func"): + model.ema_emb_func.load_state_dict(model.emb_func.state_dict()) if len(msg.missing_keys) != 0: print("Missing keys:{}".format(msg.missing_keys), level="warning") @@ -496,7 +505,9 @@ def _init_model(self, config): self.config["resume_path"], "checkpoints", "model_last.pth" ) print("load the resume model checkpoints dict from {}.".format(resume_path)) - state_dict = torch.load(resume_path, map_location="cpu")["model"] + state_dict = torch.load( + resume_path, map_location="cpu", weights_only=False + )["model"] msg = model.load_state_dict(state_dict, strict=False) if len(msg.missing_keys) != 0: @@ -659,7 +670,9 @@ def _init_optim(self, config): resume_path ) ) - all_state_dict = torch.load(resume_path, map_location="cpu") + all_state_dict = torch.load( + resume_path, map_location="cpu", weights_only=False + ) state_dict = all_state_dict["optimizer"] optimizer.load_state_dict(state_dict) state_dict = all_state_dict["lr_scheduler"] diff --git a/core/utils/utils.py b/core/utils/utils.py index 2bf88777..69e2c9a4 100644 --- a/core/utils/utils.py +++ b/core/utils/utils.py @@ -49,26 +49,27 @@ def __init__(self, name, keys, writer=None): self.reset() def reset(self): - for col in self._data.columns: - self._data[col].values[:] = 0 + self._data.loc[:, :] = 0 def update(self, key, value, n=1): if self.writer is not None: tag = "{}/{}".format(self.name, key) self.writer.add_scalar(tag, value) - self._data.last_value[key] = value - self._data.total[key] += value * n - self._data.counts[key] += n - self._data.average[key] = self._data.total[key] / self._data.counts[key] + self._data.loc[key, "last_value"] = value + self._data.loc[key, "total"] += value * n + self._data.loc[key, "counts"] += n + self._data.loc[key, "average"] = ( + self._data.loc[key, "total"] / self._data.loc[key, "counts"] + ) def avg(self, key): - return self._data.average[key] + return self._data.loc[key, "average"] def result(self): - return dict(self._data.average) + return dict(self._data["average"]) def last(self, key): - return self._data.last_value[key] + return self._data.loc[key, "last_value"] def get_local_time(): diff --git a/reproduce/LDC_CLIP/README.md b/reproduce/LDC_CLIP/README.md new file mode 100644 index 00000000..34e06a11 --- /dev/null +++ b/reproduce/LDC_CLIP/README.md @@ -0,0 +1,66 @@ +# LDC reproduction + +This integration reproduces **Logits DeConfusion with CLIP for Few-Shot +Learning (CVPR 2025)** inside the LibFewShot workspace. LDC uses fixed dataset +splits and CLIP adapters rather than LibFewShot's usual episodic training loop, +so it has a dedicated `run_ldc.py` entry point. + +## External dependency + +The original `clip_ldc` package is not copied into this repository. Clone the +authors' repository next to LibFewShot: + +```bash +git clone https://github.com/LiShuo1001/LDC.git ../LDC +``` + +If it is stored elsewhere, either edit `ldc_root` in `config/ldc.yaml` or set +the `LDC_ROOT` environment variable. Dataset paths, downloaded CLIP weights, +and generated checkpoints remain local and are ignored by Git. + +## Configuration + +The example [`config/ldc.yaml`](../../config/ldc.yaml) targets EuroSAT 1-shot. +Set at least these values for a different experiment: + +```yaml +ldc_root: ../LDC +data_root: ./data/EuroSAT +dataset_name: eurosat +path_prefix: "2750" +shot_num: 1 +``` + +By default, the CoOp-style split JSON and CLIP cache are resolved below +`ldc_root`. You may override `split_json`, `clip_cache_dir`, or `output_dir` +explicitly in the YAML file. + +## Commands + +Run all commands from the LibFewShot repository root. + +Check dataset and split paths without loading CLIP: + +```bash +python run_ldc.py --config ./config/ldc.yaml --check-data +``` + +Run a single forward-pass smoke test: + +```bash +python run_ldc.py --config ./config/ldc.yaml --smoke-forward +``` + +Train: + +```bash +python run_ldc.py --config ./config/ldc.yaml +``` + +Evaluate the best checkpoint from the configured output directory: + +```bash +python run_ldc.py --config ./config/ldc.yaml --eval-only +``` + +To evaluate another checkpoint, add `--checkpoint PATH`. diff --git a/run_ldc.py b/run_ldc.py new file mode 100644 index 00000000..8f1e8c25 --- /dev/null +++ b/run_ldc.py @@ -0,0 +1,66 @@ +# -*- coding: utf-8 -*- +"""Train, validate, or test the paper-style LDC integration.""" + +import argparse +import sys + +sys.dont_write_bytecode = True + +from core.config import Config +from core.model.metric.ldc import is_ldc_paper_config, run_ldc_from_config + + +def parse_args(): + parser = argparse.ArgumentParser(description="Run the LDC paper reproduction.") + parser.add_argument( + "-c", + "--config", + default="./config/ldc.yaml", + help="Path to the LDC YAML config.", + ) + mode = parser.add_mutually_exclusive_group() + mode.add_argument( + "--eval-only", + "--eval_only", + dest="eval_only", + action="store_true", + help="Evaluate a trained checkpoint instead of training.", + ) + parser.add_argument( + "--checkpoint", + default=None, + help="Checkpoint path for --eval-only (defaults to OUTPUT_DIR/model_best.pth).", + ) + mode.add_argument( + "--check-data", + "--check_data", + dest="check_data", + action="store_true", + help="Validate dataset paths and exit.", + ) + mode.add_argument( + "--smoke-forward", + "--smoke_forward", + dest="smoke_forward", + action="store_true", + help="Run one forward pass and exit.", + ) + args, config_args = parser.parse_known_args() + if args.checkpoint and not args.eval_only: + parser.error("--checkpoint can only be used with --eval-only") + sys.argv = [sys.argv[0]] + config_args + return args + + +if __name__ == "__main__": + cli_args = parse_args() + ldc_config = Config(cli_args.config).get_config_dict() + if not is_ldc_paper_config(ldc_config): + raise ValueError("The selected config must use classifier.name=LDC and paper_mode=true.") + run_ldc_from_config( + ldc_config, + eval_only=cli_args.eval_only, + checkpoint=cli_args.checkpoint, + check_data=cli_args.check_data, + smoke_forward=cli_args.smoke_forward, + ) diff --git a/run_test.py b/run_test.py index 958c87f2..43860ddb 100644 --- a/run_test.py +++ b/run_test.py @@ -1,34 +1,95 @@ # -*- coding: utf-8 -*- +"""Evaluate a trained LibFewShot experiment.""" + +import argparse +import os import sys sys.dont_write_bytecode = True -import os import torch -from core.config import Config + from core import Test +from core.config import Config + + +def parse_args(): + parser = argparse.ArgumentParser(description="Evaluate a LibFewShot result directory.") + parser.add_argument( + "-r", + "--result-path", + "--result_path", + dest="result_path", + required=True, + help="Training result directory containing config.yaml and checkpoints/.", + ) + parser.add_argument( + "-c", + "--config", + default=None, + help="Optional config file. Defaults to RESULT_PATH/config.yaml.", + ) + parser.add_argument( + "--checkpoint", + default=None, + help="Checkpoint filename inside RESULT_PATH/checkpoints/.", + ) + parser.add_argument("--data-root", "--data_root", dest="data_root") + parser.add_argument("--test-way", "--test_way", dest="test_way", type=int) + parser.add_argument("--test-shot", "--test_shot", dest="test_shot", type=int) + parser.add_argument("--test-query", "--test_query", dest="test_query", type=int) + parser.add_argument("--test-episode", "--test_episode", dest="test_episode", type=int) + parser.add_argument("--test-epoch", "--test_epoch", dest="test_epoch", type=int) + parser.add_argument("--episode-size", "--episode_size", dest="episode_size", type=int) + parser.add_argument("--workers", type=int) + parser.add_argument("--device-ids", "--device_ids", dest="device_ids") + parser.add_argument("--n-gpu", "--n_gpu", dest="n_gpu", type=int) + parser.add_argument("--seed", type=int) + + args, config_args = parser.parse_known_args() + # Config owns the framework-wide command-line options that remain. + sys.argv = [sys.argv[0]] + config_args + return args + +def build_config(args): + config_path = args.config or os.path.join(args.result_path, "config.yaml") + if not os.path.isfile(config_path): + raise FileNotFoundError("Config file not found: {}".format(config_path)) -PATH = "./results/DN4-miniImageNet--ravi-Conv64F-5-1-Dec-01-2021-06-05-20" -VAR_DICT = { - "test_epoch": 5, - "device_ids": "4,5", - "n_gpu": 2, - "test_episode": 600, - "episode_size": 2, -} + overrides = { + "data_root": args.data_root, + "test_way": args.test_way, + "test_shot": args.test_shot, + "test_query": args.test_query, + "test_episode": args.test_episode, + "test_epoch": args.test_epoch, + "episode_size": args.episode_size, + "workers": args.workers, + "device_ids": args.device_ids, + "n_gpu": args.n_gpu, + "seed": args.seed, + "test_checkpoint": args.checkpoint, + } + overrides = {key: value for key, value in overrides.items() if value is not None} + return Config(config_path, overrides).get_config_dict() -def main(rank, config): - test = Test(rank, config, PATH) - test.test_loop() +def main(rank, config, result_path): + tester = Test(rank, config, result_path) + tester.test_loop() if __name__ == "__main__": - config = Config(os.path.join(PATH, "config.yaml"), VAR_DICT).get_config_dict() + cli_args = parse_args() + test_config = build_config(cli_args) - if config["n_gpu"] > 1: - os.environ["CUDA_VISIBLE_DEVICES"] = config["device_ids"] - torch.multiprocessing.spawn(main, nprocs=config["n_gpu"], args=(config,)) + if test_config["n_gpu"] > 1: + os.environ["CUDA_VISIBLE_DEVICES"] = test_config["device_ids"] + torch.multiprocessing.spawn( + main, + nprocs=test_config["n_gpu"], + args=(test_config, cli_args.result_path), + ) else: - main(0, config) + main(0, test_config, cli_args.result_path) diff --git a/run_trainer.py b/run_trainer.py index 70d56b27..03eac8c0 100644 --- a/run_trainer.py +++ b/run_trainer.py @@ -1,12 +1,30 @@ # -*- coding: utf-8 -*- +"""Train a LibFewShot model from a YAML config.""" + +import argparse +import os import sys sys.dont_write_bytecode = True import torch -import os -from core.config import Config + from core import Trainer +from core.config import Config + + +def parse_args(): + parser = argparse.ArgumentParser(description="Train a LibFewShot model.") + parser.add_argument( + "-c", + "--config", + required=True, + help="Path to a LibFewShot YAML config.", + ) + args, config_args = parser.parse_known_args() + # Config owns the framework-wide command-line options that remain. + sys.argv = [sys.argv[0]] + config_args + return args def main(rank, config): @@ -15,10 +33,11 @@ def main(rank, config): if __name__ == "__main__": - config = Config("./reproduce/Baseline/Baseline-miniImageNet--ravi-Conv64F-5-Reproduce.yaml").get_config_dict() + cli_args = parse_args() + train_config = Config(cli_args.config).get_config_dict() - if config["n_gpu"] > 1: - os.environ["CUDA_VISIBLE_DEVICES"] = config["device_ids"] - torch.multiprocessing.spawn(main, nprocs=config["n_gpu"], args=(config,)) + if train_config["n_gpu"] > 1: + os.environ["CUDA_VISIBLE_DEVICES"] = train_config["device_ids"] + torch.multiprocessing.spawn(main, nprocs=train_config["n_gpu"], args=(train_config,)) else: - main(0, config) \ No newline at end of file + main(0, train_config) diff --git a/run_trainer_resume.py b/run_trainer_resume.py index 142b02c8..8ea32436 100644 --- a/run_trainer_resume.py +++ b/run_trainer_resume.py @@ -1,14 +1,32 @@ # -*- coding: utf-8 -*- +"""Resume a LibFewShot experiment from its result directory.""" + +import argparse +import os import sys sys.dont_write_bytecode = True import torch -import os -from core.config import Config + from core import Trainer +from core.config import Config -PATH = "./results/Baseline-miniImageNet--ravi-Conv64F-5-5-Dec-01-2021-06-52-41" + +def parse_args(): + parser = argparse.ArgumentParser(description="Resume LibFewShot training.") + parser.add_argument( + "-r", + "--result-path", + "--result_path", + dest="result_path", + required=True, + help="Training result directory containing config.yaml and checkpoints/.", + ) + args, config_args = parser.parse_known_args() + # Config owns the framework-wide command-line options that remain. + sys.argv = [sys.argv[0]] + config_args + return args def main(rank, config): @@ -17,10 +35,14 @@ def main(rank, config): if __name__ == "__main__": - config = Config(os.path.join(PATH, "config.yaml"), is_resume=True).get_config_dict() - - if config["n_gpu"] > 1: - os.environ["CUDA_VISIBLE_DEVICES"] = config["device_ids"] - torch.multiprocessing.spawn(main, nprocs=config["n_gpu"], args=(config,)) + cli_args = parse_args() + config_path = os.path.join(cli_args.result_path, "config.yaml") + if not os.path.isfile(config_path): + raise FileNotFoundError("Config file not found: {}".format(config_path)) + + resume_config = Config(config_path, is_resume=True).get_config_dict() + if resume_config["n_gpu"] > 1: + os.environ["CUDA_VISIBLE_DEVICES"] = resume_config["device_ids"] + torch.multiprocessing.spawn(main, nprocs=resume_config["n_gpu"], args=(resume_config,)) else: - main(0, config) + main(0, resume_config)