model_fusion_optimal_transport_multilayer.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformer import TransformerClassifier\n",
    "from dataloader import *\n",
    "from CV.utils import * \n",
    "import torch\n",
    "import ot\n",
    "from copy import deepcopy\n",
    "from tqdm import tqdm\n",
    "import optuna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset initializing start\n",
      "Length of data after first step of preprocessing:  35832\n",
      "Tokenizing the data...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\atace\\OneDrive\\Desktop\\ETH\\9.Semester\\Deep Learning\\project\\Exploring-Model-Fusion-with-Optimal-Transport-on-Transformers\\dataloader.py:80: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  data[\"len\"] = data.iloc[:, 0].apply(lambda x : len(self.tokenize(x)))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of the data :  29544\n",
      "0\n",
      "review       [[CLS], one, of, the, many, silent, comedies, ...\n",
      "sentiment                                                    0\n",
      "len                                                        186\n",
      "Name: 46539, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 23635/23635 [00:01<00:00, 14257.29it/s]\n",
      "100%|██████████| 2954/2954 [00:00<00:00, 17048.35it/s]\n",
      "100%|██████████| 2955/2955 [00:00<00:00, 16056.68it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset initializing done\n",
      "Vocabulary Size :  23050\n"
     ]
    }
   ],
   "source": [
    "#init\n",
    "tokenizer = Tokenizer()\n",
    "loader = DataLoader(tokenize = tokenizer.tokenize)\n",
    "\n",
    "# import data (combine train/test as we split afterwards anyways)\n",
    "data = pd.read_csv(\"./Data/IMDB Dataset.csv\", encoding='ISO-8859-1')\n",
    "# convert string label to binary (int) label (spam:1, non-spam:0)\n",
    "data[\"sentiment\"] = data['sentiment'].apply(lambda x : int(x == \"positive\"))\n",
    "\n",
    "# train, test, val split\n",
    "train, valid, test = loader.make_dataset(data)\n",
    "vocab = loader.get_vocab(train.iloc[:, 0])\n",
    "train_iter, valid_iter, test_iter = loader.make_iter(train, valid, test,\n",
    "                                                     batch_size=512,\n",
    "                                                     device=device)\n",
    "\n",
    "# NLP stuff\n",
    "pad_idx = vocab['__PAD__']\n",
    "voc_size = len(vocab)\n",
    "print(\"Vocabulary Size : \", voc_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the embedding matrix\n",
    "embedding = torch.load(\"Models/embedding_16_trained_multilayer.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[-0.3602, -2.3004,  1.2823,  ...,  1.3947, -0.0689, -0.4801],\n",
       "        [-0.1015,  0.2015, -1.9556,  ..., -0.1140,  0.2189,  0.8465],\n",
       "        [ 1.2555,  0.2540,  1.1998,  ..., -0.5179, -0.0536,  1.1762],\n",
       "        ...,\n",
       "        [-0.2308,  0.4539,  0.3592,  ..., -0.7208, -0.6409, -0.3728],\n",
       "        [ 1.0316,  1.4529,  0.5553,  ..., -0.8593,  0.1438, -1.7995],\n",
       "        [-0.0379, -1.8211, -0.7899,  ...,  0.9734,  0.4521,  0.8370]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding.weight"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Model Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TransformerClassifier(\n",
       "  (encoder): Encoder(\n",
       "    (emb): TransformerEmbedding(\n",
       "      (tok_emb): TokenEmbedding(\n",
       "        (embedding): Embedding(23050, 16)\n",
       "      )\n",
       "      (pos_emb): PositionalEncoding()\n",
       "      (drop_out): Dropout(p=0.5, inplace=False)\n",
       "    )\n",
       "    (layers): ModuleList(\n",
       "      (0): EncoderLayer(\n",
       "        (attention): MultiHeadAttention(\n",
       "          (attention): ScaleDotProductAttention(\n",
       "            (softmax): Softmax(dim=-1)\n",
       "          )\n",
       "          (w_q): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_k): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_v): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_concat): Linear(in_features=16, out_features=16, bias=True)\n",
       "        )\n",
       "        (norm1): LayerNorm()\n",
       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
       "        (ffn): PositionwiseFeedForward(\n",
       "          (linear1): Linear(in_features=16, out_features=32, bias=True)\n",
       "          (linear2): Linear(in_features=32, out_features=16, bias=True)\n",
       "          (relu): ReLU()\n",
       "          (dropout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "        (norm2): LayerNorm()\n",
       "        (dropout2): Dropout(p=0.5, inplace=False)\n",
       "      )\n",
       "      (1): EncoderLayer(\n",
       "        (attention): MultiHeadAttention(\n",
       "          (attention): ScaleDotProductAttention(\n",
       "            (softmax): Softmax(dim=-1)\n",
       "          )\n",
       "          (w_q): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_k): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_v): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_concat): Linear(in_features=16, out_features=16, bias=True)\n",
       "        )\n",
       "        (norm1): LayerNorm()\n",
       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
       "        (ffn): PositionwiseFeedForward(\n",
       "          (linear1): Linear(in_features=16, out_features=32, bias=True)\n",
       "          (linear2): Linear(in_features=32, out_features=16, bias=True)\n",
       "          (relu): ReLU()\n",
       "          (dropout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "        (norm2): LayerNorm()\n",
       "        (dropout2): Dropout(p=0.5, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (linear): Linear(in_features=4096, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "modelA = TransformerClassifier(src_pad_idx = pad_idx,\n",
    "                              embedding=embedding,\n",
    "                              enc_voc_size = voc_size,\n",
    "                              max_len = 256,\n",
    "                              d_model = 16,\n",
    "                              ffn_hidden = 32,\n",
    "                              n_head = 1,\n",
    "                              n_layers = 2,\n",
    "                              drop_prob = 0.5,\n",
    "                              device = device)\n",
    "\n",
    "modelA.load_state_dict(torch.load(\"Models\\modelA_IMDB_256_multilayer\"))\n",
    "modelA.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TransformerClassifier(\n",
       "  (encoder): Encoder(\n",
       "    (emb): TransformerEmbedding(\n",
       "      (tok_emb): TokenEmbedding(\n",
       "        (embedding): Embedding(23050, 16)\n",
       "      )\n",
       "      (pos_emb): PositionalEncoding()\n",
       "      (drop_out): Dropout(p=0.5, inplace=False)\n",
       "    )\n",
       "    (layers): ModuleList(\n",
       "      (0): EncoderLayer(\n",
       "        (attention): MultiHeadAttention(\n",
       "          (attention): ScaleDotProductAttention(\n",
       "            (softmax): Softmax(dim=-1)\n",
       "          )\n",
       "          (w_q): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_k): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_v): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_concat): Linear(in_features=16, out_features=16, bias=True)\n",
       "        )\n",
       "        (norm1): LayerNorm()\n",
       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
       "        (ffn): PositionwiseFeedForward(\n",
       "          (linear1): Linear(in_features=16, out_features=32, bias=True)\n",
       "          (linear2): Linear(in_features=32, out_features=16, bias=True)\n",
       "          (relu): ReLU()\n",
       "          (dropout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "        (norm2): LayerNorm()\n",
       "        (dropout2): Dropout(p=0.5, inplace=False)\n",
       "      )\n",
       "      (1): EncoderLayer(\n",
       "        (attention): MultiHeadAttention(\n",
       "          (attention): ScaleDotProductAttention(\n",
       "            (softmax): Softmax(dim=-1)\n",
       "          )\n",
       "          (w_q): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_k): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_v): Linear(in_features=16, out_features=16, bias=True)\n",
       "          (w_concat): Linear(in_features=16, out_features=16, bias=True)\n",
       "        )\n",
       "        (norm1): LayerNorm()\n",
       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
       "        (ffn): PositionwiseFeedForward(\n",
       "          (linear1): Linear(in_features=16, out_features=32, bias=True)\n",
       "          (linear2): Linear(in_features=32, out_features=16, bias=True)\n",
       "          (relu): ReLU()\n",
       "          (dropout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "        (norm2): LayerNorm()\n",
       "        (dropout2): Dropout(p=0.5, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (linear): Linear(in_features=4096, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "modelB = TransformerClassifier(src_pad_idx = pad_idx,\n",
    "                              embedding=embedding,\n",
    "                              enc_voc_size = voc_size,\n",
    "                              max_len = 256,\n",
    "                              d_model = 16,\n",
    "                              ffn_hidden = 32,\n",
    "                              n_head = 1,\n",
    "                              n_layers = 2,\n",
    "                              drop_prob = 0.5,\n",
    "                              device = device)\n",
    "\n",
    "modelB.load_state_dict(torch.load(\"Models\\modelB_IMDB_256_multilayer\"))\n",
    "modelB.eval()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OT Functions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getSupport(model, trainloader, l, alignment = \"acts\", numOfBatches= 10):\n",
    "    '''\n",
    "    Get the support matrices using Activation-based (\"acts\") or Weight-based (\"wts\") alignment \n",
    "    '''\n",
    "    if alignment == \"acts\":\n",
    "        activation = None\n",
    "        for i, data in enumerate(trainloader, 0):\n",
    "            if i >= numOfBatches:\n",
    "                break\n",
    "            \n",
    "            inputs, targets = data\n",
    "            outputs = model(inputs)\n",
    "\n",
    "            if activation is None:\n",
    "                activation = model.actMatrix[l]\n",
    "            else:\n",
    "                activation = torch.cat((activation, model.actMatrix[l]))\n",
    "\n",
    "        return activation\n",
    "    elif alignment == \"wts\":\n",
    "        return model.state_dict()[l]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fusion(nameA, nameB, weightA, weightB, transport_matrix, beta):\n",
    "    support_y = getSupport(modelB, train_iter, nameB, alignment=\"wts\")\n",
    "    # Get the weights at layer \"idx\" from the first model\n",
    "    W_A = weightA\n",
    "    W_B = weightB\n",
    "    # Align the weights from the first model\n",
    "    aligned_W = torch.matmul(W_A, torch.matmul(transport_matrix, torch.diag(1 / beta)))\n",
    "    # Get the X-Support\n",
    "    n = W_A.shape[0]\n",
    "    alpha = torch.ones(n) * (1/n)\n",
    "    support_x = getSupport(modelA, train_iter, nameA, alignment=\"wts\")\n",
    "    # Calculate the euclidean distance between the supports\n",
    "    distance = ot.dist(support_x, support_y)\n",
    "    # Calculate beta\n",
    "    m = W_B.shape[0]\n",
    "    beta = torch.ones(m) * (1/m)\n",
    "    # Calculate the transport matrix using optimal transport\n",
    "    transport_matrix = torch.from_numpy(ot.emd(alpha.numpy(), beta.numpy(), distance.detach().numpy())).float().reshape((n, m))\n",
    "    # Align model neurons\n",
    "    aligned_model = torch.matmul(torch.diag(1 / beta), torch.matmul(transport_matrix.T, aligned_W))\n",
    "    # Get the weights at layer \"idx\" from the second model\n",
    "    fused = (aligned_model + W_B) / 2\n",
    "    return  fused, transport_matrix, beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fusion_multilayer(nameA, nameB, weightA, weightB, transport_matrix, beta, layers):\n",
    "    layer = int(nameA.partition(\"layers.\")[2][0])\n",
    "    \n",
    "    W_A = torch.empty((weightA.shape[0] * layers, weightA.shape[1]))\n",
    "    W_B = torch.empty((weightB.shape[0] * layers, weightB.shape[1]))\n",
    "\n",
    "    for l in range(layers):\n",
    "        W_A[l * weightB.shape[0] : l * weightB.shape[0] + weightB.shape[0], :] = modelA.state_dict()[nameA.replace(f\"layers.{layer}\", f\"layers.{l}\")]\n",
    "        W_B[l * weightB.shape[0] : l * weightB.shape[0] + weightB.shape[0], :] = modelB.state_dict()[nameA.replace(f\"layers.{layer}\", f\"layers.{l}\")]\n",
    "\n",
    "    support_y = W_B\n",
    "    support_x = W_A\n",
    "\n",
    "    # Initialize the fused model and transport matrix\n",
    "    fused = torch.empty(W_B.shape)\n",
    "    transport_matrix_new = torch.zeros((W_A.shape[0], W_B.shape[0]))\n",
    "    # Align the weights from the first model\n",
    "    aligned_W = torch.matmul(W_A, torch.matmul(transport_matrix, torch.diag(1 / beta)))\n",
    "    # Get the X-Support\n",
    "    n = W_A.shape[0]\n",
    "    alpha = torch.ones(n) * (1/n)\n",
    "    # Calculate the euclidean distance between the supports\n",
    "    distance = ot.dist(support_x, support_y)\n",
    "    # Calculate beta\n",
    "    m = W_B.shape[0]\n",
    "    beta = torch.ones(m) * (1/m)\n",
    "    # Calculate the transport matrix using optimal transport\n",
    "    transport_matrix = torch.from_numpy(ot.emd(alpha.numpy(), beta.numpy(), distance.detach().numpy())).float().reshape((n, m))\n",
    "    # Align model neurons\n",
    "    aligned_model = torch.matmul(torch.diag(1 / beta), torch.matmul(transport_matrix.T, aligned_W))\n",
    "    # Get the weights at layer \"idx\" from the second model\n",
    "    fused = (aligned_model + W_B) / 2\n",
    "    dim = layer * weightB.shape[0]\n",
    "    dim_to = layer * weightB.shape[0] + weightB.shape[0]\n",
    "    return  fused[dim:dim_to, :], transport_matrix_new[dim:dim_to, dim:dim_to], beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fusion_multihead_multilayer(nameA, nameB, weightA, weightB, transport_matrix, beta, heads, layers):\n",
    "    layer = int(nameA.partition(\"layers.\")[2][0])\n",
    "    \n",
    "    W_A = torch.empty((weightA.shape[0] * layers, weightA.shape[1]))\n",
    "    W_B = torch.empty((weightB.shape[0] * layers, weightB.shape[1]))\n",
    "    \n",
    "    stride = weightB.shape[0] // heads\n",
    "    for l in range(layers):\n",
    "        W_A[l * weightB.shape[0] : l * weightB.shape[0] + weightB.shape[0], :] = modelA.state_dict()[nameA.replace(f\"layers.{layer}\", f\"layers.{l}\")]\n",
    "        W_B[l * weightB.shape[0] : l * weightB.shape[0] + weightB.shape[0], :] = modelB.state_dict()[nameA.replace(f\"layers.{layer}\", f\"layers.{l}\")]\n",
    "\n",
    "    support_y = W_B\n",
    "    support_x = W_A\n",
    "\n",
    "    # Initialize the fused model and transport matrix\n",
    "    fused = torch.empty(W_B.shape)\n",
    "    transport_matrix_new = torch.zeros((W_A.shape[0], W_B.shape[0]))\n",
    "    stride = W_B.shape[0] // (heads * layers)\n",
    "    for i in range(0, W_B.shape[0], stride):\n",
    "        # Align the weights from the first model\n",
    "        aligned_W = torch.matmul(W_A[i:i+stride, :], torch.matmul(transport_matrix, torch.diag(1 / beta)))\n",
    "        # Get the X-Support\n",
    "        n = W_A.shape[0] // (heads * layers)\n",
    "        alpha = torch.ones(n) * (1/n)\n",
    "        # Calculate the euclidean distance between the supports\n",
    "        distance = ot.dist(support_x[i:i+stride, :], support_y[i:i+stride, :])\n",
    "        # Calculate beta\n",
    "        m = W_B.shape[0] // (heads * layers)\n",
    "        beta_new = torch.ones(m) * (1/m)\n",
    "        # Calculate the transport matrix using optimal transport\n",
    "        transport_matrix_new[i:i+stride, i:i+stride] = torch.from_numpy(ot.emd(alpha.numpy(), beta_new.numpy(), distance.detach().numpy())).float().reshape((n, m))\n",
    "        # Align model neurons\n",
    "        aligned_model = torch.matmul(torch.diag(1 / beta_new), torch.matmul(transport_matrix_new[i:i+stride, i:i+stride].T, aligned_W))\n",
    "        # Get the weights at layer \"idx\" from the second model\n",
    "        fused[i:i+stride, :] = (aligned_model + W_B[i:i+stride, :]) / 2 \n",
    "    dim = layer * weightB.shape[0]\n",
    "    dim_to = layer * weightB.shape[0] + weightB.shape[0]\n",
    "    return  fused[dim:dim_to, :], transport_matrix_new[dim:dim_to, dim:dim_to], beta_new"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fusion via Optimal Transport"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = {'loss': {'A_test_set': [], 'B_test_set': [], 'OT_pre': [], 'OT_calibr': [], 'OT_last_layer': [], 'OT_test_set' : []},\n",
    "          'accuracy': {'A_test_set': [], 'B_test_set': [], 'OT_pre': [], 'OT_calibr': [], 'OT_last_layer': [], 'OT_test_set' : []},\n",
    "          'f1': {'A_test_set': [], 'B_test_set': [], 'OT_pre': [], 'OT_calibr': [], 'OT_last_layer': [], 'OT_test_set' : []}}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "fusedModel = TransformerClassifier(src_pad_idx = pad_idx,\n",
    "                              embedding = embedding,\n",
    "                              enc_voc_size = voc_size,\n",
    "                              max_len = 256,\n",
    "                              d_model = 16,\n",
    "                              ffn_hidden = 32,\n",
    "                              n_head = 1,\n",
    "                              n_layers = 2,\n",
    "                              drop_prob = 0.2,\n",
    "                              device = device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def obj(trial):\n",
    "    a = trial.suggest_float('a', 0, 1)\n",
    "    \n",
    "    # Create the fused weights matrix\n",
    "    W_fusion = dict.fromkeys(list(modelA.state_dict().keys()))\n",
    "    # Initialize the algorithm\n",
    "    m = list(modelB.state_dict().items())[1][1].shape[1]\n",
    "    beta = torch.ones(m) * (1/m)\n",
    "    transport_matrix = torch.matmul(torch.diag(beta), torch.eye(m))\n",
    "\n",
    "    # Fusion via Optimal Transport\n",
    "    for (nameA, weightA), (nameB, weightB) in zip(modelA.named_parameters(), modelB.named_parameters()):\n",
    "        if nameA == \"encoder.emb.tok_emb.embedding.weight\":\n",
    "            W_fusion[nameA] = weightA\n",
    "        else:   \n",
    "            if \"weight\" in nameA:\n",
    "                if \"encoder\" in nameA:\n",
    "                    if \"concat\" not in nameA and \"linear\" not in nameA: \n",
    "                        transport_matrix_triplet = torch.empty((weightA.shape[0], weightB.shape[0]))\n",
    "                        W_fusion[nameA], transport_matrix_triplet, _ = fusion_multilayer(nameA, nameB, weightA, weightB, transport_matrix, beta, 2)\n",
    "                    else:\n",
    "                        W_fusion[nameA], transport_matrix, beta = fusion(nameA, nameB, weightA, weightB, transport_matrix, beta)\n",
    "\n",
    "                else:\n",
    "                    W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "            elif \"bias\" in nameA:\n",
    "                if \"encoder\" in nameA: \n",
    "                    if \"concat\" not in nameA and \"linear\" not in nameA: \n",
    "                        m = weightB.shape[0]\n",
    "                        beta_bias = torch.ones(m) * (1/m)\n",
    "                        W_A_bias = weightA.reshape(m, 1)\n",
    "                        aligned_bias = torch.matmul(torch.diag(1 / beta_bias), torch.matmul(transport_matrix_triplet.T, W_A_bias))\n",
    "                        aligned_bias = aligned_bias.reshape(m)\n",
    "                        W_fusion[nameA] = (aligned_bias + weightB) / 2\n",
    "                    else:\n",
    "                        m = weightB.shape[0]\n",
    "                        beta_bias = torch.ones(m) * (1/m)\n",
    "                        W_A_bias = weightA.reshape(m, 1)\n",
    "                        aligned_bias = torch.matmul(torch.diag(1 / beta_bias), torch.matmul(transport_matrix.T, W_A_bias))\n",
    "                        aligned_bias = aligned_bias.reshape(m)\n",
    "                        W_fusion[nameA] = (aligned_bias + weightB) / 2\n",
    "                else:\n",
    "                    W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "            else:\n",
    "                W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "\n",
    "    # Assign the weights\n",
    "    with torch.no_grad():\n",
    "        for name, param in fusedModel.named_parameters():\n",
    "            param.data = torch.nn.Parameter(W_fusion[name])\n",
    "\n",
    "    # Validate the fused model\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    val_loss, val_acc, val_f1 = validation(fusedModel, valid_iter, criterion, device)\n",
    "    return val_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[I 2023-01-14 11:53:57,171]\u001b[0m A new study created in memory with name: no-name-ce8cac9a-d997-4937-b0fb-8e09021aa32f\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:54:04,032]\u001b[0m Trial 0 finished with value: 0.8183954656124115 and parameters: {'a': 0.723256023267953}. Best is trial 0 with value: 0.8183954656124115.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:54:12,721]\u001b[0m Trial 1 finished with value: 0.4955558280150096 and parameters: {'a': 0.03416997858031012}. Best is trial 1 with value: 0.4955558280150096.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:54:22,662]\u001b[0m Trial 2 finished with value: 0.5235801140467325 and parameters: {'a': 0.20763937634628304}. Best is trial 1 with value: 0.4955558280150096.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:54:31,604]\u001b[0m Trial 3 finished with value: 0.6849202513694763 and parameters: {'a': 0.550675425060867}. Best is trial 1 with value: 0.4955558280150096.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:54:40,478]\u001b[0m Trial 4 finished with value: 0.7962923844655355 and parameters: {'a': 0.6934538825629025}. Best is trial 1 with value: 0.4955558280150096.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:54:49,431]\u001b[0m Trial 5 finished with value: 0.526525283853213 and parameters: {'a': 0.21499434139200568}. Best is trial 1 with value: 0.4955558280150096.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:54:57,811]\u001b[0m Trial 6 finished with value: 1.025560329357783 and parameters: {'a': 0.9407074119044411}. Best is trial 1 with value: 0.4955558280150096.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:55:07,728]\u001b[0m Trial 7 finished with value: 0.5039957265059153 and parameters: {'a': 0.11859858420564673}. Best is trial 1 with value: 0.4955558280150096.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:55:16,918]\u001b[0m Trial 8 finished with value: 0.525639737645785 and parameters: {'a': 0.20984408813318756}. Best is trial 1 with value: 0.4955558280150096.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:55:25,360]\u001b[0m Trial 9 finished with value: 0.54738716284434 and parameters: {'a': 0.2771681536492231}. Best is trial 1 with value: 0.4955558280150096.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:55:33,334]\u001b[0m Trial 10 finished with value: 0.49431393543879193 and parameters: {'a': 0.01902530211177894}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:55:41,327]\u001b[0m Trial 11 finished with value: 0.4947035213311513 and parameters: {'a': 0.032069761201378784}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:55:50,250]\u001b[0m Trial 12 finished with value: 0.5912463466326395 and parameters: {'a': 0.3832792382869353}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:55:58,627]\u001b[0m Trial 13 finished with value: 0.49542353053887683 and parameters: {'a': 0.04024162684909677}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:56:07,257]\u001b[0m Trial 14 finished with value: 0.5751198828220367 and parameters: {'a': 0.3498308040509291}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:56:17,914]\u001b[0m Trial 15 finished with value: 0.49565283954143524 and parameters: {'a': 0.023241467596566402}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:56:26,438]\u001b[0m Trial 16 finished with value: 0.650277336438497 and parameters: {'a': 0.49325200203457176}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:56:34,807]\u001b[0m Trial 17 finished with value: 1.0751781463623047 and parameters: {'a': 0.9887828739944314}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:56:43,806]\u001b[0m Trial 18 finished with value: 0.5112801790237427 and parameters: {'a': 0.1575654267576905}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:56:51,785]\u001b[0m Trial 19 finished with value: 0.5668520331382751 and parameters: {'a': 0.3266451356968844}. Best is trial 10 with value: 0.49431393543879193.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "study = optuna.create_study()\n",
    "study.optimize(obj, n_trials=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Loss:      0.49432618419329327\n",
      "Validation Accuracy:      0.7745908682318104\n",
      "Validation F+:      0.7742044925689697\n"
     ]
    }
   ],
   "source": [
    "a = study.best_params[\"a\"]\n",
    "\n",
    "# Create the fused weights matrix\n",
    "W_fusion = dict.fromkeys(list(modelA.state_dict().keys()))\n",
    "# Initialize the algorithm\n",
    "m = list(modelB.state_dict().items())[1][1].shape[1]\n",
    "beta = torch.ones(m) * (1/m)\n",
    "transport_matrix = torch.matmul(torch.diag(beta), torch.eye(m))\n",
    "# Fusion via Optimal Transport\n",
    "for (nameA, weightA), (nameB, weightB) in zip(modelA.named_parameters(), modelB.named_parameters()):\n",
    "    if nameA == \"encoder.emb.tok_emb.embedding.weight\":\n",
    "        W_fusion[nameA] = weightA\n",
    "    else:   \n",
    "        if \"weight\" in nameA:\n",
    "            if \"encoder\" in nameA:\n",
    "                if \"concat\" not in nameA and \"linear\" not in nameA: \n",
    "                    transport_matrix_triplet = torch.empty((weightA.shape[0], weightB.shape[0]))\n",
    "                    W_fusion[nameA], transport_matrix_triplet, _ = fusion_multilayer(nameA, nameB, weightA, weightB, transport_matrix, beta, 2)\n",
    "                else:\n",
    "                    W_fusion[nameA], transport_matrix, beta = fusion(nameA, nameB, weightA, weightB, transport_matrix, beta)\n",
    "            else:\n",
    "                W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "        elif \"bias\" in nameA:\n",
    "            if \"encoder\" in nameA: \n",
    "                if \"concat\" not in nameA and \"linear\" not in nameA: \n",
    "                    m = weightB.shape[0]\n",
    "                    beta_bias = torch.ones(m) * (1/m)\n",
    "                    W_A_bias = weightA.reshape(m, 1)\n",
    "                    aligned_bias = torch.matmul(torch.diag(1 / beta_bias), torch.matmul(transport_matrix_triplet.T, W_A_bias))\n",
    "                    aligned_bias = aligned_bias.reshape(m)\n",
    "                    W_fusion[nameA] = (aligned_bias + weightB) / 2\n",
    "                else:\n",
    "                    m = weightB.shape[0]\n",
    "                    beta_bias = torch.ones(m) * (1/m)\n",
    "                    W_A_bias = weightA.reshape(m, 1)\n",
    "                    aligned_bias = torch.matmul(torch.diag(1 / beta_bias), torch.matmul(transport_matrix.T, W_A_bias))\n",
    "                    aligned_bias = aligned_bias.reshape(m)\n",
    "                    W_fusion[nameA] = (aligned_bias + weightB) / 2\n",
    "            else:\n",
    "                W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "        else:\n",
    "            W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "# Assign the weights\n",
    "with torch.no_grad():\n",
    "    for name, param in fusedModel.named_parameters():\n",
    "        param.data = torch.nn.Parameter(W_fusion[name])\n",
    "# Validate the fused model\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "val_loss, val_acc, val_f1 = validation(fusedModel, valid_iter, criterion, device)\n",
    "scores['loss']['OT_pre'], scores['accuracy']['OT_pre'], scores['f1']['OT_pre'] = val_loss, val_acc, val_f1\n",
    "print(\"Validation Loss:     \", val_loss)\n",
    "print(\"Validation Accuracy:     \", val_acc.item())\n",
    "print(\"Validation F+:     \", val_f1.item())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "fusedModel = TransformerClassifier(src_pad_idx = pad_idx,\n",
    "                              embedding = embedding,\n",
    "                              enc_voc_size = voc_size,\n",
    "                              max_len = 256,\n",
    "                              d_model = 16,\n",
    "                              ffn_hidden = 32,\n",
    "                              n_head = 1,\n",
    "                              n_layers = 2,\n",
    "                              drop_prob = 0.2,\n",
    "                              device = device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def obj(trial):\n",
    "    a = trial.suggest_float('a', 0, 1)\n",
    "    \n",
    "    # Create the fused weights matrix\n",
    "    W_fusion = dict.fromkeys(list(modelA.state_dict().keys()))\n",
    "    # Initialize the algorithm\n",
    "    m = list(modelB.state_dict().items())[1][1].shape[1]\n",
    "    beta = torch.ones(m) * (1/m)\n",
    "    transport_matrix = torch.matmul(torch.diag(beta), torch.eye(m))\n",
    "\n",
    "    # Fusion via Optimal Transport\n",
    "    for (nameA, weightA), (nameB, weightB) in zip(modelA.named_parameters(), modelB.named_parameters()):\n",
    "        if nameA == \"encoder.emb.tok_emb.embedding.weight\":\n",
    "            W_fusion[nameA] = weightA\n",
    "        else:   \n",
    "            if \"weight\" in nameA:\n",
    "                if \"encoder\" in nameA:\n",
    "                    if \"concat\" not in nameA and \"linear\" not in nameA: \n",
    "                        transport_matrix_triplet = torch.empty((weightA.shape[0], weightB.shape[0]))\n",
    "                        W_fusion[nameA], transport_matrix_triplet, _ = fusion_multihead_multilayer(nameA, nameB, weightA, weightB, transport_matrix, beta, fusedModel.n_head, 2)\n",
    "                    else:\n",
    "                        W_fusion[nameA], transport_matrix, beta = fusion(nameA, nameB, weightA, weightB, transport_matrix, beta)\n",
    "\n",
    "                else:\n",
    "                    W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "            elif \"bias\" in nameA:\n",
    "                if \"encoder\" in nameA: \n",
    "                    if \"concat\" not in nameA and \"linear\" not in nameA: \n",
    "                        m = weightB.shape[0]\n",
    "                        beta_bias = torch.ones(m) * (1/m)\n",
    "                        W_A_bias = weightA.reshape(m, 1)\n",
    "                        aligned_bias = torch.matmul(torch.diag(1 / beta_bias), torch.matmul(transport_matrix_triplet.T, W_A_bias))\n",
    "                        aligned_bias = aligned_bias.reshape(m)\n",
    "                        W_fusion[nameA] = (aligned_bias + weightB) / 2\n",
    "                    else:\n",
    "                        m = weightB.shape[0]\n",
    "                        beta_bias = torch.ones(m) * (1/m)\n",
    "                        W_A_bias = weightA.reshape(m, 1)\n",
    "                        aligned_bias = torch.matmul(torch.diag(1 / beta_bias), torch.matmul(transport_matrix.T, W_A_bias))\n",
    "                        aligned_bias = aligned_bias.reshape(m)\n",
    "                        W_fusion[nameA] = (aligned_bias + weightB) / 2\n",
    "                else:\n",
    "                    W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "            else:\n",
    "                W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "\n",
    "    # Assign the weights\n",
    "    with torch.no_grad():\n",
    "        for name, param in fusedModel.named_parameters():\n",
    "            param.data = torch.nn.Parameter(W_fusion[name])\n",
    "\n",
    "    # Validate the fused model\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    val_loss, val_acc, val_f1 = validation(fusedModel, valid_iter, criterion, device)\n",
    "    return val_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[I 2023-01-14 11:44:22,037]\u001b[0m A new study created in memory with name: no-name-6b09a687-c6f9-4b9b-83ef-48f0d66a1684\u001b[0m\n",
      "c:\\Users\\atace\\OneDrive\\Desktop\\ETH\\9.Semester\\Deep Learning\\project\\Exploring-Model-Fusion-with-Optimal-Transport-on-Transformers\\CV\\utils.py:249: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  src, trg = torch.tensor(src).to(device), torch.tensor(trg).to(device)  # put to cpu/gpu\n",
      "\u001b[32m[I 2023-01-14 11:44:31,048]\u001b[0m Trial 0 finished with value: 0.7893851002057394 and parameters: {'a': 0.4174552341685299}. Best is trial 0 with value: 0.7893851002057394.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:44:40,141]\u001b[0m Trial 1 finished with value: 0.6804643074671427 and parameters: {'a': 0.2509478447327288}. Best is trial 1 with value: 0.6804643074671427.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:44:48,585]\u001b[0m Trial 2 finished with value: 0.9434540271759033 and parameters: {'a': 0.6104170066096141}. Best is trial 1 with value: 0.6804643074671427.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:44:57,550]\u001b[0m Trial 3 finished with value: 1.1019093990325928 and parameters: {'a': 0.7998780503285836}. Best is trial 1 with value: 0.6804643074671427.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:45:06,376]\u001b[0m Trial 4 finished with value: 1.2766624490420024 and parameters: {'a': 0.9751243450152994}. Best is trial 1 with value: 0.6804643074671427.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:45:14,955]\u001b[0m Trial 5 finished with value: 0.9723222057024637 and parameters: {'a': 0.6576113331621465}. Best is trial 1 with value: 0.6804643074671427.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:45:23,605]\u001b[0m Trial 6 finished with value: 0.6688589553038279 and parameters: {'a': 0.2253368827222828}. Best is trial 6 with value: 0.6688589553038279.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:45:33,172]\u001b[0m Trial 7 finished with value: 1.1571231385072072 and parameters: {'a': 0.8566487779933387}. Best is trial 6 with value: 0.6688589553038279.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:45:43,943]\u001b[0m Trial 8 finished with value: 0.6639886001745859 and parameters: {'a': 0.21815695331284946}. Best is trial 8 with value: 0.6639886001745859.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:45:54,548]\u001b[0m Trial 9 finished with value: 0.9450706839561462 and parameters: {'a': 0.6180796462046848}. Best is trial 8 with value: 0.6639886001745859.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:46:05,729]\u001b[0m Trial 10 finished with value: 0.5933017333348592 and parameters: {'a': 0.07203535182163812}. Best is trial 10 with value: 0.5933017333348592.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:46:16,070]\u001b[0m Trial 11 finished with value: 0.5871751209100088 and parameters: {'a': 0.04802379139273577}. Best is trial 11 with value: 0.5871751209100088.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:46:26,462]\u001b[0m Trial 12 finished with value: 0.5711204210917155 and parameters: {'a': 0.008379738135282089}. Best is trial 12 with value: 0.5711204210917155.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:46:36,094]\u001b[0m Trial 13 finished with value: 0.5839745104312897 and parameters: {'a': 0.03957160082937276}. Best is trial 12 with value: 0.5711204210917155.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:46:45,724]\u001b[0m Trial 14 finished with value: 0.7513589262962341 and parameters: {'a': 0.37420784219244596}. Best is trial 12 with value: 0.5711204210917155.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:46:54,594]\u001b[0m Trial 15 finished with value: 0.5738290051619211 and parameters: {'a': 0.01993286841932975}. Best is trial 12 with value: 0.5711204210917155.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:47:03,927]\u001b[0m Trial 16 finished with value: 0.5720842182636261 and parameters: {'a': 0.007705009712307542}. Best is trial 12 with value: 0.5711204210917155.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:47:12,600]\u001b[0m Trial 17 finished with value: 0.6380487283070883 and parameters: {'a': 0.17112041373900413}. Best is trial 12 with value: 0.5711204210917155.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:47:20,132]\u001b[0m Trial 18 finished with value: 0.7270094652970632 and parameters: {'a': 0.32822151993796794}. Best is trial 12 with value: 0.5711204210917155.\u001b[0m\n",
      "\u001b[32m[I 2023-01-14 11:47:27,547]\u001b[0m Trial 19 finished with value: 0.8300996323426565 and parameters: {'a': 0.4810611827370507}. Best is trial 12 with value: 0.5711204210917155.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "study = optuna.create_study()\n",
    "study.optimize(obj, n_trials=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Loss:      0.5696245928605398\n",
      "Validation Accuracy:      0.7417479642554992\n",
      "Validation Accuracy:      0.7413676381111145\n"
     ]
    }
   ],
   "source": [
    "a = study.best_params[\"a\"]\n",
    "\n",
    "# Create the fused weights matrix\n",
    "W_fusion = dict.fromkeys(list(modelA.state_dict().keys()))\n",
    "# Initialize the algorithm\n",
    "m = list(modelB.state_dict().items())[1][1].shape[1]\n",
    "beta = torch.ones(m) * (1/m)\n",
    "transport_matrix = torch.matmul(torch.diag(beta), torch.eye(m))\n",
    "# Fusion via Optimal Transport\n",
    "for (nameA, weightA), (nameB, weightB) in zip(modelA.named_parameters(), modelB.named_parameters()):\n",
    "    if nameA == \"encoder.emb.tok_emb.embedding.weight\":\n",
    "        W_fusion[nameA] = weightA\n",
    "    else:   \n",
    "        if \"weight\" in nameA:\n",
    "            if \"encoder\" in nameA:\n",
    "                if \"concat\" not in nameA and \"linear\" not in nameA: \n",
    "                    transport_matrix_triplet = torch.empty((weightA.shape[0], weightB.shape[0]))\n",
    "                    W_fusion[nameA], transport_matrix_triplet, _ = fusion_multihead_multilayer(nameA, nameB, weightA, weightB, transport_matrix, beta, fusedModel.n_head, 2)\n",
    "                else:\n",
    "                    W_fusion[nameA], transport_matrix, beta = fusion(nameA, nameB, weightA, weightB, transport_matrix, beta)\n",
    "            else:\n",
    "                W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "        elif \"bias\" in nameA:\n",
    "            if \"encoder\" in nameA: \n",
    "                if \"concat\" not in nameA and \"linear\" not in nameA: \n",
    "                    m = weightB.shape[0]\n",
    "                    beta_bias = torch.ones(m) * (1/m)\n",
    "                    W_A_bias = weightA.reshape(m, 1)\n",
    "                    aligned_bias = torch.matmul(torch.diag(1 / beta_bias), torch.matmul(transport_matrix_triplet.T, W_A_bias))\n",
    "                    aligned_bias = aligned_bias.reshape(m)\n",
    "                    W_fusion[nameA] = (aligned_bias + weightB) / 2\n",
    "                else:\n",
    "                    m = weightB.shape[0]\n",
    "                    beta_bias = torch.ones(m) * (1/m)\n",
    "                    W_A_bias = weightA.reshape(m, 1)\n",
    "                    aligned_bias = torch.matmul(torch.diag(1 / beta_bias), torch.matmul(transport_matrix.T, W_A_bias))\n",
    "                    aligned_bias = aligned_bias.reshape(m)\n",
    "                    W_fusion[nameA] = (aligned_bias + weightB) / 2\n",
    "            else:\n",
    "                W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "        else:\n",
    "            W_fusion[nameA] = a * weightA + (1-a) * weightB\n",
    "# Assign the weights\n",
    "with torch.no_grad():\n",
    "    for name, param in fusedModel.named_parameters():\n",
    "        param.data = torch.nn.Parameter(W_fusion[name])\n",
    "# Validate the fused model\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "val_loss, val_acc, val_f1 = validation(fusedModel, valid_iter, criterion, device)\n",
    "scores['loss']['OT_pre'], scores['accuracy']['OT_pre'], scores['f1']['OT_pre'] = val_loss, val_acc, val_f1\n",
    "print(\"Validation Loss:     \", val_loss)\n",
    "print(\"Validation Accuracy:     \", val_acc.item())\n",
    "print(\"Validation F1:     \", val_f1.item())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recalibrate the normalization layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/47 [00:00<?, ?it/s]C:\\Users\\atace\\AppData\\Local\\Temp/ipykernel_17584/1843898165.py:14: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  src, trg = torch.tensor(src).to(device), torch.tensor(trg).to(device) # put to cpu/gpu\n",
      "100%|██████████| 47/47 [01:44<00:00,  2.22s/it]\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(fusedModel.parameters(), lr=0.01)\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "clip = 1\n",
    "\n",
    "fusedModel.train()\n",
    "for name, param in fusedModel.named_parameters():\n",
    "    if \"weight\" in name or \"bias\" in name:\n",
    "        param.requires_grad = False\n",
    "    \n",
    "for i, batch in enumerate(tqdm(train_iter)):\n",
    "            src = batch[0] # X\n",
    "            trg = batch[1] # y\n",
    "            src, trg = torch.tensor(src).to(device), torch.tensor(trg).to(device) # put to cpu/gpu\n",
    "            optimizer.zero_grad() # reset optimizer\n",
    "            output = fusedModel(src) # predict\n",
    "            y_pred = torch.argmax(output, dim=-1) # logits -> labels\n",
    "            output_reshape = output.contiguous().view(-1, output.shape[-1])\n",
    "            trg = trg.to(torch.int64)\n",
    "            loss = criterion(output_reshape, trg) # calculate loss\n",
    "            agreements = torch.eq(y_pred, trg)\n",
    "            accuracy = torch.mean(agreements.double()) # calculate accuracy\n",
    "            loss.backward() # backward pass\n",
    "\n",
    "            torch.nn.utils.clip_grad_norm_(fusedModel.parameters(), clip)\n",
    "            optimizer.step() # optimize model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Loss:      0.5758781035741171\n",
      "Validation Accuracy:      0.7533989662648054\n",
      "Validation F1:      0.7518618702888489\n"
     ]
    }
   ],
   "source": [
    "# Validate the fused model\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "val_loss, val_acc, val_f1 = validation(fusedModel, valid_iter, criterion, device)\n",
    "scores['loss']['OT_pre'], scores['accuracy']['OT_pre'], scores['f1']['OT_pre'] = val_loss, val_acc, val_f1\n",
    "scores['loss']['OT_calibr'], scores['accuracy']['OT_calibr'], scores['f1']['OT_calibr'] = val_loss, val_acc, val_f1\n",
    "print(\"Validation Loss:     \", val_loss)\n",
    "print(\"Validation Accuracy:     \", val_acc.item())\n",
    "print(\"Validation F1:     \", val_f1.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, param in fusedModel.named_parameters():\n",
    "    param.requires_grad = True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the last layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/47 [00:00<?, ?it/s]C:\\Users\\atace\\AppData\\Local\\Temp/ipykernel_17584/2985598703.py:14: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  src, trg = torch.tensor(src).to(device), torch.tensor(trg).to(device) # put to cpu/gpu\n",
      "100%|██████████| 47/47 [01:11<00:00,  1.52s/it]\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(fusedModel.parameters(), lr=0.001)\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "clip = 1\n",
    "\n",
    "fusedModel.train()\n",
    "for name, param in fusedModel.named_parameters():\n",
    "    if \"encoder\" in name:\n",
    "        param.requires_grad = False\n",
    "    \n",
    "for i, batch in enumerate(tqdm(train_iter)):\n",
    "            src = batch[0] # X\n",
    "            trg = batch[1] # y\n",
    "            src, trg = torch.tensor(src).to(device), torch.tensor(trg).to(device) # put to cpu/gpu\n",
    "            optimizer.zero_grad() # reset optimizer\n",
    "            output = fusedModel(src) # predict\n",
    "            y_pred = torch.argmax(output, dim=-1) # logits -> labels\n",
    "            output_reshape = output.contiguous().view(-1, output.shape[-1])\n",
    "            trg = trg.to(torch.int64)\n",
    "            loss = criterion(output_reshape, trg) # calculate loss\n",
    "            agreements = torch.eq(y_pred, trg)\n",
    "            accuracy = torch.mean(agreements.double()) # calculate accuracy\n",
    "            loss.backward() # backward pass\n",
    "\n",
    "            torch.nn.utils.clip_grad_norm_(fusedModel.parameters(), clip)\n",
    "            optimizer.step() # optimize model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Loss:      0.46658990780512494\n",
      "Validation Accuracy:      0.8097768612521151\n",
      "Validation F1:      0.8090723752975464\n"
     ]
    }
   ],
   "source": [
    "# Validate the fused model\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "val_loss, val_acc, val_f1 = validation(fusedModel, valid_iter, criterion, device)\n",
    "scores['loss']['OT_last_layer'], scores['accuracy']['OT_last_layer'], scores['f1']['OT_last_layer'] = val_loss, val_acc, val_f1\n",
    "print(\"Validation Loss:     \", val_loss)\n",
    "print(\"Validation Accuracy:     \", val_acc.item())\n",
    "print(\"Validation F1:     \", val_f1.item())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss Fused Model:      0.43884028991063434\n",
      "Test Accuracy Fused Model:      0.8128881526898734\n",
      "Test F1 Fused Model:      0.8131979703903198\n",
      "Test Loss Model A:      0.4548064172267914\n",
      "Test Accuracy Model A:      0.8403274986814346\n",
      "Test F1 Model A:      0.838917076587677\n",
      "Test Loss Model B:      0.41694880028565723\n",
      "Test Accuracy Model B:      0.8556393393987342\n",
      "Test F1 Model B:      0.8558376431465149\n"
     ]
    }
   ],
   "source": [
    "# Test the models\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "test_loss_fused, test_acc_fused, test_f1_fused  = validation(fusedModel, test_iter, criterion, device)\n",
    "scores['loss']['OT_test_set'], scores['accuracy']['OT_test_set'], scores['f1']['OT_test_set'] = test_loss_fused, test_acc_fused, test_f1_fused\n",
    "test_loss_A, test_acc_A, test_f1_A  = validation(modelA, test_iter, criterion, device)\n",
    "scores['loss']['A_test_set'], scores['accuracy']['A_test_set'], scores['f1']['A_test_set'] = test_loss_A, test_acc_A, test_f1_A\n",
    "test_loss_B, test_acc_B, test_f1_B  = validation(modelB, test_iter, criterion, device)\n",
    "scores['loss']['B_test_set'], scores['accuracy']['B_test_set'], scores['f1']['B_test_set'] = test_loss_A, test_acc_A, test_f1_A\n",
    "\n",
    "print(\"Test Loss Fused Model:     \", test_loss_fused)\n",
    "print(\"Test Accuracy Fused Model:     \", test_acc_fused.item())\n",
    "print(\"Test F1 Fused Model:     \", test_f1_fused.item())\n",
    "print(\"Test Loss Model A:     \", test_loss_A)\n",
    "print(\"Test Accuracy Model A:     \", test_acc_A.item())\n",
    "print(\"Test F1 Model A:     \", test_f1_A.item())\n",
    "print(\"Test Loss Model B:     \", test_loss_B)\n",
    "print(\"Test Accuracy Model B:     \", test_acc_B.item())\n",
    "print(\"Test F1 Model B:     \", test_f1_B.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame.from_dict(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>loss</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>f1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>A_test_set</th>\n",
       "      <td>0.473577</td>\n",
       "      <td>tensor(0.8281, dtype=torch.float64)</td>\n",
       "      <td>tensor(0.8281)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>B_test_set</th>\n",
       "      <td>0.473577</td>\n",
       "      <td>tensor(0.8281, dtype=torch.float64)</td>\n",
       "      <td>tensor(0.8281)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>OT_pre</th>\n",
       "      <td>0.485952</td>\n",
       "      <td>tensor(0.7773, dtype=torch.float64)</td>\n",
       "      <td>tensor(0.7776)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>OT_calibr</th>\n",
       "      <td>0.485952</td>\n",
       "      <td>tensor(0.7773, dtype=torch.float64)</td>\n",
       "      <td>tensor(0.7776)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>OT_last_layer</th>\n",
       "      <td>0.470886</td>\n",
       "      <td>tensor(0.7895, dtype=torch.float64)</td>\n",
       "      <td>tensor(0.7891)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>OT_test_set</th>\n",
       "      <td>0.479690</td>\n",
       "      <td>tensor(0.7863, dtype=torch.float64)</td>\n",
       "      <td>tensor(0.7868)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   loss                             accuracy              f1\n",
       "A_test_set     0.473577  tensor(0.8281, dtype=torch.float64)  tensor(0.8281)\n",
       "B_test_set     0.473577  tensor(0.8281, dtype=torch.float64)  tensor(0.8281)\n",
       "OT_pre         0.485952  tensor(0.7773, dtype=torch.float64)  tensor(0.7776)\n",
       "OT_calibr      0.485952  tensor(0.7773, dtype=torch.float64)  tensor(0.7776)\n",
       "OT_last_layer  0.470886  tensor(0.7895, dtype=torch.float64)  tensor(0.7891)\n",
       "OT_test_set    0.479690  tensor(0.7863, dtype=torch.float64)  tensor(0.7868)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rename rows, cols\n",
    "df.columns = ['Loss', 'Accuracy', 'F1 score']\n",
    "df.index = ['Model A on test set', 'Model B on test set', 'Optimal transport + weighted fusion', 'OT + weighted fusion (recalibrated)',\n",
    "            'OT + weighted fusion (last layer retrained)', 'OT + weighted fusion (last layer retrained) on test set']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\atace\\AppData\\Local\\Temp/ipykernel_13764/1818736071.py:2: FutureWarning: In future versions `DataFrame.to_latex` is expected to utilise the base implementation of `Styler.to_latex` for formatting and rendering. The arguments signature may therefore change. It is recommended instead to use `DataFrame.style.to_latex` which also contains additional functionality.\n",
      "  latex = df.to_latex(index=True,\n"
     ]
    }
   ],
   "source": [
    "# convert to latex\n",
    "latex = df.to_latex(index=True,\n",
    "                    bold_rows=True,\n",
    "                    caption='Model performance (5-fold CV)',\n",
    "                    position='H').replace('BOLD\\\\', r'\\textbf').replace('\\}', '}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\begin{table}[H]\n",
      "\\centering\n",
      "\\caption{Model performance (5-fold CV)}\n",
      "\\begin{tabular}{lrll}\n",
      "\\toprule\n",
      "{} &      Loss &                             Accuracy &        F1 score \\\\\n",
      "\\midrule\n",
      "\\textbf{Model A on test set                               } &  0.473577 &  tensor(0.8281, dtype=torch.float64) &  tensor(0.8281) \\\\\n",
      "\\textbf{Model B on test set                               } &  0.473577 &  tensor(0.8281, dtype=torch.float64) &  tensor(0.8281) \\\\\n",
      "\\textbf{Optimal transport + weighted fusion               } &  0.485952 &  tensor(0.7773, dtype=torch.float64) &  tensor(0.7776) \\\\\n",
      "\\textbf{OT + weighted fusion (recalibrated)               } &  0.485952 &  tensor(0.7773, dtype=torch.float64) &  tensor(0.7776) \\\\\n",
      "\\textbf{OT + weighted fusion (last layer retrained)       } &  0.470886 &  tensor(0.7895, dtype=torch.float64) &  tensor(0.7891) \\\\\n",
      "\\textbf{OT + weighted fusion (last layer retrained) on ...} &  0.479690 &  tensor(0.7863, dtype=torch.float64) &  tensor(0.7868) \\\\\n",
      "\\bottomrule\n",
      "\\end{tabular}\n",
      "\\end{table}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(latex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5 (default, Sep  3 2020, 21:29:08) [MSC v.1916 64 bit (AMD64)]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "b1744cd9dc0832a8d503a2c77e6bee76d4493b3bf33a738cf38afd0bb2e60262"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}