{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### Online usage","metadata":{}},{"cell_type":"code","source":"# Online\n# !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## !pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.7-cp37-cp37m-linux_x86_64.whl","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Offline usage","metadata":{}},{"cell_type":"code","source":"# Offline\n!pip install -U ../input/torchxla17wheel/torch_xla-1.7-cp37-cp37m-linux_x86_64.whl","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset","metadata":{}},{"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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"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.utils.serialization as xser\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"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\"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df=pd.read_csv(train_df_path)\ntrain_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"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))])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HPATestDataset(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        \n        data_file = cv2.imread(self.path + '/' + ID + '_green.png')\n            \n        img = cv2.resize(data_file, (self.img_size, self.img_size))\n        X = img/255.             \n        \n        return X          \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        if \"train\" in self.path:\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\n        y = self.img_ids[index]\n        return x, y","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_split, eval_split = train_test_split(train_df, test_size=0.2, random_state=42)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"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\nresnet50 = get_model()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"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 = resnet50.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)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"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\n# xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=1, start_method='fork')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Submission\nTest with one epoch's worth of training","metadata":{}},{"cell_type":"code","source":"state_dict = xser.load('../input/model-epoch-1pth/model_epoch_1.pth')\nresnet50.load_state_dict(state_dict)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x = [name.rstrip('green.png').rstrip('_') for name in (os.listdir(test_images_path)) if '_green.png' in name]\ny = np.zeros(len(x))\nz = zip(x, y)\ntest_df = pd.DataFrame(list(z), columns = ['ID', 'Label'])\ntest_dataset = HPATestDataset(test_images_path, test_df, FLAGS['IMG_SIZE'], Transform)\ntest_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test(model):\n    s_ls = []\n    with torch.no_grad():\n        model.eval()\n        DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n        for image, fname in test_loader:\n            image = image.float().to(DEVICE)\n            model.to(DEVICE)\n            logits = model(image)                     \n            prob = F.softmax(logits, dim=1)\n            p, top_class = prob.topk(1, dim=1)\n            sp = ' '.join(str(e) for e in [top_class[0][0].item(), p[0][0].item()])      \n            img = cv2.imread(test_images_path + \"/\" + fname[0] + '_green.png')\n\n            if img.shape[0] == 2048:\n                sp = sp + ' eNoLCAgIMAEABJkBdQ=='\n            elif img.shape[0] == 1728:\n                sp = sp + ' eNoLCAjJNgIABNkBkg=='\n            else:\n                sp = sp + ' eNoLCAgIsAQABJ4Beg=='\n\n            s_ls.append([fname[0], img.shape[1], img.shape[0], sp])\n    return s_ls\n            \nresults = test(resnet50)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub = pd.DataFrame.from_records(results, columns=['ID', 'ImageWidth', 'ImageHeight', 'PredictionString'])\nsub","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub.to_csv('submission.csv', index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}