FRTrain_clean.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FR-Train on clean synthetic data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "from torch.autograd import Variable\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from argparse import Namespace\n",
    "\n",
    "from FRTrain_arch import Generator, DiscriminatorF, DiscriminatorR, weights_init_normal, test_model\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and process data (using clean y train label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a namespace object which contains some of the hyperparameters\n",
    "opt = Namespace(num_train=2000, num_val1=200, num_val2=500, num_test=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_train = opt.num_train\n",
    "num_val1 = opt.num_val1\n",
    "num_val2 = opt.num_val2\n",
    "num_test = opt.num_test\n",
    "\n",
    "X = np.load('X_synthetic.npy') # Input features\n",
    "y = np.load('y_synthetic.npy') # Labels\n",
    "s1 = np.load('s1_synthetic.npy') # Sensitive features\n",
    "\n",
    "X = torch.FloatTensor(X)\n",
    "y = torch.FloatTensor(y)\n",
    "s1 = torch.FloatTensor(s1)\n",
    "\n",
    "X_train = X[:num_train - num_val1]\n",
    "y_train = y[:num_train - num_val1] # Clean label\n",
    "s1_train = s1[:num_train - num_val1]\n",
    "\n",
    "\n",
    "X_val = X[num_train: num_train + num_val1]\n",
    "y_val = y[num_train: num_train + num_val1]\n",
    "s1_val = s1[num_train: num_train + num_val1]\n",
    "\n",
    "# Currently not used\n",
    "# X_val2 = X[num_train + num_val1 : num_train + num_val1 + num_val2]\n",
    "# y_val2 = y[num_train + num_val1 : num_train + num_val1 + num_val2]\n",
    "# s1_val2 = s1[num_train + num_val1 : num_train + num_val1 + num_val2]\n",
    "\n",
    "X_test = X[num_train + num_val1 + num_val2 : num_train + num_val1 + num_val2 + num_test]\n",
    "y_test = y[num_train + num_val1 + num_val2 : num_train + num_val1 + num_val2 + num_test]\n",
    "s1_test = s1[num_train + num_val1 + num_val2 : num_train + num_val1 + num_val2 + num_test]\n",
    "\n",
    "XS_train = torch.cat([X_train, s1_train.reshape((s1_train.shape[0], 1))], dim=1)\n",
    "XS_val = torch.cat([X_val, s1_val.reshape((s1_val.shape[0], 1))], dim=1)\n",
    "XS_test = torch.cat([X_test, s1_test.reshape((s1_test.shape[0], 1))], dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------- Number of Data -------------------------\n",
      "Train data : 1800, Validation data : 200, Test data : 1000 \n",
      "--------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "print(\"--------------------- Number of Data -------------------------\" )\n",
    "print(\n",
    "    \"Train data : %d, Validation data : %d, Test data : %d \"\n",
    "    % (len(y_train), len(y_val), len(y_test))\n",
    ")       \n",
    "print(\"--------------------------------------------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training with clean data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def train_model(train_tensors, val_tensors, test_tensors, train_opt, lambda_f, lambda_r, seed):\n",
    "    \"\"\"\n",
    "      Trains FR-Train by using the classes in FRTrain_arch.py.\n",
    "      \n",
    "      Args:\n",
    "        train_tensors: Training data.\n",
    "        val_tensors: Clean validation data.\n",
    "        test_tensors: Test data.\n",
    "        train_opt: Options for the training. It currently contains size of validation set, \n",
    "                number of epochs, generator/discriminator update ratio, and learning rates.\n",
    "        lambda_f: The tuning knob for L_2 (ref: FR-Train paper, Section 3.3).\n",
    "        lambda_r: The tuning knob for L_3 (ref: FR-Train paper, Section 3.3).\n",
    "        seed: An integer value for specifying torch random seed.\n",
    "        \n",
    "      Returns:\n",
    "        Information about the tuning knobs (lambda_f, lambda_r),\n",
    "        the test accuracy of the trained model, and disparate impact of the trained model.\n",
    "    \"\"\"\n",
    "    \n",
    "    XS_train = train_tensors.XS_train\n",
    "    y_train = train_tensors.y_train\n",
    "    s1_train = train_tensors.s1_train\n",
    "    \n",
    "    XS_val = val_tensors.XS_val\n",
    "    y_val = val_tensors.y_val\n",
    "    s1_val = val_tensors.s1_val\n",
    "    \n",
    "    XS_test = test_tensors.XS_test\n",
    "    y_test = test_tensors.y_test\n",
    "    s1_test = test_tensors.s1_test\n",
    "    \n",
    "    # Saves return values here\n",
    "    test_result = []\n",
    "    \n",
    "    val = train_opt.val # Number of data points in validation set\n",
    "    k = train_opt.k     # Updates ratio of generator and discriminator (1:k training).\n",
    "    n_epochs = train_opt.n_epochs  # Number of training epoch\n",
    "    \n",
    "    # Changes the input validation data to an appropriate shape for the training\n",
    "    XSY_val = torch.cat([XS_val, y_val.reshape((y_val.shape[0], 1))], dim=1)  \n",
    "\n",
    "    # The loss values of each component will be saved in the following lists. \n",
    "    # We can draw epoch-loss graph by the following lists, if necessary.\n",
    "    g_losses =[]\n",
    "    d_f_losses = []\n",
    "    d_r_losses = []\n",
    "    clean_test_result = []\n",
    "\n",
    "    bce_loss = torch.nn.BCELoss()\n",
    "\n",
    "    # Initializes generator and discriminator\n",
    "    generator = Generator()\n",
    "    discriminator_F = DiscriminatorF()\n",
    "    discriminator_R = DiscriminatorR()\n",
    "\n",
    "    # Initializes weights\n",
    "    torch.manual_seed(seed)\n",
    "    generator.apply(weights_init_normal)\n",
    "    discriminator_F.apply(weights_init_normal)\n",
    "    discriminator_R.apply(weights_init_normal)\n",
    "\n",
    "    optimizer_G = torch.optim.Adam(generator.parameters(), lr=train_opt.lr_g)\n",
    "    optimizer_D_F = torch.optim.SGD(discriminator_F.parameters(), lr=train_opt.lr_f)\n",
    "    optimizer_D_R = torch.optim.SGD(discriminator_R.parameters(), lr=train_opt.lr_r)\n",
    "\n",
    "    XSY_val_data = XSY_val[:val]\n",
    "\n",
    "    train_len = XS_train.shape[0]\n",
    "    val_len = XSY_val.shape[0]\n",
    "\n",
    "    # Ground truths using in Disriminator_R\n",
    "    Tensor = torch.FloatTensor\n",
    "    valid = Variable(Tensor(train_len, 1).fill_(1.0), requires_grad=False)\n",
    "    generated = Variable(Tensor(train_len, 1).fill_(0.0), requires_grad=False)\n",
    "    fake = Variable(Tensor(train_len, 1).fill_(0.0), requires_grad=False)\n",
    "    clean = Variable(Tensor(val_len, 1).fill_(1.0), requires_grad=False)\n",
    "    \n",
    "    r_weight = torch.ones_like(y_train, requires_grad=False).float()\n",
    "    r_ones = torch.ones_like(y_train, requires_grad=False).float()\n",
    "\n",
    "    for epoch in range(n_epochs):\n",
    "\n",
    "        # -------------------\n",
    "        #  Forwards Generator\n",
    "        # -------------------\n",
    "        if epoch % k == 0 or epoch < 500:\n",
    "            optimizer_G.zero_grad()\n",
    "\n",
    "        gen_y = generator(XS_train)\n",
    "        gen_data = torch.cat([XS_train, gen_y.detach().reshape((gen_y.shape[0], 1))], dim=1)\n",
    "        \n",
    "        # -----------------------------\n",
    "        #  Trains Fairness Discriminator\n",
    "        # -----------------------------\n",
    "\n",
    "        optimizer_D_F.zero_grad()\n",
    "\n",
    "        # Discriminator_F tries to distinguish the sensitive groups by using the output of the generator.\n",
    "        d_f_loss= bce_loss(discriminator_F(gen_y.detach()), s1_train)\n",
    "        d_f_loss.backward()\n",
    "        d_f_losses.append(d_f_loss)\n",
    "        optimizer_D_F.step()\n",
    "            \n",
    "\n",
    "        # ---------------------------------\n",
    "        #  Trains Robustness Discriminator\n",
    "        # ---------------------------------\n",
    "        optimizer_D_R.zero_grad()\n",
    "\n",
    "        # Discriminator_R tries to distinguish whether the input is from the validation data or the generated data from generator.\n",
    "        clean_loss = bce_loss(discriminator_R(XSY_val_data), clean)\n",
    "        poison_loss = bce_loss(discriminator_R(gen_data.detach()), fake)\n",
    "        d_r_loss = 0.5 * (clean_loss + poison_loss)\n",
    "\n",
    "        d_r_loss.backward()\n",
    "        d_r_losses.append(d_r_loss)\n",
    "        optimizer_D_R.step()\n",
    "\n",
    "\n",
    "        # ---------------------\n",
    "        #  Updates Generator\n",
    "        # ---------------------\n",
    "\n",
    "        # Loss measures generator's ability to fool the discriminators\n",
    "        if epoch < 500 :\n",
    "            g_loss = bce_loss((F.tanh(gen_y)+1)/2, (y_train+1)/2)\n",
    "            g_loss.backward()\n",
    "            g_losses.append(g_loss)\n",
    "            optimizer_G.step()\n",
    "    \n",
    "        elif epoch % k == 0:\n",
    "            r_decision = discriminator_R(gen_data)\n",
    "            r_gen = bce_loss(r_decision, generated)\n",
    "            \n",
    "            # ------------------------------\n",
    "            #  Re-weights using output of D_R\n",
    "            # ------------------------------\n",
    "            \n",
    "            if epoch % 100 == 0:\n",
    "                loss_ratio = (g_losses[-1]/d_r_losses[-1]).detach()\n",
    "                a = 1/(1+torch.exp(-(loss_ratio-3)))\n",
    "                b = 1-a\n",
    "                r_weight_tmp = r_decision.detach().squeeze()\n",
    "                r_weight = a * r_weight_tmp + b * r_ones\n",
    "\n",
    "            f_cost = F.binary_cross_entropy(discriminator_F(gen_y), s1_train, reduction=\"none\").squeeze()\n",
    "            g_cost = F.binary_cross_entropy_with_logits(gen_y.squeeze(), (y_train.squeeze()+1)/2, reduction=\"none\").squeeze()\n",
    "\n",
    "            f_gen = torch.mean(f_cost*r_weight)\n",
    "            g_loss = (1-lambda_f-lambda_r) * torch.mean(g_cost*r_weight) -  lambda_r * r_gen - lambda_f * f_gen\n",
    "\n",
    "            g_loss.backward()\n",
    "            optimizer_G.step()\n",
    "\n",
    "        g_losses.append(g_loss)\n",
    "\n",
    "        if epoch % 200 == 0:\n",
    "            print(\n",
    "                \"[Lambda_f: %1f] [Epoch %d/%d] [D_F loss: %f] [D_R loss: %f] [G loss: %f]\"\n",
    "                % (lambda_f, epoch, n_epochs, d_f_losses[-1], d_r_losses[-1], g_losses[-1])\n",
    "            )\n",
    "\n",
    "#     torch.save(generator.state_dict(), './FR-Train_on_clean_synthetic.pth')\n",
    "    tmp = test_model(generator, XS_test, y_test, s1_test)\n",
    "    test_result.append([lambda_f, lambda_r, tmp[0].item(), tmp[1]])\n",
    "\n",
    "    return test_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Lambda_f: 0.200000] [Epoch 0/4000] [D_F loss: 0.686341] [D_R loss: 0.716233] [G loss: 0.465272]\n",
      "[Lambda_f: 0.200000] [Epoch 200/4000] [D_F loss: 0.584361] [D_R loss: 0.680139] [G loss: 0.312735]\n",
      "[Lambda_f: 0.200000] [Epoch 400/4000] [D_F loss: 0.577822] [D_R loss: 0.673043] [G loss: 0.312282]\n",
      "[Lambda_f: 0.200000] [Epoch 600/4000] [D_F loss: 0.641887] [D_R loss: 0.658369] [G loss: 0.034481]\n",
      "[Lambda_f: 0.200000] [Epoch 800/4000] [D_F loss: 0.610650] [D_R loss: 0.642437] [G loss: 0.035496]\n",
      "[Lambda_f: 0.200000] [Epoch 1000/4000] [D_F loss: 0.605026] [D_R loss: 0.631439] [G loss: 0.035850]\n",
      "[Lambda_f: 0.200000] [Epoch 1200/4000] [D_F loss: 0.605569] [D_R loss: 0.623288] [G loss: 0.036699]\n",
      "[Lambda_f: 0.200000] [Epoch 1400/4000] [D_F loss: 0.606820] [D_R loss: 0.616102] [G loss: 0.037531]\n",
      "[Lambda_f: 0.200000] [Epoch 1600/4000] [D_F loss: 0.608355] [D_R loss: 0.609425] [G loss: 0.038328]\n",
      "[Lambda_f: 0.200000] [Epoch 1800/4000] [D_F loss: 0.609794] [D_R loss: 0.603060] [G loss: 0.039060]\n",
      "[Lambda_f: 0.200000] [Epoch 2000/4000] [D_F loss: 0.611048] [D_R loss: 0.596871] [G loss: 0.039729]\n",
      "[Lambda_f: 0.200000] [Epoch 2200/4000] [D_F loss: 0.612065] [D_R loss: 0.590783] [G loss: 0.040371]\n",
      "[Lambda_f: 0.200000] [Epoch 2400/4000] [D_F loss: 0.612848] [D_R loss: 0.584764] [G loss: 0.040978]\n",
      "[Lambda_f: 0.200000] [Epoch 2600/4000] [D_F loss: 0.613452] [D_R loss: 0.578767] [G loss: 0.041541]\n",
      "[Lambda_f: 0.200000] [Epoch 2800/4000] [D_F loss: 0.613902] [D_R loss: 0.572769] [G loss: 0.042094]\n",
      "[Lambda_f: 0.200000] [Epoch 3000/4000] [D_F loss: 0.614227] [D_R loss: 0.566771] [G loss: 0.042629]\n",
      "[Lambda_f: 0.200000] [Epoch 3200/4000] [D_F loss: 0.614463] [D_R loss: 0.560784] [G loss: 0.043138]\n",
      "[Lambda_f: 0.200000] [Epoch 3400/4000] [D_F loss: 0.614629] [D_R loss: 0.554812] [G loss: 0.043645]\n",
      "[Lambda_f: 0.200000] [Epoch 3600/4000] [D_F loss: 0.614743] [D_R loss: 0.548843] [G loss: 0.044144]\n",
      "[Lambda_f: 0.200000] [Epoch 3800/4000] [D_F loss: 0.614820] [D_R loss: 0.542872] [G loss: 0.044634]\n",
      "Test accuracy: 0.8669999837875366\n",
      "P(y_hat=1 | z=0) = 0.353, P(y_hat=1 | z=1) = 0.717\n",
      "P(y_hat=1 | y=1, z=0) = 0.799, P(y_hat=1 | y=1, z=1) = 0.900\n",
      "Disparate Impact ratio = 0.492\n",
      "[Lambda_f: 0.300000] [Epoch 0/4000] [D_F loss: 0.686341] [D_R loss: 0.716233] [G loss: 0.465272]\n",
      "[Lambda_f: 0.300000] [Epoch 200/4000] [D_F loss: 0.584361] [D_R loss: 0.680139] [G loss: 0.312735]\n",
      "[Lambda_f: 0.300000] [Epoch 400/4000] [D_F loss: 0.577822] [D_R loss: 0.673043] [G loss: 0.312282]\n",
      "[Lambda_f: 0.300000] [Epoch 600/4000] [D_F loss: 0.657889] [D_R loss: 0.659524] [G loss: -0.061782]\n",
      "[Lambda_f: 0.300000] [Epoch 800/4000] [D_F loss: 0.628343] [D_R loss: 0.644665] [G loss: -0.055544]\n",
      "[Lambda_f: 0.300000] [Epoch 1000/4000] [D_F loss: 0.621546] [D_R loss: 0.635200] [G loss: -0.055713]\n",
      "[Lambda_f: 0.300000] [Epoch 1200/4000] [D_F loss: 0.623105] [D_R loss: 0.627479] [G loss: -0.055150]\n",
      "[Lambda_f: 0.300000] [Epoch 1400/4000] [D_F loss: 0.624589] [D_R loss: 0.620657] [G loss: -0.054601]\n",
      "[Lambda_f: 0.300000] [Epoch 1600/4000] [D_F loss: 0.626184] [D_R loss: 0.614504] [G loss: -0.054003]\n",
      "[Lambda_f: 0.300000] [Epoch 1800/4000] [D_F loss: 0.627639] [D_R loss: 0.608623] [G loss: -0.053412]\n",
      "[Lambda_f: 0.300000] [Epoch 2000/4000] [D_F loss: 0.628897] [D_R loss: 0.602859] [G loss: -0.052875]\n",
      "[Lambda_f: 0.300000] [Epoch 2200/4000] [D_F loss: 0.629923] [D_R loss: 0.597171] [G loss: -0.052302]\n",
      "[Lambda_f: 0.300000] [Epoch 2400/4000] [D_F loss: 0.630724] [D_R loss: 0.591525] [G loss: -0.051740]\n",
      "[Lambda_f: 0.300000] [Epoch 2600/4000] [D_F loss: 0.631354] [D_R loss: 0.585882] [G loss: -0.051229]\n",
      "[Lambda_f: 0.300000] [Epoch 2800/4000] [D_F loss: 0.631835] [D_R loss: 0.580231] [G loss: -0.050697]\n",
      "[Lambda_f: 0.300000] [Epoch 3000/4000] [D_F loss: 0.632193] [D_R loss: 0.574576] [G loss: -0.050183]\n",
      "[Lambda_f: 0.300000] [Epoch 3200/4000] [D_F loss: 0.632463] [D_R loss: 0.568918] [G loss: -0.049712]\n",
      "[Lambda_f: 0.300000] [Epoch 3400/4000] [D_F loss: 0.632661] [D_R loss: 0.563290] [G loss: -0.049214]\n",
      "[Lambda_f: 0.300000] [Epoch 3600/4000] [D_F loss: 0.632804] [D_R loss: 0.557681] [G loss: -0.048721]\n",
      "[Lambda_f: 0.300000] [Epoch 3800/4000] [D_F loss: 0.632905] [D_R loss: 0.552067] [G loss: -0.048259]\n",
      "Test accuracy: 0.8659999966621399\n",
      "P(y_hat=1 | z=0) = 0.368, P(y_hat=1 | z=1) = 0.698\n",
      "P(y_hat=1 | y=1, z=0) = 0.809, P(y_hat=1 | y=1, z=1) = 0.894\n",
      "Disparate Impact ratio = 0.528\n",
      "[Lambda_f: 0.400000] [Epoch 0/4000] [D_F loss: 0.686341] [D_R loss: 0.716233] [G loss: 0.465272]\n",
      "[Lambda_f: 0.400000] [Epoch 200/4000] [D_F loss: 0.584361] [D_R loss: 0.680139] [G loss: 0.312735]\n",
      "[Lambda_f: 0.400000] [Epoch 400/4000] [D_F loss: 0.577822] [D_R loss: 0.673043] [G loss: 0.312282]\n",
      "[Lambda_f: 0.400000] [Epoch 600/4000] [D_F loss: 0.673050] [D_R loss: 0.660671] [G loss: -0.161856]\n",
      "[Lambda_f: 0.400000] [Epoch 800/4000] [D_F loss: 0.644337] [D_R loss: 0.646927] [G loss: -0.149258]\n",
      "[Lambda_f: 0.400000] [Epoch 1000/4000] [D_F loss: 0.636801] [D_R loss: 0.639172] [G loss: -0.150326]\n",
      "[Lambda_f: 0.400000] [Epoch 1200/4000] [D_F loss: 0.639130] [D_R loss: 0.632800] [G loss: -0.150139]\n",
      "[Lambda_f: 0.400000] [Epoch 1400/4000] [D_F loss: 0.641093] [D_R loss: 0.626902] [G loss: -0.149952]\n",
      "[Lambda_f: 0.400000] [Epoch 1600/4000] [D_F loss: 0.642729] [D_R loss: 0.621424] [G loss: -0.149628]\n",
      "[Lambda_f: 0.400000] [Epoch 1800/4000] [D_F loss: 0.644147] [D_R loss: 0.616179] [G loss: -0.149245]\n",
      "[Lambda_f: 0.400000] [Epoch 2000/4000] [D_F loss: 0.645335] [D_R loss: 0.611019] [G loss: -0.148872]\n",
      "[Lambda_f: 0.400000] [Epoch 2200/4000] [D_F loss: 0.646292] [D_R loss: 0.605914] [G loss: -0.148421]\n",
      "[Lambda_f: 0.400000] [Epoch 2400/4000] [D_F loss: 0.647040] [D_R loss: 0.600810] [G loss: -0.147958]\n",
      "[Lambda_f: 0.400000] [Epoch 2600/4000] [D_F loss: 0.647638] [D_R loss: 0.595671] [G loss: -0.147533]\n",
      "[Lambda_f: 0.400000] [Epoch 2800/4000] [D_F loss: 0.648104] [D_R loss: 0.590516] [G loss: -0.147064]\n",
      "[Lambda_f: 0.400000] [Epoch 3000/4000] [D_F loss: 0.648463] [D_R loss: 0.585375] [G loss: -0.146602]\n",
      "[Lambda_f: 0.400000] [Epoch 3200/4000] [D_F loss: 0.648744] [D_R loss: 0.580220] [G loss: -0.146188]\n",
      "[Lambda_f: 0.400000] [Epoch 3400/4000] [D_F loss: 0.648959] [D_R loss: 0.575061] [G loss: -0.145746]\n",
      "[Lambda_f: 0.400000] [Epoch 3600/4000] [D_F loss: 0.649122] [D_R loss: 0.569926] [G loss: -0.145305]\n",
      "[Lambda_f: 0.400000] [Epoch 3800/4000] [D_F loss: 0.649244] [D_R loss: 0.564802] [G loss: -0.144900]\n",
      "Test accuracy: 0.8569999933242798\n",
      "P(y_hat=1 | z=0) = 0.375, P(y_hat=1 | z=1) = 0.677\n",
      "P(y_hat=1 | y=1, z=0) = 0.819, P(y_hat=1 | y=1, z=1) = 0.866\n",
      "Disparate Impact ratio = 0.554\n",
      "[Lambda_f: 0.500000] [Epoch 0/4000] [D_F loss: 0.686341] [D_R loss: 0.716233] [G loss: 0.465272]\n",
      "[Lambda_f: 0.500000] [Epoch 200/4000] [D_F loss: 0.584361] [D_R loss: 0.680139] [G loss: 0.312735]\n",
      "[Lambda_f: 0.500000] [Epoch 400/4000] [D_F loss: 0.577822] [D_R loss: 0.673043] [G loss: 0.312282]\n",
      "[Lambda_f: 0.500000] [Epoch 600/4000] [D_F loss: 0.687245] [D_R loss: 0.661841] [G loss: -0.265561]\n",
      "[Lambda_f: 0.500000] [Epoch 800/4000] [D_F loss: 0.657802] [D_R loss: 0.649252] [G loss: -0.245780]\n",
      "[Lambda_f: 0.500000] [Epoch 1000/4000] [D_F loss: 0.651085] [D_R loss: 0.642768] [G loss: -0.248019]\n",
      "[Lambda_f: 0.500000] [Epoch 1200/4000] [D_F loss: 0.653262] [D_R loss: 0.638528] [G loss: -0.248325]\n",
      "[Lambda_f: 0.500000] [Epoch 1400/4000] [D_F loss: 0.655903] [D_R loss: 0.634185] [G loss: -0.248550]\n",
      "[Lambda_f: 0.500000] [Epoch 1600/4000] [D_F loss: 0.657621] [D_R loss: 0.629968] [G loss: -0.248579]\n",
      "[Lambda_f: 0.500000] [Epoch 1800/4000] [D_F loss: 0.658911] [D_R loss: 0.625692] [G loss: -0.248469]\n",
      "[Lambda_f: 0.500000] [Epoch 2000/4000] [D_F loss: 0.659922] [D_R loss: 0.621362] [G loss: -0.248301]\n",
      "[Lambda_f: 0.500000] [Epoch 2200/4000] [D_F loss: 0.660699] [D_R loss: 0.616973] [G loss: -0.248014]\n",
      "[Lambda_f: 0.500000] [Epoch 2400/4000] [D_F loss: 0.661295] [D_R loss: 0.612527] [G loss: -0.247685]\n",
      "[Lambda_f: 0.500000] [Epoch 2600/4000] [D_F loss: 0.661772] [D_R loss: 0.608033] [G loss: -0.247368]\n",
      "[Lambda_f: 0.500000] [Epoch 2800/4000] [D_F loss: 0.662150] [D_R loss: 0.603502] [G loss: -0.247001]\n",
      "[Lambda_f: 0.500000] [Epoch 3000/4000] [D_F loss: 0.662449] [D_R loss: 0.598971] [G loss: -0.246630]\n",
      "[Lambda_f: 0.500000] [Epoch 3200/4000] [D_F loss: 0.662693] [D_R loss: 0.594418] [G loss: -0.246283]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Lambda_f: 0.500000] [Epoch 3400/4000] [D_F loss: 0.662887] [D_R loss: 0.589849] [G loss: -0.245900]\n",
      "[Lambda_f: 0.500000] [Epoch 3600/4000] [D_F loss: 0.663042] [D_R loss: 0.585281] [G loss: -0.245516]\n",
      "[Lambda_f: 0.500000] [Epoch 3800/4000] [D_F loss: 0.663166] [D_R loss: 0.580718] [G loss: -0.245163]\n",
      "Test accuracy: 0.8579999804496765\n",
      "P(y_hat=1 | z=0) = 0.396, P(y_hat=1 | z=1) = 0.651\n",
      "P(y_hat=1 | y=1, z=0) = 0.854, P(y_hat=1 | y=1, z=1) = 0.847\n",
      "Disparate Impact ratio = 0.608\n",
      "[Lambda_f: 0.600000] [Epoch 0/4000] [D_F loss: 0.686341] [D_R loss: 0.716233] [G loss: 0.465272]\n",
      "[Lambda_f: 0.600000] [Epoch 200/4000] [D_F loss: 0.584361] [D_R loss: 0.680139] [G loss: 0.312735]\n",
      "[Lambda_f: 0.600000] [Epoch 400/4000] [D_F loss: 0.577822] [D_R loss: 0.673043] [G loss: 0.312282]\n",
      "[Lambda_f: 0.600000] [Epoch 600/4000] [D_F loss: 0.700302] [D_R loss: 0.663081] [G loss: -0.373158]\n",
      "[Lambda_f: 0.600000] [Epoch 800/4000] [D_F loss: 0.668566] [D_R loss: 0.651781] [G loss: -0.345874]\n",
      "[Lambda_f: 0.600000] [Epoch 1000/4000] [D_F loss: 0.664284] [D_R loss: 0.645515] [G loss: -0.348699]\n",
      "[Lambda_f: 0.600000] [Epoch 1200/4000] [D_F loss: 0.665559] [D_R loss: 0.643799] [G loss: -0.349784]\n",
      "[Lambda_f: 0.600000] [Epoch 1400/4000] [D_F loss: 0.668383] [D_R loss: 0.641819] [G loss: -0.350503]\n",
      "[Lambda_f: 0.600000] [Epoch 1600/4000] [D_F loss: 0.670203] [D_R loss: 0.639402] [G loss: -0.350891]\n",
      "[Lambda_f: 0.600000] [Epoch 1800/4000] [D_F loss: 0.671251] [D_R loss: 0.636526] [G loss: -0.351044]\n",
      "[Lambda_f: 0.600000] [Epoch 2000/4000] [D_F loss: 0.671944] [D_R loss: 0.633285] [G loss: -0.351069]\n",
      "[Lambda_f: 0.600000] [Epoch 2200/4000] [D_F loss: 0.672429] [D_R loss: 0.629797] [G loss: -0.350957]\n",
      "[Lambda_f: 0.600000] [Epoch 2400/4000] [D_F loss: 0.672784] [D_R loss: 0.626162] [G loss: -0.350779]\n",
      "[Lambda_f: 0.600000] [Epoch 2600/4000] [D_F loss: 0.673067] [D_R loss: 0.622405] [G loss: -0.350585]\n",
      "[Lambda_f: 0.600000] [Epoch 2800/4000] [D_F loss: 0.673295] [D_R loss: 0.618589] [G loss: -0.350336]\n",
      "[Lambda_f: 0.600000] [Epoch 3000/4000] [D_F loss: 0.673482] [D_R loss: 0.614749] [G loss: -0.350063]\n",
      "[Lambda_f: 0.600000] [Epoch 3200/4000] [D_F loss: 0.673640] [D_R loss: 0.610853] [G loss: -0.349796]\n",
      "[Lambda_f: 0.600000] [Epoch 3400/4000] [D_F loss: 0.673772] [D_R loss: 0.606930] [G loss: -0.349495]\n",
      "[Lambda_f: 0.600000] [Epoch 3600/4000] [D_F loss: 0.673881] [D_R loss: 0.603008] [G loss: -0.349192]\n",
      "[Lambda_f: 0.600000] [Epoch 3800/4000] [D_F loss: 0.673975] [D_R loss: 0.599080] [G loss: -0.348899]\n",
      "Test accuracy: 0.847000002861023\n",
      "P(y_hat=1 | z=0) = 0.421, P(y_hat=1 | z=1) = 0.635\n",
      "P(y_hat=1 | y=1, z=0) = 0.879, P(y_hat=1 | y=1, z=1) = 0.825\n",
      "Disparate Impact ratio = 0.663\n",
      "[Lambda_f: 0.700000] [Epoch 0/4000] [D_F loss: 0.686341] [D_R loss: 0.716233] [G loss: 0.465272]\n",
      "[Lambda_f: 0.700000] [Epoch 200/4000] [D_F loss: 0.584361] [D_R loss: 0.680139] [G loss: 0.312735]\n",
      "[Lambda_f: 0.700000] [Epoch 400/4000] [D_F loss: 0.577822] [D_R loss: 0.673043] [G loss: 0.312282]\n",
      "[Lambda_f: 0.700000] [Epoch 600/4000] [D_F loss: 0.711665] [D_R loss: 0.664331] [G loss: -0.485118]\n",
      "[Lambda_f: 0.700000] [Epoch 800/4000] [D_F loss: 0.676495] [D_R loss: 0.654294] [G loss: -0.450321]\n",
      "[Lambda_f: 0.700000] [Epoch 1000/4000] [D_F loss: 0.675458] [D_R loss: 0.647703] [G loss: -0.452497]\n",
      "[Lambda_f: 0.700000] [Epoch 1200/4000] [D_F loss: 0.675968] [D_R loss: 0.646940] [G loss: -0.454257]\n",
      "[Lambda_f: 0.700000] [Epoch 1400/4000] [D_F loss: 0.677761] [D_R loss: 0.647567] [G loss: -0.455588]\n",
      "[Lambda_f: 0.700000] [Epoch 1600/4000] [D_F loss: 0.679239] [D_R loss: 0.647669] [G loss: -0.456415]\n",
      "[Lambda_f: 0.700000] [Epoch 1800/4000] [D_F loss: 0.680005] [D_R loss: 0.646760] [G loss: -0.456814]\n",
      "[Lambda_f: 0.700000] [Epoch 2000/4000] [D_F loss: 0.680382] [D_R loss: 0.645002] [G loss: -0.456969]\n",
      "[Lambda_f: 0.700000] [Epoch 2200/4000] [D_F loss: 0.680609] [D_R loss: 0.642713] [G loss: -0.456970]\n",
      "[Lambda_f: 0.700000] [Epoch 2400/4000] [D_F loss: 0.680764] [D_R loss: 0.640099] [G loss: -0.456896]\n",
      "[Lambda_f: 0.700000] [Epoch 2600/4000] [D_F loss: 0.680880] [D_R loss: 0.637276] [G loss: -0.456801]\n",
      "[Lambda_f: 0.700000] [Epoch 2800/4000] [D_F loss: 0.680973] [D_R loss: 0.634329] [G loss: -0.456658]\n",
      "[Lambda_f: 0.700000] [Epoch 3000/4000] [D_F loss: 0.681051] [D_R loss: 0.631308] [G loss: -0.456493]\n",
      "[Lambda_f: 0.700000] [Epoch 3200/4000] [D_F loss: 0.681119] [D_R loss: 0.628210] [G loss: -0.456320]\n",
      "[Lambda_f: 0.700000] [Epoch 3400/4000] [D_F loss: 0.681177] [D_R loss: 0.625066] [G loss: -0.456120]\n",
      "[Lambda_f: 0.700000] [Epoch 3600/4000] [D_F loss: 0.681227] [D_R loss: 0.621894] [G loss: -0.455908]\n",
      "[Lambda_f: 0.700000] [Epoch 3800/4000] [D_F loss: 0.681272] [D_R loss: 0.618669] [G loss: -0.455694]\n",
      "Test accuracy: 0.8370000123977661\n",
      "P(y_hat=1 | z=0) = 0.436, P(y_hat=1 | z=1) = 0.614\n",
      "P(y_hat=1 | y=1, z=0) = 0.894, P(y_hat=1 | y=1, z=1) = 0.800\n",
      "Disparate Impact ratio = 0.711\n",
      "[Lambda_f: 0.800000] [Epoch 0/4000] [D_F loss: 0.686341] [D_R loss: 0.716233] [G loss: 0.465272]\n",
      "[Lambda_f: 0.800000] [Epoch 200/4000] [D_F loss: 0.584361] [D_R loss: 0.680139] [G loss: 0.312735]\n",
      "[Lambda_f: 0.800000] [Epoch 400/4000] [D_F loss: 0.577822] [D_R loss: 0.673043] [G loss: 0.312282]\n",
      "[Lambda_f: 0.800000] [Epoch 600/4000] [D_F loss: 0.721308] [D_R loss: 0.665539] [G loss: -0.601194]\n",
      "[Lambda_f: 0.800000] [Epoch 800/4000] [D_F loss: 0.681845] [D_R loss: 0.656688] [G loss: -0.559816]\n",
      "[Lambda_f: 0.800000] [Epoch 1000/4000] [D_F loss: 0.682954] [D_R loss: 0.649545] [G loss: -0.560539]\n",
      "[Lambda_f: 0.800000] [Epoch 1200/4000] [D_F loss: 0.683494] [D_R loss: 0.647005] [G loss: -0.561527]\n",
      "[Lambda_f: 0.800000] [Epoch 1400/4000] [D_F loss: 0.683910] [D_R loss: 0.646956] [G loss: -0.562598]\n",
      "[Lambda_f: 0.800000] [Epoch 1600/4000] [D_F loss: 0.684278] [D_R loss: 0.647870] [G loss: -0.563626]\n",
      "[Lambda_f: 0.800000] [Epoch 1800/4000] [D_F loss: 0.684573] [D_R loss: 0.649012] [G loss: -0.564478]\n",
      "[Lambda_f: 0.800000] [Epoch 2000/4000] [D_F loss: 0.684780] [D_R loss: 0.649886] [G loss: -0.565074]\n",
      "[Lambda_f: 0.800000] [Epoch 2200/4000] [D_F loss: 0.684909] [D_R loss: 0.650201] [G loss: -0.565458]\n",
      "[Lambda_f: 0.800000] [Epoch 2400/4000] [D_F loss: 0.684996] [D_R loss: 0.649865] [G loss: -0.565656]\n",
      "[Lambda_f: 0.800000] [Epoch 2600/4000] [D_F loss: 0.685054] [D_R loss: 0.648976] [G loss: -0.565724]\n",
      "[Lambda_f: 0.800000] [Epoch 2800/4000] [D_F loss: 0.685094] [D_R loss: 0.647629] [G loss: -0.565696]\n",
      "[Lambda_f: 0.800000] [Epoch 3000/4000] [D_F loss: 0.685121] [D_R loss: 0.645964] [G loss: -0.565614]\n",
      "[Lambda_f: 0.800000] [Epoch 3200/4000] [D_F loss: 0.685141] [D_R loss: 0.644061] [G loss: -0.565508]\n",
      "[Lambda_f: 0.800000] [Epoch 3400/4000] [D_F loss: 0.685158] [D_R loss: 0.642004] [G loss: -0.565374]\n",
      "[Lambda_f: 0.800000] [Epoch 3600/4000] [D_F loss: 0.685171] [D_R loss: 0.639836] [G loss: -0.565227]\n",
      "[Lambda_f: 0.800000] [Epoch 3800/4000] [D_F loss: 0.685182] [D_R loss: 0.637571] [G loss: -0.565078]\n",
      "Test accuracy: 0.8199999928474426\n",
      "P(y_hat=1 | z=0) = 0.449, P(y_hat=1 | z=1) = 0.585\n",
      "P(y_hat=1 | y=1, z=0) = 0.894, P(y_hat=1 | y=1, z=1) = 0.766\n",
      "Disparate Impact ratio = 0.766\n",
      "[Lambda_f: 0.850000] [Epoch 0/4000] [D_F loss: 0.686341] [D_R loss: 0.716233] [G loss: 0.465272]\n",
      "[Lambda_f: 0.850000] [Epoch 200/4000] [D_F loss: 0.584361] [D_R loss: 0.680139] [G loss: 0.312735]\n",
      "[Lambda_f: 0.850000] [Epoch 400/4000] [D_F loss: 0.577822] [D_R loss: 0.673043] [G loss: 0.312282]\n",
      "[Lambda_f: 0.850000] [Epoch 600/4000] [D_F loss: 0.725627] [D_R loss: 0.666131] [G loss: -0.660628]\n",
      "[Lambda_f: 0.850000] [Epoch 800/4000] [D_F loss: 0.683724] [D_R loss: 0.657923] [G loss: -0.618142]\n",
      "[Lambda_f: 0.850000] [Epoch 1000/4000] [D_F loss: 0.685031] [D_R loss: 0.649860] [G loss: -0.617109]\n",
      "[Lambda_f: 0.850000] [Epoch 1200/4000] [D_F loss: 0.685526] [D_R loss: 0.645590] [G loss: -0.616939]\n",
      "[Lambda_f: 0.850000] [Epoch 1400/4000] [D_F loss: 0.685710] [D_R loss: 0.643427] [G loss: -0.617099]\n",
      "[Lambda_f: 0.850000] [Epoch 1600/4000] [D_F loss: 0.685770] [D_R loss: 0.642253] [G loss: -0.617415]\n",
      "[Lambda_f: 0.850000] [Epoch 1800/4000] [D_F loss: 0.685802] [D_R loss: 0.641764] [G loss: -0.617861]\n",
      "[Lambda_f: 0.850000] [Epoch 2000/4000] [D_F loss: 0.685839] [D_R loss: 0.641981] [G loss: -0.618369]\n",
      "[Lambda_f: 0.850000] [Epoch 2200/4000] [D_F loss: 0.685879] [D_R loss: 0.642562] [G loss: -0.618910]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Lambda_f: 0.850000] [Epoch 2400/4000] [D_F loss: 0.685914] [D_R loss: 0.643272] [G loss: -0.619421]\n",
      "[Lambda_f: 0.850000] [Epoch 2600/4000] [D_F loss: 0.685946] [D_R loss: 0.644025] [G loss: -0.619845]\n",
      "[Lambda_f: 0.850000] [Epoch 2800/4000] [D_F loss: 0.685973] [D_R loss: 0.644627] [G loss: -0.620204]\n",
      "[Lambda_f: 0.850000] [Epoch 3000/4000] [D_F loss: 0.685993] [D_R loss: 0.644949] [G loss: -0.620467]\n",
      "[Lambda_f: 0.850000] [Epoch 3200/4000] [D_F loss: 0.686009] [D_R loss: 0.645002] [G loss: -0.620636]\n",
      "[Lambda_f: 0.850000] [Epoch 3400/4000] [D_F loss: 0.686022] [D_R loss: 0.644735] [G loss: -0.620731]\n",
      "[Lambda_f: 0.850000] [Epoch 3600/4000] [D_F loss: 0.686031] [D_R loss: 0.644141] [G loss: -0.620760]\n",
      "[Lambda_f: 0.850000] [Epoch 3800/4000] [D_F loss: 0.686039] [D_R loss: 0.643264] [G loss: -0.620733]\n",
      "Test accuracy: 0.8069999814033508\n",
      "P(y_hat=1 | z=0) = 0.473, P(y_hat=1 | z=1) = 0.578\n",
      "P(y_hat=1 | y=1, z=0) = 0.910, P(y_hat=1 | y=1, z=1) = 0.753\n",
      "Disparate Impact ratio = 0.818\n"
     ]
    }
   ],
   "source": [
    "train_result = []\n",
    "train_tensors = Namespace(XS_train = XS_train, y_train = y_train, s1_train = s1_train)\n",
    "val_tensors = Namespace(XS_val = XS_val, y_val = y_val, s1_val = s1_val) \n",
    "test_tensors = Namespace(XS_test = XS_test, y_test = y_test, s1_test = s1_test)\n",
    "\n",
    "train_opt = Namespace(val=len(y_val), n_epochs=4000, k=3, lr_g=0.005, lr_f=0.01, lr_r=0.001)      \n",
    "seed = 1 \n",
    "\n",
    "lambda_f_set = [0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.85] # Lambda value for the fairness discriminator of FR-Train.\n",
    "lambda_r = 0.1 # Lambda value for the robustness discriminator of FR-Train.\n",
    "\n",
    "for lambda_f in lambda_f_set:\n",
    "    train_result.append(train_model(train_tensors, val_tensors, test_tensors, train_opt, lambda_f = lambda_f, lambda_r = lambda_r, seed=seed))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "------------------ Training Results of FR-Train on clean data ------------------\n",
      "[Lambda_f: 0.20] [Lambda_r: 0.10] Accuracy : 0.867, Disparate Impact : 0.492 \n",
      "[Lambda_f: 0.30] [Lambda_r: 0.10] Accuracy : 0.866, Disparate Impact : 0.528 \n",
      "[Lambda_f: 0.40] [Lambda_r: 0.10] Accuracy : 0.857, Disparate Impact : 0.554 \n",
      "[Lambda_f: 0.50] [Lambda_r: 0.10] Accuracy : 0.858, Disparate Impact : 0.608 \n",
      "[Lambda_f: 0.60] [Lambda_r: 0.10] Accuracy : 0.847, Disparate Impact : 0.663 \n",
      "[Lambda_f: 0.70] [Lambda_r: 0.10] Accuracy : 0.837, Disparate Impact : 0.711 \n",
      "[Lambda_f: 0.80] [Lambda_r: 0.10] Accuracy : 0.820, Disparate Impact : 0.766 \n",
      "[Lambda_f: 0.85] [Lambda_r: 0.10] Accuracy : 0.807, Disparate Impact : 0.818 \n",
      "--------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "print(\"--------------------------------------------------------------------------------\")\n",
    "print(\"------------------ Training Results of FR-Train on clean data ------------------\" )\n",
    "for i in range(len(train_result)):\n",
    "    print(\n",
    "        \"[Lambda_f: %.2f] [Lambda_r: %.2f] Accuracy : %.3f, Disparate Impact : %.3f \"\n",
    "        % (train_result[i][0][0], train_result[i][0][1], train_result[i][0][2], train_result[i][0][3])\n",
    "    )       \n",
    "print(\"--------------------------------------------------------------------------------\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyTorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}