Add aug doc and reformat aug code

lizekai-richard · lizekai-richard · commit 782030b28d6d · 2024-12-29T12:48:04.000+08:00
diff --git a/configs/Demo_Hard_Label.yaml b/configs/Demo_Hard_Label.yaml
@@ -15,14 +15,14 @@ use_torchvision: false
 # data augmentation
 data_aug_func: "dsa"
 aug_params:
-  prob_flip: 0.5
-  ratio_rotate: 15.0
+  flip: 0.5
+  rotate: 15.0
   saturation: 2.0
   brightness: 1.0
   contrast: 0.5
-  ratio_scale: 1.2
-  ratio_crop_pad: 0.125
-  ratio_cutout: 0.5
+  scale: 1.2
+  crop: 0.125
+  cutout: 0.5
 use_zca: false
 
 custom_train_trans: null
diff --git a/configs/Demo_Soft_Label.yaml b/configs/Demo_Soft_Label.yaml
@@ -16,14 +16,14 @@ teacher_dir: ./teacher_models
 # syntheticdata augmentation
 data_aug_func: dsa
 aug_params:
-    prob_flip: 0.5
-    ratio_rotate: 15.0
+    flip: 0.5
+    rotate: 15.0
     saturation: 2.0
     brightness: 1.0
     contrast: 0.5
-    ratio_scale: 1.2
-    ratio_crop_pad: 0.125
-    ratio_cutout: 0.5
+    scale: 1.2
+    crop: 0.125
+    cutout: 0.5
 use_zca: true
 custom_train_trans: null
 
diff --git a/dd_ranking/__init__.py b/dd_ranking/__init__.py
@@ -1,5 +1,5 @@
-import dd_ranking.aug as aug
-import dd_ranking.config as config
-import dd_ranking.loss as loss
-import dd_ranking.metrics as metrics
-import dd_ranking.utils as utils
+from .aug import DSA, Mixup, Cutmix, ZCAWhitening
+from .config import Config
+from .loss import KLDivergenceLoss, SoftCrossEntropyLoss
+from .metrics import HardScoreEvaluator, SoftLabelEvaluator, GeneralEvaluator
+from .utils import get_dataset, build_model, get_convnet, get_lenet, get_resnet, get_vgg, get_alexnet
diff --git a/dd_ranking/aug/__init__.py b/dd_ranking/aug/__init__.py
@@ -1,4 +1,4 @@
-from .dsa import DSAugmentation
-from .mixup import MixupAugmentation
-from .cutmix import CutmixAugmentation
-from .zca import ZCAWhiteningAugmentation
+from .dsa import DSA
+from .mixup import Mixup
+from .cutmix import Cutmix
+from .zca import ZCAWhitening
diff --git a/dd_ranking/aug/cutmix.py b/dd_ranking/aug/cutmix.py
@@ -3,9 +3,9 @@
 import kornia
 
 
-class CutmixAugmentation:
+class Cutmix:
     def __init__(self, params: dict):
-        self.cutmix_p = params["cutmix_p"]
+        self.beta = params["beta"]
 
     def rand_bbox(self, size, lam):
         W = size[2]
@@ -27,7 +27,7 @@ def rand_bbox(self, size, lam):
 
     def cutmix(self, images):
         rand_index = torch.randperm(images.size()[0]).to(images.device)
-        lam = np.random.beta(self.cutmix_p, self.cutmix_p)
+        lam = np.random.beta(self.beta, self.beta)
         bbx1, bby1, bbx2, bby2 = self.rand_bbox(images.size(), lam)
 
         images[:, :, bbx1:bbx2, bby1:bby2] = images[rand_index, :, bbx1:bbx2, bby1:bby2]
diff --git a/dd_ranking/aug/dsa.py b/dd_ranking/aug/dsa.py
@@ -3,7 +3,7 @@
 import torch.nn.functional as F
 
 
-class DSAugmentation:
+class DSA:
 
     def __init__(self, params: dict, seed: int=-1, aug_mode: str='S'):        
         self.params = params
