Upload model

config.json

@@ -2,7 +2,11 @@
   "architectures": [
     "CetaceanClassifierModelForImageClassification"
   ],
-  "model_type": "efficientnet",
+  "auto_map": {
+    "AutoConfig": "configuration_cetacean_classifier.CetaceanClassifierConfig",
+    "AutoModelForImageClassification": "modeling_cetacean_classifier.CetaceanClassifierModelForImageClassification"
+  },
+  "model_type": "cetaceanet",
   "torch_dtype": "float32",
   "transformers_version": "4.46.0"
 }

config.py

@@ -0,0 +1,28 @@
+from typing import Optional
+
+import yaml
+
+
+class Config(dict):
+    def __getattr__(self, key):
+        try:
+            val = self[key]
+        except KeyError:
+            return super().__getattr__(key)
+        if isinstance(val, dict):
+            return Config(val)
+        return val
+
+
+def load_config(path: str, default_path: Optional[str]) -> Config:
+    with open(path) as f:
+        cfg = Config(yaml.full_load(f))
+    if default_path is not None:
+        # set keys not included in `path` by default
+        with open(default_path) as f:
+            default_cfg = Config(yaml.full_load(f))
+        for key, val in default_cfg.items():
+            if key not in cfg:
+                print(f"used default config {key}: {val}")
+                cfg[key] = val
+    return cfg

configuration_cetacean_classifier.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+from typing import List
+
+
+class CetaceanClassifierConfig(PretrainedConfig):
+    model_type = "cetaceanet"
+
+    def __init__(
+        self,
+        # block_type="bottleneck",
+        # layers: List[int] = [3, 4, 6, 3],
+        # num_classes: int = 1000,
+        # input_channels: int = 3,
+        # cardinality: int = 1,
+        # base_width: int = 64,
+        # stem_width: int = 64,
+        # stem_type: str = "",
+        # avg_down: bool = False,
+        **kwargs,
+    ):
+        # if block_type not in ["basic", "bottleneck"]:
+        #     raise ValueError(f"`block_type` must be 'basic' or bottleneck', got {block_type}.")
+        # if stem_type not in ["", "deep", "deep-tiered"]:
+        #     raise ValueError(f"`stem_type` must be '', 'deep' or 'deep-tiered', got {stem_type}.")
+
+        # self.block_type = block_type
+        # self.layers = layers
+        # self.num_classes = num_classes
+        # self.input_channels = input_channels
+        # self.cardinality = cardinality
+        # self.base_width = base_width
+        # self.stem_width = stem_width
+        # self.stem_type = stem_type
+        # self.avg_down = avg_down
+        super().__init__(**kwargs)

metric_learning.py

@@ -0,0 +1,59 @@
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class GeM(nn.Module):
+    def __init__(self, p=3, eps=1e-6, requires_grad=False):
+        super().__init__()
+        self.p = nn.Parameter(torch.ones(1) * p, requires_grad=requires_grad)
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor):
+        return x.clamp(min=self.eps).pow(self.p).mean((-2, -1)).pow(1.0 / self.p)
+
+
+# Copied and modified from
+# https://github.com/ChristofHenkel/kaggle-landmark-2021-1st-place/blob/034a7d8665bb4696981698348c9370f2d4e61e35/models/ch_mdl_dolg_efficientnet.py
+class ArcMarginProductSubcenter(nn.Module):
+    def __init__(self, in_features: int, out_features: int, k: int = 3):
+        super().__init__()
+        self.weight = nn.Parameter(torch.FloatTensor(out_features * k, in_features))
+        self.reset_parameters()
+        self.k = k
+        self.out_features = out_features
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.weight.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        cosine_all = F.linear(F.normalize(features), F.normalize(self.weight))
+        cosine_all = cosine_all.view(-1, self.out_features, self.k)
+        cosine, _ = torch.max(cosine_all, dim=2)
+        return cosine
+
+
+class ArcFaceLossAdaptiveMargin(nn.modules.Module):
+    def __init__(self, margins: np.ndarray, n_classes: int, s: float = 30.0):
+        super().__init__()
+        self.s = s
+        self.margins = margins
+        self.out_dim = n_classes
+
+    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        ms = self.margins[labels.cpu().numpy()]
+        cos_m = torch.from_numpy(np.cos(ms)).float().cuda()
+        sin_m = torch.from_numpy(np.sin(ms)).float().cuda()
+        th = torch.from_numpy(np.cos(math.pi - ms)).float().cuda()
+        mm = torch.from_numpy(np.sin(math.pi - ms) * ms).float().cuda()
+        labels = F.one_hot(labels, self.out_dim).float()
+        logits = logits.float()
+        cosine = logits
+        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
+        phi = cosine * cos_m.view(-1, 1) - sine * sin_m.view(-1, 1)
+        phi = torch.where(cosine > th.view(-1, 1), phi, cosine - mm.view(-1, 1))
+        return ((labels * phi) + ((1.0 - labels) * cosine)) * self.s

