Example of using the advertrain library

[1]:
import torch
from torch import nn, optim
from pathlib import Path
import data_all
from robustAI.advertrain.models import ConvNet, ResNet, ConvNetDropblock, ResNetDropblock
from robustAI.advertrain.training.classical_training import ClassicalTraining
from robustAI.advertrain.training.adversarial_training import AdversarialTraining
from robustAI.advertrain.training.autoattack_training import AutoAttackTraining
from robustAI.advertrain.training.fire_training import FIRETraining
from robustAI.advertrain.training.trades_training import TRADESTraining
from robustAI.advertrain.transforms import DataTransformations
import sys
import os
import argparse
/home/jovyan/Maturation/env-robustml/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
[2]:
# Set here the path to the directory containing the folder c00 of the example dataset. it can be downloaded
# here http://minio-storage.apps.confianceai-public.irtsysx.fr/ml-models/robust-ml-dataset.zip

ROOT_DIR = "/home/jovyan/Maturation/robust-ai/examples/dataset/c00"

DEVICE = "cuda:1" # Select here your cuda device
ARCHITECTURE = "resnetdropblock" # Choose your architecture
LEARNING_RATE = 1e-4
MODELS_PATH = "results/" # Set here the path where trained model will be stored
[3]:
transformer = DataTransformations()
train_transform = transformer.get_train_transforms()
test_transform = transformer.get_test_transforms()
[4]:
path_data = "learning_data/"

train_dataloader1, val_dataloader1 = data_all.get_train_dataloader(
    ROOT_DIR,
    path_data+"train_dataset_all.csv",
    path_data+"val_dataset_all.csv",
    batch_size=16,
    train_transform=train_transform,
    test_transform=test_transform,
)
[5]:
if ARCHITECTURE == "convnet":
    model = ConvNet(DEVICE)
elif ARCHITECTURE == "resnet":
    model = ResNet(DEVICE)
elif ARCHITECTURE == "convnetdropblock":
    model = ConvNetDropblock(DEVICE)
elif ARCHITECTURE == "resnetdropblock":
    model = ResNetDropblock(DEVICE)
else:
    raise ValueError("Unknown architecture")
[6]:
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
criterion = nn.CrossEntropyLoss(weight=torch.Tensor([1, 1]).to(DEVICE))
[7]:
if not torch.cuda.is_available() or DEVICE =="cpu":
        print("ERROR : Learning should be done on a system with GPU")
        sys.exit()

print(torch.cuda.current_device())
print("GPU device used ", DEVICE)
0
GPU device used  cuda:1

Train class

Choose here the training method you want to use. Uncomment only the corresponding cells. For example to use FireTraining method, uncomment cells [14] and [15]

ClassicalTraining

[8]:
# METHODE="vanilla"
[9]:
# trainer = ClassicalTraining(model, optimizer, criterion, DEVICE)

AdversarialTraining

[10]:
# METHODE = "adversarial_training"
# EPSILON = 8 / 255
[11]:
# trainer = AdversarialTraining(model, optimizer, criterion, DEVICE, EPSILON)

AutoAttackTraining

[12]:
# METHODE = "autoattack_training"
# EPSILON = 8 / 255
[13]:
# trainer = AutoAttackTraining(model, optimizer, criterion, DEVICE, "ce", EPSILON)

FireTraining

[14]:
METHODE = "fire_training"
EPSILON = 8 / 255
BETA = 1
[15]:
trainer = FIRETraining(model, optimizer, DEVICE, EPSILON, BETA)

TradesTraining

[16]:
# METHODE = "trades_training"
# EPSILON = 8 / 255
# BETA = 1
[17]:
# trainer = TRADESTraining(model, optimizer, DEVICE, EPSILON, BETA)

Fit the model

[18]:
path = os.path.join(MODELS_PATH,
                    f"{ARCHITECTURE}"+"_"+f"{METHODE}")

if not os.path.exists(path):
    try:
        os.makedirs(path)
    except FileExistsError:
        pass

print(f"Using '{path}' directory")

trainer.fit(epochs=2, train_dataloader=train_dataloader1, val_dataloader=val_dataloader1, patience=2, checkpoint=path)
Using 'results/resnetdropblock_fire_training' directory
Epochs 1/2 : Training: 100%|██████████| 750/750 [01:55<00:00,  6.49it/s]
Validation: 100%|██████████| 33/33 [00:06<00:00,  5.41it/s]
Epoch 1/2
Train Loss:  nan, Acc:  0.502, Recall:  0.031, Precision:  0.573,F1 Score:  0.058
Validation Loss:  nan, Acc:  0.544, Recall:  0.000, Precision:  0.000,F1 Score:  0.000
Epochs 2/2 : Training: 100%|██████████| 750/750 [01:55<00:00,  6.51it/s]
Validation: 100%|██████████| 33/33 [00:06<00:00,  5.16it/s]
Epoch 2/2
Train Loss:  nan, Acc:  0.499, Recall:  0.000, Precision:  0.000,F1 Score:  0.000
Validation Loss:  nan, Acc:  0.544, Recall:  0.000, Precision:  0.000,F1 Score:  0.000

[18]:
{'loss': [nan, nan],
 'acc': [0.5024999976158142, 0.49941667914390564],
 'val_loss': [nan, nan],
 'val_acc': [0.5435606241226196, 0.5435606241226196],
 'optimizer': 'Adam (\nParameter Group 0\n    amsgrad: False\n    betas: (0.9, 0.999)\n    capturable: False\n    differentiable: False\n    eps: 1e-08\n    foreach: None\n    fused: False\n    lr: 0.0001\n    maximize: False\n    weight_decay: 0\n)',
 'loss_func': 'None',
 'epochs': 2,
 'patience': 2}
[ ]:

[ ]: