{"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":"# Setup","metadata":{"id":"5tPFJ89V3BFT"}},{"cell_type":"code","source":"!pip install timm\n!pip install pytorch-metric-learning","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2022-05-20T07:12:36.906632Z","iopub.execute_input":"2022-05-20T07:12:36.907324Z","iopub.status.idle":"2022-05-20T07:12:55.612291Z","shell.execute_reply.started":"2022-05-20T07:12:36.907223Z","shell.execute_reply":"2022-05-20T07:12:55.611446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Imports","metadata":{"id":"MyC4gTwZ3MKJ"}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport random\nimport os\n\nfrom PIL import Image\nfrom tqdm import tqdm\n\nimport timm\nfrom timm.optim import Lookahead\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nfrom pytorch_metric_learning.distances import CosineSimilarity","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:12:55.614095Z","iopub.execute_input":"2022-05-20T07:12:55.614351Z","iopub.status.idle":"2022-05-20T07:12:59.04701Z","shell.execute_reply.started":"2022-05-20T07:12:55.614323Z","shell.execute_reply":"2022-05-20T07:12:59.046113Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Global","metadata":{"id":"tirOg6jm3aIB"}},{"cell_type":"code","source":"IMG_SIZE = 512\nDATA_FOLDER  = '/kaggle/input/creation-dataset-512x512/'\nTRAIN_FOLDER = DATA_FOLDER + 'train_images/'","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:12:59.048444Z","iopub.execute_input":"2022-05-20T07:12:59.04953Z","iopub.status.idle":"2022-05-20T07:12:59.055492Z","shell.execute_reply.started":"2022-05-20T07:12:59.04948Z","shell.execute_reply":"2022-05-20T07:12:59.053936Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"random.seed(0)\nnp.random.seed(0)\ntorch.manual_seed(0)\ntorch.cuda.manual_seed(0)\ntorch.backends.cudnn.deterministic = True\nos.environ['PYTHONHASHSEED'] = str(0)","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:12:59.057892Z","iopub.execute_input":"2022-05-20T07:12:59.058714Z","iopub.status.idle":"2022-05-20T07:12:59.074368Z","shell.execute_reply.started":"2022-05-20T07:12:59.058675Z","shell.execute_reply":"2022-05-20T07:12:59.073463Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset and transformations","metadata":{"id":"8V_xuoN73lON"}},{"cell_type":"code","source":"import albumentations as A\nimport albumentations.pytorch as APT\nimport cv2","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:12:59.076191Z","iopub.execute_input":"2022-05-20T07:12:59.076583Z","iopub.status.idle":"2022-05-20T07:13:00.379838Z","shell.execute_reply.started":"2022-05-20T07:12:59.076543Z","shell.execute_reply":"2022-05-20T07:13:00.379076Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_transform = A.Compose([\n    A.HorizontalFlip(p=0.5),\n    A.ShiftScaleRotate(p=0.25, border_mode=cv2.BORDER_CONSTANT),\n    A.OpticalDistortion(p=0.25),\n    A.Perspective(p=0.25),\n    A.CoarseDropout(p=0.25),\n    A.RandomBrightnessContrast(p=0.25),\n    A.ToFloat(),\n    APT.transforms.ToTensorV2(),\n])\n\nval_transform = A.Compose([\n    A.HorizontalFlip(p=0.5),\n    A.ToFloat(),\n    APT.transforms.ToTensorV2(),\n])","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.381296Z","iopub.execute_input":"2022-05-20T07:13:00.381555Z","iopub.status.idle":"2022-05-20T07:13:00.387706Z","shell.execute_reply.started":"2022-05-20T07:13:00.381521Z","shell.execute_reply":"2022-05-20T07:13:00.387028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HotelTrainDataset:\n    def __init__(self, data, data_path, transform):\n        self.data = data\n        self.data_path = data_path\n        self.transform = transform\n        self.fake_load = False\n\n    def __len__(self):\n        return len(self.data)\n    \n    def __getitem__(self, idx):\n        record = self.data.iloc[idx]\n        path = self.data_path + record['image_id']\n        \n        if self.fake_load:\n            image = np.random.randint(0, 255, (32, 32, 3)).astype(np.uint8)\n        else:\n            image = np.array(Image.open(path)).astype(np.uint8)\n        \n        image = self.transform(image=image)\n        return {\n            'image': image['image'],\n            'target': record['hotel_id_code'],\n        }","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.389894Z","iopub.execute_input":"2022-05-20T07:13:00.390506Z","iopub.status.idle":"2022-05-20T07:13:00.40052Z","shell.execute_reply.started":"2022-05-20T07:13:00.390467Z","shell.execute_reply":"2022-05-20T07:13:00.399571Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{"id":"FpR2HfK93pvS"}},{"cell_type":"code","source":"class EmbeddingNet(nn.Module):\n    def __init__(self, n_classes, embed_size, backbone_name):\n        super(EmbeddingNet, self).__init__()\n        \n        self.embed_size = embed_size\n        self.backbone = timm.create_model(backbone_name, pretrained=True)\n        in_features = self.backbone.get_classifier().in_features\n\n        fc_name, _ = list(self.backbone.named_modules())[-1]\n        if fc_name == 'classifier':\n            self.backbone.classifier = nn.Identity()\n        elif fc_name == 'head.fc':\n            self.backbone.head.fc = nn.Identity()\n        elif fc_name == 'fc':\n            self.backbone.fc = nn.Identity()\n        else:\n            raise Exception('Unknown classifier layer: ' + fc_name)\n        \n        self.post = nn.Sequential(\n            nn.utils.weight_norm(nn.Linear(in_features, self.embed_size * 2), dim=None),\n            nn.BatchNorm1d(self.embed_size * 2),\n            nn.Dropout(0.2),\n            nn.utils.weight_norm(nn.Linear(self.embed_size * 2, self.embed_size)),\n        )\n\n        self.classifier = nn.Sequential(\n            nn.BatchNorm1d(self.embed_size),\n            nn.Dropout(0.2),\n            nn.Linear(self.embed_size, n_classes),\n        )\n        \n    def embed_and_classify(self, x):\n        x = self.forward(x)\n        return x, self.classifier(x)\n\n    def forward(self, x):\n        x = self.backbone(x)\n        x = x.view(x.size(0), -1)\n        x = self.post(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.402345Z","iopub.execute_input":"2022-05-20T07:13:00.402652Z","iopub.status.idle":"2022-05-20T07:13:00.414902Z","shell.execute_reply.started":"2022-05-20T07:13:00.402613Z","shell.execute_reply":"2022-05-20T07:13:00.414019Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model helper functions","metadata":{"id":"mTFCinps35ci"}},{"cell_type":"code","source":"def save_checkpoint(model, scheduler, optimizer, epoch, loss=None, score=None):\n    checkpoint = {\n        'epoch': epoch,\n        'model': model.state_dict(),\n        'scheduler': scheduler.state_dict(),\n        'optimizer': optimizer.state_dict(),\n        'loss': loss,\n        'score': score,\n    }\n    torch.save(checkpoint, f'classification-model-latest.pt')\n\ndef load_checkpoint(model, scheduler):\n    checkpoint = torch.load(f'/kaggle/input/classification-training/classification-model-latest.pt')\n    model.load_state_dict(checkpoint['model'])\n    scheduler.load_state_dict(checkpoint['scheduler'])\n    return model, scheduler, checkpoint['epoch']","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.416341Z","iopub.execute_input":"2022-05-20T07:13:00.416767Z","iopub.status.idle":"2022-05-20T07:13:00.428338Z","shell.execute_reply.started":"2022-05-20T07:13:00.416687Z","shell.execute_reply":"2022-05-20T07:13:00.427659Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_epoch(args, model, loader, criterion, optimizer, scheduler, epoch):\n    losses = []\n    targets_all = []\n    outputs_all = []\n    \n    model.train()\n    t = tqdm(loader)\n    for i, sample in enumerate(t):\n        optimizer.zero_grad()\n        \n        images = sample['image'].to('cuda')\n        targets = sample['target'].to('cuda')\n        \n        embeds, outputs = model.embed_and_classify(images)\n        loss = criterion(outputs, targets)\n        \n        loss.backward()\n        optimizer.step()\n        scheduler.step()\n                \n        losses.append(loss.item())\n        targets_all.extend(targets.cpu().numpy())\n        outputs_all.extend(torch.sigmoid(outputs).detach().cpu().numpy())\n\n        score = np.mean(targets_all == np.argmax(outputs_all, axis=1))\n        t.set_description(f'Epoch {epoch}/{args.epochs} - Train loss:{loss:0.4f}, score: {score:0.4f}')\n        \n    return np.mean(losses), score","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.431513Z","iopub.execute_input":"2022-05-20T07:13:00.431935Z","iopub.status.idle":"2022-05-20T07:13:00.441884Z","shell.execute_reply.started":"2022-05-20T07:13:00.431893Z","shell.execute_reply":"2022-05-20T07:13:00.441093Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_embeds(loader, model, bar_desc):\n    targets_all = []\n    outputs_all = []\n    \n    model.eval()\n    with torch.no_grad():\n        t = tqdm(loader, desc=bar_desc)\n        for i, sample in enumerate(t):\n            input = sample['image'].to('cuda')\n            target = sample['target'].to('cuda')\n            output = model(input)\n            \n            targets_all.extend(target.cpu().numpy())\n            outputs_all.extend(output.detach().cpu().numpy())\n            \n    return targets_all, outputs_all","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.443228Z","iopub.execute_input":"2022-05-20T07:13:00.443911Z","iopub.status.idle":"2022-05-20T07:13:00.452114Z","shell.execute_reply.started":"2022-05-20T07:13:00.443872Z","shell.execute_reply":"2022-05-20T07:13:00.451355Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_distance_matrix(embeds, base_embeds, distance_func):\n    distance_matrix = []\n    base_embeds = torch.Tensor(base_embeds)\n    embeds_dataset = torch.utils.data.TensorDataset(torch.Tensor(embeds))\n    embeds_dataloader = DataLoader(embeds_dataset, num_workers=2, batch_size=1024, shuffle=False)\n    \n    t = tqdm(embeds_dataloader)\n    for i, sample in enumerate(t): \n        distances = distance_func(sample[0], base_embeds)\n        distance_matrix.extend(distances.numpy())\n        \n    return np.array(distance_matrix)","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.453418Z","iopub.execute_input":"2022-05-20T07:13:00.453858Z","iopub.status.idle":"2022-05-20T07:13:00.465064Z","shell.execute_reply.started":"2022-05-20T07:13:00.453819Z","shell.execute_reply":"2022-05-20T07:13:00.464295Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_closest_match(base_df, base_embeds, valid_targets, valid_embeds, model, distance_func):\n    distance_matrix = get_distance_matrix(valid_embeds, base_embeds, distance_func)\n\n    preds = []\n    N_val = len(valid_embeds)\n    for i in tqdm(range(N_val), total=N_val, desc='Getting closest match'):\n        tmp_df = base_df.copy()\n        tmp_df['distance'] = distance_matrix[i]\n        tmp_df = tmp_df.sort_values(by=['distance', 'hotel_id'], ascending=False).reset_index(drop=True)\n        preds.extend([tmp_df['hotel_id_code'].unique()[:5]])\n\n    y = np.repeat([valid_targets], repeats=5, axis=0).T\n    preds = np.array(preds)\n    acc_top_1 = (preds[:, 0] == valid_targets).mean()\n    acc_top_5 = (preds == y).any(axis=1).mean()\n    print(f'Accuracy: {acc_top_1:0.4f}, top 5 accuracy: {acc_top_5:0.4f}')\n    return preds, distance_matrix\n\n\ndef test(base_loader, valid_loader, model, distance_func):\n    base_targets, base_embeds = get_embeds(base_loader, model, 'Generating embeds for train')\n    valid_targets, valid_embeds = get_embeds(valid_loader, model, 'Generating embeds for test')\n    val_preds, distance_matrix = test_closest_match(base_loader.dataset.data, base_embeds, valid_targets, valid_embeds, model, distance_func)\n\n    return base_embeds, valid_embeds, base_targets, valid_targets, val_preds, distance_matrix","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.466212Z","iopub.execute_input":"2022-05-20T07:13:00.4665Z","iopub.status.idle":"2022-05-20T07:13:00.478402Z","shell.execute_reply.started":"2022-05-20T07:13:00.466461Z","shell.execute_reply":"2022-05-20T07:13:00.477619Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"def iterate_loader(loader, epochs):\n    loader.dataset.fake_load = True\n    for i in range(epochs):\n        with torch.no_grad():\n            t = tqdm(loader, desc=f'Iterating loader {i+1}/{epochs}')\n            for j, sample in enumerate(t):\n                images = sample['image']\n                targets = sample['target']\n\n    loader.dataset.fake_load = False","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.479631Z","iopub.execute_input":"2022-05-20T07:13:00.480379Z","iopub.status.idle":"2022-05-20T07:13:00.490174Z","shell.execute_reply.started":"2022-05-20T07:13:00.480338Z","shell.execute_reply":"2022-05-20T07:13:00.489354Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_and_validate(args, data_df):\n    model_name = f'classification-model-{args.backbone_name}-{args.embed_size}embeds-{args.n_classes}hotels'\n    print(model_name)\n\n    val_df   = data_df.groupby('hotel_id').sample(args.val_samples, random_state=0)\n    train_df = data_df[~data_df['image_id'].isin(val_df['image_id'])]\n\n    train_dataset = HotelTrainDataset(train_df, TRAIN_FOLDER, train_transform)\n    train_loader  = DataLoader(train_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True, drop_last=True)\n    base_dataset  = HotelTrainDataset(train_df, TRAIN_FOLDER, val_transform)\n    base_loader   = DataLoader(base_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False)\n    val_dataset   = HotelTrainDataset(val_df, TRAIN_FOLDER, val_transform)\n    valid_loader  = DataLoader(val_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False)\n\n    print('Base:', len(base_dataset))\n    print('Validation:', len(val_dataset))\n\n    model = EmbeddingNet(args.n_classes, args.embed_size, args.backbone_name)\n    model = model.to('cuda')\n\n    distance  = CosineSimilarity()\n    criterion = nn.CrossEntropyLoss()\n    optimizer = Lookahead(torch.optim.AdamW(model.parameters(), lr=args.lr), k=3)\n    scheduler = torch.optim.lr_scheduler.OneCycleLR(\n            optimizer,\n            max_lr=args.lr,\n            epochs=args.epochs,\n            steps_per_epoch=len(train_loader),\n            div_factor=10,\n            final_div_factor=1,\n            pct_start=0.1,\n            anneal_strategy='cos',\n        )\n    \n    start_epoch = 1\n    if args.continue_from_checkpoint:\n        model, scheduler, last_epoch = load_checkpoint(model, scheduler)\n        iterate_loader(train_loader, last_epoch)\n        start_epoch = start_epoch + last_epoch\n\n    torch.cuda.empty_cache()\n\n    for epoch in range(start_epoch, args.epochs + 1):\n        train_loss, train_score = train_epoch(args, model, train_loader, criterion, optimizer, scheduler, epoch)\n    \n    save_checkpoint(model, scheduler, optimizer, epoch, train_loss, train_score)\n    base_embeds, valid_embeds, base_targets, valid_targets, val_preds, distance_matrix = test(base_loader, valid_loader, model, distance)","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.491539Z","iopub.execute_input":"2022-05-20T07:13:00.492185Z","iopub.status.idle":"2022-05-20T07:13:00.508251Z","shell.execute_reply.started":"2022-05-20T07:13:00.492143Z","shell.execute_reply":"2022-05-20T07:13:00.507322Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prepare data","metadata":{"id":"F2xgmwBW4LjC"}},{"cell_type":"code","source":"data_df = pd.read_csv(DATA_FOLDER + 'train_df.csv').drop(['Unnamed: 0'], axis=1)\ndata_df = data_df.sample(frac=0.9, random_state=0)\ndata_df['hotel_id_code'] = data_df['hotel_id'].astype('category').cat.codes.values.astype(np.int64)\ndata_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.509488Z","iopub.execute_input":"2022-05-20T07:13:00.50993Z","iopub.status.idle":"2022-05-20T07:13:00.564881Z","shell.execute_reply.started":"2022-05-20T07:13:00.50989Z","shell.execute_reply":"2022-05-20T07:13:00.564029Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Evalaute","metadata":{}},{"cell_type":"code","source":"%%time \n\nclass args:\n    epochs = 10\n    lr = 1e-3\n    batch_size = 16\n    num_workers = 2\n    embed_size = 4096\n    val_samples = 1\n    continue_from_checkpoint = False\n    backbone_name = 'efficientnet_b0'\n    n_classes = data_df['hotel_id_code'].nunique()\n\ntrain_and_validate(args, data_df)","metadata":{"execution":{"iopub.status.busy":"2022-05-20T07:13:00.566232Z","iopub.execute_input":"2022-05-20T07:13:00.566555Z","iopub.status.idle":"2022-05-20T07:14:45.783312Z","shell.execute_reply.started":"2022-05-20T07:13:00.566516Z","shell.execute_reply":"2022-05-20T07:14:45.782487Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}