### Main Execution Script

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/TherEncoding.html

Setup for logging, data loading, and the training loop execution.

```python
if __name__ =='__main__':

    logger = logging.getLogger('Thermometer Encoding')

    handler = logging.StreamHandler()  # Handler for the logger
    handler.setFormatter(logging.Formatter('%(asctime)s'))
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)

    logger.info('Start attack.')

    torch.manual_seed(100)
    device = torch.device("cuda")

    #ipdb.set_trace()

    logger.info('Load trainset.')
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('deeprobust/image/data', train=True, download=True,
                     transform=transforms.Compose([transforms.ToTensor()])),
        batch_size=100,
        shuffle=True)

    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('deeprobust/image/data', train=False,
                    transform=transforms.Compose([transforms.ToTensor()])),
        batch_size=1000,
        shuffle=True)

    #ipdb.set_trace()

    #TODO: change the channel according to the dataset.
    LEVELS = 10
    channel = 1
    model = Net(in_channel1 = channel * LEVELS, out_channel1= 32 * LEVELS, out_channel2= 64 * LEVELS).to(device)
    optimizer = optim.SGD(model.parameters(), lr = 0.0001, momentum = 0.2)
    logger.info('Load model.')

    save_model = True
    for epoch in range(1, 50 + 1):     ## 5 batches
        print('Running epoch ', epoch)

        train(model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)

        if (save_model):
            torch.save(model.state_dict(), "deeprobust/image/save_models/thermometer_encoding.pt")
```

--------------------------------

### TRADES Training Loop Integration

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/trades.html

Example of how to integrate the TRADES loss calculation within a standard PyTorch training loop.

```python
for batch_idx, (data, target) in enumerate(train_loader):
            
            optimizer.zero_grad()

            data, target = data.to(self.device), target.to(self.device)

            # calculate robust loss
            loss = self.trades_loss(model = self.model,
                            x_natural = data,
                            y = target,
                            optimizer = optimizer,
                            step_size = self.step_size,
                            epsilon = self.epsilon,
                            perturb_steps = self.num_steps,
                            beta = self.beta)

            loss.backward()
            optimizer.step()

            # print progress
            if batch_idx % self.log_interval == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                    100. * batch_idx / len(train_loader), loss.item()))
```

--------------------------------

### Evaluate Model Performance

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/trades.html

Performs evaluation on natural examples and returns loss and accuracy.

```python
    def test(self, model, device, test_loader):
        model.eval()
        test_loss = 0
        correct = 0
        
        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                test_loss += F.cross_entropy(output, target, size_average=False).item()
                pred = output.max(1, keepdim=True)[1]
                correct += pred.eq(target.view_as(pred)).sum().item()
        test_loss /= len(test_loader.dataset)
        print('Test: Clean loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))
        test_accuracy = correct / len(test_loader.dataset)
        
        return test_loss, test_accuracy
```

--------------------------------

### FGSMtraining Class Implementation

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/fgsmtraining.html

The core class for performing FGSM adversarial training, inheriting from BaseDefense. It includes methods for generating adversarial examples, training, and testing.

