main_quantMIA.py

import argparse
import os
import yaml
from utils.loader import SeedManager
import random as rnd
import time
from tqdm import tqdm

import train
from utils.evaluate import save_accuracy
from utils.modelutils import model_manager 
from quantize import main_AdaRound, main_OBC, main_BRECQ
from lira import attack

def load_global_config(config_path='config.yaml'):
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"Configuration file {config_path} not found.")
    
    with open(config_path, 'r') as file:
        config = yaml.safe_load(file)
    return config
    

def load_global_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--wquant', action='store_true', help='Weight quantization')
    parser.add_argument('--aquant', action='store_true', help='Activation quantization')
    parser.add_argument('--root', type=str)
    parser.add_argument('--seed', type=int, nargs='+')
    parser.add_argument('--legend', type=str)
    parser.add_argument('--n_shadows', type=int)
    parser.add_argument('--batch_size', type=int)
    parser.add_argument('--n_amples', type=int)

    parser.add_argument('--wbit', type=float, nargs='+', help='Weight bits')
    parser.add_argument('--abit', type=float, nargs='+', help='Activation bits')
    
    parser.add_argument('--set_last_layer', type=int)

    parser.add_argument('--skip-target', action='store_true')
    parser.add_argument('--skip-shadows', action='store_true')
    parser.add_argument('--skip-quant', action='store_true')
    parser.add_argument('--skip-attack', action='store_true')

    parser.add_argument('--abort-AdaRound', action='store_true')
    parser.add_argument('--abort-BRECQ', action='store_true')
    parser.add_argument('--abort-OBC', action='store_true')
    
    parser.add_argument('--key_params', type=str, nargs='+', help='Key parameters for the model')
    parser.add_argument('--cycles', type=int)
    parser.add_argument('--auto-quant', action='store_true', help='skip quantization if already exists')
    
    nargs = vars(parser.parse_args())
    args = load_global_config()
    for key, value in nargs.items():
        if value is not None:
            if type(value) == bool and not value: continue
            if 'skip' in key:
                args_key = key.replace('skip_', '')
                assert args_key in args['Global']['stage'], f"Key {args_key} not found in config"
                args['Global']['stage'][args_key] = not value
            elif 'abort' in key:
                args_key = key.replace('abort_', '')
                assert args_key in args['Global']['quant'], f"Key {args_key} not found in config"
                args['Global']['quant'][args_key] = not value
            else:
                assert key in args['Global'], f"Key {key} not found in config"
                args['Global'][key] = value
    return args            


def target_model_trainer(args, seed_manager):
    model_id = time.strftime(r"%d%H%M%S")
    
    train_args = {
        'root': args['Global']['root'],
        'legend': args['Global']['legend'],
        'n_samples': args['Global']['n_samples'],
        'batch_size': args['Global']['batch_size'],
        'model': args['Global']['model'],
        'seed': seed_manager,
        'lr': args['Train']['lr'],
        'epochs': args['Train']['epochs'],
        'n_queries': args['Global']['n_queries'],
        'pkeep': args['Train']['pkeep'],
        'n_shadows': None,
        'id': model_id
    }
    
    key_dict, remarks = get_key_params('Train', args, train_args)      
    accuracy = train.train_model(train_args, key_dict)      
    save_accuracy(train_args['id'], str(key_dict), str(remarks), accuracy)
    
def shadow_model_trainer(args, seed_manager):
    shadow_args = {
        'root': args['Global']['root'],
        'legend': args['Global']['legend'],
        'n_samples': args['Global']['n_samples'],
        'batch_size': args['Global']['batch_size'],
        'model': args['Global']['model'],
        'seed': seed_manager,
        'lr': args['Shadow']['lr'],
        'epochs': args['Shadow']['epochs'],
        'n_queries': args['Global']['n_queries'],
        'pkeep': args['Shadow']['pkeep'],
        'n_shadows': args['Global']['n_shadows'],
        'shadow_id': None,
        'id': None
    }
    
    for shadow_id in range(args['Global']['n_shadows']):
        seed_manager.random_seed()
        shadow_args['shadow_id'] = shadow_id
        shadow_args['id'] = shadow_id
             
        key_dict, remarks = get_key_params('Shadow', args, shadow_args)   
        accuracy = train.train_model(shadow_args, key_dict)   
        save_accuracy(shadow_args['shadow_id'], str(key_dict), str(remarks), accuracy)
        
