model_fusion_optimal_transport_basic_fixed.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",
    "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validation(model, iterator, criterion, device):\n",
    "    # set model into evaluation mode\n",
    "    model.eval()\n",
    "\n",
    "    # validation\n",
    "    # loss, metrics for current epoch\n",
    "    val_epoch_loss = 0\n",
    "    val_epoch_accuracy = 0\n",
    "\n",
    "    with torch.no_grad(): # stop graph\n",
    "        # batches\n",
    "        for i, batch in enumerate(iterator):\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",
    "            output = model(src)\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",
    "\n",
    "            loss = criterion(output_reshape, trg) # calculate loss\n",
    "            agreements = torch.eq(y_pred, trg)\n",
    "            accuracy = torch.mean(agreements.double()) # calculate accuracy\n",
    "\n",
    "            val_epoch_loss += loss.item()\n",
    "            val_epoch_accuracy += accuracy\n",
    "\n",
    "    # return mean loss w.r.t. batches\n",
    "    return val_epoch_loss / len(iterator), val_epoch_accuracy / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the embedding matrix\n",
    "embedding = torch.load(\"Models/embedding_16.pt\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Model Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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 = 1,\n",
    "                              drop_prob = 0.5,\n",
    "                              device = device)\n",
    "\n",
    "modelA.load_state_dict(torch.load(\"Models\\modelA_IMDB_256\"))\n",
    "modelA.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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 = 1,\n",
    "                              drop_prob = 0.5,\n",
    "                              device = device)\n",
    "\n",
    "modelB.load_state_dict(torch.load(\"Models\\modelB_IMBD_256\"))\n",
    "modelB.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, param in modelA.named_parameters():\n",
    "    print(name, param.shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OT Functions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fusion via Optimal Transport"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 = 1,\n",
    "                              drop_prob = 0.2,\n",
    "                              device = device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "                        W_fusion[nameA], transport_matrix_triplet, _ = fusion(nameA, nameB, weightA, weightB, transport_matrix, beta)\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 = validation(fusedModel, valid_iter, criterion, device)\n",
    "    return val_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "study = optuna.create_study()\n",
    "study.optimize(obj, n_trials=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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",
    "                    W_fusion[nameA], transport_matrix_triplet, _ = fusion(nameA, nameB, weightA, weightB, transport_matrix, beta)\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",
    "                \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 = validation(fusedModel, valid_iter, criterion, device)\n",
    "print(\"Validation Loss:     \", val_loss)\n",
    "print(\"Validation Accuracy:     \", val_acc.item())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recalibrate the normalization layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Validate the fused model\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "val_loss, val_acc = validation(fusedModel, valid_iter, criterion, device)\n",
    "print(\"Validation Loss:     \", val_loss)\n",
    "print(\"Validation Accuracy:     \", val_acc.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Validate the fused model\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "val_loss, val_acc = validation(fusedModel, valid_iter, criterion, device)\n",
    "print(\"Validation Loss:     \", val_loss)\n",
    "print(\"Validation Accuracy:     \", val_acc.item())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the models\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "test_loss_fused, test_acc_fused = validation(fusedModel, test_iter, criterion, device)\n",
    "test_loss_A, test_acc_A = validation(modelA, test_iter, criterion, device)\n",
    "test_loss_B, test_acc_B = validation(modelB, test_iter, criterion, device)\n",
    "\n",
    "print(\"Validation Loss Fused Model:     \", test_loss_fused)\n",
    "print(\"Validation Accuracy Fused Model:     \", test_acc_fused.item())\n",
    "print(\"Validation Loss Model A:     \", test_loss_A)\n",
    "print(\"Validation Accuracy Model A:     \", test_acc_A.item())\n",
    "print(\"Validation Loss Model B:     \", test_loss_B)\n",
    "print(\"Validation Accuracy Model B:     \", test_acc_B.item())"
   ]
  },
  {
   "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
}