深層距離学習の学習コード

import numpy as np
import pandas as pd
np.random.seed(42)
df_train_csv = pd.read_csv('./train.csv')
df_val = df_train_csv[['gid','path','cite_gid','cite_path']]
df_val = df_val[~df_val.gid.isin([
1000415998,
1000098072,
1000389161,
1000374219,
1000591092,
1000657233,
1000186487,
1000769939,
1000277425,
1000151025,
1000189511,
1000718061,
1000704867,
1000824535,
1000644060,
1000442455,
1000145820,
1000666393,
1000331547,
1000319682,
1000436708,
1000439850,
1000487334,
1000304843,
1000127290,
1000290876,
1000826902,
1000283549,
1000241266,
1000274048,
1000266800,
1000177036,
1000437165,
1000053344,
1000408709,
1000079785,
1000355096,
1000548903,
1000687203,
1000233450,
1000761260,
1000192698,
1000569557,
1000373900,
1000654732,
1000133351,
1000690702,
1000694012,
1000423053,
1000560746,
1000155115,
1000440139,
1000839319,
1000485424,
1000056833,
1000296420,
1000304411,
1000458735,
1000435047,
1000720130,
1000604660,
1000630337,
1000596073,
1000764671,
1000161575,
1000467984,
1000035376,
1000105009,
1000099468,
1000462740,
1000757067,
1000796804,
1000137820,
1000269117,
1000008260,
1000099526,
1000307721,
1000561255,
1000466440,
1000544605,
1000816852,
1000840580,
1000194991,
1000813973,
1000744402,
1000658628,
1000076839,
1000322446,
1000835850,
1000821840,
1000591291,
1000628088,
1000636840,
1000668049,
1000305059,
1000013095,
1000172118,
1000800426,
1000752197,
1000258851,
1000039624,
1000801015,
1000504341,
1000647447,
1000597408,
1000489125,
1000134779,
1000687795,
1000375551,
1000617907,
1000356327,
1000497265,
1000089103,
1000126953,
1000371922,
1000336053,
1000482451,
1000720504,
1000807922,
1000244134,
1000176371,
1000220919,
1000474265,
1000471111,
1000356019,
1000724593,
1000615127,
1000551389,
1000793814,
1000537535,
1000490495,
1000735052,
1000243275,
1000060337,
1000527966,
1000011504,
1000458042,
1000569494,
1000382056,
1000724035,
1000252992,
1000245743,
1000102672,
1000257363,
1000297115,
1000833194,
1000582039,
1000737641,
1000207020,
1000093704,
1000103655,
1000361134,
1000373487,
1000355692,
1000161457,
1000385915,
1000602792,
1000383692,
1000338666,
1000388536,
1000364613,
1000806965,
1000694970,
1000725445,
1000412133,
1000666680,
1000029234,
1000054906,
1000112680,
1000646625
])]
df_val['path']=''
df_val['cite_path']=''
df_val['path'] = df_val['gid'].apply(lambda x: 'apply_images_yolocrop/'+str(x)+'.jpg')
df_val['cite_path'] = df_val['cite_gid'].apply(lambda x: 'resize/'+str(x)[:6]+'_yolocrop/'+str(x)+'.jpg')
df_val['split']='val'

df_cite_csv = pd.read_csv('./cite_v2.csv')
#df_cite = df_train_csv[['cite_gid']]
df_cite = df_cite_csv[['gid']]
df_cite = df_cite.rename(columns={'cite_gid':'gid'})

#df_cite_add = df_cite_csv[~df_cite_csv['gid'].isin(df_cite['gid'].values)][['gid']]
#df_cite_add = df_cite_add.iloc[np.random.randint(1,len(df_cite_add),100000-len(df_cite))].reset_index(drop=True) #後で足してちょうど100,000行になるように
#
#
#df_cite_merge = pd.concat([df_cite, df_cite_add])
#df_cite_merge['path'] = df_cite_merge['gid'].apply(lambda x: 'resize/'+str(x)[:6]+'_yolocrop/'+str(x)+'.jpg')
#df_cite_merge['cite_gid'] = df_cite_merge['gid']
#df_cite_merge['cite_path'] = df_cite_merge['path']
#df_cite_merge['split']='train'

