{"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":"This notebook uses pytorch xla package to implement the training and validation of effb7 model with tpu. Some settings refer to [Alien's notebook](https://www.kaggle.com/h053473666/siim-covid19-efnb7-train-study/data). You can also see some differences between pytorch tpu code and [tensorflow tpu code](https://www.kaggle.com/h053473666/siim-covid19-efnb7-train-study/data).\n\nFor more about PyTorch/XLA please see its [Github](https://github.com/pytorch/xla) or its [documentation](http://pytorch.org/xla/).","metadata":{}},{"cell_type":"markdown","source":"## Kaggle or Colab\nYou can comment out one of below two cells for running code on kaggle or [colab](https://drive.google.com/drive/my-drive).","metadata":{"id":"9yl8tOu7G5nw"}},{"cell_type":"code","source":"# Run this code with kaggle TPU\nENVIRONMENT = \"kaggle\"\ndir = \"/kaggle/working\"","metadata":{"id":"mNgdeub3G-CK","execution":{"iopub.status.busy":"2021-07-03T10:40:44.582751Z","iopub.execute_input":"2021-07-03T10:40:44.583269Z","iopub.status.idle":"2021-07-03T10:40:44.593846Z","shell.execute_reply.started":"2021-07-03T10:40:44.583159Z","shell.execute_reply":"2021-07-03T10:40:44.592662Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# # Run this code with colab TPU\n# # You can use larger batch_size, image_size and model with colab high memory mode\n# ENVIRONMENT = \"colab\"\n# dir = \"/content/gdrive/MyDrive/siim-effnets-classification-train\"  # directory of this code\n\n# from google.colab import drive\n# drive.mount('/content/gdrive/')\n\n# !pip install kaggle > /dev/null 2>&1\n\n# # Upload your kaggle.json file to connect with your kaggle account\n# from google.colab import files\n# files.upload()  \n# !mkdir -p ~/.kaggle\n# !cp kaggle.json ~/.kaggle/\n# !chmod 600 ~/.kaggle/kaggle.json\n\n# # Download specific kaggle dataset which contains .npy data processed from original .dcm data \n# !mkdir -p '/content/siim-dicom-to-npy-v1'  \n# !kaggle datasets download -d shangweichen/siim-dicom-to-npy-v1 -p '../content/gdrive/MyDrive'\n# !unzip -n '/content/gdrive/MyDrive/siim-dicom-to-npy-v1.zip' -d '/content/siim-dicom-to-npy-v1' > /dev/null 2>&1\n# !rm '/content/gdrive/MyDrive/siim-dicom-to-npy-v1.zip'\n\n# # Download timm package\n# !mkdir -p '/content/timm-pytorch-image-models'\n# !kaggle datasets download -d kozodoi/timm-pytorch-image-models -p '../content/gdrive/MyDrive'\n# !unzip -n '/content/gdrive/MyDrive/timm-pytorch-image-models.zip' -d '/content/timm-pytorch-image-models' > /dev/null 2>&1\n# !rm '/content/gdrive/MyDrive/timm-pytorch-image-models.zip'\n\n# !pip install pydicom > /dev/null 2>&1\n# !pip install --upgrade albumentations > /dev/null 2>&1","metadata":{"id":"ZPmLPzPSF3PM","executionInfo":{"status":"ok","timestamp":1625219711448,"user_tz":-480,"elapsed":92878,"user":{"displayName":"陈尚伟","photoUrl":"","userId":"07754673877521011076"}},"outputId":"4590f540-34bc-4fc0-90ab-18e923daa009","execution":{"iopub.status.busy":"2021-07-03T10:40:44.604451Z","iopub.execute_input":"2021-07-03T10:40:44.60478Z","iopub.status.idle":"2021-07-03T10:40:44.6095Z","shell.execute_reply.started":"2021-07-03T10:40:44.60475Z","shell.execute_reply":"2021-07-03T10:40:44.608511Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if ENVIRONMENT not in [\"kaggle\", \"colab\"]:\n    raise ValueError(\"ENVIRONMENT Wrong!\")","metadata":{"id":"2nN4N0HkMZj2","execution":{"iopub.status.busy":"2021-07-03T10:40:44.611041Z","iopub.execute_input":"2021-07-03T10:40:44.611576Z","iopub.status.idle":"2021-07-03T10:40:44.621871Z","shell.execute_reply.started":"2021-07-03T10:40:44.611464Z","shell.execute_reply":"2021-07-03T10:40:44.620864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Import","metadata":{"id":"19YnjmasHNvQ"}},{"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 2>&1\n!python pytorch-xla-env-setup.py --version 20210331 --apt-packages libomp5 libopenblas-dev > /dev/null 2>&1","metadata":{"_kg_hide-output":true,"scrolled":true,"id":"Cgo4Bg8_1Rfx","executionInfo":{"status":"ok","timestamp":1625219782622,"user_tz":-480,"elapsed":71181,"user":{"displayName":"陈尚伟","photoUrl":"","userId":"07754673877521011076"}},"outputId":"525ef45e-d67d-4c87-ee9a-f40992f276e5","execution":{"iopub.status.busy":"2021-07-03T10:40:44.623447Z","iopub.execute_input":"2021-07-03T10:40:44.623913Z","iopub.status.idle":"2021-07-03T10:42:03.453455Z","shell.execute_reply.started":"2021-07-03T10:40:44.623868Z","shell.execute_reply":"2021-07-03T10:42:03.452303Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nif ENVIRONMENT == \"kaggle\":\n    sys.path.append(\"/kaggle/input/timm-pytorch-image-models/pytorch-image-models-master\")\nelif ENVIRONMENT == \"colab\":\n    sys.path.append(\"/content/timm-pytorch-image-models/pytorch-image-models-master\")\n\nimport platform\nimport numpy as np\nimport pandas as pd\nimport os\nfrom tqdm.notebook import tqdm\nimport cv2\nimport pydicom\nimport random\nimport glob\nimport gc\nfrom math import ceil\nimport albumentations as A\nimport matplotlib.pyplot as plt\nfrom pydicom.pixel_data_handlers.util import apply_voi_lut\nfrom sklearn.metrics import roc_auc_score, confusion_matrix\nfrom sklearn.model_selection import StratifiedKFold, train_test_split\nimport torch\nimport timm\nimport time\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.cuda.amp import GradScaler, autocast\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau, LambdaLR, StepLR\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.test.test_utils as test_utils\nimport warnings\n\nwarnings.simplefilter('ignore')\nnp.set_printoptions(suppress=True)\nos.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'","metadata":{"id":"surgical-consultancy","executionInfo":{"status":"error","timestamp":1625219784481,"user_tz":-480,"elapsed":1870,"user":{"displayName":"陈尚伟","photoUrl":"","userId":"07754673877521011076"}},"outputId":"97decb1b-0939-429f-98c1-5afee40ac3ca","execution":{"iopub.status.busy":"2021-07-03T10:52:01.461059Z","iopub.execute_input":"2021-07-03T10:52:01.461699Z","iopub.status.idle":"2021-07-03T10:52:01.473556Z","shell.execute_reply.started":"2021-07-03T10:52:01.46166Z","shell.execute_reply":"2021-07-03T10:52:01.472654Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if ENVIRONMENT == \"kaggle\":\n    print(np.load(f\"/kaggle/input/siim-dicom-to-npy-v1/trainset/000a312787f2.npy\").shape)\nelif ENVIRONMENT == \"colab\":\n    print(np.load(f\"/content/siim-dicom-to-npy-v1/trainset/000a312787f2.npy\").shape)","metadata":{"id":"AJge5f_o9AJq","execution":{"iopub.status.busy":"2021-07-03T10:52:01.474849Z","iopub.execute_input":"2021-07-03T10:52:01.475259Z","iopub.status.idle":"2021-07-03T10:52:01.512262Z","shell.execute_reply.started":"2021-07-03T10:52:01.475227Z","shell.execute_reply":"2021-07-03T10:52:01.511116Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Necessary Settings\nIf choose effb7 model, gradient will explode if set lr to 1e-3/8, and batch size cannot reach 16 * 8 otherwise will raise OOM(out of memory) error, this is the most intuitive difference between this notebook and [tensorflow tpu code](https://www.kaggle.com/h053473666/siim-covid19-efnb7-train-study/data).","metadata":{"id":"speaking-bloom"}},{"cell_type":"code","source":"class Config:\n    train_pcent = 0.80\n    model_name = 'tf_efficientnet_b7'\n    image_size = (400, 400)\n    batch_size = 8 * 8\n    epochs = 20\n    seed = 2021\n    lr = 1e-4 / 8  \n    workers = 8\n    drop_last = True\n    augments = A.Compose([\n                   A.augmentations.crops.transforms.RandomResizedCrop(height=image_size[1], \n                                                                      width=image_size[0], \n                                                                      scale=(0.88*0.88, 1), \n                                                                      ratio=(0.8, 1.2), \n                                                                      p=0.5),\n                   A.augmentations.transforms.HorizontalFlip(p=0.5),\n                   A.augmentations.transforms.VerticalFlip(p=0.5),\n                   A.augmentations.geometric.rotate.Rotate(p=0.5),\n                   A.OneOf([\n                       A.augmentations.transforms.Blur(),\n                       A.augmentations.transforms.GlassBlur(),\n                       A.augmentations.transforms.GaussianBlur(),\n                       A.augmentations.transforms.GaussNoise(),\n                       A.augmentations.transforms.RandomGamma(),\n                       A.augmentations.transforms.InvertImg(),\n                       A.augmentations.transforms.RandomFog()\n                   ], p=0.5)\n    ])\n\n    def get_loss_fn():\n        return nn.CrossEntropyLoss()\n\n    def get_optimizer(model, learning_rate):\n        return torch.optim.Adam(model.parameters(), lr=learning_rate)\n\n    def get_scheduler(optimizer):\n        return ReduceLROnPlateau(optimizer, \n                                 mode='min', \n                                 factor=0.1, \n                                 patience=5, \n                                 verbose=False, \n                                 min_lr=1e-5)\n    \n\n# Make results reproducible\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.benchmark = False\n    torch.backends.cudnn.deterministic = True\n    \nseed_everything(Config.seed)","metadata":{"id":"distinct-attraction","execution":{"iopub.status.busy":"2021-07-03T10:52:01.514074Z","iopub.execute_input":"2021-07-03T10:52:01.514437Z","iopub.status.idle":"2021-07-03T10:52:01.530335Z","shell.execute_reply.started":"2021-07-03T10:52:01.514402Z","shell.execute_reply":"2021-07-03T10:52:01.529211Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Preprocessing","metadata":{"id":"civilian-monte"}},{"cell_type":"code","source":"# Read file\nif ENVIRONMENT == \"kaggle\":\n    train_image = pd.read_csv(f\"/kaggle/input/siim-covid19-detection/train_image_level.csv\")\n    train_study = pd.read_csv(f\"/kaggle/input/siim-covid19-detection/train_study_level.csv\")\nelif ENVIRONMENT == \"colab\":\n    train_image = pd.read_csv(f\"{dir}/SIIM-FISABIO-RSNA COVID-19 Detection/train_image_level.csv\")\n    train_study = pd.read_csv(f\"{dir}/SIIM-FISABIO-RSNA COVID-19 Detection/train_study_level.csv\")\n\n# Merge study_level and image_level\ntrain_study['StudyInstanceUID'] = train_study['id'].apply(lambda x: x.replace('_study', ''))\ntrain_study.drop(['id'], axis=1, inplace=True)\ntrain = train_image.merge(train_study, on='StudyInstanceUID', how='left')\ntrain[\"class\"] = np.argmax(train[['Negative for Pneumonia', 'Typical Appearance',\n                  'Indeterminate Appearance', 'Atypical Appearance']].values, axis=1)\n\ntrain.head()","metadata":{"id":"indoor-hypothetical","execution":{"iopub.status.busy":"2021-07-03T10:52:01.532018Z","iopub.execute_input":"2021-07-03T10:52:01.532483Z","iopub.status.idle":"2021-07-03T10:52:01.659175Z","shell.execute_reply.started":"2021-07-03T10:52:01.532445Z","shell.execute_reply":"2021-07-03T10:52:01.657709Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SIIMData(Dataset):\n    def __init__(self, df, augments=True):\n        super().__init__()\n        self.df = df.sample(frac=1).reset_index(drop=True)\n        \n        if augments:\n            self.augments = Config.augments\n        else:\n            self.augments = None\n        \n    def __getitem__(self, idx):\n        index = self.df.loc[idx, \"id\"].split(\"_\")[0]\n        \n        if ENVIRONMENT == \"kaggle\":\n            image = np.load(f\"/kaggle/input/siim-dicom-to-npy-v1/trainset/{index}.npy\")\n        elif ENVIRONMENT == \"colab\":\n            image = np.load(f\"/content/siim-dicom-to-npy-v1/trainset/{index}.npy\")\n            \n        if self.augments:\n            image = torch.from_numpy(self.augments(image=image)['image'])\n        else:\n            image = torch.tensor(image)\n            \n        image = image.permute(2, 0, 1)\n        label = torch.tensor(self.df.loc[idx][\"class\"])\n        return image / 255.0, label\n    \n    def __len__(self):\n        return len(self.df)","metadata":{"id":"framed-municipality","execution":{"iopub.status.busy":"2021-07-03T10:52:01.660736Z","iopub.execute_input":"2021-07-03T10:52:01.661118Z","iopub.status.idle":"2021-07-03T10:52:01.670671Z","shell.execute_reply.started":"2021-07-03T10:52:01.661083Z","shell.execute_reply":"2021-07-03T10:52:01.669884Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_data, valid_data = train_test_split(train, test_size=1-Config.train_pcent, \n                       stratify=train[\"class\"].values, random_state=Config.seed)\nprint(f\"Training on {train_data.shape[0]} samples and Validation on {valid_data.shape[0]} samples\")\n\ntrain_set = SIIMData(df=train_data, augments=True)\nvalid_set = SIIMData(df=valid_data, augments=False)","metadata":{"id":"asm9JURz1Rf7","execution":{"iopub.status.busy":"2021-07-03T10:52:01.672423Z","iopub.execute_input":"2021-07-03T10:52:01.673105Z","iopub.status.idle":"2021-07-03T10:52:01.703069Z","shell.execute_reply.started":"2021-07-03T10:52:01.673062Z","shell.execute_reply":"2021-07-03T10:52:01.701866Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{"id":"informational-prison"}},{"cell_type":"code","source":"class EfficientNetModel(nn.Module):\n    \"\"\"\n    Model Class for EfficientNet Model\n    \"\"\"\n    def __init__(self, num_classes=4, model_name=Config.model_name, pretrained=True):\n        super(EfficientNetModel, self).__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained, in_chans=3)\n        self.model.classifier = nn.Linear(self.model.classifier.in_features, num_classes)\n        \n    def forward(self, x):\n        x = self.model(x)\n        return x\n    \n\nclass NFNetModel(nn.Module):\n    \"\"\"\n    Model Class for EfficientNet Model\n    \"\"\"\n    def __init__(self, num_classes=4, model_name=Config.model_name, pretrained=True):\n        super(NFNetModel, self).__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained, in_chans=3)\n        self.model.head.fc = nn.Linear(self.model.head.fc.in_features, num_classes)\n        \n    def forward(self, x):\n        x = self.model(x)\n        return x","metadata":{"id":"analyzed-opportunity","execution":{"iopub.status.busy":"2021-07-03T10:52:01.704571Z","iopub.execute_input":"2021-07-03T10:52:01.704901Z","iopub.status.idle":"2021-07-03T10:52:01.714607Z","shell.execute_reply.started":"2021-07-03T10:52:01.704869Z","shell.execute_reply":"2021-07-03T10:52:01.71355Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if \"efficient\" in Config.model_name or \"eff\" in Config.model_name: \n    model_ = EfficientNetModel()\nelif \"nfnet\" in Config.model_name:\n    model_ = NFNetModel()\nelse:\n    raise RuntimeError(\"Must specify a valid model type to train.\")\nprint(f\"Training Model: {Config.model_name}\")","metadata":{"id":"dietary-pressing","execution":{"iopub.status.busy":"2021-07-03T10:52:01.716284Z","iopub.execute_input":"2021-07-03T10:52:01.716586Z","iopub.status.idle":"2021-07-03T10:52:03.22358Z","shell.execute_reply.started":"2021-07-03T10:52:01.716544Z","shell.execute_reply":"2021-07-03T10:52:03.222475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train & Valid","metadata":{"id":"hDX-3Hi01Rf-"}},{"cell_type":"code","source":"class Record:\n    '''\n    Records labels and predictions within one epoch\n    '''\n    def __init__(self):\n        self.labels = []\n        self.preds = []\n        \n    def update(self, cur_labels, cur_logits):\n        cur_labels = cur_labels.detach().cpu().numpy()\n        cur_logits = np.exp(cur_logits.detach().cpu().numpy())\n        cur_preds = cur_logits / np.sum(cur_logits, axis=1, keepdims=True)\n        self.labels.append(cur_labels)\n        self.preds.append(cur_preds)\n\n    def get_labels(self):\n        return np.concatenate(self.labels) # (n, )\n\n    def get_preds(self):\n        return np.concatenate(self.preds, axis=0) # (n, 4)\n\n    @staticmethod\n    def get_acc(confusion_mat):\n        return round(np.sum(np.eye(4) * confusion_mat) / np.sum(confusion_mat) * 100, 2)\n    \n    @staticmethod\n    def get_auc(labels, preds):\n        return round(roc_auc_score(labels, preds, average='weighted', multi_class='ovr'), 2)\n\n\nclass Trainer:\n    def __init__(self, model, optimizer, loss_fn, device):\n        \"\"\"\n        Constructor for Trainer class\n        \"\"\"\n        self.model = model\n        self.optimizer = optimizer\n        self.loss_fn = loss_fn\n        self.device = device\n    \n    def train_one_cycle(self, train_loader):\n        \"\"\"\n        Runs one epoch of training, backpropagation and optimization\n        \"\"\"\n        self.model.train()\n        total_loss = 0\n        total_nums = 0\n        record = Record()\n\n        for idx, (xtrain, ytrain) in enumerate(train_loader):\n            xtrain = xtrain.to(self.device, dtype=torch.float)\n            ytrain = ytrain.to(self.device, dtype=torch.long)\n\n            self.optimizer.zero_grad()\n            outputs = self.model(xtrain)\n            loss = self.loss_fn(outputs, ytrain)\n            \n            total_loss += (loss.detach().item() * ytrain.size(0))\n            total_nums += ytrain.size(0)\n            record.update(ytrain, outputs)\n            \n            loss.backward()\n            # The step() function now not only propagates gradients, but uses the Cloud TPU context \n            # to synchronize gradient updates across each processes' copy of the network. \n            # This ensures that each processes' network copy stays \"in sync\" (they are all identical).\n            xm.optimizer_step(self.optimizer)\n            \n        self.model.eval()\n        return total_loss / total_nums, record.get_labels(), record.get_preds()\n\n    def valid_one_cycle(self, valid_loader):\n        \"\"\"\n        Runs one epoch of prediction\n        \"\"\"\n        self.model.eval()\n        total_loss = 0\n        total_nums = 0\n        record = Record()\n        \n        for idx, (xval, yval) in enumerate(valid_loader):\n            with torch.no_grad():\n                xval = xval.to(self.device, dtype=torch.float)\n                yval = yval.to(self.device, dtype=torch.long)\n        \n                outputs = self.model(xval)\n                loss = self.loss_fn(outputs, yval)\n\n                total_loss += (loss.detach().item() * yval.size(0))\n                total_nums += yval.size(0)\n                record.update(yval, outputs)\n        \n        return total_loss / total_nums, record.get_labels(), record.get_preds()","metadata":{"id":"spiritual-mineral","execution":{"iopub.status.busy":"2021-07-03T10:52:03.225284Z","iopub.execute_input":"2021-07-03T10:52:03.225597Z","iopub.status.idle":"2021-07-03T10:52:03.244804Z","shell.execute_reply.started":"2021-07-03T10:52:03.225562Z","shell.execute_reply":"2021-07-03T10:52:03.243663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def _mp_fn(rank, flags):\n    '''\n    Train and valid\n    '''\n    torch.set_default_tensor_type('torch.FloatTensor')\n\n    # Sets a common random seed both for initialization and ensuring graph is the same\n    torch.manual_seed(Config.seed)\n\n    # Acquires the (unique) Cloud TPU core corresponding to this process's index\n    device = xm.xla_device()\n    \n    # load the model into each tpu core\n    model = model_.to(device)\n    \n    # Creates the (distributed) train sampler\n    # which let this process only access its portion of the training dataset.  \n    train_sampler = DistributedSampler(\n        train_set,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True,\n    )\n    train_loader = DataLoader(\n        train_set,\n        batch_size=int(Config.batch_size/xm.xrt_world_size()),\n        sampler=train_sampler,\n        drop_last=Config.drop_last,\n        num_workers=Config.workers,\n    )\n    valid_sampler = DistributedSampler(\n        valid_set,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False,\n    )\n    valid_loader = DataLoader(\n        valid_set,\n        batch_size=int(Config.batch_size/xm.xrt_world_size()),\n        sampler=valid_sampler,\n        drop_last=Config.drop_last,\n        num_workers=Config.workers,\n    )\n\n    optimizer = Config.get_optimizer(model, Config.lr * xm.xrt_world_size())\n    loss_fn = Config.get_loss_fn()\n    scheduler = Config.get_scheduler(optimizer)\n    \n    trainer = Trainer(\n        model=model,\n        optimizer=optimizer,\n        loss_fn=loss_fn,\n        device=device,\n    )\n\n    if rank == 0:\n        train_losses = []\n        valid_losses = []\n        train_accs = []\n        valid_accs = []\n        train_aucs = []\n        valid_aucs = []\n        lr_scheduler = []\n    \n    gc.collect()\n    for epoch in range(Config.epochs):\n        xm.master_print(f\"{'-'*30} EPOCH: {epoch+1}/{Config.epochs} {'-'*30}\")\n        \n        # Run one training epoch\n        para_loader = pl.ParallelLoader(train_loader, [device])\n        train_loss, train_labels, train_preds = trainer.train_one_cycle(para_loader.per_device_loader(device))\n\n        # Compute training metrics\n        train_loss_avg = xm.mesh_reduce('train_loss_reduce', train_loss, lambda alist: sum(alist) / len(alist))\n        train_labels_concat = xm.mesh_reduce('train_labels_concat', train_labels, lambda alist: np.concatenate(alist))  # (8n, )\n        train_preds_concat = xm.mesh_reduce('train_preds_concat', train_preds, lambda alist: np.concatenate(alist, axis=0))  # (8n, 4)\n        train_confusion_mat = confusion_matrix(train_labels_concat, np.argmax(train_preds_concat, axis=1))\n        train_acc = Record.get_acc(train_confusion_mat)\n        train_auc = Record.get_auc(train_labels_concat, train_preds_concat)\n        xm.master_print(f\"Train Loss: {train_loss_avg:.4f}  Train Acc: {train_acc}%  Train AUC: {train_auc}\")\n        xm.master_print(train_confusion_mat)\n\n        # Run one validation epoch\n        para_loader = pl.ParallelLoader(valid_loader, [device])\n        valid_loss, valid_labels, valid_preds = trainer.valid_one_cycle(para_loader.per_device_loader(device))\n        \n        # Compute validation metrics\n        valid_loss_avg = xm.mesh_reduce('valid_loss_reduce', valid_loss, lambda alist: sum(alist) / len(alist))\n        valid_labels_concat = xm.mesh_reduce('valid_labels_concat', valid_labels, lambda alist: np.concatenate(alist))\n        valid_preds_concat = xm.mesh_reduce('valid_preds_concat', valid_preds, lambda alist: np.concatenate(alist, axis=0))\n        valid_confusion_mat = confusion_matrix(valid_labels_concat, np.argmax(valid_preds_concat, axis=1))\n        valid_acc = Record.get_acc(valid_confusion_mat)\n        valid_auc = Record.get_auc(valid_labels_concat, valid_preds_concat)\n        xm.master_print(f\"Valid Loss: {valid_loss_avg:.4f}  Valid Acc: {valid_acc}%  Valid AUC: {valid_auc}\")\n        xm.master_print(valid_confusion_mat)\n\n        scheduler.step(valid_loss_avg)\n        gc.collect()\n        \n        if rank == 0:\n            train_losses.append(train_loss_avg)\n            train_accs.append(train_acc)\n            train_aucs.append(train_auc)\n            valid_losses.append(valid_loss_avg)\n            valid_accs.append(valid_acc)\n            valid_aucs.append(valid_auc)\n            lr_scheduler.append(optimizer.param_groups[0]['lr'])\n            \n    if rank == 0:\n        # All tpu cores need to do below operations if without \"if rank == 0\"\n        np.save(f\"{dir}/train_losses\", np.array(train_losses))\n        np.save(f\"{dir}/train_accs\", np.array(train_accs))\n        np.save(f\"{dir}/train_aucs\", np.array(train_aucs))\n        np.save(f\"{dir}/valid_losses\", np.array(valid_losses))\n        np.save(f\"{dir}/valid_accs\", np.array(valid_accs))\n        np.save(f\"{dir}/valid_aucs\", np.array(valid_aucs))\n        np.save(f\"{dir}/lr_scheduler\", np.array(lr_scheduler))\n        \n    # Only main process save model weights\n    xm.save(model.state_dict(), f\"{dir}/pretrained_model.bin\")","metadata":{"id":"sitting-addiction","execution":{"iopub.status.busy":"2021-07-03T10:52:03.246264Z","iopub.execute_input":"2021-07-03T10:52:03.246567Z","iopub.status.idle":"2021-07-03T10:52:03.272903Z","shell.execute_reply.started":"2021-07-03T10:52:03.246537Z","shell.execute_reply":"2021-07-03T10:52:03.271784Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"spawn() takes a function (the \"map function\"), a tuple of arguments (the placeholder flags dict), the number of processes to create, and whether to create these new processes by \"forking\" or \"spawning.\" While spawning new processes is generally recommended, Colab only supports forking.\n\nspawn() will create eight processes, one for each Cloud TPU core, and call _mp_fn() -- the map function -- on each process. The inputs to _mp_fn() are an index (zero through seven, process id) and the placeholder flags. When the processes acquire their device they actually acquire their corresponding Cloud TPU core automatically.","metadata":{}},{"cell_type":"code","source":"%%time\n\nFLAGS = {}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","metadata":{"id":"wOcf51O-1RgA","execution":{"iopub.status.busy":"2021-07-03T10:52:03.274712Z","iopub.execute_input":"2021-07-03T10:52:03.275275Z","iopub.status.idle":"2021-07-03T11:00:58.807294Z","shell.execute_reply.started":"2021-07-03T10:52:03.275238Z","shell.execute_reply":"2021-07-03T11:00:58.79599Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Plot","metadata":{"id":"C0QCVVBx1RgA"}},{"cell_type":"code","source":"def get_first_occur(learning_rates_list):\n    learning_rates_list = learning_rates_list.tolist()\n    unique_lrs = np.sort(np.unique(learning_rates_list))\n    first_occur = []\n    for lr in unique_lrs[:-1]:\n        first_occur.append(learning_rates_list.index(lr))\n    return np.array(first_occur) - 1\n\n\ntrain_losses = np.load(f\"{dir}/train_losses.npy\")\ntrain_accs = np.load(f\"{dir}/train_accs.npy\")\ntrain_aucs = np.load(f\"{dir}/train_aucs.npy\")\nvalid_losses = np.load(f\"{dir}/valid_losses.npy\")\nvalid_accs = np.load(f\"{dir}/valid_accs.npy\")\nvalid_aucs = np.load(f\"{dir}/valid_aucs.npy\")\nlr_scheduler= np.load(f\"{dir}/lr_scheduler.npy\")\nfirst_occur = get_first_occur(lr_scheduler)\n\nfig, ax = plt.subplots(2, 2, figsize=(20, 16))\n\nbest_train_loss = train_losses[np.argmin(valid_losses)]\nbest_valid_loss = np.min(valid_losses)\nax[0][0].plot(range(1, len(train_losses)+1), train_losses, \"bo-\", label=\"train_loss\")\nax[0][0].plot(range(1, len(train_losses)+1), valid_losses, \"go-\", label=\"valid_loss\")\nax[0][0].scatter(first_occur+1, train_losses[first_occur], s=300, c=\"r\", marker=\"*\")\nax[0][0].scatter(first_occur+1, valid_losses[first_occur], s=300, c=\"r\", marker=\"*\")\nax[0][0].legend()\nax[0][0].grid()\nax[0][0].set_title(f\"Best Train Loss: {np.round(best_train_loss, 4)}  Best Valid Loss: {np.round(best_valid_loss, 4)}\")\n\nbest_train_acc = train_accs[np.argmax(valid_accs)]\nbest_valid_acc = np.max(valid_accs)\nax[0][1].plot(range(1, len(train_accs)+1), train_accs, \"bo-\", label=\"train_acc\")\nax[0][1].plot(range(1, len(train_accs)+1), valid_accs, \"go-\", label=\"valid_acc\")\nax[0][1].scatter(first_occur+1, train_accs[first_occur], s=300, c=\"r\", marker=\"*\")\nax[0][1].scatter(first_occur+1, valid_accs[first_occur], s=300, c=\"r\", marker=\"*\")\nax[0][1].legend()\nax[0][1].grid()\nax[0][1].set_title(f\"Best Train Acc: {np.round(best_train_acc, 3)}  Best Valid Acc: {np.round(best_valid_acc, 3)}\")\n\nbest_train_auc = train_aucs[np.argmax(valid_aucs)]\nbest_valid_auc = np.max(valid_aucs)\nax[1][0].plot(range(1, len(train_aucs)+1), train_aucs, \"bo-\", label=\"train_auc\")\nax[1][0].plot(range(1, len(train_aucs)+1), valid_aucs, \"go-\", label=\"valid_auc\")\nax[1][0].scatter(first_occur+1, train_aucs[first_occur], s=300, c=\"r\", marker=\"*\")\nax[1][0].scatter(first_occur+1, valid_aucs[first_occur], s=300, c=\"r\", marker=\"*\")\nax[1][0].legend()\nax[1][0].grid()\nax[1][0].set_title(f\"Best Train Auc: {np.round(best_train_auc, 3)}  Best Valid Auc: {np.round(best_valid_auc, 3)}\")\n\nax[1][1].plot(range(1, len(lr_scheduler)+1), lr_scheduler, \"ro-\", label=\"learning_rate\")\nax[1][1].legend()\nax[1][1].grid()\nax[1][1].set_title(f\"MinLR={np.round(min(lr_scheduler), 6)}\")\n\nos.remove(f'{dir}/train_losses.npy')\nos.remove(f'{dir}/train_accs.npy')\nos.remove(f'{dir}/train_aucs.npy')\nos.remove(f'{dir}/valid_losses.npy')\nos.remove(f'{dir}/valid_accs.npy')\nos.remove(f'{dir}/valid_aucs.npy')\nos.remove(f'{dir}/lr_scheduler.npy')","metadata":{"id":"cosmetic-colleague","execution":{"iopub.status.busy":"2021-07-03T11:00:58.812983Z","iopub.execute_input":"2021-07-03T11:00:58.817429Z","iopub.status.idle":"2021-07-03T11:00:59.027553Z","shell.execute_reply.started":"2021-07-03T11:00:58.817365Z","shell.execute_reply":"2021-07-03T11:00:59.023717Z"},"trusted":true},"execution_count":null,"outputs":[]}]}