ml_theory_mi.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FqAz7EJJbg0a"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "#%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "id": "1HyvLqsCi-re"
   },
   "outputs": [],
   "source": [
    "#@title Mutual Information Functions\n",
    "\n",
    "# These functions take LOG-probabilities\n",
    "\n",
    "# Samples have shape (n, C) where n is the number of samples and C it the number of classes\n",
    "def joint(samples_x, samples_y, samples_z):\n",
    "    return tf.math.reduce_logsumexp(  \n",
    "                samples_x[:, :, tf.newaxis, tf.newaxis]\n",
    "              + samples_y[:, tf.newaxis, :, tf.newaxis]\n",
    "              + samples_z[:, tf.newaxis, tf.newaxis, :],\n",
    "           axis=0) - tf.math.log(tf.cast(samples_x.shape[0], tf.float32))\n",
    "\n",
    "# Computes I(X ; Y)\n",
    "# p_xy has shape (C, C)\n",
    "def I_XY(p_xy):\n",
    "    p_x = tf.math.reduce_logsumexp(p_xy, 1, keepdims=True)\n",
    "    p_y = tf.math.reduce_logsumexp(p_xy, 0, keepdims=True)\n",
    "    return tf.reduce_sum(tf.math.exp(p_xy) \n",
    "                         * tf.clip_by_value(p_xy - p_x - p_y, -1e20, 1e20)\n",
    "                         / np.log(2))\n",
    "\n",
    "# Computes I(X ; Y | Z)\n",
    "# p_xyz has shape (C, C, C)\n",
    "def I_XY_Z(p_xyz):\n",
    "    p_z = tf.math.reduce_logsumexp(p_xyz, [0, 1], keepdims=True)\n",
    "    p_xy_z = p_xyz - p_z\n",
    "    p_x_z = tf.math.reduce_logsumexp(p_xyz, 1, keepdims=True) - p_z\n",
    "    p_y_z = tf.math.reduce_logsumexp(p_xyz, 0, keepdims=True) - p_z\n",
    "    return tf.reduce_sum(tf.math.exp(p_xyz) \n",
    "                         * tf.clip_by_value(p_xy_z - p_x_z - p_y_z, -1e10, 1e10)\n",
    "                         / np.log(2))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "1OgovNfCApvM"
   },
   "outputs": [],
   "source": [
    "#@title Network Architectures\n",
    "\n",
    "# Pre-activation version of ResNet\n",
    "def resnetBlock(inputs, filters, stride):\n",
    "    \n",
    "    x = tf.keras.layers.BatchNormalization(axis=1)(inputs)\n",
    "    x = tf.keras.layers.Activation('relu')(x)\n",
    "    if stride != 1:\n",
    "        shortcut = tf.keras.layers.Conv2D(filters, 1, stride, 'same')(x)\n",
    "    else:\n",
    "        shortcut = inputs\n",
    "    x = tf.keras.layers.Conv2D(filters, 3, stride, 'same')(x)\n",
    "    x = tf.keras.layers.BatchNormalization(axis=1)(x)\n",
    "    x = tf.keras.layers.Activation('relu')(x)\n",
    "    x = tf.keras.layers.Conv2D(filters, 3, 1, 'same')(x)\n",
    "    x = tf.keras.layers.add([x, shortcut])\n",
    "    return x\n",
    "\n",
    "def resnet18(inputs, C, n_per_block=2, n_layers=4):\n",
    "    \n",
    "    x = tf.keras.layers.Conv2D(64, 3, 1, 'same')(inputs)\n",
    "    for downsize, n_blocks, filters in [(1, n_per_block, 64), \n",
    "                                        (2, n_per_block, 128), \n",
    "                                        (2, n_per_block, 256), \n",
    "                                        (2, n_per_block, 512)][:n_layers]:\n",
    "        strides = [downsize] + [1] * (n_blocks - 1)\n",
    "        for stride in strides:\n",
    "            x = resnetBlock(x, filters, stride)\n",
    "    if n_layers == 0:\n",
    "        x = tf.keras.layers.Activation('relu')(x)\n",
    "    #x = tf.keras.layers.GlobalAveragePooling2D()(x)\n",
    "    x = tf.keras.layers.AveragePooling2D(4, padding='valid')(x)\n",
    "    x = tf.keras.layers.Flatten()(x)\n",
    "    x = tf.keras.layers.Dense(C)(x)\n",
    "    return x  \n",
    "    \n",
    "\n",
    "def simpleCNN(inputs, C, n_layers=3):\n",
    "    for _ in range(n_layers):\n",
    "        x = tf.keras.layers.Conv2D(32, 3, 1, 'same')(inputs)\n",
    "        x = tf.keras.layers.BatchNormalization(axis=1)(x)\n",
    "        x = tf.keras.layers.Activation('relu')(x)\n",
    "        x = tf.keras.layers.Conv2D(64, 3, 2, 'same')(x)\n",
    "        x = tf.keras.layers.BatchNormalization(axis=1)(x)\n",
    "        x = tf.keras.layers.Activation('relu')(x)\n",
    "    x = tf.keras.layers.AveragePooling2D(4, padding='valid')(x)\n",
    "    x = tf.keras.layers.Flatten()(x)\n",
    "    x = tf.keras.layers.Dense(C)(x)\n",
    "    return x\n",
    "\n",
    "\n",
    "def simpleCNNGal(inputs, C):\n",
    "    x = tf.keras.layers.Conv2D(32, 3, 2, 'same', activation='relu')(inputs)\n",
    "    x = tf.keras.layers.Conv2D(32, 3, 2, 'same', activation='relu')(x)\n",
    "    x = tf.keras.layers.Conv2D(32, 3, 1, 'same', activation='relu')(x)\n",
    "    x = tf.keras.layers.Conv2D(32, 3, 1, 'same', activation='relu')(x)\n",
    "    x = tf.keras.layers.Flatten()(x)\n",
    "    x = tf.keras.layers.Dense(512, activation='relu')(x)\n",
    "    x = tf.keras.layers.Dense(512, activation='relu')(x)\n",
    "    x = tf.keras.layers.Dense(C)(x)\n",
    "    return x\n",
    "\n",
    "def linear(x, C):\n",
    "    x = tf.keras.layers.Flatten()(x)\n",
    "    x = tf.keras.layers.Dense(C)(x)\n",
    "    return x\n",
    "\n",
    "def small(x, C):\n",
    "    x = tf.keras.layers.Flatten()(x)\n",
    "    x = tf.keras.layers.Dense(10, activation='relu')(x)\n",
    "    x = tf.keras.layers.Dense(C)(x)\n",
    "    return x\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "cellView": "code",
    "colab": {},
    "colab_type": "code",
    "id": "bavJnJ77KF65"
   },
   "outputs": [],
   "source": [
    "#@title Training and Evaluation Loops\n",
    "\n",
    "def transform(x, training=True):\n",
    "    x = tf.transpose(x, perm=[0, 2, 3, 1])\n",
    "    if training:\n",
    "        x = tf.image.resize_image_with_crop_or_pad(x, 40, 40)\n",
    "        x = tf.image.random_crop(x, [x.shape[0], 32, 32, 3])\n",
    "        x = tf.image.random_flip_left_right(x)\n",
    "    x = tf.image.per_image_standardization(x)\n",
    "    x = tf.transpose(x, perm=[0, 3, 1, 2])\n",
    "    return x\n",
    "\n",
    "# n is number of samples in the dataset, C is number of classes\n",
    "\n",
    "def evaluate_simple(C, data, model, model_simple, hard=True):\n",
    "    \n",
    "    print('Testing mutual information')\n",
    "    \n",
    "    batches = data.batch(1000)\n",
    "    \n",
    "    y_true = []\n",
    "    y_model = []\n",
    "    y_simple = []\n",
    "    for x, y in batches:\n",
    "        \n",
    "        x = transform(x, training=False)\n",
    "        \n",
    "        y_true += [tf.math.log(tf.one_hot(tf.squeeze(y, axis=1), C))]\n",
    "        if hard:\n",
    "            y_model += [tf.math.log(tf.one_hot(tf.math.argmax(model(x, training=False), axis=1), C))]\n",
    "            y_simple += [tf.math.log(tf.one_hot(tf.math.argmax(model_simple(x), axis=1), C))]\n",
    "        else:\n",
    "            y_model += [tf.nn.log_softmax(model(x))]\n",
    "            y_simple += [tf.nn.log_softmax(model_simple(x))]\n",
    "            \n",
    "    y_true = tf.concat(y_true, 0)\n",
    "    y_model = tf.concat(y_model, 0)\n",
    "    y_simple = tf.concat(y_simple, 0)\n",
    "    \n",
    "    p_joint = joint(y_true, y_model, y_simple)\n",
    "    I_model_true_simple = I_XY_Z(p_joint)\n",
    "        \n",
    "    I_model_true = I_XY(tf.math.reduce_logsumexp(p_joint, axis=2))\n",
    "    I_model_simple = I_XY_Z(tf.transpose(p_joint, perm=[2, 0, 1]))\n",
    "    \n",
    "    I_simple_true = I_XY(tf.math.reduce_logsumexp(p_joint, axis=1))\n",
    "    print(f'I(L ; Y) = {I_simple_true}')\n",
    "    \n",
    "    print(f'I(M ; Y) = {I_model_true}, '\n",
    "          f'I(M ; Y | L) = {I_model_true_simple}')\n",
    "        \n",
    "    print('Done')\n",
    "        \n",
    "    return I_model_true_simple, I_model_true, I_model_simple\n",
    "        \n",
    "def train_simple(n, C, data, model, model_simple, epochs, batch_size, learning_rate, savename=None, save_frequency=100):\n",
    "    \n",
    "    #optimizer = tf.optimizers.SGD(learning_rate=learning_rate, momentum=0.9)\n",
    "    optimizer = tf.optimizers.Adam(learning_rate=1e-4)\n",
    "    \n",
    "    avg_loss = tf.metrics.Mean()\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "        \n",
    "        if savename is not None and epoch % save_frequency == 0:\n",
    "            model_simple.save_weights(savename)\n",
    "       \n",
    "        print(f'Training simple model: epoch {epoch}')\n",
    "        batches = data.shuffle(n).batch(batch_size)\n",
    "        \n",
    "        for i, (x, y) in enumerate(batches):\n",
    "            \n",
    "            x = transform(x, training=True)\n",
    "            \n",
    "            y_true = tf.math.log(tf.one_hot(tf.squeeze(y, axis=1), C))\n",
    "            y_model = tf.nn.log_softmax(model(x, training=False))\n",
    "            with tf.GradientTape() as tape:\n",
    "                y_simple = tf.nn.log_softmax(model_simple(x))\n",
    "                p_joint = joint(y_true, y_model, y_simple)\n",
    "                #p_joint = joint(y_simple, y_true, y_model)\n",
    "                loss = I_XY_Z(p_joint)\n",
    "                #loss += I_XY_Z(tf.transpose(p_joint, perm=[0, 2, 1]))\n",
    "            gradients = tape.gradient(loss, model_simple.trainable_variables)\n",
    "            optimizer.apply_gradients(zip(gradients, model_simple.trainable_variables))\n",
    "            \n",
    "            avg_loss(loss)\n",
    "            \n",
    "            if (i + 1) % 10 == 0:\n",
    "                tf.print(f'Epoch {epoch}, '\n",
    "                         f'{i * batch_size/n * 100:.1f}%, '\n",
    "                         f'Loss {avg_loss.result():.6f}')\n",
    "                \n",
    "        avg_loss.reset_states()\n",
    "              \n",
    "    print('Done')\n",
    "    \n",
    "def train_imitate(n, C, data_train, data_test, model, model_simple,\n",
    "                 epochs=200, batch_size=128, learning_rate=0.1):\n",
    "    \n",
    "    criterion = tf.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "    optimizer = tf.optimizers.SGD(learning_rate=learning_rate, momentum=0.9)\n",
    "    \n",
    "    avg_loss = tf.metrics.Mean()\n",
    "    accuracy = tf.metrics.SparseCategoricalAccuracy()\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "    \n",
    "        batches = data_train.shuffle(n).batch(batch_size)\n",
    "        \n",
    "        for i, (x, y) in enumerate(batches):\n",
    "            \n",
    "            x = transform(x, training=True)\n",
    "            y_model = tf.argmax(model(x), axis=1)\n",
    "            \n",
    "            with tf.GradientTape() as tape:\n",
    "                y_simple = model_simple(x)\n",
    "                loss = criterion(y_model, y_simple)\n",
    "            gradients = tape.gradient(loss, model_simple.trainable_variables)\n",
    "            optimizer.apply_gradients(zip(gradients, model_simple.trainable_variables))\n",
    "            \n",
    "            avg_loss(loss)\n",
    "            accuracy(y_model, y_simple)\n",
    "              \n",
    "            if (i + 1) % 10 == 0:\n",
    "                tf.print(f'Epoch {epoch}, '\n",
    "                         f'{i * batch_size/n * 100:.1f}%, '\n",
    "                         f'Loss {avg_loss.result():.3f}, '\n",
    "                         f'Accuracy {accuracy.result() * 100:.2f}')\n",
    "                \n",
    "        test_accuracy = evaluate(data_test, model_simple)\n",
    "        print(f'Test accuracy {test_accuracy * 100:.2f}%')\n",
    "    \n",
    "    \n",
    "\n",
    "def evaluate(data, model):\n",
    "    \n",
    "    print('Testing accuracy')\n",
    "    \n",
    "    batches = data.batch(1000)\n",
    "    accuracy = tf.metrics.SparseCategoricalAccuracy()\n",
    "    \n",
    "    for x, y in batches:\n",
    "        x = transform(x, training=False)\n",
    "        y_model = model(x, training=False)\n",
    "        accuracy(y, y_model)\n",
    "        \n",
    "    print('Done')\n",
    "        \n",
    "    return accuracy.result()\n",
    "    \n",
    "import datetime            \n",
    "\n",
    "def train(n, C, data_train, data_test, model, model_simple, \n",
    "          epochs=200, batch_size=128, learning_rate=0.1,\n",
    "          save_dir='checkpoints',\n",
    "          simple_epochs=2, simple_batch_size=1000, simple_learning_rate=0.1,\n",
    "          simple_frequency=20,\n",
    "          save_frequency=100):\n",
    "    \n",
    "    current_time = datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "    log_dir = 'logs/' + current_time\n",
    "    summary_I_model_true_simple = tf.summary.create_file_writer(log_dir + '/I_model_true_simple')\n",
    "    summary_I_model_true = tf.summary.create_file_writer(log_dir + '/I_model_true')\n",
    "    summary_I_model_simple = tf.summary.create_file_writer(log_dir + '/I_model_simple')\n",
    "    \n",
    "    criterion = tf.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "    optimizer = tf.optimizers.SGD(learning_rate=learning_rate, momentum=0.9)\n",
    "    #optimizer = tf.optimizers.SGD(learning_rate=learning_rate)\n",
    "    \n",
    "    avg_loss = tf.metrics.Mean()\n",
    "    accuracy = tf.metrics.SparseCategoricalAccuracy()\n",
    "    \n",
    "    steps = 0\n",
    "    for epoch in range(epochs):\n",
    "    \n",
    "        batches = data_train.shuffle(n).batch(batch_size)\n",
    "        \n",
    "        for i, (x, y) in enumerate(batches):\n",
    "            \n",
    "            x = transform(x, training=True)\n",
    "              \n",
    "            with tf.GradientTape() as tape:\n",
    "                y_model = model(x)\n",
    "                loss = criterion(y, y_model)\n",
    "            gradients = tape.gradient(loss, model.trainable_variables)\n",
    "            optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "            \n",
    "            avg_loss(loss)\n",
    "            accuracy(y, y_model)\n",
    "              \n",
    "            if (i + 1) % 10 == 0:\n",
    "                tf.print(f'Epoch {epoch}, '\n",
    "                         f'{i * batch_size/n * 100:.1f}%, '\n",
    "                         f'Loss {avg_loss.result():.3f}, '\n",
    "                         f'Accuracy {accuracy.result() * 100:.2f}')\n",
    "                \n",
    "            if simple_frequency is not None and (i + 1) % simple_frequency == 0:\n",
    "                \n",
    "                train_simple(n, C, data_train, model, model_simple, \n",
    "                             epochs=simple_epochs, batch_size=simple_batch_size, \n",
    "                             learning_rate=simple_learning_rate)\n",
    "                \n",
    "                I_model_true_simple, I_model_true, I_model_simple \\\n",
    "                            = evaluate_simple(C, data_test, model, model_simple, hard=True)\n",
    "                with summary_I_model_true_simple.as_default():\n",
    "                        tf.summary.scalar(f'Mutual Information', \n",
    "                                          I_model_true_simple, step=steps)\n",
    "                with summary_I_model_true.as_default():\n",
    "                        tf.summary.scalar(f'Mutual Information', \n",
    "                                          I_model_true, step=steps)\n",
    "                with summary_I_model_simple.as_default():\n",
    "                        tf.summary.scalar(f'Mutual Information', \n",
    "                                          I_model_simple, step=steps)\n",
    "            if steps % save_frequency == 0 and save_dir is not None:\n",
    "                model.save_weights(f'{save_dir}/{steps}.h5')\n",
    "            \n",
    "            steps += 1\n",
    "        \n",
    "        test_accuracy = evaluate(data_test, model)\n",
    "        print(f'Test accuracy {test_accuracy * 100:.2f}%')\n",
    "        \n",
    "        avg_loss.reset_states()\n",
    "        accuracy.reset_states()\n",
    "\n",
    "def train_from_checkpoints(n, C, data_train, data_test, model, models_simple, \n",
    "          checkpoint_dir=['checkpoints'],\n",
    "          simple_epochs=0, simple_batch_size=1000, simple_learning_rate=0.1,\n",
    "          simple_frequency=20,\n",
    "          save_frequency=100,\n",
    "          hard=True):\n",
    "    \n",
    "    current_time = datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "    log_dir = 'logs/' + current_time\n",
    "    summary_I_model_true_simple = tf.summary.create_file_writer(log_dir + '/I_model_true_simple')\n",
    "    summary_I_model_true = tf.summary.create_file_writer(log_dir + '/I_model_true')\n",
    "    summary_I_model_simple = tf.summary.create_file_writer(log_dir + '/I_model_simple')\n",
    "    \n",
    "    for steps in range(0, 20000, save_frequency):\n",
    "        \n",
    "        if simple_epochs > 0:\n",
    "            for i, model_simple in enumerate(models_simple):\n",
    "                model.load_weights(f'{checkpoint_dir[i]}/{steps}.h5')\n",
    "                train_simple(n, C, data_train, model, model_simple, \n",
    "                             epochs=simple_epochs, batch_size=simple_batch_size, \n",
    "                             learning_rate=simple_learning_rate)\n",
    "                \n",
    "        for i, model_simple in enumerate(models_simple):\n",
    "            \n",
    "            model.load_weights(f'{checkpoint_dir[i]}/{steps}.h5')\n",
    "                \n",
    "            I_model_true_simple, I_model_true, I_model_simple \\\n",
    "                    = evaluate_simple(C, data_test, model, model_simple, hard=hard)\n",
    "            with summary_I_model_true_simple.as_default():\n",
    "                tf.summary.scalar(f'Mutual Information (simple model {i})', \n",
    "                                    I_model_true_simple, step=steps)\n",
    "            with summary_I_model_true.as_default():\n",
    "                tf.summary.scalar(f'Mutual Information (simple model {i})', \n",
    "                                    I_model_true, step=steps)\n",
    "            with summary_I_model_simple.as_default():\n",
    "                tf.summary.scalar(f'Mutual Information (simple model {i})', \n",
    "                                    I_model_simple, step=steps)\n",
    "            #model_simple.save_weights(f'{checkpoint_dir}/simple{i}/{steps}.h5')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "cellView": "both",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 54
    },
    "colab_type": "code",
    "id": "t0RO_nSjoVKV",
    "outputId": "6d220873-8ca2-48b0-fd74-3b3ffb87f608"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10000, 3, 32, 32)\n",
      "(50000, 1)\n"
     ]
    }
   ],
   "source": [
    "#@title CIFAR Processing\n",
    "\n",
    "C = 2\n",
    "\n",
    "def change_classes(x, y):\n",
    "    for i in [4, 6]:\n",
    "        index = y.flatten() != i\n",
    "        x, y = x[index], y[index]\n",
    "    for i in range(len(y)):\n",
    "        if y[i, 0] in [2, 3, 5, 7]:\n",
    "            y[i, 0] = 0\n",
    "        elif y[i, 0] in [0, 1, 8, 9]:\n",
    "            y[i, 0] = 1\n",
    "    return x, y\n",
    "\n",
    "# planes, cars, birds, cats, deer vs\n",
    "# ships, trucks, frogs, dogs, horses\n",
    "def change_classes_hard(x, y):\n",
    "    for i in range(len(y)):\n",
    "        if y[i, 0] < 5:\n",
    "            y[i, 0] = 0\n",
    "        else:\n",
    "            y[i, 0] = 1\n",
    "    return x, y\n",
    "\n",
    "tf.keras.backend.set_image_data_format('channels_first')\n",
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()\n",
    "x_train, y_train = change_classes_hard(x_train, y_train)\n",
    "x_test, y_test = change_classes_hard(x_test, y_test)\n",
    "n = len(x_train)\n",
    "x_train = np.divide(x_train, 255.0, dtype=np.float32)\n",
    "x_test = np.divide(x_test, 255.0, dtype=np.float32)\n",
    "print(x_test.shape)\n",
    "print(y_train.shape)\n",
    "data_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
    "data_test = tf.data.Dataset.from_tensor_slices((x_test, y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 109632
    },
    "colab_type": "code",
    "id": "HEExShH_VMjg",
    "outputId": "e06ead08-3d13-4fbd-e794-a0633fbe213a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing mutual information\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0521 19:19:23.914018 140424419575552 deprecation.py:323] From /home/ubuntu/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/ops/image_ops_impl.py:1444: div (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Deprecated in favor of operator or tf.math.divide.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I(L ; Y) = 0.5341752767562866\n",
      "I(M ; Y) = -8.599133138886828e-08, I(M ; Y | L) = 0.0\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.048815175890922546, I(M ; Y | L) = 0.0029210764914751053\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341752767562866\n",
      "I(M ; Y) = 0.07309546321630478, I(M ; Y | L) = 0.004723785445094109\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755151748657\n",
      "I(M ; Y) = 0.08379804342985153, I(M ; Y | L) = 0.004938877187669277\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755151748657\n",
      "I(M ; Y) = 0.10661014914512634, I(M ; Y | L) = 0.0073980167508125305\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341752767562866\n",
      "I(M ; Y) = 0.14322562515735626, I(M ; Y | L) = 0.012084584683179855\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755151748657\n",
      "I(M ; Y) = 0.14822795987129211, I(M ; Y | L) = 0.012810492888092995\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.534175455570221\n",
      "I(M ; Y) = 0.1581883728504181, I(M ; Y | L) = 0.011463411152362823\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.17031848430633545, I(M ; Y | L) = 0.015135511755943298\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755151748657\n",
      "I(M ; Y) = 0.2080085426568985, I(M ; Y | L) = 0.0145859494805336\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.19320304691791534, I(M ; Y | L) = 0.014452068135142326\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341752767562866\n",
      "I(M ; Y) = 0.24050968885421753, I(M ; Y | L) = 0.01811784692108631\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.534175455570221\n",
      "I(M ; Y) = 0.28653401136398315, I(M ; Y | L) = 0.021437302231788635\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.3148871064186096, I(M ; Y | L) = 0.028575513511896133\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755151748657\n",
      "I(M ; Y) = 0.30502402782440186, I(M ; Y | L) = 0.025337742641568184\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.31672126054763794, I(M ; Y | L) = 0.029485926032066345\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.3408617377281189, I(M ; Y | L) = 0.029478907585144043\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341752171516418\n",
      "I(M ; Y) = 0.3651660680770874, I(M ; Y | L) = 0.036779988557100296\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.534175455570221\n",
      "I(M ; Y) = 0.3547340929508209, I(M ; Y | L) = 0.028866076841950417\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341756343841553\n",
      "I(M ; Y) = 0.3770671784877777, I(M ; Y | L) = 0.03229144588112831\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.4183160364627838, I(M ; Y | L) = 0.044170528650283813\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341752767562866\n",
      "I(M ; Y) = 0.41313621401786804, I(M ; Y | L) = 0.0410517118871212\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341751575469971\n",
      "I(M ; Y) = 0.42937731742858887, I(M ; Y | L) = 0.046529654413461685\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341756343841553\n",
      "I(M ; Y) = 0.43859195709228516, I(M ; Y | L) = 0.051029272377491\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755151748657\n",
      "I(M ; Y) = 0.4518405497074127, I(M ; Y | L) = 0.05456014350056648\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341752171516418\n",
      "I(M ; Y) = 0.4300680160522461, I(M ; Y | L) = 0.046489521861076355\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.45363494753837585, I(M ; Y | L) = 0.048546407371759415\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.4562300145626068, I(M ; Y | L) = 0.056135933846235275\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341752767562866\n",
      "I(M ; Y) = 0.476650595664978, I(M ; Y | L) = 0.06294035911560059\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753363609314\n",
      "I(M ; Y) = 0.5028970241546631, I(M ; Y | L) = 0.07032409310340881\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755151748657\n",
      "I(M ; Y) = 0.4637395143508911, I(M ; Y | L) = 0.058695580810308456\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.4943780303001404, I(M ; Y | L) = 0.06374270468950272\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.534175455570221\n",
      "I(M ; Y) = 0.4932459890842438, I(M ; Y | L) = 0.0669635459780693\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341752767562866\n",
      "I(M ; Y) = 0.49054014682769775, I(M ; Y | L) = 0.0644359439611435\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341753959655762\n",
      "I(M ; Y) = 0.4911257326602936, I(M ; Y | L) = 0.06244949996471405\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755747795105\n",
      "I(M ; Y) = 0.5123159289360046, I(M ; Y | L) = 0.07737322151660919\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755151748657\n",
      "I(M ; Y) = 0.503166675567627, I(M ; Y | L) = 0.0716012716293335\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.5341755151748657\n",
      "I(M ; Y) = 0.49213913083076477, I(M ; Y | L) = 0.07214400917291641\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.534175455570221\n",
      "I(M ; Y) = 0.5034264326095581, I(M ; Y | L) = 0.07753753662109375\n",
      "Done\n",
      "Testing mutual information\n",
      "I(L ; Y) = 0.534175455570221\n",
      "I(M ; Y) = 0.5075928568840027, I(M ; Y | L) = 0.07660433650016785\n",
      "Done\n"
     ]
    }
   ],
   "source": [
    "# Full model training hyperparameters:\n",
    "EPOCHS = 60\n",
    "BATCH_SIZE = 128\n",
    "LEARNING_RATE = 0.01\n",
    "\n",
    "# Simple model training hyperparameters\n",
    "SIMPLE_EPOCHS = 2\n",
    "SIMPLE_BATCH_SIZE = 512\n",
    "SIMPLE_LEARNING_RATE = 0.1\n",
    "\n",
    "# How often to evaluate mutual information in units of gradient steps\n",
    "SIMPLE_FREQUENCY = 2\n",
    "\n",
    "\n",
    "tf.keras.backend.clear_session()\n",
    "\n",
    "shape = (3, 32, 32)\n",
    "inputs = tf.keras.Input(shape=shape)\n",
    "\n",
    "model = tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C))\n",
    "#model = tf.keras.applications.ResNet50(include_top=False, input_tensor=inputs, weights=None)\n",
    "#model = tf.keras.Model(inputs=inputs, outputs=dense(model.output, C))\n",
    "#model_simple = tf.keras.Model(inputs=inputs, outputs=linear(inputs, C))\n",
    "#model_simple = tf.keras.Model(inputs=inputs, outputs=small(inputs, C))\n",
    "#models_simple = [tf.keras.Model(inputs=inputs, outputs=linear(inputs, C))] + \\\n",
    "#                [tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C, n_layers=i)) for i in range(1, 4)]\n",
    "\n",
    "#model = tf.keras.Model(inputs=inputs, outputs=simpleCNNGal(inputs, C))\n",
    "#model_simple = tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C, n_layers=1))\n",
    "#model_simple.load_weights('resnet6-0/23400.h5')\n",
    "#model_simple = tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C, n_layers=0, n_per_block=1))\n",
    "#model_simple = tf.keras.Model(inputs=inputs, outputs=linear(inputs, C))\n",
    "model_simple = tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C))\n",
    "#model_simple.load_weights('cnn6-10.h5')\n",
    "#model.load_weights('resnet18-0/23400.h5')\n",
    "#model.load_weights('resnet18-0/234000.h5')\n",
    "model_simple.load_weights('resnet18-0/23400.h5')\n",
    "\n",
    "#train_imitate(n, C, data_train, data_test, model, model_simple, \n",
    "#              epochs=10, batch_size=128, learning_rate=0.01)\n",
    "#model_simple.load_weights('linear1/23400.h5')\n",
    "#train_simple(n, C, data_train, model, model_simple, \n",
    "#               epochs=100, batch_size=512, learning_rate=0.1,\n",
    "#               savename=None, save_frequency=None)\n",
    "\n",
    "#model_simple.save_weights('linear-10-test.h5')\n",
    "\n",
    "\n",
    "'''\n",
    "models_simple = [tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C, n_layers=1)) for _ in range(4)] + \\\n",
    "                [tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C, n_layers=1, n_per_block=1)) for _ in range(4)] + \\\n",
    "                [tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C, n_layers=0)) for _ in range(4)]\n",
    "    \n",
    "names = [f'cnn6-{i}' for i in range(4)] + \\\n",
    "        [f'cnn4-{i}' for i in range(4)] + \\\n",
    "        [f'cnn2-{i}' for i in range(4)]\n",
    "for model_simple, name in zip(models_simple, names):\n",
    "    model_simple.load_weights(f'{name}/23400.h5')\n",
    "    \n",
    "#for model_simple in models_simple:\n",
    "#    evaluate_simple(C, data_test, model, model_simple, hard=True)\n",
    "\n",
    "models = [tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C)) for _ in range(4)]\n",
    "\n",
    "checkpoints = ['resnet18-0', 'resnet18-1', 'resnet18-2', 'resnet18-3'] * 3\n",
    "'''\n",
    "\n",
    "'''\n",
    "n_runs = 4\n",
    "models = [tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C, n_layers=3)) for _ in range(n_runs)] + \\\n",
    "         [tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C, n_layers=2)) for _ in range(n_runs)] + \\\n",
    "         [tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C, n_layers=1)) for _ in range(n_runs)]\n",
    "\n",
    "names = [f'resnet14-{i}' for i in range(n_runs)] + \\\n",
    "        [f'resnet10-{i}' for i in range(n_runs)] + \\\n",
    "        [f'resnet6-{i}' for i in range(n_runs)]\n",
    "\n",
    "import os\n",
    "for name in names:\n",
    "    os.makedirs(name, exist_ok=True)\n",
    "\n",
    "for model, name in zip(models, names):\n",
    "    train(n, C, data_train, data_test, model, None,\n",
    "      epochs=EPOCHS, batch_size=BATCH_SIZE, learning_rate=LEARNING_RATE,\n",
    "      save_dir=name,\n",
    "      simple_frequency=None,\n",
    "      save_frequency=100)\n",
    "'''\n",
    "\n",
    "#model.load_weights('resnet18-0/23400.h5')\n",
    "\n",
    "#train_imitate(n, C, data_train, data_test, model, model_simple, \n",
    "#              epochs=10, batch_size=128, learning_rate=0.01)\n",
    "#train_simple(n, C, data_train, model, model_simple, \n",
    "#              epochs=10, batch_size=1000, learning_rate=0.1)\n",
    "\n",
    "#model = tf.keras.Model(inputs=inputs, outputs=resnet18(inputs, C))\n",
    "#model.load_weights('resnet18-0/0.h5')\n",
    "\n",
    "'''\n",
    "train(n, C, data_train, data_test, model, model_simple,\n",
    "      epochs=EPOCHS, batch_size=BATCH_SIZE, learning_rate=LEARNING_RATE,\n",
    "      save_dir='simpleCNN-fine',\n",
    "      simple_epochs=0, simple_batch_size=SIMPLE_BATCH_SIZE, \n",
    "      simple_learning_rate=SIMPLE_LEARNING_RATE, \n",
    "      simple_frequency=None,\n",
    "      save_frequency=10)\n",
    "\n",
    "#'''\n",
    "#'''\n",
    "train_from_checkpoints(n, C, data_train, data_test, model, [model_simple],\n",
    "      checkpoint_dir=['resnet18-1'],\n",
    "      simple_epochs=0, simple_batch_size=SIMPLE_BATCH_SIZE, \n",
    "      simple_learning_rate=SIMPLE_LEARNING_RATE, \n",
    "      simple_frequency=SIMPLE_FREQUENCY,\n",
    "      save_frequency=500,\n",
    "      hard=True)\n",
    "#'''\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "ml-theory-mi.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}