if __name__ == "__main__": from keras.models import Model from keras.layers import Input, Dense import keras.callbacks import logging logging.getLogger('keras').setLevel(logging.INFO) import argparse parser = argparse.ArgumentParser(description='Run Variational IB', formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('beta' , type=float, default=0.0, help='beta hyperparameter value') parser.add_argument('--epoch_report_mi', action='store_true', default=False, help='Report MI values every epoch?') args = parser.parse_args() print("Running variational IB with beta=%.5f" % args.beta) from buildmodel import get_mnist trn, tst = get_mnist() # Build model micalculator = MICalculatorVIB(args.beta) noiselayer = NoiseLayerVIB(mean_dims=256, test_phase_noise=True) micalculator.set_noiselayer(noiselayer) inputs = Input(shape=(784,)) c_layer = Dense(1024, kernel_initializer='he_uniform', activation='relu')(inputs) c_layer = Dense(1024, kernel_initializer='he_uniform', activation='relu')(c_layer) c_layer = Dense(512 , kernel_initializer='he_uniform', activation='linear', activity_regularizer=micalculator)(c_layer) c_layer = noiselayer(c_layer) predictions = Dense(trn.nb_classes, init='he_uniform', activation='softmax')(c_layer) model = Model(input=inputs, output=predictions)
# Example that demonstrates how to generate a figure similar to Fig.2 in Artemy Kolchinsky, Brendan D. Tracey, David H. Wolpert, "Nonlinear Information Bottleneck", https://arxiv.org/abs/1705.02436 # Minimal example of how to create a model with nonlinearIB layers import numpy as np from keras.models import Model from keras.layers import Input, Dense import buildmodel, layers, training, reporting BETA_VAL = 0.4 trn, tst = buildmodel.get_mnist() input_layer = Input((trn.X.shape[1], )) # hidden_layers_to_add should include a list of all layers that will get added before the nonlinearIB layers hidden_layers_to_add = [ Dense(800, activation='relu'), Dense(800, activation='relu'), Dense(2, activation='linear'), ] # *** The following creates the layers and callbacks necessary to run nonlinearIB *** micalculator = layers.MICalculator(BETA_VAL, model_layers=hidden_layers_to_add, data=trn.X, miN=1000) noiselayer = layers.NoiseLayer(logvar_trainable=True, test_phase_noise=False) micalculator.set_noiselayer(noiselayer) # Start hooking up the layers together
else: import os os.environ['KERAS_BACKEND'] = 'tensorflow' logging.getLogger('keras').setLevel(logging.INFO) import reporting import buildmodel import keras.callbacks arg_dict = vars(args) VALIDATE_ON_TEST = True arg_dict['noise_logvar_grad_trainable'] = True trn, tst = buildmodel.get_mnist(args.trainN, args.testN) # *************************** arg_dict['INPUT_DIM'] = trn.X.shape[1] print '# ARGS:', arg_dict model, cbs, noiselayer, micalculator = buildmodel.buildmodel(arg_dict, trn=trn) # Reports MI and Cross Entropy Values reporter = reporting.Reporter(trn=trn, tst=tst, noiselayer=noiselayer, micalculator=micalculator, on_epoch_report_mi=args.epoch_report_mi) cbs.append(reporter)