{"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":"\n- Basically a 1D CNN starter with bandpass. Filter size hard-coded from [https://www.kaggle.com/kit716/grav-wave-detection](https://www.kaggle.com/kit716/grav-wave-detection) which uses the simple architecture from https://journals.aps.org/prl/pdf/10.1103/PhysRevLett.120.141103 \n- Added inference to @hidehisaarai1213 's PyTorch starter, iteration order changed from Y.Nakama's pipeline: \"iter on loader first then load model\" to \"load model first then iter the loader\"\n","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":"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 kaggle_datasets import KaggleDatasets\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import KFold\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2021-09-16T04:44:15.710519Z","iopub.execute_input":"2021-09-16T04:44:15.710916Z","iopub.status.idle":"2021-09-16T04:44:22.392701Z","shell.execute_reply.started":"2021-09-16T04:44:15.710811Z","shell.execute_reply":"2021-09-16T04:44:22.39185Z"},"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-16T04:44:22.394086Z","iopub.execute_input":"2021-09-16T04:44:22.394439Z","iopub.status.idle":"2021-09-16T04:44:22.441016Z","shell.execute_reply.started":"2021-09-16T04:44:22.3944Z","shell.execute_reply":"2021-09-16T04:44:22.440007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CFG","metadata":{}},{"cell_type":"code","source":"class CFG:\n    debug = False\n    print_freq = 1000\n    num_workers = 4\n    scheduler = \"CosineAnnealingLR\"\n    model_name = \"1dcnn\"\n    epochs = 5\n    T_max = 3\n    lr = 1e-4\n    min_lr = 1e-7\n    batch_size = 64\n    val_batch_size = 100\n    weight_decay = 1e-5\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, 1, 2, 3]  # [0, 1, 2, 3, 4]\n    train = True\n    bandpass_params = dict(lf=20, \n                           hf=500)","metadata":{"execution":{"iopub.status.busy":"2021-09-16T04:44:22.443588Z","iopub.execute_input":"2021-09-16T04:44:22.443985Z","iopub.status.idle":"2021-09-16T04:44:22.45243Z","shell.execute_reply.started":"2021-09-16T04:44:22.443943Z","shell.execute_reply":"2021-09-16T04:44:22.451503Z"},"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-16T04:44:22.453966Z","iopub.execute_input":"2021-09-16T04:44:22.454317Z","iopub.status.idle":"2021-09-16T04:44:22.466594Z","shell.execute_reply.started":"2021-09-16T04:44:22.45428Z","shell.execute_reply":"2021-09-16T04:44:22.465737Z"},"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-16T04:44:22.46848Z","iopub.execute_input":"2021-09-16T04:44:22.468763Z","iopub.status.idle":"2021-09-16T04:44:25.142142Z","shell.execute_reply.started":"2021-09-16T04:44:22.468713Z","shell.execute_reply":"2021-09-16T04:44:25.141269Z"},"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-16T04:44:25.143446Z","iopub.execute_input":"2021-09-16T04:44:25.143772Z","iopub.status.idle":"2021-09-16T04:44:25.150509Z","shell.execute_reply.started":"2021-09-16T04:44:25.143735Z","shell.execute_reply":"2021-09-16T04:44:25.148449Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Bandpass\n\nModified from various notebooks and https://www.kaggle.com/c/g2net-gravitational-wave-detection/discussion/261721#1458564","metadata":{}},{"cell_type":"code","source":"def bandpass(x, lf=20, hf=500, order=8, sr=2048):\n    '''\n    Cell 33 of https://www.gw-openscience.org/LVT151012data/LOSC_Event_tutorial_LVT151012.html\n    https://scipy-cookbook.readthedocs.io/items/ButterworthBandpass.html\n    '''\n    sos = signal.butter(order, [lf, hf], btype=\"bandpass\", output=\"sos\", fs=sr)\n    normalization = np.sqrt((hf - lf) / (sr / 2))\n    if x.ndim ==2:\n        for i in range(3):\n            x[i] = signal.sosfilt(sos, x[i]) * normalization\n    elif x.ndim == 3: # batch\n        for i in range(x.shape[0]):\n            for j in range(3):\n                x[i, j] = signal.sosfilt(sos, x[i, j]) * normalization\n    return x","metadata":{"execution":{"iopub.status.busy":"2021-09-16T04:44:25.152193Z","iopub.execute_input":"2021-09-16T04:44:25.152574Z","iopub.status.idle":"2021-09-16T04:44:25.161142Z","shell.execute_reply.started":"2021-09-16T04:44:25.152539Z","shell.execute_reply":"2021-09-16T04:44:25.159966Z"},"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-16T04:44:25.164147Z","iopub.execute_input":"2021-09-16T04:44:25.164689Z","iopub.status.idle":"2021-09-16T04:44:25.179583Z","shell.execute_reply.started":"2021-09-16T04:44:25.164651Z","shell.execute_reply":"2021-09-16T04:44:25.17851Z"},"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-16T04:44:25.182021Z","iopub.execute_input":"2021-09-16T04:44:25.182273Z","iopub.status.idle":"2021-09-16T04:44:25.194739Z","shell.execute_reply.started":"2021-09-16T04:44:25.182246Z","shell.execute_reply":"2021-09-16T04:44:25.193898Z"},"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 CNN1d(nn.Module):\n    \"\"\"1D convolutional neural network. Classifier of the gravitational waves.\n    Architecture from there https://journals.aps.org/prl/pdf/10.1103/PhysRevLett.120.141103\n    \"\"\"\n\n    def __init__(self, debug=False):\n        super().__init__()\n        self.cnn1 = nn.Sequential(\n            nn.Conv1d(3, 64, kernel_size=64),\n            nn.BatchNorm1d(64),\n            nn.ELU(),\n        )\n        self.cnn2 = nn.Sequential(\n            nn.Conv1d(64, 64, kernel_size=32),\n            nn.AvgPool1d(kernel_size=8),\n            nn.BatchNorm1d(64),\n            nn.ELU(),\n        )\n        self.cnn3 = nn.Sequential(\n            nn.Conv1d(64, 128, kernel_size=32),\n            nn.BatchNorm1d(128),\n            nn.ELU(),\n        )\n        self.cnn4 = nn.Sequential(\n            nn.Conv1d(128, 128, kernel_size=16),\n            nn.AvgPool1d(kernel_size=6),\n            nn.BatchNorm1d(128),\n            nn.ELU(),\n        )\n        self.cnn5 = nn.Sequential(\n            nn.Conv1d(128, 256, kernel_size=16),\n            nn.BatchNorm1d(256),\n            nn.ELU(),\n        )\n        self.cnn6 = nn.Sequential(\n            nn.Conv1d(256, 256, kernel_size=16),\n            nn.MaxPool1d(kernel_size=4),\n            nn.BatchNorm1d(256),\n            nn.ELU(),\n        )\n        self.fc1 = nn.Sequential(\n            nn.Linear(256 * 11, 64),\n            nn.BatchNorm1d(64),\n            nn.Dropout(0.4),\n            nn.ELU(),\n        )\n        self.fc2 = nn.Sequential(\n            nn.Linear(64, 64),\n            nn.BatchNorm1d(64),\n            nn.Dropout(0.4),\n            nn.ELU(),\n        )\n        self.fc3 = nn.Sequential(\n            nn.Linear(64, 1),\n        )\n        self.debug = debug\n\n    def forward(self, x, pos=None):\n        x = self.cnn1(x)\n        x = self.cnn2(x)\n        x = self.cnn3(x)\n        x = self.cnn4(x)\n        x = self.cnn5(x)\n        x = self.cnn6(x)\n        x = x.flatten(start_dim=1)\n        x = self.fc1(x)\n        x = self.fc2(x)\n        x = self.fc3(x)\n        return x\n","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-16T04:44:25.198151Z","iopub.execute_input":"2021-09-16T04:44:25.198413Z","iopub.status.idle":"2021-09-16T04:44:25.214884Z","shell.execute_reply.started":"2021-09-16T04:44:25.198387Z","shell.execute_reply":"2021-09-16T04:44:25.214202Z"},"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-16T04:44:25.216068Z","iopub.execute_input":"2021-09-16T04:44:25.216477Z","iopub.status.idle":"2021-09-16T04:44:25.228666Z","shell.execute_reply.started":"2021-09-16T04:44:25.216415Z","shell.execute_reply":"2021-09-16T04:44:25.227671Z"},"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 = bandpass(d[0], **CFG.bandpass_params)\n        x = torch.from_numpy(x).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 = bandpass(d[0], **CFG.bandpass_params)\n        x = torch.from_numpy(x).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-16T04:44:25.230104Z","iopub.execute_input":"2021-09-16T04:44:25.230442Z","iopub.status.idle":"2021-09-16T04:44:25.253097Z","shell.execute_reply.started":"2021-09-16T04:44:25.230407Z","shell.execute_reply":"2021-09-16T04:44:25.252304Z"},"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 = CNN1d()\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-16T04:44:25.254259Z","iopub.execute_input":"2021-09-16T04:44:25.254651Z","iopub.status.idle":"2021-09-16T04:44:25.273588Z","shell.execute_reply.started":"2021-09-16T04:44:25.25458Z","shell.execute_reply":"2021-09-16T04:44:25.272728Z"},"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-16T04:44:25.274812Z","iopub.execute_input":"2021-09-16T04:44:25.275167Z","iopub.status.idle":"2021-09-16T05:21:40.709461Z","shell.execute_reply.started":"2021-09-16T04:44:25.275133Z","shell.execute_reply":"2021-09-16T05:21:40.707924Z"},"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":{"iopub.status.busy":"2021-09-16T05:21:40.710356Z","iopub.status.idle":"2021-09-16T05:21:40.710755Z"},"trusted":true},"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-16T05:21:40.711751Z","iopub.status.idle":"2021-09-16T05:21:40.712252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model= CNN1d()\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 = bandpass(d[0], **CFG.bandpass_params)\n        x = torch.from_numpy(x).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-16T05:21:40.713281Z","iopub.status.idle":"2021-09-16T05:21:40.713843Z"},"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-16T05:21:40.714867Z","iopub.status.idle":"2021-09-16T05:21:40.715449Z"},"trusted":true},"execution_count":null,"outputs":[]}]}