DropAttack-BiGRU-PHEME.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "import string\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "from nltk import word_tokenize\n",
    "import matplotlib.pyplot as plt\n",
    "from nltk.corpus import stopwords\n",
    "from sklearn import metrics\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "enstop = stopwords.words('english')\n",
    "punct = string.punctuation\n",
    "\n",
    "def tokenizer(sent):\n",
    "    sent = sent.lower()\n",
    "    tmp = word_tokenize(sent)\n",
    "    res = []\n",
    "    for word in tmp:\n",
    "        if word not in enstop and word not in punct:\n",
    "            res.append(word)\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IMDB data read:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nynlluqzE_2N"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchtext import data\n",
    "from torchtext import vocab\n",
    "\n",
    "text_field = data.Field(tokenize=tokenizer, lower=True, include_lengths=True, fix_length=256)\n",
    "label_field = data.Field(sequential=False, use_vocab=False, dtype=torch.long)\n",
    "train, valid, test = data.TabularDataset.splits(path='',\n",
    "                                                train='./PHEME/2018_train.csv',\n",
    "                                                validation='./PHEME/2018_val.csv',\n",
    "                                                test='./PHEME/2018_test.csv',\n",
    "                                                format='csv', skip_header=True,\n",
    "                                                fields=[('id', None),('sentence', text_field), ('label', label_field)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read the GloVe word vector:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BoI5yPFIVEUG"
   },
   "outputs": [],
   "source": [
    "vec = vocab.Vectors(name='glove.6B.300d.txt')\n",
    "text_field.build_vocab(train, valid, test, max_size=250000, vectors=vec,\n",
    "                       unk_init=torch.Tensor.normal_)\n",
    "label_field.build_vocab(train, valid, test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wzwZVu7EbZRY"
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "train_iter, valid_iter, test_iter = data.BucketIterator.splits((train, valid, test), batch_sizes=(128, 128, 128),\n",
    "                                                               sort_key=lambda x: len(x.sentence),\n",
    "                                                               sort_within_batch=True,\n",
    "                                                               repeat=False, shuffle=True,\n",
    "                                                               device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OfDvZOj9w9_A"
   },
   "source": [
    "### Model training function (standard training and DropAttack adversarial training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "C4eRtD_gxA3l"
   },
   "outputs": [],
   "source": [
    "def train(model, train_iter, dev_iter, num_epoch, opt, criterion, eva, out_model_file):\n",
    "    print(\"Training begin!\")\n",
    "    model.train()\n",
    "    loss_list = []\n",
    "    dev_acc = []\n",
    "    train_acc = []\n",
    "    best_dev_acc = 0.\n",
    "    for epoch in range(num_epoch):\n",
    "        total_loss = 0.\n",
    "        for batch in train_iter:\n",
    "            output = model(batch.sentence)\n",
    "            loss = criterion(output, batch.label)\n",
    "            opt.zero_grad()\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "            total_loss += loss.item()\n",
    "        loss_list.append(total_loss)\n",
    "        dev_acc.append(eva(model, dev_iter))\n",
    "        train_acc.append(eva(model,train_iter))\n",
    "        print(f\"Epoch: {epoch+1}/{num_epoch}. Total loss: {total_loss:.3f}.Train_Acc: {train_acc[-1]:.3%}. Validation Set Acc: {dev_acc[-1]:.3%}.\")\n",
    "        if dev_acc[-1] > best_dev_acc:\n",
    "            best_dev_acc = dev_acc[-1]\n",
    "            torch.save(model.state_dict(), out_model_file)\n",
    "    return loss_list, dev_acc\n",
    "# import torch\n",
    "class DropAttack():\n",
    "    def __init__(self, model):\n",
    "        self.model = model\n",
    "        self.param_backup = {}\n",
    "        self.grad_backup = {}\n",
    "        self.mask_backup = {}\n",
    "\n",
    "    def attack(self, epsilon=5.0, p_attack =0.5, param_name='embed.weight', is_first_attack=False):\n",
    "        # The emb_name parameter should be replaced with the name of the parameter to be attacked in your model\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad and param_name == name:\n",
    "                if is_first_attack:\n",
    "                    self.param_backup[name] = param.data.clone()\n",
    "                    mask = np.random.binomial(n=1, p=p_attack, size= param.grad.shape)\n",
    "                    mask = torch.from_numpy(mask).float()  # attack mask\n",
    "                    self.mask_backup['mask'] = mask.clone()\n",
    "                else: mask = self.mask_backup['mask']\n",
    "                norm = torch.norm(param.grad)\n",
    "                if norm != 0 and not torch.isnan(norm):\n",
    "                    r_at = epsilon * param.grad / norm\n",
    "                    r_at *= mask.cuda()   # Randomly attack some of the parameters\n",
    "                    param.data.add_(r_at)    \n",
    "\n",
    "    def restore(self, param_name='embed.weight'):\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad and param_name == name: \n",
    "                assert name in self.param_backup\n",
    "                param.data = self.param_backup[name]\n",
    "                param_backup = {}\n",
    "\n",
    "    def backup_grad(self):\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                self.grad_backup[name] = param.grad.clone()\n",
    "\n",
    "    def restore_grad(self):\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                param.grad = self.grad_backup[name]\n",
    "                grad_backup = {}\n",
    "        \n",
    "def train_DA(model, train_iter, dev_iter, num_epoch, opt, criterion, eva, out_model_file):\n",
    "    K = 1\n",
    "    print(f'Adversarial training begin! (DropAttack-{K})')\n",
    "    model.train()\n",
    "    dropattack = DropAttack(model)\n",
    "    loss_list = []\n",
    "    dev_acc = []\n",
    "    best_dev_acc = 0.\n",
    "    for epoch in range(num_epoch):\n",
    "        total_loss = 0.\n",
    "        model.train()\n",
    "        for batch in train_iter:\n",
    "            output = model(batch.sentence)\n",
    "            loss = criterion(output, batch.label)\n",
    "            loss.backward(retain_graph=True)  # Calculate the original gradient\n",
    "            dropattack.backup_grad()    # Backup the initial gradient\n",
    "            # Attack the embedding layer\n",
    "            for t in range(K):\n",
    "                dropattack.attack(5, 0.5, 'embed.weight', is_first_attack=(t==0))  # Add adversarial disturbance to the parameters, backup param.data for the first attack\n",
    "                output = model(batch.sentence)\n",
    "                loss_adv1 = criterion(output, batch.label)/K\n",
    "                loss_adv1.backward(retain_graph=True) # # Backpropagation, and accumulate the gradient of the adversarial training based on the normal grad\n",
    "                loss += loss_adv1\n",
    "            dropattack.restore('embed.weight') # # Restore the disturbed parameters\n",
    "            \n",
    "            dropattack.restore_grad() \n",
    "            # Attack the hidden layer\n",
    "            for t in range(K):\n",
    "                dropattack.attack(5, 0.5, 'rnn.rnn.weight_hh_l0', is_first_attack=(t==0)) # Add adversarial disturbance to the parameters, backup param.data for the first attack\n",
    "                output = model(batch.sentence)\n",
    "                loss_adv2 = criterion(output, batch.label)/K\n",
    "                loss_adv2.backward(retain_graph=True) # Backpropagation, and accumulate the gradient of the adversarial training based on the normal grad\n",
    "                loss += loss_adv2\n",
    "            dropattack.restore('rnn.rnn.weight_hh_l0') # Restore the disturbed parameters\n",
    "            opt.zero_grad()\n",
    "            \n",
    "            loss.backward()\n",
    "            opt.step()  # Update parameters\n",
    "            # loss = loss + loss_adv1 + loss_adv2\n",
    "            total_loss += loss.item()\n",
    "            opt.zero_grad()\n",
    "        loss_list.append(total_loss)\n",
    "        dev_acc.append(eva(model, dev_iter))\n",
    "        print(f\"Epoch: {epoch+1}/{num_epoch}. Total loss: {total_loss:.3f}. Validation Set Acc: {dev_acc[-1]:.3%}.\")\n",
    "        if dev_acc[-1] > best_dev_acc:\n",
    "            best_dev_acc = dev_acc[-1]\n",
    "            torch.save(model.state_dict(), out_model_file)\n",
    "    return loss_list, dev_acc\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5dY4gBmd3Xuv"
   },
   "source": [
    "### Model evaluation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2lSdhhl03Zav"
   },
   "outputs": [],
   "source": [
    "def eva(model, data_iter):\n",
    "    correct, count = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for batch in data_iter:\n",
    "            pred = model(batch.sentence)\n",
    "            pred = torch.argmax(pred, dim=-1)\n",
    "            correct += (pred == batch.label).sum().item()\n",
    "            count += len(pred)\n",
    "    return correct / count"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LSTM model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "okrtrPcHemUR"
   },
   "outputs": [],
   "source": [
    "class LSTM(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, num_layers, bidirectional):\n",
    "        super(LSTM, self).__init__()\n",
    "        self.rnn = nn.LSTM(input_size=input_size, hidden_size=hidden_size,\n",
    "                           num_layers=num_layers, bidirectional=bidirectional)\n",
    "            \n",
    "    def forward(self, x, length):\n",
    "        packed_x = nn.utils.rnn.pack_padded_sequence(x, length)\n",
    "        packed_output, (hidden, cell) = self.rnn(packed_x)\n",
    "        output, output_lengths = nn.utils.rnn.pad_packed_sequence(packed_output)\n",
    "        return hidden, output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GRU(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, num_layers, bidirectional):\n",
    "        super(GRU, self).__init__()\n",
    "        self.rnn = nn.GRU(input_size=input_size, hidden_size=hidden_size,\n",
    "                          num_layers=num_layers, bidirectional=bidirectional)\n",
    "            \n",
    "    def forward(self, x, length):\n",
    "        packed_x = nn.utils.rnn.pack_padded_sequence(x, length)\n",
    "        packed_output, hidden = self.rnn(packed_x)\n",
    "        output, output_lengths = nn.utils.rnn.pad_packed_sequence(packed_output)\n",
    "        return hidden, output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextRNN(nn.Module):\n",
    "    def __init__(self, embed_size, hidden_size, num_layers, bidirectional, out_dim,\n",
    "                 pretrained_embed, use_gru=False, freeze=True,\n",
    "                 random_embed=False, vocab_size=None):\n",
    "        super(TextRNN, self).__init__()\n",
    "        if random_embed:\n",
    "            self.embed = nn.Embedding(vocab_size, embed_size)\n",
    "        else:\n",
    "            self.embed = nn.Embedding.from_pretrained(pretrained_embed, freeze=False)\n",
    "        if use_gru:\n",
    "            self.rnn = GRU(embed_size, hidden_size, num_layers, bidirectional)\n",
    "        else:\n",
    "            self.rnn = LSTM(embed_size, hidden_size, num_layers, bidirectional)\n",
    "        self.proj = nn.Linear(2*hidden_size, out_dim)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        text, text_length = x # text: [seq_len, bs]\n",
    "        text = text.permute(1, 0) # text: [bs, seq_len]\n",
    "        embed_x = self.embed(text) # embed_x: [bs, seq_len, embed_dim]\n",
    "        embed_x = embed_x.permute(1, 0, 2) # embed_x: [seq_len, bs, embed_dim]\n",
    "        hidden, _ = self.rnn(embed_x, text_length) # hidden: [2*num_layers, bs, hidden_size]\n",
    "        hidden = torch.cat((hidden[-1,:,:], hidden[-2,:,:]), dim=1)\n",
    "        return self.proj(hidden)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adversarial training begin! (DropAttack-1)\n",
      "Epoch: 1/5. Total loss: 65.645. Validation Set Acc: 83.516%.\n",
      "Epoch: 2/5. Total loss: 46.544. Validation Set Acc: 86.028%.\n"
     ]
    }
   ],
   "source": [
    "embed_size = 300\n",
    "hidden_size = 300\n",
    "num_layers = 1\n",
    "bidirectional = True\n",
    "out_dim = 2\n",
    "pretrained_embed = text_field.vocab.vectors\n",
    "lr = 0.001\n",
    "num_epoch = 5\n",
    "freeze = False\n",
    "use_gru = True\n",
    "random_embed = False\n",
    "vocab_size = len(text_field.vocab.stoi)\n",
    "out_model_file = 'textrnn_PHEME_DA.pt'\n",
    "# ————————————————————————————————————————————————————————\n",
    "use_dropattack = True   # Whether to use DropAttack\n",
    "# ————————————————————————————————————————————————————————\n",
    "model = TextRNN(embed_size, hidden_size, num_layers, bidirectional, out_dim,\n",
    "                                pretrained_embed, use_gru=use_gru, freeze=freeze,\n",
    "                               random_embed=random_embed, vocab_size=None).to(device)\n",
    "opt = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "if use_dropattack:\n",
    "    loss_list, dev_acc_list = train_DA(model, train_iter, valid_iter, num_epoch, opt, criterion, eva, out_model_file)\n",
    "else:\n",
    "    loss_list, dev_acc_list = train(model, train_iter, valid_iter, num_epoch, opt, criterion, eva, out_model_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test set acc: 85.185%\n"
     ]
    }
   ],
   "source": [
    "model = TextRNN(embed_size, hidden_size, num_layers, bidirectional, out_dim,\n",
    "                             pretrained_embed, use_gru=use_gru, freeze=freeze,\n",
    "                               random_embed=random_embed, vocab_size=None).to(device)\n",
    "model.load_state_dict(torch.load('textrnn_PHEME_DA.pt'))\n",
    "print(f\"Test set acc: {eva(model, test_iter):.3%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "文本分类项目.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}