Update demo

Author: lizekai-richard

demo_aug.py

@@ -1,28 +1,79 @@
 import os
 import argparse
 import torch
-from ddranking.metrics import AugmentationRobustnessEvaluator
+from ddranking.metrics import AugmentationRobustScore
 from ddranking.config import Config
 
-def main(args):
-    root = args.root
-    method_name = args.method_name
-    dataset = args.dataset
-    ipc = args.ipc
+""" Use config file to specify the arguments (Recommended) """
+config = Config.from_file("./configs/Demo_ARS.yaml")
+aug_evaluator = AugmentationRobustScore(config)
 
-    print(f"Evaluating {method_name} on {dataset} with ipc{ipc}")
-    syn_image_dir = os.path.join(root, f"DD-Ranking/baselines/{method_name}/{dataset}/IPC{ipc}/")
-    config = Config.from_file(f"./{method_name}/{dataset}/IPC{ipc}_Aug.yaml")
-    aug_obj = AugmentationRobustnessEvaluator(config)
-    aug_obj.compute_metrics(image_path=syn_image_dir, syn_lr=args.syn_lr)
+syn_data_dir = "./baselines/SRe2L/ImageNet1K/IPC10/"
+print(aug_evaluator.compute_metrics(image_path=syn_data_dir, syn_lr=0.001))
 
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--root", type=str, default="/home/wangkai/")
-    parser.add_argument("--method_name", type=str, default="SRe2L")
-    parser.add_argument("--dataset", type=str, default="ImageNet1K")
-    parser.add_argument("--ipc", type=int, default=10)
-    parser.add_argument("--syn_lr", type=float, default=0.001)
-    args = parser.parse_args()
-    main(args)
+""" Use keyword arguments """
+from torchvision import transforms
+device = "cuda"
+method_name = "SRe2L"                    # Specify your method name
+ipc = 10                                # Specify your IPC
+dataset = "ImageNet1K"                     # Specify your dataset name
+syn_data_dir = "./SRe2L/ImageNet1K/IPC10/"  # Specify your synthetic data path
+data_dir = "./datasets"                 # Specify your dataset path
+model_name = "ResNet-18-BN"                # Specify your model name
+im_size = (224, 224)                      # Specify your image size
+cutmix_params = {                          # Specify your data augmentation parameters
+    "beta": 1.0
+}
+
+syn_images = torch.load(os.path.join(syn_data_dir, f"images.pt"), map_location='cpu')
+soft_labels = torch.load(os.path.join(syn_data_dir, f"labels.pt"), map_location='cpu')
+syn_lr = torch.load(os.path.join(syn_data_dir, f"lr.pt"), map_location='cpu')
+save_path = f"./results/{dataset}/{model_name}/IPC{ipc}/dm_hard_scores.csv"
+
+custom_train_trans = transforms.Compose([
+    transforms.RandomResizedCrop(224, scale=(0.08, 1.0)),
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+custom_val_trans = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+aug_evaluator = AugmentationRobustScore(
+    dataset=dataset,
+    real_data_path=data_dir, 
+    ipc=ipc, 
+    model_name=model_name,
+    label_type='soft',
+    soft_label_criterion='kl',  # Use Soft Cross Entropy Loss
+    soft_label_mode='M',         # Use one-to-one image to soft label mapping
+    loss_fn_kwargs={'temperature': 1.0, 'scale_loss': False},
+    optimizer='adamw',             # Use SGD optimizer
+    lr_scheduler='cosine',         # Use StepLR learning rate scheduler
+    weight_decay=0.01,         
+    momentum=0.9,                
+    num_eval=5,                  
+    data_aug_func='cutmix',         # Use DSA data augmentation
+    aug_params=cutmix_params,       # Specify dsa parameters
+    im_size=im_size,
+    num_epochs=300,
+    num_workers=4,
+    stu_use_torchvision=True,
+    tea_use_torchvision=True,
+    random_data_format='tensor',
+    random_data_path='./random_data',
+    custom_train_trans=custom_train_trans,
+    custom_val_trans=custom_val_trans,
+    batch_size=256,
+    teacher_dir='./teacher_models',
+    teacher_model_name=['ResNet-18-BN'],
+    device=device,
+    dist=True,
+    save_path=save_path
+)
+print(aug_evaluator.compute_metrics(image_path=syn_data_dir, syn_lr=0.001))

demo_hard.py

@@ -1,14 +1,14 @@
 import os
 import torch
 import warnings
-from ddranking.metrics import HardLabelEvaluator
+from ddranking.metrics import LabelRobustScoreHard
 from ddranking.config import Config
 warnings.filterwarnings("ignore")
 
 
 """ Use config file to specify the arguments (Recommended) """
-config = Config.from_file("./configs/Demo_Hard_Label.yaml")
-hard_label_evaluator = HardLabelEvaluator(config)
+config = Config.from_file("./configs/Demo_LRS_Hard_Label.yaml")
+hard_label_evaluator = LabelRobustScoreHard(config)
 
 syn_data_dir = "./baselines/DM/CIFAR10/IPC10/"
 syn_images = torch.load(os.path.join(syn_data_dir, f"images.pt"), map_location='cpu')
@@ -39,12 +39,11 @@
 
 syn_images = torch.load(os.path.join(syn_data_dir, f"images.pt"), map_location='cpu')
 save_path = f"./results/{dataset}/{model_name}/IPC{ipc}/dm_hard_scores.csv"
-hard_label_evaluator = HardLabelEvaluator(
+hard_label_evaluator = LabelRobustScoreHard(
     dataset=dataset, 
     real_data_path=data_dir, 
     ipc=ipc, 
     model_name=model_name,
-    default_lr=0.01,
     optimizer='sgd',             # Use SGD optimizer
     lr_scheduler='step',         # Use StepLR learning rate scheduler
     weight_decay=0.0005,
@@ -57,10 +56,15 @@
     num_epochs=1000,
     num_workers=4,
     use_torchvision=False,
-    syn_batch_size=128,
+    syn_batch_size=256,
     real_batch_size=256,
+    custom_train_trans=None,
     custom_val_trans=None,
     device=device,
-    save_path=save_path
+    dist=True,
+    save_path=save_path,
+    random_data_format='tensor',
+    random_data_path='./random_data',
+    eval_full_data=True,
 )
 print(hard_label_evaluator.compute_metrics(image_tensor=syn_images, syn_lr=0.01))

demo_soft.py

@@ -1,14 +1,14 @@
 import os
 import torch
 import warnings
-from ddranking.metrics import SoftLabelEvaluator
+from ddranking.metrics import LabelRobustScoreSoft
 from ddranking.config import Config
 warnings.filterwarnings("ignore")
 
 
 """ Use config file to specify the arguments (Recommended) """
-config = Config.from_file("./configs/Demo_Soft_Label.yaml")
-soft_label_evaluator = SoftLabelEvaluator(config)
+config = Config.from_file("./configs/Demo_LRS_Soft_Label.yaml")
+soft_label_evaluator = LabelRobustScoreSoft(config)
 
 syn_data_dir = "./baselines/DATM/CIFAR10/IPC10/"
 syn_images = torch.load(os.path.join(syn_data_dir, f"images.pt"), map_location='cpu')
@@ -41,16 +41,17 @@
 soft_labels = torch.load(os.path.join(syn_data_dir, f"labels.pt"), map_location='cpu')
 syn_lr = torch.load(os.path.join(syn_data_dir, f"lr.pt"), map_location='cpu')
 save_path = f"./results/{dataset}/{model_name}/IPC{ipc}/dm_hard_scores.csv"
-soft_label_evaluator = SoftLabelEvaluator(
+soft_label_evaluator = LabelRobustScoreSoft(
     dataset=dataset,
     real_data_path=data_dir, 
     ipc=ipc, 
     model_name=model_name,
     soft_label_criterion='sce',  # Use Soft Cross Entropy Loss
     soft_label_mode='S',         # Use one-to-one image to soft label mapping
-    default_lr=0.01,
+    loss_fn_kwargs={'temperature': 1.0, 'scale_loss': False},
     optimizer='sgd',             # Use SGD optimizer
     lr_scheduler='step',         # Use StepLR learning rate scheduler
+    step_size=500,
     weight_decay=0.0005,         
     momentum=0.9,                
     use_zca=True,                # Use ZCA whitening (please disable it if you didn't use it to distill synthetic data)
@@ -60,13 +61,19 @@
     im_size=im_size,
     num_epochs=1000,
     num_workers=4,
+    eval_full_data=True,
     stu_use_torchvision=False,
     tea_use_torchvision=False,
+    random_data_format='tensor',
+    random_data_path='./random_data',
+    custom_train_trans=None,
     custom_val_trans=None,
-    syn_batch_size=128,
+    syn_batch_size=256,
     real_batch_size=256,
     teacher_dir='./teacher_models',
+    teacher_model_name=['ConvNet-3'],
     device=device,
+    dist=True,
     save_path=save_path
 )
 print(soft_label_evaluator.compute_metrics(image_tensor=syn_images, soft_labels=soft_labels, syn_lr=syn_lr))

doc/metrics/ars.md

@@ -25,8 +25,7 @@ dd_ranking.metrics.AugmentationRobustScore(config: Optional[Config] = None,
     use_zca: bool=False,
     random_data_format: str='image',
     random_data_path: str=None,
-    real_batch_size: int=256, 
-    syn_batch_size: int=256, 
+    batch_size: int=256,
     save_path: str=None, 
     stu_use_torchvision: bool=False, 
     tea_use_torchvision: bool=False, 
@@ -65,8 +64,7 @@ A class for evaluating the performance of a dataset distillation method with sof
 - **num_eval**(<span style="color:#FF6B00;">int</span>): Number of evaluations to perform.
 - **im_size**(<span style="color:#FF6B00;">tuple</span>): Size of the images.
 - **num_epochs**(<span style="color:#FF6B00;">int</span>): Number of epochs to train.
-- **real_batch_size**(<span style="color:#FF6B00;">int</span>): Batch size for the real dataset.
-- **syn_batch_size**(<span style="color:#FF6B00;">int</span>): Batch size for the synthetic dataset.
+- **batch_size**(<span style="color:#FF6B00;">int</span>): Batch size for the model training.
 - **stu_use_torchvision**(<span style="color:#FF6B00;">bool</span>): Whether to use torchvision to initialize the student model.
 - **tea_use_torchvision**(<span style="color:#FF6B00;">bool</span>): Whether to use torchvision to initialize the teacher model.
 - **teacher_dir**(<span style="color:#FF6B00;">str</span>): Path to the teacher model.