{"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":"## About\n\nIn this notebook, I will show a way to load data from TFRecord, and train a model implemented in PyTorch. In this way, you can get some advantages like\n\n1. You can speed up your training. TFRecord will allow you to load data faster than reading from npy file in your PyTorch Dataset and DataLoader.\n2. If you are using Google Colab (Pro/Pro+), you don't need to put your data on Google Drive. What you only need is the GCS path of the TFRecord dataset, which can be obtained through kaggle_datasets API.\n\nOf course, Tensorflow + TPU can speed up training even more, but some of you may want to use PyTorch maybe because it has `timm` or maybe because you are more used to use PyTorch that Tensorflow. This notebook is for those people.\n\nThe whole pipeline is based on [Y.Nakama's notebook](https://www.kaggle.com/yasufuminakama/g2net-efficientnet-b7-baseline-training).","metadata":{}},{"cell_type":"code","source":"!pip install -q nnAudio\n!pip install -q ../input/pytorch-image-models/pytorch-image-models-master/","metadata":{"execution":{"iopub.status.busy":"2021-08-22T07:45:18.889404Z","iopub.execute_input":"2021-08-22T07:45:18.889768Z","iopub.status.idle":"2021-08-22T07:45:36.722737Z","shell.execute_reply.started":"2021-08-22T07:45:18.889689Z","shell.execute_reply":"2021-08-22T07:45:36.721708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Libraries","metadata":{"papermill":{"duration":0.014815,"end_time":"2021-05-11T16:05:42.265815","exception":false,"start_time":"2021-05-11T16:05:42.251","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import os\nimport time\nimport math\nimport random\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport scipy as sp\nimport tensorflow as tf  # for reading TFRecord Dataset\nimport tensorflow_datasets as tfds  # for making tf.data.Dataset to return numpy arrays\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nimport timm\nfrom kaggle_datasets import KaggleDatasets\nfrom nnAudio.Spectrogram import CQT1992v2\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import KFold\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2021-08-22T08:10:09.927347Z","iopub.execute_input":"2021-08-22T08:10:09.927692Z","iopub.status.idle":"2021-08-22T08:10:09.934347Z","shell.execute_reply.started":"2021-08-22T08:10:09.927661Z","shell.execute_reply":"2021-08-22T08:10:09.933398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SAVEDIR = Path(\"./\")\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2021-08-22T07:45:44.968455Z","iopub.execute_input":"2021-08-22T07:45:44.968778Z","iopub.status.idle":"2021-08-22T07:45:45.020546Z","shell.execute_reply.started":"2021-08-22T07:45:44.968748Z","shell.execute_reply":"2021-08-22T07:45:45.019474Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CFG","metadata":{}},{"cell_type":"code","source":"class CFG:\n    debug = False\n    print_freq = 50\n    num_workers = 4\n    model_name = \"tf_efficientnet_b0_ns\"\n    qtransform_params = {\"sr\": 2048, \"fmin\": 20, \"fmax\": 1024, \"hop_length\": 32, \"bins_per_octave\": 8}\n    scheduler = \"CosineAnnealingLR\"\n    epochs = 3\n    T_max = 3\n    lr = 1e-4\n    min_lr = 1e-7\n    batch_size = 64\n    weight_decay = 1e-3\n    gradient_accumulation_steps = 1\n    max_grad_norm = 1000\n    seed = 42\n    target_size = 1\n    target_col = \"target\"\n    n_fold = 5\n    trn_fold = [0]  # [0, 1, 2, 3, 4]\n    train = True\n\nif CFG.debug:\n    CFG.epochs = 1","metadata":{"execution":{"iopub.status.busy":"2021-08-22T07:45:47.258304Z","iopub.execute_input":"2021-08-22T07:45:47.258626Z","iopub.status.idle":"2021-08-22T07:45:47.264851Z","shell.execute_reply.started":"2021-08-22T07:45:47.258595Z","shell.execute_reply":"2021-08-22T07:45:47.263882Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utils","metadata":{"papermill":{"duration":0.028465,"end_time":"2021-05-11T16:05:49.544956","exception":false,"start_time":"2021-05-11T16:05:49.516491","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# ====================================================\n# Utils\n# ====================================================\ndef get_score(y_true, y_pred):\n    score = roc_auc_score(y_true, y_pred)\n    return score\n\n\ndef init_logger(log_file=SAVEDIR / 'train.log'):\n    from logging import getLogger, INFO, FileHandler,  Formatter,  StreamHandler\n    logger = getLogger(__name__)\n    logger.setLevel(INFO)\n    handler1 = StreamHandler()\n    handler1.setFormatter(Formatter(\"%(message)s\"))\n    handler2 = FileHandler(filename=log_file)\n    handler2.setFormatter(Formatter(\"%(message)s\"))\n    logger.addHandler(handler1)\n    logger.addHandler(handler2)\n    return logger\n\nLOGGER = init_logger()\n\n\ndef seed_torch(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.backends.cudnn.deterministic = True\n\nseed_torch(seed=CFG.seed)","metadata":{"papermill":{"duration":0.041144,"end_time":"2021-05-11T16:05:49.614272","exception":false,"start_time":"2021-05-11T16:05:49.573128","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-08-22T07:45:49.687577Z","iopub.execute_input":"2021-08-22T07:45:49.687911Z","iopub.status.idle":"2021-08-22T07:45:49.700080Z","shell.execute_reply.started":"2021-08-22T07:45:49.687879Z","shell.execute_reply":"2021-08-22T07:45:49.698663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TFRecord Loader\n\nThis is the heart of this notebook. Instead of using PyTorch's Dataset and DataLoader, here I define custom Loader that reads samples from TFRecords.\n\nFYI, there's a library that does the same thing, but its implementation is not optimized, so it's slower.\n\nhttps://github.com/vahidk/tfrecord","metadata":{}},{"cell_type":"code","source":"gcs_paths = []\nfor i, j in [(0, 4), (5, 9), (10, 14), (15, 19)]:\n    path = f\"g2net-waveform-tfrecords-train-{i}-{j}\"\n    n_trial = 0\n    while True:\n        try:\n            gcs_path = KaggleDatasets().get_gcs_path(path)\n            gcs_paths.append(gcs_path)\n            print(gcs_path)\n            break\n        except:\n            if n_trial > 10:\n                break\n            n_trial += 1\n            continue\n            \nall_files = []\nfor path in gcs_paths:\n    all_files.extend(np.sort(np.array(tf.io.gfile.glob(path + \"/train*.tfrecords\"))))\n    \nprint(\"train_files: \", len(all_files))\nall_files = np.array(all_files)","metadata":{"execution":{"iopub.status.busy":"2021-08-22T07:51:25.548151Z","iopub.execute_input":"2021-08-22T07:51:25.548478Z","iopub.status.idle":"2021-08-22T07:51:27.826545Z","shell.execute_reply.started":"2021-08-22T07:51:25.548449Z","shell.execute_reply":"2021-08-22T07:51:27.825532Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def count_data_items(fileids):\n    \"\"\"\n    Count the number of samples.\n    Each of the TFRecord datasets is designed to contain 28000 samples.\n    \"\"\"\n    return len(fileids) * 28000\n\n\nAUTO = tf.data.experimental.AUTOTUNE","metadata":{"execution":{"iopub.status.busy":"2021-08-22T07:53:24.992593Z","iopub.execute_input":"2021-08-22T07:53:24.992950Z","iopub.status.idle":"2021-08-22T07:53:24.996865Z","shell.execute_reply.started":"2021-08-22T07:53:24.992903Z","shell.execute_reply":"2021-08-22T07:53:24.996057Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def prepare_wave(wave):\n    wave = tf.reshape(tf.io.decode_raw(wave, tf.float64), (3, 4096))\n    normalized_waves = []\n    for i in range(3):\n        normalized_wave = wave[i] / tf.math.reduce_max(wave[i])\n        normalized_waves.append(normalized_wave)\n    wave = tf.stack(normalized_waves, axis=0)\n    wave = tf.cast(wave, tf.float32)\n    return wave\n\n\ndef read_labeled_tfrecord(example):\n    tfrec_format = {\n        \"wave\": tf.io.FixedLenFeature([], tf.string),\n        \"wave_id\": tf.io.FixedLenFeature([], tf.string),\n        \"target\": tf.io.FixedLenFeature([], tf.int64)\n    }\n    example = tf.io.parse_single_example(example, tfrec_format)\n    return prepare_wave(example[\"wave\"]), tf.reshape(tf.cast(example[\"target\"], tf.float32), [1]), example[\"wave_id\"]\n\n\ndef read_unlabeled_tfrecord(example, return_image_id):\n    tfrec_format = {\n        \"wave\": tf.io.FixedLenFeature([], tf.string),\n        \"wave_id\": tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrec_format)\n    return prepare_wave(example[\"wave\"]), example[\"wave_id\"] if return_image_id else 0\n\n\ndef get_dataset(files, batch_size=16, repeat=False, cache=False, shuffle=False, labeled=True, return_image_ids=True):\n    ds = tf.data.TFRecordDataset(files, num_parallel_reads=AUTO, compression_type=\"GZIP\")\n    if cache:\n        # You'll need around 15GB RAM if you'd like to cache val dataset, and 50~60GB RAM for train dataset.\n        ds = ds.cache()\n\n    if repeat:\n        ds = ds.repeat()\n\n    if shuffle:\n        ds = ds.shuffle(1024 * 2)\n        opt = tf.data.Options()\n        opt.experimental_deterministic = False\n        ds = ds.with_options(opt)\n\n    if labeled:\n        ds = ds.map(read_labeled_tfrecord, num_parallel_calls=AUTO)\n    else:\n        ds = ds.map(lambda example: read_unlabeled_tfrecord(example, return_image_ids), num_parallel_calls=AUTO)\n\n    ds = ds.batch(batch_size)\n    ds = ds.prefetch(AUTO)\n    return tfds.as_numpy(ds)","metadata":{"execution":{"iopub.status.busy":"2021-08-22T08:11:32.069988Z","iopub.execute_input":"2021-08-22T08:11:32.070359Z","iopub.status.idle":"2021-08-22T08:11:32.083200Z","shell.execute_reply.started":"2021-08-22T08:11:32.070329Z","shell.execute_reply":"2021-08-22T08:11:32.082257Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TFRecordDataLoader:\n    def __init__(self, files, batch_size=32, cache=False, train=True, repeat=False, shuffle=False, labeled=True, return_image_ids=True):\n        self.ds = get_dataset(\n            files, \n            batch_size=batch_size,\n            cache=cache,\n            repeat=repeat,\n            shuffle=shuffle,\n            labeled=labeled,\n            return_image_ids=return_image_ids)\n        \n        if train:\n            self.num_examples = count_data_items(files)\n        else:\n            self.num_examples = count_data_items_test(files)\n\n        self.batch_size = batch_size\n        self.labeled = labeled\n        self.return_image_ids = return_image_ids\n        self._iterator = None\n    \n    def __iter__(self):\n        if self._iterator is None:\n            self._iterator = iter(self.ds)\n        else:\n            self._reset()\n        return self._iterator\n\n    def _reset(self):\n        self._iterator = iter(self.ds)\n\n    def __next__(self):\n        batch = next(self._iterator)\n        return batch\n\n    def __len__(self):\n        n_batches = self.num_examples // self.batch_size\n        if self.num_examples % self.batch_size == 0:\n            return n_batches\n        else:\n            return n_batches + 1","metadata":{"execution":{"iopub.status.busy":"2021-08-22T07:55:59.008667Z","iopub.execute_input":"2021-08-22T07:55:59.009025Z","iopub.status.idle":"2021-08-22T07:55:59.018050Z","shell.execute_reply.started":"2021-08-22T07:55:59.008992Z","shell.execute_reply":"2021-08-22T07:55:59.017144Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## MODEL","metadata":{"papermill":{"duration":0.029891,"end_time":"2021-05-11T16:05:50.085695","exception":false,"start_time":"2021-05-11T16:05:50.055804","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# ====================================================\n# MODEL\n# ====================================================\nclass CustomModel(nn.Module):\n    def __init__(self, cfg, pretrained=False):\n        super().__init__()\n        self.cfg = cfg\n        self.wave_transform = CQT1992v2(**CFG.qtransform_params)\n        self.model = timm.create_model(self.cfg.model_name, pretrained=pretrained, in_chans=3)\n        self.n_features = self.model.classifier.in_features\n        self.model.classifier = nn.Linear(self.n_features, self.cfg.target_size)\n\n    def forward(self, x):\n        waves = []\n        for i in range(3):\n            waves.append(self.wave_transform(x[:, i]))\n        x = torch.stack(waves, dim=1)\n        output = self.model(x)\n        return output","metadata":{"papermill":{"duration":0.037539,"end_time":"2021-05-11T16:05:50.153617","exception":false,"start_time":"2021-05-11T16:05:50.116078","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-08-22T07:57:46.799678Z","iopub.execute_input":"2021-08-22T07:57:46.800015Z","iopub.status.idle":"2021-08-22T07:57:46.806742Z","shell.execute_reply.started":"2021-08-22T07:57:46.799982Z","shell.execute_reply":"2021-08-22T07:57:46.805896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Helper functions","metadata":{"papermill":{"duration":0.029417,"end_time":"2021-05-11T16:05:50.211711","exception":false,"start_time":"2021-05-11T16:05:50.182294","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\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\n\ndef asMinutes(s):\n    m = math.floor(s / 60)\n    s -= m * 60\n    return '%dm %ds' % (m, s)\n\n\ndef timeSince(since, percent):\n    now = time.time()\n    s = now - since\n    es = s / (percent)\n    rs = es - s\n    return '%s (remain %s)' % (asMinutes(s), asMinutes(rs))\n\n\ndef max_memory_allocated():\n    MB = 1024.0 * 1024.0\n    mem = torch.cuda.max_memory_allocated() / MB\n    return f\"{mem:.0f} MB\"","metadata":{"execution":{"iopub.status.busy":"2021-08-22T07:58:27.741704Z","iopub.execute_input":"2021-08-22T07:58:27.742042Z","iopub.status.idle":"2021-08-22T07:58:27.750000Z","shell.execute_reply.started":"2021-08-22T07:58:27.742011Z","shell.execute_reply":"2021-08-22T07:58:27.749031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Trainer","metadata":{}},{"cell_type":"code","source":"def train_fn(files, model, criterion, optimizer, epoch, scheduler, device):\n    batch_time = AverageMeter()\n    data_time = AverageMeter()\n    losses = AverageMeter()\n    scores = AverageMeter()\n\n    # switch to train mode\n    model.train()\n    start = end = time.time()\n    global_step = 0\n\n    train_loader = TFRecordDataLoader(\n        files, batch_size=CFG.batch_size, shuffle=True)\n    for step, d in enumerate(train_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n        images = torch.from_numpy(d[0]).to(device)\n        labels = torch.from_numpy(d[1]).to(device)\n\n        batch_size = labels.size(0)\n        y_preds = model(images)\n        loss = criterion(y_preds.view(-1), labels.view(-1))\n        # record loss\n        losses.update(loss.item(), batch_size)\n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n        loss.backward()\n        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), CFG.max_grad_norm)\n        if (step + 1) % CFG.gradient_accumulation_steps == 0:\n            optimizer.step()\n            optimizer.zero_grad()\n            global_step += 1\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n        if step % CFG.print_freq == 0:\n            print('Epoch: [{0}/{1}][{2}/{3}] '\n                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '\n                  'Grad: {grad_norm:.4f}  '\n                  'LR: {lr:.6f}  '\n                  'Elapsed: {remain:s} '\n                  'Max mem: {mem:s}'\n                  .format(\n                   epoch+1, CFG.epochs, step, len(train_loader),\n                   loss=losses,\n                   grad_norm=grad_norm,\n                   lr=scheduler.get_lr()[0],\n                   remain=timeSince(start, float(step + 1) / len(train_loader)),\n                   mem=max_memory_allocated()))\n    return losses.avg\n\n\ndef valid_fn(files, model, criterion, device):\n    batch_time = AverageMeter()\n    data_time = AverageMeter()\n    losses = AverageMeter()\n    scores = AverageMeter()\n    # switch to evaluation mode\n    model.eval()\n    filenames = []\n    targets = []\n    preds = []\n    start = end = time.time()\n    valid_loader = TFRecordDataLoader(\n        files, batch_size=CFG.batch_size * 2, shuffle=False)\n    for step, d in enumerate(valid_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n        \n        targets.extend(d[1].reshape(-1).tolist())\n        filenames.extend([f.decode(\"UTF-8\") for f in d[2]])\n        \n        images = torch.from_numpy(d[0]).to(device)\n        labels = torch.from_numpy(d[1]).to(device)\n\n        batch_size = labels.size(0)\n        # compute loss\n        with torch.no_grad():\n            y_preds = model(images)\n        loss = criterion(y_preds.view(-1), labels.view(-1))\n        losses.update(loss.item(), batch_size)\n\n        preds.append(y_preds.sigmoid().to('cpu').numpy())\n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n        if step % CFG.print_freq == 0:\n            print('EVAL: [{0}/{1}] '\n                  'Data {data_time.val:.3f} ({data_time.avg:.3f}) '\n                  'Elapsed {remain:s} '\n                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '\n                  .format(\n                   step, len(valid_loader), batch_time=batch_time,\n                   data_time=data_time, loss=losses,\n                   remain=timeSince(start, float(step+1)/len(valid_loader)),\n                   ))\n    predictions = np.concatenate(preds).reshape(-1)\n    return losses.avg, predictions, np.array(targets), np.array(filenames)","metadata":{"papermill":{"duration":0.053284,"end_time":"2021-05-11T16:05:50.294128","exception":false,"start_time":"2021-05-11T16:05:50.240844","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-08-22T08:43:00.412120Z","iopub.execute_input":"2021-08-22T08:43:00.412458Z","iopub.status.idle":"2021-08-22T08:43:00.432494Z","shell.execute_reply.started":"2021-08-22T08:43:00.412430Z","shell.execute_reply":"2021-08-22T08:43:00.431494Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train loop","metadata":{"papermill":{"duration":0.030241,"end_time":"2021-05-11T16:05:50.353874","exception":false,"start_time":"2021-05-11T16:05:50.323633","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# ====================================================\n# Train loop\n# ====================================================\ndef train_loop(train_tfrecords: np.ndarray, val_tfrecords: np.ndarray, fold: int):\n    \n    LOGGER.info(f\"========== fold: {fold} training ==========\")\n    \n    # ====================================================\n    # scheduler \n    # ====================================================\n    def get_scheduler(optimizer):\n        if CFG.scheduler=='ReduceLROnPlateau':\n            scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=CFG.factor, patience=CFG.patience, verbose=True, eps=CFG.eps)\n        elif CFG.scheduler=='CosineAnnealingLR':\n            scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=CFG.T_max, eta_min=CFG.min_lr, last_epoch=-1)\n        elif CFG.scheduler=='CosineAnnealingWarmRestarts':\n            scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=CFG.T_0, T_mult=1, eta_min=CFG.min_lr, last_epoch=-1)\n        return scheduler\n\n    # ====================================================\n    # model & optimizer\n    # ====================================================\n    model = CustomModel(CFG, pretrained=True)\n    model.to(device)\n\n    optimizer = optim.Adam(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay, amsgrad=False)\n    scheduler = get_scheduler(optimizer)\n\n    # ====================================================\n    # loop\n    # ====================================================\n    criterion = nn.BCEWithLogitsLoss()\n\n    best_score = 0.\n    best_loss = np.inf\n    \n    for epoch in range(CFG.epochs):\n        \n        start_time = time.time()\n        \n        # train\n        avg_loss = train_fn(train_tfrecords, model, criterion, optimizer, epoch, scheduler, device)\n\n        # eval\n        avg_val_loss, preds, targets, files = valid_fn(val_tfrecords, model, criterion, device)\n        valid_result_df = pd.DataFrame({\"target\": targets, \"preds\": preds, \"id\": files})\n        \n        if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):\n            scheduler.step(avg_val_loss)\n        elif isinstance(scheduler, optim.lr_scheduler.CosineAnnealingLR):\n            scheduler.step()\n        elif isinstance(scheduler, optim.lr_scheduler.CosineAnnealingWarmRestarts):\n            scheduler.step()\n\n        # scoring\n        score = get_score(targets, preds)\n\n        elapsed = time.time() - start_time\n\n        LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  avg_val_loss: {avg_val_loss:.4f}  time: {elapsed:.0f}s')\n        LOGGER.info(f'Epoch {epoch+1} - Score: {score:.4f}')\n\n        if score > best_score:\n            best_score = score\n            LOGGER.info(f'Epoch {epoch+1} - Save Best Score: {best_score:.4f} Model')\n            torch.save({'model': model.state_dict(), \n                        'preds': preds},\n                        SAVEDIR / f'{CFG.model_name}_fold{fold}_best_score.pth')\n        \n        if avg_val_loss < best_loss:\n            best_loss = avg_val_loss\n            LOGGER.info(f'Epoch {epoch+1} - Save Best Loss: {best_loss:.4f} Model')\n            torch.save({'model': model.state_dict(), \n                        'preds': preds},\n                        SAVEDIR / f'{CFG.model_name}_fold{fold}_best_loss.pth')\n    \n    valid_result_df[\"preds\"] = torch.load(SAVEDIR / f\"{CFG.model_name}_fold{fold}_best_loss.pth\",\n                                          map_location=\"cpu\")[\"preds\"]\n\n    return valid_result_df","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":0.049518,"end_time":"2021-05-11T16:05:50.433341","exception":false,"start_time":"2021-05-11T16:05:50.383823","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-08-22T09:23:07.020406Z","iopub.execute_input":"2021-08-22T09:23:07.020739Z","iopub.status.idle":"2021-08-22T09:23:07.035890Z","shell.execute_reply.started":"2021-08-22T09:23:07.020705Z","shell.execute_reply":"2021-08-22T09:23:07.034918Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def main():\n    def get_result(result_df):\n        preds = result_df['preds'].values\n        labels = result_df[CFG.target_col].values\n        score = get_score(labels, preds)\n        LOGGER.info(f'Score: {score:<.4f}')\n    \n    if CFG.train:\n        # train \n        oof_df = pd.DataFrame()\n        kf = KFold(n_splits=CFG.n_fold, shuffle=True, random_state=CFG.seed)\n\n        folds = list(kf.split(all_files))\n        for fold in range(CFG.n_fold):\n            if fold in CFG.trn_fold:\n                trn_idx, val_idx = folds[fold]\n                train_files = all_files[trn_idx]\n                valid_files = all_files[val_idx]\n                _oof_df = train_loop(train_files, valid_files, fold)\n                oof_df = pd.concat([oof_df, _oof_df])\n                LOGGER.info(f\"========== fold: {fold} result ==========\")\n                get_result(_oof_df)\n        # CV result\n        LOGGER.info(f\"========== CV ==========\")\n        get_result(oof_df)\n        # save result\n        oof_df.to_csv(SAVEDIR / 'oof_df.csv', index=False)","metadata":{"papermill":{"duration":0.038415,"end_time":"2021-05-11T16:05:50.501397","exception":false,"start_time":"2021-05-11T16:05:50.462982","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-08-22T09:23:07.971552Z","iopub.execute_input":"2021-08-22T09:23:07.971860Z","iopub.status.idle":"2021-08-22T09:23:07.979729Z","shell.execute_reply.started":"2021-08-22T09:23:07.971832Z","shell.execute_reply":"2021-08-22T09:23:07.978694Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if __name__ == '__main__':\n    main()","metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","papermill":{"duration":13853.248037,"end_time":"2021-05-11T19:56:43.779136","exception":false,"start_time":"2021-05-11T16:05:50.531099","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-08-22T09:23:08.684900Z","iopub.execute_input":"2021-08-22T09:23:08.685250Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}