{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Summary\nThis notebook can be used for training, validation, and testing of multiple PyTorch models on anything from CPU to 8 TPU cores. It makes use of some useful additional features like\n\n* class balancing for highly undersampled classes\n* warmstarting from models with fewer to models with more class\n* model and stage switching depending on needs\n\nSupported models are ResNets and EfficientNets for specific choices of layers but feel free to extend this to your needs if you like.\n\nMajor parts of this notebook were inspired by\n\n* https://www.kaggle.com/andypenrose/pytorch-training-inference-efficientnet-b4\n* https://www.kaggle.com/rhtsingh/pytorch-training-inference-efficientnet-baseline/notebook\n\nlike the dataset class and the submission handling."},{"metadata":{},"cell_type":"markdown","source":"# Imports and dependencies\n\nWe start with the necessary imports and dependencies, among which general python imports, special libary imports, and PyTorch imports. The first two are"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport glob\nimport gc\ngc.enable()\nimport time\nimport random\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nimport multiprocessing\n\nimport cv2\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport albumentations as A\nfrom PIL import Image\nfrom sklearn.preprocessing import LabelEncoder","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Before other imports however we have to define the main parameters for the run (to remain constants throughout) which will influcence further imports:"},{"metadata":{"trusted":true},"cell_type":"code","source":"# one of: resnet, efficientnet\nMODEL = \"resnet\"\n# what stages should be performed (test here is simply generating the submission for the test sets)\nKERNEL_STAGES = [\"train\", \"validate\", \"test\"]\n# minimum required samples per class, could be 0 for all classes (will be balanced during sampling)\nMIN_SAMPLES_PER_CLASS = 0\n# split between training and validation, typically 20% or 10% of the training is used for validation\nVALIDATION_SPLIT = .1\n# recommentded: 128/256 for one TPU core, 16/32 for 8 TPU cores, 128/256 for one GPU\nBATCH_SIZE = 256\nNUM_EPOCHS = 1\n# number of CPU cores and TPU cores if the TPU accelerator is turned on\nCPU_CORES = multiprocessing.cpu_count()\nTPU_CORES = 0\n# log loss and accuracy (and save model) on a given number of minibatches\nLOG_FREQ = 32\nSAVE_FREQ = 128\n# number of previous classes for warmstarting (0 means off)\nWARMSTARTING = 0\n# number of top predictions\nNUM_TOP_PREDICTS = 1\n# minimum confidence below which to output nothing (no class)\nMIN_CONF = 0.1\n# use this switch for proper submission\nENABLE_FAST_SKIP = False\n# image width and height in pixes (assumed squares)\nIMG_SIZE = 128","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can now continue and finish with all imports:"},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nSteps for 8x TPU acceleration to run on entire chip:\n  - set TPU_CORES > 0 (1 or 8)\n  - turn on TPU accelerator\n  - turn on internet switch\n  - exclude the test stage if included (not available for TPUs)\n  - load model from a dataset (which might need an update)\n  - balance CPU_CORES with TPU_CORES (memory trade-off between the two)\nTODO: consider xm.mesh_reduce and xm.master_print to print just one reduced loss or accuracy\nTODO: consider whether to scale learning rate as lr * xm.xrt_world_size()\n\"\"\"\nif TPU_CORES > 0:\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 --apt-packages libomp5 libopenblas-dev\n    # use the mantisa-reducing float type\n    !export XLA_USE_BF16=1\n    import torch_xla\n    import torch_xla.core.xla_model as xm\n    if TPU_CORES > 1:\n        import torch_xla.distributed.parallel_loader as pl\n        import torch_xla.distributed.xla_multiprocessing as xmp\n\nimport torch\nfrom torch import Tensor\nfrom torch import nn\nfrom torch.optim import lr_scheduler, Adam, SGD\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SubsetRandomSampler, WeightedRandomSampler\nfrom torchvision import transforms\nfrom torchvision import models\n\nfrom catalyst.data.sampler import DistributedSamplerWrapper\nfrom tqdm import tqdm\n# somewhat fancier\n#from tqdm.notebook import tqdm\n\n# uncomment for offline install (needs a minimal dataset)\n#!mkdir -p /tmp/pip/cache/\n#!cp ../input/glr-efficientnets-ext/efficientnet_pytorch-0.6.3-py3-none-any.whl /tmp/pip/cache/\n#!pip install --no-index --find-links /tmp/pip/cache/ efficientnet_pytorch\n!pip install efficientnet_pytorch\nimport efficientnet_pytorch","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now define all paths to be used for the datasets and models:"},{"metadata":{"trusted":true},"cell_type":"code","source":"pretrain_df = pd.read_csv('../input/landmark-recognition-2020/train.csv')\npretest_df = pd.read_csv('../input/landmark-recognition-2020/sample_submission.csv')\ntrain_dir = '../input/landmark-recognition-2020/train/'\ntest_dir = '../input/landmark-recognition-2020/test/'\nmodel_dir = \"../input/pytorch-pretrained-image-models/\"\nmodel_load_file = \"\"  # e.g. \"../input/models/glr_resnet_81313.pth\"\nmodel_save_file = \"glr_resnet.pth\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"While seeding is good for reproducibility, considering the constrained resources on Kaggle and the possibility of running just one epoch, I would rather recommend commenting out the `seed_randoms()` line in order to be able to properly sample all classes across multiple sessions."},{"metadata":{"trusted":true},"cell_type":"code","source":"# seed everything to avoid non-determinism\ndef seed_randoms(seed=2020):\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.deterministic = True\n# comment this out to not oversample some classes due to short deterministic sessions\nseed_randoms()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The dataset class is fairly standard in PyTorch:"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class ImageDataset(Dataset):\n    \"\"\"\n    Standard class sourced from:\n\n    https://www.kaggle.com/rhtsingh/pytorch-training-inference-efficientnet-baseline\n    \"\"\"\n\n    def __init__(self, dataframe: pd.DataFrame, image_dir:str, mode: str):\n        self.df = dataframe\n        self.mode = mode\n        self.image_dir = image_dir\n\n        transforms_list = []\n        if self.mode == 'train':\n            # Increase image size from (64,64) to higher resolution,\n            # Make sure to change in RandomResizedCrop as well.\n            transforms_list = [\n                transforms.Resize((IMG_SIZE,IMG_SIZE)),\n                transforms.RandomHorizontalFlip(),\n                transforms.RandomChoice([\n                    transforms.RandomResizedCrop(IMG_SIZE),\n                    transforms.ColorJitter(0.2, 0.2, 0.2, 0.2),\n                    transforms.RandomAffine(degrees=15, translate=(0.2, 0.2),\n                                            scale=(0.8, 1.2), shear=15,\n                                            resample=Image.BILINEAR)\n                ]),\n                transforms.ToTensor(),\n                transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                                      std=[0.229, 0.224, 0.225]),\n            ]\n        else:\n            transforms_list.extend([\n                # Keep this resize same as train\n                transforms.Resize((IMG_SIZE,IMG_SIZE)),\n                transforms.ToTensor(),\n                transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                                      std=[0.229, 0.224, 0.225]),\n            ])\n        self.transforms = transforms.Compose(transforms_list)\n\n    def __getitem__(self, index: int):\n        image_id = self.df.iloc[index].id\n        image_path = f\"{self.image_dir}/{image_id[0]}/{image_id[1]}/{image_id[2]}/{image_id}.jpg\"\n        image = Image.open(image_path)\n        image = self.transforms(image)\n\n        if self.mode == 'test':\n            return {'image':image}\n        else:\n            return {'image':image, \n                    'target':self.df.iloc[index].landmark_id}\n\n    def __len__(self) -> int:\n        return self.df.shape[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Filter out classes with too few samples (used mostly due to resource constraints):"},{"metadata":{"trusted":true},"cell_type":"code","source":"class_sample_totals = pretrain_df.landmark_id.value_counts()\nselected_classes = class_sample_totals[class_sample_totals >= MIN_SAMPLES_PER_CLASS]\nnum_classes = len(selected_classes)\ntrain_df = pretrain_df.loc[pretrain_df.landmark_id.isin(selected_classes.index)]\nprint(f'Classes with at least N={MIN_SAMPLES_PER_CLASS} samples: {num_classes}')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Reset the IDs used for the landmarks back to ones starting from zero and compatible with the network outputs."},{"metadata":{"trusted":true},"cell_type":"code","source":"# map to 0-reset class indices to also index the logits of the network\nlabel_encoder = LabelEncoder()\nlabel_encoder.fit(train_df.landmark_id.values)\nassert len(label_encoder.classes_) == num_classes\ntrain_df.landmark_id = label_encoder.transform(train_df.landmark_id)\ntrain_df.reset_index(drop=True, inplace=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Split the training samples and labels in order to be able to validate the model.\n\nWe will always reuse a fixed seed just for this particular splitting in order to make sure the validation set is never leaked into the training set depending on the order of random generator calls or other events. In this way, even if we use full randomness and shorter sessions the split between the training and validation set will remains the same and will only depend on the percentage for validation."},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_indices = train_df.index.to_numpy()\n# perform reprucible split without interrupting the random flow\nrandom_state = np.random.get_state()\nnp.random.seed(2020)\nnp.random.shuffle(sample_indices)\nnp.random.set_state(random_state)\nsplit = int(VALIDATION_SPLIT * len(sample_indices))\ntrain_indices, val_indices = sample_indices[split:], sample_indices[:split]\nassert len(train_indices) + len(val_indices) == len(sample_indices)\nprint('Train dataframe (at least N samples):', train_df.shape)\nprint('  for training:', len(train_indices))\nprint('  for validation:', len(val_indices))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"For the rest of the dataset preparation we validate the test set for existing files, initialize the datasets, and validation and training samplers. The validation sampler is a normal random subset sampler over the validation indices but the training sampler is more elaborate in order to balance the sampled classes.\n\nWe use a somewhat standard approach in PyTorch through a weighted random sampler, giving zero weights for the validation indices which should never be sampled within trianing. There is also the possibilty to split the dataset itself but relying on sampling is a cleaner approach."},{"metadata":{"trusted":true},"cell_type":"code","source":"# filter non-existing test images\nexists = lambda img: os.path.exists(f'{test_dir}/{img[0]}/{img[1]}/{img[2]}/{img}.jpg')\ntest_df = pretest_df.loc[pretest_df.id.apply(exists)]\nprint('Test dataframe (existing files):', test_df.shape)\n\ntrain_dataset = ImageDataset(train_df, train_dir, mode='train')\ntest_dataset = ImageDataset(test_df, test_dir, mode='test')\n\nval_subsampler = SubsetRandomSampler(val_indices)\n\n# balance the classes due to high polarity in the number of samples per class\nclass_sample_count = train_df.landmark_id.value_counts(sort=False).sort_index().values\nclass_weights = 1 / torch.Tensor(class_sample_count)\nsample_weights = class_weights[train_df.landmark_id.values]\n# cannot draw from the validation data\nsample_weights[val_indices] = 0.0\ntrain_subsampler = WeightedRandomSampler(sample_weights, len(sample_weights))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Models that need separate classes can be defined here:"},{"metadata":{"trusted":true},"cell_type":"code","source":"class EfficientNetEncoderHead(nn.Module):\n    \"\"\"Head for an EfficientNet encoder.\"\"\"\n\n    def __init__(self, depth, num_classes):\n        super(EfficientNetEncoderHead, self).__init__()\n        self.depth = depth\n        self.base = efficientnet_pytorch.EfficientNet.from_name(f'efficientnet-b{self.depth}')\n        #pretrained_file = glob.glob(f'../input/efficientnet-pytorch/efficientnet-b{self.depth}*')[0]\n        #self.base.load_state_dict(torch.load(pretrained_file, map_location=device))\n        # disable training of pre-trained layers\n        #for param in self.base.parameters():\n        #    param.requires_grad = False\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.classifier = nn.Linear(self.base._fc.in_features, num_classes)\n\n    def forward(self, x):\n        \"\"\"Forward method for the full model.\"\"\"\n        x = self.base.extract_features(x)\n        x = self.avg_pool(x).squeeze(-1).squeeze(-1)\n        x = self.classifier(x)\n        return x\n\n    def numel(self, only_trainable=False):\n        \"\"\"\n        Returns the total number of parameters used by the model\n        (only counting shared parameters once).\n\n        :param bool only_trainable: whether to only includes parameters\n                                    with ``requires_grad = True``\n        :returns: total number of parameters\n        :rtype: int\n        \"\"\"\n        parameters = self.parameters()\n        if only_trainable:\n            parameters = list(p for p in parameters if p.requires_grad)\n        unique = dict((p.data_ptr(), p) for p in parameters).values()\n        return sum(p.numel() for p in unique)\n\n    def summary(self):\n        \"\"\"Print a summary of the loaded model.\"\"\"\n        print()\n        print(f'model summary\\n----------------------------------------------')\n        for name, param in self.named_parameters():\n            print(f'module {name} of shape {list(param.shape)} in the range [{torch.min(param).item():.4f},'\n                f' {torch.max(param).item():.4f}]\\n----------------------------------------------')\n            #if \"fc1\" in name and \"bias\" in name:\n            #    param.register_hook(lambda grad: self.log.add_video(f\"grad/fc1\", grad))\n            #if \"conv1\" in name and \"bias\" in name:\n            #    param.register_hook(lambda grad: self.log.add_video(f\"grad/conv1\", grad))\n        print(\"Trainable parameters:\", self.numel(only_trainable=True))\n        print(\"Total parameters:\", self.numel())\n        print()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can now initialize the model of choice."},{"metadata":{"trusted":true},"cell_type":"code","source":"if MODEL == \"resnet\":\n    # load a pretrained model\n    net = models.resnet50(pretrained=False)\n    #net.load_state_dict(torch.load(os.path.join(model_dir, f'resnet{50}.pth'),\n    #                               map_location=device))\n    # disable training of pre-trained layers\n    #for param in net.parameters():\n    #    param.requires_grad = False\n    # replace the fully connected layer with one for our classes to actually be trained\n    net.fc = nn.Linear(net.fc.in_features, num_classes)\n\n    print()\n    def numel(self, only_trainable=False):\n        \"\"\"\n        Returns the total number of parameters used by the model\n        (only counting shared parameters once).\n\n        :param bool only_trainable: whether to only includes parameters\n                                    with ``requires_grad = True``\n        :returns: total number of parameters\n        :rtype: int\n        \"\"\"\n        parameters = self.parameters()\n        if only_trainable:\n            parameters = list(p for p in parameters if p.requires_grad)\n        unique = dict((p.data_ptr(), p) for p in parameters).values()\n        return sum(p.numel() for p in unique)\n    print(\"Trainable parameters:\", numel(net, only_trainable=True))\n    print(\"Total parameters:\", numel(net))\n    print()\n\nelif MODEL == \"efficientnet\":\n    net = EfficientNetEncoderHead(depth=4, num_classes=num_classes)\n    net.summary()\n\nelse:\n    raise ValueError(f\"Inappropriate chocie for model {MODEL}, only ResNets and EfficientNets supported\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can now load a previously saved model or warmstart it from an otherwise incompatible model.\n\nThe model is only loaded if there is a provided path and is always performed on CPU to avoid limitations from loading TPU models into CPU/GPU models and other accelerator switching. Once the model is loaded on CPU, it will eventually be sent to the device that is currently used.\n\nThe warmstarting matches all previously learned classes from a smaller model to the new wider class selection."},{"metadata":{"trusted":true},"cell_type":"code","source":"def warmstarting(model, filename, prev_min_delimiter=10):\n    \"\"\"Perform wormstarting of a new model from pretrained parameters of a different one.\"\"\"\n    print(\"Warmstarting the model\")\n    saved_state = torch.load(filename, map_location=\"cpu\")\n    net_state = net.state_dict()\n\n    prev_class_sample_totals = pretrain_df.landmark_id.value_counts()\n    prev_selected_classes = prev_class_sample_totals[prev_class_sample_totals >= prev_min_delimiter]\n    prev_num_classes = len(prev_selected_classes)\n    prev_train_df = pretrain_df.loc[pretrain_df.landmark_id.isin(prev_selected_classes.index)]\n    print(f'Previous classes with at least N={prev_min_delimiter} samples: {prev_num_classes}')\n\n    prev_label_encoder = LabelEncoder()\n    prev_label_encoder.fit(prev_train_df.landmark_id.values)\n    assert len(prev_label_encoder.classes_) == prev_num_classes\n\n    for key in saved_state:\n        if key in net_state:\n            if net_state[key].shape == saved_state[key].shape:\n                print(f\"Restoring {key} with matching shape {saved_state[key].shape}\")\n                net_state[key] = saved_state[key]\n                continue\n\n            print(f\"Class matching {key} from {saved_state[key].shape} to {net_state[key].shape}\")\n            # assert that the new model is larger than the old model\n            assert len(net_state[key].shape) == len(saved_state[key].shape) <= 2\n            assert net_state[key].shape[0] >= saved_state[key].shape[0]\n\n            saved_idx = np.arange(saved_state[key].shape[0])\n            net_idx = label_encoder.transform(prev_label_encoder.inverse_transform(saved_idx))\n            for i, j in zip(net_idx, saved_idx):\n                net_state[key][i] = saved_state[key][j]\n\nif model_load_file != \"\":\n    if not os.path.exists(model_load_file):\n        print(\"Warning: Previous model file was not found, creating a new one\")\n    else:\n        if WARMSTARTING > 0:\n            warmstarting(net, model_load_file, WARMSTARTING)\n        else:\n            net.load_state_dict(torch.load(model_load_file, map_location=\"cpu\"))\nnet.eval()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"What follows is some helper functions for average meters and GAP evaluating function."},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter:\n    ''' Computes and stores the average and current value '''\n    def __init__(self) -> None:\n        self.reset()\n\n    def reset(self) -> None:\n        self.val = 0.0\n        self.avg = 0.0\n        self.sum = 0.0\n        self.count = 0\n\n    def update(self, val: float, n: int = 1) -> None:\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\ndef GAP(y_true, y_pred):\n    \"\"\"Compute Global Average Precision score (GAP)\"\"\"\n    indexes = list(y_pred.keys())\n    indexes.sort(\n        key=lambda x: -y_pred[x][1],\n    )\n    queries_with_target = len([i for i in y_true.values() if i is not None])\n    correct_predictions = 0\n    total_score = 0.\n    for i, k in enumerate(indexes, 1):\n        relevance_of_prediction_i = 0\n        if y_true[k] == y_pred[k][0]:\n            correct_predictions += 1\n            relevance_of_prediction_i = 1\n        precision_at_rank_i = correct_predictions / i\n        total_score += precision_at_rank_i * relevance_of_prediction_i\n    return 1 / queries_with_target * total_score","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We are now ready to defien the training step which also has some conditionals for the not yet fully PyTorch integrated TPU accelerators."},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_step(train_loader, model, device, criterion, optimizer, epoch, lr_scheduler, filename):\n    print(f'Epoch {epoch} with total batches: {len(train_loader)}')\n    batch_time = AverageMeter()\n    losses = AverageMeter()\n    avg_score = AverageMeter()\n\n    model.train()\n\n    end = time.time()\n    lr = None\n\n    for i, data in enumerate(train_loader):\n        x = data['image']\n        y = data['target']\n\n        y_ = model(x.to(device))\n        loss = criterion(y_, y.to(device))\n\n        confs, preds = torch.max(y_.detach(), dim=1)\n        tuple_pred = list(zip(preds.cpu().numpy(), confs.cpu().numpy()))\n        true_labels = y.cpu().numpy()\n        y_true, y_pred = {}, {}\n        for j in range(len(tuple_pred)):\n            y_true[f'{j}'] = true_labels[j]\n            y_pred[f'{j}'] = tuple_pred[j]\n        avg_score.update(GAP(y_true, y_pred))\n        losses.update(loss.data.item(), x.size(0))\n\n        optimizer.zero_grad()\n        loss.backward()\n        if TPU_CORES > 0:\n            xm.optimizer_step(optimizer, barrier=True)\n        else:\n            optimizer.step()\n        lr_scheduler.step()\n        lr = optimizer.param_groups[0]['lr']\n\n        batch_time.update(time.time() - end)\n        end = time.time()\n\n        if i % LOG_FREQ == 0:\n            print(f'{device}/{epoch} [{i}/{len(train_loader)}]\\t'\n                    f'time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n                    f'loss {losses.val:.4f} ({losses.avg:.4f})\\t'\n                    f'GAP {avg_score.val:.4f} ({avg_score.avg:.4f})\\t'\n                    f'-> rate: {lr}')\n        if i % SAVE_FREQ == 0:\n            if filename != \"\":\n                if TPU_CORES > 0:\n                    xm.save(model.state_dict(), filename)\n                else:\n                    torch.save(model.state_dict(), filename)\n                torch.save(optimizer.state_dict(), filename + \"_opt.tar\")\n\n    print(f'Average GAP on train: {avg_score.avg:.4f}')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Next are the evaluation step and submission generation where the evaluation step is a bit more compact using `tqdm` updates. The training step in comparison keeps each line in order to trace the changes in the loss and GAP accuracy."},{"metadata":{"trusted":true},"cell_type":"code","source":"def eval_step(data_loader, model, device, mode=\"validation\"):\n    avg_score = AverageMeter()\n    model.eval()\n\n    softmax = nn.Softmax(dim=1)\n    all_preds, all_confs, all_targets = [], [], []\n\n    with torch.no_grad():\n        with tqdm(total=len(data_loader), ncols=100,\n                  bar_format='{l_bar}{bar}| {postfix[0]} {n_fmt}/{total_fmt} ({postfix[3]}) '\\\n                             '[{elapsed}<{remaining}, {rate_fmt}, {postfix[1]}:{postfix[2][GAP]:>.4}]',\n                  postfix=[\"Batch\", \"GAP\", dict(GAP=\"?\"), device]) as t:\n            for i, data in enumerate(data_loader):\n                # we would check `data_loader.dataset.mode` but the validation set is\n                # a subset of the training set, thus having training mode\n                if mode != 'test':\n                    x, y = data['image'], data['target']\n                else:\n                    x, y = data['image'], None\n\n                y_ = model(x.to(device))\n                y_ = softmax(y_)\n\n                confs, preds = torch.topk(y_, NUM_TOP_PREDICTS)\n\n                if y is not None:\n                    tuple_pred = list(zip(preds.cpu().numpy(), confs.cpu().numpy()))\n                    true_labels = y.cpu().numpy()\n                    y_true, y_pred = {}, {}\n                    for j in range(len(tuple_pred)):\n                        y_true[f'{j}'] = true_labels[j]\n                        y_pred[f'{j}'] = tuple_pred[j]\n                    avg_score.update(GAP(y_true, y_pred))\n                    t.postfix[2][\"GAP\"] = avg_score.avg\n\n                all_confs.append(confs)\n                all_preds.append(preds)\n                if y is not None:\n                    all_targets.append(y)\n\n                t.update()\n\n        if mode != 'test':\n            print(f'Average GAP: {avg_score.avg:.4f}')\n\n    preds = torch.cat(all_preds)\n    confs = torch.cat(all_confs)\n    targets = torch.cat(all_targets) if len(all_targets) else None\n\n    return preds, confs, targets\n\ndef generate_submission(test_loader, model, device, label_encoder):\n    sample_sub = pd.read_csv('../input/landmark-recognition-2020/sample_submission.csv')\n\n    predicts_gpu, confs_gpu, _ = eval_step(test_loader, model, device, mode=\"test\")\n    predicts, confs = predicts_gpu.cpu().numpy(), confs_gpu.cpu().numpy()\n\n    labels = [label_encoder.inverse_transform(pred) for pred in predicts]\n\n    sub = test_loader.dataset.df\n    def concat(label: np.ndarray, conf: np.ndarray) -> str:\n        return ' '.join([f'{L} {c}' if c > MIN_CONF else '' for L, c in zip(label, conf)])\n    sub['landmarks'] = [concat(label, conf) for label, conf in zip(labels, confs)]\n\n    sample_sub = sample_sub.set_index('id')\n    sub = sub.set_index('id')\n    sample_sub.update(sub)\n    sample_sub.to_csv('submission.csv')\n\n    print(sample_sub.head())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The final step to combine all previous ones is to train, test, and validate on a single (if CPU, GPU, or one TPU core is used) or multiple replicas (in the case of all 8 TPU cores). While the number of TPU cores should still be > 0 to use a TPU and only 1 or 8 cores are currently supported.\n\nThe rest of this takes care to instantiate the correct data loaders (possibly with distributed samplers on top of the previous samplers), reset or reload a previous optimizer (e.g. from previous training session), and run the selected stages possibly excluding training, validation, or testing."},{"metadata":{"trusted":true},"cell_type":"code","source":"def map_run(index):\n    device = xm.xla_device() if TPU_CORES > 0 else torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    if ENABLE_FAST_SKIP and test_df.id[0] == \"00084cdf8f600d00\":\n        # This is a run on the public data, skip it to speed up submission run on private data.\n        print(\"Skipping run on public test set.\")\n        sample_sub = pd.read_csv('../input/landmark-recognition-2020/sample_submission.csv')\n        sample_sub.to_csv('submission.csv')\n        return\n\n    if TPU_CORES > 1:\n        # TODO: we cannot subsample in an elegant way until PyTorch adds a wrapper with a subsampler\n        #train_sampler = DistributedSampler(train_dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal(), shuffle=True)\n        #val_sampler = DistributedSampler(train_dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())\n        train_sampler = DistributedSamplerWrapper(train_subsampler, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal(), shuffle=True)\n        val_sampler = DistributedSamplerWrapper(val_subsampler, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())\n\n        world_train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE,\n                                        sampler=train_sampler, num_workers=CPU_CORES, drop_last=True)\n        world_val_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE,\n                                      sampler=val_sampler, num_workers=CPU_CORES, drop_last=True)\n\n        para_train_loader = pl.ParallelLoader(world_train_loader, [device])\n        para_val_loader = pl.ParallelLoader(world_val_loader, [device])\n\n        train_loader = para_train_loader.per_device_loader(device)\n        val_loader = para_val_loader.per_device_loader(device)\n    else:\n        train_sampler = train_subsampler\n        val_sampler = val_subsampler\n\n        train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE,\n                                  sampler=train_sampler, num_workers=CPU_CORES, drop_last=True)\n        val_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE,\n                                sampler=val_sampler, num_workers=CPU_CORES, drop_last=True)\n    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE,\n                             shuffle=False, num_workers=CPU_CORES)\n\n    net.to(device)\n\n    if \"train\" in KERNEL_STAGES:\n        criterion = nn.CrossEntropyLoss()\n        optimizer = Adam(net.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-3, weight_decay=1e-4)\n        if os.path.exists(model_load_file + \"_opt.tar\"):\n            optimizer.load_state_dict(torch.load(model_load_file + \"_opt.tar\", map_location=\"cpu\"))\n        else:\n            print(\"Warning: Optimizer parameters not reloaded, training will be reset\")\n        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader)*NUM_EPOCHS, eta_min=1e-6)\n\n        for epoch in range(1, NUM_EPOCHS + 1):\n            print('-' * 50)\n            time.sleep(1)\n            train_step(train_loader, net, device, criterion, optimizer, epoch, scheduler, model_save_file)\n\n    if \"validate\" in KERNEL_STAGES:\n        print('-' * 50)\n        time.sleep(1)\n        eval_step(val_loader, net, device, mode=\"validation\")\n    if \"test\" in KERNEL_STAGES and TPU_CORES == 0:\n        print('-' * 50)\n        time.sleep(1)\n        generate_submission(test_loader, net, device, label_encoder)\n\nif TPU_CORES > 1:\n    xmp.spawn(map_run, args=(), nprocs=TPU_CORES, start_method='fork')\nelse:\n    map_run(0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"I hope this was useful, if you have any recommendations for improvements or questions I will be happy to hear!"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}