modeling_cetacean_classifier.py

@@ -0,0 +1,66 @@
+from transformers import PreTrainedModel
+from PIL import Image
+import numpy as np
+import torch
+
+from .configuration_cetacean_classifier import CetaceanClassifierConfig
+from .train import SphereClassifier
+
+
+WHALE_CLASSES = np.array(
+    [
+        "beluga",
+        "blue_whale",
+        "bottlenose_dolphin",
+        "brydes_whale",
+        "commersons_dolphin",
+        "common_dolphin",
+        "cuviers_beaked_whale",
+        "dusky_dolphin",
+        "false_killer_whale",
+        "fin_whale",
+        "frasiers_dolphin",
+        "gray_whale",
+        "humpback_whale",
+        "killer_whale",
+        "long_finned_pilot_whale",
+        "melon_headed_whale",
+        "minke_whale",
+        "pantropic_spotted_dolphin",
+        "pygmy_killer_whale",
+        "rough_toothed_dolphin",
+        "sei_whale",
+        "short_finned_pilot_whale",
+        "southern_right_whale",
+        "spinner_dolphin",
+        "spotted_dolphin",
+        "white_sided_dolphin",
+    ]
+)
+
+
+class CetaceanClassifierModelForImageClassification(PreTrainedModel):
+    config_class = CetaceanClassifierConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = SphereClassifier.load_from_checkpoint("cetacean_classifier/last.ckpt")
+        self.model.eval()
+
+    def preprocess_image(self, img: Image) -> torch.Tensor:
+        image_resized = img.resize((480, 480))
+        image_resized = np.array(image_resized)[None]
+        image_resized = np.transpose(image_resized, [0, 3, 2, 1])
+        image_tensor = torch.Tensor(image_resized)
+        return image_tensor
+
+    def forward(self, img: Image, labels=None):
+        tensor = self.preprocess_image(img)
+        head_id_logits, head_species_logits = self.model(tensor)
+        head_species_logits = head_species_logits.detach().numpy()
+        sorted_idx = head_species_logits.argsort()[0]
+        sorted_idx = np.array(list(reversed(sorted_idx)))
+        top_three_logits = sorted_idx[:3]
+        top_three_whale_preds = WHALE_CLASSES[top_three_logits]
+
+        return {"predictions": top_three_whale_preds}

train.py

