{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"STOCHATIC_DEPTH = 0.2\nMIXUP_ALPHA = 0.8\nEPOCHS = 5\nLABEL_SMOOTING = 0.1\nTRAINING_BATCH_SIZE = 256\nBASE_LEARNING_RATE = 8e-5\nWEIGHT_DECAY = 1e-2\nRANDAUG_LAYERS = 2\nRANDAUG_BLOCKS = 4\n\n# Warmup\nWARMUP_EPOCH = 5\n\nDATASET = \"CIFAR\"\n# DATASET = \"GALAXY\"","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:19.245505Z","iopub.execute_input":"2023-06-14T16:02:19.245846Z","iopub.status.idle":"2023-06-14T16:02:19.258475Z","shell.execute_reply.started":"2023-06-14T16:02:19.245816Z","shell.execute_reply":"2023-06-14T16:02:19.257581Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import DataLoader\nfrom tensorflow.python import keras\nimport matplotlib.pyplot as plt\n\nfrom torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR\nfrom timm.loss import LabelSmoothingCrossEntropy\nfrom tqdm import tqdm\nimport torchvision\nfrom torchvision.ops import stochastic_depth\nfrom torch.autograd import Variable\n\nfrom torch.utils.tensorboard import SummaryWriter","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:24.033427Z","iopub.execute_input":"2023-06-14T16:02:24.033867Z","iopub.status.idle":"2023-06-14T16:02:34.411903Z","shell.execute_reply.started":"2023-06-14T16:02:24.033837Z","shell.execute_reply":"2023-06-14T16:02:34.410906Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install einops","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:34.413644Z","iopub.execute_input":"2023-06-14T16:02:34.413975Z","iopub.status.idle":"2023-06-14T16:02:46.786649Z","shell.execute_reply.started":"2023-06-14T16:02:34.413932Z","shell.execute_reply":"2023-06-14T16:02:46.785530Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Lookhead","metadata":{}},{"cell_type":"code","source":"from collections import defaultdict\n\nimport torch\nfrom torch.optim.optimizer import Optimizer\n\n\nclass Lookahead(Optimizer):\n    r\"\"\"PyTorch implementation of the lookahead wrapper.\n\n    Lookahead Optimizer: https://arxiv.org/abs/1907.08610\n    \"\"\"\n\n    def __init__(self, optimizer, la_steps=5, la_alpha=0.8, pullback_momentum=\"none\"):\n        \"\"\"optimizer: inner optimizer\n        la_steps (int): number of lookahead steps\n        la_alpha (float): linear interpolation factor. 1.0 recovers the inner optimizer.\n        pullback_momentum (str): change to inner optimizer momentum on interpolation update\n        \"\"\"\n        self.optimizer = optimizer\n        self._la_step = 0  # counter for inner optimizer\n        self.la_alpha = la_alpha\n        self._total_la_steps = la_steps\n        pullback_momentum = pullback_momentum.lower()\n        assert pullback_momentum in [\"reset\", \"pullback\", \"none\"]\n        self.pullback_momentum = pullback_momentum\n\n        self.state = defaultdict(dict)\n\n        # Cache the current optimizer parameters\n        for group in optimizer.param_groups:\n            for p in group['params']:\n                param_state = self.state[p]\n                param_state['cached_params'] = torch.zeros_like(p.data)\n                param_state['cached_params'].copy_(p.data)\n                if self.pullback_momentum == \"pullback\":\n                    param_state['cached_mom'] = torch.zeros_like(p.data)\n\n    def __getstate__(self):\n        return {\n            'state': self.state,\n            'optimizer': self.optimizer,\n            'la_alpha': self.la_alpha,\n            '_la_step': self._la_step,\n            '_total_la_steps': self._total_la_steps,\n            'pullback_momentum': self.pullback_momentum\n        }\n\n    def zero_grad(self):\n        self.optimizer.zero_grad()\n\n    def get_la_step(self):\n        return self._la_step\n\n    def state_dict(self):\n        return self.optimizer.state_dict()\n\n    def load_state_dict(self, state_dict):\n        self.optimizer.load_state_dict(state_dict)\n\n    def _backup_and_load_cache(self):\n        \"\"\"Useful for performing evaluation on the slow weights (which typically generalize better)\n        \"\"\"\n        for group in self.optimizer.param_groups:\n            for p in group['params']:\n                param_state = self.state[p]\n                param_state['backup_params'] = torch.zeros_like(p.data)\n                param_state['backup_params'].copy_(p.data)\n                p.data.copy_(param_state['cached_params'])\n\n    def _clear_and_load_backup(self):\n        for group in self.optimizer.param_groups:\n            for p in group['params']:\n                param_state = self.state[p]\n                p.data.copy_(param_state['backup_params'])\n                del param_state['backup_params']\n\n    @property\n    def param_groups(self):\n        return self.optimizer.param_groups\n\n    def step(self, closure=None):\n        \"\"\"Performs a single Lookahead optimization step.\n        Arguments:\n            closure (callable, optional): A closure that reevaluates the model\n                and returns the loss.\n        \"\"\"\n        loss = self.optimizer.step(closure)\n        self._la_step += 1\n\n        if self._la_step >= self._total_la_steps:\n            self._la_step = 0\n            # Lookahead and cache the current optimizer parameters\n            for group in self.optimizer.param_groups:\n                for p in group['params']:\n                    param_state = self.state[p]\n                    p.data.mul_(self.la_alpha).add_(param_state['cached_params'], alpha=1.0 - self.la_alpha)  # crucial line\n                    param_state['cached_params'].copy_(p.data)\n                    if self.pullback_momentum == \"pullback\":\n                        internal_momentum = self.optimizer.state[p][\"momentum_buffer\"]\n                        self.optimizer.state[p][\"momentum_buffer\"] = internal_momentum.mul_(self.la_alpha).add_(\n                            1.0 - self.la_alpha, param_state[\"cached_mom\"])\n                        param_state[\"cached_mom\"] = self.optimizer.state[p][\"momentum_buffer\"]\n                    elif self.pullback_momentum == \"reset\":\n                        self.optimizer.state[p][\"momentum_buffer\"] = torch.zeros_like(p.data)\n\n        return loss","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:46.789602Z","iopub.execute_input":"2023-06-14T16:02:46.789974Z","iopub.status.idle":"2023-06-14T16:02:46.810337Z","shell.execute_reply.started":"2023-06-14T16:02:46.789936Z","shell.execute_reply":"2023-06-14T16:02:46.809362Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# RAdam","metadata":{}},{"cell_type":"code","source":"import math\nimport torch\nfrom torch.optim.optimizer import Optimizer, required\n\nclass RAdam(Optimizer):\n\n    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, degenerated_to_sgd=False):\n        if not 0.0 <= lr:\n            raise ValueError(\"Invalid learning rate: {}\".format(lr))\n        if not 0.0 <= eps:\n            raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n        if not 0.0 <= betas[0] < 1.0:\n            raise ValueError(\"Invalid beta parameter at index 0: {}\".format(betas[0]))\n        if not 0.0 <= betas[1] < 1.0:\n            raise ValueError(\"Invalid beta parameter at index 1: {}\".format(betas[1]))\n        \n        self.degenerated_to_sgd = degenerated_to_sgd\n        if isinstance(params, (list, tuple)) and len(params) > 0 and isinstance(params[0], dict):\n            for param in params:\n                if 'betas' in param and (param['betas'][0] != betas[0] or param['betas'][1] != betas[1]):\n                    param['buffer'] = [[None, None, None] for _ in range(10)]\n        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, buffer=[[None, None, None] for _ in range(10)])\n        super(RAdam, self).__init__(params, defaults)\n\n    def __setstate__(self, state):\n        super(RAdam, self).__setstate__(state)\n\n    def step(self, closure=None):\n\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        for group in self.param_groups:\n\n            for p in group['params']:\n                if p.grad is None:\n                    continue\n                grad = p.grad.data.float()\n                if grad.is_sparse:\n                    raise RuntimeError('RAdam does not support sparse gradients')\n\n                p_data_fp32 = p.data.float()\n\n                state = self.state[p]\n\n                if len(state) == 0:\n                    state['step'] = 0\n                    state['exp_avg'] = torch.zeros_like(p_data_fp32)\n                    state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)\n                else:\n                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)\n                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)\n\n                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n                beta1, beta2 = group['betas']\n\n#                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)\n                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n                exp_avg.mul_(beta1).add_(1 - beta1, grad)\n\n                state['step'] += 1\n                buffered = group['buffer'][int(state['step'] % 10)]\n                if state['step'] == buffered[0]:\n                    N_sma, step_size = buffered[1], buffered[2]\n                else:\n                    buffered[0] = state['step']\n                    beta2_t = beta2 ** state['step']\n                    N_sma_max = 2 / (1 - beta2) - 1\n                    N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)\n                    buffered[1] = N_sma\n\n                    # more conservative since it's an approximated value\n                    if N_sma >= 5:\n                        step_size = math.sqrt((1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / (N_sma_max - 2)) / (1 - beta1 ** state['step'])\n                    elif self.degenerated_to_sgd:\n                        step_size = 1.0 / (1 - beta1 ** state['step'])\n                    else:\n                        step_size = -1\n                    buffered[2] = step_size\n\n                # more conservative since it's an approximated value\n                if N_sma >= 5:\n                    if group['weight_decay'] != 0:\n                        p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)\n                    denom = exp_avg_sq.sqrt().add_(group['eps'])\n                    p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom)\n                    p.data.copy_(p_data_fp32)\n                elif step_size > 0:\n                    if group['weight_decay'] != 0:\n                        p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)\n                    p_data_fp32.add_(-step_size * group['lr'], exp_avg)\n                    p.data.copy_(p_data_fp32)\n\n        return loss","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:46.813166Z","iopub.execute_input":"2023-06-14T16:02:46.814398Z","iopub.status.idle":"2023-06-14T16:02:46.838641Z","shell.execute_reply.started":"2023-06-14T16:02:46.814365Z","shell.execute_reply":"2023-06-14T16:02:46.837654Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# COATnet","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\nfrom einops import rearrange\nfrom einops.layers.torch import Rearrange\n\n\ndef conv_3x3_bn(inp, oup, image_size, downsample=False):\n    stride = 1 if downsample == False else 2\n    return nn.Sequential(\n        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),\n        nn.BatchNorm2d(oup),\n        nn.GELU()\n    )\n\n\nclass PreNorm(nn.Module):\n    def __init__(self, dim, fn, norm):\n        super().__init__()\n        self.norm = norm(dim)\n        self.fn = fn\n\n    def forward(self, x, **kwargs):\n        return self.fn(self.norm(x), **kwargs)\n\n\nclass SE(nn.Module):\n    def __init__(self, inp, oup, expansion=0.25):\n        super().__init__()\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.fc = nn.Sequential(\n            nn.Linear(oup, int(inp * expansion), bias=False),\n            nn.GELU(),\n            nn.Linear(int(inp * expansion), oup, bias=False),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        b, c, _, _ = x.size()\n        y = self.avg_pool(x).view(b, c)\n        y = self.fc(y).view(b, c, 1, 1)\n        return x * y\n\n\nclass FeedForward(nn.Module):\n    def __init__(self, dim, hidden_dim, dropout=0.):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(dim, hidden_dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, dim),\n            nn.Dropout(dropout)\n        )\n\n    def forward(self, x):\n        return self.net(x)\n\n\nclass MBConv(nn.Module):\n    def __init__(self, inp, oup, image_size, downsample=False, expansion=4):\n        super().__init__()\n        self.downsample = downsample\n        stride = 1 if self.downsample == False else 2\n        hidden_dim = int(inp * expansion)\n\n        if self.downsample:\n            self.pool = nn.MaxPool2d(3, 2, 1)\n            self.proj = nn.Conv2d(inp, oup, 1, 1, 0, bias=False)\n\n        if expansion == 1:\n            self.conv = nn.Sequential(\n                # dw\n                nn.Conv2d(hidden_dim, hidden_dim, 3, stride,\n                          1, groups=hidden_dim, bias=False),\n                nn.BatchNorm2d(hidden_dim),\n                nn.GELU(),\n                # pw-linear\n                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),\n                nn.BatchNorm2d(oup),\n            )\n        else:\n            self.conv = nn.Sequential(\n                # pw\n                # down-sample in the first conv\n                nn.Conv2d(inp, hidden_dim, 1, stride, 0, bias=False),\n                nn.BatchNorm2d(hidden_dim),\n                nn.GELU(),\n                # dw\n                nn.Conv2d(hidden_dim, hidden_dim, 3, 1, 1,\n                          groups=hidden_dim, bias=False),\n                nn.BatchNorm2d(hidden_dim),\n                nn.GELU(),\n                SE(inp, hidden_dim),\n                # pw-linear\n                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),\n                nn.BatchNorm2d(oup),\n            )\n        \n        self.conv = PreNorm(inp, self.conv, nn.BatchNorm2d)\n\n    def forward(self, x):\n        if self.downsample:\n            return self.proj(self.pool(x)) + self.conv(x)\n        else:\n            return x + self.conv(x)\n\n\nclass Attention(nn.Module):\n    def __init__(self, inp, oup, image_size, heads=8, dim_head=32, dropout=0.):\n        super().__init__()\n        inner_dim = dim_head * heads\n        project_out = not (heads == 1 and dim_head == inp)\n\n        self.ih, self.iw = image_size\n\n        self.heads = heads\n        self.scale = dim_head ** -0.5\n\n        # parameter table of relative position bias\n        self.relative_bias_table = nn.Parameter(\n            torch.zeros((2 * self.ih - 1) * (2 * self.iw - 1), heads))\n\n        coords = torch.meshgrid((torch.arange(self.ih), torch.arange(self.iw)))\n        coords = torch.flatten(torch.stack(coords), 1)\n        relative_coords = coords[:, :, None] - coords[:, None, :]\n\n        relative_coords[0] += self.ih - 1\n        relative_coords[1] += self.iw - 1\n        relative_coords[0] *= 2 * self.iw - 1\n        relative_coords = rearrange(relative_coords, 'c h w -> h w c')\n        relative_index = relative_coords.sum(-1).flatten().unsqueeze(1)\n        self.register_buffer(\"relative_index\", relative_index)\n\n        self.attend = nn.Softmax(dim=-1)\n        self.to_qkv = nn.Linear(inp, inner_dim * 3, bias=False)\n\n        self.to_out = nn.Sequential(\n            nn.Linear(inner_dim, oup),\n            nn.Dropout(dropout)\n        ) if project_out else nn.Identity()\n\n    def forward(self, x):\n        qkv = self.to_qkv(x).chunk(3, dim=-1)\n        q, k, v = map(lambda t: rearrange(\n            t, 'b n (h d) -> b h n d', h=self.heads), qkv)\n\n        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale\n\n        # Use \"gather\" for more efficiency on GPUs\n        relative_bias = self.relative_bias_table.gather(\n            0, self.relative_index.repeat(1, self.heads))\n        relative_bias = rearrange(\n            relative_bias, '(h w) c -> 1 c h w', h=self.ih*self.iw, w=self.ih*self.iw)\n        dots = dots + relative_bias\n\n        attn = self.attend(dots)\n        out = torch.matmul(attn, v)\n        out = rearrange(out, 'b h n d -> b n (h d)')\n        out = self.to_out(out)\n        return out\n\n\nclass Transformer(nn.Module):\n    def __init__(self, inp, oup, image_size, heads=8, dim_head=32, downsample=False, dropout=0.):\n        super().__init__()\n        hidden_dim = int(inp * 4)\n\n        self.ih, self.iw = image_size\n        self.downsample = downsample\n\n        if self.downsample:\n            self.pool1 = nn.MaxPool2d(3, 2, 1)\n            self.pool2 = nn.MaxPool2d(3, 2, 1)\n            self.proj = nn.Conv2d(inp, oup, 1, 1, 0, bias=False)\n\n        self.attn = Attention(inp, oup, image_size, heads, dim_head, dropout)\n        self.ff = FeedForward(oup, hidden_dim, dropout)\n\n        self.attn = nn.Sequential(\n            Rearrange('b c ih iw -> b (ih iw) c'),\n            PreNorm(inp, self.attn, nn.LayerNorm),\n            Rearrange('b (ih iw) c -> b c ih iw', ih=self.ih, iw=self.iw)\n        )\n\n        self.ff = nn.Sequential(\n            Rearrange('b c ih iw -> b (ih iw) c'),\n            PreNorm(oup, self.ff, nn.LayerNorm),\n            Rearrange('b (ih iw) c -> b c ih iw', ih=self.ih, iw=self.iw)\n        )\n\n    def forward(self, x):\n        if self.downsample:\n            x = self.proj(self.pool1(x)) + self.attn(self.pool2(x))\n        else:\n            x = x + self.attn(x)\n        x = x + self.ff(x)\n        return x\n\n\nclass CoAtNet(nn.Module):\n    def __init__(self, image_size, in_channels, num_blocks, channels, num_classes=1000, block_types=['C', 'C', 'T', 'T']):\n        super().__init__()\n        ih, iw = image_size\n        block = {'C': MBConv, 'T': Transformer}\n\n        self.s0 = self._make_layer(\n            conv_3x3_bn, in_channels, channels[0], num_blocks[0], (ih // 2, iw // 2))\n        self.s1 = self._make_layer(\n            block[block_types[0]], channels[0], channels[1], num_blocks[1], (ih // 4, iw // 4))\n        self.s2 = self._make_layer(\n            block[block_types[1]], channels[1], channels[2], num_blocks[2], (ih // 8, iw // 8))\n        self.s3 = self._make_layer(\n            block[block_types[2]], channels[2], channels[3], num_blocks[3], (ih // 16, iw // 16))\n        self.s4 = self._make_layer(\n            block[block_types[3]], channels[3], channels[4], num_blocks[4], (ih // 32, iw // 32))\n\n        self.pool = nn.AvgPool2d(ih // 32, 1)\n        self.fc = nn.Linear(channels[-1], num_classes, bias=False)\n\n    def forward(self, x):\n        #https://pytorch.org/vision/main/_modules/torchvision/ops/stochastic_depth.html#stochastic_depth\n        p = STOCHATIC_DEPTH\n        \n#         x = stochastic_depth(self.s0(x), p=p, mode=\"row\", training=self.training)\n#         x = stochastic_depth(self.s1(x), p=p, mode=\"row\", training=self.training)\n#         x = stochastic_depth(self.s2(x), p=p, mode=\"row\", training=self.training)\n#         x = stochastic_depth(self.s3(x), p=p, mode=\"row\", training=self.training)\n#         x = stochastic_depth(self.s4(x), p=p, mode=\"row\", training=self.training)\n        \n        x = self.s0(x)\n        x = self.s1(x)\n        x = self.s2(x)\n        x = self.s3(x)\n        x = self.s4(x)\n\n        x = self.pool(x).view(-1, x.shape[1])\n        x = self.fc(x)\n        return x\n\n    def _make_layer(self, block, inp, oup, depth, image_size):\n        layers = nn.ModuleList([])\n        for i in range(depth):\n            if i == 0:\n                layers.append(block(inp, oup, image_size, downsample=True))\n            else:\n                layers.append(block(oup, oup, image_size))\n        return nn.Sequential(*layers)\n\n\ndef coatnet_0():\n    num_blocks = [2, 2, 3, 5, 2]            # L\n    channels = [64, 96, 192, 384, 768]      # D\n    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)\n\n\ndef coatnet_1():\n    num_blocks = [2, 2, 6, 14, 2]           # L\n    channels = [64, 96, 192, 384, 768]      # D\n    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)\n\n\ndef coatnet_2():\n    num_blocks = [2, 2, 6, 14, 2]           # L\n    channels = [128, 128, 256, 512, 1026]   # D\n    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)\n\n\ndef coatnet_3():\n    num_blocks = [2, 2, 6, 14, 2]           # L\n    channels = [192, 192, 384, 768, 1536]   # D\n    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)\n\n\ndef coatnet_4():\n    num_blocks = [2, 2, 12, 28, 2]          # L\n    channels = [192, 192, 384, 768, 1536]   # D\n    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)\n\n\ndef count_parameters(model):\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:46.840069Z","iopub.execute_input":"2023-06-14T16:02:46.840662Z","iopub.status.idle":"2023-06-14T16:02:46.901236Z","shell.execute_reply.started":"2023-06-14T16:02:46.840623Z","shell.execute_reply":"2023-06-14T16:02:46.900259Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# MIX-up","metadata":{}},{"cell_type":"code","source":"def mixup_data(x, y, alpha=1.0, use_cuda=True):\n    '''Returns mixed inputs, pairs of targets, and lambda'''\n    if alpha > 0:\n        lam = np.random.beta(alpha, alpha)\n    else:\n        lam = 1\n\n    batch_size = x.size()[0]\n\n    if device == \"cuda\":\n        index = torch.randperm(batch_size).cuda()\n    else:\n        index = torch.randperm(batch_size)\n\n    mixed_x = lam * x + (1 - lam) * x[index, :]\n    y_a, y_b = y, y[index]\n    return mixed_x, y_a, y_b, lam\n\ndef mixup_criterion(criterion, pred, y_a, y_b, lam):\n    return lam * criterion(pred, y_a.to(torch.int64)) + (1 - lam) * criterion(pred, y_b.to(torch.int64))","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:46.903049Z","iopub.execute_input":"2023-06-14T16:02:46.903693Z","iopub.status.idle":"2023-06-14T16:02:46.912996Z","shell.execute_reply.started":"2023-06-14T16:02:46.903661Z","shell.execute_reply":"2023-06-14T16:02:46.912083Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Rand augmentation","metadata":{}},{"cell_type":"code","source":"# code in this file is adpated from rpmcruz/autoaugment\n# https://github.com/rpmcruz/autoaugment/blob/master/transformations.py\nimport random\n\nimport PIL, PIL.ImageOps, PIL.ImageEnhance, PIL.ImageDraw\nimport numpy as np\nimport torch\nfrom PIL import Image\n\n\ndef ShearX(img, v):  # [-0.3, 0.3]\n    assert -0.3 <= v <= 0.3\n    if random.random() > 0.5:\n        v = -v\n    return img.transform(img.size, PIL.Image.AFFINE, (1, v, 0, 0, 1, 0))\n\n\ndef ShearY(img, v):  # [-0.3, 0.3]\n    assert -0.3 <= v <= 0.3\n    if random.random() > 0.5:\n        v = -v\n    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, v, 1, 0))\n\n\ndef TranslateX(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]\n    assert -0.45 <= v <= 0.45\n    if random.random() > 0.5:\n        v = -v\n    v = v * img.size[0]\n    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0))\n\n\ndef TranslateXabs(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]\n    assert 0 <= v\n    if random.random() > 0.5:\n        v = -v\n    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0))\n\n\ndef TranslateY(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]\n    assert -0.45 <= v <= 0.45\n    if random.random() > 0.5:\n        v = -v\n    v = v * img.size[1]\n    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v))\n\n\ndef TranslateYabs(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]\n    assert 0 <= v\n    if random.random() > 0.5:\n        v = -v\n    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v))\n\n\ndef Rotate(img, v):  # [-30, 30]\n    assert -30 <= v <= 30\n    if random.random() > 0.5:\n        v = -v\n    return img.rotate(v)\n\n\ndef AutoContrast(img, _):\n    return PIL.ImageOps.autocontrast(img)\n\n\ndef Invert(img, _):\n    return PIL.ImageOps.invert(img)\n\n\ndef Equalize(img, _):\n    return PIL.ImageOps.equalize(img)\n\n\ndef Flip(img, _):  # not from the paper\n    return PIL.ImageOps.mirror(img)\n\n\ndef Solarize(img, v):  # [0, 256]\n    assert 0 <= v <= 256\n    return PIL.ImageOps.solarize(img, v)\n\n\ndef SolarizeAdd(img, addition=0, threshold=128):\n    img_np = np.array(img).astype(int)\n    img_np = img_np + addition\n    img_np = np.clip(img_np, 0, 255)\n    img_np = img_np.astype(np.uint8)\n    img = Image.fromarray(img_np)\n    return PIL.ImageOps.solarize(img, threshold)\n\n\ndef Posterize(img, v):  # [4, 8]\n    v = int(v)\n    v = max(1, v)\n    return PIL.ImageOps.posterize(img, v)\n\n\ndef Contrast(img, v):  # [0.1,1.9]\n    assert 0.1 <= v <= 1.9\n    return PIL.ImageEnhance.Contrast(img).enhance(v)\n\n\ndef Color(img, v):  # [0.1,1.9]\n    assert 0.1 <= v <= 1.9\n    return PIL.ImageEnhance.Color(img).enhance(v)\n\n\ndef Brightness(img, v):  # [0.1,1.9]\n    assert 0.1 <= v <= 1.9\n    return PIL.ImageEnhance.Brightness(img).enhance(v)\n\n\ndef Sharpness(img, v):  # [0.1,1.9]\n    assert 0.1 <= v <= 1.9\n    return PIL.ImageEnhance.Sharpness(img).enhance(v)\n\n\ndef Cutout(img, v):  # [0, 60] => percentage: [0, 0.2]\n    assert 0.0 <= v <= 0.2\n    if v <= 0.:\n        return img\n\n    v = v * img.size[0]\n    return CutoutAbs(img, v)\n\n\ndef CutoutAbs(img, v):  # [0, 60] => percentage: [0, 0.2]\n    # assert 0 <= v <= 20\n    if v < 0:\n        return img\n    w, h = img.size\n    x0 = np.random.uniform(w)\n    y0 = np.random.uniform(h)\n\n    x0 = int(max(0, x0 - v / 2.))\n    y0 = int(max(0, y0 - v / 2.))\n    x1 = min(w, x0 + v)\n    y1 = min(h, y0 + v)\n\n    xy = (x0, y0, x1, y1)\n    color = (125, 123, 114)\n    # color = (0, 0, 0)\n    img = img.copy()\n    PIL.ImageDraw.Draw(img).rectangle(xy, color)\n    return img\n\n\ndef SamplePairing(imgs):  # [0, 0.4]\n    def f(img1, v):\n        i = np.random.choice(len(imgs))\n        img2 = PIL.Image.fromarray(imgs[i])\n        return PIL.Image.blend(img1, img2, v)\n\n    return f\n\n\ndef Identity(img, v):\n    return img\n\n\ndef augment_list():  # 16 oeprations and their ranges\n    # https://github.com/tensorflow/tpu/blob/8462d083dd89489a79e3200bcc8d4063bf362186/models/official/efficientnet/autoaugment.py#L505\n    l = [\n        (AutoContrast, 0, 1),\n        (Equalize, 0, 1),\n        (Invert, 0, 1),\n        (Rotate, 0, 30),\n        (Posterize, 0, 4),\n        (Solarize, 0, 256),\n        (SolarizeAdd, 0, 110),\n        (Color, 0.1, 1.9),\n        (Contrast, 0.1, 1.9),\n        (Brightness, 0.1, 1.9),\n        (Sharpness, 0.1, 1.9),\n        (ShearX, 0., 0.3),\n        (ShearY, 0., 0.3),\n        (CutoutAbs, 0, 40),\n        (TranslateXabs, 0., 100),\n        (TranslateYabs, 0., 100),\n    ]\n\n    return l\n\n\nclass Lighting(object):\n    \"\"\"Lighting noise(AlexNet - style PCA - based noise)\"\"\"\n\n    def __init__(self, alphastd, eigval, eigvec):\n        self.alphastd = alphastd\n        self.eigval = torch.Tensor(eigval)\n        self.eigvec = torch.Tensor(eigvec)\n\n    def __call__(self, img):\n        if self.alphastd == 0:\n            return img\n\n        alpha = img.new().resize_(3).normal_(0, self.alphastd)\n        rgb = self.eigvec.type_as(img).clone() \\\n            .mul(alpha.view(1, 3).expand(3, 3)) \\\n            .mul(self.eigval.view(1, 3).expand(3, 3)) \\\n            .sum(1).squeeze()\n\n        return img.add(rgb.view(3, 1, 1).expand_as(img))\n\n\nclass CutoutDefault(object):\n    \"\"\"\n    Reference : https://github.com/quark0/darts/blob/master/cnn/utils.py\n    \"\"\"\n    def __init__(self, length):\n        self.length = length\n\n    def __call__(self, img):\n        h, w = img.size(1), img.size(2)\n        mask = np.ones((h, w), np.float32)\n        y = np.random.randint(h)\n        x = np.random.randint(w)\n\n        y1 = np.clip(y - self.length // 2, 0, h)\n        y2 = np.clip(y + self.length // 2, 0, h)\n        x1 = np.clip(x - self.length // 2, 0, w)\n        x2 = np.clip(x + self.length // 2, 0, w)\n\n        mask[y1: y2, x1: x2] = 0.\n        mask = torch.from_numpy(mask)\n        mask = mask.expand_as(img)\n        img *= mask\n        return img\n\n\nclass RandAugment:\n    def __init__(self, n, m):\n        self.n = n\n        self.m = m      # [0, 30]\n        self.augment_list = augment_list()\n\n    def __call__(self, img):\n        ops = random.choices(self.augment_list, k=self.n)\n        for op, minval, maxval in ops:\n            val = (float(self.m) / 30) * float(maxval - minval) + minval\n            img = op(img, val)\n\n        return img","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:46.914491Z","iopub.execute_input":"2023-06-14T16:02:46.914813Z","iopub.status.idle":"2023-06-14T16:02:46.953853Z","shell.execute_reply.started":"2023-06-14T16:02:46.914783Z","shell.execute_reply":"2023-06-14T16:02:46.952948Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nfrom torchvision.io import read_image\nfrom torch.utils.data import Dataset\nimport os\nfrom PIL import Image\n\nclass Galaxy10(Dataset):\n    def __init__(self, img_dir, labels_dir, transform):\n        self.img_labels = pd.read_csv(labels_dir)\n        self.img_dir = img_dir\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.img_labels)\n\n    def __getitem__(self, idx):\n        img_path = self.img_dir + str(idx) + \".png\"\n        image = read_image(img_path).numpy()\n        image = np.transpose(image, (2, 1, 0))\n        image = Image.fromarray(image.astype('uint8'), 'RGB')\n        label = torch.Tensor(self.img_labels.iloc[idx].values)\n        \n        if self.transform:\n            image = self.transform(image)\n        return image, label\n\nif DATASET == \"GALAXY\":\n    train_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(), torchvision.transforms.ConvertImageDtype(torch.float)])\n    train_transform.transforms.insert(0, RandAugment(RANDAUG_LAYERS, RANDAUG_BLOCKS))\n    test_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(), torchvision.transforms.ConvertImageDtype(torch.float)])\n    \n    train_dataset = Galaxy10(\"/kaggle/input/galaxy10/images/images/train/image_\", \"/kaggle/input/galaxy10/train_labels.csv\", train_transform)\n    train_dataloader = DataLoader(train_dataset, batch_size=TRAINING_BATCH_SIZE, shuffle=True)\n    \n    test_dataset = Galaxy10(\"/kaggle/input/galaxy10/images/images/test/image_\", \"/kaggle/input/galaxy10/test_labels.csv\", test_transform)    \n    test_dataloader = DataLoader(test_dataset, batch_size=TRAINING_BATCH_SIZE, shuffle=False)\n    \n#     num_blocks = [2, 2, 3, 5, 2]       # Only change S2 and S3\n#     channels = [7, 14, 28, 56, 112]   # Doubles from S1 to S4\n\n    num_blocks = [2, 2, 3, 5, 2]            # L\n    channels = [64, 96, 192, 384, 768]      # D\n\n    model = CoAtNet((256, 256), 3, num_blocks, channels, num_classes=10, block_types=['C', 'C', 'C', 'T'])\n    print(\"Number of parameters:\", count_parameters(model))","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:46.955191Z","iopub.execute_input":"2023-06-14T16:02:46.955646Z","iopub.status.idle":"2023-06-14T16:02:46.969918Z","shell.execute_reply.started":"2023-06-14T16:02:46.955615Z","shell.execute_reply":"2023-06-14T16:02:46.968863Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DATASET == \"CIFAR\":\n    train_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(), torchvision.transforms.ConvertImageDtype(torch.float), torchvision.transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))])\n    train_transform.transforms.insert(0, RandAugment(RANDAUG_LAYERS, RANDAUG_BLOCKS))\n    test_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(), torchvision.transforms.ConvertImageDtype(torch.float), torchvision.transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))])\n    \n    trainset = torchvision.datasets.CIFAR10(root='./cifar10', train=True, download=True, transform=train_transform)\n    train_dataloader = torch.utils.data.DataLoader(trainset, batch_size=TRAINING_BATCH_SIZE, shuffle=True, num_workers=0)\n    \n    testset = torchvision.datasets.CIFAR10(root='./cifar10', train=False, download=True, transform=test_transform)\n    test_dataloader = torch.utils.data.DataLoader(testset, batch_size=TRAINING_BATCH_SIZE, shuffle=True, num_workers=0)\n    \n    num_blocks = [2, 2, 3, 5, 2]            # L\n    channels = [64, 96, 192, 384, 768]      # D\n\n    model = CoAtNet((32, 32), 3, num_blocks, channels, num_classes=10, block_types=['C', 'C', 'T', 'T'])\n    print(\"Number of parameters:\", count_parameters(model))","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:02:46.971421Z","iopub.execute_input":"2023-06-14T16:02:46.971805Z","iopub.status.idle":"2023-06-14T16:03:04.601618Z","shell.execute_reply.started":"2023-06-14T16:02:46.971774Z","shell.execute_reply":"2023-06-14T16:03:04.600694Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"device = 'cuda' if torch.cuda.is_available() else 'cpu'\nprint(\"Device:\", device)\n\nif device == \"cuda\":\n    model.cuda()","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:03:04.605055Z","iopub.execute_input":"2023-06-14T16:03:04.606468Z","iopub.status.idle":"2023-06-14T16:03:09.460071Z","shell.execute_reply.started":"2023-06-14T16:03:04.606432Z","shell.execute_reply":"2023-06-14T16:03:09.459151Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{"execution":{"iopub.status.busy":"2023-05-13T11:11:04.676535Z","iopub.execute_input":"2023-05-13T11:11:04.676962Z","iopub.status.idle":"2023-05-13T11:11:04.985775Z","shell.execute_reply.started":"2023-05-13T11:11:04.676934Z","shell.execute_reply":"2023-05-13T11:11:04.984886Z"}}},{"cell_type":"code","source":"def train(model, optimizer, criterion):\n    \"\"\"\n    Trains network for one epoch in batches.\n\n    Args:\n        train_iter: Data iterator for training set.\n        model: Neural network model.\n        optimizer: Optimizer (e.g. SGD).\n        lr_scheduler: Scheduler for adjusting the optimization learning rate.\n        criterion: Loss function (e.g. cross-entropy loss).\n    \"\"\"\n  \n    model.train()\n\n    train_loss = 0\n    correct = 0\n    total = 0\n    # Iterate through batches\n    for batch_idx, (inputs, targets) in enumerate(train_dataloader):\n        optimizer.zero_grad()\n        \n        if device == \"cuda\":\n            inputs, targets = inputs.cuda(), targets.cuda()\n        targets = torch.ravel(targets)\n        \n#         outputs = model(inputs)\n#         loss = criterion(outputs, targets.to(torch.int64))\n#         train_loss += loss.item()\n#         _, predicted = torch.max(outputs.data, 1)\n#         total += targets.size(0)\n#         correct += predicted.eq(targets.data).cpu().sum()\n            \n        inputs, targets_a, targets_b, lam = mixup_data(inputs, targets, MIXUP_ALPHA, device)\n        inputs, targets_a, targets_b = map(Variable, (inputs, targets_a, targets_b))        \n        outputs = model(inputs)\n        loss = mixup_criterion(criterion, outputs, targets_a, targets_b, lam)\n        train_loss += loss.item()\n        _, predicted = torch.max(outputs.data, 1)\n        total += targets.size(0)\n        correct += (lam * predicted.eq(targets_a.data).cpu().sum().float() + (1 - lam) * predicted.eq(targets_b.data).cpu().sum().float())\n        \n        loss.backward()\n        optimizer.step()\n        \n\n    return train_loss/batch_idx, 100.*correct/total","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:03:09.465339Z","iopub.execute_input":"2023-06-14T16:03:09.468306Z","iopub.status.idle":"2023-06-14T16:03:09.482591Z","shell.execute_reply.started":"2023-06-14T16:03:09.468271Z","shell.execute_reply":"2023-06-14T16:03:09.481553Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test(model, criterion):\n    \"\"\"\n    Evaluates network in batches.\n\n    Args:\n        test_iter: Data iterator for test set.\n        model: Neural network model.\n        criterion: Loss function (e.g. cross-entropy loss).\n    \"\"\"\n\n    model.eval()\n    test_loss = 0\n    correct = 0\n    total = 0\n    \n    # Use torch.no_grad to skip gradient calculation, not needed for evaluation\n    with torch.no_grad():\n        # Iterate through batches\n        for batch_idx, (inputs, targets) in enumerate(test_dataloader, 0):\n            if device == \"cuda\":\n                inputs, targets = inputs.cuda(), targets.cuda()\n            targets = torch.ravel(targets)\n            \n            # Forward pass\n            outputs = model(inputs)\n            loss = criterion(outputs, targets.to(torch.int64))\n\n            test_loss += loss.item()\n            _, predicted = torch.max(outputs.data, 1)\n            total += targets.size(0)\n            correct += predicted.eq(targets.data).cpu().sum()\n            \n            \n    acc = 100.*correct/total\n    \n    return test_loss/batch_idx, 100.*correct/total","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:03:09.487733Z","iopub.execute_input":"2023-06-14T16:03:09.491305Z","iopub.status.idle":"2023-06-14T16:03:09.504854Z","shell.execute_reply.started":"2023-06-14T16:03:09.491270Z","shell.execute_reply":"2023-06-14T16:03:09.502745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run(model, epochs=300):\n    # Create classifier model\n    model = model.to(device)\n\n    # Create loss function and optimizer\n    criterion = LabelSmoothingCrossEntropy(LABEL_SMOOTING)\n    \n    optimizer = RAdam(model.parameters(),lr=BASE_LEARNING_RATE, weight_decay=WEIGHT_DECAY)\n    optimizer = Lookahead(optimizer, la_steps=5, la_alpha=0.8)\n\n    def warmup_lr_lambda(epoch):\n        if epoch < WARMUP_EPOCH:\n            return epoch / WARMUP_EPOCH\n        return 1.0\n\n    warmup_scheduler = LambdaLR(optimizer, lr_lambda=warmup_lr_lambda)\n    base_scheduler = CosineAnnealingLR(optimizer, T_max=epochs - WARMUP_EPOCH, eta_min=BASE_LEARNING_RATE/4)\n    \n    # Training loop\n    train_losses = []\n    test_losses = []\n    train_acces = []\n    test_acces = []\n    learning_rates = []\n    \n    # Early stopping\n    best_acc = 0\n    patience_cnt = 0\n    patience = int(epochs / 10)\n    \n    for epoch in tqdm(range(epochs)):        \n        # Warm-up learning rate\n        if epoch < WARMUP_EPOCH:\n            warmup_scheduler.step()            \n        else:\n            base_scheduler.step()\n            \n        train_loss, train_acc = train(model, optimizer, criterion)\n\n        # Test on data\n        test_loss, test_acc = test(model, criterion)\n        \n        train_losses.append(train_loss)\n        test_losses.append(test_loss)\n        train_acces.append(train_acc)\n        test_acces.append(test_acc)\n        learning_rates.append(optimizer.param_groups[0]['lr'])\n        \n        if epoch % 10 == 0:\n            print(\"Epoch:\", epoch, \", train_loss:\", train_loss, \", train_acc:\", train_acc, \", test_loss:\", test_loss, \", test_acc:\", test_acc)\n        \n#         # Check for early stopping\n#         if test_acc > best_acc:\n#             best_acc = test_acc\n#             patience_cnt = 0\n            \n#         else:\n#             patience_cnt += 1\n#             if patience_cnt == patience:\n#                 print(\"Trading got stopped out due to early stopping\")\n#                 break\n\n    print(\"Training finished\")\n    return np.array(train_losses), np.array(train_acces), np.array(test_losses), np.array(test_acces), np.array(learning_rates)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:03:09.507658Z","iopub.execute_input":"2023-06-14T16:03:09.508186Z","iopub.status.idle":"2023-06-14T16:03:09.521018Z","shell.execute_reply.started":"2023-06-14T16:03:09.508155Z","shell.execute_reply":"2023-06-14T16:03:09.520237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_loss, train_acc, test_loss, test_acc, learning_rates = run(model, epochs=300)\n\nplt.plot(np.arange(train_loss.shape[0]), train_loss, label=\"Train loss\", color=\"tab:blue\")\nplt.plot(np.arange(test_loss.shape[0]), test_loss, label=\"Test loss\", color=\"tab:orange\")\nplt.xlabel(\"Epochs\")\nplt.ylabel(\"Loss\")\nplt.legend()\nplt.grid()\nplt.title(\"Loss function\")\nplt.show()\n\nplt.plot(np.arange(train_acc.shape[0]), train_acc, label=\"Train acc\", color=\"tab:blue\")\nplt.plot(np.arange(test_acc.shape[0]), test_acc, label=\"Test acc\", color=\"tab:orange\")\nplt.xlabel(\"Epochs\")\nplt.ylabel(\"Accuracy\")\nplt.legend()\nplt.grid()\nplt.title(\"Accuracy\")\nplt.show()\n\n# Plotting the learning rate over epochs\nplt.plot(np.arange(learning_rates.shape[0]), learning_rates)\nplt.xlabel('Epochs')\nplt.ylabel('Learning Rate')\nplt.title('Learning Rate Schedule')\nplt.grid()\nplt.show()\n\nprint(\"Best test acc: \", max(test_acc))\n\n# Specify the folder path to save the results\nsave_folder = \"results\"\n\n# Create the folder if it doesn't exist\nif not os.path.exists(save_folder):\n    os.makedirs(save_folder)\n\n# Save the trained model\ntorch.save(model.state_dict(), os.path.join(save_folder, \"model.pth\"))\n\n# Save the arrays as NumPy files\nnp.save(os.path.join(save_folder, \"train_acc.npy\"), train_acc)\nnp.save(os.path.join(save_folder, \"train_loss.npy\"), train_loss)\nnp.save(os.path.join(save_folder, \"test_acc.npy\"), test_acc)\nnp.save(os.path.join(save_folder, \"test_loss.npy\"), test_loss)\n\nprint(train_acc)\nprint(train_loss)\nprint(test_acc)\nprint(test_loss)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T16:03:09.523139Z","iopub.execute_input":"2023-06-14T16:03:09.523805Z","iopub.status.idle":"2023-06-14T19:36:28.002187Z","shell.execute_reply.started":"2023-06-14T16:03:09.523722Z","shell.execute_reply":"2023-06-14T19:36:28.001269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Coat0: 48.63 s/it","metadata":{"execution":{"iopub.status.busy":"2023-05-31T20:24:00.805677Z","iopub.execute_input":"2023-05-31T20:24:00.806075Z","iopub.status.idle":"2023-05-31T20:24:01.134092Z","shell.execute_reply.started":"2023-05-31T20:24:00.806043Z","shell.execute_reply":"2023-05-31T20:24:01.132883Z"}}},{"cell_type":"code","source":"# # Calculate the model memory\n# model_memory = sum(p.numel() * p.element_size() for p in model.parameters())\n\n# # Define the size of your input tensor\n# input_size = (3, 256, 256)  # Example: 3 channels, 32x32 image\n\n# # Calculate the input tensor memory\n# input_tensor = torch.randn(1, *input_size)\n# input_tensor_memory = input_tensor.element_size() * input_tensor.numel()\n\n# # Calculate the total memory requirement per sample\n# memory_per_sample = model_memory + input_tensor_memory\n\n# # Print the memory requirement per sample\n# print(f\"Memory per sample: {memory_per_sample} bytes\")","metadata":{"execution":{"iopub.status.busy":"2023-06-01T10:43:34.087761Z","iopub.status.idle":"2023-06-01T10:43:34.088534Z","shell.execute_reply.started":"2023-06-01T10:43:34.088285Z","shell.execute_reply":"2023-06-01T10:43:34.088309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# gpu_properties = torch.cuda.get_device_properties(device)\n# total_memory = gpu_properties.total_memory\n# optimal_batch_size = int(total_memory / memory_per_sample)\n# print(optimal_batch_size)","metadata":{"execution":{"iopub.status.busy":"2023-06-01T10:43:34.089923Z","iopub.status.idle":"2023-06-01T10:43:34.090690Z","shell.execute_reply.started":"2023-06-01T10:43:34.090448Z","shell.execute_reply":"2023-06-01T10:43:34.090471Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pickle\n\nfilename = 'finalized_model.sav'\npickle.dump(model, open(filename, 'wb'))","metadata":{"execution":{"iopub.status.busy":"2023-06-01T10:43:34.094247Z","iopub.status.idle":"2023-06-01T10:43:34.095016Z","shell.execute_reply.started":"2023-06-01T10:43:34.094763Z","shell.execute_reply":"2023-06-01T10:43:34.094786Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loaded_model = pickle.load(open(filename, 'rb'))","metadata":{"execution":{"iopub.status.busy":"2023-06-01T10:43:34.096320Z","iopub.status.idle":"2023-06-01T10:43:34.097105Z","shell.execute_reply.started":"2023-06-01T10:43:34.096850Z","shell.execute_reply":"2023-06-01T10:43:34.096873Z"},"trusted":true},"execution_count":null,"outputs":[]}]}