{"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":"code","source":"import numpy as np\nimport pandas as pd\nfrom scipy import signal\nimport os\nimport matplotlib.pyplot as plt\nfrom glob import glob\nimport torch\nfrom sklearn.metrics import roc_auc_score\nimport time\nimport math","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = pd.read_csv(\"../input/g2net-gravitational-wave-detection/training_labels.csv\")\ntrain_df","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"paths = [\"../input/g2net-gravitational-wave-detection/train/\" + \"/\".join(id[:3]) + \"/\" + id + \".npy\" for id in train_df.id.values]\ntrain_df['path'] = paths\ntrain_df","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def whiten(x):\n    for i in range(3):\n        spec = np.fft.rfft(x[i])\n        mag = np.sqrt(np.real(spec*np.conj(spec)))\n        norm = np.sqrt(np.array([4096/2]))\n        x[i] = np.fft.irfft(spec/mag) * norm\n    return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def apply_bandpass(x, lf=30, hf=500, order=8, sr=2048):\n    sos = signal.butter(order, [lf, hf], btype=\"bandpass\", output=\"sos\", fs=sr)\n    normalization = np.sqrt((hf - lf) / (sr / 2))\n    for i in range(3):\n        x[i] = signal.sosfiltfilt(sos, x[i]) / normalization\n    return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def preprocess(x):\n    x = x / np.max(np.abs(x), axis=-1, keepdims=True)\n    #scale = np.array([[1.5e-20], [1.5e-20], [0.5e-20]])\n    #x = x / scale\n    x *= signal.tukey(4096, 0.1)\n    #x = whiten(x)\n    x = apply_bandpass(x)\n    return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DataSet:\n    def __init__(self, paths, target, index):\n        self.paths = [paths[i] for i in index]\n        self.target = [target[i] for i in index]\n        \n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, index):\n        x = np.load(self.paths[index])\n        x = preprocess(x)\n        return x.astype(np.float32), self.target[index].astype(np.float32)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SelfAttention(torch.nn.Module):\n    def __init__(self, dim, out_dim):\n        super(SelfAttention,self).__init__()\n        self.dim = out_dim\n        self.qkv_weight = torch.nn.Linear(dim, out_dim*3, bias=False)\n    \n    def forward(self, x):\n        q, k, v = self.qkv_weight(x).chunk(3, dim=-1)\n        att_logit = torch.bmm(q, k.transpose(1,2)) * (self.dim ** -0.5) # q * k\n        att_weight = torch.softmax(att_logit, dim=-1)\n        weighted_v = torch.bmm(att_weight, v) # q*k@k*dim == q * dim\n        return weighted_v","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MultiheadSelfAttention(torch.nn.Module):\n    def __init__(self, in_dim, out_dim, heads):\n        super(MultiheadSelfAttention,self).__init__()\n        head_out = out_dim // heads\n        self.heads = torch.nn.ModuleList([SelfAttention(in_dim, head_out) for _ in range(heads)])\n        self.output = torch.nn.Linear(out_dim, out_dim)\n    def forward(self, x):\n        outs = []\n        for head in self.heads:\n            out = head(x)\n            outs.append(out)\n        outs = torch.cat(outs, dim=-1)\n        outs = self.output(outs)\n        return outs","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Residual(torch.nn.Module):\n    def __init__(self, fn, dim):\n        super(Residual,self).__init__()\n        self.fn = fn\n        self.dropout = torch.nn.Dropout(p=0.2)\n    def forward(self, x):\n        hop = self.dropout(self.fn(x))\n        y = x + hop\n        return y","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Feedforward(torch.nn.Module):\n    def __init__(self, embed_dim, hidden_dim):\n        super(Feedforward,self).__init__()\n        self.hidden = torch.nn.Sequential(torch.nn.Linear(embed_dim, hidden_dim),\n                                        torch.nn.GELU(),\n                                         torch.nn.Linear(hidden_dim, embed_dim))\n    def forward(self, x):\n        return self.hidden(x)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Encoder(torch.nn.Module):\n    def __init__(self, in_dim, embed_dim, hidden_dim, heads, depth):\n        super(Encoder,self).__init__()\n        self.feedforwards = torch.nn.ModuleList([Residual(Feedforward(embed_dim, hidden_dim), embed_dim) for _ in range(depth)])\n        self.attentions = torch.nn.ModuleList([Residual(MultiheadSelfAttention(in_dim, embed_dim, heads), in_dim) if i == 0 \n                                               else Residual(MultiheadSelfAttention(embed_dim, embed_dim, heads), embed_dim)\n                                               for i in range(depth)])\n    \n    def forward(self, x):\n        for ffn, attn in zip(self.feedforwards, self.attentions):\n            x = attn(x)\n            x = ffn(x)\n        return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Recognizer(torch.nn.Module):\n    def __init__(self, embed_dim):\n        super(Recognizer,self).__init__()\n        self.logit = torch.nn.Linear(embed_dim, 1)\n        self.out = torch.nn.Sigmoid()\n    def forward(self, x):\n        x = self.logit(x).transpose(1, 2)\n        x = x.reshape(x.shape[0], 3, -1)\n        x = torch.max(x, dim=-1)[0].mean(dim=-1)\n        out = self.out(x)\n        return out, x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CNN1d(torch.nn.Module):\n    def __init__(self, embed_dim):\n        super(CNN1d, self).__init__()\n        self.CE = torch.nn.Sequential(torch.nn.BatchNorm1d(1),\n                                    torch.nn.Conv1d(1, embed_dim//4, 16, stride=1, padding=0),\n                                    torch.nn.GELU(),\n                                    torch.nn.MaxPool1d(4, stride=4, padding=0),\n                                    torch.nn.BatchNorm1d(embed_dim//4),\n                                    torch.nn.Conv1d(embed_dim//4, embed_dim//2, 8, stride=1, padding=0),\n                                    torch.nn.GELU(),\n                                    torch.nn.MaxPool1d(4, stride=4, padding=0),\n                                    torch.nn.BatchNorm1d(embed_dim//2),\n                                    torch.nn.Conv1d(embed_dim//2, embed_dim, 8, stride=1, padding=0),\n                                    torch.nn.GELU(),\n                                    torch.nn.MaxPool1d(4, stride=4, padding=0))\n        self.length = 61\n    def forward(self, x):\n        return self.CE(x)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Embed(torch.nn.Module):\n    def __init__(self, enc_dim, pe_dim, embed_dim, pe='fix'):\n        super(Embed, self).__init__()\n        dim = enc_dim\n        self.CNN = CNN1d(dim)\n        max_len = self.CNN.length\n        if pe == 'fix':\n            position = torch.arange(max_len).unsqueeze(1)\n            div_term = torch.exp(torch.arange(0, pe_dim, 2) * (-math.log(10000.0) / pe_dim))\n            pe = torch.zeros(max_len, pe_dim)\n            pe[:, 0::2] = torch.sin(position * div_term)\n            pe[:, 1::2] = torch.cos(position * div_term)\n            pe = pe.repeat((3,1))\n            detector = torch.zeros(max_len*3, 3)\n            detector[:max_len, 0] = 1\n            detector[max_len:max_len*2, 1] = 1\n            detector[max_len*2:max_len*3, 2] = 1\n            self.pe = torch.cat([pe, detector], dim=-1)\n            out_dim = dim + pe_dim + 3\n        elif pe == 'trainable':\n            self.tpe = torch.nn.Parameter(torch.randn(max_len, pe_dim))\n            detector = torch.zeros(max_len*3, 3)\n            detector[:max_len, 0] = 1\n            detector[max_len:max_len*2, 1] = 1\n            detector[max_len*2:max_len*3, 2] = 1\n            self.pe = torch.cat([self.tpe.repeat((3,1)), detector], dim=-1)\n            out_dim = dim + pe_dim + 3\n        else:\n            self.pe = None\n            out_dim = dim\n        \n        self.embed = torch.nn.Linear(out_dim, embed_dim, bias=False)\n\n    def forward(self, x):\n        ss = []\n        for i in range(x.shape[1]):\n            s = x[:,i,:].unsqueeze(1)\n            fts = self.CNN(s).transpose(1, 2)\n            std, mean = torch.std_mean(fts, dim=(1,2), unbiased=False, keepdim=True)\n            fts = torch.div(fts-mean, std)\n            ss.append(fts)\n        x = torch.cat(ss, dim=1)\n        if self.pe is not None:\n            x = torch.cat([x, self.pe.to(x.device).repeat((x.shape[0],1,1))], dim=-1)\n        x = self.embed(x)\n        return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Model(torch.nn.Module):\n    def __init__(self, enc_dim, pe_dim, embed_dim, hidden_dim, heads, depth, pe='fix'):\n        super(Model,self).__init__()\n        self.embed = Embed(enc_dim, pe_dim, embed_dim, pe=pe)\n        self.encoder = Encoder(embed_dim, embed_dim, hidden_dim, heads, depth)\n        self.recognizer = Recognizer(embed_dim)\n    \n    def forward(self, x):\n        x = self.embed(x)\n        x = self.encoder(x)\n        out, x = self.recognizer(x)\n        return out, x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed = 2434\npe_dim = 128\nenc_dim = 128\nembed_dim = 256\nhidden_dim = 256*4\nheads = 4\ndepth = 4\npe = 'fix'\nlr = 1e-4\nweight_decay = 1e-5\ntrain_batch_size = 256\ntest_batch_size = 512\noptimizer = 'Adam'\ngrad_clip = 1000","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"np.random.seed(seed)\ntorch.manual_seed(seed)\ntrain_index = np.random.rand(train_df.shape[0]) < 0.9\nval_index = ~train_index\ntrain_index = np.nonzero(train_index)[0]\nval_index = np.nonzero(val_index)[0]\ntrain_dataset = DataSet(train_df.path.values, train_df.target.values, train_index)\nval_dataset = DataSet(train_df.path.values, train_df.target.values, val_index)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"epochs = 10\ndevice = 'cuda'\nnum_worker = os.cpu_count()\nmodel = Model(enc_dim, pe_dim, embed_dim, hidden_dim, heads, depth, pe).to(device)\nbest_model = Model(enc_dim, pe_dim, embed_dim, hidden_dim, heads, depth, pe).to(device)\nbest_model.load_state_dict(model.state_dict())\ncriterion = torch.nn.BCEWithLogitsLoss(reduction='sum')\noptim = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)\ntrain_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=train_batch_size,\n                                         shuffle=True, drop_last=True, num_workers=num_worker, pin_memory=True)\nval_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=test_batch_size,\n                                         shuffle=False, drop_last=False, num_workers=num_worker, pin_memory=True)\nstart_time = time.time()\ntrain_loss_list = []\ntrain_auc_list = []\nval_loss_list = []\nval_auc_list = []\nbest_auc = 0\nfor epoch in range(epochs):\n    model.train()\n    train_loss = 0\n    train_auc = 0\n    count = 0\n    batch_count = 0\n    preds = []\n    targets = []\n    for data, target in train_dataloader:\n        optim.zero_grad()\n        data = data.to(device)\n        target = target.to(device)\n        pred, x = model(data)\n        loss = criterion(x, target)\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n        optim.step()\n        train_loss += loss.item()\n        preds.append(pred.detach().cpu().numpy())\n        targets.append(target.cpu().numpy())\n        count += data.shape[0]\n        batch_count += 1\n    train_loss = train_loss / count\n    preds = np.concatenate(preds, axis=0)\n    targets = np.concatenate(targets, axis=0)\n    train_auc = roc_auc_score(targets, preds)\n    train_loss_list.append(train_loss)\n    train_auc_list.append(train_auc)\n    \n    model.eval()\n    with torch.no_grad():\n        val_loss = 0\n        val_auc = 0\n        count = 0\n        batch_count = 0\n        preds = []\n        targets = []\n        for data, target in val_dataloader:\n            data = data.to(device)\n            target = target.to(device)\n            pred, x = model(data)\n            loss = criterion(x, target)\n            val_loss += loss.item()\n            preds.append(pred.detach().cpu().numpy())\n            targets.append(target.cpu().numpy())\n            count += data.shape[0]\n            batch_count += 1\n        val_loss = val_loss / count\n        preds = np.concatenate(preds, axis=0)\n        targets = np.concatenate(targets, axis=0)\n        val_auc = roc_auc_score(targets, preds)\n        val_loss_list.append(val_loss)\n        val_auc_list.append(val_auc)\n    spent_time = time.time() - start_time\n    print(f'epoch: {epoch} train loss: {train_loss} train auc: {train_auc} val loss: {val_loss} val auc: {val_auc} time: {spent_time/60} min')\n    if val_auc >= best_auc:\n        best_auc = val_auc\n        best_model.load_state_dict(model.state_dict())\n    if spent_time >= 25000:\n        print('time over')\n        break","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.plot(train_loss_list)\nplt.plot(val_loss_list)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.plot(train_auc_list)\nplt.plot(val_auc_list)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(model.state_dict(), \"model\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"paths = glob(\"../input/g2net-gravitational-wave-detection/test/*/*/*/*\")\nids = [path.split(\"/\")[-1].split(\".\")[0] for path in paths]\ntest_df = pd.DataFrame({\"path\":paths,\"id\":ids})\ntest_df['target'] = 0.0\ntest_df = test_df.set_index('id')\ntest_df = test_df.sort_index()\ntest_df","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TestDataSet:\n    def __init__(self, paths, ids):\n        self.paths = paths\n        self.ids = ids\n\n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, index):\n        x = np.load(self.paths[index])\n        x = preprocess(x).astype(np.float32)\n        return x, self.ids[index]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dataset = TestDataSet(test_df.path.values, test_df.index.values)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=512,\n                                         shuffle=False, drop_last=False, num_workers=num_worker, pin_memory=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"best_model.eval()\nwith torch.no_grad():\n    for data, ids in test_dataloader:\n        data = data.to(device)\n        pred, x = best_model(data)\n        test_df.loc[list(ids),'target'] = pred.cpu().numpy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df.to_csv('submission.csv', columns=['target'])","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}