Advanced Data Augmentation for Deep Learning

This repository provides implementations for advanced data augmentation techniques for deep learning models, including CutMix and MixMo.

Installation

bash pip install torch torchvision pip install datasets # for TinyImageNet

DataHandler Usage

The DataHandler class provides easy access to datasets with various augmentation techniques.

Basic Usage

```python from data_handler import DataHandler

Initialize DataHandler

datahandler = DataHandler(dataroot='./data')

Load CIFAR-10 with standard augmentations

trainloader, testloader = datahandler.getcifar10(batch_size=128)

Load CIFAR-100 with standard augmentations

trainloader, testloader = datahandler.getcifar100(batch_size=128)

Load TinyImageNet with standard augmentations

trainloader, valloader = datahandler.gettinyimagenet(batchsize=128)

Load CIFAR-100-C for robustness test

availablecorruptions = [ 'brightness', 'contrast', 'defocusblur', 'elastictransform', 'fog', 'frost', 'gaussianblur', 'gaussiannoise', 'glassblur', 'impulsenoise', 'jpegcompression', 'motionblur', 'pixelate', 'saturate', 'shotnoise', 'snow', 'spatter', 'specklenoise', 'zoomblur' ]

corruptionname = "gaussiannoise" cifar100cloader = datahandler.getcifar100c( corruptiontype=corruptionname, batchsize=128, severity=3 # Default is already 3, but specified for clarity ) ```

Using CutMix

CutMix augmentation is implemented using the official torchvision.transforms.v2.CutMix transform:

```python

Get a dataset first

datahandler = DataHandler() traindataset, _ = datahandler.getcifar10()

Create a CutMix DataLoader

cutmixloader = datahandler.getcutmixloader( dataset=traindataset, batchsize=128, alpha=1.0, # Parameter for Beta distribution num_classes=10 # Number of classes in the dataset )

Training with CutMix

for images, labels in cutmixloader: # labels are now one-hot encoded: [batchsize, num_classes] outputs = model(images) loss = criterion(outputs, labels) # ... rest of training loop ```

Using AugMix

AugMix augmentation is implemented using the official torchvision.transforms.AugMix transform:

```python

Get a dataset first

datahandler = DataHandler() traindataset, _ = datahandler.getcifar10()

Create an AugMix DataLoader

augmixloader = datahandler.getaugmixloader( dataset=traindataset, batchsize=128, severity=3, # Severity of base augmentations mixturewidth=3, # Number of augmentation chains chaindepth=-1, # Random depth between 1-3 alpha=1.0 # Hyperparameter for probability distributions )

Training with AugMix

for images, labels in augmix_loader: outputs = model(images) loss = criterion(outputs, labels) # ... rest of training loop ```

Batch Repetition (for MixMo)

For MixMo experiments, you can use batch repetition parameter to generate batches where each sample appears multiple times:

```python

Create DataLoader with batch repetition (b=2)

trainloader, _ = datahandler.getcifar10(batchsize=128, batch_repetitions=2)

This creates batches where each unique sample appears twice

For batchsize=128 and batchrepetitions=2, you'll get 64 unique samples per batch,

each repeated 2 times consecutively

```

Running the Main Scripts

WideResNet on CIFAR Datasets

The WRN_CIFAR_Main.py script trains WideResNet models on CIFAR-10 or CIFAR-100 with different MixMo approaches:

```bash

Train a WideResNet28 with CutMixMo on CIFAR-100

python WRNCIFARMain.py --dataset cifar100 --approach cutmixmo --width 10 --batchsize 64 --batchrepetitions 2 --alpha 1.0 --dataroot ./data --savedir ./results --runnumber 1

Train a WideResNet28 with LinearMixMo on CIFAR-10

python WRNCIFARMain.py --dataset cifar10 --approach linearmixmo --width 10 --batchsize 64 --batchrepetitions 2 --alpha 1.0 --dataroot ./data --savedir ./results --runnumber 1

Train a WideResNet28 with CutMixMo combined with CutMix data augmentation

python WRNCIFARMain.py --dataset cifar100 --approach cutmixmocutmix --width 10 --batchsize 64 --batchrepetitions 2 --alpha 1.0 --dataroot ./data --savedir ./results --run_number 1 ```