```python
"""
This is the implementation of fgsm training.

References
 ----------
..[1]Szegedy, C., Zaremba, W., Sutskever, I., Estrach, J. B., Erhan, D., Goodfellow, I., & Fergus, R. (2014, January).
Intriguing properties of neural networks.
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import torch.nn.functional as F

import numpy as np
from PIL import Image
import os

from deeprobust.image.netmodels import CNN
from deeprobust.image.attack.fgsm import FGSM
from deeprobust.image.defense.base_defense import BaseDefense


[docs]class FGSMtraining(BaseDefense):
    """
    FGSM adversarial training.
    """

    def __init__(self, model, device):
        if not torch.cuda.is_available():
            print('CUDA not availiable, using cpu...')
            self.device = 'cpu'
        else:
            self.device = device

        self.model = model


[docs]    def generate(self, train_loader, test_loader, **kwargs):
        ""FGSM adversarial training process.

        Parameters
        ----------
        train_loader :
            training data loader
        test_loader :
            testing data loader
        kwargs :
            kwargs
        """
        self.parse_params(**kwargs)
        torch.manual_seed(100)
        device = torch.device(self.device)
        optimizer = optim.Adam(self.model.parameters(), self.lr_train)

        for epoch in range(1, self.epoch_num + 1):

            print(epoch, flush = True)
            self.train(self.device, train_loader, optimizer, epoch)
            self.test(self.model, self.device, test_loader)

            if (self.save_model):
                if os.path.isdir('./' + self.save_dir):
                    torch.save(self.model.state_dict(), './' + self.save_dir + "/" + self.save_name)
                    print("model saved in " + './' + self.save_dir)
                else:
                    print("make new directory and save model in " + './' + self.save_dir)
                    os.mkdir('./' + self.save_dir)
                    torch.save(self.model.state_dict(), './' + self.save_dir +"/" + self.save_name)

        return self.model


[docs]    def parse_params(self,
                     save_dir = "defense_models",
                     save_model = True,
                     save_name = "mnist_fgsmtraining_0.2.pt",
                     epsilon = 0.2,
                     epoch_num = 50,
                     lr_train = 0.005,
                     momentum = 0.1):
        ""parse_params.

        Parameters
        ----------
        save_dir :
            dir
        save_model :
            Whether to save model
        save_name :
            model name
        epsilon :
            attack perturbation constraint
        epoch_num :
            number of training epoch
        lr_train :
            training learning rate
        momentum :
            momentum for optimizor
        """
        self.save_model = True
        self.save_dir = save_dir
        self.save_name = save_name
        self.epsilon = epsilon
        self.epoch_num = epoch_num
        self.lr_train = lr_train
        self.momentum = momentum


[docs]    def train(self, device, train_loader, optimizer, epoch):
        ""
        training process.

        Parameters
        ----------
        device :
            device
        train_loader :
            training data loader
        optimizer :
            optimizer
        epoch :
            training epoch
        """
        self.model.train()
        correct = 0
        bs = train_loader.batch_size

        for batch_idx, (data, target) in enumerate(train_loader):

            optimizer.zero_grad()

            data, target = data.to(device), target.to(device)

            data_adv, output = self.adv_data(data, target, ep = self.epsilon)

            loss = self.calculate_loss(output, target)

            loss.backward()
            optimizer.step()

            pred = output.argmax(dim = 1, keepdim = True)
            correct += pred.eq(target.view_as(pred)).sum().item()

            #print every 10
            if batch_idx % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tAccuracy:{:.2f}%'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item(), 100 * correct/(10*bs)))
                correct = 0


[docs]    def test(self, model, device, test_loader):
        ""
        testing process.

        Parameters
        ----------
        model :
            model
        device :
            device
        test_loader :
            testing dataloder
        """
        model.eval()

        test_loss = 0
        correct = 0
        test_loss_adv = 0
        correct_adv = 0
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
```

--------------------------------

### Train a model using deeprobust.image.netmodels.train_model

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/netmodels/train_model.html

Example usage of the train function to initialize and train a CNN on the MNIST dataset using CUDA.

```python
>>>import deeprobust.image.netmodels.train_model as trainmodel
>>>trainmodel.train('CNN', 'MNIST', 'cuda', 20)
```

--------------------------------

### ThermEncoding Module

Source: https://deeprobust.readthedocs.io/en/latest/source/deeprobust.image.defense.html

Implementation of Thermometer Encoding for resisting adversarial examples.

```APIDOC
## deeprobust.image.defense.ThermEncoding module

This is an implementation of Thermometer Encoding.

References
[1]| Buckman, Jacob, Aurko Roy, Colin Raffel, and Ian Goodfellow. “Thermometer encoding: One hot way to resist adversarial examples.” In International Conference on Learning Representations. 2018.

### Thermometer

Applies Thermometer Encoding to the input.

### Parameters
- **x** (Tensor) - Input tensor.
- **levels** (int) - The number of levels for encoding.
- **flattened** (bool, optional) - Whether to flatten the output. Defaults to False.

### one_hot

Performs one-hot encoding on the input.

### Parameters
- **x** (Tensor) - Input tensor.
- **levels** (int) - The number of levels for encoding.

### one_hot_to_thermometer

Converts one-hot encoding to Thermometer Encoding.

### Parameters
- **x** (Tensor) - Input tensor (one-hot encoded).
- **levels** (int) - The number of levels used during encoding.
- **flattened** (bool, optional) - Whether the input is flattened. Defaults to False.

### train

Training process for the Thermometer Encoding model.

### Parameters
- **model** (nn.Module) - The model to train.
- **device** (str) - The device to use for training ('cpu' or 'cuda').
- **train_loader** (DataLoader) - DataLoader for the training data.
- **optimizer** (Optimizer) - The optimizer for training.
- **epoch** (int) - The number of epochs to train.
```

--------------------------------

