{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.10","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":5441,"databundleVersionId":38425,"sourceType":"competition"},{"sourceId":7018640,"sourceType":"datasetVersion","datasetId":3533769},{"sourceId":7018641,"sourceType":"datasetVersion","datasetId":3533752}],"dockerImageVersionId":30513,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\n#import numpy as np # linear algebra\n#import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n# import matplotlib.pyplot as plt\n# # import tensorflow as tf\n# # import tensorflow.keras as krs\n# import os\n# import zipfile\n# import glob\n# import cv2\n# import shutil\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\n# for dirname, _, filenames in os.walk('/kaggle/input'):\n#     for filename in filenames:\n#         print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:35.862769Z","iopub.execute_input":"2023-11-17T17:32:35.863158Z","iopub.status.idle":"2023-11-17T17:32:35.868770Z","shell.execute_reply.started":"2023-11-17T17:32:35.863129Z","shell.execute_reply":"2023-11-17T17:32:35.867731Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir -p /kaggle/working/models","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:35.870779Z","iopub.execute_input":"2023-11-17T17:32:35.871055Z","iopub.status.idle":"2023-11-17T17:32:36.834178Z","shell.execute_reply.started":"2023-11-17T17:32:35.871032Z","shell.execute_reply":"2023-11-17T17:32:36.832914Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir -p /kaggle/working/metrics","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:36.835756Z","iopub.execute_input":"2023-11-17T17:32:36.836119Z","iopub.status.idle":"2023-11-17T17:32:37.824492Z","shell.execute_reply.started":"2023-11-17T17:32:36.836080Z","shell.execute_reply":"2023-11-17T17:32:37.823271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install memorizing_transformers_pytorch einops_exts ","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-11-17T17:32:37.825961Z","iopub.execute_input":"2023-11-17T17:32:37.826289Z","iopub.status.idle":"2023-11-17T17:32:48.618093Z","shell.execute_reply.started":"2023-11-17T17:32:37.826259Z","shell.execute_reply":"2023-11-17T17:32:48.616981Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport gc\nimport copy\nimport random\nimport pickle\nimport zipfile\nimport multiprocessing\nimport math\nfrom functools import partial\nfrom contextlib import contextmanager\nfrom pathlib import Path\nimport glob\nfrom filelock import FileLock\nfrom itertools import chain\n\n\nimport torch\n# import torch_xla\n# import torch_xla.core.xla_model as xm\n\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch import nn, einsum\n\nfrom einops import rearrange, repeat\nfrom einops_exts import repeat_many\nfrom einops.layers.torch import Rearrange\n\nfrom memorizing_transformers_pytorch.knn_memory import KNNMemoryList, DEFAULT_KNN_MEMORY_MEMMAP_DIRECTORY\n\nfrom __future__ import print_function\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom sklearn.model_selection import train_test_split\nfrom torch.optim.lr_scheduler import StepLR\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import datasets\nimport torchvision.transforms.v2 as transforms\n\n#from tqdm.notebook import tqdm\nfrom tqdm.auto import tqdm\n\n\ntorch.set_default_dtype(torch.float32)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.620974Z","iopub.execute_input":"2023-11-17T17:32:48.621296Z","iopub.status.idle":"2023-11-17T17:32:48.631430Z","shell.execute_reply.started":"2023-11-17T17:32:48.621266Z","shell.execute_reply":"2023-11-17T17:32:48.630517Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import torch_xla\n# import torch_xla.core.xla_model as xm\n\n# device = xm.xla_device()\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.632591Z","iopub.execute_input":"2023-11-17T17:32:48.632861Z","iopub.status.idle":"2023-11-17T17:32:48.644762Z","shell.execute_reply.started":"2023-11-17T17:32:48.632839Z","shell.execute_reply":"2023-11-17T17:32:48.643921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# helper functions","metadata":{}},{"cell_type":"code","source":"def pair(t):\n    return t if isinstance(t, tuple) else (t, t)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.645835Z","iopub.execute_input":"2023-11-17T17:32:48.646121Z","iopub.status.idle":"2023-11-17T17:32:48.656229Z","shell.execute_reply.started":"2023-11-17T17:32:48.646098Z","shell.execute_reply":"2023-11-17T17:32:48.655344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def identity(t):\n    return t\n\ndef exists(val):\n    return val is not None\n\ndef unique(arr):\n    return list({el: True for el in arr}.keys())\n\ndef default(val, d):\n    return val if exists(val) else d\n\ndef cast_tuple(val, length = 1):\n    return val if isinstance(val, tuple) else ((val,) * length)\n\ndef l2norm(t):\n    return F.normalize(t, dim = -1)\n\ndef stable_softmax(t, dim = -1):\n    t = t - t.amax(dim = dim, keepdim = True).detach()\n    return F.softmax(t, dim = dim)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.657169Z","iopub.execute_input":"2023-11-17T17:32:48.657399Z","iopub.status.idle":"2023-11-17T17:32:48.667276Z","shell.execute_reply.started":"2023-11-17T17:32:48.657379Z","shell.execute_reply":"2023-11-17T17:32:48.666365Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# helper classes","metadata":{}},{"cell_type":"code","source":"def posemb_sincos_2d(patches, temperature = 10000, dtype = torch.float32):\n    _, h, w, dim, device, dtype = *patches.shape, patches.device, patches.dtype\n\n    y, x = torch.meshgrid(torch.arange(h, device = device), torch.arange(w, device = device), indexing = 'ij')\n    assert (dim % 4) == 0, 'feature dimension must be multiple of 4 for sincos emb'\n    omega = torch.arange(dim // 4, device = device) / (dim // 4 - 1)\n    omega = 1. / (temperature ** omega)\n\n    y = y.flatten()[:, None] * omega[None, :]\n    x = x.flatten()[:, None] * omega[None, :] \n    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim = 1)\n    return pe.type(dtype)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.668475Z","iopub.execute_input":"2023-11-17T17:32:48.668829Z","iopub.status.idle":"2023-11-17T17:32:48.678003Z","shell.execute_reply.started":"2023-11-17T17:32:48.668805Z","shell.execute_reply":"2023-11-17T17:32:48.677114Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class PreNormResidual(nn.Module):\n    def __init__(self, dim, fn):\n        super().__init__()\n        self.fn = fn\n        self.norm = nn.LayerNorm(dim)\n\n    def forward(self, x, **kwargs):\n        out = self.fn(self.norm(x), **kwargs)\n\n        if not isinstance(out, tuple):\n            return out + x\n\n        head, *tail = out\n        return (head + x, *tail)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.679156Z","iopub.execute_input":"2023-11-17T17:32:48.679523Z","iopub.status.idle":"2023-11-17T17:32:48.691834Z","shell.execute_reply.started":"2023-11-17T17:32:48.679493Z","shell.execute_reply":"2023-11-17T17:32:48.690928Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# t5 relative positional bias","metadata":{}},{"cell_type":"code","source":"class T5RelativePositionBias(nn.Module):\n    def __init__(\n        self,\n        scale,\n        num_buckets = 32,\n        max_distance = 128,\n        heads = 8\n    ):\n        super().__init__()\n        self.scale = scale\n        self.num_buckets = num_buckets\n        self.max_distance = max_distance\n        self.relative_attention_bias = nn.Embedding(num_buckets, heads)\n\n    @staticmethod\n    def _relative_position_bucket(\n        relative_position,\n        num_buckets = 32,\n        max_distance = 128\n    ):\n        n = -relative_position\n        n = torch.max(n, torch.zeros_like(n))\n\n        max_exact = num_buckets // 2\n        is_small = n < max_exact\n\n        val_if_large = max_exact + (torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)).long()\n        val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1))\n        return torch.where(is_small, n, val_if_large)\n\n    def forward(self, i, j, *, device):\n        q_pos = torch.arange(i, dtype = torch.long, device = device)\n        k_pos = torch.arange(j, dtype = torch.long, device = device)\n        rel_pos = rearrange(k_pos, 'j -> 1 j') - rearrange(q_pos, 'i -> i 1')\n        rp_bucket = self._relative_position_bucket(rel_pos, num_buckets = self.num_buckets, max_distance = self.max_distance)\n        values = self.relative_attention_bias(rp_bucket)\n        bias = rearrange(values, 'i j h -> () h i j')\n        return bias * self.scale","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.693000Z","iopub.execute_input":"2023-11-17T17:32:48.693276Z","iopub.status.idle":"2023-11-17T17:32:48.703942Z","shell.execute_reply.started":"2023-11-17T17:32:48.693254Z","shell.execute_reply":"2023-11-17T17:32:48.703001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# feedforward","metadata":{}},{"cell_type":"code","source":"class FeedForward(nn.Module):\n    def __init__(self, dim, mult = 4, dropout = 0.):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.LayerNorm(dim),\n            nn.Linear(dim, dim * mult),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(dim * mult, dim),\n            nn.Dropout(dropout)\n        )\n    def forward(self, x):\n        return self.net(x)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.704940Z","iopub.execute_input":"2023-11-17T17:32:48.705187Z","iopub.status.idle":"2023-11-17T17:32:48.717030Z","shell.execute_reply.started":"2023-11-17T17:32:48.705166Z","shell.execute_reply":"2023-11-17T17:32:48.716200Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Attention","metadata":{}},{"cell_type":"code","source":"class Attention(nn.Module):\n    def __init__(\n        self,\n        *,\n        dim,\n        heads = 8,\n        dim_head = 64,\n        dropout = 0.,\n        xl_max_memories = 0.,\n    ):\n        super().__init__()\n        \n        inner_dim = heads * dim_head\n        self.heads = heads\n        self.scale = dim_head ** -0.5\n        self.xl_max_memories = xl_max_memories\n\n        self.norm = nn.LayerNorm(dim) #vit\n        self.attend = nn.Softmax(dim = -1) #vit\n\n        self.dropout = nn.Dropout(dropout)\n\n        self.to_q = nn.Linear(dim, inner_dim, bias = False)\n        self.to_kv = nn.Linear(dim, dim_head * 2, bias = False)\n        self.to_out = nn.Linear(inner_dim, dim)\n\n    def forward(self, x, *, xl_memory = None, rel_pos_bias = None):\n        \n        x = self.norm(x)\n        \n        h, device = self.heads, x.device\n        q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim = -1))\n\n        q = rearrange(q, 'b n (h d) -> b h n d', h = h)\n\n        q = q * self.scale\n\n        if exists(xl_memory):\n            k_xl_mem, v_xl_mem = xl_memory.unbind(dim = -2)\n            k = torch.cat((k_xl_mem, k), dim = -2)\n            v = torch.cat((v_xl_mem, v), dim = -2)\n\n        sim = einsum('b h i d, b j d -> b h i j', q, k)\n        i, j = sim.shape[-2:]\n\n        if exists(rel_pos_bias):\n            sim = rel_pos_bias[..., -i:, -j:] + sim\n\n        causal_mask = torch.ones((i, j), dtype = torch.bool, device = device).triu(j - i + 1)\n        sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)\n\n        attn = sim.softmax(dim = -1)\n        attn = self.dropout(attn)\n\n        out = einsum('b h i j, b j d -> b h i d', attn, v)\n        out = rearrange(out, 'b h n d -> b n (h d)')\n\n        # new xl memories\n\n        new_kv_memories = torch.stack((k, v), dim = -2).detach()\n\n        if self.xl_max_memories > 0:\n            new_xl_kv_memories = new_kv_memories[:, -self.xl_max_memories:]\n        else:\n            new_xl_kv_memories = None\n\n        return self.to_out(out), new_xl_kv_memories","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.718077Z","iopub.execute_input":"2023-11-17T17:32:48.718368Z","iopub.status.idle":"2023-11-17T17:32:48.732507Z","shell.execute_reply.started":"2023-11-17T17:32:48.718345Z","shell.execute_reply":"2023-11-17T17:32:48.731249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# approximate nearest neighbor attention","metadata":{}},{"cell_type":"code","source":"class KNNAttention(nn.Module):\n    def __init__(\n        self,\n        *,\n        dim,\n        heads = 8,\n        dim_head = 64,\n        dropout = 0.,\n        num_retrieved_memories = 32,\n        xl_max_memories = 0.,\n        attn_scale_init = 20,\n        gate_output = False\n    ):\n        super().__init__()\n        self.heads = heads\n        self.scale = nn.Parameter(torch.ones(heads, 1, 1) * math.log(attn_scale_init))\n\n        inner_dim = heads * dim_head\n        self.xl_max_memories = xl_max_memories\n\n        self.num_retrieved_memories = num_retrieved_memories\n\n        self.dropout = nn.Dropout(dropout)\n        self.knn_mem_dropout = nn.Dropout(dropout)\n\n        self.to_q = nn.Linear(dim, inner_dim, bias = False)\n        self.to_kv = nn.Linear(dim, dim_head * 2, bias = False)\n        self.to_out = nn.Linear(inner_dim, dim, bias = False)\n\n        self.output_gate = nn.Parameter(torch.zeros(1)) if gate_output else None\n\n    def forward(\n        self,\n        x,\n        *,\n        knn_memory,\n        xl_memory = None,\n        add_knn_memory = True,\n        rel_pos_bias = None\n    ):\n        b, n, h, device = *x.shape[:2], self.heads, x.device\n        q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim = -1))\n\n        q = rearrange(q, 'b n (h d) -> b h n d', h = h)\n\n        # in paper, they showed normalizing of keys led to more stable training\n        # we'll just go with full cosine sim attention https://arxiv.org/abs/2010.04245\n\n        q, k = map(l2norm, (q, k))\n\n        # handle xl memory\n\n        if exists(xl_memory):\n            k_xl_mem, v_xl_mem = xl_memory.unbind(dim = -2)\n            k = torch.cat((k_xl_mem, k), dim = -2)\n            v = torch.cat((v_xl_mem, v), dim = -2)\n\n        # calculate local attention\n\n        scale = self.scale.exp()\n\n        sim = einsum('b h i d, b j d -> b h i j', q, k) * scale\n        i, j = sim.shape[-2:]\n\n        if exists(rel_pos_bias):\n            sim = rel_pos_bias[..., -i:, -j:] + sim\n\n        mask_value = -torch.finfo(sim.dtype).max\n\n        causal_mask = torch.ones((i, j), dtype = torch.bool, device = device).triu(j - i + 1)\n        sim = sim.masked_fill(causal_mask, mask_value)\n\n        # calculate knn attention over memory, if index is passed in\n\n        mem_kv, mem_mask = knn_memory.search(q, self.num_retrieved_memories)\n        mem_k, mem_v = mem_kv.unbind(dim = -2)\n\n        sim_mem = einsum('b h i d, b h i j d -> b h i j', q, mem_k) * scale\n        sim_mem = sim_mem.masked_fill(~mem_mask, mask_value)\n\n        # calculate new XL memories, as well as memories to be discarded\n\n        new_kv_memories = torch.stack((k, v), dim = -2).detach()\n\n        if self.xl_max_memories > 0:\n            new_kv_memories_discarded, new_xl_kv_memories = new_kv_memories[:, :-self.xl_max_memories], new_kv_memories[:, -self.xl_max_memories:]\n        else:\n            new_kv_memories_discarded, new_xl_kv_memories = new_kv_memories, None\n\n        # add memories to be discarded into KNN memory\n\n        if add_knn_memory and new_kv_memories_discarded.numel() > 0:\n            knn_memory.add(new_kv_memories_discarded)\n\n        # attention (combining local and distant)\n\n        sim = torch.cat((sim_mem, sim), dim = -1)\n        attn = sim.softmax(dim = -1)\n        attn = self.dropout(attn)\n\n        local_attn, mem_attn = attn[..., self.num_retrieved_memories:], attn[..., :self.num_retrieved_memories]\n        local_out = einsum('b h i j, b j d -> b h i d', local_attn, v)\n        mem_out = einsum('b h i j, b h i j d -> b h i d', mem_attn, mem_v)\n\n        out = local_out + mem_out\n\n        # combine heads and project out\n\n        out = rearrange(out, 'b h n d -> b n (h d)')\n        out = self.to_out(out)\n\n        # use flamingo styled gating of output, so that memorizing transformers can be gated into an existing LLM\n        # preparation to add this to block-recurrent-transformer-pytorch, for the pinnacle of long context attention network\n\n        if exists(self.output_gate):\n            out = out * self.output_gate.tanh()\n\n        return out, new_xl_kv_memories","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.737043Z","iopub.execute_input":"2023-11-17T17:32:48.737303Z","iopub.status.idle":"2023-11-17T17:32:48.756929Z","shell.execute_reply.started":"2023-11-17T17:32:48.737282Z","shell.execute_reply":"2023-11-17T17:32:48.756185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# main class","metadata":{}},{"cell_type":"code","source":"class SPT(nn.Module):\n    def __init__(self, *, dim, patch_size, channels = 3):\n        super().__init__()\n        patch_dim = patch_size * patch_size * 5 * channels\n\n        self.to_patch_tokens = nn.Sequential(\n            Rearrange('b c (h p1) (w p2) -> b h w (p1 p2 c)', p1 = patch_size, p2 = patch_size),\n            nn.LayerNorm(patch_dim),\n            nn.Linear(patch_dim, dim)\n        )\n        \n        # self.to_patch_embedding = nn.Sequential(\n        #     Rearrange('b c (h p1) (w p2) -> b h w (p1 p2 c)', p1 = patch_height, p2 = patch_width),\n        #     nn.LayerNorm(patch_dim),\n        #     nn.Linear(patch_dim, dim),\n        #     nn.LayerNorm(dim),\n        # )\n\n    def forward(self, x):\n        shifts = ((1, -1, 0, 0), (-1, 1, 0, 0), (0, 0, 1, -1), (0, 0, -1, 1))\n        shifted_x = list(map(lambda shift: F.pad(x, shift), shifts))\n        x_with_shifts = torch.cat((x, *shifted_x), dim = 1)\n        return self.to_patch_tokens(x_with_shifts)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.757846Z","iopub.execute_input":"2023-11-17T17:32:48.758135Z","iopub.status.idle":"2023-11-17T17:32:48.772078Z","shell.execute_reply.started":"2023-11-17T17:32:48.758112Z","shell.execute_reply":"2023-11-17T17:32:48.771329Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MemorizingSPTViT(nn.Module):\n    def __init__(\n        self,\n        *,\n\n        image_size,\n        patch_size,\n        num_classes,\n\n        dim,\n        depth,\n\n        channels = 3,\n        \n        dim_head = 64,\n        heads = 8,\n        knn_attn_heads = None,\n        attn_dropout = 0.,\n        ff_mult = 4,\n        ff_dropout = 0.,\n        memorizing_layers = None,\n        max_knn_memories = 250000,\n        num_retrieved_memories = 32,\n        clear_memories_on_sos_token_id = None,\n        clear_memories_on_eos_token_id = None,\n        knn_memories_directory = DEFAULT_KNN_MEMORY_MEMMAP_DIRECTORY,\n        shift_knn_memories_down = 0.,\n        xl_max_memories = 0,\n        xl_memory_layers = None,\n        shift_xl_memories_down = 0.,\n        knn_memory_multiprocessing = False\n    ):\n        super().__init__()\n\n        image_height, image_width = pair(image_size)\n        patch_height, patch_width = pair(patch_size)\n\n        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'\n\n        num_patches = (image_height // patch_height) * (image_width // patch_width)\n        patch_dim = channels * patch_height * patch_width\n\n        self.to_patch_embedding = SPT(dim = dim, patch_size = patch_size, channels = channels)\n        \n        # self.to_patch_embedding = nn.Sequential(\n        #     Rearrange('b c (h p1) (w p2) -> b h w (p1 p2 c)', p1 = patch_height, p2 = patch_width),\n        #     nn.LayerNorm(patch_dim),\n        #     nn.Linear(patch_dim, dim),\n        #     nn.LayerNorm(dim),\n        # )\n\n        block_wrapper = partial(PreNormResidual, dim)\n        valid_layers = set(range(1, depth + 1))\n\n        memorizing_layers = default(memorizing_layers, (depth // 2,)) # default KNN attention layer to midpoint of transformer\n        memorizing_layers = cast_tuple(memorizing_layers)\n        memorizing_layers = tuple(filter(lambda i: i in valid_layers, memorizing_layers))\n\n        self.dim_head = dim_head\n\n        knn_attn_heads = default(knn_attn_heads, heads)\n\n        # xl memory hyperparameter\n\n        if xl_max_memories > 0:\n            xl_memory_layers = default(xl_memory_layers, tuple(range(1, depth + 1)))\n            xl_memory_layers = unique(xl_memory_layers)\n            self.xl_memory_layers = tuple(filter(lambda i: i in valid_layers, xl_memory_layers))            \n            self.num_xl_memory_layers = len(self.xl_memory_layers)\n        else:\n            self.xl_memory_layers = tuple()\n            self.num_xl_memory_layers = 0\n\n        # knn memory hyperparameters\n\n        self.max_knn_memories = max_knn_memories\n        self.knn_memories_directory = knn_memories_directory\n        self.memorizing_layers = unique(memorizing_layers)\n        self.num_memory_layers = len(memorizing_layers)\n\n        self.clear_memories_on_sos_token_id = clear_memories_on_sos_token_id\n        self.clear_memories_on_eos_token_id = clear_memories_on_eos_token_id\n\n        # relative positional bias\n\n        self.rel_pos_bias = T5RelativePositionBias(scale = dim_head ** 0.5, heads = heads)\n        self.knn_rel_pos_bias = T5RelativePositionBias(scale = dim_head ** 0.5, heads = heads)\n\n        # layers\n\n        self.layers = nn.ModuleList([])\n        for idx in range(depth):\n            layer_num = idx + 1\n\n            use_xl_memories = layer_num in self.xl_memory_layers\n            use_knn_attention = layer_num in memorizing_layers\n            xl_max_memories_layer = 0 if not use_xl_memories else xl_max_memories\n\n            if use_knn_attention:\n                attn = KNNAttention(dim = dim, dim_head = dim_head, heads = knn_attn_heads, dropout = attn_dropout, num_retrieved_memories = num_retrieved_memories, xl_max_memories = xl_max_memories_layer)\n            else:\n                attn = Attention(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout, xl_max_memories = xl_max_memories_layer)\n\n            self.layers.append(nn.ModuleList([\n                block_wrapper(attn),\n                block_wrapper(FeedForward(dim = dim, mult = ff_mult, dropout = ff_dropout)),\n            ]))\n\n        # memory layer shifting\n        # from a little known paper https://arxiv.org/abs/2012.15688\n\n        self.shift_knn_memories_down = shift_knn_memories_down\n        self.shift_xl_memories_down = shift_xl_memories_down\n\n        # to logits\n        self.to_latent = nn.Identity()\n        self.to_logits = nn.Sequential(\n            nn.LayerNorm(dim),\n            nn.Linear(dim, num_classes)\n        )\n\n        # knn memories init\n\n        self.knn_mem_kwargs = dict(\n            dim = self.dim_head,\n            max_memories = self.max_knn_memories,\n            multiprocessing = knn_memory_multiprocessing\n        )\n\n    def create_knn_memories(\n        self,\n        *,\n        batch_size\n    ):\n        return KNNMemoryList.create_memories(\n            batch_size = batch_size,\n            num_memory_layers = self.num_memory_layers,\n            memories_directory = self.knn_memories_directory,\n        )(**self.knn_mem_kwargs)\n\n    @contextmanager\n    def knn_memories_context(\n        self,\n        **kwargs\n    ):\n        knn_dir = Path(self.knn_memories_directory)\n        knn_dir.mkdir(exist_ok = True, parents = True)\n        lock = FileLock(str(knn_dir / 'mutex'))\n\n        with lock:\n            knn_memories = self.create_knn_memories(**kwargs)\n            yield knn_memories\n            knn_memories.cleanup()\n\n    def clear_memory(self, x, token_id, knn_memories):\n        \"\"\" clears the KNN memories based on if the batch row contains the specified token id \"\"\"\n        \"\"\" for auto-clearing KNN memories based on start and end of strings \"\"\"\n\n        clear_memory = (x == token_id).any(dim = -1)\n        batch_indices, _ = clear_memory.nonzero(as_tuple = True)\n        batch_indices_to_clear = batch_indices.tolist()\n\n        if len(batch_indices_to_clear) == 0:\n            return\n\n        knn_memories.clear_memory(batch_indices_to_clear)\n\n    def forward(\n        self,\n        img,\n        knn_memories,\n        xl_memories    = None,\n        add_knn_memory = True\n    ):\n        \n        x = self.to_patch_embedding(img)\n        \n        *_, h, w, dtype = *x.shape, x.dtype\n\n        pe = posemb_sincos_2d(x)\n        x = rearrange(x, 'b ... d -> b (...) d') + pe\n        \n        batch_size, seq_len, *_, device = *x.shape, x.device\n\n        # validate KNN memories to have enough indices for batch size\n\n        assert all([memory.num_indices == batch_size for memory in knn_memories]), f'you passed in an input with batch size {batch_size} but your memories were not instantiated with that number of KNN indices'\n\n        # if KNN memories are passed in, and researcher wants memories auto-cleared on <sos> token detection\n        # do the appropriate logic\n\n        if exists(self.clear_memories_on_sos_token_id):\n            self.clear_memory(x, self.clear_memories_on_sos_token_id, knn_memories)\n\n        # handle XL memories\n\n        xl_memories = default(xl_memories, (None,) * self.num_xl_memory_layers)\n        assert len(xl_memories) == self.num_xl_memory_layers\n        has_xl_memories = len(xl_memories) > 0\n\n        # shifting memories a number of layers down, little known technique shown to enhance memories from Ernie-Doc paper\n\n        if len(knn_memories) > 0 and self.shift_knn_memories_down > 0:\n            knn_memories = [*knn_memories[self.shift_knn_memories_down:], *knn_memories[:self.shift_knn_memories_down]]\n\n        if len(xl_memories) > 0 and self.shift_xl_memories_down > 0:\n            xl_memories = [*xl_memories[self.shift_xl_memories_down:], *xl_memories[:self.shift_xl_memories_down]]\n\n        # iterate through the memories in order of the ascending layers that contain KNNAttention\n\n        xl_memories_iter = iter(xl_memories)\n        knn_memories_iter = iter(knn_memories)\n\n        # positional bias\n\n        max_context_len = max([seq_len, *map(lambda t: (t.shape[-3] if exists(t) else 0) + seq_len, xl_memories)])\n\n        rel_pos_bias = self.rel_pos_bias(seq_len, max_context_len, device = device)\n        knn_rel_pos_bias = self.knn_rel_pos_bias(seq_len, max_context_len, device = device)\n\n        # keep track of new xl memories\n\n        new_xl_memories = [] if has_xl_memories else None\n\n        # go through all layers\n\n        for ind, (attn, ff) in enumerate(self.layers):\n            layer_num = ind + 1\n\n            is_memorizing_layer = layer_num in self.memorizing_layers\n            is_xl_memory_layer = layer_num in self.xl_memory_layers\n\n            attn_kwargs = dict(rel_pos_bias = rel_pos_bias if not is_memorizing_layer else knn_rel_pos_bias)\n\n            if is_memorizing_layer:\n                attn_kwargs = {**attn_kwargs, 'knn_memory': next(knn_memories_iter), 'add_knn_memory': add_knn_memory}\n\n            if is_xl_memory_layer:\n                attn_kwargs = {**attn_kwargs, 'xl_memory': next(xl_memories_iter)}\n\n            # attention\n\n            x, xl_mem = attn(x, **attn_kwargs)\n\n            # add new XL memories if needed\n\n            if exists(xl_mem):\n                new_xl_memories.append(xl_mem)\n\n            # feedforward\n\n            x = ff(x)\n\n        x = x.mean(dim = 1)\n        #x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]\n        #x = x[:, 0]\n        x = self.to_latent(x)\n        logits = self.to_logits(x)\n\n        # auto-clear KNN memories on end of token\n\n        if exists(self.clear_memories_on_eos_token_id):\n            self.clear_memory(x, self.clear_memories_on_eos_token_id, knn_memories)\n\n        if exists(new_xl_memories):\n            return logits, new_xl_memories\n\n        return logits","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.773337Z","iopub.execute_input":"2023-11-17T17:32:48.773629Z","iopub.status.idle":"2023-11-17T17:32:48.807134Z","shell.execute_reply.started":"2023-11-17T17:32:48.773606Z","shell.execute_reply":"2023-11-17T17:32:48.806224Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"# from __future__ import print_function\n\n# import glob\n# from itertools import chain\n# import os\n# import random\n# import zipfile\n\n# import matplotlib.pyplot as plt\n# import numpy as np\n# import pandas as pd\n# import torch\n# import torch.nn as nn\n# import torch.nn.functional as F\n# import torch.optim as optim\n# from PIL import Image\n# from sklearn.model_selection import train_test_split\n# from torch.optim.lr_scheduler import StepLR\n# from torch.utils.data import DataLoader, Dataset\n# from torchvision import datasets, transforms\n# from tqdm.notebook import tqdm","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.808560Z","iopub.execute_input":"2023-11-17T17:32:48.808827Z","iopub.status.idle":"2023-11-17T17:32:48.820831Z","shell.execute_reply.started":"2023-11-17T17:32:48.808805Z","shell.execute_reply":"2023-11-17T17:32:48.820000Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"Torch: {torch.__version__}\")","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.821886Z","iopub.execute_input":"2023-11-17T17:32:48.822148Z","iopub.status.idle":"2023-11-17T17:32:48.835594Z","shell.execute_reply.started":"2023-11-17T17:32:48.822126Z","shell.execute_reply":"2023-11-17T17:32:48.834770Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training settings","metadata":{}},{"cell_type":"code","source":"batch_size = 256\n# batch_size = 16\nsizes      = 224\nepochs     = 50\nlr         = 3e-5\ngamma      = 0.9\nseed       = 42","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.836525Z","iopub.execute_input":"2023-11-17T17:32:48.836797Z","iopub.status.idle":"2023-11-17T17:32:48.846621Z","shell.execute_reply.started":"2023-11-17T17:32:48.836766Z","shell.execute_reply":"2023-11-17T17:32:48.845742Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(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    torch.backends.cudnn.deterministic = True\n\nseed_everything(seed)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.847703Z","iopub.execute_input":"2023-11-17T17:32:48.848040Z","iopub.status.idle":"2023-11-17T17:32:48.858077Z","shell.execute_reply.started":"2023-11-17T17:32:48.848016Z","shell.execute_reply":"2023-11-17T17:32:48.857317Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# device = \"cuda\" if torch.cuda.is_available() else \"cpu\"","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.859129Z","iopub.execute_input":"2023-11-17T17:32:48.859391Z","iopub.status.idle":"2023-11-17T17:32:48.867832Z","shell.execute_reply.started":"2023-11-17T17:32:48.859364Z","shell.execute_reply":"2023-11-17T17:32:48.867038Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Data","metadata":{}},{"cell_type":"code","source":"train_dir = 'train'\n\nwith zipfile.ZipFile('/kaggle/input/dogs-vs-cats-redux-kernels-edition/train.zip', 'r') as train_zip:\n    train_zip.extractall('')\n\ntrain_list = glob.glob(os.path.join(train_dir,'*.jpg'))\nprint(f\"Train Data: {len(train_list)}\")\n\nbase_dir = '/kaggle/working'\ntrain_dir = os.path.join(base_dir, 'train')","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:32:48.868735Z","iopub.execute_input":"2023-11-17T17:32:48.868983Z","iopub.status.idle":"2023-11-17T17:33:00.672803Z","shell.execute_reply.started":"2023-11-17T17:32:48.868962Z","shell.execute_reply":"2023-11-17T17:33:00.671812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dir = 'test'\n\nwith zipfile.ZipFile('/kaggle/input/dogs-vs-cats-redux-kernels-edition/test.zip', 'r') as test_zip:\n    test_zip.extractall('')\n\ntest_list = glob.glob(os.path.join(test_dir,'*.jpg'))\nprint(f\"Test Data: {len(test_list)}\")\n\nbase_dir = '/kaggle/working'\ntest_dir = os.path.join(base_dir, 'test')","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:00.673899Z","iopub.execute_input":"2023-11-17T17:33:00.674188Z","iopub.status.idle":"2023-11-17T17:33:06.785364Z","shell.execute_reply.started":"2023-11-17T17:33:00.674164Z","shell.execute_reply":"2023-11-17T17:33:06.784480Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"labels = [path.split('/')[-1].split('.')[0] for path in train_list]","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:06.786514Z","iopub.execute_input":"2023-11-17T17:33:06.786832Z","iopub.status.idle":"2023-11-17T17:33:06.804767Z","shell.execute_reply.started":"2023-11-17T17:33:06.786808Z","shell.execute_reply":"2023-11-17T17:33:06.803943Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Random Plots","metadata":{}},{"cell_type":"code","source":"random_idx = np.random.randint(1, len(train_list), size=9)\nfig, axes = plt.subplots(3, 3, figsize=(16, 12))\n\nfor idx, ax in enumerate(axes.ravel()):\n    img = Image.open(train_list[idx])\n    ax.set_title(labels[idx])\n    ax.imshow(img)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:06.805854Z","iopub.execute_input":"2023-11-17T17:33:06.806103Z","iopub.status.idle":"2023-11-17T17:33:09.356114Z","shell.execute_reply.started":"2023-11-17T17:33:06.806082Z","shell.execute_reply":"2023-11-17T17:33:09.355117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Split","metadata":{}},{"cell_type":"code","source":"train_list, valid_list = train_test_split(train_list, \n                                          test_size=0.2,\n                                          stratify=labels,\n                                          random_state=seed)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:09.357582Z","iopub.execute_input":"2023-11-17T17:33:09.357921Z","iopub.status.idle":"2023-11-17T17:33:09.391952Z","shell.execute_reply.started":"2023-11-17T17:33:09.357893Z","shell.execute_reply":"2023-11-17T17:33:09.391019Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"Train Data: {len(train_list)}\")\nprint(f\"Validation Data: {len(valid_list)}\")\nprint(f\"Test Data: {len(test_list)}\")","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:09.393170Z","iopub.execute_input":"2023-11-17T17:33:09.393447Z","iopub.status.idle":"2023-11-17T17:33:09.398699Z","shell.execute_reply.started":"2023-11-17T17:33:09.393423Z","shell.execute_reply":"2023-11-17T17:33:09.397596Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Image Augmentation","metadata":{}},{"cell_type":"code","source":"val_transforms = transforms.Compose(\n    [\n        transforms.Resize(sizes),\n        transforms.CenterCrop(sizes),\n        transforms.ToImageTensor(),\n        transforms.ConvertImageDtype(torch.float),\n        transforms.Normalize(\n            (0.485, 0.456, 0.406),\n            (0.229, 0.224, 0.225)\n        )\n    ]\n)\n\ntrain_transforms = transforms.Compose(\n    [\n        transforms.RandomResizedCrop(sizes),\n        transforms.RandomChoice([\n            transforms.GaussianBlur(kernel_size=(11,11), sigma=(0.05, 2.)),\n            transforms.RandomAffine(0, shear=10,scale=(0.8, 1.2)),\n            transforms.RandomPhotometricDistort(),\n            transforms.RandomRotation(45),\n            transforms.RandomHorizontalFlip()\n        ]),\n        transforms.ToImageTensor(),\n        transforms.ConvertImageDtype(torch.float),\n        transforms.Normalize(\n            (0.485, 0.456, 0.406),\n            (0.229, 0.224, 0.225)\n        ),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n        transforms.RandomErasing(),\n    ]\n)\n\ntest_transforms = transforms.Compose(\n    [\n        transforms.Resize(sizes),\n        transforms.CenterCrop(sizes),\n        transforms.ToImageTensor(),\n        transforms.ConvertImageDtype(torch.float)\n    ]\n)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:09.400075Z","iopub.execute_input":"2023-11-17T17:33:09.400464Z","iopub.status.idle":"2023-11-17T17:33:09.470885Z","shell.execute_reply.started":"2023-11-17T17:33:09.400431Z","shell.execute_reply":"2023-11-17T17:33:09.470114Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Datasets","metadata":{}},{"cell_type":"code","source":"class CatsDogsDataset(Dataset):\n    def __init__(self, file_list, transform=None):\n        self.file_list = file_list\n        self.transform = transform\n\n    def __len__(self):\n        self.filelength = len(self.file_list)\n        return self.filelength\n\n    def __getitem__(self, idx):\n        img_path = self.file_list[idx]\n        img = Image.open(img_path)\n        img_transformed = self.transform(img)\n\n        label = img_path.split(\"/\")[-1].split(\".\")[0]\n        label = 1 if label == \"dog\" else 0\n\n        return img_transformed, label\n","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:09.471839Z","iopub.execute_input":"2023-11-17T17:33:09.472084Z","iopub.status.idle":"2023-11-17T17:33:09.478382Z","shell.execute_reply.started":"2023-11-17T17:33:09.472063Z","shell.execute_reply":"2023-11-17T17:33:09.477516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_data = CatsDogsDataset(train_list, transform=train_transforms)\nvalid_data = CatsDogsDataset(valid_list, transform=test_transforms)\ntest_data = CatsDogsDataset(test_list, transform=test_transforms)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:09.479691Z","iopub.execute_input":"2023-11-17T17:33:09.480019Z","iopub.status.idle":"2023-11-17T17:33:09.491853Z","shell.execute_reply.started":"2023-11-17T17:33:09.479989Z","shell.execute_reply":"2023-11-17T17:33:09.491096Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_loader = DataLoader(dataset = train_data, batch_size=batch_size, shuffle=True )\nvalid_loader = DataLoader(dataset = valid_data, batch_size=batch_size, shuffle=True )\ntest_loader  = DataLoader(dataset = test_data,  batch_size=batch_size, shuffle=True )","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:09.493020Z","iopub.execute_input":"2023-11-17T17:33:09.493612Z","iopub.status.idle":"2023-11-17T17:33:09.502399Z","shell.execute_reply.started":"2023-11-17T17:33:09.493581Z","shell.execute_reply":"2023-11-17T17:33:09.501457Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(len(train_data), len(train_loader))\nprint(len(valid_data), len(valid_loader))","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:09.503405Z","iopub.execute_input":"2023-11-17T17:33:09.503669Z","iopub.status.idle":"2023-11-17T17:33:09.512801Z","shell.execute_reply.started":"2023-11-17T17:33:09.503647Z","shell.execute_reply":"2023-11-17T17:33:09.511875Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train(\n    model,\n    criterion,\n    optimizer,\n    scheduler,\n    loader,\n    metric,\n    epochs,\n    phase_train = 'train',\n    max_grad_clip_norm = 0.5\n):\n    for epoch in range(int(len(metric)/2), int(len(metric)/2)+epochs):\n        for phase in loader:\n            \n            epoch_loss = 0\n            epoch_accuracy = 0\n\n            if phase == phase_train:\n                model.train()  # Set model to training mode\n            else:\n                model.eval()   # Set model to evaluate mode\n\n            for data, label in loader[phase]:\n                data = data.to(device, dtype=torch.float)\n                label = label.to(device)\n                \n                if phase == phase_train:\n                    optimizer.zero_grad()\n                \n                with torch.set_grad_enabled(phase == phase_train):\n                    with model.knn_memories_context(batch_size = data.size()[0]) as knn_memories:\n                        output = model(\n                            data,\n                            knn_memories = knn_memories\n                        )\n                        del data, knn_memories\n                        gc.collect()\n                    \n                    loss = criterion(output, label)\n                    if phase == phase_train:\n                        #torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_clip_norm)\n                        loss.backward()\n                        torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_clip_norm)\n                        optimizer.step()\n#                       xm.mark_step()\n#                       xm.optimizer_step(optimizer, barrier=True)\n\n                    acc = (output.argmax(dim=1) == label).float().mean()\n\n                    epoch_accuracy += float(acc) / len(loader[phase])\n                    epoch_loss += float(loss) / len(loader[phase])\n\n                    # clear memory\n                    del label, output, loss\n                    gc.collect()\n                \n            print(\n                f\"Epoch: {epoch+1} Phase: {phase} - loss: {epoch_loss:.4f} - acc: {epoch_accuracy:.4f}\"\n            )\n            metric.append({\n                \"epoch\": epoch,\n                \"phase\": phase,\n                \"loss\": epoch_loss,\n                \"accuracy\": epoch_accuracy\n            })\n            \n    return epoch","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:09.514110Z","iopub.execute_input":"2023-11-17T17:33:09.514395Z","iopub.status.idle":"2023-11-17T17:33:09.526378Z","shell.execute_reply.started":"2023-11-17T17:33:09.514373Z","shell.execute_reply":"2023-11-17T17:33:09.525407Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Test SimpleViT with Memorizing","metadata":{}},{"cell_type":"code","source":"model = MemorizingSPTViT(\n    image_size = sizes,                 # Image size. If you have rectangular images, make sure your image size is the maximum of the width and height\n    patch_size = 32,                    # Number of patches. image_size must be divisible by patch_size. The number of patches is:  n = (image_size // patch_size) ** 2 and n must be greater than 16.\n    num_classes = 2,                    # Number of classes to classify.\n    heads = 16,\n    dim = 1024,                         # Last dimension of output tensor after linear transformation nn.Linear(..., dim).\n    dim_head = 32,                      # dimension per attention head\n    depth = 8,                          # number of layers\n    \n    #pool = 'cls',\n    \n    memorizing_layers = (),             # which layers to have ANN memories\n    max_knn_memories = 512*64,          # maximum ANN memories to keep (once it hits this capacity, it will be reset for now, due to limitations in faiss' ability to remove entries)\n    num_retrieved_memories = 64,        # number of ANN memories to retrieve\n    clear_memories_on_sos_token_id = None, # clear passed in ANN memories automatically for batch indices which contain this specified SOS token id - otherwise, you can also manually iterate through the ANN memories and clear the indices before the next iteration\n \n    xl_memory_layers = None,\n    xl_max_memories = 0,\n\n    shift_knn_memories_down = 0,\n    shift_xl_memories_down = 0,\n\n    knn_attn_heads = None,\n    attn_dropout = 0.1,\n    ff_mult = 4,\n    ff_dropout = 0.1,\n    clear_memories_on_eos_token_id = None,\n    knn_memories_directory = DEFAULT_KNN_MEMORY_MEMMAP_DIRECTORY,\n    \n    knn_memory_multiprocessing = True\n\n).to(device)\n\nmodel","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:09.527657Z","iopub.execute_input":"2023-11-17T17:33:09.527979Z","iopub.status.idle":"2023-11-17T17:33:13.526058Z","shell.execute_reply.started":"2023-11-17T17:33:09.527950Z","shell.execute_reply":"2023-11-17T17:33:13.525310Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.load_state_dict(\n    torch.load(\n        \"/kaggle/input/models/pretrained-net-without-mem-1000.pt\",\n        map_location='cpu'\n    ), \n    strict=False)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:13.527255Z","iopub.execute_input":"2023-11-17T17:33:13.527624Z","iopub.status.idle":"2023-11-17T17:33:18.047214Z","shell.execute_reply.started":"2023-11-17T17:33:13.527591Z","shell.execute_reply":"2023-11-17T17:33:18.046471Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# loss function\ncriterion = nn.CrossEntropyLoss()\n# optimizer\noptimizer = optim.AdamW(model.parameters(), lr=lr)\n# scheduler\nscheduler = StepLR(optimizer, step_size=3, gamma=gamma)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:18.048351Z","iopub.execute_input":"2023-11-17T17:33:18.048724Z","iopub.status.idle":"2023-11-17T17:33:18.055016Z","shell.execute_reply.started":"2023-11-17T17:33:18.048691Z","shell.execute_reply":"2023-11-17T17:33:18.054000Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loader={\n    'train':train_loader,\n    'val':valid_loader\n}\n\n# net_without_mem=[]","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:18.056117Z","iopub.execute_input":"2023-11-17T17:33:18.056376Z","iopub.status.idle":"2023-11-17T17:33:18.069612Z","shell.execute_reply.started":"2023-11-17T17:33:18.056354Z","shell.execute_reply":"2023-11-17T17:33:18.068729Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open(\"/kaggle/input/metrics/net-without-mem-1000.pickle\", \"rb\") as handle:\n    net_without_mem = pickle.load(handle)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:18.070775Z","iopub.execute_input":"2023-11-17T17:33:18.071084Z","iopub.status.idle":"2023-11-17T17:33:18.091831Z","shell.execute_reply.started":"2023-11-17T17:33:18.071062Z","shell.execute_reply":"2023-11-17T17:33:18.091113Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for _ in range(1):\n    epoch=train(\n        model=model,\n        criterion=criterion,\n        optimizer=optimizer,\n        scheduler=scheduler,\n        loader=loader,\n        metric=net_without_mem,\n        epochs=epochs,\n    )\n    torch.save(model.state_dict(), f\"/kaggle/working/models/pretrained-net-without-mem-{epoch+1}.pt\")\n    with open(f\"/kaggle/working/metrics/net-without-mem-{epoch+1}.pickle\", \"wb\") as handle:\n        pickle.dump(net_without_mem, handle, protocol=pickle.HIGHEST_PROTOCOL)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:18.092762Z","iopub.execute_input":"2023-11-17T17:33:18.093009Z","iopub.status.idle":"2023-11-17T17:33:33.056277Z","shell.execute_reply.started":"2023-11-17T17:33:18.092988Z","shell.execute_reply":"2023-11-17T17:33:33.053655Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# with open(\"/kaggle/working/metrics/net-without-mem-2400.pickle\", \"rb\") as handle:\n#     net_without_mem = pickle.load(handle)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.057116Z","iopub.status.idle":"2023-11-17T17:33:33.057443Z","shell.execute_reply.started":"2023-11-17T17:33:33.057286Z","shell.execute_reply":"2023-11-17T17:33:33.057301Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loss_net_without_mem_val, acc_net_without_mem_val = zip(*[(k[\"loss\"], k[\"accuracy\"]) for k in net_without_mem if k[\"phase\"]=='val'])\nloss_net_without_mem_train, acc_net_without_mem_train = zip(*[(k[\"loss\"], k[\"accuracy\"]) for k in net_without_mem if k[\"phase\"]=='train'])","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.059071Z","iopub.status.idle":"2023-11-17T17:33:33.059407Z","shell.execute_reply.started":"2023-11-17T17:33:33.059249Z","shell.execute_reply":"2023-11-17T17:33:33.059264Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.title(\"Net without memory\")\n\nplt.plot(loss_net_without_mem_val, label = \"Loss Val\")\n\nplt.plot(acc_net_without_mem_val, label = \"Accuracy Val\")\n\nplt.plot(loss_net_without_mem_train, label = \"Loss Train\")\n\nplt.plot(acc_net_without_mem_train, label = \"Accuracy Train\")\n\nplt.legend()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.060924Z","iopub.status.idle":"2023-11-17T17:33:33.061364Z","shell.execute_reply.started":"2023-11-17T17:33:33.061136Z","shell.execute_reply":"2023-11-17T17:33:33.061157Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Test SPTViT Finetuning with Memorizing ","metadata":{}},{"cell_type":"code","source":"# model = MemorizingSPTViT(\n#     image_size = sizes,                 # Image size. If you have rectangular images, make sure your image size is the maximum of the width and height\n#     patch_size = 32,                    # Number of patches. image_size must be divisible by patch_size. The number of patches is:  n = (image_size // patch_size) ** 2 and n must be greater than 16.\n#     num_classes = 2,                    # Number of classes to classify.\n#     heads = 16,\n#     dim = 1024,                         # Last dimension of output tensor after linear transformation nn.Linear(..., dim).\n#     dim_head = 32,                      # dimension per attention head\n#     depth = 8,                          # number of layers\n    \n#     #pool = 'cls',\n    \n#     memorizing_layers = (1, 2, 3, 4, 5, 6, 7, 8),         # which layers to have ANN memories\n#     max_knn_memories = 512*64,          # maximum ANN memories to keep (once it hits this capacity, it will be reset for now, due to limitations in faiss' ability to remove entries)\n#     num_retrieved_memories = 64,        # number of ANN memories to retrieve\n#     clear_memories_on_sos_token_id = None, # clear passed in ANN memories automatically for batch indices which contain this specified SOS token id - otherwise, you can also manually iterate through the ANN memories and clear the indices before the next iteration\n    \n#     xl_memory_layers = None,\n#     xl_max_memories = 0,\n\n#     shift_knn_memories_down = 0,\n#     shift_xl_memories_down = 0,\n\n#     knn_attn_heads = None,\n#     attn_dropout = 0.1,\n#     ff_mult = 4,\n#     ff_dropout = 0.1,\n#     clear_memories_on_eos_token_id = None,\n#     knn_memories_directory = DEFAULT_KNN_MEMORY_MEMMAP_DIRECTORY,\n    \n#     knn_memory_multiprocessing = False\n\n# ).to(device)\n\n# model","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.062890Z","iopub.status.idle":"2023-11-17T17:33:33.063236Z","shell.execute_reply.started":"2023-11-17T17:33:33.063064Z","shell.execute_reply":"2023-11-17T17:33:33.063079Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model.load_state_dict(\n#     torch.load(\n#         \"/kaggle/input/models/pretrained-net-with-mem-104.pt\",\n#         map_location=device\n#     ),\n#     strict=False)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.064145Z","iopub.status.idle":"2023-11-17T17:33:33.064452Z","shell.execute_reply.started":"2023-11-17T17:33:33.064298Z","shell.execute_reply":"2023-11-17T17:33:33.064312Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# # loss function\n# criterion = nn.CrossEntropyLoss()\n# # optimizer\n# optimizer = optim.AdamW(model.parameters(), lr=lr)\n# # scheduler\n# scheduler = StepLR(optimizer, step_size=3, gamma=gamma)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.065415Z","iopub.status.idle":"2023-11-17T17:33:33.065874Z","shell.execute_reply.started":"2023-11-17T17:33:33.065644Z","shell.execute_reply":"2023-11-17T17:33:33.065665Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# loader={\n#     'train':train_loader,\n#     'val':valid_loader\n# }","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.067377Z","iopub.status.idle":"2023-11-17T17:33:33.067749Z","shell.execute_reply.started":"2023-11-17T17:33:33.067569Z","shell.execute_reply":"2023-11-17T17:33:33.067586Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# #net_with_mem = copy.deepcopy(net_without_mem[:200])\n# with open('/kaggle/input/metrics/net-with-mem-104.pickle', 'rb') as handle:\n#     net_with_mem = pickle.load(handle)\n# epochs       = 1","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.069109Z","iopub.status.idle":"2023-11-17T17:33:33.069440Z","shell.execute_reply.started":"2023-11-17T17:33:33.069280Z","shell.execute_reply":"2023-11-17T17:33:33.069295Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# for _ in range(5):\n#     epoch=train(\n#         model=model,\n#         criterion=criterion,\n#         optimizer=optimizer,\n#         scheduler=scheduler,\n#         loader=loader,\n#         metric=net_with_mem,\n#         epochs=epochs,\n#     )\n#     torch.save(model.state_dict(), f'/kaggle/working/models/pretrained-net-with-mem-{epoch+1}.pt')\n#     with open(f'/kaggle/working/metrics/net-with-mem-{epoch+1}.pickle', 'wb') as handle:\n#         pickle.dump(net_with_mem, handle, protocol=pickle.HIGHEST_PROTOCOL)","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.070957Z","iopub.status.idle":"2023-11-17T17:33:33.071267Z","shell.execute_reply.started":"2023-11-17T17:33:33.071113Z","shell.execute_reply":"2023-11-17T17:33:33.071128Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# loss_net_with_mem_val, acc_net_with_mem_val = zip(*[(k[\"loss\"], k[\"accuracy\"]) for k in net_with_mem if k[\"phase\"]=='val'])\n# loss_net_with_mem_train, acc_net_with_mem_train = zip(*[(k[\"loss\"], k[\"accuracy\"]) for k in net_with_mem if k[\"phase\"]=='train'])","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.073059Z","iopub.status.idle":"2023-11-17T17:33:33.073482Z","shell.execute_reply.started":"2023-11-17T17:33:33.073265Z","shell.execute_reply":"2023-11-17T17:33:33.073285Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# plt.title(\"Net with memory\")\n\n# plt.plot(loss_net_with_mem_val, label = \"LossVal+mem\")\n\n# plt.plot(acc_net_with_mem_val, label = \"AccuracyVal+mem\")\n\n# plt.plot(loss_net_with_mem_train, label = \"LossTrain+mem\")\n\n# plt.plot(acc_net_with_mem_train, label = \"AccuracyTrain+mem\")\n\n# # loss_net_without_mem_val\n\n# plt.plot(loss_net_without_mem_val, label = \"Loss Val\")\n\n# plt.plot(acc_net_without_mem_val, label = \"Accuracy Val\")\n\n# plt.plot(loss_net_without_mem_train, label = \"Loss Train\")\n\n# plt.plot(acc_net_without_mem_train, label = \"Accuracy Train\")\n\n# plt.axvline(x = 100, color = 'lightgray', linestyle=\"dashed\", label = 'Start mem train')\n\n# plt.legend(loc=\"lower right\", ncol=1)\n\n# #plt.ylim(0, 1)\n# plt.xlim(0, 300)\n\n# plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-11-17T17:33:33.076288Z","iopub.status.idle":"2023-11-17T17:33:33.076663Z","shell.execute_reply.started":"2023-11-17T17:33:33.076466Z","shell.execute_reply":"2023-11-17T17:33:33.076482Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}