{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import torch\nfrom torch import nn\nfrom torch.nn import Parameter\nimport torch.nn.functional as F\n\nclass GATConvLayer_new(nn.Module):\n    def __init__(self,\n                 node_in_fts,\n                 node_out_fts,\n                 ed_in_fts,\n                 ed_out_fts,\n                 state_dim,\n                 concat=True,\n                 negative_slope=0.1,\n                 dropout=0.1,\n                 bias=True,\n                 repeat_edge = 1\n                 ):\n        super(GATConvLayer_new, self).__init__()\n \n        self.node_in_fts = node_in_fts\n        self.node_out_fts = node_out_fts\n        self.ed_out_fts = ed_out_fts\n        self.ed_in_fts = ed_in_fts\n        self.concat = concat\n        self.negative_slope = negative_slope\n        self.dropout = dropout\n        self.repeat_edge = repeat_edge\n        self.state_dim = state_dim\n \n        self.node_update_layer = nn.Linear(self.node_in_fts, node_out_fts)\n        self.ed_trans_layer = nn.Linear(2 * self.node_in_fts + self.ed_in_fts, 2 * self.state_dim)\n        self.ed_hidden_1 = nn.Linear(2 * self.state_dim, 2 * state_dim)\n        self.ed_hidden_2 = nn.Linear(2 * self.state_dim, state_dim)\n        self.ed_out_layer = nn.Linear(self.state_dim, self.ed_out_fts)\n        self.e = Parameter(torch.zeros(size=(2 * self.node_out_fts + self.ed_out_fts, 1)))\n\n        self.batch_norm_n = nn.BatchNorm1d(self.node_out_fts)\n        self.batch_norm_e = nn.BatchNorm1d(self.ed_out_fts)\n        \n        if bias and concat:\n            self.bias = Parameter(torch.Tensor(node_out_fts))\n        elif bias and not concat:\n            self.bias = Parameter(torch.Tensor(node_out_fts))\n        else:\n            self.register_parameter('bias', None)\n \n        self.reset_parameters()\n \n    def reset_parameters(self):\n        nn.init.xavier_uniform_(self.e.data, gain=1.414)\n \n    def forward(self, x, edge_attr, mask):\n        return self.propagate( x=x, num_nodes=x.size(0), edge_attr=edge_attr, mask = mask)\n \n    def propagate(self, **kwargs):\n        edge_attr = kwargs['edge_attr']\n        x = kwargs['x']\n        N = kwargs['num_nodes']\n        mask = kwargs['mask']\n\n        triplet_e = torch.cat([x.repeat(1, self.repeat_edge*N).view(self.repeat_edge*N*N, -1), edge_attr, x.repeat(self.repeat_edge*N, 1)], dim=-1)\n        new_edge_attr = F.leaky_relu(self.ed_trans_layer(triplet_e), self.negative_slope)\n        new_edge_attr = F.leaky_relu(self.ed_hidden_1(triplet_e), self.negative_slope)\n        new_edge_attr = F.leaky_relu(self.ed_hidden_2(triplet_e), self.negative_slope)\n        new_edge_attr = F.leaky_relu(self.ed_out_layer(new_edge_attr), self.negative_slope)\n        new_edge_attr = new_edge_attr * mask.repeat(1, new_edge_attr.shape[1])\n\n        x = F.leaky_relu(self.node_update_layer(x), self.negative_slope)\n        triplet_n = torch.cat([x.repeat(1, self.repeat_edge*N).view(self.repeat_edge*N*N, -1), new_edge_attr, x.repeat(self.repeat_edge*N, 1)], dim=-1)\n        triplet_n = triplet_n.view(N, -1, 2 * self.node_out_fts + self.ed_out_fts)\n        e = F.leaky_relu(torch.matmul(triplet_n, self.e).squeeze(2), self.negative_slope)\n        attention = F.softmax(e, dim = 1)\n        attention = attention.view(N * N, -1)\n        attention = attention * mask\n        attention = attention.view(N, N)\n        attention = F.dropout(attention, self.dropout, training=self.training)\n        new_x = torch.matmul(attention, x)\n\n        new_x = self.batch_norm_n(new_x)\n        new_edge_attr = self.batch_norm_e(new_edge_attr)\n\n        return [new_x, new_edge_attr]\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import math\nimport torch\nfrom torch.optim.optimizer import Optimizer, required\n\nclass RAdam(Optimizer):\n\n    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):\n        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)\n        self.buffer = [[None, None, None] for ind in range(10)]\n        super(RAdam, self).__init__(params, defaults)\n\n    def __setstate__(self, state):\n        super(RAdam, self).__setstate__(state)\n\n    def step(self, closure=None):\n\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        for group in self.param_groups:\n\n            for p in group['params']:\n                if p.grad is None:\n                    continue\n                grad = p.grad.data.float()\n                if grad.is_sparse:\n                    raise RuntimeError('RAdam does not support sparse gradients')\n\n                p_data_fp32 = p.data.float()\n\n                state = self.state[p]\n\n                if len(state) == 0:\n                    state['step'] = 0\n                    state['exp_avg'] = torch.zeros_like(p_data_fp32)\n                    state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)\n                else:\n                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)\n                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)\n\n                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n                beta1, beta2 = group['betas']\n\n                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n                exp_avg.mul_(beta1).add_(1 - beta1, grad)\n\n                state['step'] += 1\n                buffered = self.buffer[int(state['step'] % 10)]\n                if state['step'] == buffered[0]:\n                    N_sma, step_size = buffered[1], buffered[2]\n                else:\n                    buffered[0] = state['step']\n                    beta2_t = beta2 ** state['step']\n                    N_sma_max = 2 / (1 - beta2) - 1\n                    N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)\n                    buffered[1] = N_sma\n\n                    # more conservative since it's an approximated value\n                    if N_sma >= 5:\n                        step_size = math.sqrt((1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / (N_sma_max - 2)) / (1 - beta1 ** state['step'])\n                    else:\n                        step_size = 1.0 / (1 - beta1 ** state['step'])\n                    buffered[2] = step_size\n\n                if group['weight_decay'] != 0:\n                    p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)\n\n                # more conservative since it's an approximated value\n                if N_sma >= 5:            \n                    denom = exp_avg_sq.sqrt().add_(group['eps'])\n                    p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom)\n                else:\n                    p_data_fp32.add_(-step_size * group['lr'], exp_avg)\n\n                p.data.copy_(p_data_fp32)\n\n        return loss\n\nclass PlainRAdam(Optimizer):\n\n    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):\n        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)\n\n        super(PlainRAdam, self).__init__(params, defaults)\n\n    def __setstate__(self, state):\n        super(PlainRAdam, self).__setstate__(state)\n\n    def step(self, closure=None):\n\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        for group in self.param_groups:\n\n            for p in group['params']:\n                if p.grad is None:\n                    continue\n                grad = p.grad.data.float()\n                if grad.is_sparse:\n                    raise RuntimeError('RAdam does not support sparse gradients')\n\n                p_data_fp32 = p.data.float()\n\n                state = self.state[p]\n\n                if len(state) == 0:\n                    state['step'] = 0\n                    state['exp_avg'] = torch.zeros_like(p_data_fp32)\n                    state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)\n                else:\n                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)\n                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)\n\n                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n                beta1, beta2 = group['betas']\n\n                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n                exp_avg.mul_(beta1).add_(1 - beta1, grad)\n\n                state['step'] += 1\n                beta2_t = beta2 ** state['step']\n                N_sma_max = 2 / (1 - beta2) - 1\n                N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)\n\n                if group['weight_decay'] != 0:\n                    p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)\n\n                # more conservative since it's an approximated value\n                if N_sma >= 5:                    \n                    step_size = group['lr'] * math.sqrt((1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / (N_sma_max - 2)) / (1 - beta1 ** state['step'])\n                    denom = exp_avg_sq.sqrt().add_(group['eps'])\n                    p_data_fp32.addcdiv_(-step_size, exp_avg, denom)\n                else:\n                    step_size = group['lr'] / (1 - beta1 ** state['step'])\n                    p_data_fp32.add_(-step_size, exp_avg)\n\n                p.data.copy_(p_data_fp32)\n\n        return loss\n\n\nclass AdamW(Optimizer):\n\n    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, warmup = 0):\n        defaults = dict(lr=lr, betas=betas, eps=eps,\n                        weight_decay=weight_decay, warmup = warmup)\n        super(AdamW, self).__init__(params, defaults)\n\n    def __setstate__(self, state):\n        super(AdamW, self).__setstate__(state)\n\n    def step(self, closure=None):\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        for group in self.param_groups:\n\n            for p in group['params']:\n                if p.grad is None:\n                    continue\n                grad = p.grad.data.float()\n                if grad.is_sparse:\n                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')\n\n                p_data_fp32 = p.data.float()\n\n                state = self.state[p]\n\n                if len(state) == 0:\n                    state['step'] = 0\n                    state['exp_avg'] = torch.zeros_like(p_data_fp32)\n                    state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)\n                else:\n                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)\n                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)\n\n                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n                beta1, beta2 = group['betas']\n\n                state['step'] += 1\n\n                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n                exp_avg.mul_(beta1).add_(1 - beta1, grad)\n\n                denom = exp_avg_sq.sqrt().add_(group['eps'])\n                bias_correction1 = 1 - beta1 ** state['step']\n                bias_correction2 = 1 - beta2 ** state['step']\n                \n                if group['warmup'] > state['step']:\n                    scheduled_lr = 1e-8 + state['step'] * group['lr'] / group['warmup']\n                else:\n                    scheduled_lr = group['lr']\n\n                step_size = scheduled_lr * math.sqrt(bias_correction2) / bias_correction1\n                \n                if group['weight_decay'] != 0:\n                    p_data_fp32.add_(-group['weight_decay'] * scheduled_lr, p_data_fp32)\n\n                p_data_fp32.addcdiv_(-step_size, exp_avg, denom)\n\n                p.data.copy_(p_data_fp32)\n\n        return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np \n\nclass GAT_skip_connect(nn.Module):\n    def __init__(self, n_node_feat, n_edge_feat, node_embedding_dim, edge_embedding_dim, state_dim, num_layers, concat=True, negative_slope=0.2, dropout=0, bias=True, repeat_edge = 1):\n        super(GAT_skip_connect, self).__init__()   \n        self.dropout = dropout\n        self.negative_slope = negative_slope\n        self.n_node_feat = n_node_feat\n        self.n_edge_feat = n_edge_feat\n        self.state_dim = state_dim\n        self.num_layers = num_layers\n        self.concat = concat\n        self.node_embedding_dim = node_embedding_dim\n        self.edge_embedding_dim = edge_embedding_dim\n        self.bias = bias\n        self.repeat_edge = repeat_edge\n        self.Node_embedd = nn.Linear(self.n_node_feat, self.node_embedding_dim, bias = False)\n        self.Edge_embedd = nn.Linear(self.n_edge_feat, self.edge_embedding_dim, bias = False)\n\n        self.GATConvs = [GATConvLayer_new(self.node_embedding_dim, self.node_embedding_dim, self.edge_embedding_dim, self.edge_embedding_dim, self.state_dim) for _ in range(num_layers)]\n\n        self.fc1 = nn.Linear(2 * self.node_embedding_dim + self.edge_embedding_dim, 2 * self.node_embedding_dim + self.edge_embedding_dim)\n        self.fc2 = nn.Linear(2 * self.node_embedding_dim + self.edge_embedding_dim, self.node_embedding_dim + self.edge_embedding_dim)\n        self.fc3 = nn.Linear(self.node_embedding_dim + self.edge_embedding_dim, self.node_embedding_dim + self.edge_embedding_dim)\n        self.fc4 = nn.Linear(self.node_embedding_dim + self.edge_embedding_dim, self.edge_embedding_dim)\n        self.reg = nn.Linear(self.edge_embedding_dim, 1)\n\n        self.dropout = nn.Dropout(0.1)\n\n    \n    def forward(self, x, edge_feat, mask_out):\n        len_edges = edge_feat.shape[-1]\n        idx_edge = edge_feat[:,len_edges - 1: len_edges]\n        mask = torch.where(idx_edge == 0, idx_edge, torch.ones_like(idx_edge).cuda())\n        N = x.size(0)\n        x = self.Node_embedd(x)\n        edge_feat = self.Edge_embedd(edge_feat)\n       \n        x_ = x\n        edge_feat_ = edge_feat\n\n        for i, GATConv in enumerate(self.GATConvs):\n            GATConv = GATConv.cuda()\n            x, edge_feat = GATConv(x, edge_feat, mask)\n            x = F.leaky_relu(x, self.negative_slope)\n            edge_feat = F.leaky_relu(edge_feat, self.negative_slope)\n\n            x = x + x_\n            edge_feat = edge_feat + edge_feat_\n\n            edge_feat_ = edge_feat\n            x_ = x\n        \n        # concat edge feature and node feature\n        out = torch.cat([edge_feat, x.repeat(1, N).view(N * N,  -1), x.repeat(N, 1)], dim = 1)        \n        out = F.leaky_relu(self.fc1(out), self.negative_slope) \n        out = F.leaky_relu(self.fc2(out), self.negative_slope)  \n        out = F.leaky_relu(self.fc3(out), self.negative_slope)  \n        out = F.leaky_relu(self.fc4(out), self.negative_slope)  \n\n        # compute output\n        out = self.reg(out)\n        out = out * mask\n        out = out.reshape(29*29)\n        out = out[mask_out]\n        return out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from __future__ import division\nfrom __future__ import print_function\n\nimport os\nimport glob\nimport time\nimport random\nimport argparse\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nfrom torch.autograd import Variable\n\nimport torch.multiprocessing\ntorch.multiprocessing.set_sharing_strategy('file_system')\nimport random \nimport pandas as pd\nimport pickle\n\n\ndef log_mae(output, lable):\n    output = output.view(-1)\n    lable = lable.view(-1)\n    loss = torch.abs(output - lable)\n    loss = loss.mean()\n    loss = torch.log(loss)\n    return loss\n\n\n\n# Model and optimizer\nmodel = GAT_skip_connect(n_node_feat = 11, n_edge_feat = 26, node_embedding_dim = 256, edge_embedding_dim = 512, state_dim=512, num_layers=3)\n# model.load_state_dict(torch.load('../input/maskoutmodel/new_GATConv_3layer_4fcEdgeoutlayer_normout_no_norm_nonzero_maskout_Kagglellpv_v1.pkl'))\n\nmodel = model.cuda()\nmodel.share_memory()\n\n# optimizer = optim.Adam(model.parameters(), \n#                        lr=0.00001, \n#                        weight_decay=5e-4, amsgrad=False)\noptimizer = RAdam(model.parameters(), lr=0.00001, weight_decay=5e-4)\ntotal_loss = {}\nbest_loss = torch.Tensor([10000000000]).cuda()\nbest_epoch = 10000\n\n\n# # Load data, features are tensor 2708, 1433; lables are tensor 2708,; adj is adjacency matrix 2708 * 2708;\nin_edges_train = np.load('../input/internalgraphdata/in_edges_train.npz')['arr_0']\nnodes_train = np.load('../input/internalgraphdata/nodes_train.npz')['arr_0']\nout_edges_train = np.load('../input/extradata/out_edges_train_1.npz', allow_pickle=True)['arr_0']\n#train_ = pd.read_csv(\"../input/champs-scalar-coupling/train.csv\")\n\nnum_molecules_train = 85003\nrandom.seed(12)\ntrain_total_size = len(nodes_train)\ntrain_size = int(train_total_size * 0.9)\nnodes_val = nodes_train[train_size:]\nin_edges_val = in_edges_train[train_size:]\nout_edges_val = out_edges_train[train_size:]\n\ndef train(epoch):\n    global best_epoch, best_loss\n    model.train()\n    loss_epoch = torch.zeros(1).cuda()\n\n    print('--train--')\n    for i, node_train in enumerate(nodes_train[:train_size]):\n        optimizer.zero_grad()\n        \n        node_train = torch.Tensor(node_train).cuda()\n        node_train = node_train.view(29, -1)\n        edge_train = torch.Tensor(in_edges_train[i]).cuda()\n        edge_train = edge_train.view(29 * 29, -1)\n\n        out_edge_train = torch.Tensor(out_edges_train[i]).cuda()\n        out_edge_train = out_edge_train.view(29 * 29)\n        mask_out = out_edge_train.nonzero()\n        out_edge_train = out_edge_train[mask_out]\n        \n\n        out_edge_train  = Variable(out_edge_train)\n        edge_train = Variable(edge_train)\n        node_train = Variable(node_train)\n\n        output = model(node_train, edge_train, mask_out)\n        loss_train = log_mae(out_edge_train, output)\n        loss_train.backward()\n        optimizer.step()\n        loss_epoch += loss_train.item()\n        if i % 1000 == 0:\n            print(\"Train: \")\n            print('epoch ', epoch, i, loss_epoch/(i+1))\n    loss_epoch /= train_size\n    if loss_epoch < best_loss:\n        best_loss = loss_epoch\n        torch.save(model.state_dict(), 'new_GATConv_3layer_4fcEdgeoutlayer_normout_no_norm_nonzero_maskout_Kagglellpv_v1.pkl')\n        best_epoch = epoch\n\n    \n    total_loss[epoch] = loss_epoch\n    print('Loss epoch ', loss_epoch)\n\n    print('Done ' + str(epoch))\n\n\n\ndef val(epoch):\n    model.eval()\n    loss_epoch = torch.zeros(1).cuda()\n    print('--validation--')\n    with torch.no_grad():\n        for i, node_train in enumerate(nodes_val):\n            node_train = torch.Tensor(node_train).cuda()\n            node_train = node_train.view(29, -1)\n            edge_train = torch.Tensor(in_edges_val[i]).cuda()\n            edge_train = edge_train.view(29 * 29, -1)\n            out_edge_train = torch.Tensor(out_edges_val[i]).cuda()\n            out_edge_train = out_edge_train.view(29 * 29)\n            mask_out = out_edge_train.nonzero()\n            out_edge_train = out_edge_train[mask_out]\n            \n                \n            out_edge_train = Variable(out_edge_train)\n            edge_train = Variable(edge_train)\n            node_train = Variable(node_train)\n\n            output = model(node_train, edge_train, mask_out)\n            loss_train = log_mae(out_edge_train, output)\n            loss_epoch += loss_train.item()\n            if i % 2000 == 0:\n                print('validattion: ')\n                print('epoch ', epoch, i, loss_epoch/(i+1))\n    loss_epoch /= (train_total_size - train_size)\n    print('Loss epoch ', loss_epoch)\n\nfor i in range(4):\n    train(i)\n    val(i)\n# train(0)\n","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":1}