### Initialize FastGradientLayerOneTrainer

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Sets up the FastGradientLayerOneTrainer for adversarial training. Requires a Hamiltonian function, optimizer, inner iterations, sigma, and epsilon.

```python
LayerOneTrainer = FastGradientLayerOneTrainer(Hamiltonian_func, layer_one_optimizer,
                                                inner_iters, sigma, eps)
```

--------------------------------

### Main Training Pipeline Configuration

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Sets up hyperparameters, attack methods, and model components for the training loop.

```python
def main():
    num_epochs = 40
    val_interval = 1
    weight_decay = 5e-4

    inner_iters = 10
    K = 5
    sigma = 0.01
    eps = 0.3
    lr = 1e-2
    momentum = 0.9
    create_optimizer = SGDOptimizerMaker(lr =1e-2 / K, momentum = 0.9, weight_decay = weight_decay)

    create_lr_scheduler = PieceWiseConstantLrSchedulerMaker(milestones = [30, 35, 39], gamma = 0.1)

    create_loss_function = None

    create_attack_method = None

    create_evaluation_attack_method = IPGDAttackMethodMaker(eps = 0.3, sigma = 0.01, nb_iters = 40, norm = np.inf,
                                      mean=torch.tensor(np.array([0]).astype(np.float32)[np.newaxis, :, np.newaxis, np.newaxis]),
                                      std=torch.tensor(np.array([1]).astype(np.float32)[np.newaxis, :, np.newaxis, np.newaxis]))

    parser = argparse.ArgumentParser()

    parser.add_argument('--model_dir',default = "./trained_models")
    parser.add_argument('--resume', default=None, type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
    parser.add_argument('-b', '--batch_size', default=256, type=int,
                    metavar='N', help='mini-batch size')
    parser.add_argument('-d', type=int, default=0, help='Which gpu to use')
    parser.add_argument('-adv_coef', default=1.0, type = float,
                        help = 'Specify the weight for adversarial loss')
    parser.add_argument('--auto-continue', default=False, action = 'store_true',
                        help = 'Continue from the latest checkpoint')
    args = parser.parse_args()

    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    net = YOPOCNN.Net()
    net.to(DEVICE)
    criterion = CrossEntropyWithWeightPenlty(net.other_layers, DEVICE, weight_decay)#.to(DEVICE)
    optimizer = create_optimizer(net.other_layers.parameters())
    lr_scheduler = create_lr_scheduler(optimizer)

    Hamiltonian_func = Hamiltonian(net.layer_one, weight_decay)
    layer_one_optimizer = optim.SGD(net.layer_one.parameters(), lr = lr_scheduler.get_lr()[0], momentum=0.9, weight_decay=5e-4)
```

--------------------------------

### GET /utils/torch_accuracy

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Calculates the top-k accuracy for a given model output and target labels.

```APIDOC
## GET /utils/torch_accuracy

### Description
Calculates the top-k accuracy of model predictions compared to target labels.

### Parameters
#### Request Body
- **output** (torch.Tensor) - Required - The model output predictions.
- **target** (torch.Tensor) - Required - The ground truth labels.
- **topk** (tuple) - Optional - The top-k values to calculate (default: (1,)).

### Response
#### Success Response (200)
- **accuracy** (List[torch.Tensor]) - A list of accuracy values for each k in topk.
```

--------------------------------

### Initialize ResNet Variants

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/netmodels/resnet.html

Helper functions to instantiate specific ResNet architectures.

```python
def ResNet18():
    return Net(BasicBlock, [2,2,2,2])

def ResNet34():
    return Net(BasicBlock, [3,4,6,3])

def ResNet50():
    return Net(Bottleneck, [3,4,6,3])

def ResNet101():
    return Net(Bottleneck, [3,4,23,3])

def ResNet152():
    return Net(Bottleneck, [3,8,36,3])
```

--------------------------------

### Evaluate Model Robustness

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Evaluates a model's accuracy on clean and adversarial examples over an epoch.

