google-research
1614 строк · 64.8 Кб
1{2"cells": [3{4"cell_type": "markdown",5"metadata": {6"id": "ozc9I4X1C4nX"7},8"source": [9"Copyright 2022 Google LLC.\n",10"\n",11"Licensed under the Apache License, Version 2.0 (the \"License\");"12]13},14{15"cell_type": "code",16"execution_count": null,17"metadata": {18"cellView": "form",19"id": "aIZExrKYCqRi"20},21"outputs": [],22"source": [23"#@title License\n",24"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",25"# you may not use this file except in compliance with the License.\n",26"# You may obtain a copy of the License at\n",27"#\n",28"# https://www.apache.org/licenses/LICENSE-2.0\n",29"#\n",30"# Unless required by applicable law or agreed to in writing, software\n",31"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",32"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",33"# See the License for the specific language governing permissions and\n",34"# limitations under the License."35]36},37{38"cell_type": "markdown",39"metadata": {40"id": "4mdq2mi2Y7_N"41},42"source": [43"# Shortcut testing with synthetic images\n",44"\n",45"We illustrate the method proposed in Brown et al. 2022 (https://arxiv.org/abs/2207.10384) to detect model shortcutting when a model relies on both signals related to the label Y and to an auxiliary/sensitive attribute A.\n",46"\n",47"To demonstrate ShorT, we need a dataset such that the main (Y) and auxiliary (A) tasks are not too easy and \"help\" each other. We refer to the MNIST dataset, with the aim to discriminate between numbers smaller than 5 and larger than 5. We add a confounder to each image, i.e. a small colored squared. The color of the square (red or green) can then be correlated to the label in the image to create a spurious signal. We add significnat amounts of noise to both signals to avoid a \"binary\" behavior of the model (i.e. when it fully relies on one signal or the other, but never on both).\n",48"\n",49"We demonstrate how ShorT identifies the confounding, when present."50]51},52{53"cell_type": "markdown",54"metadata": {55"id": "-UcYpkXxZApt"56},57"source": [58"# Imports"59]60},61{62"cell_type": "code",63"execution_count": null,64"metadata": {65"id": "z9NNjxN4XXD3"66},67"outputs": [],68"source": [69"import os\n",70"os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' \n",71"\n",72"import importlib\n",73"import random\n",74"\n",75"from typing import *\n",76"\n",77"import attr\n",78"import matplotlib as mpl\n",79"import matplotlib.pyplot as plt\n",80"import math\n",81"import numpy as np\n",82"import seaborn as sns\n",83"import scipy\n",84"\n",85"from matplotlib.cm import get_cmap\n",86"from matplotlib import gridspec\n",87"from matplotlib import pyplot as plt\n",88"from matplotlib import cm\n",89"\n",90"import sklearn\n",91"from sklearn.utils import shuffle\n",92"\n",93"import tensorflow as tf\n",94"from tensorflow import keras\n",95"from keras import backend as K\n",96"import numpy as np\n"97]98},99{100"cell_type": "code",101"execution_count": null,102"metadata": {103"id": "LGs2__K3Y4ew"104},105"outputs": [],106"source": [107"!pip install ml_collections"108]109},110{111"cell_type": "code",112"execution_count": null,113"metadata": {114"id": "SUFJth7Qby8R"115},116"outputs": [],117"source": [118"import ml_collections as mlc"119]120},121{122"cell_type": "markdown",123"metadata": {124"id": "Hsk4c0jwcLa8"125},126"source": [127"# Synthetic dataset from MNIST\n",128"Images are generated by adding a red or green square on an MNIST image ([Deng et al., 2012](https://ieeexplore.ieee.org/document/6296535)). We select our labels to be an image with a number smaller than 5 (`label=0`) or larger than 5 (`label=1`). Two parameters specify the proportion (in %) of each label being red. A larger number (\u003e 50%) means that the class label will be correlated with the red color, while a smaller number (\u003c 50%) means that the class label will be correlated with the green color.\n",129"\n",130"Refer to the config dictionary (`cfg_data`) to specify all the data generation parameters. Please note that it is important to check that the tasks (predicting the label and predicting the attribute) are not trivial or too simple."131]132},133{134"cell_type": "code",135"execution_count": null,136"metadata": {137"cellView": "form",138"id": "eZZr6SzLdHdl"139},140"outputs": [],141"source": [142"#@title Data utils\n",143"\n",144"def add_square_to_image(image, color, square_size, noise_thresh = 0.1):\n",145" image_size = image.shape\n",146" start_x = np.random.choice(image_size[0]-square_size)\n",147" end_x = min(image_size[0], start_x+square_size)\n",148" start_y = np.random.choice(image_size[1]-square_size)\n",149" end_y = min(image_size[1], start_y+square_size)\n",150" square = np.random.uniform(\n",151" size=(square_size, square_size, image_size[-1])) \u003e= noise_thresh\n",152" image[start_x:end_x, start_y:end_y, :] = square * np.array(color) #color\n",153" return image, square, start_x, start_y\n",154"\n",155"def create_images_dataset(corr_a_y0, corr_a_y1, n_images_per_class, n_labels,\n",156" select_labels, n_colors, colors, color_map,\n",157" noise_thresh, conf_noise=0.1, mnist_split=\"train\"):\n",158"\n",159" cls0_corr = float(corr_a_y0) / 100\n",160" cls1_corr = float(corr_a_y1) / 100\n",161"\n",162" conprob_colors_orie = np.array([[cls0_corr, 1 - cls0_corr],\n",163" [cls1_corr, 1 - cls1_corr]])\n",164" \n",165" \n",166" # Load MNIST\n",167" (train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",168" if mnist_split == \"train\":\n",169" images = train_images\n",170" labels = train_labels\n",171" else:\n",172" images = test_images\n",173" labels = test_labels\n",174" \n",175" n_images = n_images_per_class * n_labels\n",176" n_pixels = images[0].shape[0]\n",177" images_colored = np.zeros((n_images, n_pixels, n_pixels, 3))\n",178" flip_images = np.zeros((n_images, n_pixels, n_pixels, 3))\n",179" image_labels = np.zeros((n_images, 1), dtype=int)\n",180" color_labels = np.zeros((n_images, 1), dtype=int)\n",181"\n",182" for i, label in enumerate(select_labels):\n",183" selected = labels == label[0]\n",184" j = 1\n",185" while j\u003clen(label):\n",186" selected = np.logical_or(selected, labels == label[j])\n",187" j += 1\n",188" selected_label = labels[selected]\n",189" selected_images = images[selected, :,:]\n",190" for class_image_index in range(n_images_per_class):\n",191" image_index = i * n_images_per_class + class_image_index\n",192" color = np.random.choice(n_colors, p=conprob_colors_orie[i])\n",193" flip_color = 1 if color == 0 else 0\n",194" # randomly select image from MNIST\n",195" idx = np.random.choice(selected_label.shape[0], replace=True)\n",196" image = np.reshape(selected_images[idx,:,:], (n_pixels,n_pixels, 1))\n",197" channels = np.concatenate((image, image, image), axis=2)\n",198" # Add colored square to the image\n",199" images_colored[image_index], square, x, y = add_square_to_image(channels/255,\n",200" color_map[colors[color]],\n",201" int(n_pixels/7),\n",202" noise_thresh=conf_noise)\n",203" flip_images[image_index] = channels/255\n",204" sq_szx, sq_szy, _ = square.shape\n",205" flip_images[image_index,x:x+sq_szx,y:y+sq_szy, :] = square * color_map[colors[flip_color]]\n",206"\n",207" noise_img = np.random.uniform(\n",208" size=[n_pixels, n_pixels, 1]) \u003c noise_thresh # white noise\n",209" noise_img = np.concatenate((noise_img, noise_img, noise_img), axis=2)\n",210" noise_img = noise_img.astype('float64')\n",211" noise_img[x:x+sq_szx,y:y+sq_szy, :] = square * (0.,0.,0.) # remove noise from added square\n",212"\n",213" images_colored[image_index] = np.clip(\n",214" images_colored[image_index] + noise_img, 0.0, 1.0)\n",215" flip_images[image_index] = np.clip(\n",216" flip_images[image_index] + noise_img, 0.0, 1.0)\n",217"\n",218" image_labels[image_index,:] = i\n",219" color_labels[image_index,:] = color\n",220"\n",221" return images_colored, image_labels, color_labels, flip_images\n",222"\n",223"\n",224"def gen_dataset(cfg_data):\n",225" \"\"\"Generates bars dataset.\n",226"\n",227" Each image contains a single bar which is either horizontal (class 0) or\n",228" vertical (class 1)\n",229" The bar is either red (concept 0) or green (concept 1)\n",230" \"\"\"\n",231"\n",232" tr_im, tr_lab, tr_col_lab, tr_fl_im = create_images_dataset(**cfg_data)\n",233" \n",234" train_images, train_labels, train_color_labels, train_flip_images = \\\n",235" sklearn.utils.shuffle(tr_im, tr_lab, tr_col_lab, tr_fl_im)\n",236"\n",237" dataset = dict(\n",238" train=dict(\n",239" image=train_images,\n",240" label=train_labels,\n",241" color_label=train_color_labels,\n",242" flip_image=train_flip_images),\n",243" )\n",244" return dataset"245]246},247{248"cell_type": "code",249"execution_count": null,250"metadata": {251"id": "A6qfNB64dU5s"252},253"outputs": [],254"source": [255"#@title Synthetic data parameters\n",256"\n",257"cfg_data = mlc.ConfigDict()\n",258"\n",259"# These must show a large difference to encourage shortcutting\n",260"cfg_data.corr_a_y0 = 20 #@param Proportion of images with label 0 being red\n",261"cfg_data.corr_a_y1 = 95 #@param Proportion of images with label 1 being red\n",262"\n",263"# Note: asymmetrical, high correlations (i.e. corr_y0 + corr_y1 != 100) lead to \n",264"# unfairness in terms equalized odds\n",265"\n",266"cfg_data.n_images_per_class = 5000 #@param\n",267"\n",268"cfg_data.noise_thresh = 0.5 # noise added to the number\n",269"cfg_data.conf_noise = 0.6 # noise added to the confounder\n",270"\n",271"cfg_data.colors = ['red', 'green', 'blue']\n",272"cfg_data.color_map = mlc.ConfigDict()\n",273"cfg_data.color_map.red = (1., 0., 0.)\n",274"cfg_data.color_map.green = (0., 1., 0.)\n",275"cfg_data.color_map.blue = (0., 0., 1.)\n",276"\n",277"cfg_data.n_colors = 2\n",278"cfg_data.n_labels = 2\n",279"\n",280"cfg_data.select_labels = [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]\n",281"label_names = ['\u003c5', '\u003e=5']\n",282"cfg_data.mnist_split = \"train\"\n",283"\n",284"# set the seed\n",285"seed = 123\n",286"tf.random.set_seed(seed)\n",287"np.random.seed(seed)"288]289},290{291"cell_type": "markdown",292"metadata": {293"id": "wwElA1s9ebiC"294},295"source": [296"Based on these parameters, we generate train, validation and test data. In addition, we can generate \"counterfactuals\" (referred to as 'flipped_images') that flip the sampled color label before generating the same image."297]298},299{300"cell_type": "code",301"execution_count": null,302"metadata": {303"id": "yayPrkNueTCD"304},305"outputs": [],306"source": [307"#@title Train and validation data\n",308"dataset = gen_dataset(cfg_data)"309]310},311{312"cell_type": "code",313"execution_count": null,314"metadata": {315"id": "iz3Y5yBle7TR"316},317"outputs": [],318"source": [319"#@markdown Explore the dataset\n",320"dataset[\"train\"].keys()"321]322},323{324"cell_type": "code",325"execution_count": null,326"metadata": {327"id": "Isq1LxhTfA5x"328},329"outputs": [],330"source": [331"#@markdown Compute proportions of the data to check specifications\n",332"# The output numbers should approximate the desired correlations between Y and A\n",333"\n",334"n_r1 = np.sum(np.logical_and(dataset[\"train\"][\"label\"] == 0, \n",335" dataset[\"train\"][\"color_label\"] == 0)) / np.sum(dataset[\"train\"][\"label\"] == 0)\n",336"print(\"Proportion of 0 labels that are red: {}\".format(n_r1))\n",337"\n",338"n_r1 = np.sum(np.logical_and(dataset[\"train\"][\"label\"] == 1,\n",339" dataset[\"train\"][\"color_label\"] == 0)) / np.sum(dataset[\"train\"][\"label\"] == 1)\n",340"print(\"Proportion of 1 labels that are red: {}\".format(n_r1))"341]342},343{344"cell_type": "code",345"execution_count": null,346"metadata": {347"id": "BeLWotfXfFuZ"348},349"outputs": [],350"source": [351"#@markdown Visualize the data (5 random samples)\n",352"\n",353"split = \"train\"\n",354"NUM_TEST_SAMPLES = len(dataset[split][\"label\"])\n",355"print(NUM_TEST_SAMPLES, dataset[split][\"label\"])\n",356"\n",357"for idx in random.choices(range(NUM_TEST_SAMPLES), k=5):\n",358" plt.figure()\n",359" print(dataset[split][\"label\"][idx], dataset[split][\"color_label\"][idx])\n",360" plt.imshow(dataset[split][\"image\"][idx, ...])\n",361" plt.show()\n",362"n_pixels, n_pixels, n_channels = dataset[split][\"image\"][idx, ...].shape"363]364},365{366"cell_type": "code",367"execution_count": null,368"metadata": {369"id": "A7XPU8_ufpbw"370},371"outputs": [],372"source": [373"#@title Test data\n",374"\n",375"cfg_data.n_images_per_class = 2000 #@param\n",376"cfg_data.mnist_split = \"test\" # We select images from the \"test\" MNIST split\n",377"\n",378"test_dataset = gen_dataset(cfg_data)"379]380},381{382"cell_type": "code",383"execution_count": null,384"metadata": {385"id": "qQ_ZhHVHf9-W"386},387"outputs": [],388"source": [389"#@title Transform data into TF Datasets\n",390"def tf_data(dataset, batch_size = 64, valid_prop = 0.2, is_train = True):\n",391"\n",392" # Original images\n",393" split = \"train\"\n",394" training_data = tf.data.Dataset.from_tensor_slices((dataset[split][\"image\"],\n",395" dataset[split][\"label\"],\n",396" dataset[split][\"color_label\"]))\n",397" n_images = dataset[split][\"label\"].shape[0]\n",398" n_valid = math.floor(n_images * valid_prop)\n",399" valid_data = training_data.take(n_valid).cache().repeat().batch(batch_size)\n",400" ds = training_data.skip(n_valid).take(n_images - n_valid).cache()\n",401" # Flipped color on the same images\n",402" data_flipped = tf.data.Dataset.from_tensor_slices((dataset[split][\"flip_image\"],\n",403" dataset[split][\"label\"],\n",404" np.logical_not(dataset[split][\"color_label\"]).astype(float)))\n",405" \n",406" if is_train:\n",407" # Cannot be used as pairs due to shuffling\n",408" data = ds.repeat().shuffle(10000).batch(batch_size)\n",409" data_flipped = data_flipped.repeat().shuffle(10000).batch(batch_size)\n",410" else:\n",411" # Can be used for paired comparisons\n",412" data = ds.repeat().batch(batch_size)\n",413" data_flipped = data_flipped.repeat().batch(batch_size)\n",414" \n",415" return data, data_flipped, valid_data"416]417},418{419"cell_type": "code",420"execution_count": null,421"metadata": {422"id": "Ma-tMJPogEwe"423},424"outputs": [],425"source": [426"pos = mlc.ConfigDict()\n",427"pos.x, pos.y, pos.a = 0, 1, 2 # Associates a position in the tf dataset to a variable\n",428"\n",429"# Train and validation data\n",430"train_data, _, valid_data = tf_data(dataset, 64)"431]432},433{434"cell_type": "code",435"execution_count": null,436"metadata": {437"id": "kZChfj7igeMd"438},439"outputs": [],440"source": [441"# Test data - avoinding to shuffle so we can compare the images and their counterfactuals\n",442"test_data, test_data_flipped, _ = tf_data(test_dataset, 64, valid_prop=0.0, is_train=False)"443]444},445{446"cell_type": "markdown",447"metadata": {448"id": "xujPFv_tgv_R"449},450"source": [451"# Models\n",452"\n",453"## Baseline model - SingleHead class\n",454"We build simple MLP models that predict one binary label. This model can be used to assess the performance of the baseline model f(X) -\u003e Y.\n",455"\n",456"## Attribute encoding - SingleHead with frozen feature extractor\n",457"In addition, this model can take as an extra input a frozen feature extractor. It then adds a linear layer to this feature extractor to perform the binary classification. Only the weights of this layer are tuned during training. This architecture is used to assess the level of attribute encoding of another model.\n",458"\n",459"## Multi-task - MultiHead with Gradient Reversal\n",460"The architecture of the multi-task model first comprises a series of layers that represent a \"feature extractor\". On top of the feature extractor, we add two heads: \n",461"- one head for the label, which is predicted from a single linear layer,\n",462"- one head for an auxiliary task (here the prediction of the color of the square). This head includes at least a non-linear layer to allow a gradient scaling operation.\n",463"\n",464"The weight of the auxiliary loss in the total loss is a hyper-parameter in the method (weight of label loss is fixed to 1.0). A gradient reversal head is added to the auxiliary task. Gradient scaling can then be performed (positive or negative, extra hyper-parameter). Setting both hyper-parameters to 0 corresponds to the baseline single head model.\n",465"\n",466"\n",467"A config (`cfg`) dictionary specifies the parameters of the MLP architecture."468]469},470{471"cell_type": "code",472"execution_count": null,473"metadata": {474"cellView": "form",475"id": "A41cqIbNi_mJ"476},477"outputs": [],478"source": [479"#@title Model utils\n",480"\n",481"class GradientReversal(tf.keras.layers.Layer):\n",482"\n",483" @tf.custom_gradient\n",484" def grad_reverse(self, x):\n",485" y = tf.identity(x)\n",486" def custom_grad(dy):\n",487" return self.hp_lambda * dy\n",488" return y, custom_grad\n",489"\n",490" def __init__(self, hp_lambda, **kwargs):\n",491" super(GradientReversal, self).__init__(**kwargs)\n",492" self.hp_lambda = K.variable(hp_lambda, dtype='float', name='hp_lambda')\n",493"\n",494" def call(self, x, mask=None):\n",495" return self.grad_reverse(x)\n",496"\n",497" def set_hp_lambda(self,hp_lambda):\n",498" K.set_value(self.hp_lambda, hp_lambda)\n",499" \n",500" def increment_hp_lambda_by(self,increment):\n",501" new_value = float(K.get_value(self.hp_lambda)) + increment\n",502" K.set_value(self.hp_lambda, new_value)\n",503"\n",504" def get_hp_lambda(self):\n",505" return float(K.get_value(self.hp_lambda))\n",506"\n",507"\n",508"class BaselineArch():\n",509" \"\"\"Superclass for multihead training.\"\"\"\n",510"\n",511" def __init__(self, main=\"y\", aux=None, dtype=tf.float32, pos=None):\n",512" \"\"\"Initializer.\n",513"\n",514" Args:\n",515" main: name of variable for the main task\n",516" aux: nema of the variable for the auxiliary task\n",517" dtype: desired dtype (e.g. tf.float32).\n",518" pos: ConfigDict that specifies the index of x, y, c, w, u in data tuple.\n",519" Default: data is of the form (x, y, c, w, u).\n",520" \"\"\"\n",521" self.model = None\n",522" self.inputs = \"x\"\n",523" self.main = main\n",524" self.aux = aux\n",525" self.dtype = dtype\n",526" if pos is None:\n",527" pos = mlc.ConfigDict()\n",528" pos.x, pos.y, pos.a = 0, 1, 2\n",529" self.pos = pos\n",530"\n",531" def get_input(self, *batch):\n",532" \"\"\"Fetch model input from the batch.\"\"\"\n",533" # first input\n",534" stack = tf.cast(batch[self.pos[self.inputs[0]]], self.dtype)\n",535" # fetch remaining ones\n",536" for c in self.inputs[1:]:\n",537" stack = tf.concat([stack, tf.cast(batch[self.pos[c]], self.dtype)],\n",538" axis=1)\n",539" return stack\n",540"\n",541" def get_output(self, *batch):\n",542" \"\"\"Fetch outputs from the batch.\"\"\"\n",543" if self.aux:\n",544" return (tf.cast(batch[self.pos[self.main]],self.dtype),\n",545" tf.cast(batch[self.pos[self.aux]], self.dtype))\n",546" else:\n",547" return (tf.cast(batch[self.pos[self.main]],self.dtype))\n",548"\n",549" def split_batch(self, *batch):\n",550" \"\"\"Split batch into input and output.\"\"\"\n",551" return self.get_input(*batch), self.get_output(*batch)\n",552"\n",553" def fit(self, data: tf.data.Dataset, **kwargs):\n",554" \"\"\"Fit model on data.\"\"\"\n",555" ds = data.map(self.split_batch)\n",556" self.model.fit(ds, **kwargs)\n",557"\n",558" def predict(self, model_input, **kwargs):\n",559" \"\"\"Predict target Y given the model input. See also: predict_mult().\"\"\"\n",560" y_pred = self.model.predict(model_input, **kwargs)\n",561" return y_pred\n",562"\n",563" def predict_mult(self, data: tf.data.Dataset, num_batches: int, **kwargs):\n",564" \"\"\"Predict target Y from the TF dataset directly. See also: predict().\"\"\"\n",565" y_true = []\n",566" y_pred = []\n",567" ds_iter = iter(data)\n",568" for _ in range(num_batches):\n",569" batch = next(ds_iter)\n",570" model_input, y = self.split_batch(*batch)\n",571" y_true.extend(y)\n",572" y_pred.extend(self.predict(model_input, **kwargs))\n",573" return np.array(y_true), np.array(y_pred)\n",574"\n",575" def score(self, data: tf.data.Dataset, num_batches: int, \n",576" metric: tf.keras.metrics.Metric , **kwargs):\n",577" \"\"\"Evaluate model on data.\n",578"\n",579" Args:\n",580" data: TF dataset.\n",581" num_batches: number of batches fetched from the dataset.\n",582" metric: which metric to evaluate (schrouf not be instantiated).\n",583" **kwargs: arguments passed to predict() method.\n",584"\n",585" Returns:\n",586" score: evaluation score.\n",587" \"\"\"\n",588" y_true, y_pred = self.predict_mult(data, num_batches, **kwargs)\n",589" return metric()(y_true, y_pred).numpy()\n",590"\n",591"\n",592"class MultiHead(BaselineArch):\n",593" \"\"\"Multihead training.\"\"\"\n",594"\n",595" def __init__(self, cfg, main, aux, dtype=tf.float32, pos=None): \n",596" \"\"\"Initializer.\n",597"\n",598" Args:\n",599" cfg: A config that describes the MLP architecture.\n",600" main: variable for the main task\n",601" aux: variable for the auxialiary task\n",602" dtype: desired dtype (e.g. tf.float32) for casting data.\n",603" \"\"\"\n",604" super(MultiHead, self).__init__(main, aux, dtype, pos)\n",605" self.main = \"y\"\n",606" self.aux = \"a\"\n",607" self.cfg = cfg\n",608" # build architecture\n",609" self.model, self.feat_extract = self.build()\n",610"\n",611" def build(self):\n",612" \"\"\"Build model.\"\"\"\n",613" cfg = self.cfg\n",614" input_shape = cfg.model.x_dim\n",615"\n",616" # set config params to defaults if missing\n",617" use_bias = cfg.model.get(\"use_bias\", True)\n",618" activation = cfg.model.get(\"activation\", \"relu\")\n",619" output_activation = cfg.model.get(\"output_activation\", \"sigmoid\")\n",620"\n",621" model_input = tf.keras.Input(shape=input_shape)\n",622" flatten_input = tf.keras.layers.Flatten()(model_input)\n",623" if cfg.model.depth:\n",624" x = tf.keras.layers.Dense(cfg.model.width, use_bias=use_bias,\n",625" activation=activation,\n",626" kernel_regularizer=cfg.model.regularizer)(flatten_input)\n",627" for _ in range(cfg.model.depth - 1):\n",628" x = tf.keras.layers.Dense(cfg.model.width, use_bias=use_bias,\n",629" activation=activation,\n",630" kernel_regularizer=cfg.model.regularizer)(x)\n",631" else:\n",632" x = flatten_input\n",633" feature_extractor = tf.keras.models.Model(inputs=flatten_input,\n",634" outputs=x)\n",635" # output layer - a single linear layer\n",636" y = tf.keras.layers.Dense(cfg.model.output_dim,\n",637" use_bias=cfg.model.use_bias,\n",638" name=\"output\",\n",639" activation=output_activation,\n",640" kernel_regularizer=cfg.model.regularizer)(x)\n",641" # attribute layer - an extra dense layer is required for gradients to flow back\n",642" attr_activation = cfg.model.get(\"attr_activation\", \"sigmoid\")\n",643" input_branch_a = GradientReversal(hp_lambda=cfg.model.attr_grad_updates)(x)\n",644" a_branch = tf.keras.layers.Dense(cfg.model.branch_dim,\n",645" use_bias=cfg.model.use_bias,\n",646" name=\"attr_branch\",\n",647" activation=activation,\n",648" kernel_regularizer=cfg.model.regularizer)(input_branch_a)\n",649" a = tf.keras.layers.Dense(cfg.model.attr_dim,\n",650" use_bias=cfg.model.use_bias,\n",651" name=\"attribute\",\n",652" activation=attr_activation,\n",653" kernel_regularizer=cfg.model.regularizer)(a_branch)\n",654" \n",655"\n",656"\n",657" # choose optimizer\n",658" if cfg.opt.name == \"sgd\":\n",659" opt = tf.keras.optimizers.SGD(learning_rate=cfg.opt.learning_rate,\n",660" momentum=cfg.opt.get(\"momentum\", 0.9))\n",661" elif cfg.opt.name == \"adam\":\n",662" opt = tf.keras.optimizers.Adam(learning_rate=cfg.opt.learning_rate)\n",663" else:\n",664" raise ValueError(\"Unrecognized optimizer type.\"\n",665" \"Please select either 'sgd' or 'adam'.\")\n",666"\n",667" # define losses\n",668" losses = {\n",669" \"output\": cfg.model.get(\"output_loss\", \"binary_crossentropy\"),\n",670" \"attribute\": cfg.model.get(\"attribute_loss\", \"binary_crossentropy\")\n",671" }\n",672" loss_weights = {\"output\": 1.0,\n",673" \"attribute\": cfg.get(\"attr_loss_weight\", 1.0)}\n",674" metrics = {\"output\": tf.keras.metrics.AUC(),\n",675" \"attribute\": tf.keras.metrics.AUC()}\n",676"\n",677" # build model\n",678" model = tf.keras.models.Model(inputs=model_input, outputs=[y,a])\n",679" model.build(input_shape)\n",680" # model.compile(optimizer=opt, loss=tf.keras.losses.BinaryCrossentropy(),\n",681" # metrics=tf.keras.metrics.BinaryAccuracy())\n",682" model.compile(optimizer=opt, loss=losses, loss_weights=loss_weights,\n",683" metrics=metrics)\n",684" return model, feature_extractor\n",685"\n",686" def predict_mult(self, data: tf.data.Dataset, num_batches: int, **kwargs):\n",687" \"\"\"Predict from the TF dataset directly. See also: predict().\"\"\"\n",688" # infer dimensions\n",689" pos = self.pos\n",690" batch = next(iter(data))\n",691" y_dim = batch[pos.y].shape[1]\n",692" a_dim = batch[pos.a].shape[1]\n",693"\n",694" # begin\n",695" data_iter = iter(data)\n",696" a_true_all = np.array([]).reshape((0, a_dim))\n",697" a_pred_all = np.array([]).reshape((0, a_dim))\n",698" y_true_all = np.array([]).reshape((0, y_dim))\n",699" y_pred_all = np.array([]).reshape((0, y_dim))\n",700"\n",701" for _ in range(num_batches):\n",702" batch = next(data_iter)\n",703" x, y_true, a_true = batch[pos.x], batch[pos.y], batch[pos.a]\n",704" y_pred, a_pred = self.predict(x, **kwargs)\n",705" a_true_all = np.append(a_true_all, a_true, axis=0)\n",706" a_pred_all = np.append(a_pred_all, a_pred, axis=0)\n",707" y_true_all = np.append(y_true_all, y_true, axis=0)\n",708" y_pred_all = np.append(y_pred_all, y_pred, axis=0)\n",709"\n",710" return (y_true_all, a_true_all), (y_pred_all, a_pred_all)\n",711"\n",712" def score(self, data: tf.data.Dataset, num_batches: int, \n",713" metric: tf.keras.metrics.Metric, **kwargs):\n",714" \"\"\"Evaluate model on data.\n",715"\n",716" Args:\n",717" data: TF dataset.\n",718" num_batches: number of batches fetched from the dataset.\n",719" metric: which metric to evaluate (should not be instantiated).\n",720" **kwargs: arguments passed to predict() method.\n",721"\n",722" Returns:\n",723" score: evaluation score.\n",724" \"\"\"\n",725" out_true, out_pred = self.predict_mult(data, num_batches, **kwargs)\n",726" scores = []\n",727" for head in range(len(out_true)):\n",728" score = metric()(out_true[head], out_pred[head])\n",729" scores.append(score.numpy())\n",730" return scores\n",731"\n",732"\n",733"# Can be used as a single task model fully trained or from a pre-trained\n",734"# feature extractor\n",735"\n",736"class SingleHead(BaselineArch):\n",737" \"\"\"Singlehead training.\"\"\"\n",738"\n",739" def __init__(self, cfg, main, dtype=tf.float32, pos=None, feat_extract=None): \n",740" \"\"\"Initializer.\n",741"\n",742" Args:\n",743" cfg: A config that describes the MLP architecture.\n",744" main: variable for the main task\n",745" aux: variable for the auxialiary task\n",746" dtype: desired dtype (e.g. tf.float32) for casting data.\n",747" \"\"\"\n",748" super(SingleHead, self).__init__(main, None, dtype, pos)\n",749" self.main = \"a\"\n",750" self.cfg = cfg\n",751" # build architecture\n",752" self.model = self.build(feat_extract)\n",753"\n",754" def build(self, feat_extract=None):\n",755" \"\"\"Build model.\"\"\"\n",756" cfg = self.cfg\n",757" input_shape = cfg.model.x_dim\n",758"\n",759" # set config params to defaults if missing\n",760" use_bias = cfg.model.get(\"use_bias\", True)\n",761" activation = cfg.model.get(\"activation\", \"relu\")\n",762" output_activation = cfg.model.get(\"output_activation\", \"sigmoid\")\n",763"\n",764" model_input = tf.keras.Input(shape=input_shape)\n",765" flatten_input = tf.keras.layers.Flatten()(model_input)\n",766" if not feat_extract:\n",767" if cfg.model.depth:\n",768" x = tf.keras.layers.Dense(cfg.model.width, use_bias=use_bias,\n",769" activation=activation,\n",770" kernel_regularizer=cfg.model.regularizer)(flatten_input)\n",771" for _ in range(cfg.model.depth - 1):\n",772" x = tf.keras.layers.Dense(cfg.model.width, use_bias=use_bias,\n",773" activation=activation,\n",774" kernel_regularizer=cfg.model.regularizer)(x)\n",775" else:\n",776" x = flatten_input\n",777" feature_extractor = x\n",778" else:\n",779" feat_extract.trainable = False\n",780" feature_extractor = feat_extract(flatten_input, training=False)\n",781" \n",782" # output layer\n",783" y = tf.keras.layers.Dense(cfg.model.output_dim,\n",784" use_bias=cfg.model.use_bias,\n",785" name=\"output\",\n",786" activation=output_activation,\n",787" kernel_regularizer=cfg.model.regularizer)(feature_extractor) \n",788"\n",789" # choose optimizer\n",790" if cfg.opt.name == \"sgd\":\n",791" opt = tf.keras.optimizers.SGD(learning_rate=cfg.opt.learning_rate,\n",792" momentum=cfg.opt.get(\"momentum\", 0.9))\n",793" elif cfg.opt.name == \"adam\":\n",794" opt = tf.keras.optimizers.Adam(learning_rate=cfg.opt.learning_rate)\n",795" else:\n",796" raise ValueError(\"Unrecognized optimizer type.\"\n",797" \"Please select either 'sgd' or 'adam'.\")\n",798"\n",799" # build model\n",800" model = tf.keras.models.Model(inputs=model_input, outputs=y)\n",801" model.build(input_shape)\n",802" model.compile(optimizer=opt,\n",803" loss=cfg.model.get(\"output_loss\", \"binary_crossentropy\"),\n",804" metrics=tf.keras.metrics.AUC())\n",805"\n",806" return model"807]808},809{810"cell_type": "code",811"execution_count": null,812"metadata": {813"id": "VPeJWiIajIOX"814},815"outputs": [],816"source": [817"#@markdown Model parameters \n",818"cfg = mlc.ConfigDict()\n",819"\n",820"cfg.model = mlc.ConfigDict()\n",821"cfg.model.width = 10 #@param architecture width\n",822"cfg.model.depth = 3 #@param model depth\n",823"cfg.model.use_bias = True # whether we add biases to activations.\n",824"cfg.model.activation = 'relu'\n",825"cfg.model.x_dim = (n_pixels, n_pixels, n_channels)\n",826"cfg.model.branch_dim = 2 # architecture width within each branch\n",827"cfg.model.regularizer = None # replace with e.g. 'l2' for weight decay\n",828"\n",829"# output head\n",830"cfg.model.output_activation = 'sigmoid'\n",831"cfg.model.output_dim = 1\n",832"\n",833"# attribute head\n",834"cfg.model.attr_activation = 'sigmoid'\n",835"cfg.model.attr_grad_updates = float(-0.05)\n",836"cfg.model.attr_dim = 1\n",837"# this is a tradeoff between the loss on A and the loss on target Y.\n",838"# If it's zero, we ignore the attribute loss completely.\n",839"cfg.attr_loss_weight = float(1.0)\n",840"\n",841"cfg.opt = mlc.ConfigDict()\n",842"cfg.opt.name = 'adam'\n",843"cfg.opt.learning_rate = 0.001\n",844"cfg.opt.momentum = 0.9"845]846},847{848"cell_type": "markdown",849"metadata": {850"id": "cnJowASOjra_"851},852"source": [853"# Baseline task model"854]855},856{857"cell_type": "code",858"execution_count": null,859"metadata": {860"id": "gd5vA7vwjvSu"861},862"outputs": [],863"source": [864"metric = tf.keras.metrics.AUC\n",865"\n",866"baseline = []\n",867"for seed in [0,1,2,3,4]:\n",868" tf.random.set_seed(seed)\n",869" np.random.seed(seed)\n",870" enc = SingleHead(cfg, main=\"y\", dtype=tf.float32, feat_extract=None)\n",871" kwargs = {'epochs': 100, 'steps_per_epoch':20, 'verbose': False}\n",872" enc.fit(train_data, **kwargs)\n",873" kwargs = {'verbose': False}\n",874" num_batches = 30\n",875" sc = enc.score(valid_data, num_batches, metric, **kwargs)\n",876" baseline.append(sc)\n",877"\n",878"print(\"Baseline model: %1.2f +- %.2f\" % (np.mean(baseline),np.std(baseline)))"879]880},881{882"cell_type": "markdown",883"metadata": {884"id": "FMewMxgikOmD"885},886"source": [887"# Bounds on attribute encoding (i.e. LEB and UEB in paper)"888]889},890{891"cell_type": "code",892"execution_count": null,893"metadata": {894"id": "1kWfJz8tkNfu"895},896"outputs": [],897"source": [898"#@title Upper bound\n",899"metric = tf.keras.metrics.AUC\n",900"\n",901"upper_bound = []\n",902"for seed in [0,1,2,3,4]:\n",903" tf.random.set_seed(seed)\n",904" np.random.seed(seed)\n",905" enc = SingleHead(cfg, main=\"a\", dtype=tf.float32, feat_extract=None)\n",906" kwargs = {'epochs': 100, 'steps_per_epoch':20, 'verbose': False}\n",907" enc.fit(train_data, **kwargs)\n",908" kwargs = {'verbose': False}\n",909" num_batches = 30\n",910" sc = enc.score(valid_data, num_batches, metric, **kwargs)\n",911" upper_bound.append(sc)\n",912"\n",913"print(\"Upper bound: %1.2f +- %.2f\" % (np.mean(upper_bound),np.std(upper_bound)))"914]915},916{917"cell_type": "code",918"execution_count": null,919"metadata": {920"id": "zCavAI_Fki2W"921},922"outputs": [],923"source": [924"#@title Lower bound\n",925"# Predict a color randomly selected from the train set (with replacement)\n",926"\n",927"n_train, a_dim = dataset[\"train\"][\"color_label\"].shape\n",928"\n",929"lower_bound = []\n",930"for seed in [0,1,2,3,4]:\n",931" tf.random.set_seed(seed)\n",932" np.random.seed(seed)\n",933" data_iter = iter(train_data)\n",934" a_true_all = np.array([]).reshape((0, a_dim))\n",935" a_pred_all = np.array([]).reshape((0, a_dim))\n",936"\n",937" for _ in range(num_batches):\n",938" batch = next(data_iter)\n",939" a_true = batch[pos.a]\n",940" a_pred = dataset[\"train\"][\"color_label\"][np.random.choice(n_train,a_true.shape[0]), ...]\n",941" preds = a_pred.reshape((-1,a_dim))\n",942" a_true_all = np.append(a_true_all, a_true, axis=0)\n",943" a_pred_all = np.append(a_pred_all, preds, axis=0)\n",944"\n",945" sc = tf.keras.metrics.AUC()(a_true_all, a_pred_all).numpy()\n",946" lower_bound.append(sc)\n",947"\n",948"print(\"Lower bound: %1.2f +- %.2f\" % (np.mean(lower_bound),np.std(lower_bound)))"949]950},951{952"cell_type": "markdown",953"metadata": {954"id": "9eY_qcOPkyo5"955},956"source": [957"# Example of Multi-Task model"958]959},960{961"cell_type": "code",962"execution_count": null,963"metadata": {964"id": "IHFinRviksIU"965},966"outputs": [],967"source": [968"#@title Initialize from config\n",969"cfg.model.attr_grad_updates = float(0.1) # if both are set to 0, this is a baseline model (i.e. the same as with no extra head)\n",970"cfg.attr_loss_weight = float(0.75)\n",971"clf = MultiHead(cfg, main=\"y\", aux=\"a\", dtype=tf.float32, pos=None)"972]973},974{975"cell_type": "code",976"execution_count": null,977"metadata": {978"id": "LrriExy3k8cc"979},980"outputs": [],981"source": [982"#@markdown Explore the architecture\n",983"clf.model.summary()"984]985},986{987"cell_type": "code",988"execution_count": null,989"metadata": {990"id": "rHyqFeJelCQM"991},992"outputs": [],993"source": [994"#@title Train the model\n",995"kwargs = {'epochs': 100, 'steps_per_epoch':20, 'verbose': True}\n",996"clf.fit(train_data, **kwargs)\n",997"clf.trainable=False"998]999},1000{1001"cell_type": "code",1002"execution_count": null,1003"metadata": {1004"id": "n1jZXkiilipk"1005},1006"outputs": [],1007"source": [1008"#@title Define the decision threshold by maximizing the F1-score on validation data\n",1009"\n",1010"from sklearn.metrics import precision_recall_curve\n",1011"def f1_curve(truth, prediction_scores, e=1e-6):\n",1012" precision, recall, thresholds = precision_recall_curve(truth, prediction_scores)\n",1013" f1 = 2*recall*precision/(recall+precision+e)\n",1014" return thresholds, f1[:-1]\n",1015"\n",1016"def threshold_at_max_f1_score(truth, prediction_scores):\n",1017" thresholds, f1 = f1_curve(truth, prediction_scores)\n",1018" peak_idx = np.argmax(f1)\n",1019" return thresholds[peak_idx]"1020]1021},1022{1023"cell_type": "code",1024"execution_count": null,1025"metadata": {1026"id": "CmlOsACMltYq"1027},1028"outputs": [],1029"source": [1030"kwargs = {'verbose': False}\n",1031"num_batches = 20\n",1032"metric = tf.keras.metrics.AUC\n",1033"scores_val = clf.score(valid_data, num_batches, metric, **kwargs)\n",1034"print(scores_val)"1035]1036},1037{1038"cell_type": "code",1039"execution_count": null,1040"metadata": {1041"id": "KjztJ34UlxIq"1042},1043"outputs": [],1044"source": [1045"kwargs = {'verbose': False}\n",1046"yt, yp = clf.predict_mult(valid_data, num_batches, **kwargs)\n",1047"threshold = threshold_at_max_f1_score(yt[0],yp[0])\n",1048"threshold"1049]1050},1051{1052"cell_type": "code",1053"execution_count": null,1054"metadata": {1055"cellView": "form",1056"id": "ZviOthNFlycK"1057},1058"outputs": [],1059"source": [1060"#@title Fairness metrics\n",1061"\n",1062"# As per the work of Alabdulmohsin et al., 2021\n",1063"\n",1064"def fairness_metrics(y_pred, y_true, sens_attr):\n",1065" eps = 1e-5\n",1066" groups = np.unique(sens_attr).tolist()\n",1067"\n",1068" max_error = 0\n",1069" min_error = 1\n",1070"\n",1071" max_mean_y = 0\n",1072" min_mean_y = 1\n",1073"\n",1074" max_mean_y0 = 0 # conditioned on y = 0\n",1075" min_mean_y0 = 1\n",1076"\n",1077" max_mean_y1 = 0\n",1078" min_mean_y1 = 1\n",1079"\n",1080" for group in groups:\n",1081" yt = y_true[sens_attr == group].astype('int32')\n",1082" ypt = (y_pred[sens_attr == group]).astype('int32')\n",1083" err = -np.mean(yt * np.log(ypt+eps) + (1-yt)*np.log(1-ypt+eps))\n",1084" mean_y = np.mean(y_pred[sens_attr == group])\n",1085" neg = np.logical_and(sens_attr == group, y_true == 0)\n",1086" pos = np.logical_and(sens_attr == group, y_true == 1)\n",1087" mean_y0 = np.mean(y_pred[neg])\n",1088" mean_y1 = np.mean(y_pred[pos])\n",1089"\n",1090" if err \u003e max_error:\n",1091" max_error = err\n",1092" if err \u003c min_error:\n",1093" min_error = err\n",1094"\n",1095" if mean_y \u003e max_mean_y:\n",1096" max_mean_y = mean_y\n",1097" if mean_y \u003c min_mean_y:\n",1098" min_mean_y = mean_y\n",1099"\n",1100" if mean_y0 \u003e max_mean_y0:\n",1101" max_mean_y0 = mean_y0\n",1102" if mean_y0 \u003c min_mean_y0:\n",1103" min_mean_y0 = mean_y0\n",1104"\n",1105" if mean_y1 \u003e max_mean_y1:\n",1106" max_mean_y1 = mean_y1\n",1107" if mean_y1 \u003c min_mean_y1:\n",1108" min_mean_y1 = mean_y1\n",1109" \n",1110" eo = 0.5*(max_mean_y0 - min_mean_y0 + max_mean_y1 - min_mean_y1)\n",1111" dp = max_mean_y - min_mean_y\n",1112" err_parity = max_error - min_error\n",1113"\n",1114" return eo, dp, err_parity"1115]1116},1117{1118"cell_type": "code",1119"execution_count": null,1120"metadata": {1121"id": "YxTljd9CmD6S"1122},1123"outputs": [],1124"source": [1125"#@title Evaluate the model on test data\n",1126"kwargs = {'verbose': False}\n",1127"num_batches = 20\n",1128"metric = tf.keras.metrics.AUC\n",1129"scores = clf.score(test_data, num_batches, metric, **kwargs)\n",1130"scores"1131]1132},1133{1134"cell_type": "code",1135"execution_count": null,1136"metadata": {1137"id": "CYTlaLromSji"1138},1139"outputs": [],1140"source": [1141"yt, yp = clf.predict_mult(test_data, num_batches, **kwargs)\n",1142"eo, dp, ep = fairness_metrics(yp[0]\u003e=threshold,\n",1143" yt[0],yt[1].astype('int32'))"1144]1145},1146{1147"cell_type": "code",1148"execution_count": null,1149"metadata": {1150"id": "rr9GtWwjmjHx"1151},1152"outputs": [],1153"source": [1154"#@title Evaluate model on counterfactual test data\n",1155"kwargs = {'verbose': False}\n",1156"num_batches = 20\n",1157"metric = tf.keras.metrics.AUC\n",1158"\n",1159"scores_flipped = clf.score(test_data_flipped, num_batches,metric, **kwargs)\n",1160"scores_flipped"1161]1162},1163{1164"cell_type": "code",1165"execution_count": null,1166"metadata": {1167"id": "r4YSRAk6mrow"1168},1169"outputs": [],1170"source": [1171"#@title Evaluate attribute encoding based on a frozen feature extractor from the multi-task model\n",1172"enc = SingleHead(cfg, main=\"a\", dtype=tf.float32, feat_extract=clf.feat_extract)\n",1173"kwargs = {'epochs': 30, 'steps_per_epoch':20, 'verbose': False}\n",1174"enc.fit(train_data, **kwargs)\n",1175"kwargs = {'verbose': False}\n",1176"num_batches = 20\n",1177"metric = tf.keras.metrics.AUC\n",1178"sc = enc.score(test_data, num_batches, metric, **kwargs)\n",1179"sc"1180]1181},1182{1183"cell_type": "markdown",1184"metadata": {1185"id": "UczpCETAmyOu"1186},1187"source": [1188"# Piecing the elements together: ShorT"1189]1190},1191{1192"cell_type": "code",1193"execution_count": null,1194"metadata": {1195"id": "jGsb0Tpsm37w"1196},1197"outputs": [],1198"source": [1199"#@title ShorT method\n",1200"def shortcut_testing(cfg, train_data, test_data, counterfact_data, val_data,\n",1201" range_grads, seeds = [0,1,2,3,4], num_epochs_train = 30,\n",1202" num_batch_test = 5):\n",1203" \"\"\"Shortcut testing.\n",1204"\n",1205" Requires:\n",1206" - cfg: a config dictionary of model specifications\n",1207" - train_data: tf dataset of train data respecting the positions for x,y,a\n",1208" - test_data: tf dataset of test data\n",1209" - counterfact_data: tf dataset of counterfactual test images\n",1210" - val_data: validation data to compute decision threshold\n",1211" - range_grads: np array of gradient scalings\n",1212" - seeds: list of seed numbers to fix, per gradient scaling\n",1213" - num_epoch_train: int, number of epochs for training\n",1214" - num_batch_test: number of batches to test the model on\n",1215"\n",1216" Outputs:\n",1217" scores_y: np array of size (# gradient scalings, # seeds) of test performance on Y\n",1218" encoding_a: np array of model's attribute encoding\n",1219" equ_odds: np array of equalized odds on test data\n",1220" dem_par: similar, demographic parity\n",1221" err_par: error parity\n",1222" count_fair: counterfactual fairness (computed locally then averaged)\n",1223" scores_c: global performance on Y on counterfactual data\n",1224" models: tf.Keras.model instances\n",1225" \"\"\" \n",1226"\n",1227"\n",1228" scores_y = np.zeros((len(range_grads), len(seeds)))\n",1229" scores_c = np.zeros((len(range_grads), len(seeds)))\n",1230" encoding_a = np.zeros((len(range_grads), len(seeds)))\n",1231" equ_odds = np.zeros((len(range_grads), len(seeds)))\n",1232" dem_par = np.zeros((len(range_grads), len(seeds)))\n",1233" err_par = np.zeros((len(range_grads), len(seeds)))\n",1234" count_fair = np.zeros((len(range_grads), len(seeds)))\n",1235" models = []\n",1236" metric = tf.keras.metrics.AUC\n",1237" loss_weight = cfg.attr_loss_weight\n",1238" for g, grad_scaling in enumerate(range_grads):\n",1239" cfg.model.attr_grad_updates = float(grad_scaling)\n",1240"\n",1241" if grad_scaling == 0: ## baseline model with no other head\n",1242" cfg.attr_loss_weight = float(0.0)\n",1243" else:\n",1244" cfg.attr_loss_weight = float(loss_weight)\n",1245" \n",1246" for s, seed in enumerate(seeds):\n",1247" tf.random.set_seed(seed)\n",1248" np.random.seed(seed)\n",1249" # instantiate\n",1250" clf = MultiHead(cfg, main=\"y\", aux=\"a\", dtype=tf.float32, pos=None)\n",1251" # train the multi-head model\n",1252" kwargs = {'epochs': num_epochs_train, 'steps_per_epoch':20, 'verbose': False}\n",1253" clf.fit(train_data, **kwargs)\n",1254" clf.trainable = False\n",1255"\n",1256" # estimate attribute encoding by freezing the weights of the feature extractor\n",1257" enc = SingleHead(cfg, main=\"a\", dtype=tf.float32, feat_extract=clf.feat_extract)\n",1258" kwargs = {'epochs': num_epochs_train, 'steps_per_epoch':20, 'verbose': False}\n",1259" enc.fit(train_data, **kwargs)\n",1260" kwargs = {'verbose': False}\n",1261" encoding_a[g,s] = enc.score(test_data, num_batch_test,\n",1262" metric, **kwargs)\n",1263"\n",1264" # Model performance on test\n",1265" scores = clf.score(test_data, num_batch_test, metric, **kwargs)\n",1266" scores_y[g,s] = scores[0]\n",1267"\n",1268" # from validation data, obtain threshold for decision making\n",1269" out_true, out_pred = clf.predict_mult(val_data,\n",1270" num_batches=num_batch_test,\n",1271" **kwargs)\n",1272" threshold = threshold_at_max_f1_score(out_true[0],out_pred[0])\n",1273"\n",1274" # Model statistical fairness metric\n",1275" out_true, out_pred = clf.predict_mult(test_data, num_batches=num_batch_test, **kwargs)\n",1276" eo, dp, ep = fairness_metrics(out_pred[0]\u003e=threshold,\n",1277" out_true[0],out_true[1].astype('int32'))\n",1278" equ_odds[g,s] = eo\n",1279" dem_par[g,s] = dp\n",1280" err_par[g,s] = ep\n",1281" # global shortcutting\n",1282" scores = clf.score(counterfact_data, num_batch_test, metric, **kwargs)\n",1283" scores_c[g,s] = scores[0]\n",1284" # counterfactual fairness metric\n",1285" cf_true, cf_pred = clf.predict_mult(counterfact_data, num_batches=num_batch_test,\n",1286" **kwargs)\n",1287" count_fair[g,s] = np.mean(np.abs(\n",1288" (out_pred[0]\u003e=threshold).astype(float) - (cf_pred[0]\u003e=threshold).astype(float)))\n",1289" models.append(clf)\n",1290" del clf\n",1291" return scores_y, encoding_a, equ_odds, dem_par, err_par, count_fair, scores_c, models"1292]1293},1294{1295"cell_type": "code",1296"execution_count": null,1297"metadata": {1298"id": "cQy-hQvqorsx"1299},1300"outputs": [],1301"source": [1302"#@title Run ShorT\n",1303"\n",1304"# Ensure that these parameters cover the attribute encoding from LEB to UEB\n",1305"range_grads = [-0.09, -0.07, -0.05, -0.03, -0.02, -0.01, -0.005,0.0, 0.005, 0.01, 0.02, 0.03, 0.05,0.07, 0.09]\n",1306"\n",1307"# This should take some time to run\n",1308"acc_y, acc_a, eo, dp, ep, cf, counter_y, models = shortcut_testing(cfg, train_data,\n",1309" test_data, valid_data,\n",1310" test_data_flipped,\n",1311" range_grads,\n",1312" num_epochs_train=100,\n",1313" num_batch_test=20)\n"1314]1315},1316{1317"cell_type": "code",1318"execution_count": null,1319"metadata": {1320"id": "Aiw0SdbbpKw8"1321},1322"outputs": [],1323"source": [1324"#@markdown Derive the correlation\n",1325"\n",1326"filt = acc_y \u003e= 0.8 #@param User-defined performance threshold. The goal is to discard trivial models\n",1327"corr, p = scipy.stats.spearmanr(acc_a[filt],cf[filt])\n",1328"print(corr)\n",1329"print(p)\n"1330]1331},1332{1333"cell_type": "code",1334"execution_count": null,1335"metadata": {1336"cellView": "form",1337"id": "WK3lk32PpfoT"1338},1339"outputs": [],1340"source": [1341"#@title Plot utils\n",1342"\n",1343"SMALL_SIZE = 10\n",1344"MEDIUM_SIZE = 14\n",1345"BIGGER_SIZE = 18\n",1346"\n",1347"plt.rc('font', size=SMALL_SIZE) # controls default text sizes\n",1348"plt.rc('axes', titlesize=BIGGER_SIZE) # fontsize of the axes title\n",1349"plt.rc('axes', labelsize=MEDIUM_SIZE) # fontsize of the x and y labels\n",1350"plt.rc('xtick', labelsize=MEDIUM_SIZE) # fontsize of the tick labels\n",1351"plt.rc('ytick', labelsize=MEDIUM_SIZE) # fontsize of the tick labels\n",1352"plt.rc('legend', fontsize=MEDIUM_SIZE) # legend fontsize\n",1353"plt.rc('figure', titlesize=BIGGER_SIZE) # fontsize of the figure title\n",1354"\n",1355"palette = sns.color_palette('tab20b')\n",1356"\n",1357"def plot_scale_gradients_encoding(scale_gradients, encoding_a, \n",1358" upper_bound, lower_bound):\n",1359" \"\"\"Plots the intervention compared to attribute encoding.\n",1360"\n",1361" Inputs:\n",1362" scale_gradients: numpy array of scale gradients used\n",1363" encoding_a: numpy array of attribute encoding values, as measure by transfer learning\n",1364" upper_bound: list or numpy array of maximum attribute encoding\n",1365" lower_bound: list or numpy array of minimum attribute encoding\n",1366" \"\"\"\n",1367" \n",1368" sg_mean = np.mean(encoding_a, axis=1)\n",1369" sg_std = np.std(encoding_a, axis=1)\n",1370"\n",1371" fig = plt.figure(figsize=(5,4))\n",1372" ax = fig.add_axes([0,0,1,1])\n",1373" ax.errorbar(scale_gradients, sg_mean, \n",1374" yerr=sg_std, \n",1375" fmt='x', \n",1376" color='tab:blue',\n",1377" ecolor='tab:blue')\n",1378" plt.hlines(np.mean(upper_bound),np.min(scale_gradients), np.max(scale_gradients),\n",1379" colors=[0.4,0.4,0.4],linestyles='dashed')\n",1380" plt.hlines(np.mean(lower_bound),np.min(scale_gradients), np.max(scale_gradients),\n",1381" colors=[0.4,0.4,0.4],linestyles='dashed')\n",1382" ax.set_xlabel('Scale Gradient')\n",1383" ax.set_ylabel('Attribute encoding')\n",1384" plt.show\n",1385"\n",1386"def plot_fairness_encoding(encoding_m, fair_m, perf_m, perf_thresh = 0):\n",1387" \"\"\"Plots fairness results vs attribute encoding.\n",1388" \n",1389" Inputs:\n",1390" encoding_m: encoding metric result, as numpy array\n",1391" fair_m: fairness metric result, as numpy array\n",1392" perf_m: model performance on the output label, as numpy array\n",1393" perf_thresh: what minimum performance to consider\n",1394" \"\"\"\n",1395"\n",1396" filt = perf_m \u003c= perf_thresh\n",1397"\n",1398" fig = plt.figure(figsize=(5,4))\n",1399" plt.scatter(encoding_m, fair_m,color=[0.6,0,0.2],alpha=0.5)\n",1400" plt.scatter(encoding_m[filt], fair_m[filt],color=[0.2,0.2,0.2])\n",1401" plt.xlabel('Attribute Accuracy')\n",1402" plt.ylabel('Fairness')\n",1403" plt.show()\n",1404"\n",1405"def point_z_order(c, midpoint):\n",1406" deviation = np.zeros_like(c)\n",1407" for i in range(c.shape[0]):\n",1408" if c[i] \u003e midpoint:\n",1409" deviation[i] = (c[i]-midpoint) / (np.max(c)-midpoint)\n",1410" else:\n",1411" deviation[i] = (midpoint-c[i]) / (midpoint-np.min(c))\n",1412" return np.argsort(deviation)\n",1413"\n",1414"def performance_fairness_age_frontier_plot(encoding_m, fair_m, perf_m,\n",1415" scale_gradients, cmap='PRGn'):\n",1416" \n",1417" class MidpointNormalize(mpl.colors.Normalize):\n",1418" \"\"\"\n",1419" class to help renormalize the color scale\n",1420" \"\"\"\n",1421" def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):\n",1422" self.midpoint = midpoint\n",1423" mpl.colors.Normalize.__init__(self, vmin, vmax, clip)\n",1424"\n",1425" def __call__(self, value, clip=None):\n",1426" x, y = [self.vmin, self.midpoint, self.vmax], [0, 0.5, 1]\n",1427" return np.ma.masked_array(np.interp(value, x, y))\n",1428"\n",1429" baseline_models = np.argwhere(np.array(scale_gradients)==0.0)[0]\n",1430" if baseline_models.shape[0] == 0:\n",1431" baseline_models = int(len(scale_gradients)/2)\n",1432" midpoint = encoding_m[baseline_models,:].mean()\n",1433" baseline_model_perf = perf_m[baseline_models,:].mean()\n",1434" baseline_model_fair = fair_m[baseline_models,:].mean()\n",1435"\n",1436" print(f'Baseline model Attribute encoding: {midpoint:.2f}')\n",1437" print(f'Baseline model Performance: {baseline_model_perf:.4f}')\n",1438" print(f'Baseline model Fairness: {baseline_model_fair:.4f}')\n",1439"\n",1440" norm = MidpointNormalize(midpoint = midpoint)\n",1441"\n",1442" attr = encoding_m.flatten()\n",1443" z_order = point_z_order(attr, midpoint)\n",1444"\n",1445" fair =fair_m.flatten()\n",1446" perf = perf_m.flatten()\n",1447" fig = plt.figure(figsize=(5,4))\n",1448" ax = fig.add_axes([0,0,1,1])\n",1449" plt.scatter(fair[z_order], perf[z_order], s=30, c=attr[z_order], cmap=cmap, norm=norm)\n",1450" # overplot the baseline models in red\n",1451" plt.scatter(fair_m[baseline_models,:], perf_m[baseline_models,:], s=30,\n",1452" color=(0.8, 0.2, 0.2))\n",1453" plt.colorbar(cm.ScalarMappable(norm=norm, cmap=cmap),label='Attribute Encoding')\n",1454" plt.ylabel('Performance')\n",1455" plt.xlabel('Fairness')\n",1456" plt.show()"1457]1458},1459{1460"cell_type": "code",1461"execution_count": null,1462"metadata": {1463"id": "0HJjIltJpFbr"1464},1465"outputs": [],1466"source": [1467"#@title Plot scale gradient vs age encoding\n",1468"plot_scale_gradients_encoding(range_grads, acc_a, upper_bound, lower_bound)"1469]1470},1471{1472"cell_type": "code",1473"execution_count": null,1474"metadata": {1475"id": "sHORVCqBpnFT"1476},1477"outputs": [],1478"source": [1479"#@title Plot Fairness w.r.t. attribute encoding\n",1480"suffix = \"{}_EO\".format(cfg_data.corr_a_y1)\n",1481"plot_fairness_encoding(acc_a, cf, acc_y, perf_thresh = 0.8)"1482]1483},1484{1485"cell_type": "code",1486"execution_count": null,1487"metadata": {1488"id": "Gfx8TY70pt_D"1489},1490"outputs": [],1491"source": [1492"#@title Pareto plot of fairness-performance colored by attribute encoding\n",1493"performance_fairness_age_frontier_plot(acc_a,eo,acc_y,range_grads)"1494]1495},1496{1497"cell_type": "code",1498"execution_count": null,1499"metadata": {1500"id": "QZIM0AucYMhX"1501},1502"outputs": [],1503"source": [1504"#@title Selecting a model and dropping the attribute head for deployment\n",1505"\n",1506"fairness_threshold = 0.09 #@param\n",1507"gradient_parameters = np.array(range_grads).reshape(-1,1).repeat(5,axis=1)\n",1508"\n",1509"possible_models = acc_y[eo\u003c=fairness_threshold]\n",1510"possible_gradients = gradient_parameters[eo\u003c=fairness_threshold]\n",1511"filt = eo\u003c=fairness_threshold\n",1512"\n",1513"ind = np.argmax(possible_models)\n",1514"selected_gradient = float(possible_gradients[ind])\n",1515"selected_model = models[np.argwhere(filt.flatten()).tolist()[ind][0]]\n",1516"print(\"Selected gradient parameter: {}\".format(selected_gradient))\n",1517"print(\"Selected model performance: {}\".format(possible_models[ind]))\n",1518"print(\"Selected model fairness: {}\".format(eo[eo\u003c=fairness_threshold][ind]))"1519]1520},1521{1522"cell_type": "code",1523"execution_count": null,1524"metadata": {1525"id": "9LRDs_1eewqb"1526},1527"outputs": [],1528"source": [1529"selected_model.model.summary()"1530]1531},1532{1533"cell_type": "code",1534"execution_count": null,1535"metadata": {1536"id": "_NjgGzo6aVOD"1537},1538"outputs": [],1539"source": [1540"# Create a new model by selecting all layers except the ones related to the attribute\n",1541"remove = ['attribute', 'attr_branch', 'gradient'] # corresponds to the names of the layer (partial matches ok), as defined in the MultiHead class\n",1542"\n",1543"final_model = tf.keras.Sequential()\n",1544"for layer in clf.model.layers:\n",1545" match = [to_pop in layer.name for to_pop in remove]\n",1546" if not any(match):\n",1547" final_model.add(layer)\n",1548" layer.trainable = False\n"1549]1550},1551{1552"cell_type": "code",1553"execution_count": null,1554"metadata": {1555"id": "waoOroZOhiF0"1556},1557"outputs": [],1558"source": [1559"final_model.summary()"1560]1561},1562{1563"cell_type": "code",1564"execution_count": null,1565"metadata": {1566"id": "mT0Ti50ChoIj"1567},1568"outputs": [],1569"source": [1570"x = test_dataset['train']['image']\n",1571"y_pred = final_model.predict_on_batch(x)"1572]1573},1574{1575"cell_type": "code",1576"execution_count": null,1577"metadata": {1578"id": "PlZ_04T9imOp"1579},1580"outputs": [],1581"source": [1582"sklearn.metrics.roc_auc_score(test_dataset['train']['label'], y_pred)"1583]1584}1585],1586"metadata": {1587"colab": {1588"collapsed_sections": [1589"xujPFv_tgv_R",1590"cnJowASOjra_",1591"FMewMxgikOmD",1592"9eY_qcOPkyo5",1593"UczpCETAmyOu"1594],1595"private_outputs": true,1596"provenance": [1597{1598"file_id": "1YbpbR1q6asdqS_ZfwuKl1FM__OLiGtXE",1599"timestamp": 16650776354491600}1601],1602"toc_visible": true1603},1604"kernelspec": {1605"display_name": "Python 3",1606"name": "python3"1607},1608"language_info": {1609"name": "python"1610}1611},1612"nbformat": 4,1613"nbformat_minor": 01614}1615