df_cite['path'] = df_cite['gid'].apply(lambda x: 'resize/'+str(x)[:6]+'_yolocrop/'+str(x)+'.jpg')
df_cite['cite_gid'] = df_cite['gid']
df_cite['cite_path'] = df_cite['path']
df_cite['split']='train'

#df_train = pd.concat([df_cite_merge, df_val]).reset_index(drop=True)
df_train = pd.concat([df_cite, df_val]).reset_index(drop=True)
df_train.to_csv('./train_preprocessed.csv')

for i in range(86):
  bno = '1000' + ('00'+str(i))[-2:]
  !cp -p /content/drive/MyDrive/{bno}_yolocrop.zip ./
  !unzip {bno}_yolocrop.zip > /dev/null
  !mv {bno}_yolocrop/ resize
!cp -p /content/drive/MyDrive/apply_images_yolocrop.zip ./
!unzip apply_images_yolocrop.zip > /dev/null

"""
類似画像検索の特徴量抽出のモデルを、深層距離学習により構築する
学習の手法として、Arcfaceを採用
参考コード: https://www.kaggle.com/tanulsingh077/pytorch-metric-learning-pipeline-only-images
"""


# Import Libraries

import sys
sys.path.append('./pytorch-image-models')

import math
import os
import random
from datetime import datetime as dt
from tqdm import tqdm

import numpy as np
import pandas as pd
from sklearn.preprocessing import LabelEncoder

import cv2

import albumentations
from albumentations.pytorch.transforms import ToTensorV2

import timm
import torch
import torch.nn as nn
from torch.nn import Parameter
from torch.nn import functional as F
from torch.utils.data import Dataset, DataLoader
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau, _LRScheduler


# Configuration

#DIM = (512, 512)
DIM = (224,224)
NUM_WORKERS = 1
#NUM_WORKERS = 0
TRAIN_BATCH_SIZE = 8
VALID_BATCH_SIZE = 8
EPOCHS = 20
SEED = 43

device = torch.device('cuda')
#device = torch.device('cpu')

## Model
model_name = 'efficientnet_b3' #efficientnet_b0-b7
#model_name = 'efficientnet_b0' #efficientnet_b0-b7

## Model Path
#IMG_MODEL_PATH = f'./models/model_efficientnet_b0_IMG_SIZE_224_arcface.bin' #***
IMG_MODEL_PATH = f'./models/model_efficientnet_b3_IMG_SIZE_224_arcface.bin' #***

## Metric Loss and its params
loss_module = 'arcface' #'cosface' #'adacos'
s = 30.0
m = 0.5
ls_eps = 0.0
easy_margin = False

## Scheduler and its params
scheduler_params = {
    "lr_start": 1e-5,
    "lr_max": 1e-5 * TRAIN_BATCH_SIZE,
    "lr_min": 1e-6,
    "lr_ramp_ep": 5,
    "lr_sus_ep": 0,
    "lr_decay": 0.8,
    }

## Model Params
df = pd.read_csv(f'./train_preprocessed.csv')

model_params = {
    'n_classes': df['cite_gid'].nunique(),
    'model_name': 'efficientnet_b3',
#    'model_name': 'efficientnet_b0',
#    'use_fc': False,
#    'fc_dim': 512,
    'use_fc': True,
    'fc_dim': 300,
    'dropout': 0.0,
    'loss_module': loss_module,
    's': 30.0,
    'margin': 0.50,
    'ls_eps': 0.0,
    'theta_zero': 0.785,
#    'pretrained': True
    'pretrained': False #***
    }


# Utils

def seed_torch(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

seed_torch(SEED)


class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def fetch_scheduler(optimizer):
    if SCHEDULER == 'ReduceLROnPlateau':
        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=factor, patience=patience, verbose=True, eps=eps)
    elif SCHEDULER == 'CosineAnnealingLR':
        scheduler = CosineAnnealingLR(optimizer, T_max=T_max, eta_min=min_lr, last_epoch=-1)
    elif SCHEDULER == 'CosineAnnealingWarmRestarts':
        scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=T_0, T_mult=1, eta_min=min_lr, last_epoch=-1)
    return scheduler


def fetch_loss():
    loss = nn.CrossEntropyLoss()
    return loss


# Augmentations

