{"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":"!pip -q install vit_pytorch linformer","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-03-06T14:30:44.517930Z","iopub.execute_input":"2022-03-06T14:30:44.518540Z","iopub.status.idle":"2022-03-06T14:30:54.028359Z","shell.execute_reply.started":"2022-03-06T14:30:44.518444Z","shell.execute_reply":"2022-03-06T14:30:54.027507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from __future__ import print_function\n\nimport glob\nfrom itertools import chain\nimport os\nimport random\nimport zipfile\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom linformer import Linformer\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, transforms\nfrom tqdm.notebook import tqdm\nimport torch\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt\nfrom torch import Tensor\n\nfrom torchvision.transforms import Compose, Resize, ToTensor\nfrom einops import rearrange, reduce, repeat\nfrom einops.layers.torch import Rearrange, Reduce\nfrom vit_pytorch.efficient import ViT","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:30:54.030596Z","iopub.execute_input":"2022-03-06T14:30:54.031136Z","iopub.status.idle":"2022-03-06T14:30:56.472375Z","shell.execute_reply.started":"2022-03-06T14:30:54.031067Z","shell.execute_reply":"2022-03-06T14:30:56.471648Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"Torch: {torch.__version__}\")","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:30:56.475119Z","iopub.execute_input":"2022-03-06T14:30:56.475524Z","iopub.status.idle":"2022-03-06T14:30:56.480015Z","shell.execute_reply.started":"2022-03-06T14:30:56.475483Z","shell.execute_reply":"2022-03-06T14:30:56.479342Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Training settings\nbatch_size = 64\nepochs = 20\nlr = 3e-5\ngamma = 0.7\nseed = 42","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:30:56.482701Z","iopub.execute_input":"2022-03-06T14:30:56.483205Z","iopub.status.idle":"2022-03-06T14:30:56.489002Z","shell.execute_reply.started":"2022-03-06T14:30:56.483167Z","shell.execute_reply":"2022-03-06T14:30:56.488271Z"},"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":"2022-03-06T14:30:56.490805Z","iopub.execute_input":"2022-03-06T14:30:56.490977Z","iopub.status.idle":"2022-03-06T14:30:56.500931Z","shell.execute_reply.started":"2022-03-06T14:30:56.490956Z","shell.execute_reply":"2022-03-06T14:30:56.500120Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:30:56.502393Z","iopub.execute_input":"2022-03-06T14:30:56.502779Z","iopub.status.idle":"2022-03-06T14:30:56.548730Z","shell.execute_reply.started":"2022-03-06T14:30:56.502744Z","shell.execute_reply":"2022-03-06T14:30:56.547901Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs('data', exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:30:56.551855Z","iopub.execute_input":"2022-03-06T14:30:56.552047Z","iopub.status.idle":"2022-03-06T14:30:56.557371Z","shell.execute_reply.started":"2022-03-06T14:30:56.552025Z","shell.execute_reply":"2022-03-06T14:30:56.556768Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dir = 'data/train'\ntest_dir = 'data/test'","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:30:56.559988Z","iopub.execute_input":"2022-03-06T14:30:56.561887Z","iopub.status.idle":"2022-03-06T14:30:56.566257Z","shell.execute_reply.started":"2022-03-06T14:30:56.561861Z","shell.execute_reply":"2022-03-06T14:30:56.565570Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with zipfile.ZipFile('../input/dogs-vs-cats-redux-kernels-edition/train.zip') as train_zip:\n    train_zip.extractall('data')\n    \nwith zipfile.ZipFile('../input/dogs-vs-cats-redux-kernels-edition/test.zip') as test_zip:\n    test_zip.extractall('data')","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:30:56.567168Z","iopub.execute_input":"2022-03-06T14:30:56.570422Z","iopub.status.idle":"2022-03-06T14:31:13.962703Z","shell.execute_reply.started":"2022-03-06T14:30:56.570387Z","shell.execute_reply":"2022-03-06T14:31:13.961979Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_list = glob.glob(os.path.join(train_dir,'*.jpg'))\ntest_list = glob.glob(os.path.join(test_dir, '*.jpg'))","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:13.966355Z","iopub.execute_input":"2022-03-06T14:31:13.966796Z","iopub.status.idle":"2022-03-06T14:31:14.092337Z","shell.execute_reply.started":"2022-03-06T14:31:13.966755Z","shell.execute_reply":"2022-03-06T14:31:14.091567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"Train Data: {len(train_list)}\")\nprint(f\"Test Data: {len(test_list)}\")","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:14.093575Z","iopub.execute_input":"2022-03-06T14:31:14.093816Z","iopub.status.idle":"2022-03-06T14:31:14.099496Z","shell.execute_reply.started":"2022-03-06T14:31:14.093782Z","shell.execute_reply":"2022-03-06T14:31:14.098796Z"},"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":"2022-03-06T14:31:14.100746Z","iopub.execute_input":"2022-03-06T14:31:14.101659Z","iopub.status.idle":"2022-03-06T14:31:14.119634Z","shell.execute_reply.started":"2022-03-06T14:31:14.101619Z","shell.execute_reply":"2022-03-06T14:31:14.119001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":"2022-03-06T14:31:14.120969Z","iopub.execute_input":"2022-03-06T14:31:14.121331Z","iopub.status.idle":"2022-03-06T14:31:15.734553Z","shell.execute_reply.started":"2022-03-06T14:31:14.121290Z","shell.execute_reply":"2022-03-06T14:31:15.733308Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":"2022-03-06T14:31:15.735539Z","iopub.execute_input":"2022-03-06T14:31:15.735759Z","iopub.status.idle":"2022-03-06T14:31:15.773931Z","shell.execute_reply.started":"2022-03-06T14:31:15.735731Z","shell.execute_reply":"2022-03-06T14:31:15.773288Z"},"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":"2022-03-06T14:31:15.775080Z","iopub.execute_input":"2022-03-06T14:31:15.775405Z","iopub.status.idle":"2022-03-06T14:31:15.781460Z","shell.execute_reply.started":"2022-03-06T14:31:15.775368Z","shell.execute_reply":"2022-03-06T14:31:15.780546Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_transforms = transforms.Compose(\n    [\n        transforms.Resize((224, 224)),\n        transforms.RandomResizedCrop(224),\n        transforms.RandomHorizontalFlip(),\n        transforms.ToTensor(),\n    ]\n)\n\nval_transforms = transforms.Compose(\n    [\n        transforms.Resize((224, 224)),\n        transforms.ToTensor(),\n    ]\n)\n\n\ntest_transforms = transforms.Compose(\n    [\n        transforms.Resize((224, 224)),\n        transforms.ToTensor(),\n    ]\n)","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.782851Z","iopub.execute_input":"2022-03-06T14:31:15.783262Z","iopub.status.idle":"2022-03-06T14:31:15.790525Z","shell.execute_reply.started":"2022-03-06T14:31:15.783227Z","shell.execute_reply":"2022-03-06T14:31:15.789752Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.791820Z","iopub.execute_input":"2022-03-06T14:31:15.792232Z","iopub.status.idle":"2022-03-06T14:31:15.800299Z","shell.execute_reply.started":"2022-03-06T14:31:15.792186Z","shell.execute_reply":"2022-03-06T14:31:15.799459Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_data = CatsDogsDataset(train_list, transform=train_transforms)\nvalid_data = CatsDogsDataset(valid_list, transform=val_transforms)\ntest_data = CatsDogsDataset(test_list, transform=test_transforms)","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.801971Z","iopub.execute_input":"2022-03-06T14:31:15.802609Z","iopub.status.idle":"2022-03-06T14:31:15.808411Z","shell.execute_reply.started":"2022-03-06T14:31:15.802571Z","shell.execute_reply":"2022-03-06T14:31:15.807744Z"},"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":"2022-03-06T14:31:15.809550Z","iopub.execute_input":"2022-03-06T14:31:15.809886Z","iopub.status.idle":"2022-03-06T14:31:15.818028Z","shell.execute_reply.started":"2022-03-06T14:31:15.809823Z","shell.execute_reply":"2022-03-06T14:31:15.817298Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(len(train_data), len(train_loader))","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.821191Z","iopub.execute_input":"2022-03-06T14:31:15.821681Z","iopub.status.idle":"2022-03-06T14:31:15.831270Z","shell.execute_reply.started":"2022-03-06T14:31:15.821648Z","shell.execute_reply":"2022-03-06T14:31:15.830443Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(len(valid_data), len(valid_loader))","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.832371Z","iopub.execute_input":"2022-03-06T14:31:15.832998Z","iopub.status.idle":"2022-03-06T14:31:15.839664Z","shell.execute_reply.started":"2022-03-06T14:31:15.832916Z","shell.execute_reply":"2022-03-06T14:31:15.838867Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class PatchEmbedding(nn.Module):\n    def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768, img_size: int = 224):\n        self.patch_size = patch_size\n        super().__init__()\n        self.projection = nn.Sequential(\n            # using a conv layer instead of a linear one -> performance gains\n            nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),\n            Rearrange('b e (h) (w) -> b (h w) e'),\n        )\n        self.cls_token = nn.Parameter(torch.randn(1,1, emb_size))\n        self.positions = nn.Parameter(torch.randn((img_size // patch_size) **2 + 1, emb_size))\n\n        \n    def forward(self, x: Tensor) -> Tensor:\n        b, _, _, _ = x.shape\n        x = self.projection(x)\n        cls_tokens = repeat(self.cls_token, '() n e -> b n e', b=b)\n        # prepend the cls token to the input\n        x = torch.cat([cls_tokens, x], dim=1)\n        # add position embedding\n        x += self.positions\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.841024Z","iopub.execute_input":"2022-03-06T14:31:15.841641Z","iopub.status.idle":"2022-03-06T14:31:15.851528Z","shell.execute_reply.started":"2022-03-06T14:31:15.841599Z","shell.execute_reply":"2022-03-06T14:31:15.850677Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MultiHeadAttention(nn.Module):\n    def __init__(self, emb_size: int = 768, num_heads: int = 8, dropout: float = 0):\n        super().__init__()\n        self.emb_size = emb_size\n        self.num_heads = num_heads\n        # fuse the queries, keys and values in one matrix\n        self.qkv = nn.Linear(emb_size, emb_size * 3)\n        self.att_drop = nn.Dropout(dropout)\n        self.projection = nn.Linear(emb_size, emb_size)\n        \n    def forward(self, x : Tensor, mask: Tensor = None) -> Tensor:\n        # split keys, queries and values in num_heads\n        qkv = rearrange(self.qkv(x), \"b n (h d qkv) -> (qkv) b h n d\", h=self.num_heads, qkv=3)\n        queries, keys, values = qkv[0], qkv[1], qkv[2]\n        # sum up over the last axis\n        energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys) # batch, num_heads, query_len, key_len\n        if mask is not None:\n            fill_value = torch.finfo(torch.float32).min\n            energy.mask_fill(~mask, fill_value)\n            \n        scaling = self.emb_size ** (1/2)\n        att = F.softmax(energy, dim=-1) / scaling\n        att = self.att_drop(att)\n        # sum up over the third axis\n        out = torch.einsum('bhal, bhlv -> bhav ', att, values)\n        out = rearrange(out, \"b h n d -> b n (h d)\")\n        out = self.projection(out)\n        return out","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.852953Z","iopub.execute_input":"2022-03-06T14:31:15.853434Z","iopub.status.idle":"2022-03-06T14:31:15.864344Z","shell.execute_reply.started":"2022-03-06T14:31:15.853399Z","shell.execute_reply":"2022-03-06T14:31:15.863611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResidualAdd(nn.Module):\n    def __init__(self, fn):\n        super().__init__()\n        self.fn = fn\n        \n    def forward(self, x, **kwargs):\n        res = x\n        x = self.fn(x, **kwargs)\n        x += res\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.865806Z","iopub.execute_input":"2022-03-06T14:31:15.866076Z","iopub.status.idle":"2022-03-06T14:31:15.876176Z","shell.execute_reply.started":"2022-03-06T14:31:15.866033Z","shell.execute_reply":"2022-03-06T14:31:15.875351Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class FeedForwardBlock(nn.Sequential):\n    def __init__(self, emb_size: int, expansion: int = 4, drop_p: float = 0.):\n        super().__init__(\n            nn.Linear(emb_size, expansion * emb_size),\n            nn.GELU(),\n            nn.Dropout(drop_p),\n            nn.Linear(expansion * emb_size, emb_size),\n        )","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.877351Z","iopub.execute_input":"2022-03-06T14:31:15.877739Z","iopub.status.idle":"2022-03-06T14:31:15.885993Z","shell.execute_reply.started":"2022-03-06T14:31:15.877702Z","shell.execute_reply":"2022-03-06T14:31:15.885253Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TransformerEncoderBlock(nn.Sequential):\n    def __init__(self,\n                 emb_size: int = 768,\n                 drop_p: float = 0.,\n                 forward_expansion: int = 4,\n                 forward_drop_p: float = 0.,\n                 ** kwargs):\n        super().__init__(\n            ResidualAdd(nn.Sequential(\n                nn.LayerNorm(emb_size),\n                MultiHeadAttention(emb_size, **kwargs),\n                nn.Dropout(drop_p)\n            )),\n            ResidualAdd(nn.Sequential(\n                nn.LayerNorm(emb_size),\n                FeedForwardBlock(\n                    emb_size, expansion=forward_expansion, drop_p=forward_drop_p),\n                nn.Dropout(drop_p)\n            )\n            ))","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.887642Z","iopub.execute_input":"2022-03-06T14:31:15.887920Z","iopub.status.idle":"2022-03-06T14:31:15.895118Z","shell.execute_reply.started":"2022-03-06T14:31:15.887885Z","shell.execute_reply":"2022-03-06T14:31:15.894156Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TransformerEncoder(nn.Sequential):\n    def __init__(self, depth: int = 12, **kwargs):\n        super().__init__(*[TransformerEncoderBlock(**kwargs) for _ in range(depth)])","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.896998Z","iopub.execute_input":"2022-03-06T14:31:15.897340Z","iopub.status.idle":"2022-03-06T14:31:15.905278Z","shell.execute_reply.started":"2022-03-06T14:31:15.897303Z","shell.execute_reply":"2022-03-06T14:31:15.904585Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ClassificationHead(nn.Sequential):\n    def __init__(self, emb_size: int = 768, n_classes: int = 1000):\n        super().__init__(\n            Reduce('b n e -> b e', reduction='mean'),\n            nn.LayerNorm(emb_size), \n            nn.Linear(emb_size, n_classes))","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.909489Z","iopub.execute_input":"2022-03-06T14:31:15.909717Z","iopub.status.idle":"2022-03-06T14:31:15.916270Z","shell.execute_reply.started":"2022-03-06T14:31:15.909693Z","shell.execute_reply":"2022-03-06T14:31:15.915425Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ViT(nn.Sequential):\n    def __init__(self,     \n                in_channels: int = 3,\n                patch_size: int = 32,\n                emb_size: int = 768,\n                img_size: int = 224,\n                depth: int = 12,\n                n_classes: int = 2,\n                **kwargs):\n        super().__init__(\n            PatchEmbedding(in_channels, patch_size, emb_size, img_size),\n            TransformerEncoder(depth, emb_size=emb_size, **kwargs),\n            ClassificationHead(emb_size, n_classes)\n        )","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.917631Z","iopub.execute_input":"2022-03-06T14:31:15.917884Z","iopub.status.idle":"2022-03-06T14:31:15.925228Z","shell.execute_reply.started":"2022-03-06T14:31:15.917850Z","shell.execute_reply":"2022-03-06T14:31:15.924282Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = ViT().to(device)","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:15.926564Z","iopub.execute_input":"2022-03-06T14:31:15.926889Z","iopub.status.idle":"2022-03-06T14:31:19.965374Z","shell.execute_reply.started":"2022-03-06T14:31:15.926854Z","shell.execute_reply":"2022-03-06T14:31:19.964632Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# loss function\ncriterion = nn.CrossEntropyLoss()\n# optimizer\noptimizer = optim.Adam(model.parameters(), lr=lr)\n# scheduler\nscheduler = StepLR(optimizer, step_size=1, gamma=gamma)","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:31:19.966672Z","iopub.execute_input":"2022-03-06T14:31:19.966910Z","iopub.status.idle":"2022-03-06T14:31:19.972723Z","shell.execute_reply.started":"2022-03-06T14:31:19.966877Z","shell.execute_reply":"2022-03-06T14:31:19.972073Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for epoch in range(epochs):\n    epoch_loss = 0\n    epoch_accuracy = 0\n\n    for data, label in tqdm(train_loader):\n        data = data.to(device)\n        label = label.to(device)\n\n        output = model(data)\n        loss = criterion(output, label)\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        acc = (output.argmax(dim=1) == label).float().mean()\n        epoch_accuracy += acc / len(train_loader)\n        epoch_loss += loss / len(train_loader)\n\n    with torch.no_grad():\n        epoch_val_accuracy = 0\n        epoch_val_loss = 0\n        for data, label in tqdm(valid_loader):\n            data = data.to(device)\n            label = label.to(device)\n\n            val_output = model(data)\n            val_loss = criterion(val_output, label)\n\n            acc = (val_output.argmax(dim=1) == label).float().mean()\n            epoch_val_accuracy += acc / len(valid_loader)\n            epoch_val_loss += val_loss / len(valid_loader)\n\n    print(\n        f\"Epoch : {epoch+1} - loss : {epoch_loss:.4f} - acc: {epoch_accuracy:.4f} - val_loss : {epoch_val_loss:.4f} - val_acc: {epoch_val_accuracy:.4f}\\n\"\n    )","metadata":{"execution":{"iopub.status.busy":"2022-03-06T14:44:36.875971Z","iopub.execute_input":"2022-03-06T14:44:36.876598Z","iopub.status.idle":"2022-03-06T14:46:02.637848Z","shell.execute_reply.started":"2022-03-06T14:44:36.876556Z","shell.execute_reply":"2022-03-06T14:46:02.636674Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}