{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":1677248,"sourceType":"datasetVersion","datasetId":993628},{"sourceId":1695720,"sourceType":"datasetVersion","datasetId":1003815}],"dockerImageVersionId":30299,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!cp -r ../input/vittutorialillustrations/* ./ \n\n!pip install nb_black\n%load_ext nb_black","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:07.340289Z","iopub.execute_input":"2024-12-23T12:21:07.341309Z","iopub.status.idle":"2024-12-23T12:21:23.617247Z","shell.execute_reply.started":"2024-12-23T12:21:07.341165Z","shell.execute_reply":"2024-12-23T12:21:23.61581Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version 1.7\n!pip install timm","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:23.62002Z","iopub.execute_input":"2024-12-23T12:21:23.620449Z","iopub.status.idle":"2024-12-23T12:21:38.933994Z","shell.execute_reply.started":"2024-12-23T12:21:23.620408Z","shell.execute_reply":"2024-12-23T12:21:38.932262Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nplt.style.use(\"ggplot\")\n\nimport torch\nimport torch.nn as nn\nimport torchvision.transforms as transforms\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.distributed.parallel_loader as pl\n\nimport timm\n\nimport gc\nimport os\nimport time\nimport random\nfrom datetime import datetime\n\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\nfrom sklearn import model_selection, metrics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:38.936159Z","iopub.execute_input":"2024-12-23T12:21:38.936614Z","iopub.status.idle":"2024-12-23T12:21:41.294189Z","shell.execute_reply.started":"2024-12-23T12:21:38.936571Z","shell.execute_reply":"2024-12-23T12:21:41.291041Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# For parallelization in TPUs\nos.environ[\"XLA_USE_BF16\"] = \"1\"\nos.environ[\"XLA_TENSOR_ALLOCATOR_MAXSIZE\"] = \"100000000\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.295099Z","iopub.status.idle":"2024-12-23T12:21:41.295552Z","shell.execute_reply.started":"2024-12-23T12:21:41.295291Z","shell.execute_reply":"2024-12-23T12:21:41.295308Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def seed_everything(seed):\n    \"\"\"\n    Seeds basic parameters for reproductibility of results\n    \n    Arguments:\n        seed {int} -- Number of the seed\n    \"\"\"\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.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n\nseed_everything(1001)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.297469Z","iopub.status.idle":"2024-12-23T12:21:41.297879Z","shell.execute_reply.started":"2024-12-23T12:21:41.297689Z","shell.execute_reply":"2024-12-23T12:21:41.297707Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# general global variables\nDATA_PATH = \"../input/cassava-leaf-disease-classification\"\nTRAIN_PATH = \"../input/cassava-leaf-disease-classification/train_images/\"\nTEST_PATH = \"../input/cassava-leaf-disease-classification/test_images/\"\nMODEL_PATH = (\n    \"../input/vit-base-models-pretrained-pytorch/jx_vit_base_p16_224-80ecf9dd.pth\"\n)\n\n# model specific global variables\nIMG_SIZE = 224\nBATCH_SIZE = 16\nLR = 2e-05\nGAMMA = 0.7\nN_EPOCHS = 10","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.299388Z","iopub.status.idle":"2024-12-23T12:21:41.299794Z","shell.execute_reply.started":"2024-12-23T12:21:41.29961Z","shell.execute_reply":"2024-12-23T12:21:41.299629Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df = pd.read_csv(os.path.join(DATA_PATH, \"train.csv\"))\ndf.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.30109Z","iopub.status.idle":"2024-12-23T12:21:41.30163Z","shell.execute_reply.started":"2024-12-23T12:21:41.301322Z","shell.execute_reply":"2024-12-23T12:21:41.301348Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df.info()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.30456Z","iopub.status.idle":"2024-12-23T12:21:41.304917Z","shell.execute_reply.started":"2024-12-23T12:21:41.304748Z","shell.execute_reply":"2024-12-23T12:21:41.304764Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df.label.value_counts().plot(kind=\"bar\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.306055Z","iopub.status.idle":"2024-12-23T12:21:41.306439Z","shell.execute_reply.started":"2024-12-23T12:21:41.306227Z","shell.execute_reply":"2024-12-23T12:21:41.306243Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df, valid_df = model_selection.train_test_split(\n    df, test_size=0.1, random_state=42, stratify=df.label.values\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.30834Z","iopub.status.idle":"2024-12-23T12:21:41.308749Z","shell.execute_reply.started":"2024-12-23T12:21:41.308574Z","shell.execute_reply":"2024-12-23T12:21:41.308591Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CassavaDataset(torch.utils.data.Dataset):\n    \"\"\"\n    Helper Class to create the pytorch dataset\n    \"\"\"\n\n    def __init__(self, df, data_path=DATA_PATH, mode=\"train\", transforms=None):\n        super().__init__()\n        self.df_data = df.values\n        self.data_path = data_path\n        self.transforms = transforms\n        self.mode = mode\n        self.data_dir = \"train_images\" if mode == \"train\" else \"test_images\"\n\n    def __len__(self):\n        return len(self.df_data)\n\n    def __getitem__(self, index):\n        img_name, label = self.df_data[index]\n        img_path = os.path.join(self.data_path, self.data_dir, img_name)\n        img = Image.open(img_path).convert(\"RGB\")\n\n        if self.transforms is not None:\n            image = self.transforms(img)\n\n        return image, label","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.310034Z","iopub.status.idle":"2024-12-23T12:21:41.310573Z","shell.execute_reply.started":"2024-12-23T12:21:41.310277Z","shell.execute_reply":"2024-12-23T12:21:41.310301Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# create image augmentations\ntransforms_train = transforms.Compose(\n    [\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.RandomHorizontalFlip(p=0.3),\n        transforms.RandomVerticalFlip(p=0.3),\n        transforms.RandomResizedCrop(IMG_SIZE),\n        transforms.ToTensor(),\n        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n    ]\n)\n\ntransforms_valid = transforms.Compose(\n    [\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.ToTensor(),\n        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n    ]\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.312416Z","iopub.status.idle":"2024-12-23T12:21:41.312979Z","shell.execute_reply.started":"2024-12-23T12:21:41.312699Z","shell.execute_reply":"2024-12-23T12:21:41.312724Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Available Vision Transformer Models: \")\ntimm.list_models(\"vit*\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.315485Z","iopub.status.idle":"2024-12-23T12:21:41.316031Z","shell.execute_reply.started":"2024-12-23T12:21:41.315753Z","shell.execute_reply":"2024-12-23T12:21:41.315777Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ViTBase16(nn.Module):\n    def __init__(self, n_classes, pretrained=False):\n\n        super(ViTBase16, self).__init__()\n\n        self.model = timm.create_model(\"vit_base_patch16_224\", pretrained=False)\n        if pretrained:\n            self.model.load_state_dict(torch.load(MODEL_PATH))\n\n        self.model.head = nn.Linear(self.model.head.in_features, n_classes)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x\n\n    def train_one_epoch(self, train_loader, criterion, optimizer, device):\n        # keep track of training loss\n        epoch_loss = 0.0\n        epoch_accuracy = 0.0\n\n        ###################\n        # train the model #\n        ###################\n        self.model.train()\n        for i, (data, target) in enumerate(train_loader):\n            # move tensors to GPU if CUDA is available\n            if device.type == \"cuda\":\n                data, target = data.cuda(), target.cuda()\n            elif device.type == \"xla\":\n                data = data.to(device, dtype=torch.float32)\n                target = target.to(device, dtype=torch.int64)\n\n            # clear the gradients of all optimized variables\n            optimizer.zero_grad()\n            # forward pass: compute predicted outputs by passing inputs to the model\n            output = self.forward(data)\n            # calculate the batch loss\n            loss = criterion(output, target)\n            # backward pass: compute gradient of the loss with respect to model parameters\n            loss.backward()\n            # Calculate Accuracy\n            accuracy = (output.argmax(dim=1) == target).float().mean()\n            # update training loss and accuracy\n            epoch_loss += loss\n            epoch_accuracy += accuracy\n\n            # perform a single optimization step (parameter update)\n            if device.type == \"xla\":\n                xm.optimizer_step(optimizer)\n\n                if i % 20 == 0:\n                    xm.master_print(f\"\\tBATCH {i+1}/{len(train_loader)} - LOSS: {loss}\")\n\n            else:\n                optimizer.step()\n\n        return epoch_loss / len(train_loader), epoch_accuracy / len(train_loader)\n\n    def validate_one_epoch(self, valid_loader, criterion, device):\n        # keep track of validation loss\n        valid_loss = 0.0\n        valid_accuracy = 0.0\n\n        ######################\n        # validate the model #\n        ######################\n        self.model.eval()\n        for data, target in valid_loader:\n            # move tensors to GPU if CUDA is available\n            if device.type == \"cuda\":\n                data, target = data.cuda(), target.cuda()\n            elif device.type == \"xla\":\n                data = data.to(device, dtype=torch.float32)\n                target = target.to(device, dtype=torch.int64)\n\n            with torch.no_grad():\n                # forward pass: compute predicted outputs by passing inputs to the model\n                output = self.model(data)\n                # calculate the batch loss\n                loss = criterion(output, target)\n                # Calculate Accuracy\n                accuracy = (output.argmax(dim=1) == target).float().mean()\n                # update average validation loss and accuracy\n                valid_loss += loss\n                valid_accuracy += accuracy\n\n        return valid_loss / len(valid_loader), valid_accuracy / len(valid_loader)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.318853Z","iopub.status.idle":"2024-12-23T12:21:41.31944Z","shell.execute_reply.started":"2024-12-23T12:21:41.319117Z","shell.execute_reply":"2024-12-23T12:21:41.319142Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def fit_tpu(\n    model, epochs, device, criterion, optimizer, train_loader, valid_loader=None\n):\n\n    valid_loss_min = np.Inf  # track change in validation loss\n\n    # keeping track of losses as it happen\n    train_losses = []\n    valid_losses = []\n    train_accs = []\n    valid_accs = []\n\n    for epoch in range(1, epochs + 1):\n        gc.collect()\n        para_train_loader = pl.ParallelLoader(train_loader, [device])\n\n        xm.master_print(f\"{'='*50}\")\n        xm.master_print(f\"EPOCH {epoch} - TRAINING...\")\n        train_loss, train_acc = model.train_one_epoch(\n            para_train_loader.per_device_loader(device), criterion, optimizer, device\n        )\n        xm.master_print(\n            f\"\\n\\t[TRAIN] EPOCH {epoch} - LOSS: {train_loss}, ACCURACY: {train_acc}\\n\"\n        )\n        train_losses.append(train_loss)\n        train_accs.append(train_acc)\n        gc.collect()\n\n        if valid_loader is not None:\n            gc.collect()\n            para_valid_loader = pl.ParallelLoader(valid_loader, [device])\n            xm.master_print(f\"EPOCH {epoch} - VALIDATING...\")\n            valid_loss, valid_acc = model.validate_one_epoch(\n                para_valid_loader.per_device_loader(device), criterion, device\n            )\n            xm.master_print(f\"\\t[VALID] LOSS: {valid_loss}, ACCURACY: {valid_acc}\\n\")\n            valid_losses.append(valid_loss)\n            valid_accs.append(valid_acc)\n            gc.collect()\n\n            # save model if validation loss has decreased\n            if valid_loss <= valid_loss_min and epoch != 1:\n                xm.master_print(\n                    \"Validation loss decreased ({:.4f} --> {:.4f}).  Saving model ...\".format(\n                        valid_loss_min, valid_loss\n                    )\n                )\n            #                 xm.save(model.state_dict(), 'best_model.pth')\n\n            valid_loss_min = valid_loss\n\n    return {\n        \"train_loss\": train_losses,\n        \"valid_losses\": valid_losses,\n        \"train_acc\": train_accs,\n        \"valid_acc\": valid_accs,\n    }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.321748Z","iopub.status.idle":"2024-12-23T12:21:41.322298Z","shell.execute_reply.started":"2024-12-23T12:21:41.322008Z","shell.execute_reply":"2024-12-23T12:21:41.322033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = ViTBase16(n_classes=5, pretrained=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.323579Z","iopub.status.idle":"2024-12-23T12:21:41.324102Z","shell.execute_reply.started":"2024-12-23T12:21:41.323829Z","shell.execute_reply":"2024-12-23T12:21:41.323854Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def _run():\n    train_dataset = CassavaDataset(train_df, transforms=transforms_train)\n    valid_dataset = CassavaDataset(valid_df, transforms=transforms_valid)\n\n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n        train_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True,\n    )\n\n    valid_sampler = torch.utils.data.distributed.DistributedSampler(\n        valid_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False,\n    )\n\n    train_loader = torch.utils.data.DataLoader(\n        dataset=train_dataset,\n        batch_size=BATCH_SIZE,\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=8,\n    )\n\n    valid_loader = torch.utils.data.DataLoader(\n        dataset=valid_dataset,\n        batch_size=BATCH_SIZE,\n        sampler=valid_sampler,\n        drop_last=True,\n        num_workers=8,\n    )\n\n    criterion = nn.CrossEntropyLoss()\n    #     device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    device = xm.xla_device()\n    model.to(device)\n\n    lr = LR * xm.xrt_world_size()\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n\n    xm.master_print(f\"INITIALIZING TRAINING ON {xm.xrt_world_size()} TPU CORES\")\n    start_time = datetime.now()\n    xm.master_print(f\"Start Time: {start_time}\")\n\n    logs = fit_tpu(\n        model=model,\n        epochs=N_EPOCHS,\n        device=device,\n        criterion=criterion,\n        optimizer=optimizer,\n        train_loader=train_loader,\n        valid_loader=valid_loader,\n    )\n\n    xm.master_print(f\"Execution time: {datetime.now() - start_time}\")\n\n    xm.master_print(\"Saving Model\")\n    xm.save(\n        model.state_dict(), f'model_5e_{datetime.now().strftime(\"%Y%m%d-%H%M\")}.pth'\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.326472Z","iopub.status.idle":"2024-12-23T12:21:41.327032Z","shell.execute_reply.started":"2024-12-23T12:21:41.326743Z","shell.execute_reply":"2024-12-23T12:21:41.326767Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Start training processes\ndef _mp_fn(rank, flags):\n    torch.set_default_tensor_type(\"torch.FloatTensor\")\n    a = _run()\n\n\n# _run()\nFLAGS = {}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method=\"fork\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-23T12:21:41.329501Z","iopub.status.idle":"2024-12-23T12:21:41.330023Z","shell.execute_reply.started":"2024-12-23T12:21:41.32979Z","shell.execute_reply":"2024-12-23T12:21:41.329815Z"}},"outputs":[],"execution_count":null}]}