from __future__ import print_function import mxnet as mx import mxnet.ndarray as nd import numpy import logging import matplotlib.pyplot as plt from scipy.stats import gaussian_kde import argparse from algos import * from data_loader import * from utils import * class CrossEntropySoftmax(mx.operator.NumpyOp): def __init__(self): super(CrossEntropySoftmax, self).__init__(False) def list_arguments(self): return ['data', 'label'] def list_outputs(self): return ['output'] def infer_shape(self, in_shape): data_shape = in_shape[0] label_shape = in_shape[0] output_shape = in_shape[0] return [data_shape, label_shape], [output_shape] def forward(self, in_data, out_data): x = in_data[0] y = out_data[0] y[:] = numpy.exp(x - x.max(axis=1).reshape((x.shape[0], 1))).astype('float32') y /= y.sum(axis=1).reshape((x.shape[0], 1)) def backward(self, out_grad, in_data, out_data, in_grad): l = in_data[1] y = out_data[0] dx = in_grad[0] dx[:] = (y - l) class LogSoftmax(mx.operator.NumpyOp): def __init__(self): super(LogSoftmax, self).__init__(False) def list_arguments(self): return ['data', 'label'] def list_outputs(self): return ['output'] def infer_shape(self, in_shape): data_shape = in_shape[0] label_shape = in_shape[0] output_shape = in_shape[0] return [data_shape, label_shape], [output_shape] def forward(self, in_data, out_data): x = in_data[0] y = out_data[0] y[:] = (x - x.max(axis=1, keepdims=True)).astype('float32') y -= numpy.log(numpy.exp(y).sum(axis=1, keepdims=True)).astype('float32') # y[:] = numpy.exp(x - x.max(axis=1).reshape((x.shape[0], 1))) # y /= y.sum(axis=1).reshape((x.shape[0], 1)) def backward(self, out_grad, in_data, out_data, in_grad): l = in_data[1] y = out_data[0] dx = in_grad[0] dx[:] = (numpy.exp(y) - l).astype('float32') def classification_student_grad(student_outputs, teacher_pred): return [student_outputs[0] - teacher_pred] def regression_student_grad(student_outputs, teacher_pred, teacher_noise_precision): student_mean = student_outputs[0] student_var = student_outputs[1] grad_mean = nd.exp(-student_var) * (student_mean - teacher_pred) grad_var = (1 - nd.exp(-student_var) * (nd.square(student_mean - teacher_pred) + 1.0 / teacher_noise_precision)) / 2 return [grad_mean, grad_var] def get_mnist_sym(output_op=None, num_hidden=400): net = mx.symbol.Variable('data') net = mx.symbol.FullyConnected(data=net, name='mnist_fc1', num_hidden=num_hidden) net = mx.symbol.Activation(data=net, name='mnist_relu1', act_type="relu") net = mx.symbol.FullyConnected(data=net, name='mnist_fc2', num_hidden=num_hidden) net = mx.symbol.Activation(data=net, name='mnist_relu2', act_type="relu") net = mx.symbol.FullyConnected(data=net, name='mnist_fc3', num_hidden=10) if output_op is None: net = mx.symbol.SoftmaxOutput(data=net, name='softmax') else: net = output_op(data=net, name='softmax') return net def synthetic_grad(X, theta, sigma1, sigma2, sigmax, rescale_grad=1.0, grad=None): if grad is None: grad = nd.empty(theta.shape, theta.context) theta1 = theta.asnumpy()[0] theta2 = theta.asnumpy()[1] v1 = sigma1 ** 2 v2 = sigma2 ** 2 vx = sigmax ** 2 denominator = numpy.exp(-(X - theta1) ** 2 / (2 * vx)) + numpy.exp( -(X - theta1 - theta2) ** 2 / (2 * vx)) grad_npy = numpy.zeros(theta.shape) grad_npy[0] = -rescale_grad * ((numpy.exp(-(X - theta1) ** 2 / (2 * vx)) * (X - theta1) / vx + numpy.exp(-(X - theta1 - theta2) ** 2 / (2 * vx)) * ( X - theta1 - theta2) / vx) / denominator).sum() \ + theta1 / v1 grad_npy[1] = -rescale_grad * ((numpy.exp(-(X - theta1 - theta2) ** 2 / (2 * vx)) * ( X - theta1 - theta2) / vx) / denominator).sum() \ + theta2 / v2 grad[:] = grad_npy return grad def get_toy_sym(teacher=True, teacher_noise_precision=None): if teacher: net = mx.symbol.Variable('data') net = mx.symbol.FullyConnected(data=net, name='teacher_fc1', num_hidden=100) net = mx.symbol.Activation(data=net, name='teacher_relu1', act_type="relu") net = mx.symbol.FullyConnected(data=net, name='teacher_fc2', num_hidden=1) net = mx.symbol.LinearRegressionOutput(data=net, name='teacher_output', grad_scale=teacher_noise_precision) else: net = mx.symbol.Variable('data') net = mx.symbol.FullyConnected(data=net, name='student_fc1', num_hidden=100) net = mx.symbol.Activation(data=net, name='student_relu1', act_type="relu") student_mean = mx.symbol.FullyConnected(data=net, name='student_mean', num_hidden=1) student_var = mx.symbol.FullyConnected(data=net, name='student_var', num_hidden=1) net = mx.symbol.Group([student_mean, student_var]) return net def dev(): return mx.gpu() def run_mnist_SGD(training_num=50000): X, Y, X_test, Y_test = load_mnist(training_num) minibatch_size = 100 net = get_mnist_sym() data_shape = (minibatch_size,) + X.shape[1::] data_inputs = {'data': nd.zeros(data_shape, ctx=dev()), 'softmax_label': nd.zeros((minibatch_size,), ctx=dev())} initializer = mx.init.Xavier(factor_type="in", magnitude=2.34) exe, exe_params, _ = SGD(sym=net, dev=dev(), data_inputs=data_inputs, X=X, Y=Y, X_test=X_test, Y_test=Y_test, total_iter_num=1000000, initializer=initializer, lr=5E-6, prior_precision=1.0, minibatch_size=100) def run_mnist_SGLD(training_num=50000): X, Y, X_test, Y_test = load_mnist(training_num) minibatch_size = 100 net = get_mnist_sym() data_shape = (minibatch_size,) + X.shape[1::] data_inputs = {'data': nd.zeros(data_shape, ctx=dev()), 'softmax_label': nd.zeros((minibatch_size,), ctx=dev())} initializer = mx.init.Xavier(factor_type="in", magnitude=2.34) exe, sample_pool = SGLD(sym=net, dev=dev(), data_inputs=data_inputs, X=X, Y=Y, X_test=X_test, Y_test=Y_test, total_iter_num=1000000, initializer=initializer, learning_rate=4E-6, prior_precision=1.0, minibatch_size=100, thin_interval=100, burn_in_iter_num=1000) def run_mnist_DistilledSGLD(training_num=50000): X, Y, X_test, Y_test = load_mnist(training_num) minibatch_size = 100 if training_num >= 10000: num_hidden = 800 total_iter_num = 1000000 teacher_learning_rate = 1E-6 student_learning_rate = 0.0001 teacher_prior = 1 student_prior = 0.1 perturb_deviation = 0.1 else: num_hidden = 400 total_iter_num = 20000 teacher_learning_rate = 4E-5 student_learning_rate = 0.0001 teacher_prior = 1 student_prior = 0.1 perturb_deviation = 0.001 teacher_net = get_mnist_sym(num_hidden=num_hidden) logsoftmax = LogSoftmax() student_net = get_mnist_sym(output_op=logsoftmax, num_hidden=num_hidden) data_shape = (minibatch_size,) + X.shape[1::] teacher_data_inputs = {'data': nd.zeros(data_shape, ctx=dev()), 'softmax_label': nd.zeros((minibatch_size,), ctx=dev())} student_data_inputs = {'data': nd.zeros(data_shape, ctx=dev()), 'softmax_label': nd.zeros((minibatch_size, 10), ctx=dev())} teacher_initializer = BiasXavier(factor_type="in", magnitude=1) student_initializer = BiasXavier(factor_type="in", magnitude=1) student_exe, student_params, _ = \ DistilledSGLD(teacher_sym=teacher_net, student_sym=student_net, teacher_data_inputs=teacher_data_inputs, student_data_inputs=student_data_inputs, X=X, Y=Y, X_test=X_test, Y_test=Y_test, total_iter_num=total_iter_num, student_initializer=student_initializer, teacher_initializer=teacher_initializer, student_optimizing_algorithm="adam", teacher_learning_rate=teacher_learning_rate, student_learning_rate=student_learning_rate, teacher_prior_precision=teacher_prior, student_prior_precision=student_prior, perturb_deviation=perturb_deviation, minibatch_size=100, dev=dev()) def run_toy_SGLD(): X, Y, X_test, Y_test = load_toy() minibatch_size = 1 teacher_noise_precision = 1.0 / 9.0 net = get_toy_sym(True, teacher_noise_precision) data_shape = (minibatch_size,) + X.shape[1::] data_inputs = {'data': nd.zeros(data_shape, ctx=dev()), 'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev())} initializer = mx.init.Uniform(0.07) exe, params, _ = \ SGLD(sym=net, data_inputs=data_inputs, X=X, Y=Y, X_test=X_test, Y_test=Y_test, total_iter_num=50000, initializer=initializer, learning_rate=1E-4, # lr_scheduler=mx.lr_scheduler.FactorScheduler(100000, 0.5), prior_precision=0.1, burn_in_iter_num=1000, thin_interval=10, task='regression', minibatch_size=minibatch_size, dev=dev()) def run_toy_DistilledSGLD(): X, Y, X_test, Y_test = load_toy() minibatch_size = 1 teacher_noise_precision = 1.0 teacher_net = get_toy_sym(True, teacher_noise_precision) student_net = get_toy_sym(False) data_shape = (minibatch_size,) + X.shape[1::] teacher_data_inputs = {'data': nd.zeros(data_shape, ctx=dev()), 'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev())} student_data_inputs = {'data': nd.zeros(data_shape, ctx=dev())} # 'softmax_label': nd.zeros((minibatch_size, 10), ctx=dev())} teacher_initializer = mx.init.Uniform(0.07) student_initializer = mx.init.Uniform(0.07) student_grad_f = lambda student_outputs, teacher_pred: \ regression_student_grad(student_outputs, teacher_pred, teacher_noise_precision) student_exe, student_params, _ = \ DistilledSGLD(teacher_sym=teacher_net, student_sym=student_net, teacher_data_inputs=teacher_data_inputs, student_data_inputs=student_data_inputs, X=X, Y=Y, X_test=X_test, Y_test=Y_test, total_iter_num=80000, teacher_initializer=teacher_initializer, student_initializer=student_initializer, teacher_learning_rate=1E-4, student_learning_rate=0.01, # teacher_lr_scheduler=mx.lr_scheduler.FactorScheduler(100000, 0.5), student_lr_scheduler=mx.lr_scheduler.FactorScheduler(8000, 0.8), student_grad_f=student_grad_f, teacher_prior_precision=0.1, student_prior_precision=0.001, perturb_deviation=0.1, minibatch_size=minibatch_size, task='regression', dev=dev()) def run_toy_HMC(): X, Y, X_test, Y_test = load_toy() minibatch_size = Y.shape[0] noise_precision = 1 / 9.0 net = get_toy_sym(True, noise_precision) data_shape = (minibatch_size,) + X.shape[1::] data_inputs = {'data': nd.zeros(data_shape, ctx=dev()), 'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev())} initializer = mx.init.Uniform(0.07) sample_pool = HMC(net, data_inputs=data_inputs, X=X, Y=Y, X_test=X_test, Y_test=Y_test, sample_num=300000, initializer=initializer, prior_precision=1.0, learning_rate=1E-3, L=10, dev=dev()) def run_synthetic_SGLD(): theta1 = 0 theta2 = 1 sigma1 = numpy.sqrt(10) sigma2 = 1 sigmax = numpy.sqrt(2) X = load_synthetic(theta1=theta1, theta2=theta2, sigmax=sigmax, num=100) minibatch_size = 1 total_iter_num = 1000000 lr_scheduler = SGLDScheduler(begin_rate=0.01, end_rate=0.0001, total_iter_num=total_iter_num, factor=0.55) optimizer = mx.optimizer.create('sgld', learning_rate=None, rescale_grad=1.0, lr_scheduler=lr_scheduler, wd=0) updater = mx.optimizer.get_updater(optimizer) theta = mx.random.normal(0, 1, (2,), mx.cpu()) grad = nd.empty((2,), mx.cpu()) samples = numpy.zeros((2, total_iter_num)) start = time.time() for i in xrange(total_iter_num): if (i + 1) % 100000 == 0: end = time.time() print("Iter:%d, Time spent: %f" % (i + 1, end - start)) start = time.time() ind = numpy.random.randint(0, X.shape[0]) synthetic_grad(X[ind], theta, sigma1, sigma2, sigmax, rescale_grad= X.shape[0] / float(minibatch_size), grad=grad) updater('theta', grad, theta) samples[:, i] = theta.asnumpy() plt.hist2d(samples[0, :], samples[1, :], (200, 200), cmap=plt.cm.jet) plt.colorbar() plt.show() if __name__ == '__main__': numpy.random.seed(100) mx.random.seed(100) parser = argparse.ArgumentParser( description="Examples in the paper [NIPS2015]Bayesian Dark Knowledge and " "[ICML2011]Bayesian Learning via Stochastic Gradient Langevin Dynamics") parser.add_argument("-d", "--dataset", type=int, default=1, help="Dataset to use. 0 --> TOY, 1 --> MNIST, 2 --> Synthetic Data in " "the SGLD paper") parser.add_argument("-l", "--algorithm", type=int, default=2, help="Type of algorithm to use. 0 --> SGD, 1 --> SGLD, other-->DistilledSGLD") parser.add_argument("-t", "--training", type=int, default=50000, help="Number of training samples") args = parser.parse_args() training_num = args.training if args.dataset == 1: if 0 == args.algorithm: run_mnist_SGD(training_num) elif 1 == args.algorithm: run_mnist_SGLD(training_num) else: run_mnist_DistilledSGLD(training_num) elif args.dataset == 0: if 1 == args.algorithm: run_toy_SGLD() elif 2 == args.algorithm: run_toy_DistilledSGLD() elif 3 == args.algorithm: run_toy_HMC() else: run_synthetic_SGLD()