Encoder_Transformer_Layers.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "UJaSvaVdjBz2"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import scipy\n",
        "import math"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "3VHt77rsdpXu"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "OZD-TDeJjVU2"
      },
      "outputs": [],
      "source": [
        "from torch.nn import Linear, Dropout, BatchNorm1d, LayerNorm, MultiheadAttention, TransformerEncoderLayer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Get activation function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def _get_activation_fn(activation):\n",
        "    if activation == \"relu\":\n",
        "        return F.relu\n",
        "    elif activation == \"gelu\":\n",
        "        return F.gelu\n",
        "    raise ValueError(\"activation should be relu/gelu, not {}\".format(activation))"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "rwUc2Jno9jiR"
      },
      "source": [
        "LearnablePositionalEncoding Class"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qYazeroU9d2h"
      },
      "outputs": [],
      "source": [
        "# from https://github.com/gzerveas/mvts_transformer/blob/master/src/models/ts_transformer.py\n",
        "class LearnablePositionalEncoding(nn.Module):\n",
        "    def __init__(self, d_model, max_len, scale_factor=1.0, dropout=0.1):\n",
        "        super(LearnablePositionalEncoding, self).__init__()\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "        # Each position gets its own embedding\n",
        "        # Since indices are always 0 ... max_len, we don't have to do a look-up\n",
        "        self.pe = nn.Parameter(torch.empty(max_len, 1, d_model))  # requires_grad automatically set to True\n",
        "        nn.init.uniform_(self.pe, -0.02, 0.02)\n",
        "\n",
        "    def forward(self, x):\n",
        "        r\"\"\"Inputs of forward function\n",
        "        Args:\n",
        "            x: the sequence fed to the positional encoder model (required).\n",
        "        Shape:\n",
        "            x: [sequence length, batch size, embed dim]\n",
        "            output: [sequence length, batch size, embed dim]\n",
        "        \"\"\"\n",
        "\n",
        "        x = x + self.pe[:x.size(0), :]\n",
        "        return self.dropout(x)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "6gFkP94ceLtz"
      },
      "source": [
        "Fixed positional Encoding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HPIV0UpoeQZq"
      },
      "outputs": [],
      "source": [
        "# From https://github.com/pytorch/examples/blob/master/word_language_model/model.py\n",
        "class FixedPositionalEncoding(nn.Module):\n",
        "    r\"\"\"Inject some information about the relative or absolute position of the tokens\n",
        "        in the sequence. The positional encodings have the same dimension as\n",
        "        the embeddings, so that the two can be summed. Here, we use sine and cosine\n",
        "        functions of different frequencies.\n",
        "    .. math::\n",
        "        \\text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))\n",
        "        \\text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))\n",
        "        \\text{where pos is the word position and i is the embed idx)\n",
        "    Args:\n",
        "        d_model: the embed dim (required).\n",
        "        dropout: the dropout value (default=0.1).\n",
        "        max_len: the max. length of the incoming sequence (default=1024).\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, d_model, max_len, scale_factor=1.0, dropout=0.1):\n",
        "        super(FixedPositionalEncoding, self).__init__()\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "\n",
        "        pe = torch.zeros(max_len, d_model)  # positional encoding\n",
        "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
        "        pe[:, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 1::2] = torch.cos(position * div_term)\n",
        "        pe = scale_factor * pe.unsqueeze(0).transpose(0, 1)\n",
        "        self.register_buffer('pe', pe)  # this stores the variable in the state_dict (used for non-trainable variables)\n",
        "\n",
        "    def forward(self, x):\n",
        "        r\"\"\"Inputs of forward function\n",
        "        Args:\n",
        "            x: the sequence fed to the positional encoder model (required).\n",
        "        Shape:\n",
        "            x: [sequence length, batch size, embed dim]\n",
        "            output: [sequence length, batch size, embed dim]\n",
        "        \"\"\"\n",
        "\n",
        "        x = x + self.pe[:x.size(0), :]\n",
        "        return self.dropout(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_8AV4RRiccyX"
      },
      "outputs": [],
      "source": [
        "def get_pos_encoder(pos_encoding):\n",
        "    if pos_encoding == \"learnable\":\n",
        "        return LearnablePositionalEncoding\n",
        "    elif pos_encoding == \"fixed\":\n",
        "        return FixedPositionalEncoding\n",
        "\n",
        "    raise NotImplementedError(\"pos_encoding should be 'learnable'/'fixed', not '{}'\".format(pos_encoding))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ibKwL2Ogn8Vc"
      },
      "outputs": [],
      "source": [
        "class TransformerBatchNormEncoderLayer(nn.modules.Module):\n",
        "    r\"\"\"This transformer encoder layer block is made up of self-attn and feedforward network.\n",
        "    It differs from TransformerEncoderLayer in torch/nn/modules/transformer.py in that it replaces LayerNorm\n",
        "    with BatchNorm.\n",
        "    Args:\n",
        "        d_model: the number of expected features in the input (required).\n",
        "        nhead: the number of heads in the multiheadattention models (required).\n",
        "        dim_feedforward: the dimension of the feedforward network model (default=2048).\n",
        "        dropout: the dropout value (default=0.1).\n",
        "        activation: the activation function of intermediate layer, relu or gelu (default=relu).\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, d_model, nhead, dim_feedforward=256, dropout=0.1, activation=\"relu\"):\n",
        "        super(TransformerBatchNormEncoderLayer, self).__init__()\n",
        "        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)\n",
        "        # Implementation of Feedforward model\n",
        "        self.linear1 = Linear(d_model, dim_feedforward)\n",
        "        self.dropout = Dropout(dropout)\n",
        "        self.linear2 = Linear(dim_feedforward, d_model)\n",
        "\n",
        "        self.norm1 = BatchNorm1d(d_model, eps=1e-5)  # normalizes each feature across batch samples and time steps\n",
        "        self.norm2 = BatchNorm1d(d_model, eps=1e-5)\n",
        "        self.dropout1 = Dropout(dropout)\n",
        "        self.dropout2 = Dropout(dropout)\n",
        "\n",
        "        self.activation = _get_activation_fn(activation)\n",
        "\n",
        "    def __setstate__(self, state):\n",
        "        if 'activation' not in state:\n",
        "            state['activation'] = F.relu\n",
        "        super(TransformerBatchNormEncoderLayer, self).__setstate__(state)\n",
        "\n",
        "    def forward(self, src, src_mask = None, is_causal = None,\n",
        "                src_key_padding_mask = None):\n",
        "        r\"\"\"Pass the input through the encoder layer.\n",
        "        Args:\n",
        "            src: the sequence to the encoder layer (required).\n",
        "            src_mask: the mask for the src sequence (optional).\n",
        "            src_key_padding_mask: the mask for the src keys per batch (optional).\n",
        "        Shape:\n",
        "            see the docs in Transformer class. \n",
        "             - src: :math:`(S, E)` for unbatched input, :math:`(S, N, E)` if `batch_first=False` or\n",
        "              `(N, S, E)` if `batch_first=True`.\n",
        "        - src_key_padding_mask: :math:`(S)` for unbatched input otherwise :math:`(N, S)`.\n",
        "        - output: :math:`(T, E)` for unbatched input, :math:`(T, N, E)` if `batch_first=False` or\n",
        "              `(N, T, E)` if `batch_first=True`.\n",
        "        batch_first – If True, then the input and output tensors are provided as (batch, seq, feature). Default: False (seq, batch, feature).\n",
        "        - src_mask: :math:`(S, S)` or :math:`(N\\cdot\\text{num\\_heads}, S, S)`.\n",
        "        \"\"\"\n",
        "        src2 = self.self_attn(src, src, src, attn_mask=src_mask,\n",
        "                              key_padding_mask=src_key_padding_mask)[0]   # src2 shape = (seq, batch, feature) \n",
        "                              # because batch false is set to false because we would have batched it before calling self-attention function\n",
        "        \n",
        "        # Add dropout and add x to output of multihead attention for residual...\n",
        "        src = src + self.dropout1(src2)  # (seq_len, batch_size, d_model)\n",
        "        # Reshape to be able to do batch norm and not layer norm\n",
        "        src = src.permute(1, 2, 0)  # (batch_size, d_model, seq_len)\n",
        "        # Perform batchnorm\n",
        "        src = self.norm1(src)\n",
        "        # restore shape back\n",
        "        src = src.permute(2, 0, 1)  # restore (seq_len, batch_size, d_model)\n",
        "        # apply linear layers 1 and 2 to src\n",
        "        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))\n",
        "        # Add x/src for residual efect and normalize again\n",
        "        src = src + self.dropout2(src2)  # (seq_len, batch_size, d_model)\n",
        "        src = src.permute(1, 2, 0)  # (batch_size, d_model, seq_len)\n",
        "        src = self.norm2(src)\n",
        "        src = src.permute(2, 0, 1)  # restore (seq_len, batch_size, d_model)\n",
        "        return src\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Bring your data in normally (NxTxD), it is in here you permutate it85\n",
        "# Call this class and specify whether you want pytorch's TransformerEncoder with layer norm or the one we created above with batch norm\n",
        "# Watch out for your src_mask shape!!!\n",
        "# Adapted to this problem from : https://github.com/gzerveas/mvts_transformer/blob/master/src/models/ts_transformer.py code base\n",
        "class TSTransformerEncoderClassiregressor(nn.Module):\n",
        "    \"\"\"\n",
        "    Simplest classifier/regressor. Can be either regressor or classifier because the output does not include\n",
        "    softmax. Concatenates final layer embeddings and uses 0s to ignore padding embeddings in final output layer.\n",
        "    Transfomations can either be linear, 1D-CNN or no transform\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, feat_dim, max_len, d_model, n_heads, num_layers, dim_feedforward, num_classes,\n",
        "                 dropout=0.1, pos_encoding='fixed', activation='gelu', norm='BatchNorm', freeze=False, transform='linear'):\n",
        "        super(TSTransformerEncoderClassiregressor, self).__init__()\n",
        "\n",
        "        self.max_len = max_len\n",
        "        self.d_model = d_model\n",
        "        self.n_heads = n_heads\n",
        "\n",
        "        self.project_inp = nn.Linear(feat_dim, d_model)\n",
        "\n",
        "        # CNN Model - 1D Convnet\n",
        "        self.CNN = nn.Sequential(\n",
        "            nn.Conv1d(14, 6, kernel_size=3, padding=\"same\"),\n",
        "            nn.ReLU(),\n",
        "            nn.Dropout(0.4))\n",
        "\n",
        "\n",
        "        self.pos_enc = get_pos_encoder(pos_encoding)(d_model=self.d_model, max_len=self.max_len, dropout=dropout*(1.0 - freeze))\n",
        "\n",
        "        if norm == 'LayerNorm':\n",
        "            encoder_layer = TransformerEncoderLayer(d_model, self.n_heads, dim_feedforward, dropout*(1.0 - freeze), activation=activation)\n",
        "        else:\n",
        "            encoder_layer = TransformerBatchNormEncoderLayer(d_model, self.n_heads, dim_feedforward, dropout*(1.0 - freeze), activation=activation)\n",
        "\n",
        "        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)   # \n",
        "\n",
        "        self.act = _get_activation_fn(activation)\n",
        "\n",
        "        self.dropout1 = nn.Dropout(dropout)\n",
        "\n",
        "        self.feat_dim = feat_dim\n",
        "        self.num_classes = num_classes\n",
        "        self.output_layer = self.build_output_module(d_model, max_len, num_classes)\n",
        "\n",
        "    def build_output_module(self, d_model, max_len, num_classes):\n",
        "        output_layer = nn.Linear(d_model * max_len, num_classes)\n",
        "        # no softmax (or log softmax), because CrossEntropyLoss does this internally. If probabilities are needed,\n",
        "        # add F.log_softmax and use NLLoss\n",
        "        return output_layer\n",
        "\n",
        "    def forward(self, X, padding_masks, transform):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            X: (batch_size, seq_length, feat_dim) torch tensor of masked features (input)\n",
        "            padding_masks: (batch_size, seq_length) boolean tensor, 1 means keep vector at this position, 0 means padding\n",
        "        Returns:\n",
        "            output: (batch_size, num_classes)\n",
        "        \"\"\"\n",
        "\n",
        "        # permute because pytorch convention for transformers is [seq_length, batch_size, feat_dim]. padding_masks [batch_size, feat_dim]\n",
        "        inp = X.permute(1, 0, 2)\n",
        "        if transform == 'linear':\n",
        "\n",
        "          # Linear input transformation\n",
        "          inp = self.project_inp(inp) * math.sqrt(\n",
        "              self.d_model)  # [seq_length, batch_size, d_model] project input vectors to d_model dimensional space\n",
        "\n",
        "        if transform == '1D-CNN':\n",
        "          # 1D-CNN input transformation\n",
        "          inp = X.permute(0, 2, 1)\n",
        "          inp = self.CNN(inp)\n",
        "          inp = inp.permute(2,0,1)\n",
        "\n",
        "\n",
        "        inp = self.pos_enc(inp)  # add positional encoding\n",
        "        # NOTE: logic for padding masks is reversed to comply with definition in MultiHeadAttention, TransformerEncoderLayer\n",
        "        output = self.transformer_encoder(inp, src_key_padding_mask=~padding_masks)  # (seq_length, batch_size, d_model)\n",
        "        output = self.act(output)  # the output transformer encoder/decoder embeddings don't include non-linearity\n",
        "        output = output.permute(1, 0, 2)  # (batch_size, seq_length, d_model) This particular permutation is to make sure we are working with NxTxD to compute Y\n",
        "        output = self.dropout1(output)\n",
        "\n",
        "        # Output\n",
        "        output = output * padding_masks.unsqueeze(-1)  # zero-out padding embeddings, now this is where you zero it out not actually remove d spaces\n",
        "        output = output.reshape(output.shape[0], -1)  # (batch_size, seq_length * d_model)\n",
        "        output = self.output_layer(output)  # (batch_size, num_classes) (No activation, it is in crossentropy loss)\n",
        "\n",
        "        return output\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}