{"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":"I played around with some of the model choices today and found b5 to give a much better result hence I have decided to share this. If this helps then please make sure you have upvoted Jun Koda's original notebook as he did all the work here!","metadata":{}},{"cell_type":"markdown","source":"# Basic spectrogram image classification\n\nPower spectrum is averaged over every N = 32 timesteps:\n\n$$ P(\\omega) = \\frac{1}{N} \\sum_n^N \\left|a(\\omega, t_n) \\right|^2 $$\n\nImage size becomes: 2 channels x 360 frequencies x 128 timesteps.\n\nThe run time is dominated by data loading but I am not optimizing for that.\n\nVersion 4: Fix output target value to [0, 1] with `y_pred.sigmoid()`.","metadata":{}},{"cell_type":"code","source":"# Use timm pretrained image model\n! pip3 install timm","metadata":{"execution":{"iopub.status.busy":"2022-12-04T00:06:24.263226Z","iopub.execute_input":"2022-12-04T00:06:24.263658Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport time\nimport h5py\nimport timm\nimport torch\nimport torch.nn as nn\n\nfrom tqdm.auto import tqdm\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import roc_auc_score\nfrom timm.scheduler import CosineLRScheduler\nfrom imblearn.over_sampling import KMeansSMOTE,SMOTE,SVMSMOTE\nfrom imblearn.under_sampling import EditedNearestNeighbours\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n# criterion = nn.BCEWithLogitsLoss()\n\n# Train metadata\ndi = '/kaggle/input/g2net-detecting-continuous-gravitational-waves'\ndf = pd.read_csv(di + '/train_labels.csv')\ndf = df[df.target >= 0]  # Remove 3 unknowns (target = -1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df2 = pd.DataFrame(columns=['id'], data = os.listdir('/kaggle/input/g2net-detecting-continuous-gravitational-waves/test/'))\n\ndf2['id'] = df2['id'].apply(lambda x: x.split('.')[0])\n\ndf2['target'] = 0","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class Dataset(torch.utils.data.Dataset):\n    \"\"\"\n    dataset = Dataset(data_type, df)\n\n    img, y = dataset[i]\n      img (np.float32): 2 x 360 x 128\n      y (np.float32): label 0 or 1\n    \"\"\"\n    def __init__(self, data_type, df):\n        self.data_type = data_type\n        self.df = df\n        \n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, i):\n        \"\"\"\n        i (int): get ith data\n        \"\"\"\n        r = self.df.iloc[i]\n        y = np.float32(r.target)\n        file_id = r.id\n\n        img = np.empty((2, 360, 128), dtype=np.float32)\n\n        filename = '%s/%s/%s.hdf5' % (di, self.data_type, file_id)\n        with h5py.File(filename, 'r') as f:\n            g = f[file_id]\n\n            for ch, s in enumerate(['H1', 'L1']):\n                a = g[s]['SFTs'][:, :4096] * 1e22  # Fourier coefficient complex64\n\n                p = a.real**2 + a.imag**2  # power\n                p /= np.mean(p)  # normalize\n                p = np.mean(p.reshape(360, 128, 32), axis=2)  # compress 4096 -> 128\n\n                img[ch] = p\n\n        return img, y","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install -U catalyst","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset = Dataset('train', df)\nimg, y = dataset[10]\n\nplt.figure(figsize=(8, 3))\nplt.title('Spectrogram')\nplt.xlabel('time')\nplt.ylabel('frequency')\nplt.imshow(img[0, 300:360])  # zooming in for dataset[10]\nplt.colorbar()\nplt.show()\nplt.figure(figsize=(8, 3))\nplt.title('Spectrogram')\nplt.xlabel('time')\nplt.ylabel('frequency')\nplt.imshow(img[1, 300:360])  # zooming in for dataset[10]\nplt.colorbar()\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from catalyst.data import SelfSupervisedDatasetWrapper\n\nclass AddGaussianNoise(object):\n    def __init__(self, mean=0., std=1.):\n        self.std = std\n        self.mean = mean\n        \n    def __call__(self, tensor):\n        return tensor + torch.randn(tensor.size()) * self.std + self.mean\n    \n    def __repr__(self):\n        return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)\n\nfrom torchvision import transforms","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transform=transforms.Compose([\n    transforms.ToTensor(),\n    transforms.Normalize((0.1307,), (0.3081,)),\n    AddGaussianNoise(0., 1.)\n])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"class Model(nn.Module):\n    def __init__(self, name, *, pretrained=False, output_dim = 512):\n        \"\"\"\n        name (str): timm model name, e.g. tf_efficientnet_b2_ns\n        \"\"\"\n        super().__init__()\n\n        # Use timm\n        model = timm.create_model(name, pretrained=pretrained, in_chans=2)\n\n        clsf = model.default_cfg['classifier']\n        n_features = model._modules[clsf].in_features\n        model._modules[clsf] = nn.Identity()\n\n        self.fc = nn.Linear(n_features, output_dim)\n        self.model = model\n\n    def forward(self, x):\n        x = self.model(x)\n        x = self.fc(x)\n        return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Self SupervisedRunner**","metadata":{}},{"cell_type":"code","source":"from typing import Any, Mapping\nfrom collections import OrderedDict\n\nimport torch\nfrom torch import nn\n\nfrom catalyst.callbacks.backward import BackwardCallback\nfrom catalyst.callbacks.criterion import CriterionCallback\nfrom catalyst.callbacks.optimizer import OptimizerCallback\nfrom catalyst.callbacks.scheduler import SchedulerCallback\nfrom catalyst.core.callback import (\n    Callback,\n    IBackwardCallback,\n    ICriterionCallback,\n    IOptimizerCallback,\n    ISchedulerCallback,\n)\nfrom catalyst.core.engine import Engine\nfrom catalyst.core.misc import callback_isinstance\nfrom catalyst.core.runner import IRunner\nfrom catalyst.runners.runner import Runner\nfrom catalyst.typing import RunnerModel, TorchCriterion, TorchOptimizer, TorchScheduler\n\n\nclass ISelfSupervisedRunner(IRunner):\n    \"\"\"IRunner for experiments with contrastive model.\n    Args:\n        input_key: key in ``runner.batch`` dict mapping for model input\n        target_key: key in ``runner.batch`` dict mapping for target\n        loss_key: key for ``runner.batch_metrics`` to store criterion loss output\n        augemention_prefix: key for ``runner.batch`` to sample augumentions\n        projection_prefix: key for ``runner.batch`` to store model projection\n        embedding_prefix: key for `runner.batch`` to store model embeddings\n    Abstraction, please check out implementations for more details:\n        - :py:mod:`catalyst.runners.contrastive.ContrastiveRunner`\n    .. note::\n        ISelfSupervisedRunner contains only the logic with batch handling.\n    ISelfSupervisedRunner logic pseudocode:\n    .. code-block:: python\n        batch = {\"aug1\": tensor, \"aug2\": tensor, ...}\n        _, proj1 = model(batch[\"aug1\"])\n        _, proj2 = model(batch[\"aug2\"])\n        loss = criterion(proj1, proj2)\n        batch_metrics[\"loss_key\"] = loss\n    Examples:\n    .. code-block:: python\n        # 1. loader and transforms\n        transforms = Compose(\n            [\n                ToTensor(),\n                Normalize((0.1307,), (0.3081,)),\n                torchvision.transforms.RandomCrop((28, 28)),\n                torchvision.transforms.RandomVerticalFlip(),\n                torchvision.transforms.RandomHorizontalFlip(),\n            ]\n        )\n        mnist = MNIST(\"./logdir\", train=True, download=True, transform=None)\n        contrastive_mnist = ContrastiveDataset(mnist, transforms=transforms)\n        train_loader = torch.utils.data.DataLoader(contrastive_mnist, batch_size=BATCH_SIZE)\n        # 2. model and optimizer\n        encoder = MnistSimpleNet(out_features=16)\n        projection_head = nn.Sequential(\n            nn.Linear(16, 16, bias=False), nn.ReLU(inplace=True), nn.Linear(16, 16, bias=True)\n        )\n        class ContrastiveModel(torch.nn.Module):\n            def __init__(self, model, encoder):\n                super(ContrastiveModel, self).__init__()\n                self.model = model\n                self.encoder = encoder\n            def forward(self, x):\n                emb = self.encoder(x)\n                projection = self.model(emb)\n                return emb, projection\n        model = ContrastiveModel(model=projection_head, encoder=encoder)\n        optimizer = Adam(model.parameters(), lr=LR)\n        # 3. criterion with triplets sampling\n        criterion = NTXentLoss(tau=0.1)\n        callbacks = [\n            dl.ControlFlowCallback(\n                dl.CriterionCallback(\n                    input_key=\"projection_left\",\n                    target_key=\"projection_right\",\n                    metric_key=\"loss\"\n                ),\n                loaders=\"train\",\n            ),\n            dl.SklearnModelCallback(\n                feature_key=\"embedding_left\",\n                target_key=\"target\",\n                train_loader=\"train\",\n                valid_loaders=\"valid\",\n                model_fn=RandomForestClassifier,\n                predict_method=\"predict_proba\",\n                predict_key=\"sklearn_predict\",\n                random_state=RANDOM_STATE,\n                n_estimators=10,\n            ),\n            dl.ControlFlowCallback(\n                dl.AccuracyCallback(\n                    target_key=\"target\", input_key=\"sklearn_predict\", topk=(1, 3)\n                ),\n                loaders=\"valid\",\n            ),\n        ]\n        runner = dl.ContrastiveRunner()\n        logdir = \"./logdir\"\n        runner.train(\n            model=model,\n            criterion=criterion,\n            optimizer=optimizer,\n            callbacks=callbacks,\n            loaders={\"train\": train_loader, \"valid\": train_loader},\n            verbose=True,\n            logdir=logdir,\n            valid_loader=\"train\",\n            valid_metric=\"loss\",\n            minimize_valid_metric=True,\n            num_epochs=10,\n        )\n    .. note::\n        Please follow the `minimal examples`_ sections for use cases.\n        .. _`minimal examples`: https://github.com/catalyst-team/catalyst#minimal-examples  # noqa: E501, W505\n    \"\"\"\n\n    def __init__(\n        self,\n        input_key: str = \"features\",\n        target_key: str = \"target\",\n        loss_key: str = \"loss\",\n        augemention_prefix: str = \"augment\",\n        projection_prefix: str = \"projection\",\n        embedding_prefix: str = \"embedding\",\n    ):\n        \"\"\"Init.\"\"\"\n        IRunner.__init__(self)\n\n        self._target_key = target_key\n        self._loss_key = loss_key\n        self._projection_prefix = projection_prefix\n        self._augemention_prefix = augemention_prefix\n        self._embedding_prefix = embedding_prefix\n        self._input_key = input_key\n\n    def _process_batch(self, batch):\n        if isinstance(batch, (tuple, list)):\n            assert len(batch) in [3, 4]\n\n            if len(batch) == 4:\n                batch = {\n                    self._input_key: batch[0],\n                    f\"{self._augemention_prefix}_left\": batch[1],\n                    f\"{self._augemention_prefix}_right\": batch[2],\n                    self._target_key: batch[3],\n                }\n            elif len(batch) == 3:\n                batch = {\n                    self._input_key: batch[0],\n                    f\"{self._augemention_prefix}_left\": batch[1],\n                    f\"{self._augemention_prefix}_right\": batch[2],\n                }\n\n        return batch\n\n    def on_experiment_start(self, runner: \"IRunner\"):\n        \"\"\"on_experiment_start event handler.\"\"\"\n        super().on_experiment_start(runner)\n        self.is_kv_model = False\n        if isinstance(self.model, (Mapping, nn.ModuleDict)):\n            self.is_kv_model = True\n\n    def _process_input(self, batch: Mapping[str, Any], **kwargs):\n\n        if self.is_kv_model:\n            encoders = [\n                (encoder_name, self.model[encoder_name]) for encoder_name in self.model\n            ]\n        else:\n            encoders = [(\"\", self.model)]\n\n        for (encoder_name, encoder) in encoders:\n            embedding1, projection1 = encoder(\n                batch[f\"{self._augemention_prefix}_left\"], **kwargs\n            )\n            embedding2, projection2 = encoder(\n                batch[f\"{self._augemention_prefix}_right\"], **kwargs\n            )\n            origin_embeddings, projection_origin = encoder(\n                batch[self._input_key], **kwargs\n            )\n            prefix = f\"{encoder_name}_\" if encoder_name else \"\"\n            batch = {\n                **batch,\n                f\"{prefix}{self._projection_prefix}_left\": projection1,\n                f\"{prefix}{self._projection_prefix}_right\": projection2,\n                f\"{prefix}{self._projection_prefix}_origin\": projection_origin,\n                f\"{prefix}{self._embedding_prefix}_left\": embedding1,\n                f\"{prefix}{self._embedding_prefix}_right\": embedding2,\n                f\"{prefix}{self._embedding_prefix}_origin\": origin_embeddings,\n            }\n\n        return batch\n\n    def on_batch_start(self, runner: \"IRunner\"):\n        \"\"\"Event handler.\"\"\"\n        self.batch = self._process_batch(self.batch)\n        super().on_batch_start(runner)\n\n    def forward(self, batch: Mapping[str, Any], **kwargs) -> Mapping[str, Any]:\n        \"\"\"\n        Forward method for your Runner.\n        Should not be called directly outside of runner.\n        If your model has specific interface, override this method to use it\n        Args:\n            batch (Mapping[str, Any]): dictionary with data batches\n                from DataLoaders.\n            **kwargs: additional parameters to pass to the model\n        Returns:\n            dict with model output batch\n        \"\"\"\n        return self._process_input(batch, **kwargs)\n\n    def handle_batch(self, batch: Mapping[str, Any]) -> None:\n        \"\"\"\n        Inner method to handle specified data batch.\n        Used to make a train/valid/infer stage during Experiment run.\n        Args:\n            batch: dictionary with data batches from DataLoader.\n        \"\"\"\n        self.batch = {**batch, **self.forward(batch)}\n\n\nclass SelfSupervisedRunner(ISelfSupervisedRunner, Runner):\n    \n    def __init__(\n        self,\n        model: RunnerModel = None,\n        engine: Engine = None,\n        input_key: str = \"features\",\n        target_key: str = \"target\",\n        loss_key: str = \"loss\",\n        augemention_prefix: str = \"augment\",\n        projection_prefix: str = \"projection\",\n        embedding_prefix: str = \"embedding\",\n        loss_mode_prefix: str = \"projection\",\n    ):\n        \"\"\"Init.\"\"\"\n        ISelfSupervisedRunner.__init__(\n            self,\n            input_key=input_key,\n            target_key=target_key,\n            loss_key=loss_key,\n            augemention_prefix=augemention_prefix,\n            projection_prefix=projection_prefix,\n            embedding_prefix=embedding_prefix,\n        )\n        Runner.__init__(self, model=model, engine=engine)\n        self.loss_mode_prefix = loss_mode_prefix\n\n    @torch.no_grad()\n    def predict_batch(self, batch: Mapping[str, Any], **kwargs) -> Mapping[str, Any]:\n        \"\"\"\n        Run model inference on specified data batch.\n        .. warning::\n            You should not override this method. If you need specific model\n            call, override forward() method\n        Args:\n            batch: dictionary with data batch from DataLoader.\n            **kwargs: additional kwargs to pass to the model\n        Returns:\n            Mapping[str, Any]: model output dictionary\n        \"\"\"\n        batch = self._process_batch(batch)\n        output = self.forward(batch, **kwargs)\n        return output\n\n    def get_callbacks(self) -> \"OrderedDict[str, Callback]\":\n        \"\"\"Prepares the callbacks for selected stage.\n        Args:\n            stage: stage name\n        Returns:\n            dictionary with stage callbacks\n        \"\"\"\n        callbacks = super().get_callbacks()\n        callback_exists = lambda callback_fn: any(\n            callback_isinstance(x, callback_fn) for x in callbacks.values()\n        )\n        if isinstance(self._criterion, TorchCriterion) and not callback_exists(\n            ICriterionCallback\n        ):\n            callbacks[\"_criterion\"] = CriterionCallback(\n                input_key=f\"{self.loss_mode_prefix}_left\",\n                target_key=f\"{self.loss_mode_prefix}_right\",\n                metric_key=self._loss_key,\n            )\n        if isinstance(self._optimizer, TorchOptimizer) and not callback_exists(\n            IBackwardCallback\n        ):\n            callbacks[\"_backward\"] = BackwardCallback(metric_key=self._loss_key)\n        if isinstance(self._optimizer, TorchOptimizer) and not callback_exists(\n            IOptimizerCallback\n        ):\n            callbacks[\"_optimizer\"] = OptimizerCallback(metric_key=self._loss_key)\n        if isinstance(self._scheduler, TorchScheduler) and not callback_exists(\n            ISchedulerCallback\n        ):\n            callbacks[\"_scheduler\"] = SchedulerCallback(\n                loader_key=self._valid_loader, metric_key=self._valid_metric\n            )\n        return callbacks\n\n\n__all__ = [\"ISelfSupervisedRunner\", \"SelfSupervisedRunner\"]\n\n\nclass ContrastiveModel(torch.nn.Module):\n    \"\"\"Contrastive model with projective head.\n    Args:\n        model: projective head for the train time\n        encoder: model for the future uses\n    \"\"\"\n\n    def __init__(self, model, encoder):\n        super(ContrastiveModel, self).__init__()\n        self.model = model\n        self.encoder = encoder\n\n    def forward(self, x):\n        \"\"\"Forward method.\n        Args:\n            x: input for the encoder\n        Returns:\n            (embeddings, projections)\n        \"\"\"\n        emb = self.encoder(x)\n        projection = self.model(emb)\n        return emb, projection","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"feature_dim = 100\nhidden_dim = 512\nencoder_dim = 512\nmodel_name = 'tf_efficientnet_b7_ns'\nmodel = Model(model_name, pretrained=True, output_dim=hidden_dim)\n_ = model.to(device)\n\nprojection_head = nn.Sequential(\n#         nn.Linear(encoder_dim, hidden_dim, bias=False),\n        nn.ReLU(inplace=True),\n        nn.Linear(hidden_dim, feature_dim, bias=True),\n    )\nprojection_head.to(device)\n\ncmodel = ContrastiveModel(projection_head, encoder = model)\n_ = cmodel.to(device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# m = timm.create_model(model_name, pretrained=True, in_chans=2)\n\n\n# clsf = m.default_cfg['classifier']\n# n_features = m._modules[clsf].in_features\n\n# clsf","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from catalyst import dl, utils","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"callbacks = [\n        dl.CriterionCallback(\n            input_key=\"target_projection_left\",\n            target_key=\"target_projection_right\",\n            metric_key=\"loss\",\n        ),\n        dl.BackwardCallback(metric_key=\"loss\"),\n        dl.OptimizerCallback(metric_key=\"loss\"),\n#         dl.ControlFlowCallbackWrapper(\n#             dl.SoftUpdateCallaback(\n#                 target_model=\"target\",\n#                 source_model=\"online\",\n#                 tau=0.1,\n#                 scope=\"on_batch_end\",\n#             ),\n#             loaders=\"train\",\n#         ),\n#         dl.ControlFlowCallbackWrapper(\n#             dl.AccuracyCallback(\n#                 target_key=\"target\", input_key=\"sklearn_predict\", topk=(1, 3)\n#             ),\n#             loaders=\"valid\",\n#         ),\n    ]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SSRunner(SelfSupervisedRunner):\n    def handle_batch(self, batch: Mapping[str, Any]):\n        for key in self.batch:\n            if 'augment' in key:\n#                 print(key, \":\", self.batch[key].shape)\n                self.batch[key] = self.batch[key].swapaxes(1, 3).swapaxes(1, 2)\n#                 print(key, \":\", self.batch[key].shape)\n#                 gc.collect()\n        self.batch = {**batch, **self.forward(batch)}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\nfrom catalyst.contrib.losses import NTXentLoss\nepochs = 40\nbatch_size = 9\nnum_workers = 2\nweight_decay = 1e-6\nmax_grad_norm = 1000\n\ncriterion = NTXentLoss(tau=0.5)\ndataset = Dataset('test', df2)\ncdataset = SelfSupervisedDatasetWrapper(dataset, transforms = transform)\nloader_train = torch.utils.data.DataLoader(cdataset, batch_size=batch_size,\n                     num_workers=num_workers, pin_memory=True, shuffle=True, drop_last=True)\n\n# dataset = Dataset('test', df2.tail(300))\n# cdataset = SelfSupervisedDatasetWrapper(dataset, transforms = transform)\n# loader_val = torch.utils.data.DataLoader(cdataset, batch_size=batch_size,\n#                      num_workers=num_workers, pin_memory=True, shuffle=True, drop_last=True)\nlr_max = 4e-4\nepochs_warmup = 1.0\ncmodels = nn.ModuleDict(\n        {\n#             \"online\": cmodel,\n            \"target\": cmodel\n        }\n    )\noptimizer = torch.optim.Adam(cmodel.parameters(), lr=lr_max, weight_decay=weight_decay)\ngc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"runner = SSRunner()\nrunner.train(\n        model=cmodels,\n        criterion=criterion,\n        optimizer=optimizer,\n        callbacks=callbacks,\n        loaders={'train' : loader_train},\n        num_epochs=epochs,\n        logdir=\"./\",\n        valid_loader=\"train\",\n        valid_metric=\"loss\",\n        minimize_valid_metric=True,\n        verbose=False,\n        # check=args.check,\n    )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"markdown","source":"# Predict and evaluate","metadata":{}},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"# model_name = 'tf_efficientnet_b7_ns'\n# nfold = 5\n# kfold = StratifiedKFold(n_splits=nfold, random_state=42, shuffle=True)\n\n# epochs = 16\n# batch_size = 16\n# num_workers = 2\n# weight_decay = 1e-6\n# max_grad_norm = 1000\n\n# lr_max = 4e-4\n# epochs_warmup = 1.0\n\n# for ifold, (idx_train, idx_test) in enumerate(kfold.split(dataset, df['target'])):\n#     print('Fold %d/%d' % (ifold, nfold))\n#     torch.manual_seed(42 + ifold + 1)\n\n#     # Train - val split\n#     dataset_train = Dataset('train', df.iloc[idx_train])\n#     dataset_val = Dataset('train', df.iloc[idx_test])\n \n#     loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size,\n#                      num_workers=num_workers, pin_memory=True, shuffle=True, drop_last=True)\n#     loader_val = torch.utils.data.DataLoader(dataset_val, batch_size=batch_size,\n#                      num_workers=num_workers, pin_memory=True)\n\n#     # Model and optimizer\n#     model = Model(model_name, pretrained=True)\n#     model.to(device)\n#     model.train()\n\n#     optimizer = torch.optim.Adam(model.parameters(), lr=lr_max, weight_decay=weight_decay)\n\n#     # Learning-rate schedule\n#     nbatch = len(loader_train)\n#     warmup = epochs_warmup * nbatch  # number of warmup steps\n#     nsteps = epochs * nbatch        # number of total steps\n\n#     scheduler = CosineLRScheduler(optimizer,\n#                   warmup_t=warmup, warmup_lr_init=0.0, warmup_prefix=True, # 1 epoch of warmup\n#                   t_initial=(nsteps - warmup), lr_min=1e-6)                # 3 epochs of cosine\n    \n#     time_val = 0.0\n#     lrs = []\n\n#     tb = time.time()\n#     print('Epoch   loss          score   lr')\n#     for iepoch in range(epochs):\n#         loss_sum = 0.0\n#         n_sum = 0\n\n#         # Train\n#         for ibatch, (img, y) in enumerate(loader_train):\n#             n = y.size(0)\n#             img = img.to(device)\n#             y = y.to(device)\n\n#             optimizer.zero_grad()\n\n#             y_pred = model(img)\n#             loss = criterion(y_pred.view(-1), y)\n\n#             loss_train = loss.item()\n#             loss_sum += n * loss_train\n#             n_sum += n\n\n#             loss.backward()\n\n#             grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),\n#                                                        max_grad_norm)\n#             optimizer.step()\n            \n#             scheduler.step(iepoch * nbatch + ibatch + 1)\n#             lrs.append(optimizer.param_groups[0]['lr'])            \n\n#         # Evaluate\n#         val = evaluate(model, loader_val)\n#         time_val += val['time']\n#         loss_train = loss_sum / n_sum\n#         lr_now = optimizer.param_groups[0]['lr']\n#         dt = (time.time() - tb) / 60\n#         print('Epoch %d %.4f %.4f %.4f  %.2e  %.2f min' %\n#               (iepoch + 1, loss_train, val['loss'], val['score'], lr_now, dt))\n\n#     dt = time.time() - tb\n#     print('Training done %.2f min total, %.2f min val' % (dt / 60, time_val / 60))\n\n#     # Save model\n#     ofilename = 'model%d.pytorch' % ifold\n#     torch.save(model.state_dict(), ofilename)\n#     print(ofilename, 'written')\n\n#     break  # 1 fold only","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# plt.title('LR Schedule: Cosine with linear warmup')\n# plt.xlabel('steps')\n# plt.ylabel('learning rate')\n# plt.plot(lrs)\n# plt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Predict and submit","metadata":{}},{"cell_type":"code","source":"# # Load model (if necessary)\n# model = Model(model_name, pretrained=False)\n# filename = 'model0.pytorch'\n# model.to(device)\n# model.load_state_dict(torch.load(filename, map_location=device))\n# model.eval()\n\n# # Predict\n# submit = pd.read_csv(di + '/sample_submission.csv')\n# dataset_test = Dataset('test', submit)\n# loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=64,\n#                                           num_workers=num_workers, pin_memory=True)\n\n# test = evaluate(model, loader_test, compute_score=False, pbar=len(submit))\n\n# # Write prediction\n# submit['target'] = test['y_pred']\n# submit.to_csv('submission.csv', index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# print('target range [%.2f, %.2f]' % (submit['target'].min(), submit['target'].max()))","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}