import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

import os
print('-------------------------')
print('all files:')
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))
print('-------------------------')
"""
# Create a GT Teacher network
"""
input_dim = 10
nn_depth = 4
nn_width = 4
output_dim = 1

leaky_relu_alpha = 0.3
inputs = keras.Input(shape=(None, input_dim), name='input')

x = layers.Dense(nn_width, name='FC_1')(inputs)
x = layers.LeakyReLU(alpha=leaky_relu_alpha, name='LReLU_1')(x)

for k in range(nn_depth-1):
    x = layers.Dense(nn_width, name='FC_%d' %(k+2))(x)
    x = layers.LeakyReLU(alpha=leaky_relu_alpha, name='LReLU_%d' %(k+2))(x)

outputs = layers.Dense(output_dim, name='output')(x)

model_name = 'FCN_%dx%d_GT_model' %(nn_depth, nn_width)
GT_model = keras.Model(inputs=inputs, outputs=outputs, name=model_name)
GT_model.summary()

GT_model.compile(optimizer="adam", loss="mse")
"""
# create a generator that will create data based on teacher GT network
"""
class UniformDataGenerator(keras.utils.Sequence):
    def __init__(self, teacher_model, input_dim=10, num_batches_per_epoch=1, batch_size=64):
        self.teacher_model = teacher_model
        self.input_dim = input_dim
        self.batch_size = batch_size
        self.num_batches_per_epoch = num_batches_per_epoch

    def __len__(self):
        return self.num_batches_per_epoch

    def __getitem__(self, idx):
        batch_x = -1.0 + 2.0 * np.random.rand(self.batch_size, self.input_dim)
        batch_y = self.teacher_model.predict(batch_x)
        return batch_x, batch_y
    
num_batches_per_epoch = 1
batch_size = 32

train_datagen = UniformDataGenerator(GT_model, input_dim=input_dim, num_batches_per_epoch=num_batches_per_epoch, batch_size=batch_size)
valid_datagen = UniformDataGenerator(GT_model, input_dim=input_dim, num_batches_per_epoch=num_batches_per_epoch, batch_size=batch_size)
"""
# Test the generator
"""
test_generator = False

if test_generator:
    batch_1 = train_datagen[0]
    batch_2 = train_datagen[42]

    print('batch 1 shapes: (X.shape = %s, y.shape = %s)' %(batch_1[0].shape, batch_1[1].shape))
    #print('batch 1 X: \n%s' %(batch_1[0]))
    #print('batch 1 y: \n%s' %(batch_1[1]))

    print('X batch 1 (mean, std) = \n (%s, \n  %s)' %(batch_1[0].mean(axis=0), batch_1[0].std(axis=0)))
    print('y batch 1 (mean, std) = \n (%s, %s)' %(batch_1[1].mean(axis=0), batch_1[1].std(axis=0)))
    
    print('batch 2 shapes: (X.shape = %s, y.shape = %s)' %(batch_2[0].shape, batch_2[1].shape))
    #print('batch 2 X: \n%s' %(batch_2[0]))
    #print('batch 2 y: \n%s' %(batch_2[1]))
    
    print('X batch 2 (mean, std) = \n (%s, \n  %s)' %(batch_2[0].mean(axis=0), batch_2[0].std(axis=0)))
    print('y batch 2 (mean, std) = \n (%s, %s)' %(batch_2[1].mean(axis=0), batch_2[1].std(axis=0)))

"""
# Create a student network
"""
student_nn_depth = 1
student_nn_width = 4

leaky_relu_alpha = 0.3
inputs = keras.Input(shape=(None, input_dim), name='input')

x = layers.Dense(student_nn_width, name='FC_1')(inputs)
x = layers.LeakyReLU(alpha=leaky_relu_alpha, name='LReLU_1')(x)

for k in range(student_nn_depth-1):
    x = layers.Dense(student_nn_width, name='FC_%d' %(k+2))(x)
    x = layers.LeakyReLU(alpha=leaky_relu_alpha, name='LReLU_%d' %(k+2))(x)

outputs = layers.Dense(nn_width, name='output')(x)

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
student_model = keras.Model(inputs=inputs, outputs=outputs, name=model_name)
student_model.summary()

student_model.compile(optimizer="adam", loss="mse")
"""
# fit the student network and plot learning curve
"""
num_train_iterations = 1000

history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
train_loss = history.history['loss']
valid_loss = history.history['val_loss']
batch_index = np.arange(len(train_loss))