@@ -30,7 +30,7 @@ def set_seed_DiffAug(self):
     def rand_scale(self, x):
         # x>1, max scale
         # sx, sy: (0, +oo), 1: orignial size, 0.5: enlarge 2 times
-        ratio = self.params["ratio_scale"]
+        ratio = self.params["scale"]
         self.set_seed_DiffAug()
         sx = torch.rand(x.shape[0]) * (ratio - 1.0/ratio) + 1.0/ratio
         self.set_seed_DiffAug()
@@ -45,7 +45,7 @@ def rand_scale(self, x):
         return x
 
     def rand_rotate(self, x): # [-180, 180], 90: anticlockwise 90 degree
-        ratio = self.params["ratio_rotate"]
+        ratio = self.params["rotate"]
         self.set_seed_DiffAug()
         theta = (torch.rand(x.shape[0]) - 0.5) * 2 * ratio / 180 * float(np.pi)
         theta = [[[torch.cos(theta[i]), torch.sin(-theta[i]), 0],
@@ -58,7 +58,7 @@ def rand_rotate(self, x): # [-180, 180], 90: anticlockwise 90 degree
         return x
 
     def rand_flip(self, x):
-        prob = self.params["prob_flip"]
+        prob = self.params["flip"]
         self.set_seed_DiffAug()
         randf = torch.rand(x.size(0), 1, 1, 1, device=x.device)
         if self.params["siamese"]: # Siamese augmentation:
@@ -99,7 +99,7 @@ def rand_color(self, x):
 
     def rand_crop(self, x):
         # The image is padded on its surrounding and then cropped.
-        ratio = self.params["ratio_crop_pad"]
+        ratio = self.params["crop"]
         shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
         self.set_seed_DiffAug()
         translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)
@@ -120,7 +120,7 @@ def rand_crop(self, x):
         return x
 
     def rand_cutout(self, x):
-        ratio = self.params["ratio_cutout"]
+        ratio = self.params["cutout"]
         cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
         self.set_seed_DiffAug()
         offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device)
diff --git a/dd_ranking/aug/mixup.py b/dd_ranking/aug/mixup.py
@@ -3,13 +3,13 @@
 import kornia
 
 
-class MixupAugmentation:
+class Mixup:
     def __init__(self, params: dict):
-        self.mixup_p = params["mixup_p"]
+        self.lambda_ = params["lambda"]
 
     def mixup(self, images):
         rand_index = torch.randperm(images.size()[0]).to(images.device)
-        lam = np.random.beta(self.mixup_p, self.mixup_p)
+        lam = np.random.beta(self.lambda_, self.lambda_)
 
         mixed_images = lam * images + (1 - lam) * images[rand_index]
         return mixed_images
diff --git a/dd_ranking/aug/zca.py b/dd_ranking/aug/zca.py
@@ -1,7 +1,7 @@
 import kornia
 
 
-class ZCAWhiteningAugmentation:
+class ZCAWhitening:
     def __init__(self, params: dict):
         self.transform = kornia.enhance.ZCAWhitening()
 
diff --git a/dd_ranking/metrics/general.py b/dd_ranking/metrics/general.py
@@ -14,7 +14,7 @@
 from dd_ranking.utils import set_seed, get_optimizer, get_lr_scheduler
 from dd_ranking.utils import train_one_epoch, validate
 from dd_ranking.loss import SoftCrossEntropyLoss, KLDivergenceLoss
-from dd_ranking.aug import DSAugmentation, MixupAugmentation, CutmixAugmentation, ZCAWhiteningAugmentation
+from dd_ranking.aug import DSA, Mixup, Cutmix, ZCAWhitening
 from dd_ranking.config import Config
 
 