@@ -0,0 +1,318 @@
+# import argparse
+# import os
+# import warnings
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+# import optuna
+# import pandas as pd
+import timm
+import torch
+# import wandb
+# from optuna.integration import PyTorchLightningPruningCallback
+from pytorch_lightning import LightningDataModule, LightningModule, Trainer
+# from pytorch_lightning import loggers as pl_loggers
+# from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+# from sklearn.model_selection import StratifiedKFold
+# from torch.utils.data import ConcatDataset, DataLoader
+
+from .config import Config, load_config
+# from .dataset import WhaleDataset, load_df
+from .metric_learning import ArcFaceLossAdaptiveMargin, ArcMarginProductSubcenter, GeM
+from .utils import WarmupCosineLambda, map_dict, topk_average_precision
+
+
+# def parse():
+#     parser = argparse.ArgumentParser(description="Training for HappyWhale")
+#     parser.add_argument("--out_base_dir", default="result")
+#     parser.add_argument("--in_base_dir", default="input")
+#     parser.add_argument("--exp_name", default="tmp")
+#     parser.add_argument("--load_snapshot", action="store_true")
+#     parser.add_argument("--save_checkpoint", action="store_true")
+#     parser.add_argument("--wandb_logger", action="store_true")
+#     parser.add_argument("--config_path", default="config/debug.yaml")
+#     return parser.parse_args()
+
+
+# class WhaleDataModule(LightningDataModule):
+#     def __init__(
+#         self,
+#         df: pd.DataFrame,
+#         cfg: Config,
+#         image_dir: str,
+#         val_bbox_name: str,
+#         fold: int,
+#         additional_dataset: WhaleDataset = None,
+#     ):
+#         super().__init__()
+#         self.cfg = cfg
+#         self.image_dir = image_dir
+#         self.val_bbox_name = val_bbox_name
+#         self.additional_dataset = additional_dataset
+#         if cfg.n_data != -1:
+#             df = df.iloc[: cfg.n_data]
+#         self.all_df = df
+#         if fold == -1:
+#             self.train_df = df
+#         else:
+#             skf = StratifiedKFold(n_splits=cfg.n_splits, shuffle=True, random_state=0)
+#             train_idx, val_idx = list(skf.split(df, df.individual_id))[fold]
+#             self.train_df = df.iloc[train_idx].copy()
+#             self.val_df = df.iloc[val_idx].copy()
+#             # relabel ids not included in training data as "new individual"
+#             new_mask = ~self.val_df.individual_id.isin(self.train_df.individual_id)
+#             self.val_df.individual_id.mask(new_mask, cfg.num_classes, inplace=True)
+#             print(f"new: {(self.val_df.individual_id == cfg.num_classes).sum()} / {len(self.val_df)}")
+
+#     def get_dataset(self, df, data_aug):
+#         return WhaleDataset(df, self.cfg, self.image_dir, self.val_bbox_name, data_aug)
+
+#     def train_dataloader(self):
+#         dataset = self.get_dataset(self.train_df, True)
+#         if self.additional_dataset is not None:
+#             dataset = ConcatDataset([dataset, self.additional_dataset])
+#         return DataLoader(
+#             dataset,
+#             batch_size=self.cfg.batch_size,
+#             shuffle=True,
+#             num_workers=2,
+#             pin_memory=True,
+#             drop_last=True,
+#         )
+
+#     def val_dataloader(self):
+#         if self.cfg.n_splits == -1:
+#             return None
+#         return DataLoader(
+#             self.get_dataset(self.val_df, False),
+#             batch_size=self.cfg.batch_size,
+#             shuffle=False,
+#             num_workers=2,
+#             pin_memory=True,
+#         )
+
+#     def all_dataloader(self):
+#         return DataLoader(
+#             self.get_dataset(self.all_df, False),
+#             batch_size=self.cfg.batch_size,
+#             shuffle=False,
+#             num_workers=2,
+#             pin_memory=True,
+#         )
+
+
+class SphereClassifier(LightningModule):
+    def __init__(self, cfg: dict, id_class_nums=None, species_class_nums=None):
+        super().__init__()
+        if not isinstance(cfg, Config):
+            cfg = Config(cfg)
+        self.save_hyperparameters(cfg, ignore=["id_class_nums", "species_class_nums"])
+        self.test_results_fp = None
+
+        print(cfg.model_name)
+
+        # NN architecture
+        self.backbone = timm.create_model(
+            cfg.model_name,
+            in_chans=3,
+            pretrained=cfg.pretrained,
+            num_classes=0,
+            features_only=True,
+            out_indices=cfg.out_indices,
+        )
+        feature_dims = self.backbone.feature_info.channels()
+        print(f"feature dims: {feature_dims}")
+        self.global_pools = torch.nn.ModuleList(
+            [GeM(p=cfg.global_pool.p, requires_grad=cfg.global_pool.train) for _ in cfg.out_indices]
+        )
+        self.mid_features = np.sum(feature_dims)
+        if cfg.normalization == "batchnorm":
+            self.neck = torch.nn.BatchNorm1d(self.mid_features)
+        elif cfg.normalization == "layernorm":
+            self.neck = torch.nn.LayerNorm(self.mid_features)
+        self.head_id = ArcMarginProductSubcenter(self.mid_features, cfg.num_classes, cfg.n_center_id)
+        self.head_species = ArcMarginProductSubcenter(self.mid_features, cfg.num_species_classes, cfg.n_center_species)
+        if id_class_nums is not None and species_class_nums is not None:
+            margins_id = np.power(id_class_nums, cfg.margin_power_id) * cfg.margin_coef_id + cfg.margin_cons_id
+            margins_species = (
+                np.power(species_class_nums, cfg.margin_power_species) * cfg.margin_coef_species
+                + cfg.margin_cons_species
+            )
+            print("margins_id", margins_id)
+            print("margins_species", margins_species)
+            self.margin_fn_id = ArcFaceLossAdaptiveMargin(margins_id, cfg.num_classes, cfg.s_id)
+            self.margin_fn_species = ArcFaceLossAdaptiveMargin(margins_species, cfg.num_species_classes, cfg.s_species)
+            self.loss_fn_id = torch.nn.CrossEntropyLoss()
+            self.loss_fn_species = torch.nn.CrossEntropyLoss()
+
+    def get_feat(self, x: torch.Tensor) -> torch.Tensor:
+        ms = self.backbone(x)
+        h = torch.cat([global_pool(m) for m, global_pool in zip(ms, self.global_pools)], dim=1)
+        return self.neck(h)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        feat = self.get_feat(x)
+        return self.head_id(feat), self.head_species(feat)
+
+    def training_step(self, batch, batch_idx):
+        x, ids, species = batch["image"], batch["label"], batch["label_species"]
+        logits_ids, logits_species = self(x)
+        margin_logits_ids = self.margin_fn_id(logits_ids, ids)
+        loss_ids = self.loss_fn_id(margin_logits_ids, ids)
+        loss_species = self.loss_fn_species(self.margin_fn_species(logits_species, species), species)
+        self.log_dict({"train/loss_ids": loss_ids.detach()}, on_step=False, on_epoch=True)
+        self.log_dict({"train/loss_species": loss_species.detach()}, on_step=False, on_epoch=True)
+        with torch.no_grad():
+            self.log_dict(map_dict(logits_ids, ids, "train"), on_step=False, on_epoch=True)
+            self.log_dict(
+                {"train/acc_species": topk_average_precision(logits_species, species, 1).mean().detach()},
+                on_step=False,
+                on_epoch=True,
+            )
+        return loss_ids * self.hparams.loss_id_ratio + loss_species * (1 - self.hparams.loss_id_ratio)
+
+    def validation_step(self, batch, batch_idx):
+        x, ids, species = batch["image"], batch["label"], batch["label_species"]
+        out1, out_species1 = self(x)
+        out2, out_species2 = self(x.flip(3))
+        output, output_species = (out1 + out2) / 2, (out_species1 + out_species2) / 2
+        self.log_dict(map_dict(output, ids, "val"), on_step=False, on_epoch=True)
+        self.log_dict(
+            {"val/acc_species": topk_average_precision(output_species, species, 1).mean().detach()},
+            on_step=False,
+            on_epoch=True,
+        )
+
+    def configure_optimizers(self):
+        backbone_params = list(self.backbone.parameters()) + list(self.global_pools.parameters())
+        head_params = (
+            list(self.neck.parameters()) + list(self.head_id.parameters()) + list(self.head_species.parameters())
+        )
+        params = [
+            {"params": backbone_params, "lr": self.hparams.lr_backbone},
+            {"params": head_params, "lr": self.hparams.lr_head},
+        ]
+        if self.hparams.optimizer == "Adam":
+            optimizer = torch.optim.Adam(params)
+        elif self.hparams.optimizer == "AdamW":
+            optimizer = torch.optim.AdamW(params)
+        elif self.hparams.optimizer == "RAdam":
+            optimizer = torch.optim.RAdam(params)
+
+        warmup_steps = self.hparams.max_epochs * self.hparams.warmup_steps_ratio
+        cycle_steps = self.hparams.max_epochs - warmup_steps
+        lr_lambda = WarmupCosineLambda(warmup_steps, cycle_steps, self.hparams.lr_decay_scale)
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+        return [optimizer], [scheduler]
+
+    def test_step(self, batch, batch_idx):
+        x = batch["image"]
+        feat1 = self.get_feat(x)
+        out1, out_species1 = self.head_id(feat1), self.head_species(feat1)
+        feat2 = self.get_feat(x.flip(3))
+        out2, out_species2 = self.head_id(feat2), self.head_species(feat2)
+        pred_logit, pred_idx = ((out1 + out2) / 2).cpu().sort(descending=True)
+        return {
+            "original_index": batch["original_index"],
+            "label": batch["label"],
+            "label_species": batch["label_species"],
+            "pred_logit": pred_logit[:, :1000],
+            "pred_idx": pred_idx[:, :1000],
+            "pred_species": ((out_species1 + out_species2) / 2).cpu(),
+            "embed_features1": feat1.cpu(),
+            "embed_features2": feat2.cpu(),
+        }
+
+    def test_epoch_end(self, outputs: List[Dict[str, torch.Tensor]]):
+        outputs = self.all_gather(outputs)
+        if self.trainer.global_rank == 0:
+            epoch_results: Dict[str, np.ndarray] = {}
+            for key in outputs[0].keys():
+                if torch.cuda.device_count() > 1:
+                    result = torch.cat([x[key] for x in outputs], dim=1).flatten(end_dim=1)
+                else:
+                    result = torch.cat([x[key] for x in outputs], dim=0)
+                epoch_results[key] = result.detach().cpu().numpy()
+            np.savez_compressed(self.test_results_fp, **epoch_results)
+
+
+# def train(
+#     df: pd.DataFrame,
+#     args: argparse.Namespace,
+#     cfg: Config,
+#     fold: int,
+#     do_inference: bool = False,
+#     additional_dataset: WhaleDataset = None,
+#     optuna_trial: Optional[optuna.Trial] = None,
+# ) -> Optional[float]:
+#     out_dir = f"{args.out_base_dir}/{args.exp_name}/{fold}"
+#     id_class_nums = df.individual_id.value_counts().sort_index().values
+#     species_class_nums = df.species.value_counts().sort_index().values
+#     model = SphereClassifier(cfg, id_class_nums=id_class_nums, species_class_nums=species_class_nums)
+#     data_module = WhaleDataModule(
+#         df, cfg, f"{args.in_base_dir}/train_images", cfg.val_bbox, fold, additional_dataset=additional_dataset
+#     )
+#     loggers = [pl_loggers.CSVLogger(out_dir)]
+#     if args.wandb_logger:
+#         loggers.append(
+#             pl_loggers.WandbLogger(
+#                 project="kaggle-happywhale", group=args.exp_name, name=f"{args.exp_name}/{fold}", save_dir=out_dir
+#             )
+#         )
+#     callbacks = [LearningRateMonitor("epoch")]
+#     if optuna_trial is not None:
+#         callbacks.append(PyTorchLightningPruningCallback(optuna_trial, "val/mapNone"))
+#     if args.save_checkpoint:
+#         callbacks.append(ModelCheckpoint(out_dir, save_last=True, save_top_k=0))
+#     trainer = Trainer(
+#         gpus=torch.cuda.device_count(),
+#         max_epochs=cfg["max_epochs"],
+#         logger=loggers,
+#         callbacks=callbacks,
+#         checkpoint_callback=args.save_checkpoint,
+#         precision=16,
+#         sync_batchnorm=True,
+#     )
+#     ckpt_path = f"{out_dir}/last.ckpt"
+#     if not os.path.exists(ckpt_path) or not args.load_snapshot:
+#         ckpt_path = None
+#     trainer.fit(model, ckpt_path=ckpt_path, datamodule=data_module)
+#     if do_inference:
+#         for test_bbox in cfg.test_bboxes:
+#             # all train data
+#             model.test_results_fp = f"{out_dir}/train_{test_bbox}_results.npz"
+#             trainer.test(model, data_module.all_dataloader())
+#             # test data
+#             model.test_results_fp = f"{out_dir}/test_{test_bbox}_results.npz"
+#             df_test = load_df(args.in_base_dir, cfg, "sample_submission.csv", False)
+#             test_data_module = WhaleDataModule(df_test, cfg, f"{args.in_base_dir}/test_images", test_bbox, -1)
+#             trainer.test(model, test_data_module.all_dataloader())
+
+#     if args.wandb_logger:
+#         wandb.finish()
+#     if optuna_trial is not None:
+#         return trainer.callback_metrics["val/mapNone"].item()
+#     else:
+#         return None
+
+
+# def main():
+#     args = parse()
+#     warnings.filterwarnings("ignore", ".*does not have many workers.*")
+#     cfg = load_config(args.config_path, "config/default.yaml")
+#     print(cfg)
+#     df = load_df(args.in_base_dir, cfg, "train.csv", True)
+#     pseudo_dataset = None
+#     if cfg.pseudo_label is not None:
+#         pseudo_df = load_df(args.in_base_dir, cfg, cfg.pseudo_label, False)
+#         pseudo_dataset = WhaleDataset(
+#             pseudo_df[pseudo_df.conf > cfg.pseudo_conf_threshold], cfg, f"{args.in_base_dir}/test_images", "", True
+#         )
+#     if cfg["n_splits"] == -1:
+#         train(df, args, cfg, -1, do_inference=True, additional_dataset=pseudo_dataset)
+#     else:
+#         train(df, args, cfg, 0, do_inference=True, additional_dataset=pseudo_dataset)
+
+
+# if __name__ == "__main__":
+#     main()