```python
def eval_one_epoch(net, batch_generator,  DEVICE=torch.device('cuda:0'), AttackMethod = None):
    net.eval()
    pbar = tqdm(batch_generator)
    clean_accuracy = AvgMeter()
    adv_accuracy = AvgMeter()

    pbar.set_description('Evaluating')
    for (data, label) in pbar:
        data = data.to(DEVICE)
        label = label.to(DEVICE)

        with torch.no_grad():
            pred = net(data)
            acc = torch_accuracy(pred, label, (1,))
            clean_accuracy.update(acc[0].item())

        if AttackMethod is not None:
            adv_inp = AttackMethod.attack(net, data, label)

            with torch.no_grad():
                pred = net(adv_inp)
                acc = torch_accuracy(pred, label, (1,))
                adv_accuracy.update(acc[0].item())

        pbar_dic = OrderedDict()
        pbar_dic['CleanAcc'] = '{:.2f}'.format(clean_accuracy.mean)
        pbar_dic['AdvAcc'] = '{:.2f}'.format(adv_accuracy.mean)

        pbar.set_postfix(pbar_dic)

        adv_acc = adv_accuracy.mean if AttackMethod is not None else 0
    return clean_accuracy.mean, adv_acc
```

--------------------------------

### Initialize MultiStepLR Scheduler

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Initializes a MultiStepLR scheduler for adjusting learning rates at specified milestones. Requires an optimizer and a list of milestones.

```python
lyaer_one_optimizer_lr_scheduler = optim.lr_scheduler.MultiStepLR(layer_one_optimizer,
                                                                    milestones = [15, 19], gamma = 0.1)
```

--------------------------------

### Generate Adversarial Data

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/pgdtraining.html

Uses the PGD adversary to generate adversarial examples for training or testing.

```python
    def adv_data(self, data, output, ep = 0.3, num_steps = 10, perturb_step_size = 0.01):
        """
        Generate input(adversarial) data for training.
        """

        adversary = PGD(self.model)
        data_adv = adversary.generate(data, output.flatten(), epsilon = ep, num_steps = num_steps, step_size = perturb_step_size)
        output = self.model(data_adv)

        return data_adv, output
```

--------------------------------

### Create Training and Validation Datasets

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Creates training and validation datasets using a utility function. Accepts batch size as an argument.

```python
ds_train = utils.create_train_dataset(args.batch_size)
    ds_val = utils.create_test_dataset(args.batch_size)
```

--------------------------------

### Initialize Fast Adversarial Training

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/fast.html

Initializes the Fast defense mechanism with a given model and device. It checks for CUDA availability and sets the device accordingly.

```python
class Fast(BaseDefense):
    def __init__(self, model, device):
        if not torch.cuda.is_available():
            print('CUDA not availiable, using cpu...')
            self.device = 'cpu'
        else:
            self.device = device

        self.model = model
```

--------------------------------

### Adversarial Data Generation

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/pgdtraining.html

Generates adversarial examples using the PGD (Projected Gradient Descent) attack method.

```APIDOC
## adv_data(data, output, ep, num_steps, perturb_step_size)

### Description
Generates adversarial input data using the PGD algorithm and returns the perturbed data along with the model's output on that data.

### Parameters
- **data** (tensor) - Required - The input data to perturb.
- **output** (tensor) - Required - The target labels or output for the attack.
- **ep** (float) - Optional - The epsilon value for the perturbation magnitude.
- **num_steps** (int) - Optional - The number of steps for the PGD attack.
- **perturb_step_size** (float) - Optional - The step size for each perturbation iteration.
```

--------------------------------

### Initialize TRADES Defense

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/trades.html

Instantiates the TRADES class with a target model and device configuration.

```python
class TRADES(BaseDefense):
    """TRADES.
    """

    def __init__(self, model, device = 'cuda'):
        if not torch.cuda.is_available():
            print('CUDA not available, using cpu...')
            self.device = 'cpu'
        else:
            self.device = device
        
        self.model = model.to(self.device)
```

--------------------------------

### Checkpoint and File Utilities

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Functions for loading model checkpoints, creating symbolic links, and modifying system paths.

```python
def load_checkpoint(file_name, net = None, optimizer = None, lr_scheduler = None):
    if os.path.isfile(file_name):
        print("=> loading checkpoint '{}'".format(file_name))
        check_point = torch.load(file_name)
        if net is not None:
            print('Loading network state dict')
            net.load_state_dict(check_point['state_dict'])
        if optimizer is not None:
            print('Loading optimizer state dict')
            optimizer.load_state_dict(check_point['optimizer_state_dict'])
        if lr_scheduler is not None:
            print('Loading lr_scheduler state dict')
            lr_scheduler.load_state_dict(check_point['lr_scheduler_state_dict'])

        return check_point['epoch']
    else:
        print("=> no checkpoint found at '{}'".format(file_name))


def make_symlink(source, link_name):
    if os.path.exists(link_name):
        #print("Link name already exist! Removing '{}' and overwriting".format(link_name))
        os.remove(link_name)
    if os.path.exists(source):
        os.symlink(source, link_name)
        return
    else:
        print('Source path not exists')
    #print('SymLink Wrong!')

def add_path(path):
    if path not in sys.path:
        print('Adding {}'.format(path))
        sys.path.append(path)
```