Parameters: - --dataset: Dataset to use ('cifar10' or 'cifar100') - --approach: MixMo approach ('linearmixmo', 'cutmixmo', 'linearmixmocutmix', 'cutmixmocutmix') - --width: Width factor for WideResNet (default: 10) - --batch_size: Batch size per GPU (default: 64) - --batch_repetitions: Batch repetition factor (b parameter, default: 2) - --alpha: Alpha parameter for Beta distribution (default: 1.0) - --data_root: Path to data directory - --save_dir: Directory to save results - --run_number: Run number for averaging results (1, 2, or 3) - --seed: Random seed (default: 42)

PreActResNet on TinyImageNet

The PreAct_TinyImageNet_Main.py script trains PreActResNet models on TinyImageNet with different MixMo approaches:

```bash

Train a PreActResNet18 with CutMixMo on TinyImageNet

python PreActTinyImageNetMain.py --dataset tinyimagenet --approach cutmixmo --width 2 --batchsize 100 --batchrepetitions 2 --alpha 2.0 --dataroot ./data --savedir ./results --runnumber 1

Train a PreActResNet18 with LinearMixMo on TinyImageNet

python PreActTinyImageNetMain.py --dataset tinyimagenet --approach linearmixmo --width 2 --batchsize 100 --batchrepetitions 2 --alpha 2.0 --dataroot ./data --savedir ./results --runnumber 1

Train a PreActResNet18 with CutMixMo on TinyImageNet with increased width

python PreActTinyImageNetMain.py --dataset tinyimagenet --approach cutmixmo --width 3 --batchsize 100 --batchrepetitions 2 --alpha 2.0 --dataroot ./data --savedir ./results --runnumber 1 ```

Parameters: - --dataset: Dataset to use (default: 'tinyimagenet') - --approach: MixMo approach ('linearmixmo' or 'cutmixmo') - --width: Width factor for PreActResNet (1, 2, or 3) - --batch_size: Batch size per GPU (default: 100) - --batch_repetitions: Batch repetition factor (b parameter, default: 2) - --alpha: Alpha parameter for Beta distribution (default: 1.0) - --data_root: Path to data directory - --save_dir: Directory to save results - --run_number: Run number for averaging results (1, 2, or 3) - --seed: Random seed (default: 42)

Implementation Details

CutMix vs CutMixMo

CutMix

CutMix combines two images by replacing a random patch in one image with a patch from another image:

```python from torchvision.transforms import v2 import torch

Initialize CutMix transform

cutmix = v2.CutMix(num_classes=10, alpha=1.0)

Apply CutMix to a batch

images, labels = next(iter(trainloader)) mixedimages, mixed_labels = cutmix(images, labels)

mixed_labels are now one-hot encoded with soft labels

```

CutMixMo

CutMixMo extends the CutMix concept to feature embeddings in a network. It is typically used with multi-branch networks where the outputs of parallel branches are mixed. The implementation is provided in the MixMo class and is ready to be used.

The MixMo models (WideResNet28 and PreActResNet18) provided in this repository already implement both Linear and Cut MixMo variants, so you can use them directly as shown in the examples below.

Example: Visualizing CutMix

Here's how to visualize the CutMix augmentation:

