AA_TransUNet_Training_Demo_Cloud_Cover.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/YangYimin98/AA-TransUNet/blob/main/AA_TransUNet_Training_Demo_Cloud_Cover.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Author: Yimin Yang\n",
        "\n",
        "#Date: Feb 1, 2022\n",
        "\n",
        "# Small demo for paper: AA-TransUNet: Attention Augmented TransUNet For Nowcasting Tasks.\n",
        "Note that this is just for eaiser use, some implementations may be different with regular one, like early stopping criterion is not used here.\n",
        "\n",
        "References: \n",
        "* [TransUNet](https://github.com/Beckschen/TransUNet)\n",
        "* [Self-attention-cv](https://github.com/The-AI-Summer/self-attention-cv)\n",
        "* [ViT-pytorch](https://github.com/jeonsworld/ViT-pytorch)\n",
        "* [SmaAt-UNet](https://github.com/HansBambel/SmaAt-UNet)\n"
      ],
      "metadata": {
        "id": "1oG0lYwRqFnS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-kWt6OOAKQSF",
        "outputId": "58897c36-ece1-4fdf-bf5f-60b920382428"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fKbJkKAdfByd",
        "outputId": "839c03a1-7211-4bf7-b335-43b59ed383d7"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "device(type='cuda', index=0)"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ],
      "source": [
        "import torch\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "device"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N8i46W1kfIVg"
      },
      "outputs": [],
      "source": [
        "! pip install pytorch_lightning==0.7.6\n",
        "! pip install einops"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tak9iccKfLIr"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from torch.utils.data import Dataset\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "class cloud_maps(Dataset):\n",
        "    def __init__(self, folder, train=True, input_imgs=4, output_imgs=6):\n",
        "        super(cloud_maps, self).__init__()\n",
        "        self.train = train\n",
        "\n",
        "        self.folder_name = os.path.join(folder,'train' if self.train else 'test')\n",
        "        self.input_imgs = input_imgs\n",
        "        self.output_imgs = output_imgs\n",
        "\n",
        "        # Dataset is all the images\n",
        "        self.dataset = os.listdir(self.folder_name)\n",
        "\n",
        "        self.size_dataset = len(self.dataset)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "\n",
        "        imgs = np.load(os.path.join(self.folder_name, self.dataset[index]))['arr_0']\n",
        "\n",
        "        input_img = np.transpose(imgs[:, :, :self.input_imgs], axes=[2, 0, 1]).astype(dtype=\"float32\")\n",
        "        target_imgs = np.transpose(imgs[:, :, -self.output_imgs:], axes=[2, 0, 1]).astype(dtype=\"float32\")\n",
        "\n",
        "        return input_img, target_imgs\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.size_dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sjtLWyd3fPQt"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch import nn, optim\n",
        "import torch.nn.functional\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.utils.data.sampler import SubsetRandomSampler\n",
        "import numpy as np\n",
        "import pytorch_lightning as pl\n",
        "\n",
        "\n",
        "class AA_TransUnet_base(pl.LightningModule):\n",
        "\n",
        "    def __init__(self, hparams):\n",
        "        super().__init__()\n",
        "        self.hparams = hparams\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        opt = optim.Adam(self.parameters(), lr=self.hparams['learning_rate'])\n",
        "        scheduler = {\n",
        "            'scheduler': optim.lr_scheduler.ReduceLROnPlateau(opt,\n",
        "                                                              mode=\"min\",\n",
        "                                                              factor=0.1,\n",
        "                                                              patience=self.hparams['lr_patience']),\n",
        "            'monitor': 'val_loss',  # Default: val_loss\n",
        "        }\n",
        "        return [opt], [scheduler]\n",
        "\n",
        "    def loss_func(self, y_pred, y_true):\n",
        "        return nn.functional.mse_loss(y_pred, y_true)\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        x, y = batch[0].to(device), batch[1].to(device)\n",
        "        y_pred = self(x)\n",
        "        loss = self.loss_func(y_pred.squeeze(), y)\n",
        "        return {'loss': loss}\n",
        "\n",
        "    def training_epoch_end(self, outputs):\n",
        "        loss_mean = 0.0\n",
        "        for output in outputs:\n",
        "            loss_mean += output['loss']\n",
        "\n",
        "        loss_mean /= len(outputs)\n",
        "        return {\"log\": {\"train_loss\": loss_mean},\n",
        "                \"progress_bar\": {\"train_loss\": loss_mean}}\n",
        "\n",
        "    def validation_step(self, batch, batch_idx):\n",
        "        x, y = batch[0].to(device), batch[1].to(device)\n",
        "        y_pred = self(x)\n",
        "        val_loss = self.loss_func(y_pred.squeeze(), y)\n",
        "        return {\"val_loss\": val_loss}\n",
        "\n",
        "    def validation_epoch_end(self, outputs):\n",
        "        avg_loss = 0.0\n",
        "        for output in outputs:\n",
        "            avg_loss += output[\"val_loss\"]\n",
        "        avg_loss /= len(outputs)\n",
        "        logs = {\"val_loss\": avg_loss}\n",
        "        return {\"val_loss\": avg_loss, \"log\": logs,\n",
        "                \"progress_bar\": {\"val_loss\": avg_loss}}\n",
        "\n",
        "    def test_step(self, batch, batch_idx):\n",
        "        x, y = batch[0].to(device), batch[1].to(device)\n",
        "        y_pred = self(x)\n",
        "        val_loss = self.loss_func(y_pred.squeeze(), y)\n",
        "        return {\"test_loss\": val_loss}\n",
        "\n",
        "    def test_epoch_end(self, outputs):\n",
        "        avg_loss = 0.0\n",
        "        for output in outputs:\n",
        "            avg_loss += output[\"test_loss\"]\n",
        "        avg_loss /= len(outputs)\n",
        "        logs = {\"test_loss\": avg_loss}\n",
        "        return {\"test_loss\": avg_loss, \"log\": logs,\n",
        "                \"progress_bar\": {\"test_loss\": avg_loss}}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FAbeyIBnfft1"
      },
      "outputs": [],
      "source": [
        "import torch.nn as nn\n",
        "\n",
        "def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:\n",
        "    \"\"\"3x3 convolution with padding\"\"\"\n",
        "    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
        "                     padding=dilation, groups=groups, bias=False, dilation=dilation)\n",
        "\n",
        "\n",
        "def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:\n",
        "    \"\"\"1x1 convolution\"\"\"\n",
        "    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n",
        "\n",
        "\n",
        "class Bottleneck(nn.Module):\n",
        "\n",
        "    expansion = 1\n",
        "\n",
        "    def __init__(self, inplanes, planes, stride=1, groups=1,\n",
        "                 base_width=64, dilation=1, norm_layer=None):\n",
        "        super(Bottleneck, self).__init__()\n",
        "\n",
        "        if norm_layer is None:\n",
        "            norm_layer = nn.BatchNorm2d\n",
        "\n",
        "        if stride != 1 or inplanes != planes * self.expansion:\n",
        "            self.downsample = nn.Sequential(\n",
        "                conv1x1(inplanes, planes * self.expansion, stride),\n",
        "                norm_layer(planes * self.expansion),\n",
        "            )\n",
        "        else:\n",
        "            self.downsample = nn.Identity()\n",
        "\n",
        "        width = int(planes * (base_width / 64.)) * groups\n",
        "\n",
        "        self.conv1 = conv1x1(inplanes, width)\n",
        "        self.bn1 = norm_layer(width)\n",
        "        self.conv2 = conv3x3(width, width, stride, groups, dilation)\n",
        "        self.bn2 = norm_layer(width)\n",
        "        self.conv3 = conv1x1(width, planes * self.expansion)\n",
        "        self.bn3 = norm_layer(planes * self.expansion)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.conv1(x)\n",
        "        out = self.bn1(out)\n",
        "        out = self.relu(out)\n",
        "\n",
        "        out = self.conv2(out)\n",
        "        out = self.bn2(out)\n",
        "        out = self.relu(out)\n",
        "\n",
        "        out = self.conv3(out)\n",
        "        out = self.bn3(out)\n",
        "        identity = self.downsample(x)\n",
        "        out += identity\n",
        "        out = self.relu(out)\n",
        "\n",
        "        return out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yvK8gyXkfcvQ"
      },
      "outputs": [],
      "source": [
        "class Cloud_base(AA_TransUnet_base):\n",
        "\n",
        "    def __init__(self, hparams):\n",
        "        super(Cloud_base, self).__init__(hparams=hparams)\n",
        "        self.train_dataset = None\n",
        "        self.valid_dataset = None\n",
        "        self.train_sampler = None\n",
        "        self.valid_sampler = None\n",
        "\n",
        "    def prepare_data(self):\n",
        "        self.train_dataset = cloud_maps(\n",
        "            folder=self.hparams['dataset_folder'], train=True, input_imgs=self.hparams['num_input_images'],\n",
        "            output_imgs=self.hparams['num_output_images']\n",
        "        )\n",
        "        self.valid_dataset = cloud_maps(\n",
        "            folder=self.hparams['dataset_folder'], train=True, input_imgs=self.hparams['num_input_images'],\n",
        "            output_imgs=self.hparams['num_output_images']\n",
        "        )\n",
        "        num_train = len(self.train_dataset)\n",
        "        indices = list(range(num_train))\n",
        "        split = int(np.floor(self.hparams['valid_size'] * num_train))\n",
        "\n",
        "        np.random.shuffle(indices)\n",
        "\n",
        "        train_idx, valid_idx = indices[split:], indices[:split]\n",
        "        self.train_sampler = SubsetRandomSampler(train_idx)\n",
        "        self.valid_sampler = SubsetRandomSampler(valid_idx)\n",
        "\n",
        "    def train_dataloader(self):\n",
        "        train_loader = torch.utils.data.DataLoader(\n",
        "            self.train_dataset, batch_size=self.hparams['batch_size'], sampler=self.train_sampler,\n",
        "            num_workers=4, pin_memory=True\n",
        "        )\n",
        "        return train_loader\n",
        "\n",
        "    def val_dataloader(self):\n",
        "        valid_loader = torch.utils.data.DataLoader(\n",
        "            self.valid_dataset, batch_size=self.hparams['batch_size'], sampler=self.valid_sampler,\n",
        "            num_workers=4, pin_memory=True\n",
        "        )\n",
        "        return valid_loader\n",
        "    def test_dataloader(self):\n",
        "        test_loader = torch.utils.data.DataLoader(\n",
        "            self.test_dataset, batch_size=self.hparams['batch_size'], sampler=self.test_sampler,\n",
        "            num_workers=2, pin_memory=True\n",
        "        )\n",
        "        return test_loader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gntaBeXHfmbh"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "\n",
        "class SingleConv(nn.Module):\n",
        "    \"\"\"\n",
        "    Double convolution block that keeps that spatial sizes the same\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, in_ch, out_ch, norm_layer=None):\n",
        "        super(SingleConv, self).__init__()\n",
        "\n",
        "        if norm_layer is None:\n",
        "            norm_layer = nn.BatchNorm2d\n",
        "\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1),\n",
        "            norm_layer(out_ch),\n",
        "            nn.ReLU(inplace=True))\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.conv(x)\n",
        "\n",
        "\n",
        "class DoubleConv(nn.Module):\n",
        "    \"\"\"\n",
        "    Double convolution block that keeps that spatial sizes the same\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, in_ch, out_ch, norm_layer=None):\n",
        "        super(DoubleConv, self).__init__()\n",
        "        self.conv = nn.Sequential(SingleConv(in_ch, out_ch, norm_layer),\n",
        "                                  SingleConv(out_ch, out_ch, norm_layer))\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.conv(x)\n",
        "\n",
        "\n",
        "class Up(nn.Module):\n",
        "    \"\"\"\n",
        "    Doubles spatial size with bilinear upsampling\n",
        "    Skip connections and double convs\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, in_ch, out_ch):\n",
        "        super(Up, self).__init__()\n",
        "        mode = \"bilinear\"\n",
        "        self.up = nn.Upsample(scale_factor=2, mode=mode, align_corners=True)\n",
        "        self.conv = DoubleConv(in_ch, out_ch)\n",
        "\n",
        "    def forward(self, x1, x2=None):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            x1: [b,c, h, w]\n",
        "            x2: [b,c, 2*h,2*w]\n",
        "        Returns: 2x upsampled double conv reselt\n",
        "        \"\"\"\n",
        "        x = self.up(x1)\n",
        "        if x2 is not None:\n",
        "            x = torch.cat([x2, x], dim=1)\n",
        "        return self.conv(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Czum4jHmfqY5"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import random\n",
        "from typing import List, Tuple\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "from einops import repeat\n",
        "from torch import Tensor, nn\n",
        "\n",
        "\n",
        "def expand_to_batch(tensor, desired_size):\n",
        "    tile = desired_size // tensor.shape[0]\n",
        "    return repeat(tensor, 'b ... -> (b tile) ...', tile=tile)\n",
        "\n",
        "\n",
        "def init_random_seed(seed, gpu=False):\n",
        "    random.seed(seed)\n",
        "    np.random.seed(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    torch.cuda.manual_seed_all(seed)\n",
        "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
        "    if gpu:\n",
        "        torch.backends.cudnn.deterministic = True\n",
        "\n",
        "\n",
        "def get_module_device(parameter: nn.Module):\n",
        "    try:\n",
        "        return next(parameter.parameters()).device\n",
        "    except StopIteration:\n",
        "        # For nn.DataParallel compatibility in PyTorch 1.5\n",
        "\n",
        "        def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]:\n",
        "            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]\n",
        "            return tuples\n",
        "\n",
        "        gen = parameter._named_members(get_members_fn=find_tensor_attributes)\n",
        "        first_tuple = next(gen)\n",
        "        return first_tuple[1].device"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "from einops import rearrange\n",
        "from torch import nn\n",
        "\n",
        "\n",
        "def compute_mhsa(q, k, v, scale_factor=1, mask=None):\n",
        "    # resulted shape will be: [batch, heads, tokens, tokens]\n",
        "    scaled_dot_prod = torch.einsum('... i d , ... j d -> ... i j', q, k) * scale_factor\n",
        "\n",
        "    if mask is not None:\n",
        "        assert mask.shape == scaled_dot_prod.shape[2:]\n",
        "        scaled_dot_prod = scaled_dot_prod.masked_fill(mask, -np.inf)\n",
        "\n",
        "    attention = torch.softmax(scaled_dot_prod, dim=-1)\n",
        "    # calc result per head\n",
        "    return torch.einsum('... i j , ... j d -> ... i d', attention, v)\n",
        "\n",
        "class MultiHeadSelfAttention(nn.Module):\n",
        "    def __init__(self, dim, heads=8, dim_head=None):\n",
        "        \"\"\"\n",
        "        Implementation of multi-head attention layer of the original transformer model.\n",
        "        einsum and einops.rearrange is used whenever possible\n",
        "        Args:\n",
        "            dim: token's dimension, i.e. word embedding vector size\n",
        "            heads: the number of distinct representations to learn\n",
        "            dim_head: the dim of the head. In general dim_head<dim.\n",
        "            However, it may not necessary be (dim/heads)\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.dim_head = (int(dim / heads)) if dim_head is None else dim_head\n",
        "        _dim = self.dim_head * heads\n",
        "        self.heads = heads\n",
        "        self.to_qvk = nn.Linear(dim, _dim * 3, bias=False)\n",
        "        self.W_0 = nn.Linear(_dim, dim, bias=False)\n",
        "        self.scale_factor = self.dim_head ** -0.5\n",
        "\n",
        "    def forward(self, x, mask=None):\n",
        "        assert x.dim() == 3\n",
        "        qkv = self.to_qvk(x)  # [batch, tokens, dim*3*heads ]\n",
        "\n",
        "        # decomposition to q,v,k and cast to tuple\n",
        "        # the resulted shape before casting to tuple will be: [3, batch, heads, tokens, dim_head]\n",
        "        q, k, v = tuple(rearrange(qkv, 'b t (d k h ) -> k b h t d ', k=3, h=self.heads))\n",
        "\n",
        "        out = compute_mhsa(q, k, v, mask=mask, scale_factor=self.scale_factor)\n",
        "\n",
        "        # re-compose: merge heads with dim_head\n",
        "        out = rearrange(out, \"b h t d -> b t (h d)\")\n",
        "        # Apply final linear transformation layer\n",
        "        return self.W_0(out)\n",
        "class TransformerBlock(nn.Module):\n",
        "    \"\"\"\n",
        "    Vanilla transformer block from the original paper \"Attention is all you need\"\n",
        "    Detailed analysis: https://theaisummer.com/transformer/\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dim, heads=8, dim_head=None,\n",
        "                 dim_linear_block=1024, dropout=0.1, activation=nn.GELU,\n",
        "                 mhsa=None, prenorm=False):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            dim: token's vector length\n",
        "            heads: number of heads\n",
        "            dim_head: if none dim/heads is used\n",
        "            dim_linear_block: the inner projection dim\n",
        "            dropout: probability of droppping values\n",
        "            mhsa: if provided you can change the vanilla self-attention block\n",
        "            prenorm: if the layer norm will be applied before the mhsa or after\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.mhsa = mhsa if mhsa is not None else MultiHeadSelfAttention(dim=dim, heads=heads, dim_head=dim_head)\n",
        "        self.prenorm = prenorm\n",
        "        self.drop = nn.Dropout(dropout)\n",
        "        self.norm_1 = nn.LayerNorm(dim)\n",
        "        self.norm_2 = nn.LayerNorm(dim)\n",
        "\n",
        "        self.linear = nn.Sequential(\n",
        "            nn.Linear(dim, dim_linear_block),\n",
        "            activation(),  # nn.ReLU or nn.GELU\n",
        "            nn.Dropout(dropout),\n",
        "            nn.Linear(dim_linear_block, dim),\n",
        "            nn.Dropout(dropout)\n",
        "        )\n",
        "\n",
        "    def forward(self, x, mask=None):\n",
        "        if self.prenorm:\n",
        "            y = self.drop(self.mhsa(self.norm_1(x), mask)) + x\n",
        "            out = self.linear(self.norm_2(y)) + y\n",
        "        else:\n",
        "            y = self.norm_1(self.drop(self.mhsa(x, mask)) + x)\n",
        "            out = self.norm_2(self.linear(y) + y)\n",
        "        return out\n",
        "class TransformerEncoder(nn.Module):\n",
        "\n",
        "    def __init__(self, dim, blocks=6, heads=8, dim_head=None, dim_linear_block=1024, dropout=0, prenorm=False):\n",
        "        super().__init__()\n",
        "        self.block_list = [TransformerBlock(dim, heads, dim_head,\n",
        "                                            dim_linear_block, dropout, prenorm=prenorm) for _ in range(blocks)]\n",
        "        self.layers = nn.ModuleList(self.block_list)\n",
        "\n",
        "    def forward(self, x, mask=None):\n",
        "        for layer in self.layers:\n",
        "            x = layer(x, mask)\n",
        "        return x"
      ],
      "metadata": {
        "id": "J69ffaVdMAf_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HE1UZVm6f00G"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from einops import rearrange\n",
        "\n",
        "\n",
        "class ViT(nn.Module):\n",
        "    def __init__(self, *,\n",
        "                 img_dim,\n",
        "                 in_channels=3,\n",
        "                 patch_dim=16,\n",
        "                 num_classes=1,\n",
        "                 dim=512,\n",
        "                 blocks=6,\n",
        "                 heads=4,\n",
        "                 dim_linear_block=1024,\n",
        "                 dim_head=None,\n",
        "                 dropout=0.1, transformer=None, classification=True):\n",
        "        \"\"\"\n",
        "        Minimal re-implementation of ViT\n",
        "        Args:\n",
        "            img_dim: the spatial image size\n",
        "            in_channels: number of img channels\n",
        "            patch_dim: desired patch dim\n",
        "            num_classes: classification task classes\n",
        "            dim: the linear layer's dim to project the patches for MHSA\n",
        "            blocks: number of transformer blocks\n",
        "            heads: number of heads\n",
        "            dim_linear_block: inner dim of the transformer linear block\n",
        "            dim_head: dim head in case you want to define it. defaults to dim/heads\n",
        "            dropout: for pos emb and transformer\n",
        "            transformer: in case you want to provide another transformer implementation\n",
        "            classification: creates an extra CLS token that we will index in the final classification layer\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        assert img_dim % patch_dim == 0, f'patch size {patch_dim} not divisible by img dim {img_dim}'\n",
        "        self.p = patch_dim\n",
        "        self.classification = classification\n",
        "        # tokens = number of patches\n",
        "        tokens = (img_dim // patch_dim) ** 2\n",
        "        self.token_dim = in_channels * (patch_dim ** 2)\n",
        "        self.dim = dim\n",
        "        self.dim_head = (int(self.dim / heads)) if dim_head is None else dim_head\n",
        "\n",
        "        # Projection and pos embeddings\n",
        "        self.project_patches = nn.Linear(self.token_dim, self.dim)\n",
        "\n",
        "        self.emb_dropout = nn.Dropout(dropout)\n",
        "\n",
        "        self.cls_token = nn.Parameter(torch.randn(1, 1, self.dim))\n",
        "        self.pos_emb1D = nn.Parameter(torch.randn(tokens + 1, self.dim))\n",
        "\n",
        "        if self.classification:\n",
        "            self.mlp_head = nn.Linear(self.dim, num_classes)\n",
        "\n",
        "        if transformer is None:\n",
        "            self.transformer = TransformerEncoder(self.dim, blocks=blocks, heads=heads,\n",
        "                                                  dim_head=self.dim_head,\n",
        "                                                  dim_linear_block=dim_linear_block,\n",
        "                                                  dropout=dropout)\n",
        "        else:\n",
        "            self.transformer = transformer\n",
        "\n",
        "    def forward(self, img, mask=None):\n",
        "        # Create patches\n",
        "        # from [batch, channels, h, w] to [batch, tokens , N], N=p*p*c , tokens = h/p *w/p\n",
        "        img_patches = rearrange(img,\n",
        "                                'b c (patch_x x) (patch_y y) -> b (x y) (patch_x patch_y c)',\n",
        "                                patch_x=self.p, patch_y=self.p)\n",
        "\n",
        "        batch_size, tokens, _ = img_patches.shape\n",
        "\n",
        "        # project patches with linear layer + add pos emb\n",
        "        img_patches = self.project_patches(img_patches)\n",
        "\n",
        "        img_patches = torch.cat((expand_to_batch(self.cls_token, desired_size=batch_size), img_patches), dim=1)\n",
        "\n",
        "        # add pos. embeddings. + dropout\n",
        "        # indexing with the current batch's token length to support variable sequences\n",
        "        img_patches = img_patches + self.pos_emb1D[:tokens + 1, :]\n",
        "        patch_embeddings = self.emb_dropout(img_patches)\n",
        "\n",
        "        # feed patch_embeddings and output of transformer. shape: [batch, tokens, dim]\n",
        "        y = self.transformer(patch_embeddings, mask)\n",
        "\n",
        "        # we index only the cls token for classification. nlp tricks :P\n",
        "        return self.mlp_head(y[:, 0, :]) if self.classification else y[:, 1:, :]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XBtO02mVCLxK"
      },
      "outputs": [],
      "source": [
        "import torch.nn.functional as F\n",
        "\n",
        "class DepthwiseSeparableConv(nn.Module):\n",
        "    def __init__(self, in_channels, output_channels, kernels_per_layer=1):\n",
        "        super(DepthwiseSeparableConv, self).__init__()\n",
        "        # In Tensorflow DepthwiseConv2D has depth_multiplier instead of kernels_per_layer\n",
        "        self.depthwise = nn.Conv2d(in_channels, in_channels * kernels_per_layer, groups=in_channels, kernel_size=1)\n",
        "        self.pointwise = nn.Conv2d(in_channels * kernels_per_layer, output_channels, kernel_size=1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.depthwise(x)\n",
        "        x = self.pointwise(x)\n",
        "        return x\n",
        "class DoubleConvDS(nn.Module):\n",
        "    \"\"\"(convolution => [BN] => ReLU) * 2\"\"\"\n",
        "\n",
        "    def __init__(self, in_channels, out_channels, mid_channels=None, kernels_per_layer=1):\n",
        "        super().__init__()\n",
        "        if not mid_channels:\n",
        "            mid_channels = out_channels\n",
        "        self.double_conv = nn.Sequential(\n",
        "            DepthwiseSeparableConv(in_channels, mid_channels, kernels_per_layer=kernels_per_layer),\n",
        "            nn.BatchNorm2d(mid_channels),\n",
        "            nn.ReLU(inplace=True),\n",
        "            DepthwiseSeparableConv(mid_channels, out_channels, kernels_per_layer=kernels_per_layer),\n",
        "            nn.BatchNorm2d(out_channels),\n",
        "            nn.ReLU(inplace=True)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.double_conv(x)\n",
        "class UpDS(nn.Module):\n",
        "    \"\"\"Upscaling then double conv\"\"\"\n",
        "\n",
        "    def __init__(self, in_channels, out_channels, bilinear=True, kernels_per_layer=1):\n",
        "        super().__init__()\n",
        "\n",
        "        # if bilinear, use the normal convolutions to reduce the number of channels\n",
        "        if bilinear:\n",
        "            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
        "            self.conv = DoubleConvDS(in_channels, out_channels, in_channels // 2, kernels_per_layer=kernels_per_layer)\n",
        "        else:\n",
        "            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)\n",
        "            self.conv = DoubleConvDS(in_channels, out_channels, kernels_per_layer=kernels_per_layer)\n",
        "\n",
        "    def forward(self, x1, x2=None):\n",
        "        x = self.up(x1)\n",
        "        if x2 is not None:\n",
        "            x = torch.cat([x2, x], dim=1)\n",
        "        return self.conv(x)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ypUE80JaWjQm"
      },
      "outputs": [],
      "source": [
        "class Flatten(nn.Module):\n",
        "    def forward(self, x):\n",
        "        return x.view(x.size(0), -1)\n",
        "\n",
        "\n",
        "class ChannelAttention(nn.Module):\n",
        "    def __init__(self, input_channels, reduction_ratio=16):\n",
        "        super(ChannelAttention, self).__init__()\n",
        "        self.input_channels = input_channels\n",
        "        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n",
        "        self.max_pool = nn.AdaptiveMaxPool2d(1)\n",
        "        #  https://github.com/luuuyi/CBAM.PyTorch/blob/master/model/resnet_cbam.py\n",
        "        #  uses Convolutions instead of Linear\n",
        "        self.MLP = nn.Sequential(\n",
        "            Flatten(),\n",
        "            nn.Linear(input_channels, input_channels // reduction_ratio),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(input_channels // reduction_ratio, input_channels)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Take the input and apply average and max pooling\n",
        "        avg_values = self.avg_pool(x)\n",
        "        max_values = self.max_pool(x)\n",
        "        out = self.MLP(avg_values) + self.MLP(max_values)\n",
        "        scale = x * torch.sigmoid(out).unsqueeze(2).unsqueeze(3).expand_as(x)\n",
        "        return scale\n",
        "\n",
        "\n",
        "class SpatialAttention(nn.Module):\n",
        "    def __init__(self, kernel_size=7):\n",
        "        super(SpatialAttention, self).__init__()\n",
        "        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'\n",
        "        padding = 3 if kernel_size == 7 else 1\n",
        "        self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=padding, bias=False)\n",
        "        self.bn = nn.BatchNorm2d(1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        avg_out = torch.mean(x, dim=1, keepdim=True)\n",
        "        max_out, _ = torch.max(x, dim=1, keepdim=True)\n",
        "        out = torch.cat([avg_out, max_out], dim=1)\n",
        "        out = self.conv(out)\n",
        "        out = self.bn(out)\n",
        "        scale = x * torch.sigmoid(out)\n",
        "        return scale\n",
        "\n",
        "\n",
        "class CBAM(nn.Module):\n",
        "    def __init__(self, input_channels, reduction_ratio=16, kernel_size=7):\n",
        "        super(CBAM, self).__init__()\n",
        "        self.channel_att = ChannelAttention(input_channels, reduction_ratio=reduction_ratio)\n",
        "        self.spatial_att = SpatialAttention(kernel_size=kernel_size)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.channel_att(x)\n",
        "        out = self.spatial_att(out)\n",
        "        return out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XyFr_on6gh0w"
      },
      "outputs": [],
      "source": [
        "# model\n",
        "import torch.nn as nn\n",
        "from einops import rearrange\n",
        "\n",
        "\n",
        "class AA_TransUnet(Cloud_base):\n",
        "    def __init__(self, hparams\n",
        "                 ):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            img_dim: the img dimension\n",
        "            in_channels: channels of the input\n",
        "            classes: desired segmentation classes\n",
        "            vit_blocks: MHSA blocks of ViT\n",
        "            vit_heads: number of MHSA heads\n",
        "            vit_dim_linear_mhsa_block: MHSA MLP dimension\n",
        "            vit_transformer: pass your own version of vit\n",
        "            vit_channels: the channels of your pretrained vit. default is 128*8\n",
        "            patch_dim: for image patches of the vit\n",
        "        \"\"\"\n",
        "        super(AA_TransUnet, self).__init__(hparams=hparams)\n",
        "\n",
        "        self.inplanes = 128\n",
        "        self.patch_size = hparams['patch_size']\n",
        "        self.vit_transformer_dim = hparams['vit_transformer_dim']\n",
        "        vit_channels = self.inplanes * 8 if hparams['vit_channels'] is None else hparams['vit_channels']\n",
        "\n",
        "        in_conv1 = nn.Conv2d(hparams['in_channels'], self.inplanes, kernel_size=7, stride=2, padding=3,\n",
        "                             bias=False)\n",
        "\n",
        "        bn1 = nn.BatchNorm2d(self.inplanes)\n",
        "        self.init_conv = nn.Sequential(in_conv1, bn1, nn.ReLU(inplace=True))\n",
        "        self.CBAM9 = CBAM(128)\n",
        "        self.conv1 = Bottleneck(self.inplanes, self.inplanes * 2, stride=2)\n",
        "        self.CBAM6 = CBAM(256)\n",
        "        self.conv2 = Bottleneck(self.inplanes * 2, self.inplanes * 4, stride=2)\n",
        "        self.CBAM7 = CBAM(512)\n",
        "        self.conv3 = Bottleneck(self.inplanes * 4, vit_channels, stride=2)\n",
        "        self.CBAM8 = CBAM(1024)\n",
        "\n",
        "        self.img_dim_vit = hparams['img_dim'] // 16\n",
        "\n",
        "        assert (self.img_dim_vit % hparams['patch_size'] == 0), \"Vit patch_dim not divisible\"\n",
        "        #\n",
        "        self.vit = ViT(img_dim=self.img_dim_vit,\n",
        "                       in_channels=vit_channels,  # input features' channels (encoder)\n",
        "                       patch_dim=hparams['patch_size'],\n",
        "                       # transformer inside dimension that input features will be projected\n",
        "                       # out will be [batch, dim_out_vit_tokens, dim ]\n",
        "                       dim=hparams['vit_transformer_dim'],\n",
        "                       blocks=hparams['vit_blocks'],\n",
        "                       heads=hparams['vit_heads'],\n",
        "                       dim_linear_block=hparams['vit_dim_linear_mhsa_block'],\n",
        "                       classification=False) if hparams['vit_transformer'] is None else hparams['vit_transformer']\n",
        "\n",
        "        # to project patches back - undoes vit's patchification\n",
        "        token_dim = vit_channels * (hparams['patch_size'] ** 2)\n",
        "        self.project_patches_back = nn.Linear(hparams['vit_transformer_dim'], token_dim)\n",
        "        # upsampling path\n",
        "        self.vit_conv = SingleConv(in_ch=vit_channels, out_ch=512)\n",
        "        self.cbam1 = CBAM(512)\n",
        "        self.dec1 = UpDS(vit_channels, 256)\n",
        "        self.cbam2 = CBAM(256)\n",
        "        self.dec2 = UpDS(512, 128)\n",
        "        self.cbam3 = CBAM(128)\n",
        "        self.dec3 = UpDS(256, 64)\n",
        "        self.cbam4 = CBAM(64)\n",
        "        self.dec4 = UpDS(64, 16)\n",
        "        self.cbam5 = CBAM(16)\n",
        "        self.conv1x1 = nn.Conv2d(in_channels=16, out_channels=hparams['classes'], kernel_size=1)\n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        x2 = self.init_conv(x)\n",
        "        x2 = self.CBAM9(x2)\n",
        "\n",
        "        x4 = self.conv1(x2)\n",
        "        x4 = self.CBAM6(x4)\n",
        "\n",
        "        x8 = self.conv2(x4)\n",
        "        x8 = self.CBAM7(x8)\n",
        "\n",
        "        x16 = self.conv3(x8)  # out shape of 1024, img_dim_vit, img_dim_vit\n",
        "        x16 = self.CBAM8(x16)\n",
        "\n",
        "        y = self.vit(x16)  # out shape of number_of_patches, vit_transformer_dim\n",
        "        # from [number_of_patches, vit_transformer_dim] -> [number_of_patches, token_dim]\n",
        "        y = self.project_patches_back(y)\n",
        "        # from [batch, number_of_patches, token_dim] -> [batch, channels, img_dim_vit, img_dim_vit]\n",
        "        y = rearrange(y, 'b (x y) (patch_x patch_y c) -> b c (patch_x x) (patch_y y)',\n",
        "                      x=self.img_dim_vit // self.patch_size, y=self.img_dim_vit // self.patch_size,\n",
        "                      patch_x=self.patch_size, patch_y=self.patch_size)\n",
        "\n",
        "        y = self.vit_conv(y)\n",
        "        y = self.cbam1(y)\n",
        "        y = self.dec1(y, x8)\n",
        "\n",
        "        y = self.cbam2(y)\n",
        "        y = self.dec2(y, x4)\n",
        "\n",
        "        y = self.cbam3(y)\n",
        "        y = self.dec3(y, x2)\n",
        "\n",
        "        y = self.cbam4(y)\n",
        "        y = self.dec4(y)\n",
        "\n",
        "        y = self.cbam5(y)\n",
        "        y = self.conv1x1(y)\n",
        "        return y\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "itmit1ApgnAB"
      },
      "outputs": [],
      "source": [
        "if __name__ == \"__main__\":\n",
        "    args = {\n",
        "\n",
        "        'vit_blocks': 1,\n",
        "        'vit_heads': 1,\n",
        "        'vit_dim_linear_mhsa_block':3072,\n",
        "        'patch_size': 2,\n",
        "        'vit_transformer_dim': 1024,  #768 is also a good parameter values\n",
        "        'vit_transformer': None,\n",
        "        'vit_channels': None,\n",
        "        'classes': 6,\n",
        "        'img_dim': 256,  # 288 for precipitation data set\n",
        "        \"in_channels\": 4,  \n",
        "\n",
        "        \"batch_size\": 6,\n",
        "        \"learning_rate\": 0.001,\n",
        "        'gpus': -1,\n",
        "        \"lr_patience\": 4,  # learning rate decay\n",
        "        \"es_patience\": 30,  # early stopping criterion\n",
        "        \"use_oversampled_dataset\": True,\n",
        "        \"bilinear\": True,\n",
        "        \"num_input_images\": 4, # 12 for precipitation data set\n",
        "        \"num_output_images\": 6,\n",
        "        \"valid_size\": 0.1,\n",
        "\n",
        "        \"dataset_folder\": \"/content/drive/MyDrive/AA_TransUNet_Project/Data_cloud_cover_preprocessed\"\n",
        "}\n",
        "\n",
        "    net = AA_TransUnet(hparams=args)\n",
        "    net = net.to(device)\n",
        "    trainer = pl.Trainer(gpus=-1,\n",
        "                         fast_dev_run=False,\n",
        "                         weights_summary='top',\n",
        "                         max_epochs=100)\n",
        "    trainer.fit(net)\n",
        "    trainer.save_checkpoint('/content/drive/MyDrive/AA_TransUNet_Project/results/Model_Saved/cloud_1.ckpt')\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "AA_TransUNet_Demo.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}