relaxed-attention-vision.ipynb

{"cells":[{"cell_type":"markdown","metadata":{},"source":["# Relaxed Attention for Transformer Models"]},{"cell_type":"markdown","metadata":{},"source":["paper -> https://arxiv.org/pdf/2209.09735.pdf\n","\n","### Notes\n","\n","- abstract: The powerful modeling capabilities of all-attention-based transformer architec-\n","tures often cause overfitting.\n","\n","- instead of swin, I implemented vanilla ViT with relaxed attention.\n","- followed the training receipe at the paper."]},{"cell_type":"code","execution_count":1,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T12:21:59.500476Z","iopub.status.busy":"2022-09-27T12:21:59.499279Z","iopub.status.idle":"2022-09-27T12:22:21.456031Z","shell.execute_reply":"2022-09-27T12:22:21.454720Z","shell.execute_reply.started":"2022-09-27T12:21:59.500345Z"},"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["Collecting einops\n","  Downloading einops-0.4.1-py3-none-any.whl (28 kB)\n","Installing collected packages: einops\n","Successfully installed einops-0.4.1\n","\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n","\u001b[0m"]}],"source":["import torch\n","import torch.nn as nn\n","import torchvision\n","import torch.nn.functional as F\n","!pip install einops\n","import einops\n","import torchvision\n","import torchvision.transforms as transforms\n","import cv2\n","import numpy as np\n","from torchvision import datasets\n","from torch.utils.data.sampler import SubsetRandomSampler\n","import os\n","from pathlib import Path\n","import torch\n","from PIL import Image\n","from tqdm import tqdm\n","from torchmetrics import F1Score,Accuracy\n","import matplotlib.pyplot as plt\n","import copy\n","import pytorch_lightning as pl\n","import pandas as pd\n","import seaborn as sn\n","from IPython.core.display import display\n","from torch.utils.data import random_split, DataLoader\n"]},{"cell_type":"code","execution_count":29,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:14.115278Z","iopub.status.busy":"2022-09-27T13:00:14.114825Z","iopub.status.idle":"2022-09-27T13:00:14.125251Z","shell.execute_reply":"2022-09-27T13:00:14.123693Z","shell.execute_reply.started":"2022-09-27T13:00:14.115215Z"},"trusted":true},"outputs":[],"source":["valid_per = 0.20\n","batch_size = 1\n","epochs = 20\n","beta1 = 0.9\n","beta2 = 0.990\n","weight_decay = 0.3\n","embed_size =768\n","num_heads = 8\n","patch_size = 12\n","in_channels = 3\n","num_encoders = 8\n","num_class = 12\n","relaxation_coeff = 0.03\n","device = \"cuda\"\n","batch_size = 8\n","learning_rate = 0.000125\n","epochs = 30\n","label_smooth = 0.1"]},{"cell_type":"code","execution_count":3,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T12:22:21.475535Z","iopub.status.busy":"2022-09-27T12:22:21.474138Z","iopub.status.idle":"2022-09-27T12:22:21.488064Z","shell.execute_reply":"2022-09-27T12:22:21.486692Z","shell.execute_reply.started":"2022-09-27T12:22:21.475505Z"},"trusted":true},"outputs":[],"source":["class Project(nn.Module):\n","    def __init__(self,patch_size:int,in_channels:int,embed_size:int,batch_size:int):\n","        super().__init__()\n","        self.patch_size = patch_size\n","        self.batch_size = batch_size \n","        self.embed_size = embed_size # embed size is the size of linearly projected patch of image\n","        self.in_channels = in_channels # channel size of image, 1 for grayscale, 3 for colored image\n","\n","        self.linear = nn.Linear(self.in_channels*self.patch_size**2,self.embed_size)\n","        self.class_token = nn.Parameter(torch.randn(self.batch_size,1,embed_size))\n","        self.position_embed = nn.Parameter(torch.randn(self.patch_size**2+1,self.embed_size))\n","\n","    def forward(self,x:int): # num_batch x in_channel x width x height -> num_bahch x  num_patch x embed_size\n","        out = einops.rearrange(x,\"b c (h px) (w py) -> b (h w) (c px py)\",px = self.patch_size, py = self.patch_size)\n","        out = self.linear(out)\n","        out = torch.cat([out,self.class_token],dim = 1)\n","        out = out + self.position_embed\n","        return out\n","\n"]},{"cell_type":"code","execution_count":4,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T12:22:21.494195Z","iopub.status.busy":"2022-09-27T12:22:21.492687Z","iopub.status.idle":"2022-09-27T12:22:21.504829Z","shell.execute_reply":"2022-09-27T12:22:21.503219Z","shell.execute_reply.started":"2022-09-27T12:22:21.494149Z"},"trusted":true},"outputs":[],"source":["class DotProductAttention(nn.Module):\n","    def __init__(self,relaxation_coeff,patch_size):\n","        super().__init__()\n","        self.softmax = nn.Softmax(dim = 1)\n","        self.relaxation = Relaxation(relaxation_coeff,patch_size)\n","    def forward(self,query,key,value): \n","        relaxation = self.relaxation(self.softmax(torch.matmul(query,key)/key.size(dim = 2)**(1/2)))\n","        sdp = torch.matmul(relaxation,value)\n","        return sdp"]},{"cell_type":"code","execution_count":5,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T12:22:21.508341Z","iopub.status.busy":"2022-09-27T12:22:21.506655Z","iopub.status.idle":"2022-09-27T12:22:21.523344Z","shell.execute_reply":"2022-09-27T12:22:21.521638Z","shell.execute_reply.started":"2022-09-27T12:22:21.508298Z"},"trusted":true},"outputs":[],"source":["class Relaxation(nn.Module):\n","    def __init__(self,relaxation_coeff:float,patch_size:int):\n","        super().__init__()\n","        self.relaxation_coeff = relaxation_coeff\n","        self.T = (torch.ones((patch_size**2+1,patch_size**2+1))/patch_size).to(\"cuda\")\n","\n","    def forward(self,embed:torch.Tensor):\n","        return ((1-self.relaxation_coeff)*embed + self.relaxation_coeff*self.T)\n"]},{"cell_type":"code","execution_count":30,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:29.320845Z","iopub.status.busy":"2022-09-27T13:00:29.320400Z","iopub.status.idle":"2022-09-27T13:00:29.335049Z","shell.execute_reply":"2022-09-27T13:00:29.333336Z","shell.execute_reply.started":"2022-09-27T13:00:29.320813Z"},"trusted":true},"outputs":[],"source":["class MultiHeadAttention(nn.Module):\n","    def __init__(self,embed_size,num_heads,relaxation_coeff,patch_size,dropout = 0.2):\n","        super().__init__()\n","        self.num_heads = num_heads\n","        self.embed_size = embed_size\n","        self.DotProductAttention = DotProductAttention(relaxation_coeff,patch_size)\n","       \n","\n","        self.key = nn.Linear(self.embed_size,self.num_heads* self.embed_size,bias = False)\n","        self.query = nn.Linear(self.embed_size,self.num_heads* self.embed_size,bias = False)\n","        self.value = nn.Linear(self.embed_size,self.num_heads* self.embed_size,bias = False)\n","        \n","        self.linear = nn.Linear(self.num_heads*self.embed_size,embed_size,bias = False)\n","        self.layer_norm = nn.LayerNorm(self.embed_size, eps=1e-6)\n","        self.dropout = nn.Dropout(dropout)\n","    def forward(self,embed):\n","        batch_size = embed.size(0)\n","        query = self.query(embed)\n","        key = einops.rearrange(self.query(embed),\"b n e ->b e n\")\n","        value = self.value(embed)\n","        sdp = self.DotProductAttention(query,key,value)\n","        return self.linear(self.dropout(sdp))\n","        \n"]},{"cell_type":"code","execution_count":31,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:32.853772Z","iopub.status.busy":"2022-09-27T13:00:32.853310Z","iopub.status.idle":"2022-09-27T13:00:32.865217Z","shell.execute_reply":"2022-09-27T13:00:32.863399Z","shell.execute_reply.started":"2022-09-27T13:00:32.853742Z"},"trusted":true},"outputs":[],"source":["class EncoderBlock(nn.Module):\n","    def __init__(self,embed_size,num_heads,relaxation_coeff,patch_size,dropout = 0.2):\n","        super().__init__()\n","        self.embed_size = embed_size\n","        self.num_heads = num_heads\n","        self.dropout =dropout\n","        self.mha = MultiHeadAttention(768,8,relaxation_coeff,patch_size)\n","        self.Linear1 = nn.Linear(self.embed_size,self.embed_size*4,bias=False)\n","        self.Linear2 = nn.Linear(self.embed_size*4,self.embed_size,bias = False)\n","        self.gelu = nn.GELU()\n","        self.layer_norm1 = nn.LayerNorm(self.embed_size,eps = 1e-6)\n","        self.layer_norm2 = nn.LayerNorm(self.embed_size,eps = 1e-6)\n","        self.dropout = nn.Dropout(dropout)\n","    def forward(self,embed:torch.Tensor):\n","        embed = embed + self.mha(self.layer_norm1(embed))\n","        embed = embed+ self.Linear2(self.dropout(self.gelu(self.Linear1(self.dropout(self.layer_norm2(embed))))))\n","        return embed\n"]},{"cell_type":"code","execution_count":32,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:40.318726Z","iopub.status.busy":"2022-09-27T13:00:40.318299Z","iopub.status.idle":"2022-09-27T13:00:40.334103Z","shell.execute_reply":"2022-09-27T13:00:40.332593Z","shell.execute_reply.started":"2022-09-27T13:00:40.318696Z"},"trusted":true},"outputs":[],"source":["class Transformer(nn.Module):\n","    def __init__(self,embed_size,num_heads,patch_size,in_channels,batch_size,num_encoders,num_class,device,relaxation_coeff):\n","        super().__init__()\n","        self.device = device\n","        self.num_heads = num_heads\n","        self.embed_size = embed_size\n","        self.patch_size =patch_size\n","        self.in_channels = in_channels\n","        self.batch_size = batch_size\n","        self.num_encoders = num_encoders\n","        self.num_class = num_class\n","        self.proj = Project(self.patch_size,self.in_channels,self.embed_size,self.batch_size)\n","\n","        self.tiny_block = [EncoderBlock(self.embed_size,self.num_heads,relaxation_coeff,patch_size) for i in range(self.num_encoders)]\n","        self.block_seq = nn.Sequential(*self.tiny_block)\n","        #self.block_seq = EncoderBlock(self.embed_size,self.num_heads,relaxation_coeff,patch_size)\n","        self.linear1 = nn.Linear(self.embed_size,self.num_class*4)\n","        self.linear2 = nn.Linear(self.num_class*4,self.num_class)\n","        self.layernorm1 = nn.LayerNorm(num_class*4)\n","        self.layernorm2 = nn.LayerNorm(num_class)\n","        self.logsoftmax = nn.LogSoftmax(dim = 0)\n","    def num_of_parameters(self,):\n","\n","        return sum(p.numel() for p in self.parameters())\n","    \n","    def forward(self,img):\n","        out = self.proj(img)\n","        out = self.block_seq(out)\n","        out = self.linear1(torch.squeeze(torch.index_select(out,1,torch.tensor([self.patch_size**2]).to(self.device))))\n","        out = self.layernorm1(out)\n","        out = self.linear2(out)\n","        out = self.layernorm2(out)\n","        out = self.logsoftmax(out)\n","        return out"]},{"cell_type":"code","execution_count":33,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:40.943120Z","iopub.status.busy":"2022-09-27T13:00:40.942676Z","iopub.status.idle":"2022-09-27T13:00:40.958830Z","shell.execute_reply":"2022-09-27T13:00:40.957042Z","shell.execute_reply.started":"2022-09-27T13:00:40.943089Z"},"trusted":true},"outputs":[],"source":["class Custom_Dataset():\n","\n","    def __init__(self, directory):\n","        self.path = Path(directory)\n","        Path.ls = lambda x: list(x.iterdir())\n","        try:\n","            files = os.listdir(directory)\n","        except:\n","            print(\"wrong path\")\n","        self.x = [torch.tensor(np.transpose(np.array(Image.open(img).resize((144,144)))[:, :, :3], (2, 0, 1))).type(\n","            torch.FloatTensor) for img in (self.path/files[0]).ls()]\n","        self.x = torch.stack(self.x)/255\n","        self.y = torch.tensor([0]*len((self.path/files[0]).ls()))\n","        \n","        for i in range(len(files)-1):\n","            try:\n","                self.x2 = [torch.tensor(np.transpose(np.array(Image.open(img).resize((144,144)))[:, :, :3], (2, 0, 1))).type(\n","                torch.FloatTensor) for img in (self.path/files[i+1]).ls()]\n","            except:\n","                return \n","            self.x2 = torch.stack(self.x2)/255\n","            self.x = torch.cat((self.x, self.x2), 0)\n","            self.y = torch.cat((self.y, torch.tensor(\n","                [i+1]*len((self.path/files[i+1]).ls()))))\n","        \n","    def __len__(self):\n","        return len(self.x)\n","    \n","    def __getitem__(self, i):\n","        return self.x[i], self.y[i]\n","    "]},{"cell_type":"code","execution_count":11,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T12:25:08.010708Z","iopub.status.busy":"2022-09-27T12:25:08.009843Z","iopub.status.idle":"2022-09-27T12:26:56.796851Z","shell.execute_reply":"2022-09-27T12:26:56.795455Z","shell.execute_reply.started":"2022-09-27T12:25:08.010676Z"},"trusted":true},"outputs":[],"source":["dataset = Custom_Dataset(\"../input/animal-image-classification-dataset/Animal Image Dataset\")\n","data_train, data_val = random_split(dataset, [len(dataset)-3000, 3000])"]},{"cell_type":"code","execution_count":34,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:46.344296Z","iopub.status.busy":"2022-09-27T13:00:46.343852Z","iopub.status.idle":"2022-09-27T13:00:46.350599Z","shell.execute_reply":"2022-09-27T13:00:46.349067Z","shell.execute_reply.started":"2022-09-27T13:00:46.344240Z"},"trusted":true},"outputs":[],"source":["trainloader =DataLoader(data_train, batch_size=8,drop_last=True)\n","validloader = DataLoader(data_val, batch_size=8,drop_last=True)"]},{"cell_type":"code","execution_count":35,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:46.978717Z","iopub.status.busy":"2022-09-27T13:00:46.977733Z","iopub.status.idle":"2022-09-27T13:00:49.189889Z","shell.execute_reply":"2022-09-27T13:00:49.188492Z","shell.execute_reply.started":"2022-09-27T13:00:46.978668Z"},"trusted":true},"outputs":[],"source":["model = Transformer(embed_size,num_heads,patch_size,in_channels,batch_size,num_encoders,num_class,device,relaxation_coeff).to(device)"]},{"cell_type":"code","execution_count":36,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:49.193241Z","iopub.status.busy":"2022-09-27T13:00:49.192691Z","iopub.status.idle":"2022-09-27T13:00:49.203185Z","shell.execute_reply":"2022-09-27T13:00:49.201541Z","shell.execute_reply.started":"2022-09-27T13:00:49.193180Z"},"trusted":true},"outputs":[],"source":["f1 = F1Score().to(device)"]},{"cell_type":"code","execution_count":37,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:50.879538Z","iopub.status.busy":"2022-09-27T13:00:50.878016Z","iopub.status.idle":"2022-09-27T13:00:50.889163Z","shell.execute_reply":"2022-09-27T13:00:50.887376Z","shell.execute_reply.started":"2022-09-27T13:00:50.879494Z"},"trusted":true},"outputs":[],"source":["class LabelSmoothing_NLLL(nn.Module):\n","    \"\"\"NLL loss with label smoothing.\n","    \"\"\"\n","    def __init__(self, smoothing=0.0):\n","        \"\"\"Constructor for the LabelSmoothing module.\n","        :param smoothing: label smoothing factor\n","        \"\"\"\n","        super(LabelSmoothing_NLLL, self).__init__()\n","        self.confidence = 1.0 - smoothing\n","        self.smoothing = smoothing\n","\n","    def forward(self, x, target):\n","        logprobs = torch.nn.functional.log_softmax(x, dim=-1)\n","        nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))\n","        nll_loss = nll_loss.squeeze(1)\n","        smooth_loss = -logprobs.mean(dim=-1)\n","        loss = self.confidence * nll_loss + self.smoothing * smooth_loss\n","        return loss.mean()"]},{"cell_type":"code","execution_count":38,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:51.249953Z","iopub.status.busy":"2022-09-27T13:00:51.249517Z","iopub.status.idle":"2022-09-27T13:00:51.256559Z","shell.execute_reply":"2022-09-27T13:00:51.255034Z","shell.execute_reply.started":"2022-09-27T13:00:51.249922Z"},"trusted":true},"outputs":[],"source":["criterion = LabelSmoothing_NLLL(smoothing=label_smooth)"]},{"cell_type":"code","execution_count":39,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:00:58.569625Z","iopub.status.busy":"2022-09-27T13:00:58.569214Z","iopub.status.idle":"2022-09-27T13:00:58.578614Z","shell.execute_reply":"2022-09-27T13:00:58.576871Z","shell.execute_reply.started":"2022-09-27T13:00:58.569592Z"},"trusted":true},"outputs":[],"source":["optimizer = torch.optim.Adam(model.parameters(),lr = learning_rate,betas = (beta1,beta2),weight_decay=weight_decay)"]},{"cell_type":"code","execution_count":40,"metadata":{"execution":{"iopub.execute_input":"2022-09-27T13:01:00.595611Z","iopub.status.busy":"2022-09-27T13:01:00.595164Z","iopub.status.idle":"2022-09-27T13:42:47.683506Z","shell.execute_reply":"2022-09-27T13:42:47.681448Z","shell.execute_reply.started":"2022-09-27T13:01:00.595563Z"},"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["Epoch - 1 Started\n","F1 Score for train: 0.20975783467292786, F1 Score for validation: 0.24799999594688416 \n","Epoch: 1 \tTraining Loss: 2.103555 \tValidation Loss: 2.089716\n","Validation loss decreased (inf --> 2.089716).  Saving model ...\n","Epoch - 2 Started\n","F1 Score for train: 0.1905270665884018, F1 Score for validation: 0.21833333373069763 \n","Epoch: 2 \tTraining Loss: 2.090135 \tValidation Loss: 2.088041\n","Validation loss decreased (2.089716 --> 2.088041).  Saving model ...\n","Epoch - 3 Started\n","F1 Score for train: 0.20477208495140076, F1 Score for validation: 0.21466666460037231 \n","Epoch: 3 \tTraining Loss: 2.087483 \tValidation Loss: 2.086394\n","Validation loss decreased (2.088041 --> 2.086394).  Saving model ...\n","Epoch - 4 Started\n","F1 Score for train: 0.2004985809326172, F1 Score for validation: 0.21466666460037231 \n","Epoch: 4 \tTraining Loss: 2.084855 \tValidation Loss: 2.084636\n","Validation loss decreased (2.086394 --> 2.084636).  Saving model ...\n","Epoch - 5 Started\n","F1 Score for train: 0.18660968542099, F1 Score for validation: 0.21466666460037231 \n","Epoch: 5 \tTraining Loss: 2.082845 \tValidation Loss: 2.083203\n","Validation loss decreased (2.084636 --> 2.083203).  Saving model ...\n","Epoch - 6 Started\n","F1 Score for train: 0.1844729334115982, F1 Score for validation: 0.18533332645893097 \n","Epoch: 6 \tTraining Loss: 2.081569 \tValidation Loss: 2.082677\n","Validation loss decreased (2.083203 --> 2.082677).  Saving model ...\n","Epoch - 7 Started\n","F1 Score for train: 0.18696580827236176, F1 Score for validation: 0.18833333253860474 \n","Epoch: 7 \tTraining Loss: 2.080809 \tValidation Loss: 2.080968\n","Validation loss decreased (2.082677 --> 2.080968).  Saving model ...\n","Epoch - 8 Started\n","F1 Score for train: 0.16096866130828857, F1 Score for validation: 0.09300000220537186 \n","Epoch: 8 \tTraining Loss: 2.081071 \tValidation Loss: 2.079854\n","Validation loss decreased (2.080968 --> 2.079854).  Saving model ...\n","Epoch - 9 Started\n","F1 Score for train: 0.13746438920497894, F1 Score for validation: 0.21066667139530182 \n","Epoch: 9 \tTraining Loss: 2.080024 \tValidation Loss: 2.081595\n","Epoch - 10 Started\n","F1 Score for train: 0.1467236429452896, F1 Score for validation: 0.09300000220537186 \n","Epoch: 10 \tTraining Loss: 2.080961 \tValidation Loss: 2.080041\n","Epoch - 11 Started\n","F1 Score for train: 0.13354700803756714, F1 Score for validation: 0.14933332800865173 \n","Epoch: 11 \tTraining Loss: 2.080019 \tValidation Loss: 2.080619\n","Epoch - 12 Started\n","F1 Score for train: 0.14138177037239075, F1 Score for validation: 0.09300000220537186 \n","Epoch: 12 \tTraining Loss: 2.080572 \tValidation Loss: 2.080161\n","Epoch - 13 Started\n","F1 Score for train: 0.14102564752101898, F1 Score for validation: 0.18766666948795319 \n","Epoch: 13 \tTraining Loss: 2.080509 \tValidation Loss: 2.080317\n","Epoch - 14 Started\n","F1 Score for train: 0.13354700803756714, F1 Score for validation: 0.1263333261013031 \n","Epoch: 14 \tTraining Loss: 2.080085 \tValidation Loss: 2.080500\n","Epoch - 15 Started\n","F1 Score for train: 0.1378205120563507, F1 Score for validation: 0.18799999356269836 \n","Epoch: 15 \tTraining Loss: 2.079919 \tValidation Loss: 2.080436\n","Epoch - 16 Started\n","F1 Score for train: 0.1517094075679779, F1 Score for validation: 0.2173333317041397 \n","Epoch: 16 \tTraining Loss: 2.079828 \tValidation Loss: 2.080558\n","Epoch - 17 Started\n","F1 Score for train: 0.15918803215026855, F1 Score for validation: 0.2173333317041397 \n","Epoch: 17 \tTraining Loss: 2.079780 \tValidation Loss: 2.080502\n","Epoch - 18 Started\n","F1 Score for train: 0.1755698025226593, F1 Score for validation: 0.2173333317041397 \n","Epoch: 18 \tTraining Loss: 2.079742 \tValidation Loss: 2.080483\n","Epoch - 19 Started\n","F1 Score for train: 0.1855413168668747, F1 Score for validation: 0.2173333317041397 \n","Epoch: 19 \tTraining Loss: 2.079694 \tValidation Loss: 2.080279\n","Epoch - 20 Started\n","F1 Score for train: 0.18233618140220642, F1 Score for validation: 0.18799999356269836 \n","Epoch: 20 \tTraining Loss: 2.079635 \tValidation Loss: 2.080007\n","Epoch - 21 Started\n"]},{"ename":"KeyboardInterrupt","evalue":"","output_type":"error","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_18/2032847016.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     16\u001b[0m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 18\u001b[0;31m         \u001b[0mtrain_loss\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     19\u001b[0m         \u001b[0mtrain_score\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_score\u001b[0m \u001b[0;34m+\u001b[0m\u001b[0mf1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtarget\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;31mKeyboardInterrupt\u001b[0m: "]}],"source":["device = \"cuda\"\n","valid_loss_min = np.Inf\n","for i in range(epochs):\n","    print(\"Epoch - {} Started\".format(i+1))\n","\n","    train_loss = 0.0\n","    valid_loss = 0.0\n","    train_score = 0.0\n","    val_score = 0.0\n","    model.train()\n","    for data, target in trainloader:\n","        data, target = data.to(device), target.to(device)\n","        optimizer.zero_grad()\n","        output = model(data)\n","        loss = criterion(output, target)\n","        loss.backward()\n","        optimizer.step()\n","        train_loss += loss.item()*data.size(0)\n","        train_score = train_score +f1(output,target)\n","    \n","    model.eval()\n","    \n","    for data, target in validloader:\n","        data, target = data.to(device), target.to(device)\n","        with torch.no_grad():\n","            output = model(data)\n","        loss = criterion(output, target)\n","        valid_loss += loss.item()*data.size(0)\n","        val_score = val_score + f1(output,target)\n","\n","    train_loss = train_loss/len(trainloader.sampler)\n","    valid_loss = valid_loss/len(validloader.sampler)\n","    train_score = batch_size*train_score/len(trainloader.sampler)\n","    val_score = batch_size*val_score/len(validloader.sampler)\n","\n","    print(f\"F1 Score for train: {train_score}, F1 Score for validation: {val_score} \")\n","    print('Epoch: {} \\tTraining Loss: {:.6f} \\tValidation Loss: {:.6f}'.format(\n","        i+1, train_loss, valid_loss))\n","\n","    if valid_loss <= valid_loss_min:\n","        print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(\n","            valid_loss_min,\n","            valid_loss))\n","        best_model_wts = copy.deepcopy(model.state_dict())\n","\n","        valid_loss_min = valid_loss\n","torch.save(best_model_wts, 'model.pt')"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.7.12"},"vscode":{"interpreter":{"hash":"916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"}}},"nbformat":4,"nbformat_minor":4}