@@ -147,11 +147,11 @@ def __init__(self,
         if data_aug_func is None:
             self.aug_func = None
         elif data_aug_func == 'dsa':
-            self.aug_func = DSA_Augmentation(aug_params)
+            self.aug_func = DSA(aug_params)
         elif data_aug_func == 'mixup':
-            self.aug_func = Mixup_Augmentation(aug_params)  
+            self.aug_func = Mixup(aug_params)  
         elif data_aug_func == 'cutmix':
-            self.aug_func = Cutmix_Augmentation(aug_params)
+            self.aug_func = Cutmix(aug_params)
         else:
             raise ValueError(f"Invalid data augmentation function: {data_aug_func}")
 
diff --git a/dd_ranking/metrics/hard_label.py b/dd_ranking/metrics/hard_label.py
@@ -12,7 +12,7 @@
 from dd_ranking.utils import build_model, get_pretrained_model_path
 from dd_ranking.utils import TensorDataset, get_random_images, get_dataset, save_results
 from dd_ranking.utils import set_seed, train_one_epoch, validate, get_optimizer, get_lr_scheduler
-from dd_ranking.aug import DSAugmentation, MixupAugmentation, CutmixAugmentation, ZCAWhiteningAugmentation
+from dd_ranking.aug import DSA, Mixup, Cutmix, ZCAWhitening
 from dd_ranking.config import Config
 
 
@@ -82,13 +82,13 @@ def __init__(self, config: Config=None, dataset: str='CIFAR10', real_data_path:
         self.device = device
 
         if data_aug_func == 'dsa':
-            self.aug_func = DSA_Augmentation(aug_params)
+            self.aug_func = DSA(aug_params)
         elif data_aug_func == 'zca':
-            self.aug_func = ZCA_Whitening_Augmentation(aug_params)
+            self.aug_func = ZCAWhitening(aug_params)
         elif data_aug_func == 'mixup':
-            self.aug_func = Mixup_Augmentation(aug_params)
+            self.aug_func = Mixup(aug_params)
         elif data_aug_func == 'cutmix':
-            self.aug_func = Cutmix_Augmentation(aug_params)
+            self.aug_func = Cutmix(aug_params)
         else:
             self.aug_func = None
 
diff --git a/dd_ranking/metrics/soft_label.py b/dd_ranking/metrics/soft_label.py
@@ -13,14 +13,14 @@
 from dd_ranking.utils import TensorDataset, get_random_images, get_dataset, save_results
 from dd_ranking.utils import set_seed, train_one_epoch, validate, get_optimizer, get_lr_scheduler
 from dd_ranking.loss import SoftCrossEntropyLoss, KLDivergenceLoss
-from dd_ranking.aug import DSAugmentation, MixupAugmentation, CutmixAugmentation, ZCAWhiteningAugmentation
+from dd_ranking.aug import DSA, Mixup, Cutmix, ZCAWhitening
 from dd_ranking.config import Config
 
 
 class SoftLabelEvaluator:
 
     def __init__(self, config: Config=None, dataset: str='CIFAR10', real_data_path: str='./dataset/', ipc: int=10, model_name: str='ConvNet-3', 
-                 soft_label_criterion: str='kl', data_aug_func: str='cutmix', aug_params: dict={'cutmix_p': 1.0}, soft_label_mode: str='S',
+                 soft_label_criterion: str='kl', data_aug_func: str='cutmix', aug_params: dict={'beta': 1.0}, soft_label_mode: str='S',
                  optimizer: str='sgd', lr_scheduler: str='step', temperature: float=1.0, weight_decay: float=0.0005, 
                  momentum: float=0.9, num_eval: int=5, im_size: tuple=(32, 32), num_epochs: int=300, use_zca: bool=False,
                  real_batch_size: int=256, syn_batch_size: int=256, default_lr: float=0.01, save_path: str=None, 
@@ -92,11 +92,11 @@ def __init__(self, config: Config=None, dataset: str='CIFAR10', real_data_path:
         self.device = device
 
         if data_aug_func == 'dsa':
-            self.aug_func = DSA_Augmentation(aug_params)
+            self.aug_func = DSA(aug_params)
         elif data_aug_func == 'mixup':
-            self.aug_func = Mixup_Augmentation(aug_params)
+            self.aug_func = Mixup(aug_params)
         elif data_aug_func == 'cutmix':
-            self.aug_func = Cutmix_Augmentation(aug_params)
+            self.aug_func = Cutmix(aug_params)
         else:
             self.aug_func = None
 
@@ -239,7 +239,7 @@ def compute_soft_label_metrics(self, model, image_tensor, image_path, lr, soft_l
         if image_tensor is None:
             soft_label_dataset = datasets.ImageFolder(root=image_path, transform=self.custom_train_trans)
         else:
-            soft_label_dataset = TensorDataset(image_tensor, labels)
+            soft_label_dataset = TensorDataset(image_tensor, labels, transform=self.custom_train_trans)
         train_loader = DataLoader(soft_label_dataset, batch_size=self.syn_batch_size, num_workers=self.num_workers, shuffle=True)
 
         if self.soft_label_criterion == 'sce':
diff --git a/dd_ranking/utils/__init__.py b/dd_ranking/utils/__init__.py
@@ -1,4 +1,4 @@
 from .misc import set_seed, save_results
 from .data import get_dataset, get_random_images, TensorDataset
-from .model import build_model, get_pretrained_model_path
+from .model import build_model, get_pretrained_model_path, get_convnet, get_lenet, get_resnet, get_vgg, get_alexnet
 from .train_and_eval import get_optimizer, get_lr_scheduler, train_one_epoch, validate
diff --git a/dd_ranking/utils/data.py b/dd_ranking/utils/data.py
@@ -34,16 +34,14 @@ class Config:
 
 
 class TensorDataset(torch.utils.data.Dataset):
-    def __init__(self, images: Tensor, labels: Tensor, transform: transforms.Compose=None):
+    
+    def __init__(self, images: Tensor, labels: Tensor):
         self.images = images
         self.labels = labels
-        self.transform = transform
 
     def __getitem__(self, index: int):
         image = self.images[index]
         label = self.labels[index]
-        if self.transform:
-            image = self.transform(image)
         return image, label
     
     def __len__(self):
diff --git a/demo_hard.py b/demo_hard.py
@@ -1,7 +1,9 @@
 import os
 import torch
+import warnings
 from dd_ranking.metrics import HardLabelEvaluator
 from dd_ranking.config import Config
+warnings.filterwarnings("ignore")
 
 
 """ Use config file to specify the arguments (Recommended) """
diff --git a/demo_soft.py b/demo_soft.py
@@ -3,7 +3,7 @@
 import warnings
 from dd_ranking.metrics import SoftLabelEvaluator
 from dd_ranking.config import Config
-warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore")
 
 
 """ Use config file to specify the arguments (Recommended) """
@@ -27,14 +27,14 @@
 model_name = "ConvNet-3"                # Specify your model name
 im_size = (32, 32)                      # Specify your image size
 dsa_params = {                          # Specify your data augmentation parameters
-    "prob_flip": 0.5,
-    "ratio_rotate": 15.0,
+    "flip": 0.5,
+    "rotate": 15.0,
     "saturation": 2.0,
     "brightness": 1.0,
     "contrast": 0.5,
-    "ratio_scale": 1.2,
-    "ratio_crop_pad": 0.125,
-    "ratio_cutout": 0.5
+    "scale": 1.2,
+    "crop": 0.125,
+    "cutout": 0.5
 }
 
 syn_images = torch.load(os.path.join(syn_data_dir, f"images.pt"), map_location='cpu')
diff --git a/doc/augmentations/cutmix.md b/doc/augmentations/cutmix.md
@@ -0,0 +1,25 @@
+## Cutmix
+
+Cutmix is a data augmentation technique that creates new samples by combining patches from two images while blending their labels proportionally to the area of the patches.. We follow the implementation of cutmix in [SRe2L](https://github.com/VILA-Lab/SRe2L/tree/main/SRe2L).
+
+<div style="background-color:#F7F7F7; padding:15px; border:1px solid #E0E0E0; border-top:3px solid #FF0000; font-family:monospace; font-size:14px;">
+
+<span style="color:#FF6B00;">CLASS</span> 
+dd_ranking.aug.Cutmix(params: dict)
+[**[SOURCE]**](https://github.com/NUS-HPC-AI-Lab/DD-Ranking/blob/main/dd_ranking/aug/cutmix.py)
+
+</div>
+
+### Parameters
+
+- **params**(<span style="color:#FF6B00;">dict</span>): Parameters for the cutmix augmentation. We require the parameters to be in the format of `{'param_name': param_value}`. For cutmix, only `beta` (beta distribution parameter) needs to be specified, e.g. `{'beta': 1.0}`.
+
+### Example
+    
+```python
+# When intializing an evaluator with cutmix augmentation, and cutmix object will be constructed.
+>>> self.aug_func = Cutmix(params={'beta': 1.0})
+
+# During training, the cutmix object will be used to augment the data.
+>>> images = aug_func(images)
+```
diff --git a/doc/augmentations/dsa.md b/doc/augmentations/dsa.md
@@ -0,0 +1,37 @@
+## Differentiable Siamese Augmentation (DSA)
+
+DSA is one of differentiable data augmentations, first used in the dataset distillation task by [DSA](https://github.com/VICO-UoE/DatasetCondensation). 
+Our implementation of DSA is adopted from [DSA](https://github.com/VICO-UoE/DatasetCondensation). It supports the following differentiable augmentations:
+
+- Random Flip
+- Random Rotation
+- Random Saturation
+- Random Brightness
+- Random Contrast
+- Random Scale
+- Random Crop
+- Random Cutout
+
+<div style="background-color:#F7F7F7; padding:15px; border:1px solid #E0E0E0; border-top:3px solid #FF0000; font-family:monospace; font-size:14px;">
+
+<span style="color:#FF6B00;">CLASS</span> 
+dd_ranking.aug.DSA(params: dict, seed: int, aug_mode: str)
+[**[SOURCE]**](https://github.com/NUS-HPC-AI-Lab/DD-Ranking/blob/main/dd_ranking/aug/dsa.py)
+
+</div>
+
+### Parameters
+
+- **params**(<span style="color:#FF6B00;">dict</span>): Parameters for the DSA augmentations. We require the parameters to be in the format of `{'param_name': param_value}`. For example, `{'flip': 0.5, 'rotate': 15.0, 'scale': 1.2, 'crop': 0.125, 'cutout': 0.5, 'brightness': 1.0, 'contrast': 0.5, 'saturation': 2.0}`.
+- **seed**(<span style="color:#FF6B00;">int</span>): Random seed. Default is `-1`.
+- **aug_mode**(<span style="color:#FF6B00;">str</span>): `S` for randomly selecting one augmentation for each batch. `M` for applying all augmentations for each batch.
+
+### Example
+
+```python
+# When intializing an evaluator with DSA augmentation, and DSA object will be constructed.
+>>> self.aug_func = DSA(params={'flip': 0.5, 'rotate': 15.0, 'scale': 1.2, 'crop': 0.125, 'cutout': 0.5, 'brightness': 1.0, 'contrast': 0.5, 'saturation': 2.0}, seed=-1, aug_mode='S')
+
+# During training, the DSA object will be used to augment the data.
+>>> images = aug_func(images)
+```
diff --git a/doc/augmentations/mixup.md b/doc/augmentations/mixup.md
diff --git a/doc/augmentations/overview.md b/doc/augmentations/overview.md
diff --git a/doc/augmentations/torchvision.md b/doc/augmentations/torchvision.md
diff --git a/doc/metrics/hard-label.md b/doc/metrics/hard-label.md
diff --git a/doc/metrics/soft-label.md b/doc/metrics/soft-label.md
diff --git a/setup.py b/setup.py
