ERM-IRM/IRM_AC_CMNIST_final.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instructions to run this notebook\n",
    "\n",
    "In this notebook, we present the comparisons for AC-MNIST: Anti-causal colored MNIST.\n",
    "Run all the cells sequentially from top to bottom; we have commented the cells to help the reader."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import argparse\n",
    "import IPython.display as display\n",
    "import matplotlib.pyplot as plt\n",
    "from tensorflow import keras\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "from sklearn.utils import shuffle\n",
    "import pandas as pd\n",
    "tf.compat.v1.enable_eager_execution()\n",
    "import cProfile\n",
    "from sklearn.model_selection import train_test_split\n",
    "import copy as cp\n",
    "from sklearn.model_selection import KFold\n",
    "from datetime import date\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_construct import * ## contains functions for constructing data \n",
    "from IRM_methods import *    ## contains IRM and ERM methods "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample complexity on AC-CMNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "n_trial =10\n",
    "n_tr_list = [1000, 5000, 10000, 30000, 60000] # list of training sample sizes\n",
    "\n",
    "k=0\n",
    "K = len(n_tr_list)\n",
    "ERM_model_acc = np.zeros((K,n_trial))\n",
    "ERM_model_acc_nb = np.zeros((K,n_trial))\n",
    "IRM_model_acc = np.zeros((K,n_trial))\n",
    "IRM_model_acc_v = np.zeros((K,n_trial))\n",
    "\n",
    "ERM_model_acc1 = np.zeros((K,n_trial))\n",
    "ERM_model_acc1_nb = np.zeros((K,n_trial))\n",
    "IRM_model_acc1 = np.zeros((K,n_trial))\n",
    "IRM_model_acc1_v = np.zeros((K,n_trial))\n",
    "IRM_model_ind_v = np.zeros((K,n_trial))\n",
    "\n",
    "ERM_model_acc_av = np.zeros(K)\n",
    "ERM_model_acc_av_nb = np.zeros(K)\n",
    "IRM_model_acc_av = np.zeros(K)\n",
    "IRM_model_acc_av_v = np.zeros(K)\n",
    "\n",
    "\n",
    "ERM_model_acc_av1 = np.zeros(K)\n",
    "ERM_model_acc_av1_nb = np.zeros(K)\n",
    "IRM_model_acc_av1 = np.zeros(K)\n",
    "IRM_model_acc_av1_v = np.zeros(K)\n",
    "\n",
    "list_params = []\n",
    "for n_tr in n_tr_list:\n",
    "    print (\"tr\" + str(n_tr))\n",
    "#     print (\"start\")\n",
    "    t_start = time.time()\n",
    "    for trial in range(n_trial):\n",
    "        print (\"trial \" + str(trial))\n",
    "        n_e=2\n",
    "        p_color_list = [0.2, 0.1]\n",
    "        p_label_list = [0.25]*n_e\n",
    "        D = assemble_data_mnist_child(n_tr) # initialize mnist digits data object\n",
    "\n",
    "        D.create_training_data(n_e, p_color_list, p_label_list) # creates the training environments\n",
    "\n",
    "        p_label_test = 0.25 # probability of switching pre-label in test environment\n",
    "        p_color_test = 0.9  # probability of switching the final label to obtain the color index in test environment\n",
    "\n",
    "        D.create_testing_data(p_color_test, p_label_test, n_e)  # sets up the testing environment\n",
    "        (num_examples_environment,length, width, height) = D.data_tuple_list[0][0].shape # attributes of the data\n",
    "        num_classes = len(np.unique(D.data_tuple_list[0][1])) # number of classes in the data\n",
    "\n",
    "        model_erm =  keras.Sequential([\n",
    "                keras.layers.Flatten(input_shape=(length,width,height)),\n",
    "                keras.layers.Dense(390, activation = 'relu',kernel_regularizer=keras.regularizers.l2(0.0011)),\n",
    "                keras.layers.Dense(390, activation='relu',kernel_regularizer=keras.regularizers.l2(0.0011)),\n",
    "                keras.layers.Dense(2, activation='softmax')\n",
    "        ])\n",
    "        \n",
    "        num_epochs = 100\n",
    "        batch_size = 512\n",
    "        learning_rate = 4.9e-4\n",
    "        erm_model1 = standard_erm_model(model_erm, num_epochs, batch_size, learning_rate)\n",
    "        erm_model1.fit(D.data_tuple_list)\n",
    "        erm_model1.evaluate(D.data_tuple_test)\n",
    "        print (\"Training accuracy:\" + str(erm_model1.train_acc))\n",
    "        print (\"Testing accuracy:\" + str(erm_model1.test_acc))\n",
    "        \n",
    "        ERM_model_acc[k][trial] = erm_model1.test_acc\n",
    "        ERM_model_acc1[k][trial] = erm_model1.train_acc\n",
    "\n",
    "\n",
    "        gamma_list = [10000,33000, 66000, 100000.0]\n",
    "       \n",
    "         \n",
    "        index=0\n",
    "        best_err = 1e6\n",
    "        train_list =[]\n",
    "        val_list = []\n",
    "        test_list = []\n",
    "        for gamma_new in gamma_list:\n",
    "\n",
    "            model_irm = keras.Sequential([\n",
    "                                keras.layers.Flatten(input_shape=(length,width,height)),\n",
    "                                keras.layers.Dense(390, activation = 'relu',kernel_regularizer=keras.regularizers.l2(0.0011)),\n",
    "                                keras.layers.Dense(390, activation='relu',kernel_regularizer=keras.regularizers.l2(0.0011)),\n",
    "                                keras.layers.Dense(num_classes)\n",
    "                        ])\n",
    "            batch_size       = 512\n",
    "            steps_max        = 1000\n",
    "            steps_threshold  = 190  ## threshold after which gamma_new is used\n",
    "            learning_rate    = 4.9e-4\n",
    "\n",
    "\n",
    "            irm_model1 = irm_model(model_irm, learning_rate, batch_size, steps_max, steps_threshold, gamma_new)\n",
    "            irm_model1.fit(D.data_tuple_list)\n",
    "            irm_model1.evaluate(D.data_tuple_test)\n",
    "            error_val = 1-irm_model1.val_acc\n",
    "            train_list.append(irm_model1.train_acc)\n",
    "            val_list.append(irm_model1.val_acc)\n",
    "            test_list.append(irm_model1.test_acc)\n",
    "            if(error_val<best_err):\n",
    "                index_best =index\n",
    "                best_err = error_val\n",
    "            index= index+1\n",
    "\n",
    "        print (\"Training accuracy:\" + str(train_list[index_best]))\n",
    "        print (\"Validation accuracy:\" + str(val_list[index_best]))\n",
    "        print (\"Testing accuracy:\" + str(test_list[index_best]))\n",
    "\n",
    "        IRM_model_acc_v[k][trial]  = test_list[index_best]\n",
    "        IRM_model_acc1_v[k][trial] = train_list[index_best]\n",
    "        IRM_model_ind_v[k][trial]  = index_best\n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "    IRM_model_acc_av_v[k] = np.mean(IRM_model_acc_v[k])\n",
    "    list_params.append([n_tr,\"IRMv_test\", np.mean(IRM_model_acc_v[k]),np.std(IRM_model_acc_v[k])])\n",
    "\n",
    "    ERM_model_acc_av[k] = np.mean(ERM_model_acc[k])\n",
    "    list_params.append([n_tr,\"ERM_test\", np.mean(ERM_model_acc[k]),np.std(ERM_model_acc[k])])\n",
    "\n",
    "    \n",
    "    IRM_model_acc_av1_v[k] = np.mean(IRM_model_acc1_v[k])\n",
    "    list_params.append([n_tr,\"IRMv_train\", np.mean(IRM_model_acc1_v[k]),np.std(IRM_model_acc1_v[k])])\n",
    "    \n",
    "    ERM_model_acc_av1[k] = np.mean(ERM_model_acc1[k])\n",
    "    list_params.append([n_tr, \"ERM_train\", np.mean(ERM_model_acc1[k]),np.std(ERM_model_acc1[k])])\n",
    "\n",
    "\n",
    "\n",
    "    k=k+1\n",
    "\n",
    "    t_end = time.time()\n",
    "    print(\"total time: \" + str(t_end-t_start))\n",
    "\n",
    "\n",
    "results = pd.DataFrame(list_params, columns= [\"Sample\",\"Method\", \"Performance\", \"Sdev\"])\n",
    "\n",
    "ideal_error = np.ones(5)*0.25\n",
    "print (\"end\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.xlabel(\"Number of samples\", fontsize=16)\n",
    "plt.ylabel(\"Test error\", fontsize=16)\n",
    "plt.plot(n_tr_list, 1-ERM_model_acc_av, \"-r\", marker=\"+\", label=\"ERM\")\n",
    "plt.plot(n_tr_list, 1-IRM_model_acc_av_v, \"-b\", marker=\"s\",label=\"IRMv1\")\n",
    "plt.plot(n_tr_list, ideal_error, \"-g\", marker=\"x\", label=\"Optimal invariant\")\n",
    "plt.legend(loc=\"upper left\", fontsize=18)\n",
    "plt.ylim(-0.01,0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results"
   ]
  }
 ],
 "metadata": {
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}