{"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":"!pip install tqdm\nfrom tqdm import tqdm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install efficientnet_pytorch\nfrom efficientnet_pytorch import EfficientNet","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 AdamW\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':27e-5,\n         'epochs':30,\n         'image_size': 256\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":"# from https://www.kaggle.com/nroman/melanoma-pytorch-starter-efficientnet\n\nclass Net(nn.Module):\n    def __init__(self, arch):\n        super(Net, self).__init__()\n        # 자체 기능이 있어서 그걸 쓰는 쪽으로~\n        #in_features = arch._fc.in_features\n        #arch._fc = nn.Linear(in_features=in_features, out_features=5, bias=True)\n        self.arch = arch\n        \n    def forward(self, x):\n        \"\"\"\n        No sigmoid in forward because we are going to use BCEWithLogitsLoss\n        Which applies sigmoid for us when calculating a loss\n        \n        hmmm\n        \"\"\"\n        x = self.arch(x)\n        return x","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\n\n\nskf = model_selection.StratifiedKFold(n_splits=10)\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)\n\n\n\n#sss = model_selection.StratifiedShuffleSplit(n_splits=10, test_size=0.1, random_state=2021)\n#sss.split(X=df, y=y)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#model = EfficientNet.from_pretrained('efficientnet-b4', num_classes=5)\nmodel = EfficientNet.from_pretrained('efficientnet-b5', num_classes=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"MX = xmp.MpModelWrapper(Net(arch=model))","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    \n    \n    #pbar = tqdm(enumerate(data_loader), total=len(data_loader))\n    for bi, d in enumerate(data_loader):\n    #for bi, d in pbar:\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('loss_reduce', loss, lambda x: sum(x) / len(x))\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    #pbar = tqdm(enumerate(data_loader), total=len(data_loader))\n    #for bi, d in pbar:\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}')\n    \n    return 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    best_val_acc = 0\n    \n    torch.set_default_tensor_type('torch.FloatTensor')\n    \n    xm.master_print(\"let's start!\")\n    \n\n    \n    training_data_path = '../input/cassava-jpeg-256x256/kaggle/train_images_jpeg'\n    df = pd.read_csv(\"/kaggle/working/train_folds.csv\")\n\n    xm.master_print(\"fold data ready\")\n    xm.master_print(df.shape)\n\n    device = xm.xla_device()\n\n    xm.master_print(\"device ready\")\n\n    epochs = FLAGS['epochs']\n    fold = FLAGS['fold']\n\n\n    model = MX.to(device)\n\n    xm.master_print(\"model ready\")\n\n    while fold < 5:\n\n        if fold != 0:\n            break\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        df_fold = df[df.kfold == fold].reset_index(drop=True)\n        y = df_fold.label.values\n\n        xm.master_print(df_fold.shape)\n\n        sss = model_selection.StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=2021)\n\n        for train_index, valid_index in sss.split(X=df_fold, y=y): \n            #xm.master_print(f'train index {train_index}')\n            #xm.master_print(f'valid index {valid_index}')\n            df_train = df_fold.iloc[train_index]\n            df_valid = df_fold.iloc[valid_index]\n\n        #xm.master_print(f'df_train {df_train.shape}')\n        #xm.master_print(f'df_valid {df_valid.shape}')\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\n                albumentations.Normalize(\n                    mean,\n                    std,\n                    max_pixel_value=255.0,\n                    always_apply=True\n                ),\n\n                # basic augmentation\n                #albumentations.HorizontalFlip(p=0.5),\n                albumentations.VerticalFlip(p=0.5),\n                #albumentations.RandomRotate90(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\n                # additional augmentations\n\n\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\n        #xm.master_print(f'train_loader {len(train_loader.dataset)}')\n        #xm.master_print(f'valid_loader {len(valid_loader.dataset)}')\n\n\n        mp_device_loader = pl.MpDeviceLoader(\n            train_loader,\n            device,\n            fixed_batch_size=True\n        )\n\n\n        loss_fn = nn.CrossEntropyLoss()\n        optimizer = AdamW(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        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n            optimizer,\n            T_0=10,\n            T_mult=1,\n            eta_min=1e-8\n        )\n\n        xm.master_print(f'========== training fold {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            #train_loop_fn(train_loader, loss_fn, model, optimizer, device)\n            val_acc = eval_loop_fn(valid_loader, loss_fn, model, device)\n\n            if val_acc > best_val_acc:\n                best_val_acc = val_acc\n                xm.master_print('save model')\n                xm.save(model.state_dict(), f'xla_trained_best_model.pth')\n        xm.master_print('save model')\n        xm.save(model.state_dict(), f'xla_trained_model_{FLAGS[\"epochs\"]}_epochs_fold_{fold}.pth')\n\n        fold += 1\n","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":{"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":{"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":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}