def AdaRound_trainer(args, seed_manager):
    adaround_args = {
        'legend': args['Global']['legend'],
        'seed': seed_manager,
        'root': args['Global']['root'],
        'n_samples': args['Global']['n_samples'],
        'n_queries': args['Global']['n_queries'],
        'batch_size': args['Global']['batch_size'],
        'keep_file': model_manager.keep_file,
        'wbits': None,
        'asym': args['Global']['wasym'],
        'set_last_layer': args['Global']['set_last_layer'] if 'set_last_layer' in args['Global'] else None,
        'abits': 32   
    }
    
    for bit in args['Global']['wbit']:
        adaround_args['wbits'] = bit
        if auto_quantized_check('AdaRound', bit, 32):
            continue
        
        key_dict, remarks = get_key_params('AdaRound', args, adaround_args)
        accuracy = main_AdaRound(adaround_args, key_dict)
        save_accuracy(f'{model_manager.get_target()[1]}_W{bit}A32', str(key_dict), str(remarks), accuracy)
        
def OBC_trainer(args, seed_manager):
    obc_args = {
        'seed': seed_manager,
        'legend': args['Global']['legend'],
        'root': args['Global']['root'],
        'n_samples': args['Global']['n_samples'],
        'n_queries': args['Global']['n_queries'],
        'batch_size': args['Global']['batch_size'],
        'keep_file': model_manager.keep_file,
        'aquant': args['Global']['aquant'],
        'rel_damp': args['OBC']['rel_damp'],
        'abits': 32,
        'actsym': args['OBC']['actsym'],
        'aminmax': args['OBC']['aminmax'],
        'set_last_layer': args['Global']['set_last_layer'] if 'set_last_layer' in args['Global'] else None,
        'wperweight': args['OBC']['wperweight'],
        'wasym': args['Global']['wasym'],
        'wminmax': args['OBC']['wminmax'],
        'wbits': 32,
        'wquant': args['Global']['wquant'],
        'nrounds': args['OBC']['nrounds'],
        'bnt_batches': args['OBC']['bnt_batches']    
    }
    wbit_range = args['Global']['wbit'] if args['Global']['wquant'] else [32]
    abit_range = args['Global']['abit'] if args['Global']['aquant'] else [32]
    for abit in abit_range:
        for wbit in wbit_range:
            obc_args['abits'] = abit
            obc_args['wbits'] = wbit
            
            if auto_quantized_check('OBC', wbit, abit):
                continue
            
            key_dict, remarks = get_key_params('OBC', args, obc_args)
            accuracy = main_OBC(obc_args, key_dict)
            save_accuracy(f'{model_manager.get_target()[1]}_W{wbit}A{abit}', str(key_dict), str(remarks), accuracy)
    
