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)