```python import matplotlib.pyplot as plt import torch import torchvision import torchvision.transforms as transforms from torchvision.transforms import v2 import numpy as np

Setup data

transform = transforms.Compose([ transforms.ToTensor(), ]) dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True)

Get a batch of images

images, labels = next(iter(dataloader))

Initialize CutMix

cutmix = v2.CutMix(num_classes=10, alpha=1.0)

Apply CutMix

mixedimages, mixedlabels = cutmix(images, labels)

Convert to numpy for visualization

def showimages(images, title): plt.figure(figsize=(12, 8)) for i in range(len(images)): plt.subplot(2, 2, i+1) plt.imshow(images[i].permute(1, 2, 0).numpy()) plt.title(f"Index {i}") plt.suptitle(title) plt.tightlayout() plt.show()

Show original images

show_images(images, "Original Images")

Show CutMix images

showimages(mixedimages, "CutMix Images")

Print the mixed labels (one-hot encoded)

print("Mixed Labels (One-Hot):", mixed_labels) ```

Using WideResNet and PreActResNet with Feature-Level Augmentation

The repository provides implementations of WideResNet28 and PreActResNet18 with feature-level augmentation support.

Creating and Using WideResNet28

```python from models.wideresnet import WideResNet28 import torch

Create a WideResNet28 model with CutMixMo augmentation

model = WideResNet28( widenfactor=10, # Controls network width dropoutrate=0.3, # Dropout rate for regularization numclasses=10, # Number of output classes (e.g., for CIFAR-10) augmentation_type='CutMixMo' # Options: 'none', 'LinearMixMo', 'CutMixMo' )

Example forward pass with two input batches

batchsize = 32 x1 = torch.randn(batchsize, 3, 32, 32) # First input batch x2 = torch.randn(batch_size, 3, 32, 32) # Second input batch

Forward pass with feature augmentation

out1, out2, outmix1, outmix2, kappa = model(x1, x2)

Use the outputs in your loss function

out1, out2: Original outputs from each branch

outmix1, outmix2: Outputs from mixed features

kappa: Mixing coefficients (useful for some loss functions)

```

Creating and Using PreActResNet18

```python from models.preact_resnet import PreActResNet18 import torch

Create a PreActResNet18 model with LinearMixMo augmentation

model = PreActResNet18( widenfactor=2, # Controls network width numclasses=10, # Number of output classes augmentation_type='LinearMixMo' # Options: 'none', 'LinearMixMo', 'CutMixMo' )

Example forward pass

batchsize = 32 x1 = torch.randn(batchsize, 3, 32, 32) x2 = torch.randn(batch_size, 3, 32, 32)

Forward pass with feature augmentation

out1, out2, outmix1, outmix2, lam = model(x1, x2)

lam is the mixing coefficient for LinearMixMo

```

Training a MixMo Model

When training with MixMo augmentation, you'll need to handle the dual inputs and multiple outputs:

```python import torch import torch.nn as nn import torch.optim as optim from datahandler import DataHandler from models.wideresnet import Wide_ResNet28

Create model

model = WideResNet28(widenfactor=10, dropoutrate=0.3, numclasses=10, augmentation_type='CutMixMo')

Setup optimizer

optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4) criterion = nn.CrossEntropyLoss()

Get data with batch repetition (each unique sample appears twice in a batch)

datahandler = DataHandler(dataroot='./data') trainloader, _ = datahandler.getcifar10(batchsize=128, batch_repetitions=2)

Training loop

model.train() for inputs, targets in trainloader: # Reshape the inputs and targets to separate the batch repetitions batchsize = inputs.size(0) // 2 x1, x2 = inputs[:batchsize], inputs[batchsize:] y1, y2 = targets[:batchsize], targets[batchsize:]

# Zero the parameter gradients
optimizer.zero_grad()

# Forward pass
out1, out2, out_mix1, out_mix2, kappa = model(x1, x2)

# Calculate loss
# 1. Original outputs loss
loss1 = criterion(out1, y1)
loss2 = criterion(out2, y2)

# 2. Mixed outputs loss
loss_mix1 = criterion(out_mix1, y1)
loss_mix2 = criterion(out_mix2, y2)

# Combine losses
loss = loss1 + loss2 + loss_mix1 + loss_mix2

# Backward pass and optimize
loss.backward()
optimizer.step()

```

References

• CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features
• MixMo: Mixing Multiple Inputs for Multiple Outputs via Deep Subnetworks
• AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty