[FedMix]breast_experiment.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef550afc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import copy\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "\n",
    "from torch import optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.nn import functional as F\n",
    "from unet import UNet\n",
    "from dice_loss import dice_coeff\n",
    "\n",
    "####################################################\n",
    "# for data preparation\n",
    "####################################################\n",
    "from sklearn.model_selection import train_test_split, StratifiedShuffleSplit\n",
    "from sklearn.metrics import roc_auc_score, accuracy_score, balanced_accuracy_score\n",
    "####################################################\n",
    "# for plotting\n",
    "####################################################\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output\n",
    "############################\n",
    "# Helper func\n",
    "############################\n",
    "from helper import * \n",
    "#################################\n",
    "TRAIN_RATIO = 0.8\n",
    "RS = 30448 # random state\n",
    "N_CHANNELS, N_CLASSES = 1, 1 \n",
    "bilinear = True\n",
    "BATCH_SIZE, EPOCHS = 16, 300\n",
    "IMAGE_SIZE = (256, 256)\n",
    "CROP_SIZE = (224, 224)\n",
    "#########################################\n",
    "DIR = 'dataset/breast'\n",
    "CLIENTS = ['BUS', 'BUSIS', 'UDIAT']\n",
    "CLIENTS_2 = [cl +'_2' for cl in CLIENTS]\n",
    "TOTAL_CLIENTS = len(CLIENTS)\n",
    "#####################################\n",
    "# add the classification segment ####\n",
    "#####################################\n",
    "DIR_CLASSIFICATION = DIR + '/BUS/classification'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b596b2eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONSISTENT_PATH = 'breast_5folds.npy'\n",
    "device = torch.device('cuda:0')\n",
    "LR, WD = 1e-3, 1e-4\n",
    "\n",
    "LAM, BETA, TH = 10, 1.5,0.9\n",
    "VERSION = 1\n",
    "WEIGHTS_CL = [0.0, 0.0, 0.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eed770c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CLIENTS_SUPERVISION = ['unlabeled', 'unlabeled', 'labeled']\n",
    "# CLIENTS_SUPERVISION = ['image', 'unlabeled', 'labeled']\n",
    "CLIENTS_SUPERVISION = ['image', 'bbox', 'labeled']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb798488",
   "metadata": {},
   "source": [
    "## Training path - Testing path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2522f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "consistent_path = np.load(CONSISTENT_PATH, allow_pickle=True).item()\n",
    "breast_dataset = dict()\n",
    "\n",
    "idx_ = 0\n",
    "denom_ = 0\n",
    "for client, sup in zip(CLIENTS, CLIENTS_SUPERVISION):\n",
    "    dir_of_interest = consistent_path[client][VERSION]\n",
    "    x_train = dir_of_interest['x_train']\n",
    "    x_test = dir_of_interest['x_test']\n",
    "    y_train = dir_of_interest['y_train']\n",
    "    y_test = dir_of_interest['y_test']\n",
    "    '''\n",
    "    add normal images \n",
    "    '''\n",
    "    if sup == 'unlabeled':\n",
    "#     if client != 'UDIAT': \n",
    "        DATA_TYPE = ['original', 'GT']\n",
    "        for _,_, files in os.walk(DIR + '/'+ 'BUS' +'/classification/GT'):\n",
    "            selected = [f for f in files if f[:6] =='normal']\n",
    "            # update accordingly #\n",
    "            for data in DATA_TYPE:\n",
    "                tmp = [DIR + '/'+ 'BUS'  + '/classification/' + data + '/' + f for f in selected]\n",
    "                if data == 'GT':\n",
    "                    y_train += tmp\n",
    "                else:\n",
    "                    x_train += tmp\n",
    "    \n",
    "    print('full training data')\n",
    "    WEIGHTS_CL[idx_] = len(x_train)\n",
    "    denom_ += len(x_train)\n",
    "    idx_ += 1\n",
    "    breast_dataset[client+'_train']=Cancer(x_train, y_train, train=True,\\\n",
    "                                          IMAGE_SIZE=IMAGE_SIZE\\\n",
    "                                           , CROP_SIZE=CROP_SIZE)\n",
    "\n",
    "    \n",
    "    breast_dataset[client+'_test'] =Cancer(x_test, y_test, train=False,\\\n",
    "                                          IMAGE_SIZE=IMAGE_SIZE\\\n",
    "                                           , CROP_SIZE=CROP_SIZE)\n",
    "\n",
    "for idx_ in range(len(WEIGHTS_CL)):\n",
    "    WEIGHTS_CL[idx_] = WEIGHTS_CL[idx_]/denom_\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f7a1ffd",
   "metadata": {},
   "source": [
    "# storage file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ee2c91",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_clients, testing_clients = dict(), dict()\n",
    "training_clients_pl = dict()\n",
    "\n",
    "acc_train, acc_test, loss_train, loss_test = dict(), dict(), \\\n",
    "                                            dict(), dict()\n",
    "    \n",
    "nets, optimizers = dict(), dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a69dd98",
   "metadata": {},
   "outputs": [],
   "source": [
    "nets['global'] = UNet(n_channels=N_CHANNELS, n_classes=N_CLASSES, \\\n",
    "                      bilinear=True).to(device)\n",
    "\n",
    "nets['global_2'] = UNet(n_channels=N_CHANNELS, n_classes=N_CLASSES, \\\n",
    "                      bilinear=True).to(device)\n",
    "\n",
    "for client in CLIENTS:\n",
    "    training_clients[client] = DataLoader(breast_dataset[client+'_train'], batch_size=16,\\\n",
    "                 shuffle=True, num_workers=8)\n",
    "    training_clients_pl[client] = DataLoader(breast_dataset[client+'_train'], batch_size=1, \\\n",
    "                shuffle=True, num_workers=8)\n",
    "    ###################################################################################\n",
    "    testing_clients[client] = DataLoader(breast_dataset[client+'_test'], batch_size=1,\\\n",
    "                         shuffle=False, num_workers=1)\n",
    "    \n",
    "    acc_train[client], acc_test[client] = [], []\n",
    "    loss_train[client], loss_test[client] = [], []\n",
    "        \n",
    "    nets[client] = UNet(n_channels=N_CHANNELS, n_classes=N_CLASSES, \\\n",
    "                      bilinear=True).to(device)\n",
    "    nets[client+'_2'] = UNet(n_channels=N_CHANNELS, n_classes=N_CLASSES, \\\n",
    "                      bilinear=True).to(device)\n",
    "    optimizers[client]= optim.Adam(nets[client].parameters(), \\\n",
    "                                   lr=LR,weight_decay=WD)\n",
    "    optimizers[client+'_2']= optim.Adam(nets[client+'_2'].parameters(), \\\n",
    "                                   lr=LR,weight_decay=WD)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "986ecb70",
   "metadata": {},
   "source": [
    "# FedMix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2967e949",
   "metadata": {},
   "source": [
    "### hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2f16c0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "LAMBDA_ = LAM\n",
    "BETA_ = BETA\n",
    "TH = TH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1016e6d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "WEIGHTS_CL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b718ed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "LEN_UDIAT = len(consistent_path['UDIAT'][VERSION]['x_train'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84931843",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_avg_acc, best_epoch_avg = 0, 0\n",
    "index = []\n",
    "\n",
    "for client in CLIENTS:\n",
    "    acc_train[client], acc_test[client] = [], []\n",
    "    loss_train[client], loss_test[client] = [], []\n",
    "    \n",
    "score = [0.,0.,0.]\n",
    "WEIGHTS = [0.,0.,0.]\n",
    "for epoch in range(EPOCHS):\n",
    "    index.append(epoch)\n",
    "    DATA_NUM = [0,0,LEN_UDIAT]\n",
    "    #################### copy fed model ###################\n",
    "    copy_fed(CLIENTS, nets, fed_name='global')\n",
    "    copy_fed(CLIENTS_2, nets, fed_name='global_2')\n",
    "    ######################################################\n",
    "    # generate and refine pseudo labels ##################\n",
    "    ######################################################\n",
    "    for order, client in enumerate(CLIENTS):\n",
    "        bbox, image = False, False\n",
    "        if CLIENTS_SUPERVISION[order] == 'labeled':\n",
    "            continue\n",
    "        elif CLIENTS_SUPERVISION[order] == 'bbox':\n",
    "            bbox = True\n",
    "        elif CLIENTS_SUPERVISION[order] == 'image':\n",
    "            image= True\n",
    "        ##################################################\n",
    "        # save pl ########################################\n",
    "        ##################################################\n",
    "        im_store, pl1_store, pl2_store = [], [], []\n",
    "        \n",
    "        tmp_ = select_pl_singleclass(nets['global'], nets['global_2'], device,\\\n",
    "                      training_clients_pl[client], im_store, pl1_store, \\\n",
    "                      pl2_store, TH=TH, bbox=bbox, image=image)\n",
    "        DATA_NUM[order] = tmp_\n",
    "        if len(im_store) >= 1:\n",
    "            tmp_dataset = cancer_v2(im_store, pl1_store, pl2_store)\n",
    "            training_clients[client] = DataLoader(tmp_dataset, batch_size=16,\\\n",
    "                             shuffle=True, num_workers=8)\n",
    "\n",
    "    #######################################################\n",
    "    #### Conduct training #################################\n",
    "    #######################################################\n",
    "    for order, (client, supervision_t) in enumerate(zip(CLIENTS, CLIENTS_SUPERVISION)):\n",
    "        if supervision_t == 'labeled':\n",
    "            # train network 1 #\n",
    "            train_model(training_clients[client], nets[client], optimizers[client], device, \\\n",
    "                       acc=acc_train[client], loss=loss_train[client], \\\n",
    "                        supervision_type=supervision_t)\n",
    "            \n",
    "            # train network 2 # \n",
    "            train_model(training_clients[client], nets[client+'_2'], optimizers[client+'_2'], device, \\\n",
    "                       acc=None, loss=None, \\\n",
    "                        supervision_type=supervision_t)\n",
    "            \n",
    "        else: # train using pseudo label # \n",
    "            # train network 1 #\n",
    "            train_model(training_clients[client], nets[client], optimizers[client], device, \\\n",
    "                       acc=acc_train[client], loss=loss_train[client], \\\n",
    "                        supervision_type=supervision_t, FedMix_network=1)\n",
    "            \n",
    "            # train network 2 # \n",
    "            train_model(training_clients[client], nets[client+'_2'], optimizers[client+'_2'], device, \\\n",
    "                       acc=None, loss=None, \\\n",
    "                        supervision_type=supervision_t, FedMix_network=2)\n",
    "        \n",
    "        \n",
    "        # save loss for future reweighting # \n",
    "        score[order] = loss_train[client][-1] ** BETA_\n",
    "    ###################################\n",
    "    ####### dynamic weighting #########\n",
    "    ###################################\n",
    "    denominator = sum(score)\n",
    "    score = [s/denominator for s in score]\n",
    "    \n",
    "    denominator = sum(DATA_NUM)\n",
    "    WEIGHTS_CL = [s/denominator for s in DATA_NUM]\n",
    "    \n",
    "    for order, _ in enumerate(WEIGHTS):\n",
    "        WEIGHTS[order] = WEIGHTS_CL[order] + LAMBDA_ * score[order]\n",
    "        \n",
    "    ### normalize #####################\n",
    "    denominator = sum(WEIGHTS)\n",
    "    WEIGHTS = [w/denominator for w in WEIGHTS]\n",
    "\n",
    "    ###################################\n",
    "    ####### aggregation ###############\n",
    "    ###################################\n",
    "    aggr_fed(CLIENTS, WEIGHTS, nets, fed_name='global')\n",
    "    aggr_fed(CLIENTS_2, WEIGHTS, nets, fed_name='global_2')\n",
    "    \n",
    "    ################### test ##############################\n",
    "    avg_acc = 0.0\n",
    "    for client in CLIENTS:\n",
    "        test(epoch, testing_clients[client], nets['global'], device, acc_test[client],\\\n",
    "             loss_test[client])\n",
    "        avg_acc += acc_test[client][-1]\n",
    "        \n",
    "    avg_acc = avg_acc / TOTAL_CLIENTS\n",
    "    ############################################################\n",
    "    ########################################################\n",
    "    if avg_acc > best_avg_acc:\n",
    "        best_avg_acc = avg_acc\n",
    "        best_epoch = epoch\n",
    "    \n",
    "    \n",
    "    \n",
    "    clear_output(wait=True)\n",
    "    print(avg_acc, best_avg_acc)\n",
    "    print(WEIGHTS)\n",
    "    plt.figure(0)\n",
    "    plt.plot(index, acc_train['UDIAT'], colors[0], label='UDIAT train')\n",
    "    plt.plot(index, acc_train['BUS'], colors[1], label='BUS  train')\n",
    "    plt.plot(index, acc_train['BUSIS'], colors[3], label='BUSIS  train')\n",
    "\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "    plt.figure(1)\n",
    "    plt.plot(index, loss_train['UDIAT'], colors[0], label='UDIAT loss train')\n",
    "    plt.plot(index, loss_train['BUS'], colors[1], label='BUS  loss train')\n",
    "    plt.plot(index, loss_train['BUSIS'], colors[3], label='BUSIS  loss train')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "    plot_graphs(2, CLIENTS, index, acc_test, ' acc_test')\n",
    "\n",
    "print(best_avg_acc, best_epoch)\n",
    "for client in CLIENTS:\n",
    "    print(client)\n",
    "    tmp = best_epoch\n",
    "    best_epoch = best_epoch \n",
    "    print(\"shared epoch specific\")\n",
    "    print(acc_test[client][best_epoch])\n",
    "    print(\"max client-specific\")\n",
    "    print(np.max(acc_test[client]))\n",
    "    best_epoch = tmp"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (DeepLearning)",
   "language": "python",
   "name": "python3.6"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}