def get_train_transforms():
    return albumentations.Compose(
        [
            albumentations.Resize(DIM[0], DIM[1], always_apply=True),
            albumentations.HorizontalFlip(p=0.5),
            albumentations.VerticalFlip(p=0.5),
            albumentations.Rotate(limit=120, p=0.8),
            albumentations.RandomBrightness(limit=(0.09, 0.6), p=0.5),
            albumentations.Cutout(num_holes=8, max_h_size=8, max_w_size=8, fill_value=0, always_apply=False, p=0.5),
#            albumentations.Cutout(num_holes=16, max_h_size=16, max_w_size=16, fill_value=255, always_apply=False, p=0.5),
            albumentations.ShiftScaleRotate(
                shift_limit=0.25, scale_limit=0.1, rotate_limit=0
            ),
            albumentations.Normalize(),
            ToTensorV2(p=1.0),
        ]
    )

def get_valid_transforms():

    return albumentations.Compose(
        [
         albumentations.Resize(DIM[0], DIM[1], always_apply=True),
         albumentations.Normalize(),
         ToTensorV2(p=1.0)
        ]
    )


# Dataset

class JPODataset(Dataset):
    def __init__(self, csv, transforms=None):

        self.csv = csv.reset_index()
        self.augmentations = transforms

    def __len__(self):
        return self.csv.shape[0]

    def __getitem__(self, index):
        row = self.csv.iloc[index]

        image = cv2.imread(row.path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        if self.augmentations:
            augmented = self.augmentations(image=image)
            image = augmented['image']

        return image, torch.tensor(row.cite_gid)


# Model

class JPONet(nn.Module):

    def __init__(self,
                 n_classes,
                 model_name='efficientnet_b3',
                 use_fc=False,
                 fc_dim=512,
                 dropout=0.0,
                 loss_module='softmax',
                 s=30.0,
                 margin=0.50,
                 ls_eps=0.0,
                 theta_zero=0.785,
                 pretrained=True):
        """
        :param n_classes:
        :param model_name: name of model from pretrainedmodels
            e.g. resnet50, resnext101_32x4d, pnasnet5large
        :param pooling: One of ('SPoC', 'MAC', 'RMAC', 'GeM', 'Rpool', 'Flatten', 'CompactBilinearPooling')
        :param loss_module: One of ('arcface', 'cosface', 'softmax')
        """
        super(JPONet, self).__init__()
        print('Building Model Backbone for {} model'.format(model_name))

        self.backbone = timm.create_model(model_name, pretrained=pretrained)
        final_in_features = self.backbone.classifier.in_features

        self.backbone.classifier = nn.Identity()
        self.backbone.global_pool = nn.Identity()

        self.pooling = nn.AdaptiveAvgPool2d(1)

        self.use_fc = use_fc
        if use_fc:
            self.dropout = nn.Dropout(p=dropout)
            self.fc = nn.Linear(final_in_features, fc_dim)
            self.bn = nn.BatchNorm1d(fc_dim)
            self._init_params()
            final_in_features = fc_dim

        self.loss_module = loss_module
        if loss_module == 'arcface':
#            self.final = ArcMarginProduct(final_in_features, n_classes, s=s, m=margin, easy_margin=False, ls_eps=ls_eps)
            self.final = ArcMarginProduct(final_in_features, n_classes, s=s, m=margin, easy_margin=True, ls_eps=ls_eps) #★
        elif loss_module == 'cosface':
            self.final = AddMarginProduct(final_in_features, n_classes, s=s, m=margin)
        elif loss_module == 'adacos':
            self.final = AdaCos(final_in_features, n_classes, m=margin, theta_zero=theta_zero)
        else:
            self.final = nn.Linear(final_in_features, n_classes)

    def _init_params(self):
        nn.init.xavier_normal_(self.fc.weight)
        nn.init.constant_(self.fc.bias, 0)
        nn.init.constant_(self.bn.weight, 1)
        nn.init.constant_(self.bn.bias, 0)

    def forward(self, x, label):
        feature = self.extract_feat(x)
        if self.loss_module in ('arcface', 'cosface', 'adacos'):
            logits = self.final(feature, label)
        else:
            logits = self.final(feature)
        return logits

    def extract_feat(self, x):
        batch_size = x.shape[0]
        x = self.backbone(x)
        x = self.pooling(x).view(batch_size, -1)

        if self.use_fc:
            x = self.dropout(x)
            x = self.fc(x)
            x = self.bn(x)

        return x


# Metric Learning Losses
# ref: https://github.com/lyakaap/Landmark2019-1st-and-3rd-Place-Solution/blob/master/src/modeling/metric_learning.py

class ArcMarginProduct(nn.Module):
    """
    Implement of large margin arc distance:
    Args:
        in_features: size of each input sample
        out_features: size of each output sample
        s: norm of input feature
        m: margin
        cos(theta + m)
    """
    def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False, ls_eps=0.0):
        super(ArcMarginProduct, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m
        self.ls_eps = ls_eps  # label smoothing
        self.weight = Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)

        self.easy_margin = easy_margin
        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        self.th = math.cos(math.pi - m)
        self.mm = math.sin(math.pi - m) * m

    def forward(self, input, label):
        # --------------------------- cos(theta) & phi(theta) ---------------------------
        cosine = F.linear(F.normalize(input), F.normalize(self.weight))
        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
        phi = cosine * self.cos_m - sine * self.sin_m
        if self.easy_margin:
            phi = torch.where(cosine > 0, phi, cosine)
        else:
            phi = torch.where(cosine > self.th, phi, cosine - self.mm)
        # --------------------------- convert label to one-hot ---------------------------
        # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda')
        one_hot = torch.zeros(cosine.size(), device='cuda')
