{"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":"
\n

Pytorch TPU Starter Code

\n
","metadata":{}},{"cell_type":"markdown","source":"This notebook is heavily based on the work of these two awesome notebooks:\n\n[[Pytorch] ArcFace + GeM Pooling Starter](https://www.kaggle.com/debarshichanda/pytorch-arcface-gem-pooling-starter)\n\n[The Ultimate PyTorch+TPU Tutorial](https://www.kaggle.com/tanlikesmath/the-ultimate-pytorch-tpu-tutorial-jigsaw-xlm-r#The-Ultimate-PyTorch-XLA/TPU-Tutorial-(Jigsaw-XLM-R)%F0%9F%94%A5)","metadata":{}},{"cell_type":"markdown","source":"![](https://media.istockphoto.com/illustrations/the-whale-is-blowing-illustration-id164494826?k=20&m=164494826&s=612x612&w=0&h=SGm8bwFqE7-h_ekqaXOVfIUIpKN8aW2AAMcFSbvpwYg=)","metadata":{}},{"cell_type":"markdown","source":"#

Install Required Libraries

","metadata":{}},{"cell_type":"code","source":"!pip install timm\n# !pip install --upgrade wandb","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-02-13T07:45:29.757132Z","iopub.execute_input":"2022-02-13T07:45:29.757661Z","iopub.status.idle":"2022-02-13T07:45:39.815443Z","shell.execute_reply.started":"2022-02-13T07:45:29.75756Z","shell.execute_reply":"2022-02-13T07:45:39.81437Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !pip install cloud-tpu-client==0.10 torch==1.9.0 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:45:39.817978Z","iopub.execute_input":"2022-02-13T07:45:39.818227Z","iopub.status.idle":"2022-02-13T07:45:39.822842Z","shell.execute_reply.started":"2022-02-13T07:45:39.818199Z","shell.execute_reply":"2022-02-13T07:45:39.821478Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!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 --apt-packages libomp5 libopenblas-dev","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:45:39.823728Z","iopub.execute_input":"2022-02-13T07:45:39.823979Z","iopub.status.idle":"2022-02-13T07:46:43.351929Z","shell.execute_reply.started":"2022-02-13T07:45:39.823934Z","shell.execute_reply":"2022-02-13T07:46:43.350291Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Import Required Libraries 📚

","metadata":{}},{"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport math\nimport copy\nimport time\nimport random\n\n# For data manipulation\nimport numpy as np\nimport pandas as pd\n\n# Pytorch Imports\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.optim import lr_scheduler\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda import amp\n\n# Utils\nimport joblib\nfrom tqdm import tqdm\nfrom collections import defaultdict\n\n# Sklearn Imports\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.model_selection import StratifiedKFold\n\n# For Image Models\nimport timm\n\n# Albumentations for augmentations\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n# For colored terminal text\nfrom colorama import Fore, Back, Style\nb_ = Fore.BLUE\nsr_ = Style.RESET_ALL\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nos.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:43.35626Z","iopub.execute_input":"2022-02-13T07:46:43.356686Z","iopub.status.idle":"2022-02-13T07:46:47.478676Z","shell.execute_reply.started":"2022-02-13T07:46:43.356645Z","shell.execute_reply":"2022-02-13T07:46:47.477726Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**PyTorch XLA-specific modules**","metadata":{}},{"cell_type":"code","source":"import 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\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:47.482055Z","iopub.execute_input":"2022-02-13T07:46:47.48262Z","iopub.status.idle":"2022-02-13T07:46:47.7682Z","shell.execute_reply.started":"2022-02-13T07:46:47.482572Z","shell.execute_reply":"2022-02-13T07:46:47.767191Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Training Configuration ⚙️

","metadata":{}},{"cell_type":"code","source":"# from pprint import pprint\n# model_names = timm.list_models(pretrained=True)\n# pprint(model_names)","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:47.770675Z","iopub.execute_input":"2022-02-13T07:46:47.770929Z","iopub.status.idle":"2022-02-13T07:46:47.773743Z","shell.execute_reply.started":"2022-02-13T07:46:47.770901Z","shell.execute_reply":"2022-02-13T07:46:47.773233Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CONFIG = {\"seed\": 2022,\n \"epochs\": 2,\n \"img_size\": 448,\n \"model_name\": \"mobilenetv2_050\",\n \"num_classes\": 15587,\n \"train_batch_size\": 16,\n \"valid_batch_size\": 64,\n \"learning_rate\": 1e-4,\n \"scheduler\": 'CosineAnnealingLR',\n \"min_lr\": 1e-6,\n \"T_max\": 500,\n \"weight_decay\": 1e-6,\n \"n_fold\": 5,\n \"n_accumulate\": 1,\n# \"device\": xm.xla_device(),\n # ArcFace Hyperparameters\n \"s\": 30.0, \n \"m\": 0.50,\n \"ls_eps\": 0.0,\n \"easy_margin\": False\n }","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:47.774797Z","iopub.execute_input":"2022-02-13T07:46:47.775113Z","iopub.status.idle":"2022-02-13T07:46:47.823987Z","shell.execute_reply.started":"2022-02-13T07:46:47.775087Z","shell.execute_reply":"2022-02-13T07:46:47.822987Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.set_default_tensor_type('torch.FloatTensor')","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:47.82567Z","iopub.execute_input":"2022-02-13T07:46:47.826054Z","iopub.status.idle":"2022-02-13T07:46:47.8531Z","shell.execute_reply.started":"2022-02-13T07:46:47.826008Z","shell.execute_reply":"2022-02-13T07:46:47.85219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Set Seed for Reproducibility

","metadata":{}},{"cell_type":"code","source":"def set_seed(seed=42):\n '''Sets the seed of the entire notebook so results are the same every time we run.\n This is for REPRODUCIBILITY.'''\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n # When running on the CuDNN backend, two further options must be set\n torch.backends.cudnn.deterministic = True\n torch.backends.cudnn.benchmark = False\n # Set a fixed value for the hash seed\n os.environ['PYTHONHASHSEED'] = str(seed)\n \nset_seed(CONFIG['seed'])","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:47.85465Z","iopub.execute_input":"2022-02-13T07:46:47.855116Z","iopub.status.idle":"2022-02-13T07:46:47.866251Z","shell.execute_reply.started":"2022-02-13T07:46:47.855068Z","shell.execute_reply":"2022-02-13T07:46:47.865465Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ROOT_DIR = '../input/happy-whale-and-dolphin'\nTRAIN_DIR = '../input/happy-whale-and-dolphin/train_images'\nTEST_DIR = '../input/happy-whale-and-dolphin/test_images'","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:47.867548Z","iopub.execute_input":"2022-02-13T07:46:47.867795Z","iopub.status.idle":"2022-02-13T07:46:47.876713Z","shell.execute_reply.started":"2022-02-13T07:46:47.867767Z","shell.execute_reply":"2022-02-13T07:46:47.876067Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_train_file_path(id):\n return f\"{TRAIN_DIR}/{id}\"","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:47.878011Z","iopub.execute_input":"2022-02-13T07:46:47.880182Z","iopub.status.idle":"2022-02-13T07:46:47.888841Z","shell.execute_reply.started":"2022-02-13T07:46:47.880137Z","shell.execute_reply":"2022-02-13T07:46:47.887846Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Read the Data 📖

","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(f\"{ROOT_DIR}/train.csv\")\ndf['file_path'] = df['image'].apply(get_train_file_path)\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:47.890471Z","iopub.execute_input":"2022-02-13T07:46:47.890713Z","iopub.status.idle":"2022-02-13T07:46:48.089902Z","shell.execute_reply.started":"2022-02-13T07:46:47.890685Z","shell.execute_reply":"2022-02-13T07:46:48.089163Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"encoder = LabelEncoder()\ndf['individual_id'] = encoder.fit_transform(df['individual_id'])\n\nwith open(\"le.pkl\", \"wb\") as fp:\n joblib.dump(encoder, fp)","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.090918Z","iopub.execute_input":"2022-02-13T07:46:48.091206Z","iopub.status.idle":"2022-02-13T07:46:48.136964Z","shell.execute_reply.started":"2022-02-13T07:46:48.091176Z","shell.execute_reply":"2022-02-13T07:46:48.136247Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Create Folds

","metadata":{}},{"cell_type":"code","source":"skf = StratifiedKFold(n_splits=CONFIG['n_fold'])\n\nfor fold, ( _, val_) in enumerate(skf.split(X=df, y=df.individual_id)):\n df.loc[val_ , \"kfold\"] = fold","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.139909Z","iopub.execute_input":"2022-02-13T07:46:48.140155Z","iopub.status.idle":"2022-02-13T07:46:48.840445Z","shell.execute_reply.started":"2022-02-13T07:46:48.140128Z","shell.execute_reply":"2022-02-13T07:46:48.8395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Dataset Class

","metadata":{}},{"cell_type":"code","source":"class HappyWhaleDataset(Dataset):\n def __init__(self, df, transforms=None):\n self.df = df\n self.file_names = df['file_path'].values\n self.labels = df['individual_id'].values\n self.transforms = transforms\n \n def __len__(self):\n return len(self.df)\n \n def __getitem__(self, index):\n img_path = self.file_names[index]\n img = cv2.imread(img_path)\n img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n label = self.labels[index]\n \n if self.transforms:\n img = self.transforms(image=img)[\"image\"]\n \n return {\n 'image': img,\n 'label': torch.tensor(label, dtype=torch.long)\n }","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.841785Z","iopub.execute_input":"2022-02-13T07:46:48.842058Z","iopub.status.idle":"2022-02-13T07:46:48.851089Z","shell.execute_reply.started":"2022-02-13T07:46:48.842028Z","shell.execute_reply":"2022-02-13T07:46:48.850221Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Augmentations

","metadata":{}},{"cell_type":"code","source":"data_transforms = {\n \"train\": A.Compose([\n A.Resize(CONFIG['img_size'], CONFIG['img_size']),\n A.HorizontalFlip(p=0.5),\n A.VerticalFlip(p=0.5),\n A.Rotate(limit=30, p=0.5),\n A.Normalize(\n mean=[0.485, 0.456, 0.406], \n std=[0.229, 0.224, 0.225], \n max_pixel_value=255.0, \n p=1.0\n ),\n ToTensorV2()], p=1.),\n \n \"valid\": A.Compose([\n A.Resize(CONFIG['img_size'], CONFIG['img_size']),\n A.Normalize(\n mean=[0.485, 0.456, 0.406], \n std=[0.229, 0.224, 0.225], \n max_pixel_value=255.0, \n p=1.0\n ),\n ToTensorV2()], p=1.)\n}","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.852561Z","iopub.execute_input":"2022-02-13T07:46:48.853136Z","iopub.status.idle":"2022-02-13T07:46:48.868848Z","shell.execute_reply.started":"2022-02-13T07:46:48.853101Z","shell.execute_reply":"2022-02-13T07:46:48.868052Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

GeM Pooling

\n\nCode taken from GeM Pooling Explained\n\n![](https://i.imgur.com/thTgYWG.jpg)","metadata":{}},{"cell_type":"code","source":"class GeM(nn.Module):\n def __init__(self, p=3, eps=1e-6):\n super(GeM, self).__init__()\n self.p = nn.Parameter(torch.ones(1)*p)\n self.eps = eps\n\n def forward(self, x):\n return self.gem(x, p=self.p, eps=self.eps)\n \n def gem(self, x, p=3, eps=1e-6):\n return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1./p)\n \n def __repr__(self):\n return self.__class__.__name__ + \\\n '(' + 'p=' + '{:.4f}'.format(self.p.data.tolist()[0]) + \\\n ', ' + 'eps=' + str(self.eps) + ')'\n\n\n\n\n\n","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.870889Z","iopub.execute_input":"2022-02-13T07:46:48.871193Z","iopub.status.idle":"2022-02-13T07:46:48.882223Z","shell.execute_reply.started":"2022-02-13T07:46:48.871161Z","shell.execute_reply":"2022-02-13T07:46:48.881468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

ArcFace

\n\nCode taken from Landmark2019-1st-and-3rd-Place-Solution","metadata":{}},{"cell_type":"code","source":"class ArcMarginProduct(nn.Module):\n r\"\"\"Implement of large margin arc distance: :\n Args:\n in_features: size of each input sample\n out_features: size of each output sample\n s: norm of input feature\n m: margin\n cos(theta + m)\n \"\"\"\n def __init__(self, in_features, out_features, s=30.0, \n m=0.50, easy_margin=False, ls_eps=0.0):\n super(ArcMarginProduct, self).__init__()\n self.in_features = in_features\n self.out_features = out_features\n self.s = s\n self.m = m\n self.ls_eps = ls_eps # label smoothing\n self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))\n nn.init.xavier_uniform_(self.weight)\n\n self.easy_margin = easy_margin\n self.cos_m = math.cos(m)\n self.sin_m = math.sin(m)\n self.th = math.cos(math.pi - m)\n self.mm = math.sin(math.pi - m) * m\n\n def forward(self, input, label):\n # --------------------------- cos(theta) & phi(theta) ---------------------\n cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n sine = torch.sqrt(1.0 - torch.pow(cosine, 2))\n phi = cosine * self.cos_m - sine * self.sin_m\n if self.easy_margin:\n phi = torch.where(cosine > 0, phi, cosine)\n else:\n phi = torch.where(cosine > self.th, phi, cosine - self.mm)\n # --------------------------- convert label to one-hot ---------------------\n # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda')\n one_hot = torch.zeros(cosine.size(), device=device)\n one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n if self.ls_eps > 0:\n one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features\n # -------------torch.where(out_i = {x_i if condition_i else y_i) ------------\n output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n output *= self.s\n\n return output","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.883667Z","iopub.execute_input":"2022-02-13T07:46:48.883919Z","iopub.status.idle":"2022-02-13T07:46:48.898834Z","shell.execute_reply.started":"2022-02-13T07:46:48.883891Z","shell.execute_reply":"2022-02-13T07:46:48.897866Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Create Model

","metadata":{}},{"cell_type":"code","source":"class HappyWhaleModel(nn.Module):\n def __init__(self, model_name, pretrained=True):\n super(HappyWhaleModel, self).__init__()\n self.model = timm.create_model(model_name, pretrained=pretrained)\n in_features = self.model.classifier.in_features\n self.model.classifier = nn.Identity()\n self.model.global_pool = nn.Identity()\n self.pooling = GeM()\n self.fc = ArcMarginProduct(in_features, \n CONFIG[\"num_classes\"],\n s=CONFIG[\"s\"], \n m=CONFIG[\"m\"], \n easy_margin=CONFIG[\"ls_eps\"], \n ls_eps=CONFIG[\"ls_eps\"])\n\n def forward(self, images, labels):\n features = self.model(images)\n pooled_features = self.pooling(features).flatten(1)\n output = self.fc(pooled_features, labels)\n return output\n \n","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.900126Z","iopub.execute_input":"2022-02-13T07:46:48.900388Z","iopub.status.idle":"2022-02-13T07:46:48.915784Z","shell.execute_reply.started":"2022-02-13T07:46:48.900351Z","shell.execute_reply":"2022-02-13T07:46:48.91484Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Loss Function

","metadata":{}},{"cell_type":"code","source":"def criterion(outputs, labels):\n return nn.CrossEntropyLoss()(outputs, labels)","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.916986Z","iopub.execute_input":"2022-02-13T07:46:48.917217Z","iopub.status.idle":"2022-02-13T07:46:48.932125Z","shell.execute_reply.started":"2022-02-13T07:46:48.917189Z","shell.execute_reply":"2022-02-13T07:46:48.931181Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Training Function

","metadata":{}},{"cell_type":"markdown","source":"We finally define our main functions that will be spawned by PyTorch XLA multiprocessing. These functions will be run on each of the 8 cores. There are several steps for 8-core training:\n1. We need to use a `DistributedSampler` that will appropriately distribute the dataset across the 8 cores.\n2. We are using `num_workers=0` as that decreases memory usage (only master process loading data). On higher memory VMs, this could be increased to speed up the training.\n3. The learning rate is scaled by the number of TPU cores (`xm.xrt_world_size()`)\n4. We put the model onto the TPU\n5. We use `ParallelLoader` which is a PyTorch XLA-specific DataLoader for loading data onto the TPU.\n6. We save the model at the end of training with `xm.model_save`","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"def train_one_epoch(model, optimizer, scheduler, dataloader, device, epoch):\n model.train()\n \n dataset_size = 0\n running_loss = 0.0\n \n bar = tqdm(enumerate(dataloader), total=len(dataloader))\n for step, data in bar:\n images = data['image'].to(device, dtype=torch.float)\n labels = data['label'].to(device, dtype=torch.long)\n \n batch_size = images.size(0)\n \n outputs = model(images, labels)\n loss = criterion(outputs, labels)\n\n loss = loss / CONFIG['n_accumulate']\n \n loss.backward()\n \n if (step + 1) % CONFIG['n_accumulate'] == 0:\n xm.optimizer_step(optimizer)\n\n # zero the parameter gradients\n optimizer.zero_grad()\n\n if scheduler is not None:\n scheduler.step()\n \n running_loss += (loss.item() * batch_size)\n dataset_size += batch_size\n \n epoch_loss = running_loss / dataset_size\n \n bar.set_postfix(Epoch=epoch, Train_Loss=epoch_loss,\n LR=optimizer.param_groups[0]['lr'])\n gc.collect()\n \n return epoch_loss","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.948082Z","iopub.execute_input":"2022-02-13T07:46:48.948325Z","iopub.status.idle":"2022-02-13T07:46:48.959558Z","shell.execute_reply.started":"2022-02-13T07:46:48.948297Z","shell.execute_reply":"2022-02-13T07:46:48.958937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Validation Function

","metadata":{}},{"cell_type":"code","source":"# @torch.inference_mode()\ndef valid_one_epoch(model, dataloader, device, epoch):\n model.eval()\n \n dataset_size = 0\n running_loss = 0.0\n \n bar = tqdm(enumerate(dataloader), total=len(dataloader))\n for step, data in bar: \n images = data['image'].to(device, dtype=torch.float)\n labels = data['label'].to(device, dtype=torch.long)\n \n batch_size = images.size(0)\n\n outputs = model(images, labels)\n loss = criterion(outputs, labels)\n \n running_loss += (loss.item() * batch_size)\n dataset_size += batch_size\n \n epoch_loss = running_loss / dataset_size\n \n bar.set_postfix(Epoch=epoch, Valid_Loss=epoch_loss,\n LR=optimizer.param_groups[0]['lr']) \n \n gc.collect()\n \n return epoch_loss","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.960493Z","iopub.execute_input":"2022-02-13T07:46:48.960838Z","iopub.status.idle":"2022-02-13T07:46:48.976557Z","shell.execute_reply.started":"2022-02-13T07:46:48.960808Z","shell.execute_reply":"2022-02-13T07:46:48.975892Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

Run Training

","metadata":{}},{"cell_type":"code","source":"# device = CONFIG[\"device\"]\nnum_epochs = CONFIG[\"epochs\"]","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.977663Z","iopub.execute_input":"2022-02-13T07:46:48.977999Z","iopub.status.idle":"2022-02-13T07:46:48.987892Z","shell.execute_reply.started":"2022-02-13T07:46:48.977948Z","shell.execute_reply":"2022-02-13T07:46:48.987079Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_training(model, optimizer, scheduler, device, num_epochs):\n # To automatically log gradients\n# wandb.watch(model, log_freq=100)\n \n# if torch.cuda.is_available():\n# print(\"[INFO] Using GPU: {}\\n\".format(torch.cuda.get_device_name()))\n \n start = time.time()\n best_model_wts = copy.deepcopy(model.state_dict())\n best_epoch_loss = np.inf\n history = defaultdict(list)\n \n \n for epoch in range(1, num_epochs + 1): \n gc.collect()\n # We use ParallelLoader (provided by PyTorch XLA) for TPU-core-specific dataloading:\n para_loader = pl.ParallelLoader(train_loader, [device]) \n xm.master_print('parallel loader created... training now')\n gc.collect()\n #call the training loop\n train_epoch_loss = train_one_epoch(model, optimizer, scheduler, \n dataloader = para_loader.per_device_loader(device), \n device=device, epoch=epoch)\n del para_loader\n para_loader = pl.ParallelLoader(valid_loader, [device])\n gc.collect()\n # call evaluation loop\n val_epoch_loss = valid_one_epoch(model, para_loader.per_device_loader(device), device=device, \n epoch=epoch)\n del para_loader\n gc.collect()\n \n history['Train Loss'].append(train_epoch_loss)\n history['Valid Loss'].append(val_epoch_loss)\n \n # Log the metrics\n# wandb.log({\"Train Loss\": train_epoch_loss})\n# wandb.log({\"Valid Loss\": val_epoch_loss})\n \n # deep copy the model\n if val_epoch_loss <= best_epoch_loss:\n print(f\"{b_}Validation Loss Improved ({best_epoch_loss} ---> {val_epoch_loss})\")\n best_epoch_loss = val_epoch_loss\n run.summary[\"Best Loss\"] = best_epoch_loss\n best_model_wts = copy.deepcopy(model.state_dict())\n PATH = \"Loss{:.4f}_epoch{:.0f}.bin\".format(best_epoch_loss, epoch)\n# torch.save(model.state_dict(), PATH)\n xm.save(model.state_dict(), PATH)\n # Save a model file from the current directory\n print(f\"Model Saved{sr_}\")\n \n print()\n \n end = time.time()\n time_elapsed = end - start\n print('Training complete in {:.0f}h {:.0f}m {:.0f}s'.format(\n time_elapsed // 3600, (time_elapsed % 3600) // 60, (time_elapsed % 3600) % 60))\n print(\"Best Loss: {:.4f}\".format(best_epoch_loss))\n \n # load best model weights\n# model.load_state_dict(best_model_wts)\n \n# return model, history","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:48.989186Z","iopub.execute_input":"2022-02-13T07:46:48.989571Z","iopub.status.idle":"2022-02-13T07:46:49.005786Z","shell.execute_reply.started":"2022-02-13T07:46:48.989542Z","shell.execute_reply":"2022-02-13T07:46:49.004947Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def fetch_scheduler(optimizer):\n if CONFIG['scheduler'] == 'CosineAnnealingLR':\n scheduler = lr_scheduler.CosineAnnealingLR(optimizer,T_max=CONFIG['T_max'], \n eta_min=CONFIG['min_lr'])\n elif CONFIG['scheduler'] == 'CosineAnnealingWarmRestarts':\n scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer,T_0=CONFIG['T_0'], \n eta_min=CONFIG['min_lr'])\n elif CONFIG['scheduler'] == None:\n return None\n \n return scheduler","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:49.007798Z","iopub.execute_input":"2022-02-13T07:46:49.008507Z","iopub.status.idle":"2022-02-13T07:46:49.02121Z","shell.execute_reply.started":"2022-02-13T07:46:49.00847Z","shell.execute_reply":"2022-02-13T07:46:49.02036Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def prepare_loaders(df, fold):\n df_train = df[df.kfold != fold].reset_index(drop=True)\n df_valid = df[df.kfold == fold].reset_index(drop=True)\n \n train_dataset = HappyWhaleDataset(df_train, transforms=data_transforms[\"train\"])\n valid_dataset = HappyWhaleDataset(df_valid, transforms=data_transforms[\"valid\"])\n \n train_sampler = torch.utils.data.distributed.DistributedSampler(\n train_dataset,\n num_replicas=xm.xrt_world_size(), # tell PyTorch how many devices (TPU cores) we are using for training\n rank=xm.get_ordinal(), # tell PyTorch which device (core) we are on currently\n shuffle=True)\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 train_loader = DataLoader(train_dataset, batch_size=CONFIG['train_batch_size'], sampler=train_sampler,\n num_workers=0, pin_memory=True, drop_last=True)\n \n valid_loader = DataLoader(valid_dataset, batch_size=CONFIG['valid_batch_size'], sampler=valid_sampler, \n num_workers=0, pin_memory=True)\n \n return train_loader, valid_loader","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:49.02234Z","iopub.execute_input":"2022-02-13T07:46:49.022577Z","iopub.status.idle":"2022-02-13T07:46:49.034087Z","shell.execute_reply.started":"2022-02-13T07:46:49.02255Z","shell.execute_reply":"2022-02-13T07:46:49.03326Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Prepare Dataloaders","metadata":{}},{"cell_type":"code","source":"train_loader, valid_loader = prepare_loaders(df, fold=0)","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:49.035255Z","iopub.execute_input":"2022-02-13T07:46:49.035476Z","iopub.status.idle":"2022-02-13T07:46:49.068475Z","shell.execute_reply.started":"2022-02-13T07:46:49.035451Z","shell.execute_reply":"2022-02-13T07:46:49.067499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Start Training","metadata":{}},{"cell_type":"markdown","source":"[TPU multi-processing extremely slow in comparison to GPUs](https://github.com/pytorch/xla/issues/2383)\n\n[PyTorch XLA is very slow on Google Colab](https://github.com/pytorch/xla/issues/2247)","metadata":{}},{"cell_type":"markdown","source":"![Upvote!](https://img.shields.io/badge/Upvote-If%20you%20like%20my%20work-07b3c8?style=for-the-badge&logo=kaggle)","metadata":{}},{"cell_type":"markdown","source":"***Make sure not to execute/call any of your PyTorch XLA specific function outised of the below fucntion call. Doing so will save you the head ache of unexpected errors*.**","metadata":{}},{"cell_type":"code","source":"import time\n\n# Start training processes\ndef _mp_fn(rank, flags):\n global device, optimizer,scheduler\n torch.set_default_tensor_type('torch.FloatTensor')\n \n device = xm.xla_device()\n num_epochs = CONFIG['epochs']\n model = HappyWhaleModel(CONFIG['model_name'])\n# SERIAL_EXEC = xmp.MpSerialExecutor()\n # Only instantiate model weights once in memory.\n# WRAPPED_MODEL = xmp.MpModelWrapper(HappyWhaleModel(CONFIG['model_name']))\n model = model.to(device)\n xm.master_print('done loading model')\n optimizer = optim.Adam(model.parameters(), lr=CONFIG['learning_rate'] * xm.xrt_world_size(), \n weight_decay=CONFIG['weight_decay'])\n scheduler = fetch_scheduler(optimizer)\n a = run_training(model = model, optimizer = optimizer, scheduler = scheduler, device = device, num_epochs = num_epochs)\n\nFLAGS={}\nstart_time = time.time()\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","metadata":{"execution":{"iopub.status.busy":"2022-02-13T07:46:49.07076Z","iopub.execute_input":"2022-02-13T07:46:49.071054Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}