{"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":"# CurricularFace\nCurricularFace was modified from arcface to achive the good result and rubost than arcface. \nyou can read paper [CurricularFace](https://arxiv.org/pdf/2004.00288.pdf)\n\nMy code include: \n\n**Algorithm**: arcface, CurricularFace\n\n**Augmentation**: Mixup,Cutmix\n\n**Another technique** SAM optimizer, adamp,..\n\n**training stratify**: small size (224) to big size ( 512)\n\n**learning finder**: on-going\n","metadata":{}},{"cell_type":"code","source":"!pip install adamp\n!git clone https://github.com/tks1998/pytorch-lr-finder.git\n%cd pytorch-lr-finder \n!pip install .","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-03-24T11:08:34.974799Z","iopub.execute_input":"2022-03-24T11:08:34.975084Z","iopub.status.idle":"2022-03-24T11:08:52.890579Z","shell.execute_reply.started":"2022-03-24T11:08:34.975053Z","shell.execute_reply":"2022-03-24T11:08:52.889725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cd ..","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:08:52.892717Z","iopub.execute_input":"2022-03-24T11:08:52.893004Z","iopub.status.idle":"2022-03-24T11:08:52.902692Z","shell.execute_reply.started":"2022-03-24T11:08:52.892966Z","shell.execute_reply":"2022-03-24T11:08:52.901937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ==================================================== Library ==================================================== #\n# sys\nimport sys\nimport os\nimport time\nimport random\nfrom collections import defaultdict, Counter\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport glob\n\n# dataloader\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nimport torch.utils.data as Data\nfrom sklearn.model_selection import train_test_split\nfrom torchvision.datasets.folder import *\nfrom torch.nn.parameter import Parameter\nfrom torchvision import datasets, transforms\nfrom torch.utils.data import DataLoader, Dataset\n\nimport cv2\nfrom PIL import Image\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n### import algorithm\nimport math\nfrom torch.optim import Adam, SGD\nimport torchvision.models as models\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\n\n### visulization\nimport tensorflow as tf\nfrom matplotlib import pyplot as plt\ntorch.manual_seed(1234)\n### timm module\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master/')\nimport timm\nfrom adamp import AdamP\nfrom torch_lr_finder import LRFinder\nfrom torch.cuda import amp\nfrom torch.autograd import Variable\nfrom torch.cuda.amp import autocast, GradScaler","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:08:52.905691Z","iopub.execute_input":"2022-03-24T11:08:52.905888Z","iopub.status.idle":"2022-03-24T11:08:52.916742Z","shell.execute_reply.started":"2022-03-24T11:08:52.905865Z","shell.execute_reply":"2022-03-24T11:08:52.915948Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def set_seed(seed=42):\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.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nset_seed(1213)","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:08:52.919176Z","iopub.execute_input":"2022-03-24T11:08:52.919612Z","iopub.status.idle":"2022-03-24T11:08:52.928266Z","shell.execute_reply.started":"2022-03-24T11:08:52.919572Z","shell.execute_reply":"2022-03-24T11:08:52.927448Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CFG:\n    # path\n    \n    path_info = '../input/happywhale-cropped-dataset-yolov5-ds/train.csv'\n    root_data_2022 = '../input/happywhale-cropped-dataset-yolov5-ds/train_images/'\n    \n    # dataloader\n    input_size = [224]\n    batch_size = 128\n    worker = 4\n    phase_idx = 0\n\n    # model\n    use_pretrained = True\n    model_name = \"eca_nfnet_l0\"\n    num_classes = 15587\n    embedding_size = 512\n    adaptive = False\n    dropout = 0.5\n    \n    ### metric\n    metric = 'CurricularFace' # 'arcface', 'cosface', 'adacos','adaptive_arcface'\n    use_fc = True\n    s = 30\n    margin = 0.5\n    \n    ### losss \n    MIXUP = True\n    CUTMIX = False\n    ls_eps = 0.0\n    theta_zero = 0.785\n    MIXUP_PROB = 0.5\n    MIXUP_ALPHA = 1.0\n    ### learning rate finder \n    skip_start = 10\n    skip_end = 5\n\n    # optimizer\n    SAM = False\n    lr = 1e-4\n    min_lr = 1e-6\n    weight_decay = 3e-5\n    scheduler = 'ReduceLROnPlateau'\n    epochs = [120]\n    T_max = 6\n    gradient_accumulation_steps = 1\n    max_grad_norm = 1000\n    T_MULT = 2\n    T_0 = 5\n    gamma = 1.5\n    \n    # other parameter\n    seed = 42\n    print_freq = 100\n    early_stop = 10\n    train_continue = False\n    freq_save = 10\n    valid_every = 5\n    loss_name = \"HardMining\"  # or focal loss\n    use_amp = False\n    \n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n    tensor_log = \"./log/product_\"      + model_name + \"_\" + metric + \"_\" + loss_name + \"_\" + scheduler\n    path_save = \"./out_model/product_\" + model_name + \"_\" + metric + \"_\" + loss_name + \"_\" + scheduler +\".pt\"\n\n    # make dir\n    os.makedirs('log', exist_ok=True)\n    os.makedirs('out_model', exist_ok=True)\n","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:08:52.929347Z","iopub.execute_input":"2022-03-24T11:08:52.930011Z","iopub.status.idle":"2022-03-24T11:08:52.942629Z","shell.execute_reply.started":"2022-03-24T11:08:52.929956Z","shell.execute_reply":"2022-03-24T11:08:52.941733Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class min_edge_crop(A.ImageOnlyTransform):\n    \n    def __init__(self, p: float = 0.5, always_apply=True):\n        super().__init__(always_apply, p)\n        self.position = 'center'\n\n    def apply(self, img, **params):\n        assert self.position in ['center', 'left', 'right'], \"position must either be: left, center or right\"\n\n        h, w = img.shape[:2]\n\n        if h == w:\n            return img\n\n        min_edge = min(h, w)\n        if h > min_edge:\n            if self.position == \"left\":\n                img = img[:min_edge]\n            elif self.position == \"center\":\n                d = (h - min_edge) // 2\n                img = img[d:-d] if d != 0 else img\n\n                if h % 2 != 0:\n                    img = img[1:]\n            else:\n                img = img[-min_edge:]\n\n        if w > min_edge:\n            if self.position == \"left\":\n                img = img[:, :min_edge]\n            elif self.position == \"center\":\n                d = (w - min_edge) // 2\n                img = img[:, d:-d] if d != 0 else img\n\n                if w % 2 != 0:\n                    img = img[:, 1:]\n            else:\n                img = img[:, -min_edge:]\n        return img\ndef get_augmentation(phase,input_size):\n    if phase == \"train\":\n        return  A.Compose([\n#                     min_edge_crop(),\n                    A.Resize(height=input_size, width=input_size),\n                     A.ToGray(p=0.01),\n                     A.OneOf([\n                       A.GaussNoise(var_limit=[10, 50]),\n                       A.GaussianBlur(),\n                       A.MotionBlur(),\n                       A.MedianBlur(),\n                      ], p=0.2),\n                    A.OneOf([\n                       A.OpticalDistortion(distort_limit=1.0),\n                       A.GridDistortion(num_steps=5, distort_limit=1.),\n                       A.ElasticTransform(alpha=3),\n                   ], p=0.2),\n                     A.OneOf([\n                         A.CLAHE(),\n                         A.RandomBrightnessContrast(),\n                     ], p=0.25),\n                     A.HueSaturationValue(p=0.25),\n                    A.ShiftScaleRotate(p=0.5, shift_limit=0.0625, scale_limit=0.2, rotate_limit=20),\n#                     A.Cutout(max_h_size=int(input_size * 0.1), max_w_size=int(input_size * 0.1), num_holes=5, p=0.5),\n                    A.Normalize(),\n                    ToTensorV2()\n                ])\n    elif phase in ['test','valid']:\n        return A.Compose([\n#             min_edge_crop(),\n            A.Resize(height=input_size, width=input_size),\n            A.Normalize(),\n            ToTensorV2()\n        ])\n","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:08:52.943830Z","iopub.execute_input":"2022-03-24T11:08:52.944097Z","iopub.status.idle":"2022-03-24T11:08:52.963003Z","shell.execute_reply.started":"2022-03-24T11:08:52.944060Z","shell.execute_reply":"2022-03-24T11:08:52.962205Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class whale_huback():\n    \n    def __init__(self, df, transform = None):\n        \n        self.df = df.reset_index()\n        self.transform = transform\n                \n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self,index):\n        \n        img_path, class_id = self.df.loc[index, 'image_path'], self.df.loc[index,'individual_id']\n        sample = cv2.imread(img_path)\n        if sample is None:\n            print(img_path)\n            \n        sample = cv2.cvtColor(sample, cv2.COLOR_BGR2RGB)\n        if self.transform is not None:\n            sample = self.transform(image=sample)[\"image\"]\n\n        return sample, torch.tensor(class_id)","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:08:52.964487Z","iopub.execute_input":"2022-03-24T11:08:52.964903Z","iopub.status.idle":"2022-03-24T11:08:52.974115Z","shell.execute_reply.started":"2022-03-24T11:08:52.964864Z","shell.execute_reply":"2022-03-24T11:08:52.973293Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"info = pd.read_csv(CFG.path_info)\ninfo['image'] = CFG.root_data_2022+info['image']\nprint(len(info['image']))\n\ninfo = info.sample(frac=1).reset_index(drop=True)\n\nmapping = {}\nmapping_invert = {}\nindex=-1\n\nall_class = info['individual_id'].unique()\nnp.random.shuffle(all_class)\nfor class_id in all_class:\n    if class_id not in mapping.keys():\n        index+=1\n        mapping[class_id] = index\n        mapping_invert[index] = class_id\ninfo['individual_id'] = info['individual_id'].apply(lambda class_id: mapping[class_id]) \nbalanced_data = pd.DataFrame([])\nfor class_id in info['individual_id'].unique():\n    sub_df = info[info['individual_id'] == class_id]\n    if sub_df.shape[0] < 2:\n        sub_df_repeat = pd.concat([sub_df]*3)\n        balanced_data = pd.concat([balanced_data,sub_df_repeat])\ninfo = pd.concat([balanced_data,info])\ntrain,valid = train_test_split(info, random_state=42, stratify=info['individual_id'], train_size=0.7)\ntrain.to_csv(\"train_arcface.csv\")\nvalid.to_csv(\"valid_arcface.csv\")","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:08:52.975551Z","iopub.execute_input":"2022-03-24T11:08:52.975982Z","iopub.status.idle":"2022-03-24T11:09:33.138996Z","shell.execute_reply.started":"2022-03-24T11:08:52.975945Z","shell.execute_reply":"2022-03-24T11:09:33.138111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\n\nclass SAM(torch.optim.Optimizer):\n    def __init__(self, params, base_optimizer, rho=0.05, **kwargs):\n        assert rho >= 0.0, f\"Invalid rho, should be non-negative: {rho}\"\n\n        defaults = dict(rho=rho, **kwargs)\n        super(SAM, self).__init__(params, defaults)\n\n        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)\n        self.param_groups = self.base_optimizer.param_groups\n\n    @torch.no_grad()\n    def first_step(self, zero_grad=False):\n        grad_norm = self._grad_norm()\n        for group in self.param_groups:\n            scale = group[\"rho\"] / (grad_norm + 1e-12)\n\n            for p in group[\"params\"]:\n                if p.grad is None: continue\n                e_w = p.grad * scale.to(p)\n                p.add_(e_w)  # climb to the local maximum \"w + e(w)\"\n                self.state[p][\"e_w\"] = e_w\n\n        if zero_grad: self.zero_grad()\n\n    @torch.no_grad()\n    def second_step(self, zero_grad=False):\n        for group in self.param_groups:\n            for p in group[\"params\"]:\n                if p.grad is None: continue\n                p.sub_(self.state[p][\"e_w\"])  # get back to \"w\" from \"w + e(w)\"\n\n        self.base_optimizer.step()  # do the actual \"sharpness-aware\" update\n\n        if zero_grad: self.zero_grad()\n\n    @torch.no_grad()\n    def step(self, closure=None):\n        assert closure is not None, \"Sharpness Aware Minimization requires closure, but it was not provided\"\n        closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass\n\n        self.first_step(zero_grad=True)\n        closure()\n        self.second_step()\n\n    def _grad_norm(self):\n        shared_device = self.param_groups[0][\"params\"][0].device  # put everything on the same device, in case of model parallelism\n        norm = torch.norm(\n                    torch.stack([\n                        p.grad.norm(p=2).to(shared_device)\n                        for group in self.param_groups for p in group[\"params\"]\n                        if p.grad is not None\n                    ]),\n                    p=2\n               )\n        return norm","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:09:33.143317Z","iopub.execute_input":"2022-03-24T11:09:33.143526Z","iopub.status.idle":"2022-03-24T11:09:33.156035Z","shell.execute_reply.started":"2022-03-24T11:09:33.143499Z","shell.execute_reply":"2022-03-24T11:09:33.155251Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class whale_huback():\n    \n    def __init__(self, df, transform = None):\n        \n        self.df = df.reset_index()\n        self.transform = transform\n                \n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self,index):\n        \n        img_path, class_id = self.df.loc[index, 'image_path'], self.df.loc[index,'individual_id']\n        sample = cv2.imread(img_path)\n        if sample is None:\n            print(img_path)\n            \n        sample = cv2.cvtColor(sample, cv2.COLOR_BGR2RGB)\n        if self.transform is not None:\n            sample = self.transform(image=sample)[\"image\"]\n\n        return sample, torch.tensor(class_id)\n    \n\n\n# In[7]:\n\n\nimport torch\nfrom torch import nn\nfrom tqdm import tqdm\n\nclass FocalLoss(nn.Module):\n\n    def __init__(self, gamma=0, eps=1e-7):\n        super(FocalLoss, self).__init__()\n        self.gamma = gamma\n        self.eps = eps\n        self.ce = torch.nn.CrossEntropyLoss()\n\n    def forward(self, input, target):\n        logp = self.ce(input, target)\n        p = torch.exp(-logp)\n        loss = (1 - p) ** self.gamma * logp\n        return loss.mean()\nclass HardMining(nn.Module):\n    def __init__(self, save_rate = 2):\n        super(HardMining, self).__init__()\n        self.save_rate = save_rate\n        self.ce = nn.CrossEntropyLoss()\n    \n    def forward(self, input, target):\n        batch_size = input.shape[0]\n        loss = self.ce(input, target)\n        ind_sorted = torch.argsort(-loss) # from big to small\n        num_saved = int(self.save_rate * batch_size)\n        ind_update = ind_sorted[:num_saved]\n        loss_final = torch.sum(self.ce(input[ind_update], target[ind_update]))\n        return loss_final\nclass AverageMeter(object):\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\ndef train_fn(dataloader, model, criterion, optimizer, epoch, CFG):\n    model.train()\n    loss_score = AverageMeter()\n\n    tk0 = tqdm(enumerate(dataloader), total=len(dataloader))\n    for bi, d in tk0:\n        batch_size = d[0].shape[0]\n\n        images = d[0]\n        targets = d[1]\n        \n        images = images.to(CFG.device)\n        targets = targets.to(CFG.device)\n        \n        \n        if CFG.MIXUP:\n            images, label_a, label_b, lam = mixup_data(images, targets, CFG.device, CFG.MIXUP_PROB, CFG.MIXUP_ALPHA)\n            images, label_a, label_b = map(Variable, (images, label_a, label_b))\n        elif CFG.CUTMIX:\n            images, label_a, label_b, lam = cutmix_data(images, targets, CFG.device, CFG.MIXUP_PROB, CFG.MIXUP_ALPHA)\n            images, label_a, label_b, lam = map(Variable, (inputs, label_a, label_b))\n\n        with amp.autocast(enabled=True):\n            output = model(images, targets)\n            if CFG.MIXUP or CFG.CUTMIX:\n                loss = mixup_criterion(criterion, output, label_a, label_b, lam)\n            else:\n                loss = criterion(output, targets)\n\n        optimizer.zero_grad()\n\n        if CFG.adaptive:\n            loss +=CFG.lambda_*torch.mean(model.final.m)\n        \n        loss.backward()\n        optimizer.step()\n        loss_score.update(loss.detach().item(), batch_size)\n        tk0.set_postfix(Train_Loss=loss_score.avg, Epoch=epoch, LR=optimizer.param_groups[0]['lr'])\n    return loss_score\n\n\ndef eval_fn(data_loader, model, criterion, CFG):\n    loss_score = AverageMeter()\n\n    model.eval()\n    tk0 = tqdm(enumerate(data_loader), total=len(data_loader))\n\n    with torch.no_grad():\n        for bi, d in tk0:\n            batch_size = d[0].size()[0]\n\n            image = d[0]\n            targets = d[1]\n\n            image = image.to(CFG.device)\n            targets = targets.to(CFG.device)\n\n            output = model(image, targets)\n\n            loss = criterion(output, targets)\n\n            loss_score.update(loss.detach().item(), batch_size)\n            tk0.set_postfix(Eval_Loss=loss_score.avg)\n\n    return loss_score\n\n","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:09:33.158656Z","iopub.execute_input":"2022-03-24T11:09:33.159103Z","iopub.status.idle":"2022-03-24T11:09:33.186022Z","shell.execute_reply.started":"2022-03-24T11:09:33.159069Z","shell.execute_reply":"2022-03-24T11:09:33.185275Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"### model\ndef l2_norm(input, axis = 1):\n    norm = torch.norm(input, 2, axis, True)\n    output = torch.div(input, norm)\n\n    return output\n\nclass CurricularFace(nn.Module):\n    r\"\"\"Implement of CurricularFace (https://arxiv.org/pdf/2004.00288.pdf):\n    Args:\n        in_features: size of each input sample\n        out_features: size of each output sample\n        device_id: the ID of GPU where the model will be trained by model parallel. \n                       if device_id=None, it will be trained on CPU without model parallel.\n        m: margin\n        s: scale of outputs\n    \"\"\"\n    def __init__(self, in_features, out_features, m = 0.5, s = 64.):\n        super(CurricularFace, self).__init__()\n        self.in_features = in_features\n        self.out_features = out_features\n        self.m = torch.tensor(m).to(CFG.device)\n        self.s = torch.tensor(s).to(CFG.device)\n        self.cos_m = torch.cos(self.m)\n        self.sin_m = torch.sin(self.m)\n        torch.pi = torch.acos(torch.zeros(1))* 2 \n        torch.pi = torch.pi.to(CFG.device)\n        self.threshold = torch.cos(torch.pi - self.m)\n        self.mm = torch.sin(torch.pi - self.m) * self.m\n        self.kernel = Parameter(torch.Tensor(in_features, out_features))\n        torch_t = torch.zeros(1).float() \n        self.register_buffer('t', torch_t.to(CFG.device))\n        nn.init.normal_(self.kernel, std=0.01)\n        \n    @autocast()\n    def forward(self, embbedings, label):\n        embbedings = l2_norm(embbedings, axis = 1)\n        embbedings = embbedings.to(CFG.device)\n        kernel_norm = l2_norm(self.kernel, axis = 0)\n        kernel_norm = kernel_norm.to(CFG.device)\n        cos_theta = torch.mm(embbedings, kernel_norm)\n        cos_theta = cos_theta.clamp(-1, 1)  # for numerical stability\n        cos_theta = cos_theta.to(CFG.device)\n        cos_theta = cos_theta.float()\n        with torch.no_grad():\n            origin_cos = cos_theta.clone()\n        target_logit = cos_theta[torch.arange(0, embbedings.size(0)), label].view(-1, 1)\n\n        sin_theta = torch.sqrt(1.0 - torch.pow(target_logit, 2))\n        cos_theta_m = target_logit * self.cos_m - sin_theta * self.sin_m #cos(target+margin)\n        mask = cos_theta > cos_theta_m\n        \n        final_target_logit = torch.where(target_logit > self.threshold, cos_theta_m, target_logit - self.mm)\n        hard_example = cos_theta[mask]\n#         hard_example = torch.tensor(hard_example)\n        hard_example = hard_example.to(CFG.device)\n        hard_example = hard_example.float()\n        with torch.no_grad():\n            self.t = target_logit.mean() * 0.01 + (1 - 0.01) * self.t\n        cos_theta[mask] = hard_example * (self.t + hard_example)\n        cos_theta.scatter_(1, label.view(-1, 1).long(), final_target_logit)\n        output = cos_theta * self.s\n        return output\n\nclass ArcMarginProduct(nn.Module):\n    r\"\"\"Implement of large margin arc distance: :\n        Args:\n            in_features: size of each input sample\n            out_features: size of each output sample\n            s: norm of input feature\n            m: margin\n            cos(theta + m)\n        \"\"\"\n    def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False, ls_eps=0.0):\n        super(ArcMarginProduct, self).__init__()\n        self.in_features = in_features\n        self.out_features = out_features\n        self.s = torch.tensor(s).to(CFG.device)\n        self.m = torch.tensor(m).to(CFG.device)\n        self.ls_eps = ls_eps  # label smoothing\n        self.weight = Parameter(torch.FloatTensor(out_features, in_features)).to(CFG.device)\n        nn.init.xavier_uniform_(self.weight)\n\n        self.easy_margin = easy_margin\n        self.cos_m = torch.cos(self.m).to(CFG.device)\n        self.sin_m = torch.sin(self.m).to(CFG.device)\n        torch.pi = torch.acos(torch.zeros(1))* 2 \n        torch.pi = torch.pi.to(CFG.device)\n        self.th = torch.cos(torch.pi - self.m)\n        self.mm = torch.sin(torch.pi - self.m) * self.m\n#         self.th = self.th.to(CFG.device)\n#         self.mm = self.mm.to(CFG.device)\n    @autocast()\n    def forward(self, input, label):\n        # --------------------------- cos(theta) & phi(theta) ---------------------------\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n        cosine = cosine.to(CFG.device)\n        sine = torch.sqrt(1.0 - torch.pow(cosine, 2)).to(CFG.device)\n        phi = cosine * self.cos_m - sine * self.sin_m\n        if self.easy_margin:\n            phi = torch.where(cosine > 0, phi, cosine)\n        else:\n            phi = torch.where(cosine > self.th, phi, cosine - self.mm)\n        # --------------------------- convert label to one-hot ---------------------------\n        # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda')\n        one_hot = torch.zeros(cosine.size(), device=CFG.device)\n        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n        if self.ls_eps > 0:\n            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features\n        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n        output *= self.s\n\n        return output\n    \nclass Model(nn.Module):\n\n    def __init__(self, n_classes, model_name='efficientnet_b3', use_fc=False, fc_dim=512,                 dropout=0.0, metric='softmax', s=30.0, margin=0.50, ls_eps=0.0,                 theta_zero=0.785, pretrained=False):\n        \"\"\"\n        :param n_classes:\n        :param model_name: name of model from pretrainedmodels\n            e.g. resnet50, resnext101_32x4d, pnasnet5large\n        :param pooling: One of ('SPoC', 'MAC', 'RMAC', 'GeM', 'Rpool', 'Flatten', 'CompactBilinearPooling')\n        :param metric: One of ('arcface', 'cosface', 'softmax')\n        \"\"\"\n        super(Model, self).__init__()\n        print('Building Model Backbone for {} model'.format(model_name))\n\n        self.backbone = timm.create_model(model_name, pretrained=True)\n#         self.backbone.load_state_dict(torch.load(CFG.orginal_pretrain_model, map_location=CFG.device))\n#         if 'coat' in CFG.model_name:\n#             final_in_features = self.backbone.head.in_features\n#         else:\n#             final_in_features = self.backbone.classifier.in_features\n#             self.backbone.classifier = nn.Identity()\n#             self.backbone.global_pool = nn.Identity()\n\n        if 'efficientnet' in CFG.model_name:\n            final_in_features = self.backbone.classifier.in_features\n            self.backbone.classifier = nn.Identity()\n            self.backbone.global_pool = nn.Identity()\n        \n        elif 'nfnet' in CFG.model_name:\n            final_in_features = self.backbone.head.fc.in_features\n            self.backbone.head.fc = nn.Identity()\n            self.backbone.head.global_pool = nn.Identity()\n            \n        self.backbone.head = nn.Identity()\n        self.pooling = nn.AdaptiveAvgPool2d(1)\n\n        self.use_fc = use_fc\n        if use_fc:\n            self.dropout = nn.Dropout(p=dropout)\n            self.fc = nn.Linear(final_in_features, fc_dim)\n            self.bn = nn.BatchNorm1d(fc_dim)\n            self._init_params()\n            final_in_features = fc_dim\n\n        self.metric = metric\n        if metric == 'arcface':\n            self.final = ArcMarginProduct(final_in_features, n_classes,\n                                          s=s, m=margin, easy_margin=False, ls_eps=ls_eps)\n        elif metric == 'cosface':\n            self.final = AddMarginProduct(final_in_features, n_classes, s=s, m=margin, device=CFG.device)\n        elif metric == 'adacos':\n            self.final = AdaCos(final_in_features, n_classes, m=margin, theta_zero=theta_zero)\n        elif metric == 'adaptive_arcface':\n            self.final = Arcface_adaptive_margin(final_in_features, n_classes, s=s, m=margin,                                                 device=CFG.device, easy_margin=False, ls_eps=ls_eps)\n        else:\n            self.final = nn.Linear(final_in_features, n_classes)\n\n    def _init_params(self):\n        nn.init.xavier_normal_(self.fc.weight)\n        nn.init.constant_(self.fc.bias, 0)\n        nn.init.constant_(self.bn.weight, 1)\n        nn.init.constant_(self.bn.bias, 0)\n\n    def forward(self, x, label):\n        feature = self.extract_feat(x)\n        if self.metric in ('arcface', 'cosface', 'adacos','adaptive_arcface'):\n            logits = self.final(feature, label)\n        else:\n            logits = self.final(feature)\n        return logits\n\n    def extract_feat(self, x):\n        batch_size = x.shape[0]\n        x = self.backbone(x)\n        x = self.pooling(x).view(batch_size, -1)\n\n        if self.use_fc:\n            x = self.dropout(x)\n            x = self.fc(x)\n            x = self.bn(x)\n\n        return x\n\n","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:09:33.189008Z","iopub.execute_input":"2022-03-24T11:09:33.189459Z","iopub.status.idle":"2022-03-24T11:09:33.230803Z","shell.execute_reply.started":"2022-03-24T11:09:33.189423Z","shell.execute_reply":"2022-03-24T11:09:33.230038Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:09:33.232219Z","iopub.execute_input":"2022-03-24T11:09:33.232995Z","iopub.status.idle":"2022-03-24T11:09:33.248781Z","shell.execute_reply.started":"2022-03-24T11:09:33.232958Z","shell.execute_reply":"2022-03-24T11:09:33.248041Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\ndef get_scheduler(CFG,optimizer):\n    if CFG.scheduler=='CosineAnnealingLR':\n        scheduler = CosineAnnealingLR(optimizer, T_max=CFG.T_max, eta_min=CFG.min_lr, last_epoch=-1)\n    elif CFG.scheduler=='CosineAnnealingWarmRestarts':\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n                optimizer=optimizer,\n                T_0=len(dataloader['train']) * CFG.T_0 * CFG.T_MULT,\n                T_mult=CFG.T_MULT,\n                eta_min=CFG.min_lr,\n        )\n    elif CFG.scheduler== 'ReduceLROnPlateau':\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer=optimizer,\n            min_lr=CFG.min_lr\n        )\n    return scheduler\n\n\n# In[10]:\n\n\ndef get_model(CFG):\n    model = Model(**{\n                'n_classes':CFG.num_classes,\n                'model_name':CFG.model_name,\n                'use_fc':CFG.use_fc,\n                'fc_dim':CFG.embedding_size,\n                'dropout':CFG.dropout,\n                'metric':CFG.metric,\n                's':CFG.s,\n                'margin':CFG.margin,\n                'ls_eps':CFG.ls_eps,\n                'theta_zero':CFG.theta_zero,\n                'pretrained':False})\n    return model\n\n\n# In[11]:\n\n\ndef mixup_criterion(criterion, pred, y_a, y_b, lam):\n    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)\n\n\n# In[12]:\n\n\ndef mixup_data(x, y, gpu, mixup_prob=0.5, alpha=1.0):\n    '''Returns mixed inputs, pairs of targets, and lambda'''\n    if random.uniform(0, 1) < mixup_prob:\n        lam = np.random.beta(alpha, alpha)\n    else:\n        lam = 1\n\n    batch_size = x.size()[0]\n    index = torch.randperm(batch_size).cuda(gpu)\n\n    mixed_x = lam * x + (1 - lam) * x[index, :]\n    y_a, y_b = y, y[index]\n    return mixed_x, y_a, y_b, lam\n\n","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:09:33.251752Z","iopub.execute_input":"2022-03-24T11:09:33.251947Z","iopub.status.idle":"2022-03-24T11:09:33.265051Z","shell.execute_reply.started":"2022-03-24T11:09:33.251923Z","shell.execute_reply":"2022-03-24T11:09:33.264381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for index, input_size in enumerate(CFG.input_size):\n    print(\"training size {}\",format(input_size))\n    dataset = {\n        phase: whale_huback(eval(phase),transform = get_augmentation(phase=phase,input_size=input_size)) \\\n        for phase in ['train','valid']\n    }\n\n    dataloader = { phase: Data.DataLoader(dataset=dataset[phase], num_workers=(len(CFG.input_size)-index)*CFG.worker, batch_size=CFG.batch_size,                                           shuffle=(phase=='train'),pin_memory = (phase=='train'), drop_last=True)                   for phase in ['train','valid'] }\n\n    model = get_model(CFG)\n\n    model.to(CFG.device)\n\n    criterion = FocalLoss(gamma=1.5) if CFG.loss_name == \"focal_loss\" else nn.CrossEntropyLoss()\n    optimizer = optim.Adam(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay)\n\n    scheduler = get_scheduler(CFG,optimizer)\n    \n    best_loss = np.inf\n    if CFG.train_continue==True or index == 1:\n        print(\"continue train\")\n        path_previous_save = \"./out_model/product_\" + CFG.model_name + \"_\" + CFG.metric + \"_\" + CFG.loss_name + \"_\" + CFG.scheduler + str(CFG.input_size[index-1]) +\".pt\"\n        checkpoint = torch.load(path_previous_save)\n        model.load_state_dict(checkpoint['model_state_dict']) \n        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n        criterion = checkpoint['criterion']\n        scheduler = checkpoint['scheduler']\n        best_loss = checkpoint['best_loss']\n    \n    \n    visualization = tf.summary.create_file_writer(CFG.tensor_log)\n\n    with visualization.as_default():\n        \n        for epoch in range(0, CFG.epochs[index]+1):\n            for phase in ['train','valid']:\n                if phase == 'train':\n                    loss = train_fn(dataloader['train'], model,criterion, optimizer, epoch=epoch,CFG=CFG)\n                else:\n                    loss = eval_fn(dataloader['valid'], model, criterion, CFG=CFG)\n                tf.summary.scalar('./loss/{}/{}'.format(input_size,phase), loss.avg, step=epoch)\n                print('epoch {} loss {}'.format(epoch,loss.avg)) \n                if phase == 'valid':\n                    path_save = \"./out_model/product_\" + CFG.model_name + \"_\" + CFG.metric + \"_\" + CFG.loss_name + \"_\" + CFG.scheduler + str(input_size) +\".pt\"\n                    if best_loss > loss.avg:\n                        best_loss = loss.avg\n                        torch.save({\n                            'model_state_dict': model.state_dict(),\n                            'optimizer_state_dict': optimizer.state_dict(),\n                            'criterion': criterion,\n                            'scheduler': scheduler,\n                            'best_loss':best_loss\n                        }, CFG.path_save)\n                        not_improve=0\n                    else:\n                        not_improve+=1\n            scheduler.step(loss.item())\n            if not_improve == CFG.early_stop:\n                print(\"early_stop at epoch \", epoch)\n                break\n    ### del  \n    del model, dataset, dataloader, criterion, scheduler, loss\n\n","metadata":{"execution":{"iopub.status.busy":"2022-03-24T11:09:33.266355Z","iopub.execute_input":"2022-03-24T11:09:33.267443Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}