plt.figure(figsize=(20,8))
plt.plot(batch_index, train_loss, color='k')
plt.plot(batch_index, valid_loss, color='b')
plt.xlabel('iteration', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(['train', 'valid'], fontsize=20)
"""
# Train many small student networks
"""
def create_FCN(input_dim, output_dim, nn_depth, nn_width, model_name, leaky_relu_alpha=0.3):
    
    # input
    inputs = keras.Input(shape=(None, input_dim), name='input')
    
    # first hiddent layer
    x = layers.Dense(nn_width, name='FC_1')(inputs)
    x = layers.LeakyReLU(alpha=leaky_relu_alpha, name='LReLU_1')(x)
    
    # rest of hidden layers
    for k in range(nn_depth-1):
        x = layers.Dense(nn_width, name='FC_%d' %(k+2))(x)
        x = layers.LeakyReLU(alpha=leaky_relu_alpha, name='LReLU_%d' %(k+2))(x)

    # output
    outputs = layers.Dense(output_dim, name='output')(x)

    # assemble the model
    FCN_model = keras.Model(inputs=inputs, outputs=outputs, name=model_name)
    
    return FCN_model
student_nn_depth = 2
student_nn_width = 2

num_random_inits = 20
num_train_iterations = 500

training_results = {}
model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
"""
# Show all learning curves overlaid 
"""
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

ylim_max = 1.1 * learning_curve_matrix.max()

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);

student_nn_depth = 2
student_nn_width = 3

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
    
    
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);
student_nn_depth = 3
student_nn_width = 2

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
    
    
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);
"""
# train many "correct size" student networks
"""
student_nn_depth = 3
student_nn_width = 3

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
    
    
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);

student_nn_depth = 3
student_nn_width = 4

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
    
    
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);
student_nn_depth = 4
student_nn_width = 3

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
    
    
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);
student_nn_depth = 4
student_nn_width = 4

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
    
    
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);
"""
# Train many large student networks
"""
student_nn_depth = 5
student_nn_width = 5

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
    
    
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);
student_nn_depth = 6
student_nn_width = 6

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
    
    
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);
student_nn_depth = 8
student_nn_width = 8

model_name = 'FCN_%dx%d_student_model' %(student_nn_depth, student_nn_width)
training_results[model_name] = []

for k in range(num_random_inits):
    curr_model_name = '%s_%d' %(model_name, k+1)
    print('training "%s"' %(curr_model_name))
    
    # create model
    student_model = create_FCN(input_dim, output_dim, student_nn_depth, student_nn_width, curr_model_name)
    #student_model.summary()
    student_model.compile(optimizer="adam", loss="mse")
    
    # train model
    history = student_model.fit(train_datagen, epochs=num_train_iterations, validation_data=valid_datagen, verbose=0)
    
    # store learning curves
    curr_learning_curves = {}
    curr_learning_curves['num_batches'] = num_batches_per_epoch * np.arange(num_train_iterations)
    curr_learning_curves['num_samples'] = batch_size * curr_learning_curves['num_batches']
    curr_learning_curves['train_loss']  = np.array(history.history['loss'])
    curr_learning_curves['valid_loss']  = np.array(history.history['val_loss'])

    training_results[model_name].append(curr_learning_curves)
    
    
# collect all learning curves for the same model in the same matrix
learning_curve_matrix = np.zeros((len(training_results[model_name]), training_results[model_name][0]['valid_loss'].shape[0]))
for k, curr_learning_curves in enumerate(training_results[model_name]):
    learning_curve_matrix[k, :] = curr_learning_curves['valid_loss']
    
num_batches_vec = training_results[model_name][0]['num_batches']
num_samples_vec = training_results[model_name][0]['num_samples']

# show all learning curves
plt.figure(figsize=(15,25))
plt.subplots_adjust(left=0.06, bottom=0.06, right=0.94, top=0.94, hspace=0.1, wspace=0.1)
plt.suptitle(model_name, fontsize=24)

plt.subplot(4,1,1)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)
plt.ylim(0,ylim_max)

plt.subplot(4,1,2)
plt.plot(num_batches_vec, learning_curve_matrix.T, color='b', alpha=0.6)
plt.plot(num_batches_vec, learning_curve_matrix.mean(axis=0), color='k', label='average')
plt.plot(num_batches_vec, learning_curve_matrix.max(axis=0), color='r', label='worst')
plt.plot(num_batches_vec, learning_curve_matrix.min(axis=0), color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,3)
plt.plot(num_batches_vec[:100], learning_curve_matrix[:,:100].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[:100], learning_curve_matrix.mean(axis=0)[:100], color='k', label='average')
plt.plot(num_batches_vec[:100], learning_curve_matrix.max(axis=0)[:100], color='r', label='worst')
plt.plot(num_batches_vec[:100], learning_curve_matrix.min(axis=0)[:100], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20)

plt.subplot(4,1,4)
plt.plot(num_batches_vec[-100:], learning_curve_matrix[:,-100:].T, color='b', alpha=0.6)
plt.plot(num_batches_vec[-100:], learning_curve_matrix.mean(axis=0)[-100:], color='k', label='average')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.max(axis=0)[-100:], color='r', label='worst')
plt.plot(num_batches_vec[-100:], learning_curve_matrix.min(axis=0)[-100:], color='g', label='best')
plt.xlabel('batch index', fontsize=20)
plt.ylabel('MSE', fontsize=20)
plt.legend(fontsize=20);