{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","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  > /dev/null\n!python pytorch-xla-env-setup.py --version nightly  > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"!pip install timm  > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import gc\nimport os\nimport time\nimport torch\nimport albumentations\n\nimport numpy as np\nimport pandas as pd\n\nimport cv2\nfrom PIL import Image\n\nimport torch.nn as nn\nfrom sklearn import metrics\nfrom sklearn import model_selection\nfrom torch.nn import functional as F\nfrom torch.optim import Adam\n\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\n\nimport timm\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nos.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"FLAGS = {'fold':0,\n         'model':'resnext50_32x4d',\n         'pretrained': True,\n         'batch_size':128,\n         'num_workers':4,\n         'lr':3e-4,\n         'epochs':10\n        }","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TimmModels(nn.Module):\n    def __init__(self, model_name, pretrained=True, num_classes=5):\n        super(TimmModels, self).__init__()\n        self.m = timm.create_model(model_name, pretrained=pretrained)\n        model_list = list(self.m.children())\n        model_list[-1] = nn.Linear(in_features=model_list[-1].in_features,\n                                   out_features=num_classes,\n                                   bias=True\n                                  )\n        self.m = nn.Sequential(*model_list)\n        \n    def forward(self, image):\n        out = self.m(image)\n        return out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageDataset:\n    def __init__(self,\n                 image_paths,\n                 targets,\n                 resize,\n                 augmentations=None,\n                 backend='pil',\n                 channel_first=True\n                ):\n        \"\"\"\n        :param image_paths: list of paths to images\n        :param targets: numpy array\n        :param resize: tuple or None\n        :param augmentations: albumentations augmentations\n        \"\"\"\n        \n        self.image_paths = image_paths\n        self.targets = targets\n        self.resize = resize\n        self.augmentations = augmentations\n        self.backend = backend\n        self.channel_first = channel_first\n        \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def __getitem__(self, item):\n        targets = self.targets[item]\n        \n        if self.backend == 'pil':\n            image = Image.open(self.image_paths[item])\n            \n            if self.resize is not None:\n                image = image.resize((self.resize[1], self.resize[0]),\n                                     resample=Image.BILINEAR\n                                    )\n                \n            image = np.array(image)\n            \n            if self.augmentations is not None:\n                augmented = self.augmentations(image=image)\n                image = augmented['image']\n                \n        elif self.backend == 'cv2':\n            image = cv2.imread(self.image_paths[item])\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n            \n            if self.resize is not None:\n                image = cv2.resize(image, \n                                   (self.resize[1], self.resize[0]),\n                                   interpolation=cv2.INTER_CUBIC\n                                  )\n                \n            if self.augmentations is not None:\n                augmented = self.augmentations(image=image)\n                \n            image = augmented['image']\n            \n        else:\n            raise Exception(\"Backend not implemented\")\n            \n        if self.channel_first:\n            image = np.transpose(image, (2,0,1)).astype(np.float32)\n            \n        return {\"image\":torch.tensor(image),\n                \"targets\":torch.tensor(targets)\n               }","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv('../input/cassava-leaf-disease-classification/train.csv')\ndf['kfold'] = -1\n\ndf = df.sample(frac=1).reset_index(drop=True)\n# shuffling\n\ny = df.label.values\nskf = model_selection.StratifiedKFold(n_splits=5)\n\nfor f, (t_, v_) in enumerate(skf.split(X=df, y=y)):\n    df.loc[v_, 'kfold'] = f\n    \ndf.to_csv('train_folds.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"MX = xmp.MpModelWrapper(TimmModels(\n    FLAGS['model'],\n    pretrained=FLAGS['pretrained'],\n    num_classes=5))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_loop_fn(data_loader, \n                  loss_fn, \n                  model,\n                  optimizer,\n                  device,\n                  scheduler=None):\n    model.train()\n    for bi, d in enumerate(data_loader):\n        \n        images = d['image'].to(device, dtype=torch.float32)\n        targets = d['targets'].to(device, dtype=torch.int64)\n        # why int64?\n        \n        optimizer.zero_grad()\n        outputs = model(images)\n        \n        loss = loss_fn(outputs, targets)\n        \n        loss.backward()\n        \n        # Use PyTorch XLA optimizer stepping\n        xm.optimizer_step(optimizer, barrier=True)\n        \n        if scheduler is not None: scheduler.step()\n            \n    loss_reduced = xm.mesh_reduce(\n        'loss_reduce',\n        loss,\n        lambda x: sum(x) / len(x)\n                                 )\n    \n    xm.master_print(f'bi={bi}, train loss={loss_reduced}')\n    \n    model.eval()\n    \ndef eval_loop_fn(data_loader, \n                 loss_fn, \n                 model, \n                 device):\n    \n    fin_targets = []\n    fin_outputs = []\n    \n    for bi, d in enumerate(data_loader):\n        \n        images = d['image'].to(device)\n        targets = d['targets'].to(device)\n        \n        with torch.no_grad(): outputs = model(images)\n            \n        targets_np = targets.cpu().detach().numpy().tolist()\n        outputs_np = outputs.cpu().detach().numpy().tolist()\n        \n        fin_targets.extend(targets_np)\n        fin_outputs.extend(outputs_np)\n        \n        del targets_np, outputs_np\n        \n        gc.collect()\n        \n    o, t = np.array(fin_outputs), np.array(fin_targets)\n    \n    loss = loss_fn(torch.tensor(o), torch.tensor(t))\n    \n    loss_reduced = xm.mesh_reduce('loss_reduce',\n                                  loss,\n                                  lambda x: sum(x) / len(x)\n                                 )\n    \n    xm.master_print(f'val. loss = {loss_reduced}')\n    \n    acc = metrics.accuracy_score(t,o.argmax(axis=1))\n    acc_reduced = xm.mesh_reduce('acc_reduce',\n                                  acc,\n                                  lambda x: sum(x) / len(x)\n                                 )\n    \n    xm.master_print(f'val. accuracy = {acc_reduced}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def run(rank, flags):\n    global FLAGS\n    \n    torch.set_default_tensor_type('torch.FloatTensor')\n    \n    xm.master_print(\"let's start!\")\n    training_data_path = '../input/cassava-jpeg-256x256/kaggle/train_images_jpeg'\n    df = pd.read_csv(\"/kaggle/working/train_folds.csv\")\n    \n    device = xm.xla_device()\n    \n    epochs = FLAGS['epochs']\n    fold = FLAGS['fold']\n    \n    df_train = df[df.kfold != fold].reset_index(drop=True)\n    df_valid = df[df.kfold == fold].reset_index(drop=True)\n    \n    mean = (0.485, 0.456, 0.406)\n    std = (0.229, 0.224, 0.225)\n    \n    train_aug = albumentations.Compose(\n        [\n            albumentations.Normalize(\n                mean,\n                std,\n                max_pixel_value=255.0,\n                always_apply=True\n            ),\n            albumentations.Transpose(p=0.5),\n            albumentations.HorizontalFlip(p=0.5),\n            albumentations.VerticalFlip(p=0.5),\n            albumentations.ShiftScaleRotate(p=0.5),\n            albumentations.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            albumentations.RandomBrightnessContrast(\n                brightness_limit=(-0.1,0.1),\n                contrast_limit=(-0.1,0.1),\n                p=0.5\n            ),\n            albumentations.CoarseDropout(p=0.5),\n            albumentations.Cutout(p=0.5)\n        ]\n    )\n    \n    valid_aug = albumentations.Compose(\n        [\n            albumentations.Normalize(\n                mean,\n                std,\n                max_pixel_value=255.0,\n                always_apply=True\n            )\n        ]\n    )\n    \n    train_images = df_train.image_id.values.tolist()\n    train_images = [\n        os.path.join(training_data_path, i) for i in train_images\n    ]\n    \n    train_targets = df_train.label.values\n    \n    valid_images = df_valid.image_id.values.tolist()\n    valid_images = [\n        os.path.join(training_data_path, i) for i in valid_images\n    ]\n    valid_targets = df_valid.label.values\n    \n    train_dataset = ImageDataset(\n        image_paths=train_images,\n        targets=train_targets,\n        resize=None,\n        augmentations=train_aug\n    )\n    \n    valid_dataset = ImageDataset(\n        image_paths=valid_images,\n        targets=valid_targets,\n        resize=None,\n        augmentations=valid_aug\n    )\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    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=FLAGS['batch_size'],\n        sampler=train_sampler,\n        num_workers=FLAGS['num_workers'],\n        drop_last=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    valid_loader = torch.utils.data.DataLoader(\n        valid_dataset,\n        batch_size=FLAGS['batch_size'],\n        sampler=valid_sampler,\n        num_workers=FLAGS['num_workers'],\n        drop_last=False\n    )\n    \n    mp_device_loader = pl.MpDeviceLoader(\n        train_loader,\n        device,\n        fixed_batch_size=True\n    )\n    \n    model = MX.to(device)\n    loss_fn = nn.CrossEntropyLoss()\n    optimizer = Adam(model.parameters(),\n                     lr=FLAGS['lr']*xm.xrt_world_size()\n                    )\n    \n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,\n                                                           len(train_loader)*FLAGS['epochs']\n                                                          )\n    \n    xm.master_print(f'========== training fold {FLAGS[\"fold\"]} for {FLAGS[\"epochs\"]} epochs ===========')\n    \n    for i in range(FLAGS['epochs']):\n        xm.master_print(f'EPOCH {i}:')\n        \n        train_loop_fn(train_loader, loss_fn, model, optimizer, device, scheduler)\n        \n        eval_loop_fn(valid_loader, loss_fn, model, device)\n        \n    xm.master_print('save model')\n    \n    xm.save(model.state_dict(), f'xla_trained_model_{FLAGS[\"epochs\"]}_epochs_fold_{FLAGS[\"fold\"]}.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"start_time = time.time()\nxmp.spawn(run, args=(FLAGS,), nprocs=8, start_method='fork')\nprint('time taken: ', time.time()-start_time)","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}