├── .gitignore ├── LICENSE ├── Readme.md ├── config └── log_config.json ├── data_compressed ├── FB15k-237.zip └── WN18RR.zip ├── data_loader.py ├── helper.py ├── model ├── __init__.py ├── compgcn_conv.py ├── compgcn_conv_basis.py ├── message_passing.py └── models.py ├── overview.png ├── preprocess.sh ├── requirements.txt └── run.py /.gitignore: -------------------------------------------------------------------------------- 1 | data/* 2 | checkpoints/* 3 | log/* 4 | 5 | # Byte-compiled / optimized / DLL files 6 | __pycache__/ 7 | *.py[cod] 8 | *$py.class 9 | 10 | # C extensions 11 | *.so 12 | 13 | # Distribution / packaging 14 | .Python 15 | build/ 16 | develop-eggs/ 17 | dist/ 18 | downloads/ 19 | eggs/ 20 | .eggs/ 21 | lib/ 22 | lib64/ 23 | parts/ 24 | sdist/ 25 | var/ 26 | wheels/ 27 | *.egg-info/ 28 | .installed.cfg 29 | *.egg 30 | MANIFEST 31 | 32 | # PyInstaller 33 | # Usually these files are written by a python script from a template 34 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 35 | *.manifest 36 | *.spec 37 | 38 | # Installer logs 39 | pip-log.txt 40 | pip-delete-this-directory.txt 41 | 42 | # Unit test / coverage reports 43 | htmlcov/ 44 | .tox/ 45 | .coverage 46 | .coverage.* 47 | .cache 48 | nosetests.xml 49 | coverage.xml 50 | *.cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | 63 | # Flask stuff: 64 | instance/ 65 | .webassets-cache 66 | 67 | # Scrapy stuff: 68 | .scrapy 69 | 70 | # Sphinx documentation 71 | docs/_build/ 72 | 73 | # PyBuilder 74 | target/ 75 | 76 | # Jupyter Notebook 77 | .ipynb_checkpoints 78 | 79 | # pyenv 80 | .python-version 81 | 82 | # celery beat schedule file 83 | celerybeat-schedule 84 | 85 | # SageMath parsed files 86 | *.sage.py 87 | 88 | # Environments 89 | .env 90 | .venv 91 | env/ 92 | venv/ 93 | ENV/ 94 | env.bak/ 95 | venv.bak/ 96 | 97 | # Spyder project settings 98 | .spyderproject 99 | .spyproject 100 | 101 | # Rope project settings 102 | .ropeproject 103 | 104 | # mkdocs documentation 105 | /site 106 | 107 | # mypy 108 | .mypy_cache/ 109 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /Readme.md: -------------------------------------------------------------------------------- 1 |

2 | CompGCN 3 |

4 | 5 |

Composition-Based Multi-Relational Graph Convolutional Networks

6 | 7 |

8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 |

16 | 17 | 18 |

19 | Overview of CompGCN 20 | ... 21 |

22 | Given node and relation embeddings, CompGCN performs a composition operation φ(·) over each edge in the neighborhood of a central node (e.g. Christopher Nolan above). The composed embeddings are then convolved with specific filters WO and WI for original and inverse relations respectively. We omit self-loop in the diagram for clarity. The message from all the neighbors are then aggregated to get an updated embedding of the central node. Also, the relation embeddings are transformed using a separate weight matrix. Please refer to the paper for details. 23 | 24 | ### Dependencies 25 | 26 | - Compatible with PyTorch 1.0 and Python 3.x. 27 | - Dependencies can be installed using `requirements.txt`. 28 | 29 | ### Dataset: 30 | 31 | - We use FB15k-237 and WN18RR dataset for knowledge graph link prediction. 32 | - FB15k-237 and WN18RR are included in the `data` directory. 33 | 34 | ### Training model: 35 | 36 | - Install all the requirements from `requirements.txt.` 37 | 38 | - Execute `./setup.sh` for extracting the dataset and setting up the folder hierarchy for experiments. 39 | 40 | - Commands for reproducing the reported results on link prediction: 41 | 42 | ```shell 43 | ##### with TransE Score Function 44 | # CompGCN (Composition: Subtraction) 45 | python run.py -score_func transe -opn sub -gamma 9 -hid_drop 0.1 -init_dim 200 46 | 47 | # CompGCN (Composition: Multiplication) 48 | python run.py -score_func transe -opn mult -gamma 9 -hid_drop 0.2 -init_dim 200 49 | 50 | # CompGCN (Composition: Circular Correlation) 51 | python run.py -score_func transe -opn corr -gamma 40 -hid_drop 0.1 -init_dim 200 52 | 53 | ##### with DistMult Score Function 54 | # CompGCN (Composition: Subtraction) 55 | python run.py -score_func distmult -opn sub -gcn_dim 150 -gcn_layer 2 56 | 57 | # CompGCN (Composition: Multiplication) 58 | python run.py -score_func distmult -opn mult -gcn_dim 150 -gcn_layer 2 59 | 60 | # CompGCN (Composition: Circular Correlation) 61 | python run.py -score_func distmult -opn corr -gcn_dim 150 -gcn_layer 2 62 | 63 | ##### with ConvE Score Function 64 | # CompGCN (Composition: Subtraction) 65 | python run.py -score_func conve -opn sub -ker_sz 5 66 | 67 | # CompGCN (Composition: Multiplication) 68 | python run.py -score_func conve -opn mult 69 | 70 | # CompGCN (Composition: Circular Correlation) 71 | python run.py -score_func conve -opn corr 72 | 73 | ##### Overall BEST: 74 | python run.py -name best_model -score_func conve -opn corr 75 | ``` 76 | 77 | - `-score_func` denotes the link prediction score score function 78 | - `-opn` is the composition operation used in **CompGCN**. It can take the following values: 79 | - `sub` for subtraction operation: Φ(e_s, e_r) = e_s - e_r 80 | - `mult` for multiplication operation: Φ(e_s, e_r) = e_s * e_r 81 | - `corr` for circular-correlation: Φ(e_s, e_r) = e_s ★ e_r 82 | - `-name` is some name given for the run (used for storing model parameters) 83 | - `-model` is name of the model `compgcn'. 84 | - `-gpu` for specifying the GPU to use 85 | - Rest of the arguments can be listed using `python run.py -h` 86 | ### Citation: 87 | Please cite the following paper if you use this code in your work. 88 | ```bibtex 89 | @inproceedings{ 90 | vashishth2020compositionbased, 91 | title={Composition-based Multi-Relational Graph Convolutional Networks}, 92 | author={Shikhar Vashishth and Soumya Sanyal and Vikram Nitin and Partha Talukdar}, 93 | booktitle={International Conference on Learning Representations}, 94 | year={2020}, 95 | url={https://openreview.net/forum?id=BylA_C4tPr} 96 | } 97 | ``` 98 | For any clarification, comments, or suggestions please create an issue or contact [Shikhar](http://shikhar-vashishth.github.io). 99 | -------------------------------------------------------------------------------- /config/log_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "version": 1, 3 | "disable_existing_loggers": false, 4 | "formatters": { 5 | "simple": { 6 | "format": "%(asctime)s - %(name)s - [%(levelname)s] - %(message)s" 7 | } 8 | }, 9 | 10 | "handlers": { 11 | "file_handler": { 12 | "class": "logging.FileHandler", 13 | "level": "DEBUG", 14 | "formatter": "simple", 15 | "filename": "python_logging.log", 16 | "encoding": "utf8" 17 | } 18 | }, 19 | 20 | "root": { 21 | "level": "DEBUG", 22 | "handlers": ["file_handler"] 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /data_compressed/FB15k-237.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/malllabiisc/CompGCN/3b06f5fec42526faa81afc158df9e64a8382982c/data_compressed/FB15k-237.zip -------------------------------------------------------------------------------- /data_compressed/WN18RR.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/malllabiisc/CompGCN/3b06f5fec42526faa81afc158df9e64a8382982c/data_compressed/WN18RR.zip -------------------------------------------------------------------------------- /data_loader.py: -------------------------------------------------------------------------------- 1 | from helper import * 2 | from torch.utils.data import Dataset 3 | 4 | class TrainDataset(Dataset): 5 | """ 6 | Training Dataset class. 7 | 8 | Parameters 9 | ---------- 10 | triples: The triples used for training the model 11 | params: Parameters for the experiments 12 | 13 | Returns 14 | ------- 15 | A training Dataset class instance used by DataLoader 16 | """ 17 | def __init__(self, triples, params): 18 | self.triples = triples 19 | self.p = params 20 | self.entities = np.arange(self.p.num_ent, dtype=np.int32) 21 | 22 | def __len__(self): 23 | return len(self.triples) 24 | 25 | def __getitem__(self, idx): 26 | ele = self.triples[idx] 27 | triple, label, sub_samp = torch.LongTensor(ele['triple']), np.int32(ele['label']), np.float32(ele['sub_samp']) 28 | trp_label = self.get_label(label) 29 | 30 | if self.p.lbl_smooth != 0.0: 31 | trp_label = (1.0 - self.p.lbl_smooth)*trp_label + (1.0/self.p.num_ent) 32 | 33 | return triple, trp_label, None, None 34 | 35 | @staticmethod 36 | def collate_fn(data): 37 | triple = torch.stack([_[0] for _ in data], dim=0) 38 | trp_label = torch.stack([_[1] for _ in data], dim=0) 39 | return triple, trp_label 40 | 41 | def get_neg_ent(self, triple, label): 42 | def get(triple, label): 43 | pos_obj = label 44 | mask = np.ones([self.p.num_ent], dtype=np.bool) 45 | mask[label] = 0 46 | neg_ent = np.int32(np.random.choice(self.entities[mask], self.p.neg_num - len(label), replace=False)).reshape([-1]) 47 | neg_ent = np.concatenate((pos_obj.reshape([-1]), neg_ent)) 48 | 49 | return neg_ent 50 | 51 | neg_ent = get(triple, label) 52 | return neg_ent 53 | 54 | def get_label(self, label): 55 | y = np.zeros([self.p.num_ent], dtype=np.float32) 56 | for e2 in label: y[e2] = 1.0 57 | return torch.FloatTensor(y) 58 | 59 | 60 | class TestDataset(Dataset): 61 | """ 62 | Evaluation Dataset class. 63 | 64 | Parameters 65 | ---------- 66 | triples: The triples used for evaluating the model 67 | params: Parameters for the experiments 68 | 69 | Returns 70 | ------- 71 | An evaluation Dataset class instance used by DataLoader for model evaluation 72 | """ 73 | def __init__(self, triples, params): 74 | self.triples = triples 75 | self.p = params 76 | 77 | def __len__(self): 78 | return len(self.triples) 79 | 80 | def __getitem__(self, idx): 81 | ele = self.triples[idx] 82 | triple, label = torch.LongTensor(ele['triple']), np.int32(ele['label']) 83 | label = self.get_label(label) 84 | 85 | return triple, label 86 | 87 | @staticmethod 88 | def collate_fn(data): 89 | triple = torch.stack([_[0] for _ in data], dim=0) 90 | label = torch.stack([_[1] for _ in data], dim=0) 91 | return triple, label 92 | 93 | def get_label(self, label): 94 | y = np.zeros([self.p.num_ent], dtype=np.float32) 95 | for e2 in label: y[e2] = 1.0 96 | return torch.FloatTensor(y) -------------------------------------------------------------------------------- /helper.py: -------------------------------------------------------------------------------- 1 | import numpy as np, sys, os, random, pdb, json, uuid, time, argparse 2 | from pprint import pprint 3 | import logging, logging.config 4 | from collections import defaultdict as ddict 5 | from ordered_set import OrderedSet 6 | 7 | # PyTorch related imports 8 | import torch 9 | from torch.nn import functional as F 10 | from torch.nn.init import xavier_normal_ 11 | from torch.utils.data import DataLoader 12 | from torch.nn import Parameter 13 | from torch_scatter import scatter_add 14 | 15 | np.set_printoptions(precision=4) 16 | 17 | def set_gpu(gpus): 18 | """ 19 | Sets the GPU to be used for the run 20 | 21 | Parameters 22 | ---------- 23 | gpus: List of GPUs to be used for the run 24 | 25 | Returns 26 | ------- 27 | 28 | """ 29 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 30 | os.environ["CUDA_VISIBLE_DEVICES"] = gpus 31 | 32 | def get_logger(name, log_dir, config_dir): 33 | """ 34 | Creates a logger object 35 | 36 | Parameters 37 | ---------- 38 | name: Name of the logger file 39 | log_dir: Directory where logger file needs to be stored 40 | config_dir: Directory from where log_config.json needs to be read 41 | 42 | Returns 43 | ------- 44 | A logger object which writes to both file and stdout 45 | 46 | """ 47 | config_dict = json.load(open( config_dir + 'log_config.json')) 48 | config_dict['handlers']['file_handler']['filename'] = log_dir + name.replace('/', '-') 49 | logging.config.dictConfig(config_dict) 50 | logger = logging.getLogger(name) 51 | 52 | std_out_format = '%(asctime)s - [%(levelname)s] - %(message)s' 53 | consoleHandler = logging.StreamHandler(sys.stdout) 54 | consoleHandler.setFormatter(logging.Formatter(std_out_format)) 55 | logger.addHandler(consoleHandler) 56 | 57 | return logger 58 | 59 | def get_combined_results(left_results, right_results): 60 | results = {} 61 | count = float(left_results['count']) 62 | 63 | results['left_mr'] = round(left_results ['mr'] /count, 5) 64 | results['left_mrr'] = round(left_results ['mrr']/count, 5) 65 | results['right_mr'] = round(right_results['mr'] /count, 5) 66 | results['right_mrr'] = round(right_results['mrr']/count, 5) 67 | results['mr'] = round((left_results['mr'] + right_results['mr']) /(2*count), 5) 68 | results['mrr'] = round((left_results['mrr'] + right_results['mrr'])/(2*count), 5) 69 | 70 | for k in range(10): 71 | results['left_hits@{}'.format(k+1)] = round(left_results ['hits@{}'.format(k+1)]/count, 5) 72 | results['right_hits@{}'.format(k+1)] = round(right_results['hits@{}'.format(k+1)]/count, 5) 73 | results['hits@{}'.format(k+1)] = round((left_results['hits@{}'.format(k+1)] + right_results['hits@{}'.format(k+1)])/(2*count), 5) 74 | return results 75 | 76 | def get_param(shape): 77 | param = Parameter(torch.Tensor(*shape)); 78 | xavier_normal_(param.data) 79 | return param 80 | 81 | def com_mult(a, b): 82 | r1, i1 = a[..., 0], a[..., 1] 83 | r2, i2 = b[..., 0], b[..., 1] 84 | return torch.stack([r1 * r2 - i1 * i2, r1 * i2 + i1 * r2], dim = -1) 85 | 86 | def conj(a): 87 | a[..., 1] = -a[..., 1] 88 | return a 89 | 90 | def cconv(a, b): 91 | return torch.irfft(com_mult(torch.rfft(a, 1), torch.rfft(b, 1)), 1, signal_sizes=(a.shape[-1],)) 92 | 93 | def ccorr(a, b): 94 | return torch.irfft(com_mult(conj(torch.rfft(a, 1)), torch.rfft(b, 1)), 1, signal_sizes=(a.shape[-1],)) -------------------------------------------------------------------------------- /model/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/malllabiisc/CompGCN/3b06f5fec42526faa81afc158df9e64a8382982c/model/__init__.py -------------------------------------------------------------------------------- /model/compgcn_conv.py: -------------------------------------------------------------------------------- 1 | from helper import * 2 | from model.message_passing import MessagePassing 3 | 4 | class CompGCNConv(MessagePassing): 5 | def __init__(self, in_channels, out_channels, num_rels, act=lambda x:x, params=None): 6 | super(self.__class__, self).__init__() 7 | 8 | self.p = params 9 | self.in_channels = in_channels 10 | self.out_channels = out_channels 11 | self.num_rels = num_rels 12 | self.act = act 13 | self.device = None 14 | 15 | self.w_loop = get_param((in_channels, out_channels)) 16 | self.w_in = get_param((in_channels, out_channels)) 17 | self.w_out = get_param((in_channels, out_channels)) 18 | self.w_rel = get_param((in_channels, out_channels)) 19 | self.loop_rel = get_param((1, in_channels)); 20 | 21 | self.drop = torch.nn.Dropout(self.p.dropout) 22 | self.bn = torch.nn.BatchNorm1d(out_channels) 23 | 24 | if self.p.bias: self.register_parameter('bias', Parameter(torch.zeros(out_channels))) 25 | 26 | def forward(self, x, edge_index, edge_type, rel_embed): 27 | if self.device is None: 28 | self.device = edge_index.device 29 | 30 | rel_embed = torch.cat([rel_embed, self.loop_rel], dim=0) 31 | num_edges = edge_index.size(1) // 2 32 | num_ent = x.size(0) 33 | 34 | self.in_index, self.out_index = edge_index[:, :num_edges], edge_index[:, num_edges:] 35 | self.in_type, self.out_type = edge_type[:num_edges], edge_type [num_edges:] 36 | 37 | self.loop_index = torch.stack([torch.arange(num_ent), torch.arange(num_ent)]).to(self.device) 38 | self.loop_type = torch.full((num_ent,), rel_embed.size(0)-1, dtype=torch.long).to(self.device) 39 | 40 | self.in_norm = self.compute_norm(self.in_index, num_ent) 41 | self.out_norm = self.compute_norm(self.out_index, num_ent) 42 | 43 | in_res = self.propagate('add', self.in_index, x=x, edge_type=self.in_type, rel_embed=rel_embed, edge_norm=self.in_norm, mode='in') 44 | loop_res = self.propagate('add', self.loop_index, x=x, edge_type=self.loop_type, rel_embed=rel_embed, edge_norm=None, mode='loop') 45 | out_res = self.propagate('add', self.out_index, x=x, edge_type=self.out_type, rel_embed=rel_embed, edge_norm=self.out_norm, mode='out') 46 | out = self.drop(in_res)*(1/3) + self.drop(out_res)*(1/3) + loop_res*(1/3) 47 | 48 | if self.p.bias: out = out + self.bias 49 | out = self.bn(out) 50 | 51 | return self.act(out), torch.matmul(rel_embed, self.w_rel)[:-1] # Ignoring the self loop inserted 52 | 53 | def rel_transform(self, ent_embed, rel_embed): 54 | if self.p.opn == 'corr': trans_embed = ccorr(ent_embed, rel_embed) 55 | elif self.p.opn == 'sub': trans_embed = ent_embed - rel_embed 56 | elif self.p.opn == 'mult': trans_embed = ent_embed * rel_embed 57 | else: raise NotImplementedError 58 | 59 | return trans_embed 60 | 61 | def message(self, x_j, edge_type, rel_embed, edge_norm, mode): 62 | weight = getattr(self, 'w_{}'.format(mode)) 63 | rel_emb = torch.index_select(rel_embed, 0, edge_type) 64 | xj_rel = self.rel_transform(x_j, rel_emb) 65 | out = torch.mm(xj_rel, weight) 66 | 67 | return out if edge_norm is None else out * edge_norm.view(-1, 1) 68 | 69 | def update(self, aggr_out): 70 | return aggr_out 71 | 72 | def compute_norm(self, edge_index, num_ent): 73 | row, col = edge_index 74 | edge_weight = torch.ones_like(row).float() 75 | deg = scatter_add( edge_weight, row, dim=0, dim_size=num_ent) # Summing number of weights of the edges 76 | deg_inv = deg.pow(-0.5) # D^{-0.5} 77 | deg_inv[deg_inv == float('inf')] = 0 78 | norm = deg_inv[row] * edge_weight * deg_inv[col] # D^{-0.5} 79 | 80 | return norm 81 | 82 | def __repr__(self): 83 | return '{}({}, {}, num_rels={})'.format( 84 | self.__class__.__name__, self.in_channels, self.out_channels, self.num_rels) 85 | -------------------------------------------------------------------------------- /model/compgcn_conv_basis.py: -------------------------------------------------------------------------------- 1 | from helper import * 2 | from model.message_passing import MessagePassing 3 | 4 | class CompGCNConvBasis(MessagePassing): 5 | def __init__(self, in_channels, out_channels, num_rels, num_bases, act=lambda x:x, cache=True, params=None): 6 | super(self.__class__, self).__init__() 7 | 8 | self.p = params 9 | self.in_channels = in_channels 10 | self.out_channels = out_channels 11 | self.num_rels = num_rels 12 | self.num_bases = num_bases 13 | self.act = act 14 | self.device = None 15 | self.cache = cache # Should be False for graph classification tasks 16 | 17 | self.w_loop = get_param((in_channels, out_channels)); 18 | self.w_in = get_param((in_channels, out_channels)); 19 | self.w_out = get_param((in_channels, out_channels)); 20 | 21 | self.rel_basis = get_param((self.num_bases, in_channels)) 22 | self.rel_wt = get_param((self.num_rels*2, self.num_bases)) 23 | self.w_rel = get_param((in_channels, out_channels)) 24 | self.loop_rel = get_param((1, in_channels)); 25 | 26 | self.drop = torch.nn.Dropout(self.p.dropout) 27 | self.bn = torch.nn.BatchNorm1d(out_channels) 28 | 29 | self.in_norm, self.out_norm, 30 | self.in_index, self.out_index, 31 | self.in_type, self.out_type, 32 | self.loop_index, self.loop_type = None, None, None, None, None, None, None, None 33 | 34 | if self.p.bias: self.register_parameter('bias', Parameter(torch.zeros(out_channels))) 35 | 36 | def forward(self, x, edge_index, edge_type, edge_norm=None, rel_embed=None): 37 | if self.device is None: 38 | self.device = edge_index.device 39 | 40 | rel_embed = torch.mm(self.rel_wt, self.rel_basis) 41 | rel_embed = torch.cat([rel_embed, self.loop_rel], dim=0) 42 | 43 | num_edges = edge_index.size(1) // 2 44 | num_ent = x.size(0) 45 | 46 | if not self.cache or self.in_norm == None: 47 | self.in_index, self.out_index = edge_index[:, :num_edges], edge_index[:, num_edges:] 48 | self.in_type, self.out_type = edge_type[:num_edges], edge_type [num_edges:] 49 | 50 | self.loop_index = torch.stack([torch.arange(num_ent), torch.arange(num_ent)]).to(self.device) 51 | self.loop_type = torch.full((num_ent,), rel_embed.size(0)-1, dtype=torch.long).to(self.device) 52 | 53 | self.in_norm = self.compute_norm(self.in_index, num_ent) 54 | self.out_norm = self.compute_norm(self.out_index, num_ent) 55 | 56 | in_res = self.propagate('add', self.in_index, x=x, edge_type=self.in_type, rel_embed=rel_embed, edge_norm=self.in_norm, mode='in') 57 | loop_res = self.propagate('add', self.loop_index, x=x, edge_type=self.loop_type, rel_embed=rel_embed, edge_norm=None, mode='loop') 58 | out_res = self.propagate('add', self.out_index, x=x, edge_type=self.out_type, rel_embed=rel_embed, edge_norm=self.out_norm, mode='out') 59 | out = self.drop(in_res)*(1/3) + self.drop(out_res)*(1/3) + loop_res*(1/3) 60 | 61 | if self.p.bias: out = out + self.bias 62 | if self.b_norm: out = self.bn(out) 63 | 64 | return self.act(out), torch.matmul(rel_embed, self.w_rel)[:-1] 65 | 66 | def rel_transform(self, ent_embed, rel_embed): 67 | if self.p.opn == 'corr': trans_embed = ccorr(ent_embed, rel_embed) 68 | elif self.p.opn == 'sub': trans_embed = ent_embed - rel_embed 69 | elif self.p.opn == 'mult': trans_embed = ent_embed * rel_embed 70 | else: raise NotImplementedError 71 | 72 | return trans_embed 73 | 74 | def message(self, x_j, edge_type, rel_embed, edge_norm, mode): 75 | weight = getattr(self, 'w_{}'.format(mode)) 76 | rel_emb = torch.index_select(rel_embed, 0, edge_type) 77 | xj_rel = self.rel_transform(x_j, rel_emb) 78 | out = torch.mm(xj_rel, weight) 79 | 80 | return out if edge_norm is None else out * edge_norm.view(-1, 1) 81 | 82 | def update(self, aggr_out): 83 | return aggr_out 84 | 85 | def compute_norm(self, edge_index, num_ent): 86 | row, col = edge_index 87 | edge_weight = torch.ones_like(row).float() 88 | deg = scatter_add( edge_weight, row, dim=0, dim_size=num_ent) # Summing number of weights of the edges [Computing out-degree] [Should be equal to in-degree (undireted graph)] 89 | deg_inv = deg.pow(-0.5) # D^{-0.5} 90 | deg_inv[deg_inv == float('inf')] = 0 91 | norm = deg_inv[row] * edge_weight * deg_inv[col] # D^{-0.5} 92 | 93 | return norm 94 | 95 | def __repr__(self): 96 | return '{}({}, {}, num_rels={})'.format( 97 | self.__class__.__name__, self.in_channels, self.out_channels, self.num_rels) 98 | -------------------------------------------------------------------------------- /model/message_passing.py: -------------------------------------------------------------------------------- 1 | import inspect, torch 2 | from torch_scatter import scatter 3 | 4 | def scatter_(name, src, index, dim_size=None): 5 | r"""Aggregates all values from the :attr:`src` tensor at the indices 6 | specified in the :attr:`index` tensor along the first dimension. 7 | If multiple indices reference the same location, their contributions 8 | are aggregated according to :attr:`name` (either :obj:`"add"`, 9 | :obj:`"mean"` or :obj:`"max"`). 10 | 11 | Args: 12 | name (string): The aggregation to use (:obj:`"add"`, :obj:`"mean"`, 13 | :obj:`"max"`). 14 | src (Tensor): The source tensor. 15 | index (LongTensor): The indices of elements to scatter. 16 | dim_size (int, optional): Automatically create output tensor with size 17 | :attr:`dim_size` in the first dimension. If set to :attr:`None`, a 18 | minimal sized output tensor is returned. (default: :obj:`None`) 19 | 20 | :rtype: :class:`Tensor` 21 | """ 22 | if name == 'add': name = 'sum' 23 | assert name in ['sum', 'mean', 'max'] 24 | out = scatter(src, index, dim=0, out=None, dim_size=dim_size, reduce=name) 25 | return out[0] if isinstance(out, tuple) else out 26 | 27 | 28 | class MessagePassing(torch.nn.Module): 29 | r"""Base class for creating message passing layers 30 | 31 | .. math:: 32 | \mathbf{x}_i^{\prime} = \gamma_{\mathbf{\Theta}} \left( \mathbf{x}_i, 33 | \square_{j \in \mathcal{N}(i)} \, \phi_{\mathbf{\Theta}} 34 | \left(\mathbf{x}_i, \mathbf{x}_j,\mathbf{e}_{i,j}\right) \right), 35 | 36 | where :math:`\square` denotes a differentiable, permutation invariant 37 | function, *e.g.*, sum, mean or max, and :math:`\gamma_{\mathbf{\Theta}}` 38 | and :math:`\phi_{\mathbf{\Theta}}` denote differentiable functions such as 39 | MLPs. 40 | See `here `__ for the accompanying tutorial. 42 | 43 | """ 44 | 45 | def __init__(self, aggr='add'): 46 | super(MessagePassing, self).__init__() 47 | 48 | self.message_args = inspect.getargspec(self.message)[0][1:] # In the defined message function: get the list of arguments as list of string| For eg. in rgcn this will be ['x_j', 'edge_type', 'edge_norm'] (arguments of message function) 49 | self.update_args = inspect.getargspec(self.update)[0][2:] # Same for update function starting from 3rd argument | first=self, second=out 50 | 51 | def propagate(self, aggr, edge_index, **kwargs): 52 | r"""The initial call to start propagating messages. 53 | Takes in an aggregation scheme (:obj:`"add"`, :obj:`"mean"` or 54 | :obj:`"max"`), the edge indices, and all additional data which is 55 | needed to construct messages and to update node embeddings.""" 56 | 57 | assert aggr in ['add', 'mean', 'max'] 58 | kwargs['edge_index'] = edge_index 59 | 60 | 61 | size = None 62 | message_args = [] 63 | for arg in self.message_args: 64 | if arg[-2:] == '_i': # If arguments ends with _i then include indic 65 | tmp = kwargs[arg[:-2]] # Take the front part of the variable | Mostly it will be 'x', 66 | size = tmp.size(0) 67 | message_args.append(tmp[edge_index[0]]) # Lookup for head entities in edges 68 | elif arg[-2:] == '_j': 69 | tmp = kwargs[arg[:-2]] # tmp = kwargs['x'] 70 | size = tmp.size(0) 71 | message_args.append(tmp[edge_index[1]]) # Lookup for tail entities in edges 72 | else: 73 | message_args.append(kwargs[arg]) # Take things from kwargs 74 | 75 | update_args = [kwargs[arg] for arg in self.update_args] # Take update args from kwargs 76 | 77 | out = self.message(*message_args) 78 | out = scatter_(aggr, out, edge_index[0], dim_size=size) # Aggregated neighbors for each vertex 79 | out = self.update(out, *update_args) 80 | 81 | return out 82 | 83 | def message(self, x_j): # pragma: no cover 84 | r"""Constructs messages in analogy to :math:`\phi_{\mathbf{\Theta}}` 85 | for each edge in :math:`(i,j) \in \mathcal{E}`. 86 | Can take any argument which was initially passed to :meth:`propagate`. 87 | In addition, features can be lifted to the source node :math:`i` and 88 | target node :math:`j` by appending :obj:`_i` or :obj:`_j` to the 89 | variable name, *.e.g.* :obj:`x_i` and :obj:`x_j`.""" 90 | 91 | return x_j 92 | 93 | def update(self, aggr_out): # pragma: no cover 94 | r"""Updates node embeddings in analogy to 95 | :math:`\gamma_{\mathbf{\Theta}}` for each node 96 | :math:`i \in \mathcal{V}`. 97 | Takes in the output of aggregation as first argument and any argument 98 | which was initially passed to :meth:`propagate`.""" 99 | 100 | return aggr_out 101 | -------------------------------------------------------------------------------- /model/models.py: -------------------------------------------------------------------------------- 1 | from helper import * 2 | from model.compgcn_conv import CompGCNConv 3 | from model.compgcn_conv_basis import CompGCNConvBasis 4 | 5 | class BaseModel(torch.nn.Module): 6 | def __init__(self, params): 7 | super(BaseModel, self).__init__() 8 | 9 | self.p = params 10 | self.act = torch.tanh 11 | self.bceloss = torch.nn.BCELoss() 12 | 13 | def loss(self, pred, true_label): 14 | return self.bceloss(pred, true_label) 15 | 16 | class CompGCNBase(BaseModel): 17 | def __init__(self, edge_index, edge_type, num_rel, params=None): 18 | super(CompGCNBase, self).__init__(params) 19 | 20 | self.edge_index = edge_index 21 | self.edge_type = edge_type 22 | self.p.gcn_dim = self.p.embed_dim if self.p.gcn_layer == 1 else self.p.gcn_dim 23 | self.init_embed = get_param((self.p.num_ent, self.p.init_dim)) 24 | self.device = self.edge_index.device 25 | 26 | if self.p.num_bases > 0: 27 | self.init_rel = get_param((self.p.num_bases, self.p.init_dim)) 28 | else: 29 | if self.p.score_func == 'transe': self.init_rel = get_param((num_rel, self.p.init_dim)) 30 | else: self.init_rel = get_param((num_rel*2, self.p.init_dim)) 31 | 32 | if self.p.num_bases > 0: 33 | self.conv1 = CompGCNConvBasis(self.p.init_dim, self.p.gcn_dim, num_rel, self.p.num_bases, act=self.act, params=self.p) 34 | self.conv2 = CompGCNConv(self.p.gcn_dim, self.p.embed_dim, num_rel, act=self.act, params=self.p) if self.p.gcn_layer == 2 else None 35 | else: 36 | self.conv1 = CompGCNConv(self.p.init_dim, self.p.gcn_dim, num_rel, act=self.act, params=self.p) 37 | self.conv2 = CompGCNConv(self.p.gcn_dim, self.p.embed_dim, num_rel, act=self.act, params=self.p) if self.p.gcn_layer == 2 else None 38 | 39 | self.register_parameter('bias', Parameter(torch.zeros(self.p.num_ent))) 40 | 41 | def forward_base(self, sub, rel, drop1, drop2): 42 | 43 | r = self.init_rel if self.p.score_func != 'transe' else torch.cat([self.init_rel, -self.init_rel], dim=0) 44 | x, r = self.conv1(self.init_embed, self.edge_index, self.edge_type, rel_embed=r) 45 | x = drop1(x) 46 | x, r = self.conv2(x, self.edge_index, self.edge_type, rel_embed=r) if self.p.gcn_layer == 2 else (x, r) 47 | x = drop2(x) if self.p.gcn_layer == 2 else x 48 | 49 | sub_emb = torch.index_select(x, 0, sub) 50 | rel_emb = torch.index_select(r, 0, rel) 51 | 52 | return sub_emb, rel_emb, x 53 | 54 | 55 | class CompGCN_TransE(CompGCNBase): 56 | def __init__(self, edge_index, edge_type, params=None): 57 | super(self.__class__, self).__init__(edge_index, edge_type, params.num_rel, params) 58 | self.drop = torch.nn.Dropout(self.p.hid_drop) 59 | 60 | def forward(self, sub, rel): 61 | 62 | sub_emb, rel_emb, all_ent = self.forward_base(sub, rel, self.drop, self.drop) 63 | obj_emb = sub_emb + rel_emb 64 | 65 | x = self.p.gamma - torch.norm(obj_emb.unsqueeze(1) - all_ent, p=1, dim=2) 66 | score = torch.sigmoid(x) 67 | 68 | return score 69 | 70 | class CompGCN_DistMult(CompGCNBase): 71 | def __init__(self, edge_index, edge_type, params=None): 72 | super(self.__class__, self).__init__(edge_index, edge_type, params.num_rel, params) 73 | self.drop = torch.nn.Dropout(self.p.hid_drop) 74 | 75 | def forward(self, sub, rel): 76 | 77 | sub_emb, rel_emb, all_ent = self.forward_base(sub, rel, self.drop, self.drop) 78 | obj_emb = sub_emb * rel_emb 79 | 80 | x = torch.mm(obj_emb, all_ent.transpose(1, 0)) 81 | x += self.bias.expand_as(x) 82 | 83 | score = torch.sigmoid(x) 84 | return score 85 | 86 | class CompGCN_ConvE(CompGCNBase): 87 | def __init__(self, edge_index, edge_type, params=None): 88 | super(self.__class__, self).__init__(edge_index, edge_type, params.num_rel, params) 89 | 90 | self.bn0 = torch.nn.BatchNorm2d(1) 91 | self.bn1 = torch.nn.BatchNorm2d(self.p.num_filt) 92 | self.bn2 = torch.nn.BatchNorm1d(self.p.embed_dim) 93 | 94 | self.hidden_drop = torch.nn.Dropout(self.p.hid_drop) 95 | self.hidden_drop2 = torch.nn.Dropout(self.p.hid_drop2) 96 | self.feature_drop = torch.nn.Dropout(self.p.feat_drop) 97 | self.m_conv1 = torch.nn.Conv2d(1, out_channels=self.p.num_filt, kernel_size=(self.p.ker_sz, self.p.ker_sz), stride=1, padding=0, bias=self.p.bias) 98 | 99 | flat_sz_h = int(2*self.p.k_w) - self.p.ker_sz + 1 100 | flat_sz_w = self.p.k_h - self.p.ker_sz + 1 101 | self.flat_sz = flat_sz_h*flat_sz_w*self.p.num_filt 102 | self.fc = torch.nn.Linear(self.flat_sz, self.p.embed_dim) 103 | 104 | def concat(self, e1_embed, rel_embed): 105 | e1_embed = e1_embed. view(-1, 1, self.p.embed_dim) 106 | rel_embed = rel_embed.view(-1, 1, self.p.embed_dim) 107 | stack_inp = torch.cat([e1_embed, rel_embed], 1) 108 | stack_inp = torch.transpose(stack_inp, 2, 1).reshape((-1, 1, 2*self.p.k_w, self.p.k_h)) 109 | return stack_inp 110 | 111 | def forward(self, sub, rel): 112 | 113 | sub_emb, rel_emb, all_ent = self.forward_base(sub, rel, self.hidden_drop, self.feature_drop) 114 | stk_inp = self.concat(sub_emb, rel_emb) 115 | x = self.bn0(stk_inp) 116 | x = self.m_conv1(x) 117 | x = self.bn1(x) 118 | x = F.relu(x) 119 | x = self.feature_drop(x) 120 | x = x.view(-1, self.flat_sz) 121 | x = self.fc(x) 122 | x = self.hidden_drop2(x) 123 | x = self.bn2(x) 124 | x = F.relu(x) 125 | 126 | x = torch.mm(x, all_ent.transpose(1,0)) 127 | x += self.bias.expand_as(x) 128 | 129 | score = torch.sigmoid(x) 130 | return score 131 | -------------------------------------------------------------------------------- /overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/malllabiisc/CompGCN/3b06f5fec42526faa81afc158df9e64a8382982c/overview.png -------------------------------------------------------------------------------- /preprocess.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | mkdir log 3 | mkdir checkpoints 4 | mkdir data 5 | 6 | unzip data_compressed/FB15k-237.zip -d data 7 | unzip data_compressed/WN18RR.zip -d data -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch==1.4.0 2 | ordered_set==3.1 3 | numpy==1.16.2 4 | torch_scatter==2.0.4 5 | scikit_learn==0.21.1 6 | -------------------------------------------------------------------------------- /run.py: -------------------------------------------------------------------------------- 1 | from helper import * 2 | from data_loader import * 3 | 4 | # sys.path.append('./') 5 | from model.models import * 6 | 7 | class Runner(object): 8 | 9 | def load_data(self): 10 | """ 11 | Reading in raw triples and converts it into a standard format. 12 | 13 | Parameters 14 | ---------- 15 | self.p.dataset: Takes in the name of the dataset (FB15k-237) 16 | 17 | Returns 18 | ------- 19 | self.ent2id: Entity to unique identifier mapping 20 | self.id2rel: Inverse mapping of self.ent2id 21 | self.rel2id: Relation to unique identifier mapping 22 | self.num_ent: Number of entities in the Knowledge graph 23 | self.num_rel: Number of relations in the Knowledge graph 24 | self.embed_dim: Embedding dimension used 25 | self.data['train']: Stores the triples corresponding to training dataset 26 | self.data['valid']: Stores the triples corresponding to validation dataset 27 | self.data['test']: Stores the triples corresponding to test dataset 28 | self.data_iter: The dataloader for different data splits 29 | 30 | """ 31 | 32 | ent_set, rel_set = OrderedSet(), OrderedSet() 33 | for split in ['train', 'test', 'valid']: 34 | for line in open('./data/{}/{}.txt'.format(self.p.dataset, split)): 35 | sub, rel, obj = map(str.lower, line.strip().split('\t')) 36 | ent_set.add(sub) 37 | rel_set.add(rel) 38 | ent_set.add(obj) 39 | 40 | self.ent2id = {ent: idx for idx, ent in enumerate(ent_set)} 41 | self.rel2id = {rel: idx for idx, rel in enumerate(rel_set)} 42 | self.rel2id.update({rel+'_reverse': idx+len(self.rel2id) for idx, rel in enumerate(rel_set)}) 43 | 44 | self.id2ent = {idx: ent for ent, idx in self.ent2id.items()} 45 | self.id2rel = {idx: rel for rel, idx in self.rel2id.items()} 46 | 47 | self.p.num_ent = len(self.ent2id) 48 | self.p.num_rel = len(self.rel2id) // 2 49 | self.p.embed_dim = self.p.k_w * self.p.k_h if self.p.embed_dim is None else self.p.embed_dim 50 | 51 | self.data = ddict(list) 52 | sr2o = ddict(set) 53 | 54 | for split in ['train', 'test', 'valid']: 55 | for line in open('./data/{}/{}.txt'.format(self.p.dataset, split)): 56 | sub, rel, obj = map(str.lower, line.strip().split('\t')) 57 | sub, rel, obj = self.ent2id[sub], self.rel2id[rel], self.ent2id[obj] 58 | self.data[split].append((sub, rel, obj)) 59 | 60 | if split == 'train': 61 | sr2o[(sub, rel)].add(obj) 62 | sr2o[(obj, rel+self.p.num_rel)].add(sub) 63 | 64 | self.data = dict(self.data) 65 | 66 | self.sr2o = {k: list(v) for k, v in sr2o.items()} 67 | for split in ['test', 'valid']: 68 | for sub, rel, obj in self.data[split]: 69 | sr2o[(sub, rel)].add(obj) 70 | sr2o[(obj, rel+self.p.num_rel)].add(sub) 71 | 72 | self.sr2o_all = {k: list(v) for k, v in sr2o.items()} 73 | self.triples = ddict(list) 74 | 75 | for (sub, rel), obj in self.sr2o.items(): 76 | self.triples['train'].append({'triple':(sub, rel, -1), 'label': self.sr2o[(sub, rel)], 'sub_samp': 1}) 77 | 78 | for split in ['test', 'valid']: 79 | for sub, rel, obj in self.data[split]: 80 | rel_inv = rel + self.p.num_rel 81 | self.triples['{}_{}'.format(split, 'tail')].append({'triple': (sub, rel, obj), 'label': self.sr2o_all[(sub, rel)]}) 82 | self.triples['{}_{}'.format(split, 'head')].append({'triple': (obj, rel_inv, sub), 'label': self.sr2o_all[(obj, rel_inv)]}) 83 | 84 | self.triples = dict(self.triples) 85 | 86 | def get_data_loader(dataset_class, split, batch_size, shuffle=True): 87 | return DataLoader( 88 | dataset_class(self.triples[split], self.p), 89 | batch_size = batch_size, 90 | shuffle = shuffle, 91 | num_workers = max(0, self.p.num_workers), 92 | collate_fn = dataset_class.collate_fn 93 | ) 94 | 95 | self.data_iter = { 96 | 'train': get_data_loader(TrainDataset, 'train', self.p.batch_size), 97 | 'valid_head': get_data_loader(TestDataset, 'valid_head', self.p.batch_size), 98 | 'valid_tail': get_data_loader(TestDataset, 'valid_tail', self.p.batch_size), 99 | 'test_head': get_data_loader(TestDataset, 'test_head', self.p.batch_size), 100 | 'test_tail': get_data_loader(TestDataset, 'test_tail', self.p.batch_size), 101 | } 102 | 103 | self.edge_index, self.edge_type = self.construct_adj() 104 | 105 | def construct_adj(self): 106 | """ 107 | Constructor of the runner class 108 | 109 | Parameters 110 | ---------- 111 | 112 | Returns 113 | ------- 114 | Constructs the adjacency matrix for GCN 115 | 116 | """ 117 | edge_index, edge_type = [], [] 118 | 119 | for sub, rel, obj in self.data['train']: 120 | edge_index.append((sub, obj)) 121 | edge_type.append(rel) 122 | 123 | # Adding inverse edges 124 | for sub, rel, obj in self.data['train']: 125 | edge_index.append((obj, sub)) 126 | edge_type.append(rel + self.p.num_rel) 127 | 128 | edge_index = torch.LongTensor(edge_index).to(self.device).t() 129 | edge_type = torch.LongTensor(edge_type). to(self.device) 130 | 131 | return edge_index, edge_type 132 | 133 | def __init__(self, params): 134 | """ 135 | Constructor of the runner class 136 | 137 | Parameters 138 | ---------- 139 | params: List of hyper-parameters of the model 140 | 141 | Returns 142 | ------- 143 | Creates computational graph and optimizer 144 | 145 | """ 146 | self.p = params 147 | self.logger = get_logger(self.p.name, self.p.log_dir, self.p.config_dir) 148 | 149 | self.logger.info(vars(self.p)) 150 | pprint(vars(self.p)) 151 | 152 | if self.p.gpu != '-1' and torch.cuda.is_available(): 153 | self.device = torch.device('cuda') 154 | torch.cuda.set_rng_state(torch.cuda.get_rng_state()) 155 | torch.backends.cudnn.deterministic = True 156 | else: 157 | self.device = torch.device('cpu') 158 | 159 | self.load_data() 160 | self.model = self.add_model(self.p.model, self.p.score_func) 161 | self.optimizer = self.add_optimizer(self.model.parameters()) 162 | 163 | 164 | def add_model(self, model, score_func): 165 | """ 166 | Creates the computational graph 167 | 168 | Parameters 169 | ---------- 170 | model_name: Contains the model name to be created 171 | 172 | Returns 173 | ------- 174 | Creates the computational graph for model and initializes it 175 | 176 | """ 177 | model_name = '{}_{}'.format(model, score_func) 178 | 179 | if model_name.lower() == 'compgcn_transe': model = CompGCN_TransE(self.edge_index, self.edge_type, params=self.p) 180 | elif model_name.lower() == 'compgcn_distmult': model = CompGCN_DistMult(self.edge_index, self.edge_type, params=self.p) 181 | elif model_name.lower() == 'compgcn_conve': model = CompGCN_ConvE(self.edge_index, self.edge_type, params=self.p) 182 | else: raise NotImplementedError 183 | 184 | model.to(self.device) 185 | return model 186 | 187 | def add_optimizer(self, parameters): 188 | """ 189 | Creates an optimizer for training the parameters 190 | 191 | Parameters 192 | ---------- 193 | parameters: The parameters of the model 194 | 195 | Returns 196 | ------- 197 | Returns an optimizer for learning the parameters of the model 198 | 199 | """ 200 | return torch.optim.Adam(parameters, lr=self.p.lr, weight_decay=self.p.l2) 201 | 202 | def read_batch(self, batch, split): 203 | """ 204 | Function to read a batch of data and move the tensors in batch to CPU/GPU 205 | 206 | Parameters 207 | ---------- 208 | batch: the batch to process 209 | split: (string) If split == 'train', 'valid' or 'test' split 210 | 211 | 212 | Returns 213 | ------- 214 | Head, Relation, Tails, labels 215 | """ 216 | if split == 'train': 217 | triple, label = [ _.to(self.device) for _ in batch] 218 | return triple[:, 0], triple[:, 1], triple[:, 2], label 219 | else: 220 | triple, label = [ _.to(self.device) for _ in batch] 221 | return triple[:, 0], triple[:, 1], triple[:, 2], label 222 | 223 | def save_model(self, save_path): 224 | """ 225 | Function to save a model. It saves the model parameters, best validation scores, 226 | best epoch corresponding to best validation, state of the optimizer and all arguments for the run. 227 | 228 | Parameters 229 | ---------- 230 | save_path: path where the model is saved 231 | 232 | Returns 233 | ------- 234 | """ 235 | state = { 236 | 'state_dict' : self.model.state_dict(), 237 | 'best_val' : self.best_val, 238 | 'best_epoch' : self.best_epoch, 239 | 'optimizer' : self.optimizer.state_dict(), 240 | 'args' : vars(self.p) 241 | } 242 | torch.save(state, save_path) 243 | 244 | def load_model(self, load_path): 245 | """ 246 | Function to load a saved model 247 | 248 | Parameters 249 | ---------- 250 | load_path: path to the saved model 251 | 252 | Returns 253 | ------- 254 | """ 255 | state = torch.load(load_path) 256 | state_dict = state['state_dict'] 257 | self.best_val = state['best_val'] 258 | self.best_val_mrr = self.best_val['mrr'] 259 | 260 | self.model.load_state_dict(state_dict) 261 | self.optimizer.load_state_dict(state['optimizer']) 262 | 263 | def evaluate(self, split, epoch): 264 | """ 265 | Function to evaluate the model on validation or test set 266 | 267 | Parameters 268 | ---------- 269 | split: (string) If split == 'valid' then evaluate on the validation set, else the test set 270 | epoch: (int) Current epoch count 271 | 272 | Returns 273 | ------- 274 | resutls: The evaluation results containing the following: 275 | results['mr']: Average of ranks_left and ranks_right 276 | results['mrr']: Mean Reciprocal Rank 277 | results['hits@k']: Probability of getting the correct preodiction in top-k ranks based on predicted score 278 | 279 | """ 280 | left_results = self.predict(split=split, mode='tail_batch') 281 | right_results = self.predict(split=split, mode='head_batch') 282 | results = get_combined_results(left_results, right_results) 283 | self.logger.info('[Epoch {} {}]: MRR: Tail : {:.5}, Head : {:.5}, Avg : {:.5}'.format(epoch, split, results['left_mrr'], results['right_mrr'], results['mrr'])) 284 | return results 285 | 286 | def predict(self, split='valid', mode='tail_batch'): 287 | """ 288 | Function to run model evaluation for a given mode 289 | 290 | Parameters 291 | ---------- 292 | split: (string) If split == 'valid' then evaluate on the validation set, else the test set 293 | mode: (string): Can be 'head_batch' or 'tail_batch' 294 | 295 | Returns 296 | ------- 297 | resutls: The evaluation results containing the following: 298 | results['mr']: Average of ranks_left and ranks_right 299 | results['mrr']: Mean Reciprocal Rank 300 | results['hits@k']: Probability of getting the correct preodiction in top-k ranks based on predicted score 301 | 302 | """ 303 | self.model.eval() 304 | 305 | with torch.no_grad(): 306 | results = {} 307 | train_iter = iter(self.data_iter['{}_{}'.format(split, mode.split('_')[0])]) 308 | 309 | for step, batch in enumerate(train_iter): 310 | sub, rel, obj, label = self.read_batch(batch, split) 311 | pred = self.model.forward(sub, rel) 312 | b_range = torch.arange(pred.size()[0], device=self.device) 313 | target_pred = pred[b_range, obj] 314 | pred = torch.where(label.byte(), -torch.ones_like(pred) * 10000000, pred) 315 | pred[b_range, obj] = target_pred 316 | ranks = 1 + torch.argsort(torch.argsort(pred, dim=1, descending=True), dim=1, descending=False)[b_range, obj] 317 | 318 | ranks = ranks.float() 319 | results['count'] = torch.numel(ranks) + results.get('count', 0.0) 320 | results['mr'] = torch.sum(ranks).item() + results.get('mr', 0.0) 321 | results['mrr'] = torch.sum(1.0/ranks).item() + results.get('mrr', 0.0) 322 | for k in range(10): 323 | results['hits@{}'.format(k+1)] = torch.numel(ranks[ranks <= (k+1)]) + results.get('hits@{}'.format(k+1), 0.0) 324 | 325 | if step % 100 == 0: 326 | self.logger.info('[{}, {} Step {}]\t{}'.format(split.title(), mode.title(), step, self.p.name)) 327 | 328 | return results 329 | 330 | 331 | def run_epoch(self, epoch, val_mrr = 0): 332 | """ 333 | Function to run one epoch of training 334 | 335 | Parameters 336 | ---------- 337 | epoch: current epoch count 338 | 339 | Returns 340 | ------- 341 | loss: The loss value after the completion of one epoch 342 | """ 343 | self.model.train() 344 | losses = [] 345 | train_iter = iter(self.data_iter['train']) 346 | 347 | for step, batch in enumerate(train_iter): 348 | self.optimizer.zero_grad() 349 | sub, rel, obj, label = self.read_batch(batch, 'train') 350 | 351 | pred = self.model.forward(sub, rel) 352 | loss = self.model.loss(pred, label) 353 | 354 | loss.backward() 355 | self.optimizer.step() 356 | losses.append(loss.item()) 357 | 358 | if step % 100 == 0: 359 | self.logger.info('[E:{}| {}]: Train Loss:{:.5}, Val MRR:{:.5}\t{}'.format(epoch, step, np.mean(losses), self.best_val_mrr, self.p.name)) 360 | 361 | loss = np.mean(losses) 362 | self.logger.info('[Epoch:{}]: Training Loss:{:.4}\n'.format(epoch, loss)) 363 | return loss 364 | 365 | 366 | def fit(self): 367 | """ 368 | Function to run training and evaluation of model 369 | 370 | Parameters 371 | ---------- 372 | 373 | Returns 374 | ------- 375 | """ 376 | self.best_val_mrr, self.best_val, self.best_epoch, val_mrr = 0., {}, 0, 0. 377 | save_path = os.path.join('./checkpoints', self.p.name) 378 | 379 | if self.p.restore: 380 | self.load_model(save_path) 381 | self.logger.info('Successfully Loaded previous model') 382 | 383 | kill_cnt = 0 384 | for epoch in range(self.p.max_epochs): 385 | train_loss = self.run_epoch(epoch, val_mrr) 386 | val_results = self.evaluate('valid', epoch) 387 | 388 | if val_results['mrr'] > self.best_val_mrr: 389 | self.best_val = val_results 390 | self.best_val_mrr = val_results['mrr'] 391 | self.best_epoch = epoch 392 | self.save_model(save_path) 393 | kill_cnt = 0 394 | else: 395 | kill_cnt += 1 396 | if kill_cnt % 10 == 0 and self.p.gamma > 5: 397 | self.p.gamma -= 5 398 | self.logger.info('Gamma decay on saturation, updated value of gamma: {}'.format(self.p.gamma)) 399 | if kill_cnt > 25: 400 | self.logger.info("Early Stopping!!") 401 | break 402 | 403 | self.logger.info('[Epoch {}]: Training Loss: {:.5}, Valid MRR: {:.5}\n\n'.format(epoch, train_loss, self.best_val_mrr)) 404 | 405 | self.logger.info('Loading best model, Evaluating on Test data') 406 | self.load_model(save_path) 407 | test_results = self.evaluate('test', epoch) 408 | 409 | if __name__ == '__main__': 410 | parser = argparse.ArgumentParser(description='Parser For Arguments', formatter_class=argparse.ArgumentDefaultsHelpFormatter) 411 | 412 | parser.add_argument('-name', default='testrun', help='Set run name for saving/restoring models') 413 | parser.add_argument('-data', dest='dataset', default='FB15k-237', help='Dataset to use, default: FB15k-237') 414 | parser.add_argument('-model', dest='model', default='compgcn', help='Model Name') 415 | parser.add_argument('-score_func', dest='score_func', default='conve', help='Score Function for Link prediction') 416 | parser.add_argument('-opn', dest='opn', default='corr', help='Composition Operation to be used in CompGCN') 417 | 418 | parser.add_argument('-batch', dest='batch_size', default=128, type=int, help='Batch size') 419 | parser.add_argument('-gamma', type=float, default=40.0, help='Margin') 420 | parser.add_argument('-gpu', type=str, default='0', help='Set GPU Ids : Eg: For CPU = -1, For Single GPU = 0') 421 | parser.add_argument('-epoch', dest='max_epochs', type=int, default=500, help='Number of epochs') 422 | parser.add_argument('-l2', type=float, default=0.0, help='L2 Regularization for Optimizer') 423 | parser.add_argument('-lr', type=float, default=0.001, help='Starting Learning Rate') 424 | parser.add_argument('-lbl_smooth', dest='lbl_smooth', type=float, default=0.1, help='Label Smoothing') 425 | parser.add_argument('-num_workers', type=int, default=10, help='Number of processes to construct batches') 426 | parser.add_argument('-seed', dest='seed', default=41504, type=int, help='Seed for randomization') 427 | 428 | parser.add_argument('-restore', dest='restore', action='store_true', help='Restore from the previously saved model') 429 | parser.add_argument('-bias', dest='bias', action='store_true', help='Whether to use bias in the model') 430 | 431 | parser.add_argument('-num_bases', dest='num_bases', default=-1, type=int, help='Number of basis relation vectors to use') 432 | parser.add_argument('-init_dim', dest='init_dim', default=100, type=int, help='Initial dimension size for entities and relations') 433 | parser.add_argument('-gcn_dim', dest='gcn_dim', default=200, type=int, help='Number of hidden units in GCN') 434 | parser.add_argument('-embed_dim', dest='embed_dim', default=None, type=int, help='Embedding dimension to give as input to score function') 435 | parser.add_argument('-gcn_layer', dest='gcn_layer', default=1, type=int, help='Number of GCN Layers to use') 436 | parser.add_argument('-gcn_drop', dest='dropout', default=0.1, type=float, help='Dropout to use in GCN Layer') 437 | parser.add_argument('-hid_drop', dest='hid_drop', default=0.3, type=float, help='Dropout after GCN') 438 | 439 | # ConvE specific hyperparameters 440 | parser.add_argument('-hid_drop2', dest='hid_drop2', default=0.3, type=float, help='ConvE: Hidden dropout') 441 | parser.add_argument('-feat_drop', dest='feat_drop', default=0.3, type=float, help='ConvE: Feature Dropout') 442 | parser.add_argument('-k_w', dest='k_w', default=10, type=int, help='ConvE: k_w') 443 | parser.add_argument('-k_h', dest='k_h', default=20, type=int, help='ConvE: k_h') 444 | parser.add_argument('-num_filt', dest='num_filt', default=200, type=int, help='ConvE: Number of filters in convolution') 445 | parser.add_argument('-ker_sz', dest='ker_sz', default=7, type=int, help='ConvE: Kernel size to use') 446 | 447 | parser.add_argument('-logdir', dest='log_dir', default='./log/', help='Log directory') 448 | parser.add_argument('-config', dest='config_dir', default='./config/', help='Config directory') 449 | args = parser.parse_args() 450 | 451 | if not args.restore: args.name = args.name + '_' + time.strftime('%d_%m_%Y') + '_' + time.strftime('%H:%M:%S') 452 | 453 | set_gpu(args.gpu) 454 | np.random.seed(args.seed) 455 | torch.manual_seed(args.seed) 456 | 457 | model = Runner(args) 458 | model.fit() 459 | --------------------------------------------------------------------------------