def quantize_network(arch,
dataset,
train_gen,
test_gen,
model_chkp=None,
x_bits=8,
w_bits=8,
desc=None):
# Initialize log file
name_str = '{}-{}_quantize_x-{}_w-{}'.format(arch, dataset, x_bits, w_bits)
name_str = name_str + '_{}'.format(desc) if desc is not None else name_str
name_str = name_str + '_seed-{}-{}'.format(cfg.SEED_TORCH, cfg.SEED_NP)
cfg.LOG.start_new_log(name=name_str)
cfg.LOG.write('desc={}, x_bits={}, w_bits={}'.format(desc, x_bits, w_bits))
# Initialize model
nn = NeuralNet(arch, dataset, model_chkp=model_chkp)
# Set configurations
nn.model.set_quantize(True)
nn.model.set_quantization_bits(x_bits, w_bits)
nn.model.set_min_max_update(True)
nn.model.set_unfold(False)
nn.best_top1_acc = 0
nn.next_train_epoch = 0
# Learning rate is set to zero
nn.train(train_gen,
test_gen,
epochs=1,
lr=0,
iterations=2048 / cfg.BATCH_SIZE)
cfg.LOG.close_log()
return
def eval_model(train_gen,
test_gen,
arch,
dataset,
pred_chkps_dict,
model_chkp=None,
epochs=5,
lr=0.0001,
wd=0.0005,
th_list=None,
th_dict=None,
mask_list=None,
fine_tune=True,
bn_phase=True):
"""
Evaluate and fine-tune an entire model with ZAPs.
:param train_gen: the training set used to train the predictors
:param test_gen: the test set to evaluate the predictor performance
:param arch: a string that represents the model architecture, e.g., 'alexnet'
:param dataset: a string that represents the dataset, e.g., 'cifar100'
:param pred_chkps_dict: ZAP checkpoint paths per mask
:param model_chkp: a model checkpoint path to be loaded (default: None)
:param epochs: number of epochs to fine-tune the network (default: 5)
:param lr: fine-tuning learning rate (default: 0.0001)
:param wd: fine-tuning weight decay (default: 0.0005)
:param th_list: list of thresholds/errors to evaluate (default: [0, 0.1, 0.2, 0.3, 0.4, 0.5])
:param th_dict: a path to the threasholds dictionary (default: None)
:param fine_tune: continue to fine-tuning after evaluation (default: True)
:param mask_list: specific mask list (default: everything in pred_chkps_dict)
:param bn_phase: recoup performance via BN phase (default: True)
"""
# Set default threshold values
if th_dict is not None:
thresholds = pickle.load(open(th_dict, 'rb'))
if th_list is not None:
thresholds_new = OrderedDict()
th_list_int = [round(float(i), 2) for i in th_list]
for k, v in thresholds.items():
if k in th_list_int:
thresholds_new[k] = thresholds[k]
thresholds = thresholds_new
else:
if th_list is None:
thresholds = np.around(np.arange(0, 0.52, 0.1), 2)
else:
thresholds = th_list
for mask, chkp in pred_chkps_dict.items():
if mask_list is not None:
if str(mask) not in mask_list:
continue
mask = int(mask)
pred_chkps = checkpoint.get_chkp_files([chkp])
for th in thresholds:
cfg.LOG.start_new_log(name='{}-{}_full-model_mask-{}_th-{}'.format(
arch, dataset, mask, th))
cfg.LOG.write_title("Mask={}".format(mask),
pad_width=40,
pad_symbol='=')
nn = NeuralNet(arch, dataset, model_chkp=model_chkp)
for pred_idx, values in pred_chkps.items():
nn.load_state_pred(values['filename'])
# Assuming that if its an OrderedDict that its a threshold dictionary
if thresholds.__class__.__name__ == 'OrderedDict':
cfg.LOG.write_title(
"Optimized Thresholds - Max Error={}".format(th),
pad_width=40,
pad_symbol='=')
j = 0
for i, pred_layer in enumerate(nn.model.pred_layers):
if arch == 'vgg16' and i == 0:
cfg.LOG.write('Prediction layer {} skipped'.format(i))
continue
pred_layer.threshold = thresholds[th][j]
cfg.LOG.write(
'Prediction layer {} threshold set to {}'.format(
i, pred_layer.threshold))
j += 1
else:
cfg.LOG.write_title("Uniform Threshold={}".format(th),
pad_width=40,
pad_symbol='=')
for i, pred_layer in enumerate(nn.model.pred_layers):
pred_layer.threshold = round(float(th), 2)
# Disable first layer in VGG-16 (the other models just don't have ZAP in the first layer)
if arch == 'vgg16':
cfg.LOG.write('Disabled first layer in {}'.format(arch))
nn.model.disabled_pred_layers.append(0)
nn.model.pred_layers[0].disable_layer()
nn.test(test_gen, stats=[cfg.STATS_GENERAL])
# Recouping BN running var and mean
nn.next_train_epoch = 0
nn.best_top1_acc = 0
nn.model.enable_pred_layers()
nn.model.set_train(False)
nn.model.enable_grad()
nn.model.disable_pred_layers_grad()
if bn_phase:
cfg.LOG.write('BN epoch')
nn.train(train_gen,
test_gen,
1,
lr=0.0,
wd=0.0,
stats=[cfg.STATS_GENERAL])
if fine_tune:
cfg.LOG.write_title("{}-Epoch Fine-Tuning".format(int(epochs)),
pad_width=40,
pad_symbol='.')
nn.train(train_gen,
test_gen,
epochs,
lr=lr,
wd=wd,
stats=[cfg.STATS_GENERAL])
cfg.LOG.close_log()
nn = None
torch.cuda.empty_cache()
def eval_roc(train_gen,
test_gen,
arch,
dataset,
pred_chkps_dict,
model_chkp=None,
mask_list=None,
recoup_bn=False):
"""
Evaluate ZAPs ROC.
:param train_gen: the training set used to train the predictors
:param test_gen: the test set to evaluate the predictor performance
:param arch: a string that represents the model architecture, e.g., 'alexnet'
:param dataset: a string that represents the dataset, e.g., 'cifar100'
:param model_chkp: a model checkpoint path to be loaded (default: None)
:param mask_list: specific mask list to train with (default: [6, 5, 4, 3])
:param recoup_bn: run BN recoup epoch (default: False)
"""
for mask, chkp in pred_chkps_dict.items():
if mask_list is not None:
if str(mask) not in mask_list:
continue
mask = int(mask)
pred_chkps = checkpoint.get_chkp_files([chkp])
cfg.LOG.start_new_log(
name='{}-{}_zap-analysis-roc_mask-{}'.format(arch, dataset, mask))
cfg.LOG.write_title("Mask={}".format(mask),
pad_width=40,
pad_symbol='=')
nn = NeuralNet(arch, dataset, model_chkp=model_chkp)
for pred_idx, values in pred_chkps.items():
nn.load_state_pred(values['filename'])
# Disable first layer in VGG-16 (the other models just don't have ZAP in the first layer)
if arch == 'vgg16':
cfg.LOG.write('Disabled first layer in {}'.format(arch))
nn.model.disabled_pred_layers.append(0)
nn.model.pred_layers[0].disable_layer()
for i, threshold in enumerate(
np.around(
np.arange(cfg.STATS_ROC_MIN, cfg.STATS_ROC_MAX,
cfg.STATS_ROC_STEP), 2)):
cfg.LOG.write_title("Threshold={}".format(threshold),
pad_width=40,
pad_symbol='=')
for pred_layer in nn.model.pred_layers:
pred_layer.threshold = threshold
if recoup_bn:
cfg.LOG.write('BN epoch')
nn.next_train_epoch = 0
nn.best_top1_acc = 0
nn.model.enable_pred_layers()
nn.model.set_train(False)
nn.model.enable_grad()
nn.model.disable_pred_layers_grad()
nn.train(train_gen,
test_gen,
1,
lr=0.0,
wd=0.0,
stats=[cfg.STATS_GENERAL, cfg.STATS_ROC])
else:
nn.test(test_gen, stats=[cfg.STATS_GENERAL, cfg.STATS_ROC])
cfg.LOG.close_log()
nn = None
torch.cuda.empty_cache()
def eval_model(train_gen,
test_gen,
arch,
dataset,
pred_chkps_dict,
model_chkp=None,
epochs=5,
lr=0.0001,
wd=0.0005,
th_list=None,
mask_list=None,
fine_tune=True):
"""
Evaluate and fine-tune an entire model with ZAPs.
:param train_gen: the training set used to train the predictors
:param test_gen: the test set to evaluate the predictor performance
:param arch: a string that represents the model architecture, e.g., 'alexnet'
:param dataset: a string that represents the dataset, e.g., 'cifar100'
:param pred_chkps_dict: ZAP checkpoint paths per mask
:param model_chkp: a model checkpoint path to be loaded (default: None)
:param epochs: number of epochs to fine-tune the network (default: 5)
:param lr: fine-tuning learning rate (default: 0.0001)
:param wd: fine-tuning weight decay (default: 0.0005)
:param th_list: list of thresholds to evaluate (default: [0, 0.1, 0.2, 0.3, 0.4, 0.5])
:param fine_tune: continue to fine-tuning after evaluation (default: True)
:param mask_list: specific mask list (default: everything in pred_chkps_dict)
"""
# Set default threshold values
if th_list is None:
th_list = np.around(np.arange(0, 0.52, 0.1), 2)
for mask, chkp in pred_chkps_dict.items():
if mask_list is not None:
if str(mask) not in mask_list:
continue
mask = int(mask)
pred_chkps = checkpoint.get_chkp_files([chkp])
for threshold in th_list:
cfg.LOG.start_new_log(name='{}-{}_full-model_mask-{}_th-{}'.format(
arch, dataset, mask, threshold))
cfg.LOG.write_title("Mask={}".format(mask),
pad_width=40,
pad_symbol='=')
nn = NeuralNet(arch, dataset, model_chkp=model_chkp)
for pred_idx, values in pred_chkps.items():
nn.load_state_pred(values['filename'])
cfg.LOG.write_title("Threshold={}".format(threshold),
pad_width=40,
pad_symbol='=')
for pred_layer in nn.model.pred_layers:
pred_layer.threshold = round(float(threshold), 2)
# Disable first layer in VGG-16 (the other models just don't have ZAP in the first layer)
if arch == 'vgg16':
cfg.LOG.write('Disabled first layer in {}'.format(arch))
nn.model.disabled_pred_layers.append(0)
nn.model.pred_layers[0].disable_layer()
nn.test(test_gen, stats=[cfg.STATS_GENERAL])
# Recouping BN running var and mean
cfg.LOG.write('BN epoch')
nn.next_train_epoch = 0
nn.best_top1_acc = 0
nn.model.enable_pred_layers()
nn.model.set_train(False)
nn.model.enable_grad()
nn.model.disable_pred_layers_grad()
nn.train(train_gen,
test_gen,
1,
lr=0.0,
wd=0.0,
stats=[cfg.STATS_GENERAL])
if fine_tune:
cfg.LOG.write_title("{}-Epoch Fine-Tuning".format(int(epochs)),
pad_width=40,
pad_symbol='.')
nn.train(train_gen,
test_gen,
epochs,
lr=lr,
wd=wd,
stats=[cfg.STATS_GENERAL])
cfg.LOG.close_log()
nn = None
torch.cuda.empty_cache()
