def main(mode, load):
random_seed = af.get_random_seed()
models_path = 'networks/{}'.format(random_seed)
device = af.get_pytorch_device()
create_params = [
# type, training, (prune?, keep_ratio for ics, batch size)
# ('dense', '0', (True, [0.75, 0.66, 0.58, 0.46], 128), [0, 0, 1, 0, 0, 1, 0, 1, 0])
('dense', '0', (False, [0.66, 0.46, 0.36, 0.36], 128), [1, 1, 1]),
('dense', '0', (False, [0.66, 0.46, 0.36, 0.36], 128), [1, 1, 1]),
('dense', '0', (False, [0.66, 0.46, 0.36, 0.36], 128), [1, 1, 1]),
('dense', '0', (False, [0.66, 0.46, 0.36, 0.36], 128), [1, 1, 1]),
('dense', '0', (False, [0.66, 0.46, 0.36, 0.36], 128), [1, 1, 1]),
]
create_bool = [1 if True else 0 for i in range(len(create_params))]
if load is not None:
model, param = arcs.load_model(models_path, load, -1)
arr = [(model, param)]
else:
arr = list(
multi_experiments(models_path, zip(create_params, create_bool),
device))
#af.print_acc(arr, groups=[5], extend=True)
af.print_acc(arr, extend=True)
#af.print_acc(arr, extend=False)
#af.plot_acc([m[1] for m in arr])
#print the numbers of paramters
for m in [t[0] for t in arr]:
print("")
print("flops: {}".format(af.calculate_flops(m, (3, 32, 32))))
for m, p in arr:
arcs.save_model(m, p, models_path, p['name'], -1)
print("model: {}".format(arr[0][0]))
def train_model(models_path, cr_params, device, num=0):
type, mode, pruning, ics = cr_params
model, params = arcs.create_resnet_iterative(models_path, type, mode,
pruning, ics, False)
dataset = af.get_dataset('cifar10')
params['name'] = params['base_model'] + '_{}_{}'.format(type, mode)
if model.prune:
params['name'] += "_prune_{}".format(
[x * 100 for x in model.keep_ratio])
print("prune: {}".format(model.keep_ratio))
if mode == "0":
params['epochs'] = 250
params['milestones'] = [120, 160, 180]
params['gammas'] = [0.1, 0.01, 0.01]
if mode == "1":
params['epochs'] = 300
params['milestones'] = [100, 150, 200]
params['gammas'] = [0.1, 0.1, 0.1]
if "full" in type:
params['learning_rate'] = 0.1
print("lr: {}".format(params['learning_rate']))
opti_param = (params['learning_rate'], params['weight_decay'],
params['momentum'], -1)
lr_schedule_params = (params['milestones'], params['gammas'])
model.to(device)
train_params = dict(
epochs=params['epochs'],
epoch_growth=[25, 50, 75],
epoch_prune=[10, 35, 60, 85, 110, 135, 160], #[10, 35, 60, 85],
prune_batch_size=pruning[2],
prune_type='2', # 0 skip layer, 1 normal full, 2 iterative
reinit=False,
min_ratio=[
0.3, 0.1, 0.05, 0.05
] # not needed if skip layers, minimum for the iterative pruning
)
params['epoch_growth'] = train_params['epoch_growth']
params['epoch_prune'] = train_params['epoch_prune']
optimizer, scheduler = af.get_full_optimizer(model, opti_param,
lr_schedule_params)
metrics, best_model = model.train_func(model, dataset, train_params,
optimizer, scheduler, device)
_link_metrics(params, metrics)
af.print_sparsity(best_model)
arcs.save_model(best_model, params, models_path, params['name'], epoch=-1)
print("test acc: {}, last val: {}".format(params['test_top1_acc'],
params['valid_top1_acc'][-1]))
return best_model, params
def train_sdns(models_path, networks, ic_only=False, device='cpu'):
if ic_only: # if we only train the ICs, we load a pre-trained CNN
load_epoch = -1
else: # if we train both ICs and the orig network, we load an untrained CNN
load_epoch = 0
for sdn_name in networks:
cnn_to_tune = sdn_name.replace('sdn', 'cnn')
sdn_params = arcs.load_params(models_path, sdn_name)
sdn_params = arcs.get_net_params(sdn_params['network_type'],
sdn_params['task'])
sdn_model, _ = af.cnn_to_sdn(
models_path, cnn_to_tune, sdn_params,
load_epoch) # load the CNN and convert it to a SDN
arcs.save_model(sdn_model, sdn_params, models_path, sdn_name,
epoch=0) # save the resulting SDN
train(models_path, networks, sdn=True, ic_only_sdn=ic_only, device=device)
def train_model(models_path, device):
_, sdn = arcs.create_resnet56(models_path, 'cifar10', save_type='d')
print('snd name: {}'.format(sdn))
# train_sdn(models_path, sdn, device)
print("Training model...")
trained_model, model_params = arcs.load_model(models_path, sdn, 0)
dataset = af.get_dataset(model_params['task'])
lr = model_params['learning_rate']
momentum = model_params['momentum']
weight_decay = model_params['weight_decay']
milestones = model_params['milestones']
gammas = model_params['gammas']
num_epochs = model_params['epochs']
model_params['optimizer'] = 'SGD'
opti_param = (lr, weight_decay, momentum, -1)
lr_schedule_params = (milestones, gammas)
optimizer, scheduler = af.get_full_optimizer(trained_model, opti_param,
lr_schedule_params)
trained_model_name = sdn + '_training'
print('Training: {}...'.format(trained_model_name))
trained_model.to(device)
metrics = trained_model.train_func(trained_model,
dataset,
num_epochs,
optimizer,
scheduler,
device=device)
model_params['train_top1_acc'] = metrics['train_top1_acc']
model_params['test_top1_acc'] = metrics['test_top1_acc']
model_params['train_top3_acc'] = metrics['train_top3_acc']
model_params['test_top3_acc'] = metrics['test_top3_acc']
model_params['epoch_times'] = metrics['epoch_times']
model_params['lrs'] = metrics['lrs']
total_training_time = sum(model_params['epoch_times'])
model_params['total_time'] = total_training_time
print('Training took {} seconds...'.format(total_training_time))
arcs.save_model(trained_model,
model_params,
models_path,
trained_model_name,
epoch=-1)
return trained_model, dataset
def sdn_ic_only_backdoored(device):
params = arcs.create_vgg16bn(None, 'cifar10', None, True)
path = 'backdoored_models'
backdoored_cnn_name = 'VGG16_cifar10_backdoored'
save_sdn_name = 'VGG16_cifar10_backdoored_SDN'
# Use the class VGG
backdoored_cnn = VGG(params)
backdoored_cnn.load_state_dict(torch.load('{}/{}'.format(
path, backdoored_cnn_name),
map_location='cpu'),
strict=False)
# convert backdoored cnn into a sdn
backdoored_sdn, sdn_params = af.cnn_to_sdn(
None, backdoored_cnn, params,
preloaded=backdoored_cnn) # load the CNN and convert it to a sdn
arcs.save_model(backdoored_sdn, sdn_params, path, save_sdn_name,
epoch=0) # save the resulting sdn
networks = [save_sdn_name]
train(path, networks, sdn=True, ic_only_sdn=True, device=device)
def train(models_path,
untrained_models,
sdn=False,
ic_only_sdn=False,
device='cpu',
ds=False):
print('Training models...')
for base_model in untrained_models:
trained_model, model_params = arcs.load_model(models_path, base_model,
0)
dataset = af.get_dataset(model_params['task'])
learning_rate = model_params['learning_rate']
momentum = model_params['momentum']
weight_decay = model_params['weight_decay']
milestones = model_params['milestones']
gammas = model_params['gammas']
num_epochs = model_params['epochs']
model_params['optimizer'] = 'SGD'
if ic_only_sdn: # IC-only training, freeze the original weights
learning_rate = model_params['ic_only']['learning_rate']
num_epochs = model_params['ic_only']['epochs']
milestones = model_params['ic_only']['milestones']
gammas = model_params['ic_only']['gammas']
model_params['optimizer'] = 'Adam'
trained_model.ic_only = True
else:
trained_model.ic_only = False
if ds:
trained_model.ds = True
else:
trained_model.ds = False
optimization_params = (learning_rate, weight_decay, momentum)
lr_schedule_params = (milestones, gammas)
# pdb.set_trace()
if sdn:
if ic_only_sdn:
optimizer, scheduler = af.get_sdn_ic_only_optimizer(
trained_model, optimization_params, lr_schedule_params)
trained_model_name = base_model + '_ic_only_ic{}'.format(
np.sum(model_params['add_ic']))
else:
optimizer, scheduler = af.get_full_optimizer(
trained_model, optimization_params, lr_schedule_params)
trained_model_name = base_model + '_sdn_training_ic{}'.format(
np.sum(model_params['add_ic']))
else:
optimizer, scheduler = af.get_full_optimizer(
trained_model, optimization_params, lr_schedule_params)
trained_model_name = base_model
if ds:
trained_model_name = trained_model_name + '_ds'
# pdb.set_trace()
print('Training: {}...'.format(trained_model_name))
# trained_model = nn.DataParallel(trained_model)
trained_model.to(device)
metrics = trained_model.train_func(trained_model,
dataset,
num_epochs,
optimizer,
scheduler,
device=device)
model_params['train_top1_acc'] = metrics['train_top1_acc']
model_params['test_top1_acc'] = metrics['test_top1_acc']
model_params['train_top5_acc'] = metrics['train_top5_acc']
model_params['test_top5_acc'] = metrics['test_top5_acc']
model_params['epoch_times'] = metrics['epoch_times']
model_params['lrs'] = metrics['lrs']
total_training_time = sum(model_params['epoch_times'])
model_params['total_time'] = total_training_time
print('Training took {} seconds...'.format(total_training_time))
arcs.save_model(trained_model,
model_params,
models_path,
trained_model_name,
epoch=-1)