soap-u-net.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 52,
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-07-02T17:59:48.131882Z",
     "start_time": "2024-07-02T17:59:48.129424Z"
    }
   },
   "id": "fa1552be6bbfac23"
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Self Attention\n",
    "\n",
    "Introduced in the now famous paper, *Attention Is All You Need* (2017), self-attention is a method to enhance contextual information between elements in an input. It was first proposed for U-Nets in *Attention U-Net: Learning Where to Look for the Pancreas* (2018).\n",
    "\n",
    "Let $\\boldsymbol{x}$ be a sequence of $T$ vector inputs. The query matrix $\\boldsymbol{W}_q$ can be thought of as embedding a question about the input sequence. The key matrix $\\boldsymbol{W}_q$ embeds how relevant a token in the sequence is in query space.\n",
    "\n",
    "$$\\boldsymbol{q}_i = \\boldsymbol{W}_q\\boldsymbol{x}_i$$\n",
    "$$\\boldsymbol{k}_i = \\boldsymbol{W}_k\\boldsymbol{x}_i$$\n",
    "\n",
    "The dot product between the ith query and jth key sequences indicates how well the key *answers* the query. These are called the unnormalised attention weights $\\boldsymbol{\\omega}$.\n",
    "\n",
    "$$\\boldsymbol{\\omega}_{ij} = \\boldsymbol{q}_i^\\top \\cdot \\boldsymbol{k}_j$$\n",
    "\n",
    "These weights are scaled by $\\sqrt{d_k}$ where $d_k$ is the length of $x_i$ and normalised by the softmax function to get the normalised attention scores $\\boldsymbol{\\alpha}$:\n",
    "\n",
    "$$\\boldsymbol{\\alpha}_{ij} = \\text{softmax}\\left(\\frac{\\boldsymbol{\\omega}_{ij}}{\\sqrt{d_k}}\\right)$$\n",
    "\n",
    "Scaling by $\\sqrt{d_k}$ ensures that the magnitude of the weight vectors will be similar. The last step is to compute the context vector $\\boldsymbol{z}$ which is an attention-weighted version of the original input $\\boldsymbol{x}$. First, the value matrix is multiplied to get the value sequence:\n",
    "\n",
    " $$\\boldsymbol{v}_i = \\boldsymbol{W}_v\\boldsymbol{x}_i$$\n",
    " \n",
    "And the context vector is taken to be the attention-weighted sum of this sequence:\n",
    "\n",
    "$$\\boldsymbol{z}_i = \\sum_{j = 1}^T \\boldsymbol{\\alpha}_{ij}\\boldsymbol{v}_i$$\n",
    "\n",
    "In multi-headed attention, multiple key, query, and value matrices are applied parallely."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ca6d0fe83c3f8567"
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "outputs": [],
   "source": [
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, channels, size, heads=4, activation=nn.GELU):\n",
    "        super().__init__()\n",
    "        self.channels = channels\n",
    "        self.size = size\n",
    "\n",
    "        self.attention = nn.MultiheadAttention(channels, heads, batch_first=True)\n",
    "        self.norm = nn.LayerNorm([channels])\n",
    "        self.ff_self = nn.Sequential(\n",
    "            nn.LayerNorm([channels]),\n",
    "            nn.Linear(channels, channels),\n",
    "            activation(),\n",
    "            nn.Linear(channels, channels)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(-1, self.channels, self.size ** 2).swapaxes(1, 2)\n",
    "        z = self.norm(x)\n",
    "        z, _ = self.mha(z, z, z)\n",
    "        z += x\n",
    "        z += self.ff_self(z)\n",
    "        return z.swapaxes(2, 1).view(-1, self.channels, self.size, self.size)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-07-02T18:58:59.556140Z",
     "start_time": "2024-07-02T18:58:59.550581Z"
    }
   },
   "id": "117ecc216067e92"
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "outputs": [],
   "source": [
    "class DoubleConv(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, mid_channels=None, residual=False, activation=nn.GELU):\n",
    "        super().__init__()\n",
    "        self.residual = residual\n",
    "\n",
    "        if mid_channels is None:\n",
    "            mid_channels = out_channels\n",
    "        self.double_conv = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),\n",
    "            nn.GroupNorm(1, mid_channels),\n",
    "            activation(),\n",
    "            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),\n",
    "            nn.GroupNorm(1, out_channels)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        if self.residual:\n",
    "            return F.gelu(x + self.double_conv())\n",
    "        return self.double_conv(x)\n",
    "    \n",
    "\n",
    "class DownConv(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, emb_dim=256):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.maxpool_conv = nn.Sequential(\n",
    "            nn.MaxPool2d(2),\n",
    "            DoubleConv(in_channels, in_channels, residual=True),\n",
    "            DoubleConv(in_channels, out_channels)\n",
    "        )\n",
    "        \n",
    "        self.embedding = nn.Sequential(\n",
    "            nn.SiLU(),\n",
    "            nn.Linear(emb_dim, out_channels)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x, t):\n",
    "        x = self.maxpool_conv(x)\n",
    "        embedding = self.embedding(t)[:, :, None, None].repeat(1, 1, x.shape[-2], x.shape[-1])\n",
    "        return x + embedding\n",
    "\n",
    "\n",
    "class UpConv(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, emb_dim=256):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.up = nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=True)\n",
    "        self.conv = nn.Sequential(\n",
    "            DoubleConv(in_channels, in_channels, residual=True),\n",
    "            DoubleConv(in_channels, out_channels, in_channels // 2)\n",
    "        )\n",
    "\n",
    "        self.embedding = nn.Sequential(\n",
    "            nn.SiLU(),\n",
    "            nn.Linear(emb_dim, out_channels)\n",
    "        )\n",
    "\n",
    "    def forward(self, x, skip_x, t):\n",
    "        x = self.up(x)\n",
    "        x = torch.concat([skip_x, x], dim=1)\n",
    "        x = self.conv(x)\n",
    "        embedding = self.embedding(t)[:, :, None, None].repeat(1, 1, x.shape[-2], x.shape[-1])\n",
    "        return x + embedding"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-07-02T18:59:01.284368Z",
     "start_time": "2024-07-02T18:59:01.278690Z"
    }
   },
   "id": "f98de379fb7a5cc9"
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "outputs": [],
   "source": [
    "class UNet(nn.Module):\n",
    "    def __init__(self, in_channels=3, out_channels=3, time_dim=256):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.time_dim = time_dim\n",
    "        \n",
    "        self.inc = DoubleConv(in_channels, 64)\n",
    "        self.down1 = DownConv(64, 128)\n",
    "        self.sa1 = SelfAttention(128, 32)\n",
    "        self.down2 = DownConv(128, 256)\n",
    "        self.sa2 = SelfAttention(256, 16)\n",
    "        self.down3 = DownConv(256, 256)\n",
    "        self.sa3 = SelfAttention(256, 8)\n",
    "\n",
    "        self.bot1 = DoubleConv(256, 512)\n",
    "        self.bot2 = DoubleConv(512, 512)\n",
    "        self.bot3 = DoubleConv(512, 256)\n",
    "\n",
    "        self.up1 = UpConv(512, 128)\n",
    "        self.sa4 = SelfAttention(128, 16)\n",
    "        self.up2 = UpConv(256, 64)\n",
    "        self.sa5 = SelfAttention(64, 32)\n",
    "        self.up3 = UpConv(128, 64)\n",
    "        self.sa6 = SelfAttention(64, 64)\n",
    "        \n",
    "        self.outc = nn.Conv2d(64, out_channels, kernel_size=1)\n",
    "\n",
    "    def pos_encoding(self, t, channels):\n",
    "        inv_freq = 1.0 / (10000 ** (torch.arange(0, channels, 2).float() / channels))\n",
    "        pos_enc_a = torch.sin(t.repeat(1, channels // 2) * inv_freq)\n",
    "        pos_enc_b = torch.cos(t.repeat(1, channels // 2) * inv_freq)\n",
    "        pos_enc = torch.cat([pos_enc_a, pos_enc_b], dim=-1)\n",
    "        return pos_enc\n",
    "\n",
    "    def forward(self, x, t):\n",
    "        t = t.unsqueeze(-1).type(torch.float)\n",
    "        t = self.pos_encoding(t, self.time_dim)\n",
    "\n",
    "        x1 = self.inc(x)\n",
    "        x2 = self.down1(x1, t)\n",
    "        x2 = self.sa1(x2)\n",
    "        x3 = self.down2(x2, t)\n",
    "        x3 = self.sa2(x3)\n",
    "        x4 = self.down3(x3, t)\n",
    "        x4 = self.sa3(x4)\n",
    "\n",
    "        x4 = self.bot1(x4)\n",
    "        x4 = self.bot2(x4)\n",
    "        x4 = self.bot3(x4)\n",
    "\n",
    "        x = self.up1(x4, x3, t)\n",
    "        x = self.sa4(x)\n",
    "        x = self.up2(x, x2, t)\n",
    "        x = self.sa5(x)\n",
    "        x = self.up3(x, x1, t)\n",
    "        x = self.sa6(x)\n",
    "        return self.outc(x)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-07-02T18:59:02.294393Z",
     "start_time": "2024-07-02T18:59:02.285610Z"
    }
   },
   "id": "c2adc484eb82575d"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class Diffusion:\n",
    "    def __init__(self, noise_steps=1000, beta_start=1e-4, beta_end=0.02, img_size=256):\n",
    "        self.noise_steps = noise_steps\n",
    "    \n",
    "        self.beta = torch.linspace(beta_start, beta_end, self.noise_steps)\n",
    "        self.alpha = 1. - self.beta\n",
    "        self.alpha_hat = torch.cumprod(self.alpha, dim=0)\n",
    "        \n",
    "        self.img_size = img_size\n",
    "\n",
    "    def noise_images(self, x, t):\n",
    "        sqrt_alpha_hat = torch.sqrt(self.alpha_hat[t])[:, None, None, None]\n",
    "        sqrt_one_minus_alpha_hat = torch.sqrt(1 - self.alpha_hat[t])[:, None, None, None]\n",
    "        err = torch.randn_like(x)\n",
    "        return sqrt_alpha_hat * x + sqrt_one_minus_alpha_hat * err, err\n",
    "\n",
    "    def sample_timesteps(self, n):\n",
    "        return torch.randint(low=1, high=self.noise_steps, size=(n,))\n",
    "    \n",
    "    def sample(self):\n",
    "        pass"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "fa0e82629a964655"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}