{"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":"###### CS3319-02 - Project 1 baseline","metadata":{}},{"cell_type":"markdown","source":"测试：加载环境、数据后，跳过start train下的代码，直接运行最后的test block。","metadata":{}},{"cell_type":"markdown","source":"# 环境配置","metadata":{}},{"cell_type":"code","source":"\nimport torch\nif torch.cuda.is_available():\n# GPU enabled\n    !pip install  dgl -f https://data.dgl.ai/wheels/cu113/repo.html\n    !pip install  dglgo -f https://data.dgl.ai/wheels-test/repo.html\nelse:\n# just use CPU\n    !pip install  dgl -f https://data.dgl.ai/wheels/repo.html\n    !pip install  dglgo -f https://data.dgl.ai/wheels-test/repo.html\n    \n!pip install ogb","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:27:08.738718Z","iopub.execute_input":"2023-06-10T11:27:08.739057Z","iopub.status.idle":"2023-06-10T11:28:04.335900Z","shell.execute_reply.started":"2023-06-10T11:27:08.739029Z","shell.execute_reply":"2023-06-10T11:28:04.334641Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Packages","metadata":{}},{"cell_type":"code","source":"import os\nimport pickle as pkl\nos.environ['DGLBACKEND'] = 'pytorch'\nimport numpy as np\nimport pandas as pd\nimport networkx as nx\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport dgl\nimport dgl.nn as dglnn\nimport dgl.function as fn\nfrom dgl.nn.pytorch import MetaPath2Vec\nfrom dgl.nn import HeteroGraphConv, GATConv, SAGEConv, GraphConv\nfrom dgl.dataloading import (\n    DataLoader,\n    MultiLayerFullNeighborSampler,\n    negative_sampler,\n    NeighborSampler,\n)\nimport tqdm\nimport copy\nfrom torch.utils import data","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:04.339253Z","iopub.execute_input":"2023-06-10T11:28:04.340232Z","iopub.status.idle":"2023-06-10T11:28:05.319530Z","shell.execute_reply.started":"2023-06-10T11:28:04.340201Z","shell.execute_reply":"2023-06-10T11:28:05.318544Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Preprocess","metadata":{}},{"cell_type":"markdown","source":"### Get Dataframe","metadata":{}},{"cell_type":"code","source":"\ndef normalize_tensor(t):\n    # 计算每一行的均值和标准差\n    mean = torch.mean(t, dim=1, keepdim=True)\n    std = torch.std(t, dim=1, keepdim=True)\n    # 进行均值方差归一化\n    t_normalized = (t - mean) / std\n    # 将数据缩放到 [0,1] 区间\n    t_normalized = (t_normalized - torch.min(t_normalized)) / (torch.max(t_normalized) - torch.min(t_normalized))\n    return t_normalized\n\n","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:05.320923Z","iopub.execute_input":"2023-06-10T11:28:05.321271Z","iopub.status.idle":"2023-06-10T11:28:05.329673Z","shell.execute_reply.started":"2023-06-10T11:28:05.321239Z","shell.execute_reply":"2023-06-10T11:28:05.327931Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n\ndef normalize_feature(t):\n    max_val = torch.max(t, dim=1, keepdim=True)[0]\n    min_val = torch.min(t, dim=1, keepdim=True)[0]\n    t_normalized = (t - min_val) / (max_val - min_val)\n    return t_normalized\n\n\ndef read_txt(file):\n    res_list = list()\n    with open(file, \"r\") as f:\n        line_list = f.readlines()\n    for line in line_list:\n        res_list.append(list(map(int, line.strip().split(' '))))      \n    return res_list\n\nbase_path = \"/kaggle/input/cs3319-02-project-1-graph-based-recommendation/\"\n\ncite_file = \"paper_file_ann.txt\"\ntrain_ref_file = \"bipartite_train_ann.txt\"\ntest_ref_file = \"bipartite_test_ann.txt\"\ncoauthor_file = \"author_file_ann.txt\"\nfeature_file = \"feature.pkl\"\n\ncitation = read_txt(os.path.join(base_path, cite_file))\nexisting_refs = read_txt(os.path.join(base_path, train_ref_file))\nrefs_to_pred = read_txt(os.path.join(base_path, test_ref_file))\ncoauthor = read_txt(os.path.join(base_path, coauthor_file))\n\nfeature_file = os.path.join(base_path, feature_file)\nwith open(feature_file, 'rb') as f:\n      paper_feature_origin = pkl.load(f)\n\n# print(paper_feature_origin[:5])\npaper_feature = normalize_feature(paper_feature_origin)\n\n\nprint(\"Number of citation edges: {}\\n\\\nNumber of existing references: {}\\n\\\nNumber of author-paper pairs to predict: {}\\n\\\nNumber of coauthor edges: {}\\n\\\nShape of paper features: {}\"\n.format(len(citation), len(existing_refs), len(refs_to_pred), len(coauthor), paper_feature.shape))","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:05.332297Z","iopub.execute_input":"2023-06-10T11:28:05.333003Z","iopub.status.idle":"2023-06-10T11:28:14.890300Z","shell.execute_reply.started":"2023-06-10T11:28:05.332967Z","shell.execute_reply":"2023-06-10T11:28:14.889291Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\n\ncite_edges = pd.DataFrame(citation, columns=['source', 'target'])\ncite_edges = cite_edges.set_index(\n    \"c-\" + cite_edges.index.astype(str)\n)\n\nref_edges = pd.DataFrame(existing_refs, columns=['source', 'target'])\nref_edges = ref_edges.set_index(\n    \"r-\" + ref_edges.index.astype(str)\n)\n\ncoauthor_edges = pd.DataFrame(coauthor, columns=['source', 'target'])\ncoauthor_edges = coauthor_edges.set_index(\n    \"a-\" + coauthor_edges.index.astype(str)\n)\n\npred_ref_edges = pd.DataFrame(refs_to_pred, columns=['source', 'target'])\npred_ref_edges = pred_ref_edges.set_index(\n    \"pr-\" + pred_ref_edges.index.astype(str)\n)\n\n\ncite_edges.head()","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:14.891704Z","iopub.execute_input":"2023-06-10T11:28:14.892050Z","iopub.status.idle":"2023-06-10T11:28:22.654650Z","shell.execute_reply.started":"2023-06-10T11:28:14.892024Z","shell.execute_reply":"2023-06-10T11:28:22.653574Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Paper and Author Node","metadata":{}},{"cell_type":"code","source":"node_tmp = pd.concat([cite_edges.loc[:, 'source'], cite_edges.loc[:, 'target'], ref_edges.loc[:, 'target']])\nnode_papers = pd.DataFrame(index=pd.unique(node_tmp))\n\nnode_tmp = pd.concat([ref_edges['source'], coauthor_edges['source'], coauthor_edges['target']])\nnode_authors = pd.DataFrame(index=pd.unique(node_tmp))\n\nprint(\"Number of paper nodes: {}, number of author nodes: {}\".format(len(node_papers), len(node_authors)))","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:22.657933Z","iopub.execute_input":"2023-06-10T11:28:22.658411Z","iopub.status.idle":"2023-06-10T11:28:22.717359Z","shell.execute_reply.started":"2023-06-10T11:28:22.658373Z","shell.execute_reply":"2023-06-10T11:28:22.716309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Train and Test Ref","metadata":{}},{"cell_type":"code","source":"train_refs = ref_edges.sample(frac=0.9,random_state=0,axis=0)\ntest_true_refs = ref_edges[~ref_edges.index.isin(train_refs.index)]\ntest_true_refs.loc[:, 'label'] = 1\n\nfalse_source = node_authors.sample(frac=test_true_refs.shape[0]/node_authors.shape[0],random_state=0,replace=True,axis=0)\nfalse_target = node_papers.sample(frac=test_true_refs.shape[0]/node_papers.shape[0],random_state=0,replace=True,axis=0)\nfalse_source = false_source.reset_index()\nfalse_target = false_target.reset_index()\ntest_false_refs = pd.concat([false_source, false_target], axis=1)\ntest_false_refs.columns = ['source', 'target']\ntest_false_refs = test_false_refs[test_false_refs.isin(ref_edges) == False]\ntest_false_refs.loc[:, 'label'] = 0\n\ntest_refs = pd.concat([test_true_refs, test_false_refs.iloc[:min(len(false_source), len(false_target))]])\ntest_refs = test_refs.sample(frac=1,random_state=0,axis=0)\ntest_refs.head()","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:22.718758Z","iopub.execute_input":"2023-06-10T11:28:22.719187Z","iopub.status.idle":"2023-06-10T11:28:23.246138Z","shell.execute_reply.started":"2023-06-10T11:28:22.719155Z","shell.execute_reply":"2023-06-10T11:28:23.244978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Heter Graph","metadata":{}},{"cell_type":"code","source":"device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\ntrain_ref_tensor = torch.from_numpy(train_refs.values)\ncite_tensor = torch.from_numpy(cite_edges.values)\ncoauthor_tensor = torch.from_numpy(coauthor_edges.values)\nrel_list = [('author', 'ref', 'paper'), ('paper', 'cite', 'paper'), ('author', 'coauthor', 'author'), ('paper', 'beref', 'author')]\ngraph_data = {\n   rel_list[0]: (train_ref_tensor[:, 0], train_ref_tensor[:, 1]),\n   rel_list[1]: (torch.cat([cite_tensor[:, 0], cite_tensor[:, 1]]), torch.cat([cite_tensor[:, 1], cite_tensor[:, 0]])),\n   rel_list[2]: (torch.cat([coauthor_tensor[:, 0], coauthor_tensor[:, 1]]), torch.cat([coauthor_tensor[:, 1], coauthor_tensor[:, 0]])),\n   rel_list[3]: (train_ref_tensor[:, 1], train_ref_tensor[:, 0])\n}\n\nhetero_graph = dgl.heterograph(graph_data)\nnode_features = {'author': torch.rand(6611, 512), 'paper': paper_feature}\nhetero_graph.ndata['features'] = node_features\nhetero_graph = hetero_graph.to(device)\nhetero_graph","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:23.247783Z","iopub.execute_input":"2023-06-10T11:28:23.248473Z","iopub.status.idle":"2023-06-10T11:28:23.772926Z","shell.execute_reply.started":"2023-06-10T11:28:23.248440Z","shell.execute_reply":"2023-06-10T11:28:23.772018Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Construct Model\n\nThis part constructs a model for generating embedding, and then use the embedding to predict link. please follow  `https://docs.dgl.ai/en/0.9.x/guide_cn/training-link.html`. ","metadata":{}},{"cell_type":"markdown","source":"### Heterograph Conv model\n\nFirst, define a Heterograph Conv model. It calculates conv on every type of edges, then aggregate the message of all types.\nThen define our baseline model.","metadata":{}},{"cell_type":"code","source":"class Encoder(nn.Module):\n    def __init__(self, in_feats, hid_feats, out_feats, edge_names, n_layers, activation, dropout, aggregate='mean'):\n        super().__init__()\n        self.convs = nn.ModuleList()\n        self.activation = activation\n        self.dropout = dropout\n        self.convs.append(HeteroGraphConv({\n            etype: dglnn.GraphConv(in_feats, hid_feats, norm='right')\n            for etype in edge_names\n        }, aggregate=aggregate))\n\n        for _ in range(n_layers - 1):\n            self.convs.append(HeteroGraphConv({\n                etype: dglnn.GraphConv(hid_feats, hid_feats, norm='right')\n                for etype in edge_names\n            }, aggregate=aggregate))\n\n\n\n    def forward(self, blocks, h):\n        layer_feat_list = []\n        min_nodes = {ntype: min([blocks[i].number_of_dst_nodes(ntype) for i in range(len(blocks)) if ntype in blocks[i].dsttypes])\n                     for ntype in blocks[0].dsttypes} \n        for l, (conv, block) in enumerate(zip(self.convs, blocks)):\n            h = conv(block, h)\n            if l != len(self.convs) - 1:\n                for k in h.keys():\n                    h[k] = self.activation(h[k])\n                    h[k] = F.dropout(h[k], self.dropout, self.training)\n            h_dict_cut = {}\n            for ntype in h:\n                num_nodes = block.number_of_dst_nodes(ntype)\n                if num_nodes > min_nodes[ntype]:\n                    h_dict_cut[ntype] = h[ntype][:min_nodes[ntype]]\n                else:\n                    h_dict_cut[ntype] = h[ntype]\n            layer_feat_list.append(h_dict_cut) \n        h = {ntype: torch.cat([layer_feat_list[i][ntype] for i in range(len(layer_feat_list))], dim=-1) for ntype in h}\n        for ntype in h:\n            layer_list_ntype = [layer_feat_list[i][ntype] for i in range(len(layer_feat_list))]\n            layer_feat_mean = torch.mean(torch.stack(layer_list_ntype,dim=0),dim=0,keepdim=True).squeeze()\n            h[ntype] = torch.cat(layer_list_ntype+[layer_feat_mean], dim=-1)\n        del h_dict_cut\n        del layer_feat_list\n        return h","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:23.774448Z","iopub.execute_input":"2023-06-10T11:28:23.774782Z","iopub.status.idle":"2023-06-10T11:28:23.788967Z","shell.execute_reply.started":"2023-06-10T11:28:23.774751Z","shell.execute_reply":"2023-06-10T11:28:23.787980Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 线性模型，通过两个节点的嵌入向量打分\nclass ScorePredictor(nn.Module):\n    def forward(self, edge_subgraph, x):\n        with edge_subgraph.local_scope():\n            edge_subgraph.ndata['x'] = x\n            edge_subgraph.apply_edges(dgl.function.u_dot_v('x', 'x', 'score'))\n            return edge_subgraph.edata['score']\n    def get_score(self, h_u, h_v):\n        result = torch.bmm(h_u.unsqueeze(1),h_v.unsqueeze(1).transpose(1, 2)) \n        return result.squeeze().unsqueeze(1)\n\n\nclass MLPPredictor(nn.Module):\n    def __init__(self, in_features, out_classes=1):\n        super().__init__()\n        self.predictor = nn.Sequential(\n            nn.Linear(2 * in_features, in_features),\n            nn.ReLU(),\n            nn.Linear(in_features, int(in_features/2)),\n            nn.ReLU(),\n            nn.Linear(int(in_features/2), int(in_features/3)),\n            nn.ReLU(),\n            nn.Linear(int(in_features/3), out_classes),\n        )\n        self.jump_conn = nn.Linear(2 * in_features, out_classes)\n        \n\n\n    def get_score(self, h_u, h_v):\n        h = torch.cat([h_u, h_v], 1)\n        return self.predictor(h) + self.jump_conn(h)\n\n    def apply_edges(self, edges):\n        h_u = edges.src['h']\n        h_v = edges.dst['h']\n        h = torch.cat([h_u, h_v], 1)\n        score = self.predictor(h) + self.jump_conn(h)\n        return {'score': score}\n\n    def forward(self, subgraph, h, edge_type=None):\n        # h: 对每种类型的边所计算的节点表示\n        with subgraph.local_scope():\n            subgraph.ndata['h'] = h\n            if not edge_type:\n                for etype in subgraph.canonical_etypes:\n                    subgraph.apply_edges(self.apply_edges, etype=etype)\n                return subgraph.edata['score']\n            else:\n                subgraph.apply_edges(self.apply_edges, etype=edge_type)\n                return subgraph.edges[edge_type].data['score']","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:23.793345Z","iopub.execute_input":"2023-06-10T11:28:23.793849Z","iopub.status.idle":"2023-06-10T11:28:23.807876Z","shell.execute_reply.started":"2023-06-10T11:28:23.793822Z","shell.execute_reply":"2023-06-10T11:28:23.806940Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Model(nn.Module):\n    def __init__(self, in_features, hidden_features, out_features, rel_names, n_layers=3, activation=F.relu, dropout=0.3, pred_opt='base'):\n        super().__init__()\n        self.Encoder = Encoder(in_features, hidden_features, out_features, rel_names, n_layers, activation, dropout)\n        self.pred_MLP = MLPPredictor(out_features, 1)\n        self.pred_base = ScorePredictor()\n        self.hidden_features = hidden_features\n        self.out_features = out_features\n        self.pred_opt = pred_opt\n\n    def forward(self, pos_g, neg_g, blocks, x):\n        torch.cuda.empty_cache()\n        h = self.Encoder(blocks, x)\n        if self.pred_opt == 'base':\n            pos_score = self.pred_base(pos_g, h)\n            neg_score = self.pred_base(neg_g, h)\n        else:\n            pos_score = self.pred_MLP(pos_g, h)\n            neg_score = self.pred_MLP(neg_g, h)         \n        return pos_score, neg_score","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:23.809309Z","iopub.execute_input":"2023-06-10T11:28:23.809830Z","iopub.status.idle":"2023-06-10T11:28:23.820386Z","shell.execute_reply.started":"2023-06-10T11:28:23.809796Z","shell.execute_reply":"2023-06-10T11:28:23.819472Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train and Test","metadata":{}},{"cell_type":"markdown","source":"### 损失函数\n损失函数采用间隔损失：$ L=\\sum_{v_i∼P_n(v),i=1,…,k}max(0,M−y_{u,v}+y_{u,v_i})$","metadata":{}},{"cell_type":"code","source":"from sklearn.metrics import f1_score, accuracy_score\n\n@torch.no_grad()\ndef inference(model:Model, graph, out_dim, n_layers, batch_size):\n    embedding = {ntype: torch.zeros(hetero_graph.num_nodes(ntype), out_dim).to(device) for ntype in hetero_graph.ntypes}\n    indices = {ntype: torch.arange(hetero_graph.num_nodes(ntype)).to(device) for ntype in hetero_graph.ntypes}\n    inference_sampler = dgl.dataloading.MultiLayerFullNeighborSampler(n_layers)\n    inference_dataloader = dgl.dataloading.DataLoader(\n            graph, indices, inference_sampler,\n            device=device,\n            batch_size=batch_size)\n    for step, (input_nodes, seeds, blocks) in enumerate(inference_dataloader):\n        input_features = blocks[0].srcdata['features']\n        output_features = model.Encoder(blocks,input_features)\n        for ntype, idx in seeds.items():\n            embedding[ntype][idx] = output_features[ntype].squeeze()\n    return embedding\n\n\n@torch.no_grad()\ndef test(model:Model, embedding, device, test_refs, res_path='',threshold=0.7):\n    # model = model.to(device)\n    model.eval()\n    test_arr = np.array(test_refs.values)\n    lbl_true = test_refs.label.to_numpy().flatten()\n    if model.pred_opt == 'base':\n        prob_test = model.pred_base.get_score(embedding['author'][test_arr[:, 0]], embedding['paper'][test_arr[:, 1]]).sigmoid()\n    else:\n        prob_test = model.pred_MLP.get_score(embedding['author'][test_arr[:, 0]], embedding['paper'][test_arr[:, 1]]).sigmoid()\n\n    best_pred_score = 0\n    threshold = 0\n    if not res_path:\n        ths_list = [0.1*n for n in range(1,10)] \n    else:\n        ths_list = [0.001*n for n in range(1,1000)]\n    for ths in ths_list:\n        lbl_pred = torch.where(prob_test > ths, torch.tensor(1, device=device), torch.tensor(0,device=device)).cpu().numpy().squeeze()\n        pred_score = f1_score(lbl_true, lbl_pred)\n        if pred_score > best_pred_score:\n            threshold = ths\n            best_pred_score = pred_score    \n\n    if not res_path:\n        return best_pred_score\n    \n    print(f\"best f1_score [{best_pred_score}] with threshold [{threshold}].\")\n    pred_ref_torch_edges = torch.from_numpy(pred_ref_edges.values) \n    pred_network = dgl.heterograph({\n    ('author', 'ref', 'paper'): (pred_ref_torch_edges[:, 0], pred_ref_torch_edges[:, 1])})\n    pred_network = pred_network.to(device)\n    if model.pred_opt == 'base':\n        prob = model.pred_base(pred_network, embedding, ('author', 'ref', 'paper')).sigmoid()\n    else:\n        prob = model.pred_MLP(pred_network, embedding, ('author', 'ref', 'paper')).sigmoid()\n\n    print(prob[:10])\n    pred = torch.where(prob > threshold, torch.tensor(1, device=device), torch.tensor(0,device=device)).cpu().numpy().squeeze()\n    index =  np.arange(0, len(pred))\n    df = pd.DataFrame(np.vstack((index, pred)).T, columns=[\"Index\",\"Predicted\"])\n    df.to_csv(res_path, index=False)\n    return \n","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:23.822185Z","iopub.execute_input":"2023-06-10T11:28:23.822539Z","iopub.status.idle":"2023-06-10T11:28:24.543586Z","shell.execute_reply.started":"2023-06-10T11:28:23.822478Z","shell.execute_reply":"2023-06-10T11:28:24.542645Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nSHOW_BATCH = 5\ndef compute_loss(pos_score, neg_score, etype):\n    n_edges = pos_score[etype].shape[0]\n    return (1 - pos_score[etype].unsqueeze(1) + neg_score[etype].view(n_edges, -1)).clamp(min=0).mean()\n\ndef train(model, g, optimizer, loss_fn, max_epoch, train_eid_dict, k_neg_sample, device, batch_size, \n          n_layers, out_dim=32,\n          save_path=\"./saved_model.pth\", edge_type=('author', 'ref', 'paper')):\n    sampler = dgl.dataloading.NeighborSampler([4,4,5,5])\n    sampler = dgl.dataloading.as_edge_prediction_sampler(\n                    sampler, \n                    negative_sampler=dgl.dataloading.negative_sampler.Uniform(k_neg_sample),       \n                    exclude='reverse_types',\n                    reverse_etypes={'ref':'beref', 'beref':'ref'})\n    dataloader = dgl.dataloading.DataLoader(\n        g, \n        train_eid_dict, \n        sampler, \n        batch_size=batch_size,\n        shuffle=True,\n        device=device)\n    best_f1 = 0\n    all_batch = len(dataloader)\n    for e in range(max_epoch):\n        model.train()\n        epoch_loss = 0\n        for i, (input_nodes, positive_graph, negative_graph, blocks) in enumerate(dataloader):\n                blocks = [b.to(torch.device(device)) for b in blocks]\n                positive_graph = positive_graph.to(torch.device(device))\n                negative_graph = negative_graph.to(torch.device(device))\n                node_features = blocks[0].srcdata['features']\n                pos_score, neg_score = model(positive_graph, negative_graph, blocks, node_features)\n                batch_loss = loss_fn(pos_score, neg_score, edge_type)\n                epoch_loss += batch_loss.item()\n                optimizer.zero_grad()\n                batch_loss.backward()\n                optimizer.step()\n                if i % int(all_batch / 3) == 0:\n                    print(f\"[training]. step:{i}/{all_batch}\")\n        print(f\"[training {e+1}/{max_epoch}] loss:{epoch_loss}\")\n        model.eval()\n        embedding = inference(model, g,out_dim,n_layers, batch_size)  \n        f1 = test(model, embedding, device, test_refs)\n        if f1 > best_f1:\n            print(f\"new best f1_score:{f1}, model saved.\")\n            torch.save(model, save_path)\n            best_f1 = f1\n        del blocks\n        del positive_graph\n        del negative_graph\n        del node_features\n\n    \n\n","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:24.545912Z","iopub.execute_input":"2023-06-10T11:28:24.546521Z","iopub.status.idle":"2023-06-10T11:28:24.561670Z","shell.execute_reply.started":"2023-06-10T11:28:24.546485Z","shell.execute_reply":"2023-06-10T11:28:24.560772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Run training process.","metadata":{}},{"cell_type":"code","source":"# You can change the args and see what will happen to your work\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\n# If you use GPU, the device should be cuda\nprint('Device: {}'.format(device))\ntorch.cuda.empty_cache()\nargs = {\n    'input_dim':512,\n    'device': device,\n    'hidden_dim': 32,\n    'out_dim': 160,\n    'dropout': 0.3,\n    \"n_layers\":4,\n    'lr' : 0.001,\n    'betas' : (0.9, 0.999),\n    'weight_decay' : 0.5e-4,\n    'eps' : 1e-6,\n    'epochs': 240,\n    'test_threshold':0.5,\n    'batch_size': 32768,\n    \"k_neg_sample\":5,\n}\ncheckpoint = \"./GCN_model_MLP.pth\"\ndevice = args['device']\nhetero_graph.to(device)\nargs","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:24.563237Z","iopub.execute_input":"2023-06-10T11:28:24.563647Z","iopub.status.idle":"2023-06-10T11:28:24.577553Z","shell.execute_reply.started":"2023-06-10T11:28:24.563616Z","shell.execute_reply":"2023-06-10T11:28:24.576667Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Start Train","metadata":{}},{"cell_type":"code","source":"author_feature = hetero_graph.nodes['author'].data['features']\npaper_feature = hetero_graph.nodes['paper'].data['features']\nfeature =  {'author': author_feature.to(device), 'paper': paper_feature.to(device)}\ntrain_eid_dict = {\n    etype: hetero_graph.edges(etype=etype, form='eid')\n    for etype in hetero_graph.etypes}\n\nmodel = Model(args['input_dim'], args['hidden_dim'],args['out_dim'],\n    hetero_graph.etypes, args[\"n_layers\"], F.elu, args['dropout']).to(device)\n\n\n\noptimizer_all = torch.optim.AdamW(model.parameters(), lr=args['lr'])\n\nloss_fn = compute_loss\ntrain(model, hetero_graph, optimizer_all, loss_fn, args[\"epochs\"], train_eid_dict,  args[\"k_neg_sample\"], \n      device, args['batch_size'], args[\"n_layers\"],out_dim= args[\"out_dim\"], save_path=checkpoint)\nprint(\"[base model with simple score predictor trained]\")\nmodel = torch.load(checkpoint)\nmodel.pred_opt='MLP'\nevaluator_param = list(map(id, model.pred_MLP.parameters()))\nbase_params = filter(lambda p: id(p) not in evaluator_param, model.parameters())\nmlp_train_params = [{'params': base_params, 'lr': args['lr']*0.0000001},\n          {'params': model.pred_MLP.parameters(), 'lr': args['lr']*5}]\noptimizer_MLP = torch.optim.AdamW(mlp_train_params, lr=args['lr'])\ntrain(model, hetero_graph, optimizer_MLP, loss_fn, int(args[\"epochs\"]/2), train_eid_dict,  args[\"k_neg_sample\"], \n      device, args['batch_size'], args[\"n_layers\"],out_dim= args[\"out_dim\"], save_path=checkpoint)\nprint(\"[MLP score predictor trained]\")\nmodel = torch.load(checkpoint)\noptimizer_all = torch.optim.AdamW(model.parameters(), lr=args['lr'])\ntrain(model, hetero_graph, optimizer_all, loss_fn, args[\"epochs\"], train_eid_dict,  args[\"k_neg_sample\"], \n      device, args['batch_size'], args[\"n_layers\"],out_dim= args[\"out_dim\"], save_path=checkpoint)\nprint(\"[All model trained]\")\ntorch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:28:24.579069Z","iopub.execute_input":"2023-06-10T11:28:24.579758Z","iopub.status.idle":"2023-06-10T11:41:26.606673Z","shell.execute_reply.started":"2023-06-10T11:28:24.579721Z","shell.execute_reply":"2023-06-10T11:41:26.605628Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Test\n\n可以将`checkpoint`的赋值语句取消注释后，直接运行该block，跳过训练过程。","metadata":{}},{"cell_type":"code","source":"# Test the model and print test accuracy\ntorch.cuda.empty_cache()\n\"\"\"You can uncomment command below to directly test on pre-trained model\"\"\"\n#################################################\n# checkpoint = \"/kaggle/input/resultofskipgcn/GCN_model_MLP.pth\"\n#################################################\nmodel = torch.load(checkpoint)\nmodel.pred_opt='MLP'\nmodel.eval()\n\nembedding = inference(model, hetero_graph, args['out_dim'],args['n_layers'], args['batch_size'])\ntest(model, embedding, device, test_refs,'pred.csv')\n","metadata":{"execution":{"iopub.status.busy":"2023-06-10T11:41:26.644567Z","iopub.execute_input":"2023-06-10T11:41:26.645193Z"},"trusted":true},"execution_count":null,"outputs":[]}]}