{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Introduction\n* Use ArcFace loss as criterion in order to enhance the intraclass compactness and deal the open-class test data.\n* Cross entropy loss in ArcFace is replaced by focal loss for the class imbalance.\n* Use EfficientNet-b1 as the backbone.\n* Use Xception style seperable convolution blocks to replace the head in EfficientNet.\n* Add SAM (Spatial Attention Module).\n* Pretrained model was trained with DIM (Deep Info Max).\n\nI use self-supervised DIM for 10 epoches to train a model that can get some universal features in this long tail dataset in order to prevent the model can only learn the features of major classes.  \nBased on this pretrained model, I had had several experiments with Arcface but end up failure. I think the main reason is the serious imbalance: amount 740,072 training sample, the majorest class contains 5644 samples, while the minorest class only have 23 samples.  \nAfter finding a sugestion on a kaggle notebook for training Arcface, the model had been trained with focal loss."},{"metadata":{},"cell_type":"markdown","source":"# Enviroments"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\n\nos.system('pip install ../input/landmark-recognition-2020-packages/efficientnet_pytorch-0.7.0-py3-none-any.whl')\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.utils.data as Data\nimport torch.nn.functional as F\nimport torchvision.transforms as transforms\nfrom efficientnet_pytorch import EfficientNet\nimport pandas as pd\nfrom PIL import Image\nfrom sklearn.preprocessing import LabelEncoder\nimport multiprocessing\nimport time\nimport math\nimport matplotlib.pyplot as plt\nimport itertools\nimport matplotlib.pyplot as plt\n\nimport warnings\n\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Hyperparameters"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# dir\nSELF_DIR = '../input/landmark-recognition-2020-self/'\nORIGIN_TRAIN_CSV = '../input/landmark-recognition-2020/train.csv'\nTRAIN_CSV = '../input/landmark-recognition-2020-self/train.csv'\nVALID_CSV = '../input/landmark-recognition-2020-self/valid.csv'\nTEST_CSV = '../input/landmark-recognition-2020/sample_submission.csv'\nTRAIN_DIR = '../input/landmark-recognition-2020/train/'\nTEST_DIR = '../input/landmark-recognition-2020/test/'\n\n# general\nUSE_CUDA = torch.cuda.is_available()\nCPU_NUM = multiprocessing.cpu_count()\nLOG_STEPS = 500\nCLASSES_NUM = 10000\nEPS = 1e-6\nEPS_LOG = 1e-40\n\n# training\nEPOCHS = 30\nBATCH_SIZE = 64\nINPUT_SIZE = 288\n\n# model\nCOMPOUND_COEF = 1\nFEATURES_NUM = 2048\nEMBEDDING_SIZE = 512\nDROPOUT_RATE = 0.2\n\n# criterion\nS = 30\nM = 0.15\nGAMMA_FOCAL = 18\n\n# optimizer\nWEIGHT_DECAY = 5e-4\nMOMENTUM = 0.9\nLR = 1e-3\nK = 8\n\n# lr_scheduler\nT_0 = 10\nT_MULT = 2","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Dataset(Data.Dataset):\n    def __init__(self, is_train, dataframe, data_dir):\n        self.is_train = is_train\n        self.dataframe = dataframe\n        self.data_dir = data_dir\n\n        self.backup = None\n\n        if self.is_train:\n            transforms_list = [\n                transforms.Resize((INPUT_SIZE, INPUT_SIZE)),\n                transforms.RandomApply([transforms.RandomResizedCrop(size=INPUT_SIZE)], p=0.33),\n                transforms.RandomChoice([\n                    transforms.ColorJitter(0.2, 0.2, 0.2, 0.2),\n                    transforms.RandomAffine(\n                        degrees=10, translate=(0.2, 0.2),\n                        scale=(0.8, 1.2),\n                        resample=Image.BILINEAR)\n                ]),\n                transforms.ToTensor(),\n                transforms.RandomApply([transforms.RandomErasing(p=1, scale=(0.2, 0.33), ratio=(0.5, 2))], p=0.8),\n                transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.25, 0.25, 0.25))\n            ]\n        else:\n            transforms_list = [\n                transforms.Resize((INPUT_SIZE, INPUT_SIZE)),\n                transforms.ToTensor(),\n                transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.25, 0.25, 0.25))\n            ]\n        self.transforms = transforms.Compose(transforms_list)\n\n    def __getitem__(self, index):\n        image_id = self.dataframe.iloc[index].id\n        image_path = os.path.join(self.data_dir, image_id[0], image_id[1], image_id[2], '{}.jpg'.format(image_id))\n        image = Image.open(image_path)\n        image = self.transforms(image)\n\n        if self.is_train:\n            return [image, self.dataframe.iloc[index].landmark_id]\n        else:\n            return [image_id, image]\n\n    def __len__(self):\n        return self.dataframe.shape[0]\n\n    \ndef split_and_save_data(data_csv, valid_ratio):\n    data = pd.read_csv(data_csv)\n\n    counts = data.landmark_id.value_counts()\n    selected_classes = counts[:CLASSES_NUM].index\n\n    data = data.loc[data.landmark_id.isin(selected_classes)]\n\n    valid = data.sample(frac=valid_ratio, replace=False, random_state=1)\n    train = data.loc[~data.id.isin(valid.id)]\n\n    valid.to_csv('valid.csv')\n    train.to_csv('train.csv')\n\n\ndef load_data(train_csv, valid_csv, dir):\n    train = pd.read_csv(train_csv)\n    valid = pd.read_csv(valid_csv)\n\n    label_encoder = LabelEncoder()\n    label_encoder.fit(np.hstack((train.landmark_id.values, valid.landmark_id.values)))\n\n    train.landmark_id = label_encoder.transform(train.landmark_id)\n    valid.landmark_id = label_encoder.transform(valid.landmark_id)\n\n    train_loader = Data.DataLoader(\n        dataset=Dataset(True, train, dir),\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=4,\n    )\n    valid_loader = Data.DataLoader(\n        dataset=Dataset(True, valid, dir),\n        batch_size=BATCH_SIZE,\n        shuffle=False,\n        num_workers=4,\n    )\n\n    return train_loader, valid_loader","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# β-trainable Swish activation"},{"metadata":{"trusted":true},"cell_type":"code","source":"class SwishFunc(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, x, beta):\n        y = x * torch.sigmoid(beta * x)\n        ctx.save_for_backward(x, y, beta)\n\n        return y\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        x, y, beta = ctx.saved_tensors\n\n        grad_x = grad_output * (beta * y + torch.sigmoid(beta * x) * (1 - beta * y))\n        grad_beta = grad_output * (x * y - y ** 2)\n\n        return grad_x, grad_beta\n\n    \nclass Swish(nn.Module):\n    def __init__(self):\n        super(Swish, self).__init__()\n        self.beta = nn.Parameter(torch.FloatTensor([1]))  # beta is initialized to be 1\n\n    def forward(self, x):\n        return SwishFunc.apply(x, self.beta)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Xception style seperable convolution and SAM block"},{"metadata":{"trusted":true},"cell_type":"code","source":"class SeparableConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(SeparableConv, self).__init__()\n        self.pointwise = nn.Conv2d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=1,\n        )\n        self.swish1 = Swish()\n        self.depthwise = nn.Conv2d(\n            in_channels=out_channels,\n            out_channels=out_channels,\n            kernel_size=3,\n            padding=1,\n            groups=out_channels,\n            bias=False,\n        )\n        self.bn = nn.BatchNorm2d(out_channels, momentum=0.99)\n        self.swish2 = Swish()\n\n    def forward(self, inputs):\n        x = self.pointwise(inputs)\n        x = self.swish1(x)\n        x = self.depthwise(x)\n        x = self.swish2(self.bn(x))\n\n        return x\n\n\nclass SAM(nn.Module):\n    def __init__(self):\n        super(SAM, self).__init__()\n        self.conv = nn.Conv2d(2, 1, kernel_size=7, padding=3)\n        self.sigmoid = nn.Sigmoid()\n\n    def forward(self, features):\n        avg_descriptor = torch.mean(features, dim=1, keepdim=True)\n        max_descriptor, _ = torch.max(features, dim=1, keepdim=True)\n\n        descriptor = torch.cat([avg_descriptor, max_descriptor], dim=1)\n\n        attention_map = self.sigmoid(self.conv(descriptor))\n\n        del avg_descriptor, max_descriptor, descriptor\n\n        return features * attention_map","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# ArcFace"},{"metadata":{"trusted":true},"cell_type":"code","source":"class ArcFace(nn.Module):\n    def __init__(self, embedding_size, class_num, s=64.0, m=0.50):\n        super().__init__()\n        self.in_features = embedding_size\n        self.out_features = class_num\n        self.s = s\n        self.m = m  # the angular penalty\n        self.weight = nn.Parameter(torch.FloatTensor(class_num, embedding_size))\n        nn.init.xavier_uniform_(self.weight)\n\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(m) * m\n\n    def forward(self, input, label):\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))  # cos(theta)\n        bs = len(label)\n        label_cosine = cosine[range(bs), label]\n        \n#         print(cosine)\n#         print(label_cosine)\n        \n        label_sine = ((1.0 - label_cosine.pow(2)).clamp(0, 1)).sqrt()\n        phi = label_cosine * self.cos_m - label_sine * self.sin_m  # cos(theta+m)\n\n        # if theta+m > pi use penalty of CosFace cos(theta) - self.mm\n        phi = torch.where(label_cosine > self.th, phi, label_cosine - self.mm)\n\n        output = cosine * 1.0  # make backward work\n        output[range(bs), label] = phi\n        \n        return output * self.s\n    \n    def valid(self, input, label):\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))  # cos(theta)\n        bs = len(label)\n        label_cosine = cosine[range(bs), label]\n        label_sine = ((1.0 - label_cosine.pow(2)).clamp(0, 1)).sqrt()\n        phi = label_cosine * self.cos_m - label_sine * self.sin_m  # cos(theta+m)\n\n        # if theta+m > pi use penalty of CosFace cos(theta) - self.mm\n        phi = torch.where(label_cosine > self.th, phi, label_cosine - self.mm)\n\n        output = cosine * 1.0  # make backward work\n        output[range(bs), label] = phi\n        \n        return cosine, output * self.s\n\n    def inference(self, input):\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n\n        return cosine","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Extractor(nn.Module):\n    def __init__(self, in_channels, out_features):\n        super(Extractor, self).__init__()\n\n        mid_channels = int(in_channels * math.pow(out_features / in_channels, 0.5))\n\n        self.sep1 = SeparableConv(in_channels, mid_channels)\n        self.sep2 = SeparableConv(mid_channels, out_features)\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n\n    def forward(self, input):\n        x = self.sep1(input)\n        x = self.sep2(x)\n        x = self.avg_pool(x)\n        x = torch.flatten(x, start_dim=1)\n\n        return x\n\n    \nclass Encoder(nn.Module):\n    def __init__(self, compound_coef, out_features):\n        super(Encoder, self).__init__()\n        self.compound_coef = compound_coef\n\n        self.base = EfficientNet.from_name('efficientnet-b{}'.format(self.compound_coef))\n        self.sam1 = SAM()\n        self.sam2 = SAM()\n        self.sam3 = SAM()\n\n        features_num = self.base._conv_head.in_channels\n        self.extractor = Extractor(features_num, out_features)\n\n    def forward(self, input):\n        # Stem\n        x = self.base._swish(self.base._bn0(self.base._conv_stem(input)))\n\n        # Blocks\n        x = self.sam1(x)\n        for idx, block in enumerate(self.base._blocks):\n            drop_connect_rate = self.base._global_params.drop_connect_rate\n            if drop_connect_rate:\n                drop_connect_rate *= float(idx) / len(self.base._blocks)  # scale drop connect_rate\n            x = block(x, drop_connect_rate=drop_connect_rate)\n\n        # Head\n        x = self.extractor(self.sam2(x))\n\n        return x\n\n    \nclass EncoderHead(nn.Module):\n    def __init__(self, in_features, out_features):\n        super(EncoderHead, self).__init__()\n        self.dropout = nn.Dropout(DROPOUT_RATE)\n        self.fc = nn.Linear(in_features, out_features)\n        self.bn = nn.BatchNorm1d(out_features)\n        \n    \n    def forward(self, input):\n        return self.bn(self.fc(self.dropout(input)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# One-hot Focal loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"class FocalLoss(nn.Module):\n    def __init__(self, gamma):\n        super(FocalLoss, self).__init__()\n        self.gamma = gamma\n\n        self.softmax = nn.Softmax(dim=1)\n\n    def forward(self, output, label):\n        pred = self.softmax(output)\n        y_ = pred[range(len(label)), label]\n        loss = -torch.mean(torch.pow(1 - y_, self.gamma) * torch.log(y_ + EPS_LOG))\n        \n        return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training and validation"},{"metadata":{"trusted":true},"cell_type":"code","source":"class AccCounter:\n    def __init__(self):\n        super(AccCounter, self).__init__()\n        self.total = .0\n        self.correction = .0\n\n    def update(self, logits, label):\n        _, pred = torch.max(logits, dim=1)\n\n        self.total += len(label)\n        self.correction += torch.sum(pred == label).data.item()\n\n    def show(self):\n        return self.correction / self.total\n\n\ndef train(start_epoch, max_epoch,\n          encoder, encoder_head, arc, criterion,\n          optim_encoder, optim_arc, scheduler_encoder, scheduler_arc,\n          train_loader, loss_list, valid_loader=None, valid_n_epochs=1):\n    encoder.train()\n    encoder_head.train()\n    arc.train()\n\n    for epoch_index in range(start_epoch, max_epoch):\n        print('-----{} epoch-----'.format(epoch_index))\n\n        losses = 0\n        start_time = time.time()\n\n        for batch_index, (input, label) in enumerate(train_loader):\n            if USE_CUDA:\n                input = input.cuda()\n                label = label.cuda()\n\n            features = encoder_head(encoder(input))\n            logits = arc(features, label)\n\n            loss = criterion(logits, label)\n            loss_list.append(loss.data.item())\n            losses += loss.data.item()\n\n            optim_encoder.zero_grad()\n            optim_arc.zero_grad()\n            loss.backward()\n            optim_arc.step()\n            optim_encoder.step()\n            scheduler_encoder.step()\n            scheduler_arc.step()\n                \n            if batch_index % LOG_STEPS == LOG_STEPS - 1:\n                print('{} / {} | '.format(batch_index + 1, len(train_loader)),\n                      'time {:.3f} | '.format((time.time() - start_time) / (batch_index + 1)),\n                      'avg loss {:.5f} | '.format(losses / LOG_STEPS),\n                      'lr {:.7f}'.format(optim_arc.param_groups[0]['lr'])\n                     )\n\n                losses = 0\n\n        state = {\n            'encoder': encoder.state_dict(),\n            'encoder_head': encoder_head.state_dict(),\n            'arc': arc.state_dict(),\n            'optim_encoder': optim_encoder.state_dict(),\n            'optim_arc': optim_arc.state_dict(),\n            'scheduler_encoder': scheduler_encoder.state_dict(),\n            'scheduler_arc': scheduler_arc.state_dict(),\n            'epoch': epoch_index,\n        }\n        torch.save(state, 'epoch_arc{}.csv'.format(epoch_index))\n\n        loss_list_np = np.array(loss_list)\n        np.save('loss_list_arc.npy', loss_list_np)\n\n        if epoch_index % valid_n_epochs == valid_n_epochs - 1 and valid_loader is not None:\n            valid(encoder, encoder_head, arc, criterion, valid_loader)\n\n    draw_loss_trend(loss_list)\n\n\ndef valid(encoder, encoder_head, arc, criterion, valid_loader):\n    print('validation start')\n    start_time = time.time()\n\n    encoder.eval()\n    encoder_head.eval()\n    arc.eval()\n\n    losses = 0\n    acc_counter = AccCounter()\n\n    with torch.no_grad():\n        for _, batch in enumerate(valid_loader):\n            input, label = batch\n            if USE_CUDA:\n                input = input.cuda()\n                label = label.cuda()\n\n            features = encoder_head(encoder(input))\n            logits_inf, logits = arc.valid(features, label)\n            \n            loss = criterion(logits, label)\n            losses += loss.data.item()\n            acc_counter.update(logits_inf, label)\n\n    avg_loss = losses / len(valid_loader)\n\n    encoder.train()\n    encoder_head.train()\n    arc.train()\n\n    print('validation complete, avg loss {:.5f}, accuracy {:.5f}, cost {:.3f} seconds'\n          .format(avg_loss, acc_counter.show(), time.time() - start_time))\n    \n    \ndef draw_loss_trend(loss_list):\n    interval = 1000\n    \n    loss_list_ = [np.mean(loss_list[i * interval : (i + 1) * interval]) for i in range(int(len(loss_list) / interval))]\n    \n    x = np.arange(1, len(loss_list_) + 1) * interval\n\n    plt.title('ArcFace loss trend')\n    plt.xlabel('iterations')\n    plt.ylabel('loss')\n    plt.plot(x, loss_list_)\n    plt.savefig('loss_iter{}.jpg'.format(len(loss_list_) * interval))\n    plt.clf()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Seed"},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed=2020):\n    np.random.seed(seed)\n    torch.manual_seed(seed)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Main"},{"metadata":{"trusted":true},"cell_type":"code","source":"seed_everything(int(time.time()))\n\ncheckpoint_file = 'epoch_arc27.csv'\nif USE_CUDA:\n    checkpoint = torch.load(os.path.join(SELF_DIR, checkpoint_file))\nelse:\n    checkpoint = torch.load(os.path.join(SELF_DIR, checkpoint_file), map_location='cpu')\n\ntrain_loader, valid_loader = load_data(TRAIN_CSV, VALID_CSV, TRAIN_DIR)\n\nencoder = Encoder(COMPOUND_COEF, FEATURES_NUM)\nencoder_head = EncoderHead(FEATURES_NUM, EMBEDDING_SIZE)\narc = ArcFace(EMBEDDING_SIZE, CLASSES_NUM, S, M)\n\ncriterion = FocalLoss(GAMMA_FOCAL)\n\nif USE_CUDA:\n    encoder = encoder.cuda()\n    encoder_head = encoder_head.cuda()\n    arc = arc.cuda()\n\noptim_encoder = torch.optim.Adam(\n    params=encoder.parameters(),\n    lr=LR / K,\n    weight_decay=WEIGHT_DECAY,\n)\noptim_arc = torch.optim.Adam(\n    params=itertools.chain(encoder_head.parameters(), arc.parameters()),\n    lr=LR,\n    weight_decay=WEIGHT_DECAY,\n)\nscheduler_encoder = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n    optimizer=optim_encoder,\n    T_0=len(train_loader) * T_0 * T_MULT,\n    T_mult=T_MULT,\n    eta_min=1e-6 / K,\n)\nscheduler_arc = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n    optimizer=optim_arc,\n    T_0=len(train_loader) * T_0 * T_MULT,\n    T_mult=T_MULT,\n    eta_min=1e-6,\n)\n\nencoder.load_state_dict(checkpoint['encoder'])\nencoder_head.load_state_dict(checkpoint['encoder_head'])\narc.load_state_dict(checkpoint['arc'])\noptim_encoder.load_state_dict(checkpoint['optim_encoder'])\noptim_arc.load_state_dict(checkpoint['optim_arc'])\nscheduler_encoder.load_state_dict(checkpoint['scheduler_encoder'])\nscheduler_arc.load_state_dict(checkpoint['scheduler_arc'])\n\nloss_list = np.load(os.path.join(SELF_DIR, 'loss_list_arc.npy')).tolist()\n\ntrain(\n    start_epoch=checkpoint['epoch']+1,\n    max_epoch=EPOCHS,\n    encoder=encoder,\n    encoder_head=encoder_head,\n    arc=arc,\n    criterion=criterion,\n    optim_encoder=optim_encoder,\n    optim_arc=optim_arc,\n    scheduler_encoder=scheduler_encoder,\n    scheduler_arc=scheduler_arc,\n    train_loader=train_loader,\n    loss_list=loss_list,\n    valid_loader=valid_loader,\n    valid_n_epochs=2,\n)","execution_count":null,"outputs":[]}],"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}