{"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":"# Basic Training Notebook\nWithout using One Cycle Policy, this is the most basic method of training. \n\nFirst, make sure you set `tpu_utility_1` and `tpu_cache_ds_utils` as utility script and **add** them to this notebook. Make sure the **scripts are python scripts when doing Quick Save rather than Jupyter Notebook** by changing them at File $\\rightarrow$ Editor Type $\\rightarrow$ Script, or the import will fail. ","metadata":{}},{"cell_type":"code","source":"from tpu_utility_1 import *\n\n# For jpeg file we use this library\n!apt-get install libturbojpeg\n!pip install jpeg4py\n\nfrom IPython.display import clear_output\nclear_output()","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:41:58.044133Z","iopub.execute_input":"2021-08-21T03:41:58.044573Z","iopub.status.idle":"2021-08-21T03:43:20.57832Z","shell.execute_reply.started":"2021-08-21T03:41:58.044469Z","shell.execute_reply":"2021-08-21T03:43:20.577289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Below is the starting code cell you get when creating a new Kaggle notebook. We comment out the `for dirname...` to prevent it listing an extremely long list of items (if you have lots of items). ","metadata":{}},{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n# for dirname, _, filenames in os.walk('/kaggle/input'):\n#     for filename in filenames:\n#         print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-08-21T03:43:20.580228Z","iopub.execute_input":"2021-08-21T03:43:20.580538Z","iopub.status.idle":"2021-08-21T03:43:20.586007Z","shell.execute_reply.started":"2021-08-21T03:43:20.580505Z","shell.execute_reply":"2021-08-21T03:43:20.58497Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import albumentations\n\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import f1_score, roc_auc_score\n\nimport matplotlib.pyplot as plt\nimport jpeg4py as jpeg  # will fail if you don't run the first code cell. \nimport pathlib\nfrom pathlib import Path\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.data as D\nfrom torchvision import models\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.utils.utils as xu","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:20.588021Z","iopub.execute_input":"2021-08-21T03:43:20.58839Z","iopub.status.idle":"2021-08-21T03:43:21.764351Z","shell.execute_reply.started":"2021-08-21T03:43:20.588347Z","shell.execute_reply":"2021-08-21T03:43:21.763555Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We list down some `FLAGS` (python dict containing the configurations).","metadata":{}},{"cell_type":"code","source":"FLAGS = {\n    \"data_dir\": Path(\"../input/plant-pathology-2021-fgvc8/train_images\"),\n    \"bs\": 32,  # Batch size PER TPU CORE. Total bs = bs * xrt.world_size()\n    \"num_workers\": os.cpu_count(),  # for pre-fetching data. \n    \"lr\": 0.001,  # this will multiply by xrt.world_size() for multiprocessing. \n    \"momentum\": 0.9,  # ONLY IF using SGD\n    \"num_epochs\": 10,\n    \"num_cores\": 8 if os.environ.get(\"TPU_NAME\", None) else 1,\n    \"log_steps\": 20,  # We are not using this but defined for historical reasons. \n    \"metrics_debug\": False,\n    \"seed\": 1397,\n    \"save_path\": \"/kaggle/working/model.pth\"\n}","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:21.76558Z","iopub.execute_input":"2021-08-21T03:43:21.765964Z","iopub.status.idle":"2021-08-21T03:43:21.771504Z","shell.execute_reply.started":"2021-08-21T03:43:21.765935Z","shell.execute_reply":"2021-08-21T03:43:21.770213Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"First, preprocess the data into the format we want to feed into the model. ","metadata":{}},{"cell_type":"code","source":"diseases = ['healthy', 'rust', 'scab', 'frog_eye_leaf_spot', 'powdery_mildew', 'complex']\ndf = pd.read_csv('../input/plant-pathology-2021-fgvc8/train.csv')\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:21.77261Z","iopub.execute_input":"2021-08-21T03:43:21.772927Z","iopub.status.idle":"2021-08-21T03:43:21.857529Z","shell.execute_reply.started":"2021-08-21T03:43:21.772901Z","shell.execute_reply":"2021-08-21T03:43:21.856485Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"labels = df.labels.str.split()\nocc = np.zeros((len(labels),len(diseases)))\ncount = 0\nfor i in diseases:\n    occ[:,count] = [i in list for list in labels]\n    count+=1\n\ndis_df = pd.DataFrame(occ.astype(np.float32), columns=diseases)\ndis_df.head()","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:21.858684Z","iopub.execute_input":"2021-08-21T03:43:21.858955Z","iopub.status.idle":"2021-08-21T03:43:22.08145Z","shell.execute_reply.started":"2021-08-21T03:43:21.858929Z","shell.execute_reply":"2021-08-21T03:43:22.080642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch_df = pd.concat([df[\"image\"], dis_df], axis=1)\ntorch_df.head()","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:22.082661Z","iopub.execute_input":"2021-08-21T03:43:22.082948Z","iopub.status.idle":"2021-08-21T03:43:22.103475Z","shell.execute_reply.started":"2021-08-21T03:43:22.08292Z","shell.execute_reply":"2021-08-21T03:43:22.102566Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We are lazy enough to use a StratifiedKFold to split our data (than do it ourselves). In real, training should be done with the number of folds you specify, repeatedly. Here, we will just fetch the last fold as a demonstration. You can easily modify the code to your need. ","metadata":{}},{"cell_type":"code","source":"skf = StratifiedKFold(n_splits=5)\n\ntdf_np = df.to_numpy()\nX, y = tdf_np[:, 0], tdf_np[:, 1]\nfor train_index, test_index in skf.split(X, y): pass","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:22.105662Z","iopub.execute_input":"2021-08-21T03:43:22.105952Z","iopub.status.idle":"2021-08-21T03:43:22.153911Z","shell.execute_reply.started":"2021-08-21T03:43:22.105924Z","shell.execute_reply":"2021-08-21T03:43:22.152874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"After this, we shall now have our `train_index` and `test_index` in place. Be sure to not accidentally delete them or you would need to run again. Optionally, see if you could set seed to get reproducible outcome. \n\nOur final array would look like this.","metadata":{}},{"cell_type":"code","source":"torch_df.to_numpy()[train_index]","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:22.155504Z","iopub.execute_input":"2021-08-21T03:43:22.155802Z","iopub.status.idle":"2021-08-21T03:43:22.171917Z","shell.execute_reply.started":"2021-08-21T03:43:22.155772Z","shell.execute_reply":"2021-08-21T03:43:22.170779Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Defining Dataset\nThis is the usual thing you do in PyTorch to define how to fetch your data. No code change is required. You are required to write this code yourself **unless you have a `xcd.CachedDataset()` saved** (in which case you don't need to define these below. ","metadata":{}},{"cell_type":"code","source":"class AppleDataset(D.Dataset):\n    def __init__(self, parent_path, df=None, shuffle=False,\n                 seed=None, transform=None, split=None):\n        \"\"\"\n        parent_path: (pathlib.Path) Parent path of images. \n        df: (pandas.DataFrame) y-labels dataframe. \n        shuffle: (Boolean) Shuffle dataset? Default: False. \n        seed: (integer) seed. Default: None. \n        transforms: (Albumentations) Transformation to images, default: None. \n            Requires writing more code to use PyTorch Transform on cpu. \n        split: (python list) Splits index for train test split. \n            Here we pass in train_index or test_index. \n        \"\"\"\n        self.df_np = df.to_numpy()\n        self.parent_path = Path(parent_path) if type(parent_path) == str else parent_path\n        self.seed = seed\n        self.transform = transform\n        \n        if type(split) != type(None): self.df_np = self.df_np[split]\n            \n    def __len__(self):\n        return len(self.df_np)\n    \n    def __getitem__(self, idx):\n        item = self.df_np[idx]\n        item_path = self.parent_path / item[0]  # 'image' column\n        image = jpeg.JPEG(item_path).decode()\n        \n        if self.transform is not None:  # only works for albumentations library. \n            image = self.transform(**{\"image\": image})[\"image\"]\n            \n        image = torch.from_numpy(image).permute(2, 0, 1)  # HWC -> CHW format\n        target = torch.from_numpy(item[1:].astype(np.float32))\n        \n        return image, target  ","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:22.17326Z","iopub.execute_input":"2021-08-21T03:43:22.173534Z","iopub.status.idle":"2021-08-21T03:43:22.183816Z","shell.execute_reply.started":"2021-08-21T03:43:22.173508Z","shell.execute_reply":"2021-08-21T03:43:22.182799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Then we list the albumentation transforms we want to perform on the dataset. ","metadata":{}},{"cell_type":"code","source":"size = 224  # final image size\n\ntrain_transform = albumentations.Compose([\n    albumentations.RandomResizedCrop(height=size, width=size, always_apply=True, scale=(0.5, 1.0)),\n    albumentations.Flip(),\n    albumentations.Rotate(limit=180, p=0.75),\n    albumentations.CoarseDropout(max_holes=4, p=0.1),\n    albumentations.RandomBrightnessContrast(p=0.75),\n    albumentations.Normalize(always_apply=True, p=1.0)\n])\n\nval_transform = albumentations.Compose([\n    albumentations.Resize(height=size, width=size, always_apply=True),\n    albumentations.Normalize(always_apply=True, p=1.0)\n])","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:22.185362Z","iopub.execute_input":"2021-08-21T03:43:22.185763Z","iopub.status.idle":"2021-08-21T03:43:22.198339Z","shell.execute_reply.started":"2021-08-21T03:43:22.185722Z","shell.execute_reply":"2021-08-21T03:43:22.197339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Define** our `get_dataset` magic function. ","metadata":{}},{"cell_type":"code","source":"def get_dataset():\n    parent_path = FLAGS[\"data_dir\"]\n    train_ds = AppleDataset(parent_path, torch_df, shuffle=True, seed=FLAGS[\"seed\"],\n                           transform=train_transform, split=train_index)\n    val_ds = AppleDataset(parent_path, torch_df, shuffle=False, seed=FLAGS[\"seed\"],\n                         transform=val_transform, split=test_index)\n    \n    return train_ds, val_ds  # required return format","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:22.199698Z","iopub.execute_input":"2021-08-21T03:43:22.19999Z","iopub.status.idle":"2021-08-21T03:43:22.213297Z","shell.execute_reply.started":"2021-08-21T03:43:22.199963Z","shell.execute_reply":"2021-08-21T03:43:22.212511Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model definition\nLet's define our model. ","metadata":{}},{"cell_type":"code","source":"num_classes = len(diseases)\n\nmodel = models.resnet50(pretrained=True)\nnum_features = model.fc.in_features\n\n# freeze models\nfor param in model.parameters():\n    param.requires_grad_(False)","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:22.214361Z","iopub.execute_input":"2021-08-21T03:43:22.214743Z","iopub.status.idle":"2021-08-21T03:43:28.679867Z","shell.execute_reply.started":"2021-08-21T03:43:22.214716Z","shell.execute_reply":"2021-08-21T03:43:28.679102Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Then we can define the head. Here we're using a more complex head. But you can always do `model.fc = nn.Linear(num_features, num_classes)` if you want to. ","metadata":{}},{"cell_type":"code","source":"def create_head(num_features, num_classes, dropout=0.1, act_func=nn.ReLU):\n    features_lst = [num_features, num_features // 2, num_features // 4]\n    layers = []\n    \n    for in_f, out_f in zip(features_lst[:-1], features_lst[1:]):\n        layers.append(nn.Linear(in_f, out_f))\n        layers.append(act_func())\n        layers.append(nn.BatchNorm1d(out_f))\n        if dropout != 0: layers.append(nn.Dropout(dropout))\n            \n    layers.append(nn.Linear(features_lst[-1], num_classes))\n    return nn.Sequential(*layers)\n\n\nmodel.fc = create_head(num_features, num_classes)","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:28.681089Z","iopub.execute_input":"2021-08-21T03:43:28.681666Z","iopub.status.idle":"2021-08-21T03:43:28.707866Z","shell.execute_reply.started":"2021-08-21T03:43:28.681624Z","shell.execute_reply":"2021-08-21T03:43:28.707147Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## LR Finder","metadata":{}},{"cell_type":"code","source":"opt = torch.optim.Adam(model.parameters(), lr=FLAGS[\"lr\"])\ncriterion = nn.BCEWithLogitsLoss()","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:28.709124Z","iopub.execute_input":"2021-08-21T03:43:28.70971Z","iopub.status.idle":"2021-08-21T03:43:28.718093Z","shell.execute_reply.started":"2021-08-21T03:43:28.709666Z","shell.execute_reply":"2021-08-21T03:43:28.717429Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Define the dataloader","metadata":{}},{"cell_type":"code","source":"train_ds, val_ds = get_dataset()\ndls = dataloader(train_ds, val_ds, FLAGS)\n\nlr_finder(model, opt, criterion, dls, device=xm.xla_device())","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:28.719232Z","iopub.execute_input":"2021-08-21T03:43:28.71977Z","iopub.status.idle":"2021-08-21T03:43:28.745993Z","shell.execute_reply.started":"2021-08-21T03:43:28.719726Z","shell.execute_reply":"2021-08-21T03:43:28.745185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Then Restart Notebook\nThe reason is because we've called `xm.xla_device()` so multiprocessing cannot run anymore unless we don't use multiprocessing. Comment out the previous cell. \n\n**So long you call `xm.xla_device()` ONCE PER KERNEL RESTART** you won't have this specific error (calling 8 device but using 1) (or something like that). \n\n**To NOT USE MULTIPROCESSING**, change the `FLAGS['num_cores'] = 1` forcefully. \n\nOtherwise, you can just do a normal training loop only changing `device` to `xm.xla_device()` AND `opt.step()` to `xm.optimizer_step(opt, barrier=True)`. \n\nFirst we will define `train_tpu()`, where **every of our function requires to be inside this**. ","metadata":{}},{"cell_type":"code","source":"SERIAL_EXEC = xmp.MpSerialExecutor()\nWRAPPED_MODEL = xmp.MpModelWrapper(model)","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:43:28.747544Z","iopub.execute_input":"2021-08-21T03:43:28.7482Z","iopub.status.idle":"2021-08-21T03:43:28.753118Z","shell.execute_reply.started":"2021-08-21T03:43:28.748153Z","shell.execute_reply":"2021-08-21T03:43:28.752154Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Ultimately in the `train_loop_fn` you define as normal with additional parameters for TPU. These are `tracker = xm.RateTracker()` at the beginning of the loop, using `xm.optimizer_step(opt)` **instead of `opt.step()`**, and `tracker.add(FLAGS[\"bs\"])` inside looping through dataloader. That's it! For `test_loop_fn`, you don't need to change your code. Just ensure you group whatever is required inside (validation loop) and it will work fine. \n\nWhat's more is since we're doing multiprocessing, setting `print(..., flush=True)` will flush printing. \n\nAlso notice that `train_loop_fn(loader)` and `test_loop_fn(loader)` takes one (and only one) input, which is the `train_dl` and `val_dl` respectively. These are available inside `dls` hence passing `dls` into the function `train_cycle_distrib` suffice. \n\n`train_cycle_distrib()` returns a value `returned_val` which contains all the value *grouped together* returned from the `test_loop_fn`. Currently, **it will break if `test_loop_fn` didn't return anything** (one will fix this in the future). ","metadata":{}},{"cell_type":"code","source":"def train_tpu():\n    torch.manual_seed(FLAGS[\"seed\"])\n    \n    dls = distrib_dataloader(get_dataset, FLAGS)\n    \n    device = xm.xla_device()\n    lr = FLAGS[\"lr\"] * xm.xrt_world_size()\n    model = WRAPPED_MODEL.to(device)\n    \n    # If you want scheduler. \n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=5, eta_min=0.005)\n    \n    \n    def train_loop_fn(loader):\n        tracker = xm.RateTracker()  # required in TPU training loop function multiprocessing.\n        \n        model.train()\n        running_loss, total_samples = 0.0, 0\n        \n        for data, target in tqdm(loader):\n            opt.zero_grad()\n            data, target = data.to(device), target.to(device)\n            \n            # (Optional) For image, to use channels last. \n            data = data.to(memory_format=torch.channels_last)\n            \n            output = model(data)\n            loss = criterion(output, target)\n            \n            # For one-hot encoding\n            preds = torch.sigmoid(output).data > 0.5\n            preds = preds.to(torch.float32)  # Don't do float16 since TPU supports bfloat16\n            \n            loss.backward()\n            xm.optimizer_step(opt)\n            scheduler.step()\n            \n            tracker.add(FLAGS[\"bs\"])  # required for TPU. \n            \n            running_loss += loss.item() * data.size(0)\n            total_samples += data.size(0)\n            \n        print(f\"Train loss: {running_loss / total_samples}\", flush=True)\n        \n    \n    def test_loop_fn(loader):\n        total_samples = 0\n        running_loss, f1Score = 0.0, 0.0\n        # One didn't manage to get RocAucScore to work due to some errors. \n        \n        model.eval()\n        \n        for data, target in tqdm(loader):\n            data, target = data.to(device), target.to(device)\n            data = data.to(memory_format=torch.channels_last)\n            \n            output = model(data)\n            loss = criterion(output, target)\n            \n            # One hot encoding\n            preds = torch.sigmoid(output).data > 0.5\n            preds = preds.to(torch.float32)  # use float32, TPU convert it to bfloat16.\n            \n            running_loss += loss.item() * data.size(0)\n            total_samples += data.size(0)\n            \n            # Other calculations for f1Score here\n            target = target.cpu().to(torch.int).numpy()\n            preds = preds.cpu().to(torch.int).numpy()\n            f1Score += f1_score(target, preds, average=\"weighted\") * data.size(0)\n            \n        epoch_loss = running_loss / total_samples\n        epoch_f1score = f1Score / total_samples\n        \n        print(f\"\"\"\n            |   Val loss    |   Val F1Score   |\n            |{epoch_loss:<15}|{epoch_f1score:<17}|\n        \"\"\", flush=True)\n        \n        return epoch_loss, epoch_f1score, data, preds, target\n    \n    \n    returned_val = train_cycle_distrib(dls, FLAGS, train_loop_fn, test_loop_fn, device=device)\n#     epoch_loss, epoch_f1score, data, preds, target = returned_val\n\n    return returned_val, model","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:54:50.181826Z","iopub.execute_input":"2021-08-21T03:54:50.182247Z","iopub.status.idle":"2021-08-21T03:54:50.199391Z","shell.execute_reply.started":"2021-08-21T03:54:50.182209Z","shell.execute_reply":"2021-08-21T03:54:50.198499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Multiprocessing function\nThis is the function that changes a lot and requires you to manually define yourself. Due to its flexibility and dependent on what you requires it to do inside the function, you requires to define it yourself rather than one hiding it away from you. Let's see how to define it. \n\nUsually this function is called `_mp_fn(rank)`. **The first argument must always be rank/index/whatever name you call it**. This is the argument that is always passed in by the TPU executor as a compulsory argument. Rank is the TPU index, signifying which TPU core it is allocated to. This is similar to GPU:0, GPU:1 if you are using multiple GPUs. \n\nThen, you can pass in any additional kwargs. The only thing **you cannot pass in** is `*args` and `**kwargs`. So, for example, we pass in additional argument `flags` inside. ","metadata":{}},{"cell_type":"code","source":"def _mp_fn(rank, flags):\n    global FLAGS\n    FLAGS = flags\n    torch.set_default_tensor_type(torch.FloatTensor)  # This is compulsory. \n    \n    returned_val, model = train_tpu()  # pass in the training function. \n    \n    # One is saving model here at end of epoch. You could do other stuffs. \n    # For example, printing summary, and remember this is OPTIONAL. \n    # You could safely ignore the line below if you don't have anything you want \n    # to do. \n    if rank == 0: torch.save(model.state_dict(), FLAGS[\"save_path\"])","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:49:29.924044Z","iopub.execute_input":"2021-08-21T03:49:29.924807Z","iopub.status.idle":"2021-08-21T03:49:29.932341Z","shell.execute_reply.started":"2021-08-21T03:49:29.924763Z","shell.execute_reply":"2021-08-21T03:49:29.931348Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The reason we ask for `if rank == 0` is because the TPU cores process stuffs not necessarily synchronously. Hence, some other cores aren't finished processing yet when others finished. So `rank == 0` will wait until it returns to the first TPU core (and all other TPU stopped processing) before whatever you want to do (here saving model) is done. If you didn't do this, you might raise exceptions (try it yourself and see what happens!)\n\n**Finally, let's call the function and train!**\n\nFor the function below, the first argument is always `_mp_fn` (or other name depending on how you name your function), second argument is all the arguments you pass into `_mp_fn` (**except for `rank`**) in order. Third is the number of TPU to use. For single processing, pass in `nprocs=1`. For multiprocessing, pass in `nprocs=8`. Note, you can only pass in 1 or 8, not 2, not 4. Either only use 1 or all cores. Then, `start_method=\"fork\"` is always fixed because the only choice is `fork`. \n\n## Important NOTE:\nIf using `nprocs=8`, due to required warming up of TPU such as moving data to TPU, preparing for multiprocessing, and other reasons that one doesn't figure out, you will wait **quite long** before the spawn process start. **This means waiting for up to 1 hour before the process start**. To check whether your program works fine, try setting `nprocs=1` during experimental phase. \n\nEven if setting `nprocs=1`, the process requires some time to warm up. Typically, this means the first epoch is used for warming up. So, for our case, we have 465 total batches here, and processing the first 2-5 batches will take much much longer (as in up to 100x longer) than upcoming batches. Just wait and the tqdm meter would speed up at some point.\n\nAnd concurrent waiting (where the bar not moving at certain points after the first few batches) is due to preprocessing images on CPU bottleneck. Consider *caching your dataset* to smoothen this process. This could speed up about 3-10x training (although CPU will still be a bottleneck as it is fully utilized to fetch and prefetch data from disk, which is still slower than TPU training). \n\nThis is why it isn't guaranteed using multiple TPUs will speed up training: because CPU is the bottleneck in both case, so using single TPU and multiple TPU doesn't make any big differences until you upgrade your CPU. **Caching dataset does makes a difference though**. With cached dataset, CPU is only 50% utilized to fetch and prefetch data when using a single TPU (in ones' case); while fully utilized when fetching for all 8 cores. This difference reduces training time by half when using all TPUs. If this is not the case, it doesn't make any difference (as in Colab you only have 2 vCPUs so it is always fully utilized irregardless of you using single or multiple TPU cores). ","metadata":{}},{"cell_type":"code","source":"# xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=FLAGS[\"num_cores\"], start_method=\"fork\")\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=1, start_method=\"fork\")","metadata":{"execution":{"iopub.status.busy":"2021-08-21T03:56:37.009145Z","iopub.execute_input":"2021-08-21T03:56:37.009565Z","iopub.status.idle":"2021-08-21T04:03:25.824236Z","shell.execute_reply.started":"2021-08-21T03:56:37.009529Z","shell.execute_reply":"2021-08-21T04:03:25.821642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}