--------------------------------

### Train a model using train_model utility

Source: https://deeprobust.readthedocs.io/en/latest/source/deeprobust.image.netmodels.html

Example usage of the train function to train a CNN model on the MNIST dataset using CUDA.

```python
>>>import deeprobust.image.netmodels.train_model as trainmodel >>>trainmodel.train(‘CNN’, ‘MNIST’, ‘cuda’, 20)
```

--------------------------------

### FGSMtraining.generate

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/fgsmtraining.html

FGSM adversarial training process.

```APIDOC
## POST /generate

### Description
Initiates the FGSM adversarial training process using the provided data loaders. The model is trained for a specified number of epochs, and optionally saved.

### Method
POST

### Endpoint
/generate

### Parameters
#### Query Parameters
- **train_loader** (loader) - Required - The training data loader.
- **test_loader** (loader) - Required - The testing data loader.

#### Request Body
- **save_dir** (string) - Optional - Directory to save the trained model. Defaults to "defense_models".
- **save_model** (boolean) - Optional - Whether to save the trained model. Defaults to True.
- **save_name** (string) - Optional - Filename for the saved model. Defaults to "mnist_fgsmtraining_0.2.pt".
- **epsilon** (float) - Optional - The attack perturbation constraint. Defaults to 0.2.
- **epoch_num** (integer) - Optional - The number of training epochs. Defaults to 50.
- **lr_train** (float) - Optional - The learning rate for training. Defaults to 0.005.
- **momentum** (float) - Optional - Momentum for the optimizer. Defaults to 0.1.

### Response
#### Success Response (200)
- **model** (object) - The trained model.

#### Response Example
```json
{
  "model": "trained_model_object"
}
```
```

--------------------------------

### LIDClassifier Training

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/LIDclassifier.html

This function trains the LID classifier model. It iterates through the training data, computes adversarial examples, calculates the loss, and updates the model weights.

```APIDOC
## POST /train

### Description
Trains the LID classifier model using the provided data loader, optimizer, and device.

### Method
POST

### Endpoint
/train

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **device** (string) - Required - The device to use for training (e.g., 'cpu', 'cuda').
- **train_loader** (DataLoader) - Required - The data loader for the training dataset.
- **optimizer** (Optimizer) - Required - The optimizer for training.
- **epoch** (int) - Required - The current epoch number.

### Request Example
```json
{
  "device": "cuda",
  "train_loader": "<DataLoader object>",
  "optimizer": "<Optimizer object>",
  "epoch": 1
}
```

### Response
#### Success Response (200)
- **message** (string) - Indicates successful training completion.

#### Response Example
```json
{
  "message": "Training epoch 1 completed successfully."
}
```
```

--------------------------------

### Load Checkpoint for Resuming Training

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Loads a checkpoint to resume training from a saved state. Checks if the resume path is a valid file and updates the current epoch.

```python
if args.resume is not None and os.path.isfile(args.resume):
        now_epoch = load_checkpoint(args.resume, net, optimizer,lr_scheduler)
```

--------------------------------

### PGD Adversarial Training Implementation

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/pgdtraining.html

The PGDtraining class inherits from BaseDefense and provides the core logic for adversarial training, including model generation and parameter configuration.