#        one_hot = torch.zeros(cosine.size(), device='cpu')
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        if self.ls_eps > 0:
            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features
        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
        output *= self.s

        return output


class AddMarginProduct(nn.Module):
    """
    Implement of large margin cosine distance:
    Args:
        in_features: size of each input sample
        out_features: size of each output sample
        s: norm of input feature
        m: margin
        cos(theta) - m
    """

    def __init__(self, in_features, out_features, s=30.0, m=0.40):
        super(AddMarginProduct, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m
        self.weight = Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)

    def forward(self, input, label):
        # --------------------------- cos(theta) & phi(theta) ---------------------------
        cosine = F.linear(F.normalize(input), F.normalize(self.weight))
        phi = cosine - self.m
        # --------------------------- convert label to one-hot ---------------------------
        one_hot = torch.zeros(cosine.size(), device='cuda')
#        one_hot = torch.zeros(cosine.size(), device='cpu')
        # one_hot = one_hot.cuda() if cosine.is_cuda else one_hot
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)  # you can use torch.where if your torch.__version__ is 0.4
        output *= self.s
        # print(output)

        return output


class AdaCos(nn.Module):
    def __init__(self, in_features, out_features, m=0.50, ls_eps=0, theta_zero=math.pi/4):
        super(AdaCos, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.theta_zero = theta_zero
        self.s = math.log(out_features - 1) / math.cos(theta_zero)
        self.m = m
        self.ls_eps = ls_eps # label smoothing
        self.weight = Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)

    def forward(self, input, label):
        # normalize features
        x = F.normalize(input)
        # normalize weights
        W = F.normalize(self.weight)
        # dot product
        logits = F.linear(x, W)
        # add margin
        theta = torch.acos(torch.clamp(logits, -1.0+1e-7, 1.0-1e-7))
        target_logits = torch.cos(theta + self.m)
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        if self.ls_eps > 0:
            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features
        output = logits * (1 - one_hot) + target_logits * one_hot
        # feature re-scale
        with torch.no_grad():
            B_avg = torch.where(one_hot < 1, torch.exp(self.s * logits), torch.zeros_like(logits))
            B_avg = torch.sum(B_avg) / input.size(0)
            theta_med = torch.median(theta)
            self.s = torch.log(B_avg) / torch.cos(torch.min(self.theta_zero * torch.ones_like(theta_med), theta_med))
        output *= self.s

        return output


# Custom LR

