Transformer_loss_activation_and_optimizer_functions.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "Q9P2n6H5IexU"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "\n",
        "# from https://github.com/gzerveas/mvts_transformer/blob/master/src/models/loss.py\n",
        "\n",
        "def get_loss_module(task):\n",
        "\n",
        "    task = task\n",
        "\n",
        "    if (task == \"imputation\") or (task == \"transduction\"):\n",
        "        return MaskedMSELoss(reduction='none')  # outputs loss for each batch element\n",
        "\n",
        "    if task == \"classification\":\n",
        "        return NoFussCrossEntropyLoss(reduction='none')  # outputs loss for each batch sample\n",
        "\n",
        "    if task == \"regression\":\n",
        "        return nn.MSELoss(reduction='none')  # outputs loss for each batch sample\n",
        "\n",
        "    else:\n",
        "        raise ValueError(\"Loss module for task '{}' does not exist\".format(task))\n",
        "\n",
        "\n",
        "def l2_reg_loss(model):\n",
        "    \"\"\"Returns the squared L2 norm of output layer of given model\"\"\"\n",
        "\n",
        "    for name, param in model.named_parameters():\n",
        "        if name == 'output_layer.weight':\n",
        "            return torch.sum(torch.square(param))\n",
        "\n",
        "\n",
        "class NoFussCrossEntropyLoss(nn.CrossEntropyLoss):\n",
        "    \"\"\"\n",
        "    pytorch's CrossEntropyLoss is fussy: 1) needs Long (int64) targets only, and 2) only 1D.\n",
        "    This function satisfies these requirements\n",
        "    \"\"\"\n",
        "\n",
        "    def forward(self, inp, target):\n",
        "        return F.cross_entropy(inp, target.long().squeeze(), weight=self.weight,\n",
        "                               ignore_index=self.ignore_index, reduction=self.reduction)\n",
        "\n",
        "\n",
        "class MaskedMSELoss(nn.Module):\n",
        "    \"\"\" Masked MSE Loss\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, reduction: str = 'mean'):\n",
        "\n",
        "        super().__init__()\n",
        "\n",
        "        self.reduction = reduction\n",
        "        self.mse_loss = nn.MSELoss(reduction=self.reduction)\n",
        "\n",
        "    def forward(self,\n",
        "                y_pred: torch.Tensor, y_true: torch.Tensor, mask: torch.BoolTensor) -> torch.Tensor:\n",
        "        \"\"\"Compute the loss between a target value and a prediction.\n",
        "\n",
        "        Args:\n",
        "            y_pred: Estimated values\n",
        "            y_true: Target values\n",
        "            mask: boolean tensor with 0s at places where values should be ignored and 1s where they should be considered\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        if reduction == 'none':\n",
        "            (num_active,) Loss for each active batch element as a tensor with gradient attached.\n",
        "        if reduction == 'mean':\n",
        "            scalar mean loss over batch as a tensor with gradient attached.\n",
        "        \"\"\"\n",
        "\n",
        "        # for this particular loss, one may also elementwise multiply y_pred and y_true with the inverted mask\n",
        "        masked_pred = torch.masked_select(y_pred, mask)\n",
        "        masked_true = torch.masked_select(y_true, mask)\n",
        "\n",
        "        return self.mse_loss(masked_pred, masked_true)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "EFBodWS0eoej"
      },
      "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))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "r1nDyks75Msg"
      },
      "outputs": [],
      "source": [
        "from torch.optim.optimizer import Optimizer\n",
        "\n",
        "def get_optimizer(name):\n",
        "\n",
        "    if name == \"Adam\":\n",
        "        return torch.optim.Adam\n",
        "    elif name == \"AdamW\":\n",
        "        return torch.optim.AdamW"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}