```python
"""
This is an implementation of pgd adversarial training.
References
----------
..[1]Mądry, A., Makelov, A., Schmidt, L., Tsipras, D., & Vladu, A. (2017).
Towards Deep Learning Models Resistant to Adversarial Attacks. stat, 1050, 9.
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import torch.nn.functional as F

import numpy as np
from PIL import Image
from deeprobust.image.attack.pgd import PGD
from deeprobust.image.netmodels.CNN import Net
from deeprobust.image.defense.base_defense import BaseDefense



[docs]class PGDtraining(BaseDefense):
    """
    PGD adversarial training.

    """


    def __init__(self, model, device):
        if not torch.cuda.is_available():
            print('CUDA not availiable, using cpu...')
            self.device = 'cpu'
        else:
            self.device = device

        self.model = model


[docs]    def generate(self, train_loader, test_loader, **kwargs):
        """Call this function to generate robust model.

        Parameters
        ----------
        train_loader :
            training data loader
        test_loader :
            testing data loader
        kwargs :
            kwargs
        """
        self.parse_params(**kwargs)

        torch.manual_seed(100)
        device = torch.device(self.device)

        optimizer = optim.Adam(self.model.parameters(), self.lr)
        scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[75, 100], gamma = 0.1)
        save_model = True
        for epoch in range(1, self.epoch + 1):
            print('Training epoch: ', epoch, flush = True)
            self.train(self.device, train_loader, optimizer, epoch)
            self.test(self.model, self.device, test_loader)

            if (self.save_model and epoch % self.save_per_epoch == 0):
                if os.path.isdir(str(self.save_dir)):
                    torch.save(self.model.state_dict(), str(self.save_dir) + self.save_name + '_epoch' + str(epoch) + '.pth')
                    print("model saved in " + str(self.save_dir))
                else:
                    print("make new directory and save model in " + str(self.save_dir))
                    os.mkdir('./' + str(self.save_dir))
                    torch.save(self.model.state_dict(), str(self.save_dir) + self.save_name + '_epoch' + str(epoch) + '.pth')

            scheduler.step()

        return self.model



[docs]    def parse_params(self,
                     epoch_num = 100,
                     save_dir = "./defense_models",
                     save_name = "mnist_pgdtraining_0.3",
                     save_model = True,
                     epsilon = 8.0 / 255.0,
                     num_steps = 10,
                     perturb_step_size = 0.01,
                     lr = 0.1,
                     momentum = 0.1,
                     save_per_epoch = 10):
        """Parameter parser.

        Parameters
        ----------
        epoch_num : int
            epoch
        save_dir : str
            model dir
        save_name : str
            model name
        save_model : bool
            Whether to save model
        epsilon : float
            attack constraint
        num_steps : int
            PGD attack iteration time
        perturb_step_size : float
            perturb step size
        lr : float
            learning rate for adversary training process
        momentum : float
            momentum for optimizor
        """
        self.epoch = epoch_num
        self.save_model = True
        self.save_dir = save_dir
        self.save_name = save_name
        self.epsilon = epsilon
        self.num_steps = num_steps
        self.perturb_step_size = perturb_step_size
        self.lr = lr
        self.momentum = momentum
        self.save_per_epoch = save_per_epoch



[docs]    def train(self, device, train_loader, optimizer, epoch):
        """
        training process.

        Parameters
        ----------
        device :
            device
        train_loader :
            training data loader
        optimizer :
            optimizer
        epoch :
            training epoch
        """

        self.model.train()
        correct = 0
        bs = train_loader.batch_size
        #scheduler = StepLR(optimizer, step_size = 10, gamma = 0.5)
        scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones = [70], gamma = 0.1)
        for batch_idx, (data, target) in enumerate(train_loader):

            optimizer.zero_grad()

            data, target = data.to(device), target.to(device)

            data_adv, output = self.adv_data(data, target, ep = self.epsilon, num_steps = self.num_steps, perturb_step_size = self.perturb_step_size)
            loss = self.calculate_loss(output, target)

            loss.backward()
            optimizer.step()

            pred = output.argmax(dim = 1, keepdim = True)
```

--------------------------------

### Train and Test Utilities

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/netmodels/CNN_multilayer.html

Provides standard training and testing loops for the CNN model using PyTorch data loaders and optimizers.

```python
[docs]def train(model, device, train_loader, optimizer, epoch):
    """train.

    Parameters
    ----------
    model :
        model
    device :
        device
    train_loader :
        train_loader
    optimizer :
        optimizer
    epoch :
        epoch
    """
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()

        #print every 10
        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))




[docs]def test(model, device, test_loader):
    """test.

    Parameters
    ----------
    model :
        model
    device :
        device
    test_loader :
        test_loader
    """
    model.eval()

    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
```

--------------------------------

### Fast Testing Loop

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/fast.html

Evaluates the model's performance on the test set, calculating both clean and adversarial accuracy. It uses the FGSM method to generate adversarial examples for testing.

```python
def test(self, model, device, test_loader):
        """
        Testing process.

        """
        model.eval()

        test_loss = 0
        correct = 0
        test_loss_adv = 0
        correct_adv = 0
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)

            # print clean accuracy
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim = 1, keepdim = True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

            # print adversarial accuracy
            data_adv, output_adv = self.adv_data(data, target, ep = self.epsilon)

            test_loss_adv += self.calculate_loss(output_adv, target, redmode = 'sum').item()  # sum up batch loss
            pred_adv = output_adv.argmax(dim = 1, keepdim = True)  # get the index of the max log-probability
            correct_adv += pred_adv.eq(target.view_as(pred_adv)).sum().item()

        test_loss /= len(test_loader.dataset)
        test_loss_adv /= len(test_loader.dataset)

        print('\nTest set: Clean loss: {:.3f}, Clean Accuracy: {}/{} ({:.0f}%)\n'.format(

```

