Finetuning/Tune.py

# Finetuning a selected model 
# on a selected dataset 
# using selected parameters

from torch.utils.data import DataLoader
from sklearn.metrics import classification_report
import SetUp
#from Data import *
# from datasets.Casia import *
#from IdentitySubset import IdentitySubset
from sets.Data import *
from sets.IdentitySubset import IdentitySubset
# models
from architectures.Model import Model, Architecture

from unlearning.LinearFiltration import LinearFiltration
from unlearning.CertifiedRemoval import CertifiedRemoval
from unlearning.WeightFiltration import WeightFiltration

import Util
# WeightFiltration, CertifiedRemoval

# numbre of classes 
CLASS_SIZE = 20
# batch 
BATCH_SIZE = 16

# size of images per class trainset + testset
# 30 works best, more than that and we dont have enough data
SAMPLE_SIZE = 30

# this is then (full_sample - test_sample)
TRAINING_SMPLE = 27

# learning rate
LR_RATE = 0.0001
EPOCHS = 20

# depends on model architecture
# ResNet, DenseNet = 224
# Inception = 299
RESOLUTION = 224

# model architecture options are 
#  - RESNET18
#  - RESNET50
#  - DENSENET121
#  - INCEPTION
#  - GOOGLENET
#  - EFFICIENTNET
#  - SHUFFLENET
arch = Architecture.RESNET50

# DATA PREPARATION
# load data set and prepare
dataset_name = Set_Name.CELEBA
set = Set_Name.CELEBA

dataset = get_set(set_name=dataset_name)
print(f"> {dataset.__class__.__name__} dataset loaded")
# select identities for experiment
#selected_identities = select_ids(
#    dataset = dataset,
#    sample_size = SAMPLE_SIZE,
#    class_size = CLASS_SIZE
#    )

# this selects the top 50 based on sample size
# that way repeated calls return the same classes
selected_identities = select_top_ids(
    dataset=dataset,
    class_size=CLASS_SIZE
)

print(f'> Selected {CLASS_SIZE} random identity classes from CelebA dataset.')
print(f'> A class has {TRAINING_SMPLE} train and {SAMPLE_SIZE-TRAINING_SMPLE} test samples')

# split class images to train/test indices
train_indices, test_indices = get_indices(
    dataset = dataset, 
    identities = selected_identities, 
    split_at = TRAINING_SMPLE,
    size= SAMPLE_SIZE
)

# helps map class id to index
id_map = {old_id: new_id for new_id, old_id in enumerate(selected_identities)}

# we remap identities because crossEntropyLoss requires in indices 0 -> (n-1)
# where n = class size.
tr_transform = train_transform(RESOLUTION)
train_data = IdentitySubset(
    dataset=dataset, 
    indices=train_indices, 
    id_mapping=id_map,
    transform=tr_transform)

train_loader = DataLoader(
    train_data, 
    batch_size = BATCH_SIZE, 
    shuffle = True)

print(f"> Total training images: {len(train_data)}")

print(f'> Constants : Classes = {CLASS_SIZE}, Batch = {BATCH_SIZE}, epochs = {EPOCHS}')

# MODEL PREPARATION
# cuda if exists (it does here)
device = SetUp.get_device()


for i in range(0,CLASS_SIZE):
    FORGET_CLASS_IDX = i
    # Create model using Factory
    model = Model.create(
        arch = arch, 
        device = device, 
        size = CLASS_SIZE)

    # we may need to load existing model or finetune
    model.train(
            epochs = EPOCHS, 
            loader = train_loader, 
            rate = LR_RATE)

        # save.
    model.save(filename=arch.name.lower())


    # done tuning

    # EVALUATE
    te_transform = test_transform(RESOLUTION)
    # Testing
    test_data = IdentitySubset(
        dataset = dataset, 
        indices=test_indices, 
        id_mapping=id_map,
        transform=te_transform)

    test_loader = DataLoader(
        test_data, 
        batch_size=BATCH_SIZE, 
        shuffle=False)

    print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")

    # Evaluate
    current_mode = "Finetuned"
    accuracy, report_dict = model.evaluate(
        loader = test_loader,
        mode=current_mode
    )

    Util._log_to_csv(
        arch=model.__class__.__name__, 
        mode = current_mode, 
        accuracy=accuracy, 
        report_dict=report_dict, 
        strategy="base"
    )

    # unlearning algorithms
    linear_filtration = LinearFiltration(target_class_idx=FORGET_CLASS_IDX)
    #filtration.apply(reloaded.model)

    weight_filtration = WeightFiltration(num_classes = CLASS_SIZE,target_class_idx=FORGET_CLASS_IDX)
    #weight_filtration.apply(reloaded.model)

    certified_removal = CertifiedRemoval(removal_bound=0.05, epsilon=0.5, l2_reg=0.1)
    #certified_removal.apply(reloaded.model)

    # to be unlearned
    forget_train_loader, retain_train_loader = get_unlearning_loaders(
        dataset=train_data, 
        forget_class_idx=FORGET_CLASS_IDX, 
        batch_size=BATCH_SIZE
    )

    # to evaluate
    forget_test_loader, retain_test_loader = get_unlearning_loaders(
            dataset=test_data, 
            forget_class_idx=FORGET_CLASS_IDX, 
            batch_size=BATCH_SIZE
        )


    strategies = [linear_filtration, weight_filtration, certified_removal]
    #strategies = [linear_filtration]
    for strategy in strategies:
        # test again
        reloaded = Model.create(
            arch=arch,
            device = device,
            size = CLASS_SIZE
            )
        reloaded.load(arch = arch)
        print("fine tunned model loaded")
        # reloaded.evaluate(
        #   loader = test_loader
        #)

        # Unlearning
        # train loaders passed here
        strategy.apply(reloaded.model, forget_train_loader, retain_train_loader)
        # Performance Analysis
        strategy_in_use = strategy.__class__.__name__

        # evaluation on retain Test_set
        current_mode = "retain"
        print("\n--- Performance on Retained Classes")
        accuracy, report_dict = reloaded.evaluate(loader=retain_test_loader, mode=current_mode)

        Util._log_to_csv(
            arch=reloaded.__class__.__name__, 
            mode = current_mode, 
            accuracy=accuracy, 
            report_dict=report_dict, 
            strategy=strategy_in_use
        )


        # evaluation on forget Test_set
        print("\n--- Performance on Forgotten Class")
        current_mode = "forget"
        accuracy, report_dict = reloaded.evaluate(loader=forget_test_loader,mode=current_mode)
        Util._log_to_csv(
            arch=reloaded.__class__.__name__, 
            mode = current_mode, 
            accuracy=accuracy, 
            report_dict=report_dict, 
            strategy=strategy_in_use
        )

        # evaluation on forget Train_set
        # we expect this to be equal or highr than accuracy on Forget Test_set
        current_mode = "forget_train"
        print("\n--- Performance on Forgotten Class (Train Set - Verifying Unlearning)")
        accuracy, report_dict = reloaded.evaluate(loader=forget_train_loader, mode=current_mode)
        Util._log_to_csv(
            arch= reloaded.__class__.__name__, 
            mode = current_mode, 
            accuracy=accuracy, 
            report_dict=report_dict, 
            strategy=strategy_in_use
        )