{"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":"## Learning Objective : \n### To use multi-GPU distributed training for PyTorch models on Kaggle\n\nIf this helps then please make sure you have upvoted Jun Koda's original notebook as he did all the work here!\n\nI found that `tf-efficientnet-bX-ns` family is giving great results for G2Net Competition, but higher order model of the family have more number\nof parameters, so it is exhausting Kaggle GPU quota.\n\nKeeping in mind the launch of T4x2 GPUs, it's time for Kaggle community to actually learn distributed training concept.\n\nI found people recommending [torch.nn.DataParallel](https://www.kaggle.com/discussions/product-feedback/361104) for distributed training, but it doesn't improve performance or reducing training time.\nPyTorch official documentation encourages to use [PyTorch Distributed](https://pytorch.org/tutorials/beginner/dist_overview.html) for distributed training, but the learning curve here is quite steep, and makes code untidy. I found that 🤗 [HuggingFace Accelerate](https://huggingface.co/docs/accelerate/index) offers an API which makes distributed training easy. Apart from it, accelerate allows us to train on larger batch size.","metadata":{}},{"cell_type":"markdown","source":"I've reused Jun Koda's spectrogram notebook, to demonstrate distributed training, Please upvote his notebook as the work underneath (sequential training) is his.","metadata":{}},{"cell_type":"code","source":"!python --version","metadata":{"execution":{"iopub.status.busy":"2022-11-05T13:13:17.063448Z","iopub.execute_input":"2022-11-05T13:13:17.064572Z","iopub.status.idle":"2022-11-05T13:13:18.267552Z","shell.execute_reply.started":"2022-11-05T13:13:17.064454Z","shell.execute_reply":"2022-11-05T13:13:18.266028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Use timm pretrained image model\n! pip3 install timm\n! pip install git+https://github.com/huggingface/accelerate # Install Huggingface Accelerate","metadata":{"execution":{"iopub.status.busy":"2022-11-05T13:13:18.270175Z","iopub.execute_input":"2022-11-05T13:13:18.270562Z","iopub.status.idle":"2022-11-05T13:14:03.150827Z","shell.execute_reply.started":"2022-11-05T13:13:18.270505Z","shell.execute_reply":"2022-11-05T13:14:03.149363Z"},"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 gc\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\nfrom accelerate import Accelerator, notebook_launcher # main interface, distributed launcher\nfrom accelerate.utils import set_seed # reproducability across devices\n\n# device = torch.device('cpu') => With accelerate, avoid specifying device\ncriterion = 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":{"execution":{"iopub.status.busy":"2022-11-05T13:14:03.156029Z","iopub.execute_input":"2022-11-05T13:14:03.158965Z","iopub.status.idle":"2022-11-05T13:14:11.57003Z","shell.execute_reply.started":"2022-11-05T13:14:03.158919Z","shell.execute_reply":"2022-11-05T13:14:11.568918Z"},"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":{"execution":{"iopub.status.busy":"2022-11-05T13:14:11.572782Z","iopub.execute_input":"2022-11-05T13:14:11.573224Z","iopub.status.idle":"2022-11-05T13:14:11.583779Z","shell.execute_reply.started":"2022-11-05T13:14:11.573173Z","shell.execute_reply":"2022-11-05T13:14:11.582471Z"},"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()","metadata":{"execution":{"iopub.status.busy":"2022-11-05T13:14:11.585595Z","iopub.execute_input":"2022-11-05T13:14:11.586227Z","iopub.status.idle":"2022-11-05T13:14:12.34458Z","shell.execute_reply.started":"2022-11-05T13:14:11.586188Z","shell.execute_reply":"2022-11-05T13:14:12.343537Z"},"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):\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, 1)\n        self.model = model\n\n    def forward(self, x):\n        x = self.model(x)\n        x = self.fc(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-11-05T13:14:12.345721Z","iopub.execute_input":"2022-11-05T13:14:12.34668Z","iopub.status.idle":"2022-11-05T13:14:12.355218Z","shell.execute_reply.started":"2022-11-05T13:14:12.346643Z","shell.execute_reply":"2022-11-05T13:14:12.35426Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Predict and evaluate","metadata":{}},{"cell_type":"code","source":"def evaluate(model, loader_val, accelerator=None, *, compute_score=True, pbar=None):\n    \"\"\"\n    Predict and compute loss and score\n    \"\"\"\n    tb = time.time()\n    was_training = model.training\n    model.eval()\n\n    loss_sum = 0.0\n    n_sum = 0\n    y_all = []\n    y_feat_all = []\n    id_all = []\n    y_pred_all = []\n\n    if pbar is not None:\n        pbar = tqdm(desc='Predict', nrows=78, total=pbar)\n\n    for img, y in loader_val:\n        n = y.size(0)\n        \n        if accelerator is None:\n            img = img.to(device)\n            y = y.to(device)\n\n        with torch.no_grad():\n            y_pred = model(img) #to(device)\n        loss = criterion(y_pred.view(-1), y)\n        \n        n_sum += n\n        loss_sum += n * loss.item()\n\n        y_all.append(y.cpu().detach().numpy())\n        y_pred_all.append(y_pred.sigmoid().squeeze().cpu().detach().numpy())\n        #y_feat_all.append(y_feat.squeeze().cpu().detach().numpy())\n        id_all.append(id)\n        \n        if pbar is not None:\n            pbar.update(len(img))\n        \n        del loss, y_pred, img, y\n\n    loss_val = loss_sum / n_sum\n\n    y = np.concatenate(y_all)\n    y_pred = np.concatenate(y_pred_all)\n\n    score = roc_auc_score(y, y_pred) if compute_score else None\n\n    ret = {'loss': loss_val,\n           'score': score,\n           'y': y,\n           'id': id,\n           'y_pred': y_pred,\n           'time': time.time() - tb}\n    \n    model.train(was_training)  # back to train from eval if necessary\n    gc.collect()\n    return ret","metadata":{"execution":{"iopub.status.busy":"2022-11-05T13:20:45.993128Z","iopub.execute_input":"2022-11-05T13:20:45.993537Z","iopub.status.idle":"2022-11-05T13:20:46.00666Z","shell.execute_reply.started":"2022-11-05T13:20:45.993489Z","shell.execute_reply":"2022-11-05T13:20:46.005699Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"def train(mixed_precision=\"fp16\", seed=42, batch_size=32):\n    \n    set_seed(seed)\n    accelerator = Accelerator(mixed_precision=mixed_precision)\n    \n    model_name = 'tf_efficientnet_b0_ns'\n    \n    nfold = 5\n    kfold = StratifiedKFold(n_splits=nfold, random_state=42, shuffle=True)\n\n    epochs = 1\n    batch_size = 32\n    num_workers = 2\n    weight_decay = 1e-6\n    max_grad_norm = 1000\n\n    lr_max = 4e-4\n    epochs_warmup = 0\n\n    for ifold, (idx_train, idx_test) in enumerate(kfold.split(dataset, df['target'])):\n        gc.collect()\n        torch.cuda.empty_cache()\n\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        print(\"Length of train dataset = \", len(dataset_train))\n        print(\"Length of val dataset = \", len(dataset_val))\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        with accelerator.main_process_first():\n            model = Model(model_name, pretrained=True)\n            # model = model.to(device)\n        \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        \n        model, optimizer, loader_train, loader_val, scheduler = accelerator.prepare(model, optimizer, loader_train, loader_val, scheduler)\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                gc.collect()\n                torch.cuda.empty_cache()\n\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                accelerator.backward(loss)\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, accelerator = accelerator)\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        accelerator.wait_for_everyone()\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        \n        accelerator.save(accelerator.unwrap_model(model).state_dict(), ofilename)\n        \n        print(ofilename, 'written')\n\n        del model\n\n        break  # 1 fold only","metadata":{"execution":{"iopub.status.busy":"2022-11-05T13:14:12.37315Z","iopub.execute_input":"2022-11-05T13:14:12.373757Z","iopub.status.idle":"2022-11-05T13:14:12.409067Z","shell.execute_reply.started":"2022-11-05T13:14:12.373722Z","shell.execute_reply":"2022-11-05T13:14:12.407835Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"notebook_launcher(train, (\"fp16\", 42, 64), num_processes=2)","metadata":{"execution":{"iopub.status.busy":"2022-11-05T13:14:12.410896Z","iopub.execute_input":"2022-11-05T13:14:12.411597Z","iopub.status.idle":"2022-11-05T13:15:33.362501Z","shell.execute_reply.started":"2022-11-05T13:14:12.41143Z","shell.execute_reply":"2022-11-05T13:15:33.36099Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Predict & Submit","metadata":{}},{"cell_type":"code","source":"# Predict on single GPU\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nsubmit = pd.read_csv('../input/g2net-detecting-continuous-gravitational-waves/sample_submission.csv')\n\ndataset_test = Dataset('test', submit)\nloader_test = torch.utils.data.DataLoader(dataset_test, batch_size=64,\n                                          num_workers=0, pin_memory=True)\n\ntrained_model = Model('tf_efficientnet_b0_ns').to(device)\ntrained_model.load_state_dict(torch.load('./model0.pytorch'))\n        \nprint(\"Loaded Model\")\n\ntest = evaluate(trained_model, loader_test, compute_score=False, pbar=len(submit))\n    \n# Write prediction\nsubmit['target'] = test['y_pred']\nsubmit.to_csv('submission.csv', index=False)\nsubmit.head()","metadata":{"execution":{"iopub.status.busy":"2022-11-05T13:21:29.325303Z","iopub.execute_input":"2022-11-05T13:21:29.325817Z","iopub.status.idle":"2022-11-05T14:13:00.220125Z","shell.execute_reply.started":"2022-11-05T13:21:29.325768Z","shell.execute_reply":"2022-11-05T14:13:00.213606Z"},"trusted":true},"execution_count":null,"outputs":[]}]}