{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport cv2\nimport time\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport gc; gc.collect()\nimport matplotlib.pyplot as plt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install efficientnet_pytorch > /dev/null\n!pip install albumentations > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.utils.data as D\nimport torch.nn.functional as F\n\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.core.xla_model as xm\nimport torch_xla.utils.utils as xu\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\nfrom efficientnet_pytorch import EfficientNet\n\nimport torchvision\nfrom torchvision import transforms as T\nimport torchvision.transforms as transforms\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom sklearn.model_selection import train_test_split","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df_path=\"../input/hpa-single-cell-image-classification/train.csv\"\ntrain_images_path=\"../input/hpa-single-cell-image-classification/train\"\ntest_images_path=\"../input/hpa-single-cell-image-classification/test\"\nsample_df_path=\"../input/hpa-single-cell-image-classification/sample_submission.csv\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"os.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df=pd.read_csv(train_df_path)\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"Transform = transforms.Compose(\n    [transforms.ToTensor(),\n    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class HPADataset(Dataset):\n    def __init__(self, path, df, img_size, Transform):\n        self.path = path\n        self.df = df\n        self.img_ids = df['ID'].values\n        self.labels = df['Label'].values\n        self.img_size = img_size        \n        self.transform = Transform\n        \n    def _get_image(self, ID):\n        R = cv2.imread(self.path + '/' + ID + '_red.png', cv2.IMREAD_UNCHANGED)\n        Y = cv2.imread(self.path + '/' + ID + '_yellow.png', cv2.IMREAD_UNCHANGED)\n        G = cv2.imread(self.path + '/' + ID + '_green.png', cv2.IMREAD_UNCHANGED)\n        B = cv2.imread(self.path + '/' + ID + '_blue.png', cv2.IMREAD_UNCHANGED)\n        img = np.stack((\n                R/2 + Y/2, \n                G/2 + Y/2, \n                B),-1)\n        \n        img = cv2.resize(img, (self.img_size, self.img_size))\n        img = np.divide(img, 255)\n        return img          \n        \n    def __len__(self):\n        return len(self.df) \n    \n    def __getitem__(self, index):\n        x = self._get_image(self.img_ids[index])\n        x = self.transform(x)\n        y = self.labels[index]\n        y = y.split('|')\n        y = list(map(int, y))            \n        y = np.eye(FLAGS['NUM_CLASSES'], dtype='float')[y]                                    \n        y = y.sum(axis=0)\n        return x, y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_split, eval_split = train_test_split(train_df, test_size=0.2, random_state=42)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# def get_model():  \n#     model = torchvision.models.resnet50()\n#     model.fc = nn.Linear(2048, 19, bias=True)\n#     return model\n\n# resnet50 = get_model()\nbaseModel = EfficientNet.from_pretrained('efficientnet-b5', num_classes=19)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def graph_losses(losses):\n    for phase, color in zip(['train', 'eval'], ['r--', 'b--']):\n        if not losses[phase]:\n            continue\n        epoch_count = range(1, len(losses[phase]) + 1)\n        plt.plot(epoch_count, losses[phase], color)\n        plt.legend([f'{phase.capitalize()} Loss'])\n        plt.xlabel('Epoch')\n        plt.ylabel('Loss')\n        plt.show() \n\ndef reduce_fn(vals):\n    # take average\n    return sum(vals) / len(vals)        \n        \ndef run(epochs=20, validate_every=2):\n    \n    device = xm.xla_device()\n    \n    # Init DataLoader\n    loaders = {}\n    Transform = transforms.Compose(\n        [transforms.ToTensor(),\n        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])\n\n    train_dataset = HPADataset(train_images_path, train_split, FLAGS['IMG_SIZE'], Transform)\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    train_loader = DataLoader(train_dataset, batch_size=FLAGS['BATCH_SIZE'], sampler=train_sampler, shuffle=False)\n\n    eval_dataset = HPADataset(train_images_path, eval_split, FLAGS['IMG_SIZE'], Transform)\n    eval_sampler = torch.utils.data.distributed.DistributedSampler(\n          eval_dataset,\n          num_replicas=xm.xrt_world_size(),\n          rank=xm.get_ordinal(),\n          shuffle=False)    \n    eval_loader = DataLoader(eval_dataset, batch_size=FLAGS['BATCH_SIZE'], sampler=eval_sampler, shuffle=False)\n\n    loaders['train'] = train_loader\n    loaders['eval'] = eval_loader    \n    \n    # Initialize model\n    model = baseModel.to(device)\n    learning_rate = FLAGS['LR'] * xm.xrt_world_size()\n    optimizer = torch.optim.AdamW(model.parameters(),\n                      lr=learning_rate, weight_decay=5e-4)\n    criterion = nn.BCEWithLogitsLoss()\n    \n    loaders = loaders\n    criterion = criterion\n    best_loss = 0.0\n    running_losses = {'train': [], 'eval': []}\n    \n    for epoch in range(1, epochs + 1):\n        phases = ['train']\n        if epoch % validate_every == 0:\n            phases.append('eval')\n\n        for phase in phases:\n            model.eval() if phase == 'eval' else model.train()\n            gc.collect() # prevent OOM problems\n            para_loader = pl.ParallelLoader(train_loader, [device]) \n            gc.collect()\n\n            xm.master_print(\"Epoch {}/{}\".format(epoch, epochs))\n            loader = para_loader.per_device_loader(device)\n            # loader = pl.MpDeviceLoader(self.loaders[phase], FLAGS['DEVICE'])\n            for idx, (imgs, labels) in enumerate(tqdm(loader)):\n                xm.master_print(f'Phase: {phase}, current step: {idx}')\n                imgs, labels = imgs.float().to(device), labels.float().to(device)\n                optimizer.zero_grad()\n                outputs = model(imgs)\n                loss = criterion(outputs, labels)\n                loss_reduced = xm.mesh_reduce('loss_reduce', loss, reduce_fn) \n                running_losses[phase].append(loss_reduced.item())\n                if phase == 'train':\n                    loss.backward()\n                    xm.optimizer_step(optimizer)\n\n            mean_loss = np.array(running_losses[phase]).mean()\n            if phase == 'eval':\n                xm.master_print(\"Eval running loss: \", running_losses['eval'])\n                if mean_loss < best_loss:\n                    best_loss = mean_loss\n                    xm.save(model.state_dict(), 'model_best.pth')\n            xm.master_print(epoch, mean_loss, best_loss, (time.time()-start_time)/60**1)\n    graph_losses(running_losses)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import time\n\n\nFLAGS = {}\nEPOCHS = 1 # For testing purposes\nFLAGS['LR'] = 1e-4\nFLAGS['IMG_SIZE'] = 256\nFLAGS['BATCH_SIZE'] = 64\nFLAGS['NUM_CLASSES'] = 19\n\nstart_time = time.time()\n\n# Start training processes\ndef _mp_fn(rank, flags):\n    a = run(EPOCHS)\n_mp_fn()\n# xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=1, start_method='fork')","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}