def BRECQ_trainer(args, seed_manager):
    obc_args = {
        'seed': seed_manager,
        'legend': args['Global']['legend'],
        'root': args['Global']['root'],
        'n_samples': args['Global']['n_samples'],
        'n_queries': args['Global']['n_queries'],
        'batch_size': args['Global']['batch_size'],
        'keep_file': model_manager.keep_file,
        'aquant': args['Global']['aquant'],
        'abits': 32,
        'set_last_layer': args['Global']['set_last_layer'] if 'set_last_layer' in args['Global'] else None,
        'wasym': args['Global']['wasym'],
        'wbits': 32,
        'wquant': args['Global']['wquant'],
        'channel_wise': args['BRECQ']['channel_wise'],
        'iters_w': args['BRECQ']['iters_w'],
        'weight': args['BRECQ']['weight'],
        'b_start': args['BRECQ']['b_start'],
        'b_end': args['BRECQ']['b_end'],
        'warmup': args['BRECQ']['warmup'],
        'iters_a': args['BRECQ']['iters_a'],
        'lr': args['BRECQ']['lr'],
        'p': args['BRECQ']['p']
    }
    wbit_range = args['Global']['wbit'] if args['Global']['wquant'] else [32]
    abit_range = args['Global']['abit'] if args['Global']['aquant'] else [32]
    for abit in abit_range:
        for wbit in wbit_range:
            obc_args['abits'] = abit
            obc_args['wbits'] = wbit
            
            if auto_quantized_check('BRECQ', wbit, abit):
                continue
            
            key_dict, remarks = get_key_params('BRECQ', args, obc_args)
            accuracy = main_BRECQ(obc_args, key_dict)
            save_accuracy(f'{model_manager.get_target()[1]}_W{wbit}A{abit}', str(key_dict), str(remarks), accuracy)  
            
def LiRA_attack(args):
    attack_args = {
        'legend': args['Global']['legend'],
        'n_shadows': args['Global']['n_shadows'],
        }  
    attack(attack_args)
     
    
def get_key_params(process, args1, args2):
    key_dict = {}
    remarks = {}
    key_params = args1['Global']['key_params'][:]
    key_params.extend(args1[process]['key_params'])

    for x in key_params:
        if x in args2.keys():
            key_dict[x] = args2[x]
        elif x in args1[process].keys():
            key_dict[x] = args1[process][x]
        elif x in args1['Global'].keys():
            key_dict[x] = args1['Global'][x]
        else:
            raise KeyError(f"Key {x} not found")
    for key, value in args2.items():
        if key not in key_dict.keys():
            remarks[key] = value
    return key_dict, remarks

def auto_quantized_check(method, wbit, abit) -> bool:
    if not args['Global']['auto_quant']:
        return False
    
    quant_model_dir = os.path.dirname(__file__) + '/experiments/quantized_models'
    quant_model = f'{quant_model_dir}/{method}/{model_manager.get_target()[1]}_W{wbit}A{abit}.pt'

    if os.path.exists(quant_model):
        import json
        with open('accuracy.json', 'r') as f:
            data = json.load(f)
            for key_param in data.keys():
                for id in data[key_param].keys():
                    if id == f'{model_manager.get_target()[1]}_W{wbit}A{abit}' and method in key_param:
                        model_manager.add_quant_model(quant_model, key_param)
        tqdm.write(f'skip quantization, already exists:{quant_model}')
        return True
    else:
        return False



if __name__ == "__main__":
    args = load_global_args()
    model_manager._prepare(args['Global']['legend'])

    for idx in range(args["Global"]["cycles"]):
        cycle = f"Cycle {idx + 1}/{args['Global']['cycles']}"    
        seed_manager = SeedManager()
        if len(args["Global"]["seed"]) > 0:
            seed = args["Global"]["seed"][idx % len(args["Global"]["seed"])]
            seed_manager.set_seed(seed)
        else:
            tqdm.write("No seed provided, using random seed.")
            seed_manager.set_seed(rnd.randint(0, 100000))
            
            
        if args['Global']['stage']['target']: 
            target_model_trainer(args, seed_manager)
        if args['Global']['stage']['shadows']:
            shadow_model_trainer(args, seed_manager)
            
        if args['Global']['stage']['quant']:
            if args['Global']['quant']['AdaRound']:
                AdaRound_trainer(args, seed_manager)
            if args['Global']['quant']['OBC']:
                OBC_trainer(args, seed_manager)
            if args['Global']['quant']['BRECQ']:
                BRECQ_trainer(args, seed_manager)
                
        if args['Global']['stage']['attack']:
            LiRA_attack(args)