{"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":"## Summary\n\n- Vision Transformer in PyTorch: a smaller and more flexible version than https://www.kaggle.com/scaomath/g2net-vision-transformer-starter, coded by the mighty lucidrains.\n- Caveat: dimension is hard-coded.\n\n## Reference\n- pipeline: modified from [Y.Nakama's notebook](https://www.kaggle.com/yasufuminakama/g2net-efficientnet-b7-baseline-training).\n- dataset: @hidehisaarai1213 https://www.kaggle.com/hidehisaarai1213/g2net-read-from-tfrecord-train-with-pytorch","metadata":{}},{"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":"!pip install -q vit-pytorch\n!pip install -q nnAudio","metadata":{"execution":{"iopub.status.busy":"2021-09-15T18:31:04.454544Z","iopub.execute_input":"2021-09-15T18:31:04.454934Z","iopub.status.idle":"2021-09-15T18:31:18.382995Z","shell.execute_reply.started":"2021-09-15T18:31:04.454869Z","shell.execute_reply":"2021-09-15T18:31:18.381830Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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\nfrom scipy import signal\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\nfrom nnAudio.Spectrogram import CQT1992v2\nfrom vit_pytorch import ViT\nfrom kaggle_datasets import KaggleDatasets\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import KFold\nfrom tqdm.auto import tqdm","metadata":{"execution":{"iopub.status.busy":"2021-09-15T18:31:18.384662Z","iopub.execute_input":"2021-09-15T18:31:18.385032Z","iopub.status.idle":"2021-09-15T18:31:25.396821Z","shell.execute_reply.started":"2021-09-15T18:31:18.384986Z","shell.execute_reply":"2021-09-15T18:31:25.395894Z"},"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-09-15T18:31:25.398456Z","iopub.execute_input":"2021-09-15T18:31:25.398760Z","iopub.status.idle":"2021-09-15T18:31:25.445228Z","shell.execute_reply.started":"2021-09-15T18:31:25.398730Z","shell.execute_reply":"2021-09-15T18:31:25.444117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CFG","metadata":{}},{"cell_type":"code","source":"class CFG:\n    debug = False\n    print_freq = 3000\n    num_workers = 4\n    scheduler = \"CosineAnnealingLR\"\n    model_name = \"vit\"\n    epochs = 6\n    T_max = 3\n    lr = 1e-4\n    min_lr = 1e-7\n    batch_size = 50\n    val_batch_size = 100\n    weight_decay = 1e-5\n    gradient_accumulation_steps = 1\n    max_grad_norm = 1000\n    seed = 1127802825\n    target_size = 1\n    target_col = \"target\"\n    n_fold = 5\n    trn_fold = [3]  # [0, 1, 2, 3, 4]\n    train = True\n    n_hidden = 256\n    n_heads = 4\n    n_layers = 12\n    image_size = (128, 256)\n    patch_size = (8, 16)\n    cqt_params = {\"sr\": 2048, \n                  \"fmin\": 20, \"fmax\": 530, \n                  \"hop_length\": 4,\n                  'filter_scale':0.5,\n                  \"bins_per_octave\": 27}","metadata":{"execution":{"iopub.status.busy":"2021-09-15T18:40:54.299229Z","iopub.execute_input":"2021-09-15T18:40:54.299550Z","iopub.status.idle":"2021-09-15T18:40:54.306375Z","shell.execute_reply.started":"2021-09-15T18:40:54.299523Z","shell.execute_reply":"2021-09-15T18:40:54.305301Z"},"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-09-15T18:31:33.862686Z","iopub.execute_input":"2021-09-15T18:31:33.863005Z","iopub.status.idle":"2021-09-15T18:31:33.874901Z","shell.execute_reply.started":"2021-09-15T18:31:33.862976Z","shell.execute_reply":"2021-09-15T18:31:33.873925Z"},"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-09-15T18:31:34.964667Z","iopub.execute_input":"2021-09-15T18:31:34.965048Z","iopub.status.idle":"2021-09-15T18:31:36.955383Z","shell.execute_reply.started":"2021-09-15T18:31:34.965016Z","shell.execute_reply":"2021-09-15T18:31:36.954001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def count_data_items(fileids, train=True):\n    \"\"\"\n    Count the number of samples.\n    Each of the TFRecord datasets is designed to contain 28000 samples for train\n    22500 for test.\n    \"\"\"\n    sizes = 28000 if train else 22500\n    return len(fileids) * sizes\n\n\nAUTO = tf.data.experimental.AUTOTUNE","metadata":{"execution":{"iopub.status.busy":"2021-09-15T18:31:36.956857Z","iopub.execute_input":"2021-09-15T18:31:36.957232Z","iopub.status.idle":"2021-09-15T18:31:36.961868Z","shell.execute_reply.started":"2021-09-15T18:31:36.957193Z","shell.execute_reply":"2021-09-15T18:31:36.961051Z"},"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    scaling = tf.constant([1.5e-20, 1.5e-20, 0.5e-20], dtype=tf.float64)\n    for i in range(3):\n#         normalized_wave = wave[i] / tf.math.reduce_max(wave[i])\n        normalized_wave = wave[i] / scaling[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, \n                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-09-15T18:31:36.963518Z","iopub.execute_input":"2021-09-15T18:31:36.964057Z","iopub.status.idle":"2021-09-15T18:31:36.979145Z","shell.execute_reply.started":"2021-09-15T18:31:36.964016Z","shell.execute_reply":"2021-09-15T18:31:36.978381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TFRecordDataLoader:\n    def __init__(self, files, batch_size=32, cache=False, train=True, \n                              repeat=False, shuffle=False, labeled=True, \n                              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        self.num_examples = count_data_items(files, labeled)\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-09-15T18:31:36.980470Z","iopub.execute_input":"2021-09-15T18:31:36.980803Z","iopub.status.idle":"2021-09-15T18:31:36.992945Z","shell.execute_reply.started":"2021-09-15T18:31:36.980766Z","shell.execute_reply":"2021-09-15T18:31:36.992127Z"},"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":"class CustomViT(nn.Module):\n    def __init__(self, cfg):\n        super().__init__()\n        self.cfg = cfg\n        self.cqt = CQT1992v2(**cfg.cqt_params)\n        self.image_size = cfg.image_size\n        self.model = ViT(image_size=cfg.image_size,\n                         patch_size=cfg.patch_size,\n                         num_classes=1,\n                         dim=CFG.n_hidden,\n                         depth=CFG.n_layers,\n                         heads=CFG.n_heads,\n                         mlp_dim=4*CFG.n_hidden,\n                         dropout=0.1,\n                         emb_dropout=0.1,\n                         )\n\n    def forward(self, x):\n        sps = []\n        for i in range(3):\n            s = self.cqt(x[:, i, :])\n            s = F.interpolate(s, size=256, mode='linear', align_corners=True)\n            sps.append(s)\n        x = torch.stack(sps, 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-09-15T18:42:14.212561Z","iopub.execute_input":"2021-09-15T18:42:14.212880Z","iopub.status.idle":"2021-09-15T18:42:14.221401Z","shell.execute_reply.started":"2021-09-15T18:42:14.212852Z","shell.execute_reply":"2021-09-15T18:42:14.220405Z"},"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-09-15T18:42:15.222780Z","iopub.execute_input":"2021-09-15T18:42:15.223126Z","iopub.status.idle":"2021-09-15T18:42:15.231895Z","shell.execute_reply.started":"2021-09-15T18:42:15.223080Z","shell.execute_reply":"2021-09-15T18:42:15.230858Z"},"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, \n        shuffle=True)\n    for step, d in enumerate(train_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n        x = 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(x)\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_last_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        x = 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(x)\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-09-15T18:42:16.287622Z","iopub.execute_input":"2021-09-15T18:42:16.287939Z","iopub.status.idle":"2021-09-15T18:42:16.307064Z","shell.execute_reply.started":"2021-09-15T18:42:16.287909Z","shell.execute_reply":"2021-09-15T18:42:16.306036Z"},"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, \n                                                             mode='min', \n                                                             factor=CFG.factor, \n                                                             patience=CFG.patience, \n                                                             verbose=True, \n                                                             eps=CFG.eps)\n        elif CFG.scheduler=='CosineAnnealingLR':\n            scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, \n                                                             T_max=CFG.T_max, \n                                                             eta_min=CFG.min_lr, \n                                                             last_epoch=-1)\n        elif CFG.scheduler=='CosineAnnealingWarmRestarts':\n            scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, \n                                                                       T_0=CFG.T_0, \n                                                                       T_mult=1, \n                                                                       eta_min=CFG.min_lr, \n                                                                       last_epoch=-1)\n        return scheduler\n\n    # ====================================================\n    # model & optimizer\n    # ====================================================\n    model = CustomViT(cfg=CFG)\n    model.to(device)\n\n    optimizer = optim.Adam(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay)\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        print(\"\\n\\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-09-15T18:42:18.131420Z","iopub.execute_input":"2021-09-15T18:42:18.131761Z","iopub.status.idle":"2021-09-15T18:42:18.146699Z","shell.execute_reply.started":"2021-09-15T18:42:18.131728Z","shell.execute_reply":"2021-09-15T18:42:18.145563Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"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\nif 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-09-15T18:42:21.773796Z","iopub.execute_input":"2021-09-15T18:42:21.774137Z","iopub.status.idle":"2021-09-15T18:43:13.572068Z","shell.execute_reply.started":"2021-09-15T18:42:21.774088Z","shell.execute_reply":"2021-09-15T18:43:13.570227Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inference","metadata":{}},{"cell_type":"code","source":"states = []\nfor fold  in CFG.trn_fold:\n    states.append(torch.load(os.path.join(SAVEDIR, f'{CFG.model_name}_fold{fold}_best_score.pth')))","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gcs_paths = []\nfor i, j in [(0, 4), (5, 9)]:\n    path = f\"g2net-waveform-tfrecords-test-{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 + \"/test*.tfrecords\"))))\n    \nprint(\"test_files: \", len(all_files))\nall_files = np.array(all_files)","metadata":{"execution":{"iopub.status.busy":"2021-09-15T18:47:40.581397Z","iopub.execute_input":"2021-09-15T18:47:40.581745Z","iopub.status.idle":"2021-09-15T18:47:41.488919Z","shell.execute_reply.started":"2021-09-15T18:47:40.581718Z","shell.execute_reply":"2021-09-15T18:47:41.487241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model= CustomViT(cfg=CFG)\nmodel.to(device)\n\nwave_ids = []\nprobs_all = []\n\nfor fold, state in enumerate(states):\n    tqdm.write(f\"\\n\\nFold{fold}\")\n    \n    model.load_state_dict(state['model'])\n    model.eval()\n    probs = []\n\n    test_loader = TFRecordDataLoader(all_files, batch_size=CFG.val_batch_size, \n                                     shuffle=False, labeled=False)\n\n    for i, d in tqdm(enumerate(test_loader), total=len(test_loader)):\n        x = torch.from_numpy(d[0]).to(device)\n\n        with torch.no_grad():\n            y_preds = model(x)\n        preds = y_preds.sigmoid().to('cpu').numpy()\n        probs.append(preds)\n\n        if fold==0: # same test loader, no need to do this the second time\n            wave_ids.append(d[1].astype('U13'))\n\n    probs = np.concatenate(probs)\n    probs_all.append(probs)\n\nprobs_avg = np.asarray(probs_all).mean(axis=0).flatten()\nwave_ids = np.concatenate(wave_ids)","metadata":{"execution":{"iopub.status.busy":"2021-09-15T18:49:15.250557Z","iopub.execute_input":"2021-09-15T18:49:15.250882Z","iopub.status.idle":"2021-09-15T18:49:18.392872Z","shell.execute_reply.started":"2021-09-15T18:49:15.250854Z","shell.execute_reply":"2021-09-15T18:49:18.392002Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df = pd.DataFrame({'id': wave_ids, 'target': probs_avg})\n# Save test dataframe to disk\nfolds = '_'.join([str(s) for s in CFG.trn_fold])\ntest_df.to_csv(f'{CFG.model_name}_folds_{folds}.csv', index = False)","metadata":{"execution":{"iopub.status.busy":"2021-09-15T18:49:22.748889Z","iopub.execute_input":"2021-09-15T18:49:22.749222Z","iopub.status.idle":"2021-09-15T18:49:22.983740Z","shell.execute_reply.started":"2021-09-15T18:49:22.749190Z","shell.execute_reply":"2021-09-15T18:49:22.982943Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}