--------------------------------

### Implement Model Training Loop

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/netmodels/resnet.html

Executes a single training epoch using a PyTorch model, optimizer, and data loader. Requires the model to be in training mode and data to be moved to the specified device.

```python
def train(model, device, train_loader, optimizer, epoch):
    model.train()

    # lr = util.adjust_learning_rate(optimizer, epoch, args) # don't need it if we use Adam

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = torch.tensor(data).to(device), torch.tensor(target).to(device)
        optimizer.zero_grad()
        output = model(data)
        # loss = F.nll_loss(output, target)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
```

--------------------------------

### Fast Adversarial Training Class

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/fast.html

The Fast class implements an adversarial training variant. It inherits from BaseDefense and provides methods for generating adversarial examples, training the model, and testing its performance.

```APIDOC
## Class: Fast

### Description
Implements an adversarial training variant based on the paper "Fast is better than free: Revisiting adversarial training."

### Methods

#### `__init__(self, model, device)`
Initializes the Fast defense object.

- **model** (torch.nn.Module) - The model to be defended.
- **device** (str) - The device to run the computations on ('cpu' or 'cuda').

#### `generate(self, train_loader, test_loader, **kwargs)`
Generates adversarial training data and trains the model.

- **train_loader** (DataLoader) - DataLoader for the training dataset.
- **test_loader** (DataLoader) - DataLoader for the testing dataset.
- **kwargs** - Additional parameters for `parse_params`.

Returns:
- **torch.nn.Module**: The trained model.

#### `parse_params(self, save_dir='defense_models', save_model=True, save_name='fast_mnist_fgsmtraining_0.2.pt', epsilon=0.2, epoch_num=30, lr_train=0.005, momentum=0.1)`
Sets parameters for the fast training process.

- **save_dir** (str) - Directory to save the trained model.
- **save_model** (bool) - Whether to save the trained model.
- **save_name** (str) - Filename for the saved model.
- **epsilon** (float) - Epsilon value for FGSM attack.
- **epoch_num** (int) - Number of training epochs.
- **lr_train** (float) - Learning rate for the training optimizer.
- **momentum** (float) - Momentum for the optimizer.

#### `train(self, device, train_loader, optimizer, epoch)`
Performs one epoch of training.

- **device** (str) - The device to use for training.
- **train_loader** (DataLoader) - DataLoader for the training dataset.
- **optimizer** (torch.optim.Optimizer) - The optimizer for training.
- **epoch** (int) - The current epoch number.

#### `test(self, model, device, test_loader)`
Evaluates the model on the test set, both clean and adversarial.

- **model** (torch.nn.Module) - The model to evaluate.
- **device** (str) - The device to use for evaluation.
- **test_loader** (DataLoader) - DataLoader for the test dataset.

### References
- [1] Wong, Eric, Leslie Rice, and J. Zico Kolter. "Fast is better than free: Revisiting adversarial training." arXiv preprint arXiv:2001.03994 (2020).
```

--------------------------------

### Get Local Intrinsic Dimensionality (LID)

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/LIDclassifier.html

Estimates the Local Intrinsic Dimensionality (LID) for clean, noisy, and adversarial data samples. This function is crucial for analyzing the characteristics of different data types within the model's feature space.

```APIDOC
## POST /get_lid

### Description
Estimates the Local Intrinsic Dimensionality (LID) for clean, noisy, and adversarial data samples. It returns the LID artifacts and corresponding labels.

### Method
POST

### Endpoint
/get_lid

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **model** (object) - Required - The trained model to use for activation extraction.
- **X_test** (numpy.ndarray) - Required - Clean test data.
- **X_test_noisy** (numpy.ndarray) - Required - Noisy test data.
- **X_test_adv** (numpy.ndarray) - Required - Adversarial test data.
- **k** (int) - Required - The number of nearest neighbors to consider for LID estimation.
- **batch_size** (int) - Required - The batch size for processing data during LID estimation.

### Request Example
```json
{
  "model": "<Model object>",
  "X_test": "<numpy.ndarray>",
  "X_test_noisy": "<numpy.ndarray>",
  "X_test_adv": "<numpy.ndarray>",
  "k": 20,
  "batch_size": 100
}
```

### Response
#### Success Response (200)
- **artifacts** (numpy.ndarray) - The estimated LID values.
- **labels** (numpy.ndarray) - The labels corresponding to the LID artifacts.

#### Response Example
```json
{
  "artifacts": "<numpy.ndarray>",
  "labels": "<numpy.ndarray>"
}
```
```

