{"cells":[{"metadata":{},"cell_type":"markdown","source":"Credit to https://www.kaggle.com/debarshichanda/tpu-training.\nI modified some codes and made it work."},{"metadata":{"trusted":true},"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 nightly --apt-packages libomp5 libopenblas-dev\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"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\n!pip install pretrainedmodels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport copy\nimport time\nimport random\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\n#import torch_optimizer as optim\nfrom torch.optim import lr_scheduler\nimport torchvision\nfrom torchvision import models\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.cuda import amp\n\nfrom sklearn.model_selection import train_test_split, StratifiedKFold\nfrom sklearn.metrics import accuracy_score\nfrom sklearn.utils import class_weight\n\nfrom tqdm.notebook import tqdm\nfrom collections import defaultdict\nfrom datetime import datetime\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nimport timm\nimport pretrainedmodels\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# For parallelization in TPUs\nos.environ[\"XLA_USE_BF16\"] = \"1\"\nos.environ[\"XLA_TENSOR_ALLOCATOR_MAXSIZE\"] = \"100000000\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CFG:\n    model_name = 'resnet18d'#'tf_efficientnet_b4_ns'\n    img_size = 380\n    scheduler = 'CosineAnnealingWarmRestarts'\n    T_max = 10\n    T_0 = 10\n    lr = 1e-4\n    min_lr = 1e-6\n    batch_size = 64#16*4\n    weight_decay = 1e-6\n    seed = 42\n    num_classes = 5\n    num_epochs = 10#10\n    n_fold = 5\n    smoothing = 0.2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def set_seed(seed = 42):\n    '''Sets the seed of the entire notebook so results are the same every time we run.\n    This is for REPRODUCIBILITY.'''\n    np.random.seed(seed)\n    random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    # When running on the CuDNN backend, two further options must be set\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    # Set a fixed value for the hash seed\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    \nset_seed(CFG.seed)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ROOT_DIR = \"../input/cassava-leaf-disease-classification\"\nTRAIN_DIR = \"../input/cassava-leaf-disease-classification/train_images\"\nTEST_DIR = \"../input/cassava-leaf-disease-classification/test_images\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv(f\"{ROOT_DIR}/train.csv\")\ndf","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"skf = StratifiedKFold(n_splits=CFG.n_fold)\nfor fold, ( _, val_) in enumerate(skf.split(X=df, y=df.label)):\n    df.loc[val_ , \"kfold\"] = int(fold)\n    \ndf['kfold'] = df['kfold'].astype(int)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaLeafDataset(nn.Module):\n    def __init__(self, root_dir, df, transforms=None):\n        self.root_dir = root_dir\n        self.df = df\n        self.labels=df['label'].values\n        self.image_ids=df['image_id'].values\n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, index):\n        img_path = os.path.join(self.root_dir, self.image_ids[index])\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        label = self.labels[index]\n        \n        if self.transforms:\n            img = self.transforms(image=img)[\"image\"]\n            \n        return img, label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_transforms = {\n    \"train\": A.Compose([\n        A.RandomResizedCrop(CFG.img_size, CFG.img_size),\n#         A.Transpose(p=0.5),\n        A.HorizontalFlip(p=0.5),\n#         A.VerticalFlip(p=0.5),\n#         A.ShiftScaleRotate(p=0.5),\n#         A.HueSaturationValue(\n#                 hue_shift_limit=0.2, \n#                 sat_shift_limit=0.2, \n#                 val_shift_limit=0.2, \n#                 p=0.5\n#             ),\n#         A.RandomBrightnessContrast(\n#                 brightness_limit=(-0.1,0.1), \n#                 contrast_limit=(-0.1, 0.1), \n#                 p=0.5\n#             ),\n        A.Normalize(\n                mean=[0.485, 0.456, 0.406], \n                std=[0.229, 0.224, 0.225], \n                max_pixel_value=255.0, \n                p=1.0\n            ),\n#         A.CoarseDropout(p=0.5),\n#         A.Cutout(p=0.5),\n        ToTensorV2()], p=1.),\n    \n    \"valid\": A.Compose([\n        A.CenterCrop(CFG.img_size, CFG.img_size, p=1.),\n        A.Resize(CFG.img_size, CFG.img_size),\n        A.Normalize(\n                mean=[0.485, 0.456, 0.406], \n                std=[0.229, 0.224, 0.225], \n                max_pixel_value=255.0, \n                p=1.0\n            ),\n        ToTensorV2()], p=1.)\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import tqdm\nfrom tqdm.notebook import tqdm as tqdm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# class EffNet(nn.Module):\n#     def __init__(self, n_classes, pretrained=True):\n#         super(EffNet, self).__init__()\n#         self.model = timm.create_model(CFG.model_name, pretrained=True)\n#         num_features = self.model.classifier.in_features\n#         self.model.classifier = nn.Linear(num_features, CFG.num_classes)\n\n#     def forward(self, x):\n#         x = self.model(x)\n#         return x\n\nclass EffNet(nn.Module):\n    def __init__(self, n_classes, pretrained=True):\n        super(EffNet, self).__init__()\n        self.model = timm.create_model(CFG.model_name, pretrained=True)\n        \n        self.logit = nn.Linear(512, CFG.num_classes)\n\n    def forward(self, x):\n        batch_size, C, H, W = x.shape\n        logit = self.model.forward_features(x)\n        logit = F.adaptive_avg_pool2d(logit,1).reshape(batch_size,-1)\n        \n        logit=self.logit(logit)\n        return logit  ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_one_epoch(model,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    model.train()\n    for i, (data, target) in enumerate(tqdm(train_loader)):\n        \n#         print(i)\n        data = data.to(device)\n        target = target.to(device)\n\n        # clear the gradients of all optimized variables\n        optimizer.zero_grad()\n\n        # forward pass: compute predicted outputs by passing inputs to the model\n        output = model(data)\n\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\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\n        xm.optimizer_step(optimizer)\n\n\n\n    return epoch_loss / len(train_loader), epoch_accuracy / len(train_loader)\n\ndef validate_one_epoch(model,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    model.eval()\n    for data, target in valid_loader:\n        # move tensors to GPU if CUDA is available\n\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 = 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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def fit_tpu(\n    model, epochs, device, criterion, optimizer, train_loader, valid_loader=None\n):\n\n    best_acc = 0.0  # track change in validation loss\n\n    # keeping track of losses as it happen\n    history = defaultdict(list)\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 = train_one_epoch(model,\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        history['train loss'].append(train_loss)\n        history['train acc'].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 = validate_one_epoch(model,\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            history['valid loss'].append(valid_loss)\n            history['valid acc'].append(valid_acc)\n            gc.collect()\n\n            # save model if validation loss has decreased\n            if valid_acc <= best_acc:\n#                 xm.master_print(\n#                     \"Validation loss decreased ({:.4f} --> {:.4f}).  Saving model ...\".format(\n#                         valid_loss_min, valid_loss\n#                     )\n#                 )\n                PATH = f\"Fold-{fold}_{best_acc}_epoch-{epoch}.pth\"\n                xm.save(model.state_dict(), PATH)\n\n            best_acc = valid_acc\n\n    return history","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class LabelSmoothingLoss(nn.Module):\n    def __init__(self, classes, smoothing=0.0, dim=-1): \n        super(LabelSmoothingLoss, self).__init__() \n        self.confidence = 1.0 - smoothing \n        self.smoothing = smoothing \n        self.cls = classes \n        self.dim = dim\n        \n    def forward(self, pred, target): \n        pred = pred.log_softmax(dim=self.dim) \n        with torch.no_grad(): \n            true_dist = torch.zeros_like(pred) \n            true_dist.fill_(self.smoothing / (self.cls - 1)) \n            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence) \n        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# model = EffNet(n_classes=CFG.num_classes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _run(fold):\n    valid_df = df[df.kfold == fold]\n    train_df = df[df.kfold != fold]\n    \n    train_dataset = CassavaLeafDataset(TRAIN_DIR, train_df, transforms=data_transforms[\"train\"])\n    valid_dataset = CassavaLeafDataset(TRAIN_DIR, valid_df, transforms=data_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=CFG.batch_size,\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=4,\n    )\n\n    valid_loader = torch.utils.data.DataLoader(\n        dataset=valid_dataset,\n        batch_size=CFG.batch_size*4,\n        sampler=valid_sampler,\n        drop_last=False,\n        num_workers=4,\n    )\n\n#     criterion = LabelSmoothingLoss(smoothing=CFG.smoothing, classes=CFG.num_classes)\n    criterion=torch.nn.CrossEntropyLoss()\n    \n    device = xm.xla_device()\n    model = EffNet(n_classes=CFG.num_classes)\n    model.to(device)\n\n    lr = CFG.lr * xm.xrt_world_size()\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr,weight_decay=CFG.weight_decay)\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=CFG.num_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    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Start training processes\ndef _mp_fn(rank, flags):\n    torch.set_default_tensor_type(\"torch.FloatTensor\")\n    a = _run(fold=0)\n\n\n# _run()\nFLAGS = {}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=1, start_method=\"fork\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!rm *.py\n!rm *.whl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}