utils.py

@@ -0,0 +1,41 @@
+import math
+from typing import Optional
+
+import torch
+
+
+class WarmupCosineLambda:
+    def __init__(self, warmup_steps: int, cycle_steps: int, decay_scale: float, exponential_warmup: bool = False):
+        self.warmup_steps = warmup_steps
+        self.cycle_steps = cycle_steps
+        self.decay_scale = decay_scale
+        self.exponential_warmup = exponential_warmup
+
+    def __call__(self, epoch: int):
+        if epoch < self.warmup_steps:
+            if self.exponential_warmup:
+                return self.decay_scale * pow(self.decay_scale, -epoch / self.warmup_steps)
+            ratio = epoch / self.warmup_steps
+        else:
+            ratio = (1 + math.cos(math.pi * (epoch - self.warmup_steps) / self.cycle_steps)) / 2
+        return self.decay_scale + (1 - self.decay_scale) * ratio
+
+
+def topk_average_precision(output: torch.Tensor, y: torch.Tensor, k: int):
+    score_array = torch.tensor([1.0 / i for i in range(1, k + 1)], device=output.device)
+    topk = output.topk(k)[1]
+    match_mat = topk == y[:, None].expand(topk.shape)
+    return (match_mat * score_array).sum(dim=1)
+
+
+def calc_map5(output: torch.Tensor, y: torch.Tensor, threshold: Optional[float]):
+    if threshold is not None:
+        output = torch.cat([output, torch.full((output.shape[0], 1), threshold, device=output.device)], dim=1)
+    return topk_average_precision(output, y, 5).mean().detach()
+
+
+def map_dict(output: torch.Tensor, y: torch.Tensor, prefix: str):
+    d = {f"{prefix}/acc": topk_average_precision(output, y, 1).mean().detach()}
+    for threshold in [None, 0.3, 0.4, 0.5, 0.6, 0.7]:
+        d[f"{prefix}/map{threshold}"] = calc_map5(output, y, threshold)
+    return d