--------------------------------

### POST /utils/load_checkpoint

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Loads a model checkpoint from a specified file path, including network state, optimizer, and scheduler.

```APIDOC
## POST /utils/load_checkpoint

### Description
Loads a saved checkpoint file to restore the state of the network, optimizer, and learning rate scheduler.

### Parameters
#### Request Body
- **file_name** (string) - Required - Path to the checkpoint file.
- **net** (torch.nn.Module) - Optional - The network model to load state into.
- **optimizer** (torch.optim.Optimizer) - Optional - The optimizer to load state into.
- **lr_scheduler** (torch.optim.lr_scheduler) - Optional - The scheduler to load state into.

### Response
#### Success Response (200)
- **epoch** (int) - The epoch number retrieved from the checkpoint.
```

--------------------------------

### Main Execution Block

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/YOPO.html

Standard Python entry point that calls the main function when the script is executed.

```python
if __name__ == "__main__":
    main()
```

--------------------------------

### Execute LID Characterization

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/LIDclassifier.html

Main entry point to extract LID features and labels from test data.

```python
if __name__ == "__main__":

    batch_size = 100
    k_nearest = 20

    #get LID characters
    characters, labels = get_lid(model, X_test, X_test_noisy, X_test_adv, k_nearest, batch_size)
    data = np.concatenate((characters, labels), axis = 1)
```

--------------------------------

### VGG Network Architecture and Training Utilities

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/netmodels/vgg.html

Contains the VGG class definition with dynamic layer generation, along with standard training and evaluation loops for PyTorch models.

```python
"""
This is an implementation of VGG net.

Reference
---------
..[1]Simonyan, Karen, and Andrew Zisserman. "Very deep convolutional networks for large-scale image recognition." arXiv preprint arXiv:1409.1556 (2014).
..[2]Original implementation: https://github.com/kuangliu/pytorch-cifar
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

cfg = {
    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}



[docs]class VGG(nn.Module):
    """VGG.
    """

    def __init__(self, vgg_name):
        super(VGG, self).__init__()
        self.features = self._make_layers(cfg[vgg_name])
        self.classifier = nn.Linear(512, 10)

    def forward(self, x):
        out = self.features(x)
        out = out.view(out.size(0), -1)
        out = self.classifier(out)
        return out

    def _make_layers(self, cfg):
        layers = []
        in_channels = 3
        for x in cfg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
                           nn.BatchNorm2d(x),
                           nn.ReLU(inplace=True)]
                in_channels = x
        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
        return nn.Sequential(*layers)




[docs]def test(model, device, test_loader):
    """test.

    Parameters
    ----------
    model :
        model
    device :
        device
    test_loader :
        test_loader
    """
    model.eval()

    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            #test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss

            test_loss += F.cross_entropy(output, target)
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))



[docs]def train(model, device, train_loader, optimizer, epoch):
    """train.

    Parameters
    ----------
    model :
        model
    device :
        device
    train_loader :
        train_loader
    optimizer :
        optimizer
    epoch :
        epoch
    """
    model.train()

    # lr = util.adjust_learning_rate(optimizer, epoch, args) # don't need it if we use Adam

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = torch.tensor(data).to(device), torch.tensor(target).to(device)
        optimizer.zero_grad()
        output = model(data)
        # loss = F.nll_loss(output, target)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()/data.shape[0]))
```

--------------------------------

### Parse Fast Training Parameters

Source: https://deeprobust.readthedocs.io/en/latest/_modules/deeprobust/image/defense/fast.html

Sets and parses parameters for the Fast adversarial training process. This includes saving configurations, epsilon for FGSM, epoch number, learning rate, and momentum.

```python
def parse_params(self,
                     save_dir = "defense_models",
                     save_model = True,
                     save_name = "fast_mnist_fgsmtraining_0.2.pt",
                     epsilon = 0.2,
                     epoch_num = 30,
                     lr_train = 0.005,
                     momentum = 0.1):
        # """
        # Set parameters for fast training.
        # """
        self.save_model = True
        self.save_dir = save_dir
        self.save_name = save_name
        self.epsilon = epsilon
        self.epoch_num = epoch_num
        self.lr_train = lr_train
        self.momentum = momentum
```