268 lines
7.4 KiB
Python
268 lines
7.4 KiB
Python
# 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 = 10
|
|
|
|
# depends on model architecture
|
|
# ResNet, DenseNet = 224
|
|
# Inception = 299
|
|
RESOLUTION = 224
|
|
|
|
FINETUNE = False # whether to fintune or just load finetuned model from dir
|
|
# 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,1):#CLASS_SIZE):
|
|
FORGET_CLASS_IDX = i
|
|
# Create model using Factory
|
|
|
|
model = None
|
|
|
|
if FINETUNE:
|
|
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.
|
|
file_name = f"{arch.name.lower}_{dataset_name.name.lower()}"
|
|
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"
|
|
if FINETUNE:
|
|
|
|
#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_index=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(
|
|
target_class_index=FORGET_CLASS_IDX,
|
|
s1=2,
|
|
s2=500,
|
|
unlearn_bs=2,
|
|
scale=100.0, # Drop scale to match lower s2 depth
|
|
std=0.00001)
|
|
#,removal_bound=0.05, epsilon=0.5, l2_reg=15)
|
|
#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 = [certified_removal]
|
|
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
|
|
#)
|
|
|
|
if not FINETUNE:
|
|
reloaded.evaluate(
|
|
loader = test_loader,
|
|
mode=current_mode
|
|
)
|
|
|
|
# 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
|
|
) |