{"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":"markdown","source":"# Overview\n\nThe transformer is a general and powerful neural network architecture, able to tackle many real-world problems. Vision transformer (ViT) is a transformer for computer vision tasks. In this notebook, Vision Transformer (ViT) is implemented from scratch using PyTorch for image classification. Later, we will train the model on a subset of RSNA breast cancer detection dataset.","metadata":{}},{"cell_type":"markdown","source":"# 1. About the Architecture\nTransformers found their initial applications in natural language processing (NLP) tasks. To use this NLP model for computer vision tasks, we have to divide our input image into patches. After flattening the patches, we can treat each flattened patches as single word. We add positional embeddings to the linear projection of flattened patches. An extra token is added at the beginning for classification tasks. In BERT model, this token is called [CLS] token.\n\nSo if our input image size is (512, 512), after dividing the image into patches of size (16, 16), we get 1024 (32 times 32) patches. After flattening the patches and projecting the flattened patches, we have 1024 tokens. After adding positional embeddings and concatenating classification token at the beginning, we have 1025 tokens.\n\nWe then feed our tokens into the transformer encoder. Transformer encoder is made up of self attention and feedforward network. The [original paper](https://paperswithcode.com/paper/attention-is-all-you-need) on attention is an excellent read if you want to understand the whole attention mechanism. PyTorch have `torch.nn.MultiHeadAttention` for anyone who one to use attention mechanism for their next project. This [video](https://www.youtube.com/watch?v=_UVfwBqcnbM) by AssemblyAI is explains the transformer architecture beautifully.\n\nThe number of tokens in the output of the transformer encoder is equal to number of input tokens. We take the first token from the output (corresponds to the classification token) and feed the token in a multilayer perceptron head for classification.\n\nFor more details, you can go through [original paper](https://paperswithcode.com/method/vision-transformer) on Vision Transformer.\n\n![Vision Transformer](https://production-media.paperswithcode.com/methods/Screen_Shot_2021-01-26_at_9.43.31_PM_uI4jjMq.png)","metadata":{}},{"cell_type":"markdown","source":"# 2. About the Dataset\nThe dataset was contributed by mammography screening programs in Australia and the U.S. It includes detailed labels, with radiologists’ evaluations and follow-up pathology results for suspected malignancies. \n\nThe dataset is stored in dicom formats. Converting dicom data to png/jpg just by rescaling it will harm the quality of the data. [This notebook](https://www.kaggle.com/code/raddar/convert-dicom-to-np-array-the-correct-way/notebook) is an awesome resource for anyone working with dicom files for X-Ray.","metadata":{}},{"cell_type":"markdown","source":"# 3. Implementation","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"markdown","source":"## 3.1 Utility functions\nWe write some utility functions beforehand.","metadata":{}},{"cell_type":"code","source":"import os\n\nimport cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pydicom\nimport pandas as pd\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Adam\nfrom torch.utils.data import DataLoader, Dataset, WeightedRandomSampler\nfrom ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator\nfrom ignite.metrics import Accuracy, Loss, RunningAverage\nfrom ignite.contrib.handlers import ProgressBar\nfrom sklearn.model_selection import train_test_split\nfrom torchvision import models, transforms","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:31:42.288365Z","iopub.execute_input":"2023-03-03T07:31:42.288731Z","iopub.status.idle":"2023-03-03T07:31:44.830314Z","shell.execute_reply.started":"2023-03-03T07:31:42.288682Z","shell.execute_reply":"2023-03-03T07:31:44.829367Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_xray(file_path, img_size=None):\n    \"\"\"\n    Read the dicom data and get the image\n    Args:\n        file_path: The path of the dicom file\n        img_size: Size of the output image\n    \"\"\"\n\n    dicom = pydicom.read_file(file_path)\n    img = dicom.pixel_array\n\n    if dicom.PhotometricInterpretation == \"MONOCHROME1\":\n        img = np.max(img) - img\n\n    if img_size:\n        img = cv2.resize(img, img_size)\n\n    # Add channel dim at First\n    img = img[np.newaxis]\n\n    # Converting img to float32\n    img = img / np.max(img)\n    img = img.astype(\"float32\")\n\n    return img","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:31:44.833175Z","iopub.execute_input":"2023-03-03T07:31:44.834093Z","iopub.status.idle":"2023-03-03T07:31:44.840824Z","shell.execute_reply.started":"2023-03-03T07:31:44.834055Z","shell.execute_reply":"2023-03-03T07:31:44.83966Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def patchify(batch, patch_size):\n    \"\"\"\n    Patchify the batch of images\n        \n    Shape:\n        batch: (b, h, w, c)\n        output: (b, nh, nw, ph, pw, c)\n    \"\"\"\n    b, c, h, w = batch.shape\n    ph, pw = patch_size\n    nh, nw = h // ph, w // pw\n\n    batch_patches = torch.reshape(batch, (b, c, nh, ph, nw, pw))\n    batch_patches = torch.permute(batch_patches, (0, 1, 2, 4, 3, 5))\n\n    return batch_patches","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:31:44.843969Z","iopub.execute_input":"2023-03-03T07:31:44.844247Z","iopub.status.idle":"2023-03-03T07:31:44.859436Z","shell.execute_reply.started":"2023-03-03T07:31:44.844223Z","shell.execute_reply":"2023-03-03T07:31:44.858353Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We test our `patchify` function on a single image.","metadata":{}},{"cell_type":"code","source":"FILE_PATH = ('/kaggle/input/rsna-breast-cancer-detection/'\n             'train_images/10006/1459541791.dcm')\n\nimg = read_xray(FILE_PATH, img_size=(512, 512))\n\nbatch = torch.tensor(img[None])\npatch_size = (16, 16)\nbatch_patches = patchify(batch, patch_size)\n\npatches = batch_patches[0]\nc, nh, nw, ph, pw = patches.shape\n\nplt.figure(figsize=(5, 5))\nplt.imshow(img[0], cmap=\"gray\")\nplt.axis(\"off\")\n\nplt.figure(figsize=(5, 5))\nfor i in range(nh):\n    for j in range(nw):\n        plt.subplot(nh, nw, i * nw + j + 1)\n        plt.imshow(patches[0, i, j], cmap=\"gray\")\n        plt.axis(\"off\")","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:31:45.959129Z","iopub.execute_input":"2023-03-03T07:31:45.959502Z","iopub.status.idle":"2023-03-03T07:32:36.608991Z","shell.execute_reply.started":"2023-03-03T07:31:45.959472Z","shell.execute_reply":"2023-03-03T07:32:36.607892Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_mlp(in_features, hidden_units, out_features):\n    \"\"\"\n    Returns a MLP head\n    \"\"\"\n    dims = [in_features] + hidden_units + [out_features]\n    layers = []\n    for dim1, dim2 in zip(dims[:-2], dims[1:-1]):\n        layers.append(nn.Linear(dim1, dim2))\n        layers.append(nn.ReLU())\n    layers.append(nn.Linear(dims[-2], dims[-1]))\n    return nn.Sequential(*layers)","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:32:40.236402Z","iopub.execute_input":"2023-03-03T07:32:40.236817Z","iopub.status.idle":"2023-03-03T07:32:40.243443Z","shell.execute_reply.started":"2023-03-03T07:32:40.236784Z","shell.execute_reply":"2023-03-03T07:32:40.242356Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3.2 Image to Sequence Block\nThis Block takes a batch of image as input and returns a batch of sequences. Later on we feed this sequences into the transformer encoder.","metadata":{}},{"cell_type":"code","source":"class Img2Seq(nn.Module):\n    \"\"\"\n    This layers takes a batch of images as input and\n    returns a batch of sequences\n    \n    Shape:\n        input: (b, h, w, c)\n        output: (b, s, d)\n    \"\"\"\n    def __init__(self, img_size, patch_size, n_channels, d_model):\n        super().__init__()\n        self.patch_size = patch_size\n        self.img_size = img_size\n\n        nh, nw = img_size[0] // patch_size[0], img_size[1] // patch_size[1]\n        n_tokens = nh * nw\n\n        token_dim = patch_size[0] * patch_size[1] * n_channels\n        self.linear = nn.Linear(token_dim, d_model)\n        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model))\n        self.pos_emb = nn.Parameter(torch.randn(n_tokens, d_model))\n\n    def __call__(self, batch):\n        batch = patchify(batch, self.patch_size)\n\n        b, c, nh, nw, ph, pw = batch.shape\n\n        # Flattening the patches\n        batch = torch.permute(batch, [0, 2, 3, 4, 5, 1])\n        batch = torch.reshape(batch, [b, nh * nw, ph * pw * c])\n\n        batch = self.linear(batch)\n        cls = self.cls_token.expand([b, -1, -1])\n        emb = batch + self.pos_emb\n\n        return torch.cat([cls, emb], axis=1)","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:32:41.815532Z","iopub.execute_input":"2023-03-03T07:32:41.816235Z","iopub.status.idle":"2023-03-03T07:32:41.826932Z","shell.execute_reply.started":"2023-03-03T07:32:41.816179Z","shell.execute_reply":"2023-03-03T07:32:41.82541Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3.3 Visual Transformer Module\nThis modules wraps up everything. We can divide this module into 3 parts: \n* An image to sequence encoder\n* Transformer encoder\n* Multilayer perceptron head classification\n\nWe use `torch.nn.TransformerEncoder` and `torch.nn.TransformerEncoderLayer` to implement our transformer encoder. I highly recommend to read the official documentation to learn more about the layers.","metadata":{}},{"cell_type":"code","source":"class ViT(nn.Module):\n    def __init__(\n        self,\n        img_size,\n        patch_size,\n        n_channels,\n        d_model,\n        nhead,\n        dim_feedforward,\n        blocks,\n        mlp_head_units,\n        n_classes,\n    ):\n        super().__init__()\n        \"\"\"\n        Args:\n            img_size: Size of the image\n            patch_size: Size of the patch\n            n_channels: Number of image channels\n            d_model: The number of features in the transformer encoder\n            nhead: The number of heads in the multiheadattention models\n            dim_feedforward: The dimension of the feedforward network model in the encoder\n            blocks: The number of sub-encoder-layers in the encoder\n            mlp_head_units: The hidden units of mlp_head\n            n_classes: The number of output classes\n        \"\"\"\n        self.img2seq = Img2Seq(img_size, patch_size, n_channels, d_model)\n\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model, nhead, dim_feedforward, activation=\"gelu\", batch_first=True\n        )\n        self.transformer_encoder = nn.TransformerEncoder(\n            encoder_layer, blocks\n        )\n        self.mlp = get_mlp(d_model, mlp_head_units, n_classes)\n        \n        self.output = nn.Sigmoid() if n_classes == 1 else nn.Softmax()\n\n    def __call__(self, batch):\n\n        batch = self.img2seq(batch)\n        batch = self.transformer_encoder(batch)\n        batch = batch[:, 0, :]\n        batch = self.mlp(batch)\n        output = self.output(batch)\n        return output","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:32:43.525101Z","iopub.execute_input":"2023-03-03T07:32:43.525452Z","iopub.status.idle":"2023-03-03T07:32:43.545769Z","shell.execute_reply.started":"2023-03-03T07:32:43.525424Z","shell.execute_reply":"2023-03-03T07:32:43.544431Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 4. Training\n\nHere, we design a simple training to loop to train or `ViT` model on a subset of dataset. We use an already cropped dataset.","metadata":{}},{"cell_type":"markdown","source":"## 4.1 Model Hyperparameters","metadata":{}},{"cell_type":"code","source":"img_size = (512, 512)\npatch_size = (16, 16)\nn_channels = 1\nd_model = 1024\nnhead = 4\ndim_feedforward = 2048\nblocks = 8\nmlp_head_units = [1024, 512]\nn_classes = 1\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:32:47.952667Z","iopub.execute_input":"2023-03-03T07:32:47.953063Z","iopub.status.idle":"2023-03-03T07:32:48.020092Z","shell.execute_reply.started":"2023-03-03T07:32:47.95303Z","shell.execute_reply":"2023-03-03T07:32:48.018718Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4.2. Dataset and DataLoader","metadata":{}},{"cell_type":"code","source":"class RSNADataset(Dataset):\n    \n    def __init__(self, df, img_path):\n        self.df = df\n        self.img_path = img_path\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        patient_id, image_id, cancer = self.df.iloc[idx][['patient_id', 'image_id', 'cancer']]\n        file = os.path.join(self.img_path, f'{patient_id}_{image_id}.png')\n        file = cv2.imread(file, cv2.COLOR_BGR2GRAY)\n        clahe = cv2.createCLAHE(clipLimit = 15, tileGridSize=[8, 8])\n        file = clahe.apply(file)\n        file = file / file.max()\n        X = torch.tensor(file[np.newaxis].astype('float32')).to(device)\n        y = torch.tensor([cancer]).float().to(device)\n        return X, y","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:32:51.419536Z","iopub.execute_input":"2023-03-03T07:32:51.419903Z","iopub.status.idle":"2023-03-03T07:32:51.42824Z","shell.execute_reply.started":"2023-03-03T07:32:51.419873Z","shell.execute_reply":"2023-03-03T07:32:51.427088Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv('/kaggle/input/rsna-breast-cancer-detection/train.csv')\ncounts = df['cancer'].value_counts()\ndf['weights'] = df['cancer'].apply(lambda x: 1/counts[x])\n\ntrain_df, val_df = train_test_split(df, test_size=0.25, stratify=df['cancer'])","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:32:53.796909Z","iopub.execute_input":"2023-03-03T07:32:53.797278Z","iopub.status.idle":"2023-03-03T07:32:54.186459Z","shell.execute_reply.started":"2023-03-03T07:32:53.797247Z","shell.execute_reply":"2023-03-03T07:32:54.185507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_path = '/kaggle/input/rsna-breast-cancer-512-pngs'\ntrain_samples = 1000\nval_samples = 500\n\ntrain_ds = RSNADataset(train_df, img_path)\nval_ds = RSNADataset(val_df, img_path)\n\ntrain_sampler = WeightedRandomSampler(train_df['weights'].values, train_samples)\ntrain_loader = DataLoader(train_ds, batch_size=8, sampler=train_sampler)\n\nval_sampler = WeightedRandomSampler(val_df['weights'].values, val_samples)\nval_loader = DataLoader(val_ds, batch_size=32, sampler=val_sampler)","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:33:05.1773Z","iopub.execute_input":"2023-03-03T07:33:05.17792Z","iopub.status.idle":"2023-03-03T07:33:05.185847Z","shell.execute_reply.started":"2023-03-03T07:33:05.177885Z","shell.execute_reply":"2023-03-03T07:33:05.184688Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4.3 Training and Validation","metadata":{}},{"cell_type":"code","source":"model = ViT(\n    img_size = (512, 512),\n    patch_size = (16, 16),\n    n_channels = 1,\n    d_model = 1024,\n    nhead = 4,\n    dim_feedforward = 1024,\n    blocks = 8,\n    mlp_head_units = [512, 512],\n    n_classes = 1,\n).to(device)\n\noptimizer = Adam(model.parameters())\ncriterion = nn.BCELoss()\n\ntrainer = create_supervised_trainer(model, optimizer, criterion, device=device)\nval_metrics = {\n    \"bce\": Loss(criterion)\n}\nevaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:33:07.839539Z","iopub.execute_input":"2023-03-03T07:33:07.840096Z","iopub.status.idle":"2023-03-03T07:33:10.737801Z","shell.execute_reply.started":"2023-03-03T07:33:07.840052Z","shell.execute_reply":"2023-03-03T07:33:10.73675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"log_interval = 10\nmax_epochs = 5\nbest_loss = float('inf')\n\nRunningAverage(output_transform=lambda x: x).attach(trainer, 'loss')\n\npbar = ProgressBar()\npbar.attach(trainer, ['loss'])\n    \n@trainer.on(Events.EPOCH_COMPLETED)\ndef log_validation_results(trainer):\n    global best_loss\n    evaluator.run(val_loader)\n    loss = evaluator.state.metrics['bce']\n    if loss < best_loss:\n        best_loss = loss\n        torch.save(model.state_dict(), 'best_model_vit.pt')\n    print(f\"Validation Results - Epoch: {trainer.state.epoch} Avg loss: {loss:.2f}\")\n    \noutput_state = trainer.run(train_loader, max_epochs=max_epochs)","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:33:10.739788Z","iopub.execute_input":"2023-03-03T07:33:10.740215Z","iopub.status.idle":"2023-03-03T07:41:19.589737Z","shell.execute_reply.started":"2023-03-03T07:33:10.740178Z","shell.execute_reply":"2023-03-03T07:41:19.588417Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4.4 Comparison with resnet50","metadata":{}},{"cell_type":"code","source":"resnet = models.resnet101(pretrained=True)\nin_features = resnet.fc.in_features\nresnet.fc = nn.Linear(in_features, 1)\n\nresnet_transforms= transforms.Compose([\n    transforms.Resize((228, 228)),\n    transforms.Lambda(lambda x: x.repeat([1, 3, 1, 1]))\n])\n\nclass MyResNet(nn.Module):\n    \n    def __init__(self, transforms, model):\n        super().__init__()\n        self.transforms = transforms\n        self.model = model\n        self.output = nn.Sigmoid()\n        \n    def forward(self, batch):\n        batch = self.transforms(batch)\n        batch = self.model(batch)\n        return self.output(batch)","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:41:19.592281Z","iopub.execute_input":"2023-03-03T07:41:19.592758Z","iopub.status.idle":"2023-03-03T07:41:30.50933Z","shell.execute_reply.started":"2023-03-03T07:41:19.592694Z","shell.execute_reply":"2023-03-03T07:41:30.508272Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = MyResNet(resnet_transforms, resnet).to(device)\n\noptimizer = Adam(model.parameters())\ncriterion = nn.BCELoss()\n\ntrainer = create_supervised_trainer(model, optimizer, criterion, device=device)\nval_metrics = {\n    \"bce\": Loss(criterion)\n}\nevaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:41:30.511183Z","iopub.execute_input":"2023-03-03T07:41:30.511833Z","iopub.status.idle":"2023-03-03T07:41:30.597265Z","shell.execute_reply.started":"2023-03-03T07:41:30.511792Z","shell.execute_reply":"2023-03-03T07:41:30.595996Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"RunningAverage(output_transform=lambda x: x).attach(trainer, 'loss')\n\npbar = ProgressBar()\npbar.attach(trainer, ['loss'])\n    \n@trainer.on(Events.EPOCH_COMPLETED)\ndef log_validation_results(trainer):\n    global best_loss\n    evaluator.run(val_loader)\n    loss = evaluator.state.metrics['bce']\n    if loss < best_loss:\n        best_loss = loss\n        torch.save(model.state_dict(), 'best_model_resnet.pt')\n    print(f\"Validation Results - Epoch: {trainer.state.epoch} Avg loss: {loss:.2f}\")\n    \noutput_state = trainer.run(train_loader, max_epochs=max_epochs)","metadata":{"execution":{"iopub.status.busy":"2023-03-03T07:41:30.599634Z","iopub.execute_input":"2023-03-03T07:41:30.600034Z","iopub.status.idle":"2023-03-03T07:44:25.708895Z","shell.execute_reply.started":"2023-03-03T07:41:30.599996Z","shell.execute_reply":"2023-03-03T07:44:25.707591Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Next Step\n\n* You can checkout [huggingface](httfc://huggingface.co/docs/transformers/model_doc/vit) for pretrained vision transformers.\n* For implementation of ViT in tensorflow and keras, checkout this [tutorial](https://keras.io/examples/vision/image_classification_with_vision_transformer/).\n\nUpvote the notebook if you like it. Feel free to give feedback in the comment section down below!","metadata":{}}]}