{"cells":[{"metadata":{},"cell_type":"markdown","source":"# About\n\n`InChI` string can be split by `/` into some parts (max number of parts in training data is 11). The first part is the format which is uniform string(`InChI=1S`) in this competition, and the second part is **chemical formula** which represents **the number of atoms** in each molecular.\n\nLet me take `InChI=1S/C13H20OS/c1-9(2)8-15-13-6-5-10(3)7-12(13)11(4)14/h5-7,9,11,14H,8H2,1-4H3`(`image_id`: `000011a64c74`) as an example.\nThe second part is `C13H20OS`, which means that the molecular have 13`Carbon`s, 20`Hydrogen`s, a `Oxygen` and a `Sulfur`.\n\nAll the parts of `InChI` have **variable length** except for the second one (chemical formula) because the kind of atoms in training data is **limited** to 12(`B`, `Br`, `C`, `Cl`, `F`, `H`, `I`, `N`, `O`, `P`, `S`, and `Si`). Therefore, we can represents a chemical formula by a **fixed length** vector and treat chemical formula prediction task as **multi-output regression task**.\n\n\nIn this notebook, I try to train a model which predicts chemical furmula by multi-variate regression. I think precisse prediction of chemical furmula is helpful to predict other parts of `InChI`.\n\nThere are some notebooks which forcus on atoms or chemical formulas.\n\n* [Bristol-Myers Squibb_count_atom](https://www.kaggle.com/kalfirst/bristol-myers-squibb-count-atom)\n* [Bristol-Myers Squibb: Counting Elements](https://www.kaggle.com/stainsby/bristol-myers-squibb-counting-elements)\n* [Step by Step 2: LS dist < 1 chemical formula](https://www.kaggle.com/wineplanetary/step-by-step-2-ls-dist-1-chemical-formula)\n\nThe last one already did what I want to do (Thank you for sharing @wineplanetary). I noticed that when I was almost done with this notebook😅\n\nTo make the difference, I will show you training process in this notebook, and try solving the competition task by utilizing predicted chemical formula in inference notebook."},{"metadata":{},"cell_type":"markdown","source":"# Prepare"},{"metadata":{},"cell_type":"markdown","source":"## import"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport re\nimport gc\nimport sys\nimport yaml\nimport copy\nimport random\nimport shutil\nimport typing as tp\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm.notebook import tqdm\nfrom sklearn.metrics import mean_squared_error\n\nimport cv2\nimport albumentations\nfrom albumentations.core.transforms_interface import ImageOnlyTransform, DualTransform\nfrom albumentations.pytorch import ToTensorV2\n\nimport torch\nfrom torch import nn\nfrom torch.utils import data\n\nsys.path.append(\"../input/pytorch-pfn-extras/pytorch-pfn-extras-0.3.2\")\nimport pytorch_pfn_extras as ppe\nfrom pytorch_pfn_extras.training import extensions as ppe_exts\n\nsys.path.append(\"../input/iterative-stratification/iterative-stratification-master\")\nfrom iterstrat.ml_stratifiers import MultilabelStratifiedKFold\n\nsys.path.append(\"../input/timm-pytorch-image-models/pytorch-image-models-master\")\nimport timm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ROOT = Path.cwd().parent\nINPUT = ROOT / \"input\"\nDATA = INPUT / \"bms-molecular-translation\"\nTRAIN = DATA / \"train\"\nTRAIN_224 = INPUT / \"bms-molecular-224px-jpg-padded\" / \"train\"\nTEST = DATA / \"test\"\nTEST_224 = INPUT / \"bms-molecular-224px-jpg-padded\" / \"test\"\n\nTMP = ROOT / \"tmp\"\nTMP.mkdir(exist_ok=True)\n\nRANDAM_SEED = 1086\n\nFOLDS = [0, 1, 2, 3, 4]\nN_FOLD = len(FOLDS)\n# FOLDS = [0,]\n# N_FOLD = 5\n\nTARGETS = [\n    'B', 'Br', 'C', 'Cl',\n    'F', 'H', 'I', 'N',\n    'O', 'P', 'S', 'Si']\n\nN_TARGETS = len(TARGETS)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## read label data"},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv(DATA / \"train_labels.csv\")\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## preprocess"},{"metadata":{},"cell_type":"markdown","source":"### extract chemical formula from InChI"},{"metadata":{"trusted":true},"cell_type":"code","source":"# # extract chemical formula\ntrain[\"formula\"] = train.InChI.str.extract(\"InChI=1S/([^/]+)/.+\")\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### counting number of atoms for each example"},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"elem_regex = re.compile(r\"[A-Z][a-z]?[0-9]*\")\natom_regex = re.compile(r\"[A-Z][a-z]?\")\ndgts_regex = re.compile(r\"[0-9]*\")\n\nformula_examples = [\n    \"C23H19ClIN3O\",\n    \"C33H49B2N3O4\",\n    \"C13H12BrF3N4OS\",\n    \"C5H18O2P2Si2\"]\n\nfor i, fml in enumerate(formula_examples):\n    print(f\"[example{i + 1}: {fml}]\")\n    print(\"\\tatom with digits:\", elem_regex.findall(fml))\n    print(\"\\tatom:\", atom_regex.findall(fml))\n    print(\"\\tdigits\", dgts_regex.findall(fml))","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"# # example for counting method\nfor fml in formula_examples:\n    atom_dict = dict()\n    print(f\"[formula: {fml}]\")\n    for elem in elem_regex.findall(fml):\n        atom = dgts_regex.sub(\"\", elem)\n        dgts = atom_regex.sub(\"\", elem)\n        atom_cnt = int(dgts) if len(dgts) > 0 else 1\n        atom_dict[atom] = atom_cnt\n        print(f\"\\t[elem:\\t{elem}]\\tatom: {atom},\\tdgts: {dgts},\\tatom_cnt: {atom_cnt}\")\n    print(f\"\\tresult: {atom_dict}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # apply to all train data\natom_dict_list = []\nfor fml in tqdm(train[\"formula\"].values):\n    atom_dict = dict()\n    for elem in elem_regex.findall(fml):\n        atom = dgts_regex.sub(\"\", elem)\n        dgts = atom_regex.sub(\"\", elem)\n        atom_cnt = int(dgts) if len(dgts) > 0 else 1\n        atom_dict[atom] = atom_cnt\n    atom_dict_list.append(atom_dict)\n    \natom_df = pd.DataFrame(\n    atom_dict_list).fillna(0).astype(int)\natom_df = atom_df.sort_index(axis=\"columns\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # merge\nfor atom in TARGETS:\n    train[atom] = atom_df[atom]\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del atom_df\ndel atom_dict\ndel atom_dict_list\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### check atom distribution"},{"metadata":{"trusted":true},"cell_type":"code","source":"# # total number of each atoms\ndisplay(train[TARGETS].sum(axis=0))\n_ = train[TARGETS].sum(axis=0).T.plot(kind=\"bar\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # distribution of n_atoms for each example\ntrain[\"n_atoms\"] = train[TARGETS].sum(axis=1)\n_ = train[\"n_atoms\"].hist(bins=100)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training Multi-Output Regression\n\n## summary\n\n* base model: resnet18d\n* CV Strategy: Multi-Label Stratified KFold(K=5) by atoms in each chemical formula\n* image size: 224x224 (simply resize)\n* batch size: 64\n* epoch: 15\n* optimizer: AdamW\n* schedular: OneCycleLR\n* augmentation: HorizontalFlip, VerticalFlip, ShiftScaleRotate, RandomResizedCrop\n\n**NOTE**: I use **only 4%** of training data for faster trainng."},{"metadata":{},"cell_type":"markdown","source":"## settings"},{"metadata":{"trusted":true},"cell_type":"code","source":"settings = yaml.safe_load(\"\"\"\nglobals:\n  seed: 1086\n  device: cuda\n  max_epoch: 15\n  patience: -1\n  use_amp: True\n  reduce_data: True\n  reduce_div_factor: 25\n\ndataset:\n  name: LabeledImageDataset\n  train:\n    transform_list:\n      # - [Resize, {always_apply: True, height: 224, width: 224}]\n      - [HorizontalFlip, {p: 0.5}]\n      - [VerticalFlip, {p: 0.5}]\n      - [ShiftScaleRotate, {\n          p: 0.5, shift_limit: 0.2, scale_limit: 0.2,\n          rotate_limit: 20, border_mode: 0, value: 0, mask_value: 0}]\n      - [RandomResizedCrop, {height: 224, width: 224, scale: [0.9, 1.0]}]\n      - [Normalize, {always_apply: True}]\n      - [ToTensorV2, {always_apply: True}]\n  val:\n    transform_list:\n      # - [Resize, {always_apply: True, height: 224, width: 224}]\n      - [Normalize, {always_apply: True}]\n      - [ToTensorV2, {always_apply: True}]\n\nloader:\n  train:\n    batch_size: 64\n    shuffle: True\n    num_workers: 4\n    pin_memory: True\n    drop_last: True\n  val:\n    batch_size: 128\n    shuffle: False\n    num_workers: 4\n    pin_memory: True\n    drop_last: False\n\nmodel:\n  name: BasicImageModel\n  params:\n    base_name: resnet18d\n    dims_head: [null, 12]\n    pretrained: True\n\nloss:\n  name: MSELoss\n  params: {}\n\neval:\n  - {name: MyMSELoss, report_name: loss, params: {}}\n\noptimizer:\n  name: AdamW\n  params:\n    lr: 1.0e-06\n    weight_decay: 1.0e-02\n\nscheduler:\n  name: OneCycleLR\n  params:\n    epochs: 15\n    max_lr: 1.0e-3\n    pct_start: 0.2\n    anneal_strategy: cos\n    div_factor: 1.0e+3\n    final_div_factor: 1.0e+3\n\"\"\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## definition"},{"metadata":{},"cell_type":"markdown","source":"### model"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class BasicImageModel(nn.Module):\n    \n    def __init__(\n        self, base_name, dims_head: tp, pretrained=False\n    ):\n        \"\"\"Initialize\"\"\"\n        self.base_name = base_name\n        super(BasicImageModel, self).__init__()\n        \n        # # prepare backbone\n        if hasattr(timm.models, base_name):\n            # # # load base model\n            base_model = timm.create_model(base_name, pretrained=pretrained)\n            in_features = base_model.num_features\n            # # remove head classifier\n            base_model.reset_classifier(0)\n        else:\n            raise NotImplementedError\n\n        self.backbone = base_model\n        \n        # # prepare head clasifier\n        if dims_head[0] is None:\n            dims_head[0] = in_features\n\n        layers_list = []\n        for i in range(len(dims_head) - 2):\n            in_dim, out_dim = dims_head[i: i + 2]\n            layers_list.extend([\n                nn.Linear(in_dim, out_dim),\n                nn.ReLU(), nn.Dropout(0.5),])\n        layers_list.append(\n            nn.Linear(dims_head[-2], dims_head[-1]))\n        self.head = nn.Sequential(*layers_list)\n\n    def forward(self, x):\n        \"\"\"Forward\"\"\"\n        h = self.backbone(x)\n        h = self.head(h)\n        return h","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### image dataset"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class ImageTransformBase:\n    \"\"\"\n    Base Image Transform class.\n\n    Args:\n        data_augmentations: List of tuple(method: str, params :dict), each elems pass to albumentations\n    \"\"\"\n\n    def __init__(self, data_augmentations: tp.List[tp.Tuple[str, tp.Dict]]):\n        \"\"\"Initialize.\"\"\"\n        augmentations_list = [\n            self._get_augmentation(aug_name)(**params)\n            for aug_name, params in data_augmentations]\n        self.data_aug = albumentations.Compose(augmentations_list)\n\n    def __call__(self, pair: tp.Tuple[np.ndarray]) -> tp.Tuple[np.ndarray]:\n        \"\"\"You have to implement this by task\"\"\"\n        raise NotImplementedError\n\n    def _get_augmentation(self, aug_name: str) -> tp.Tuple[ImageOnlyTransform, DualTransform]:\n        \"\"\"Get augmentations from albumentations\"\"\"\n        if hasattr(albumentations, aug_name):\n            return getattr(albumentations, aug_name)\n        else:\n            return eval(aug_name)\n\n\nclass ImageTransformForCls(ImageTransformBase):\n    \"\"\"Data Augmentor for Classification Task.\"\"\"\n\n    def __init__(self, data_augmentations: tp.List[tp.Tuple[str, tp.Dict]]):\n        \"\"\"Initialize.\"\"\"\n        super(ImageTransformForCls, self).__init__(data_augmentations)\n\n    def __call__(self, in_arrs: tp.Tuple[np.ndarray]) -> tp.Tuple[np.ndarray]:\n        \"\"\"Apply Transform.\"\"\"\n        img, label = in_arrs\n        augmented = self.data_aug(image=img)\n        img = augmented[\"image\"]\n\n        return img, label","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class LabeledImageDataset(data.Dataset):\n    \"\"\"Dataset class for (image, label) pairs\"\"\"\n\n    def __init__(\n        self,\n        file_list: tp.List[\n            tp.Tuple[tp.Union[str, Path], tp.Union[int, float, np.ndarray]]],\n        transform_list: tp.List[tp.Dict],\n    ):\n        \"\"\"Initialize\"\"\"\n        self.file_list = file_list\n        self.transform = ImageTransformForCls(transform_list)\n\n    def __len__(self):\n        \"\"\"Return Num of Images.\"\"\"\n        return len(self.file_list)\n\n    def __getitem__(self, index):\n        \"\"\"Return transformed image and mask for given index.\"\"\"\n        img_path, label = self.file_list[index]\n        img = self._read_image_as_array(img_path)\n\n        img, label = self.transform((img, label))\n        return img, label\n\n    def _read_image_as_array(self, path: str):\n        \"\"\"Read image file and convert into numpy.ndarray\"\"\"\n        img_arr = cv2.imread(str(path))\n        img_arr = cv2.cvtColor(img_arr, cv2.COLOR_BGR2RGB)\n        return img_arr","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### metric for evaluation"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class EvalFuncManager(nn.Module):\n    \"\"\"Manager Class for evaluation at the end of epoch\"\"\"\n\n    def __init__(\n        self,\n        evalfunc_dict: tp.Dict[str, nn.Module],\n        iters_per_epoch: int,\n        prefix: str = \"val\"\n    ) -> None:\n        \"\"\"Initialize\"\"\"\n        self.tmp_iter = 0\n        self.iters_per_epoch = iters_per_epoch\n        self.prefix = prefix\n        self.metric_names = []\n        super(EvalFuncManager, self).__init__()\n        for k, v in evalfunc_dict.items():\n            setattr(self, k, v)\n            self.metric_names.append(k)\n        self.reset()\n\n    def reset(self) -> None:\n        \"\"\"Reset State.\"\"\"\n        self.tmp_iter = 0\n        for name in self.metric_names:\n            getattr(self, name).reset()\n\n    def __call__(self, y: torch.Tensor, t: torch.Tensor) -> None:\n        \"\"\"Forward.\"\"\"\n        for name in self.metric_names:\n            getattr(self, name).update(y, t)\n        self.tmp_iter += 1\n\n        if self.tmp_iter == self.iters_per_epoch:\n            ppe.reporting.report({\n                \"{}/{}\".format(self.prefix, name): getattr(self, name).compute()\n                for name in self.metric_names\n            })\n            self.reset()\n            \n            \nclass MeanLoss(nn.Module):\n    \n    def __init__(self):\n        super(MeanLoss, self).__init__()\n        self.loss_sum = 0\n        self.n_examples = 0\n        \n    def forward(self, y: torch.Tensor, t: torch.Tensor):\n        \"\"\"Compute metric at once\"\"\"\n        return self.loss_func(y, t)\n\n    def reset(self):\n        \"\"\"Reset state\"\"\"\n        self.loss_sum = 0\n        self.n_examples = 0\n    \n    def update(self, y: torch.Tensor, t: torch.Tensor):\n        \"\"\"Update metric by mini batch\"\"\"\n        self.loss_sum += self(y, t).item() * y.shape[0]\n        self.n_examples += y.shape[0]\n\n    def compute(self):\n        \"\"\"Compute metric for dataset\"\"\"\n        return self.loss_sum / self.n_examples\n    \n    \nclass MyMSELoss(MeanLoss):\n\n    def __init__(self, **params):\n        super(MyMSELoss, self).__init__()\n        self.loss_func = nn.MSELoss(**params)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### get XXX"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def get_file_list(stgs, train_all):\n    \"\"\"Get file path and target info.\"\"\"\n    use_fold = stgs[\"globals\"][\"val_fold\"]\n    \n    train_df = train_all[train_all[\"fold\"] != use_fold]\n    val_df = train_all[train_all[\"fold\"] == use_fold]\n    \n#     train_data_dir = TRAIN\n    train_data_dir = TRAIN_224\n    train_file_list = list(zip([\n#         train_data_dir / f\"{img_id[0]}/{img_id[1]}/{img_id[2]}/{img_id}.png\"\n        train_data_dir / f\"{img_id}.jpg\"\n        for img_id in train_df[\"image_id\"].values\n    ], train_df[TARGETS].values.astype(\"f\")))\n\n    val_file_list = list(zip([\n#         train_data_dir / f\"{img_id[0]}/{img_id[1]}/{img_id[2]}/{img_id}.png\"\n        train_data_dir / f\"{img_id}.jpg\"\n        for img_id in val_df[\"image_id\"].values\n    ], val_df[TARGETS].values.astype(\"f\")))\n\n    if stgs[\"globals\"][\"reduce_data\"]:\n        div = stgs[\"globals\"][\"reduce_div_factor\"]\n        trn_smpl_idx = np.random.choice(range(len(train_df)), len(train_df) // div)\n        val_smpl_idx = np.random.choice(range(len(val_df)), len(val_df) // div)\n        train_file_list = [train_file_list[idx] for idx in trn_smpl_idx]\n        val_file_list = [val_file_list[idx] for idx in val_smpl_idx]\n        \n    return train_file_list, val_file_list\n\n\ndef get_dataloaders(\n    stgs: tp.Dict,\n    train_file_list: tp.List[tp.List],\n    val_file_list: tp.List[tp.List],\n    dataset_class: data.Dataset\n):\n    \"\"\"Create DataLoader\"\"\"\n    train_loader = val_loader = None\n    if train_file_list is not None:\n        train_dataset = dataset_class(\n            train_file_list, **stgs[\"dataset\"][\"train\"])\n        train_loader = data.DataLoader(\n            train_dataset, **stgs[\"loader\"][\"train\"])\n\n    if val_file_list is not None:\n        val_dataset = dataset_class(\n            val_file_list, **stgs[\"dataset\"][\"val\"])\n        val_loader = data.DataLoader(\n            val_dataset, **stgs[\"loader\"][\"val\"])\n\n    return train_loader, val_loader","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def get_model(args):\n    \"\"\"\"\"\"\n    return eval(args[\"name\"])(**args[\"params\"])\n\n\ndef get_optimizer(args, model):\n    \"\"\"\"\"\"\n    if hasattr(torch.optim, args[\"name\"]):\n        opt_class = getattr(torch.optim, args[\"name\"])\n    else:\n        opt_class = eval(args[\"name\"])\n\n    return opt_class(model.parameters(), **args[\"params\"])\n\n\ndef get_scheduler(args, optimizer, max_epoch, steps_per_epoch):\n    \"\"\"\"\"\"\n    if args[\"name\"] == \"OneCycleLR\":\n        args[\"params\"][\"epochs\"] = max_epoch\n        args[\"params\"][\"steps_per_epoch\"] = steps_per_epoch\n\n    if hasattr(torch.optim.lr_scheduler, args[\"name\"]):\n        scdr_class = getattr(torch.optim.lr_scheduler, args[\"name\"])\n    else:\n        scdr_class = eval(args[\"name\"])\n\n    return scdr_class(optimizer, **args[\"params\"])\n\n\ndef get_loss_function(args):\n    \"\"\"\"\"\"\n    if hasattr(nn, args[\"name\"]):\n        loss_class = getattr(nn, args[\"name\"])\n    else:\n        loss_class = eval(args[\"name\"])\n\n    return loss_class(**args[\"params\"])\n\n\ndef get_stepper(manager, stgs, scheduler):\n    \"\"\"\"\"\"\n    if stgs[\"scheduler\"][\"name\"] == \"CosineAnnealingWarmRestarts\":\n        def step_scheduler_by_epoch():\n            pass\n\n        def step_scheduler_by_iter():\n            scheduler.step(manager.epoch_detail)\n\n    elif stgs[\"scheduler\"][\"name\"] == \"OneCycleLR\":\n        def step_scheduler_by_epoch():\n            pass\n\n        def step_scheduler_by_iter():\n            scheduler.step()\n\n    else:\n        def step_scheduler_by_epoch():\n            scheduler.step()\n\n        def step_scheduler_by_iter():\n            pass\n\n    return step_scheduler_by_epoch, step_scheduler_by_iter","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def get_manager(\n    stgs, model, device, train_loader, val_loader, optimizer,\n    eval_manager, output_path, print_progress: bool = False,\n):\n    \"\"\"\"\"\"\n    # # initialize manager\n    if stgs[\"globals\"][\"patience\"] > 0:\n        trigger = ppe.training.triggers.EarlyStoppingTrigger(\n            check_trigger=(1, 'epoch'), monitor='val/loss', mode=\"min\",\n            patience=stgs[\"globals\"][\"patience\"], verbose=True,\n            max_trigger=(stgs[\"globals\"][\"max_epoch\"], 'epoch'))\n    else:\n        trigger = None    \n    manager = ppe.training.ExtensionsManager(\n        model, optimizer, stgs[\"globals\"][\"max_epoch\"],\n        iters_per_epoch=len(train_loader), stop_trigger=trigger, out_dir=output_path)\n\n    # # for logging\n    eval_names = [\"val/{}\".format(name) for name in eval_manager.metric_names]    \n    log_extentions = [\n        ppe_exts.observe_lr(optimizer=optimizer),\n        ppe_exts.LogReport(),\n        ppe_exts.PlotReport([\"train/loss\", \"val/loss\"], 'epoch', filename='loss.png'),\n        # ppe_exts.PlotReport([\"val/metric\"], 'epoch', filename='metric.png'),\n        ppe_exts.PlotReport([\"lr\"], 'epoch', filename='lr.png'),\n        ppe_exts.PrintReport([\n            \"epoch\", \"iteration\", \"lr\", \"train/loss\", *eval_names, \"elapsed_time\"])\n    ]\n    if print_progress:\n        log_extentions.append(ppe_exts.ProgressBar(update_interval=20))\n\n    for ext in log_extentions:\n        manager.extend(ext)\n        \n    # # for evaluation\n    def eval_func(*batch):\n        return run_eval(stgs, model, device, batch, eval_manager)\n    manager.extend(\n        ppe_exts.Evaluator(val_loader, model, eval_func=eval_func),\n        trigger=(1, \"epoch\"))\n    \n    # # for saving snapshot\n    manager.extend(\n        ppe_exts.snapshot(target=model, filename=\"snapshot_by_loss_epoch_{.epoch}.pth\"),\n        trigger=ppe.training.triggers.MinValueTrigger(key=\"val/loss\", trigger=(1, 'epoch')))\n\n    return manager","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### training utils"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def set_random_seed(seed: int = 42, deterministic: bool = False):\n    \"\"\"Set seeds\"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)  # type: ignore\n    torch.backends.cudnn.deterministic = deterministic  # type: ignore","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def run_train_loop(\n    manager, stgs, model, device, train_loader, optimizer, scheduler, loss_func\n):\n    \"\"\"Run minibatch training loop\"\"\"\n    step_scheduler_by_epoch, step_scheduler_by_iter = get_stepper(manager, stgs, scheduler)\n \n    use_amp = stgs[\"globals\"][\"use_amp\"]\n    while not manager.stop_trigger:\n        model.train()\n        scaler = torch.cuda.amp.GradScaler() if use_amp else None\n        for batch in train_loader:\n            with manager.run_iteration():\n                x = batch[0].to(device)\n                t = batch[-1].to(device)\n                optimizer.zero_grad()\n                if use_amp:\n                    with torch.cuda.amp.autocast():\n                        y = model(x)\n                        loss = loss_func(y, t)    \n                    scaler.scale(loss).backward()\n                    scaler.step(optimizer)\n                    scaler.update()\n                else:\n                    y = model(x)\n                    loss = loss_func(y, t)\n                    loss.backward()\n                    optimizer.step()\n\n                ppe.reporting.report({'train/loss': loss.item()})    \n                step_scheduler_by_iter()\n        step_scheduler_by_epoch()\n\n\ndef run_eval(stgs, model, device, batch, eval_manager):\n    \"\"\"Run evaliation for val or test. this function is applied to each batch.\"\"\"\n    model.eval()\n    x = batch[0].to(device)\n    t = batch[-1].to(device)\n    if stgs[\"globals\"][\"use_amp\"]:\n        with torch.cuda.amp.autocast(): \n            y = model(x)\n            eval_manager(y, t)\n    else:\n        y = model(x)\n        eval_manager(y, t)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def train_one_fold(stgs, train_all, output_path, print_progress=False):\n    \"\"\"train one fold\"\"\"\n    torch.backends.cudnn.benchmark = True\n    set_random_seed(stgs[\"globals\"][\"seed\"])\n\n    # # prepare train, valid paths\n    train_file_list, val_file_list = get_file_list(stgs, train_all)\n    print(\"train: {}, val: {}\".format(len(train_file_list), len(val_file_list)))\n\n    device = torch.device(stgs[\"globals\"][\"device\"])\n    # # get data_loader\n    train_loader, val_loader = get_dataloaders(\n        stgs, train_file_list, val_file_list, LabeledImageDataset)\n\n    # # get model\n    model = BasicImageModel(**stgs[\"model\"][\"params\"])\n    model.to(device)\n\n    # # get optimizer\n    optimizer = getattr(\n        torch.optim, stgs[\"optimizer\"][\"name\"]\n    )(model.parameters(), **stgs[\"optimizer\"][\"params\"])\n\n    # # get scheduler\n    if stgs[\"scheduler\"][\"name\"] == \"OneCycleLR\":\n        stgs[\"scheduler\"][\"params\"][\"epochs\"] = stgs[\"globals\"][\"max_epoch\"]\n        stgs[\"scheduler\"][\"params\"][\"steps_per_epoch\"] = len(train_loader)\n    scheduler = getattr(\n        torch.optim.lr_scheduler, stgs[\"scheduler\"][\"name\"]\n    )(optimizer, **stgs[\"scheduler\"][\"params\"])\n\n    # # get loss\n    if hasattr(nn, stgs[\"loss\"][\"name\"]):\n        loss_func = getattr(nn, stgs[\"loss\"][\"name\"])(**stgs[\"loss\"][\"params\"])\n    else:\n        loss_func = eval(stgs[\"loss\"][\"name\"])(**stgs[\"loss\"][\"params\"])\n    loss_func.to(device)\n\n    eval_manager = EvalFuncManager(\n        {\n            metric[\"report_name\"]: eval(metric[\"name\"])(**metric[\"params\"])\n            for metric in stgs[\"eval\"]\n        }, len(val_loader))\n    eval_manager.to(device)\n\n    # # get manager\n    manager = get_manager(\n        stgs, model, device, train_loader, val_loader,\n        optimizer, eval_manager, output_path, print_progress)\n\n    run_train_loop(\n        manager, stgs, model, device, train_loader,\n        optimizer, scheduler, loss_func)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## train"},{"metadata":{},"cell_type":"markdown","source":"### split fold"},{"metadata":{"trusted":true},"cell_type":"code","source":"mskf = MultilabelStratifiedKFold(n_splits=5, shuffle=True, random_state=RANDAM_SEED)\ntrain[\"fold\"] = -1\nfor fold_id, (trn_idx, val_idx) in enumerate(\n    mskf.split(train.InChI, (train[TARGETS] > 0).astype(int))\n):\n    train.loc[val_idx, \"fold\"] = fold_id","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train.groupby(\"fold\")[TARGETS].sum()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train.groupby(\"fold\")[\"n_atoms\"].agg(\n    [\"sum\", \"mean\", \"median\", \"std\", \"max\", \"min\"])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"(train.InChI.str.len()).groupby(train[\"fold\"]).agg(\n    [\"sum\", \"mean\", \"median\", \"std\", \"max\", \"min\"])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### run training"},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.cuda.empty_cache()\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"stgs_list = []\nfor fold_id in FOLDS:\n    tmp_stgs = copy.deepcopy(settings)\n    tmp_stgs[\"globals\"][\"val_fold\"] = fold_id\n    stgs_list.append(tmp_stgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for fold_id, tmp_stgs in zip(FOLDS, stgs_list):\n    train_one_fold(tmp_stgs, train, TMP / f\"fold{fold_id}\", False)\n    torch.cuda.empty_cache()\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Inference OOF"},{"metadata":{},"cell_type":"markdown","source":"### copy best model"},{"metadata":{"trusted":true},"cell_type":"code","source":"best_log_list = []\nfor fold_id, tmp_stgs in enumerate(stgs_list):\n    exp_dir_path = TMP / f\"fold{fold_id}\"\n    log = pd.read_json(exp_dir_path / \"log\")\n    best_log = log.iloc[[log[\"val/loss\"].idxmin()],]\n    best_epoch = best_log.epoch.values[0]\n    best_log_list.append(best_log)\n    \n    best_model_path = exp_dir_path / f\"snapshot_by_loss_epoch_{best_epoch}.pth\"\n    copy_to = f\"./best_loss_model_fold{fold_id}.pth\"\n    shutil.copy(best_model_path, copy_to)\n    \n    for p in exp_dir_path.glob(\"*.pth\"):\n        p.unlink()\n    \n    shutil.copytree(exp_dir_path, f\"./fold{fold_id}\")\n    \n    with open(f\"./fold{fold_id}/settings.yml\", \"w\") as fw:\n        yaml.dump(tmp_stgs, fw)\n    \npd.concat(best_log_list, axis=0, ignore_index=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### inference"},{"metadata":{"trusted":true},"cell_type":"code","source":"def run_inference_loop(stgs, model, loader, device):\n    model.to(device)\n    model.eval()\n    pred_list = []\n    with torch.no_grad():\n        for x, t in tqdm(loader):\n            if stgs[\"globals\"][\"use_amp\"]:\n                with torch.cuda.amp.autocast():\n                    y = model(x.to(device))\n            else:\n                y = model(x.to(device))\n            pred_list.append(y.detach().cpu().numpy())\n        \n    pred_arr = np.concatenate(pred_list)\n    del pred_list\n    return pred_arr","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"oof_pred_arr = np.zeros((len(train), N_TARGETS))\nlabel_arr = train[TARGETS].values\nscore_list = []\n\nfor fold_id in range(N_FOLD):\n    print(f\"[fold {fold_id}]\")\n    tmp_dir = Path(f\"./fold{fold_id}\")\n    with open(tmp_dir / \"settings.yml\", \"r\") as fr:\n        tmp_stgs = yaml.safe_load(fr)\n    device = torch.device(tmp_stgs[\"globals\"][\"device\"])\n    val_idx = train.query(\"fold == @fold_id\").index.values\n    \n    # # get data_loader\n    tmp_stgs[\"globals\"][\"reduce_data\"] = False\n    _, val_file_list = get_file_list(tmp_stgs, train)\n    _, val_loader = get_dataloaders(\n        tmp_stgs, None, val_file_list, LabeledImageDataset)\n#     val_idx = val_idx[:len(val_file_list)]\n    \n    # # get and load model\n    model_path =f\"./best_loss_model_fold{fold_id}.pth\"\n    tmp_stgs[\"model\"][\"params\"][\"pretrained\"] = False\n    model = BasicImageModel(**tmp_stgs[\"model\"][\"params\"])\n    model.load_state_dict(torch.load(model_path, map_location=device))\n\n    val_pred = run_inference_loop(tmp_stgs, model, val_loader, device)\n    val_loss = mean_squared_error(label_arr[val_idx], val_pred)\n    print(f\"[fold {fold_id}] val loss: {val_loss:.5f}\")\n    oof_pred_arr[val_idx] = val_pred\n    score_list.append([fold_id, val_loss])\n    \n    del model\n    del val_pred\n    torch.cuda.empty_cache()\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"oof_loss = mean_squared_error(label_arr, oof_pred_arr)\nscore_list.append([\"oof\", oof_loss])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pd.DataFrame(score_list, columns=[\"fold\", \"mse\"])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"oof_pred_arr.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"oof_df = train.copy()\noof_df[TARGETS] = oof_pred_arr\noof_df.to_csv(\"./oof_prediction.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train.to_pickle(\"./train_formula_mlskf_5fold.pkl\")","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}