class JPOScheduler(_LRScheduler):
    def __init__(self, optimizer, lr_start=5e-6, lr_max=1e-5,
                 lr_min=1e-6, lr_ramp_ep=5, lr_sus_ep=0, lr_decay=0.8,
                 last_epoch=-1):
        self.lr_start = lr_start
        self.lr_max = lr_max
        self.lr_min = lr_min
        self.lr_ramp_ep = lr_ramp_ep
        self.lr_sus_ep = lr_sus_ep
        self.lr_decay = lr_decay
        super(JPOScheduler, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        if not self._get_lr_called_within_step:
            warnings.warn("To get the last learning rate computed by the scheduler, "
                          "please use `get_last_lr()`.", UserWarning)

        if self.last_epoch == 0:
            self.last_epoch += 1
            return [self.lr_start for _ in self.optimizer.param_groups]

        lr = self._compute_lr_from_epoch()
        self.last_epoch += 1

        return [lr for _ in self.optimizer.param_groups]

    def _get_closed_form_lr(self):
        return self.base_lrs

    def _compute_lr_from_epoch(self):
        if self.last_epoch < self.lr_ramp_ep:
            lr = ((self.lr_max - self.lr_start) /
                  self.lr_ramp_ep * self.last_epoch +
                  self.lr_start)

        elif self.last_epoch < self.lr_ramp_ep + self.lr_sus_ep:
            lr = self.lr_max

        else:
            lr = ((self.lr_max - self.lr_min) * self.lr_decay**
                  (self.last_epoch - self.lr_ramp_ep - self.lr_sus_ep) +
                  self.lr_min)
        return lr


# Training Function

def train_fn(dataloader, model, criterion, optimizer, device, scheduler, epoch):
    model.train()
    loss_score = AverageMeter()

    tk0 = tqdm(enumerate(dataloader), total=len(dataloader))
    for bi,d in tk0:

        batch_size = d[0].shape[0]

        images = d[0]
        targets = d[1]

        images = images.to(device)
        targets = targets.to(device)

        optimizer.zero_grad()

        output = model(images, targets)

        loss = criterion(output, targets)
#        loss = criterion(output, targets.long()) #★

        loss.backward()
        optimizer.step()

        loss_score.update(loss.detach().item(), batch_size)
        tk0.set_postfix(Train_Loss=loss_score.avg, Epoch=epoch, LR=optimizer.param_groups[0]['lr'])

    if scheduler is not None:
        scheduler.step()

    return loss_score


# Evaluation Function

def eval_fn(data_loader,model,criterion,device):

    loss_score = AverageMeter()

    model.eval()
    tk0 = tqdm(enumerate(data_loader), total=len(data_loader))

    with torch.no_grad():

        for bi,d in tk0:
            batch_size = d[0].size()[0]

            image = d[0]
            targets = d[1]

            image = image.to(device)
            targets = targets.to(device)

            output = model(image,targets)

            loss = criterion(output,targets)

            loss_score.update(loss.detach().item(), batch_size)
            tk0.set_postfix(Eval_Loss=loss_score.avg)

    return loss_score

df = pd.read_csv(f'./train_preprocessed.csv')

encoder = LabelEncoder()
df['cite_gid'] = encoder.fit_transform(df['cite_gid'].map(str))

train = df[df['split']=='train'].reset_index(drop=True)
valid = df[df['split']=='val'].reset_index(drop=True)

# Defining DataSet
train_dataset = JPODataset(
    csv=train,
    transforms=get_train_transforms(),
)

valid_dataset = JPODataset(
    csv=valid,
    transforms=get_valid_transforms(),
)

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=TRAIN_BATCH_SIZE,
    pin_memory=True,
    drop_last=True,
    num_workers=NUM_WORKERS
)

valid_loader = torch.utils.data.DataLoader(
    valid_dataset,
    batch_size=VALID_BATCH_SIZE,
    num_workers=NUM_WORKERS,
    shuffle=False,
    pin_memory=True,
    drop_last=False,
)

# Defining Device
device = torch.device("cuda")
#device = torch.device("cpu")

# Defining Model for specific fold
model = JPONet(**model_params)
model.load_state_dict(torch.load(IMG_MODEL_PATH), strict=False) #***
model.to(device)

# Defining criterion
criterion = fetch_loss()
criterion.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=scheduler_params['lr_start'])
#optimizer = torch.optim.SGD([{'params':model.parameters()},{'params':}], lr=scheduler_params['lr_start']) #★

# Defining LR SCheduler
scheduler = JPOScheduler(optimizer, **scheduler_params)

best_loss = 10000
for epoch in range(EPOCHS):
    train_loss = train_fn(train_loader, model, criterion, optimizer, device, scheduler=scheduler, epoch=epoch)

    valid_loss = eval_fn(valid_loader, model, criterion,device)

    if valid_loss.avg < best_loss:
        best_loss = valid_loss.avg
        torch.save(model.state_dict(), f'./models/model_{model_name}_IMG_SIZE_{DIM[0]}_{loss_module}.bin')
        print('best model found for epoch {}'.format(epoch))