modeling script

.gitignore

@@ -0,0 +1,15 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+output/
+.vscode/
+.DS_Store
+datasets/
+ckpts/
+*.csv
+*.txt

modeling.py

@@ -1,9 +1,10 @@
 import os
 import sys
-from transformers import PretrainedConfig, PreTrainedModel, AutoModel, AutoConfig
+from transformers import PretrainedConfig, PreTrainedModel
 sys.path.append(os.path.dirname(os.path.dirname(__file__)))
 from ultra.models import Ultra
-import torch
+from ultra.datasets import WN18RR, CoDExSmall, FB15k237, FB15k237Inductive
+from ultra.eval import test
 
 
 class UltraConfig(PretrainedConfig):
@@ -58,8 +59,8 @@ class UltraLinkPrediction(PreTrainedModel):
 
 if __name__ == "__main__":
 
-    ultra_config = UltraConfig()
-    ultra_model = UltraLinkPrediction(ultra_config)
-    ultra_model = UltraLinkPrediction.from_pretrained("mgalkin/ultra_3g")
-    print(ultra_model)
-    print("Done")
+    model = UltraLinkPrediction.from_pretrained("mgalkin/ultra_3g")
+    dataset = CoDExSmall(root="./datasets/")
+    test(model, mode="test", dataset=dataset, gpus=None)
+    # mrr:      0.472035
+    # hits@10:  0.66849

ultra/eval.py

@@ -0,0 +1,153 @@
+import math
+
+import torch
+from torch import distributed as dist
+from torch.utils import data as torch_data
+from torch_geometric.data import Data
+
+from ultra import tasks, util
+
+
+TRANSDUCTIVE = ("WordNet18RR", "RelLinkPredDataset", "CoDExSmall", "CoDExMedium", "CoDExLarge",
+                "YAGO310", "NELL995", "ConceptNet100k", "DBpedia100k", "Hetionet", "AristoV4",
+                "WDsinger", "NELL23k", "FB15k237_10", "FB15k237_20", "FB15k237_50")
+
+
+def get_filtered_data(dataset, mode):
+    train_data, valid_data, test_data = dataset[0], dataset[1], dataset[2]
+    ds_name = dataset.__class__.__name__
+
+    if ds_name in TRANSDUCTIVE:
+        filtered_data = Data(edge_index=dataset._data.target_edge_index, edge_type=dataset._data.target_edge_type, num_nodes=dataset[0].num_nodes)
+    else:
+        if "ILPC" in ds_name or "Ingram" in ds_name:
+            full_inference_edges = torch.cat([valid_data.edge_index, valid_data.target_edge_index, test_data.target_edge_index], dim=1)
+            full_inference_etypes = torch.cat([valid_data.edge_type, valid_data.target_edge_type, test_data.target_edge_type])
+            filtered_data = Data(edge_index=full_inference_edges, edge_type=full_inference_etypes, num_nodes=test_data.num_nodes)
+        else:
+            # test filtering graph: inference edges + test edges
+            full_inference_edges = torch.cat([test_data.edge_index, test_data.target_edge_index], dim=1)
+            full_inference_etypes = torch.cat([test_data.edge_type, test_data.target_edge_type])
+            if mode == "test":
+                filtered_data = Data(edge_index=full_inference_edges, edge_type=full_inference_etypes, num_nodes=test_data.num_nodes)
+            else:
+                # validation filtering graph: train edges + validation edges
+                filtered_data = Data(
+                    edge_index=torch.cat([train_data.edge_index, valid_data.target_edge_index], dim=1),
+                    edge_type=torch.cat([train_data.edge_type, valid_data.target_edge_type])
+                )
+
+    return filtered_data
+
+
+@torch.no_grad()
+def test(model, mode, dataset,  batch_size=32, eval_metrics=["mrr", "hits@10"], gpus=None, return_metrics=False):
+    logger = util.get_root_logger()
+    test_data = dataset[1] if mode == "valid" else dataset[2]
+    filtered_data = get_filtered_data(dataset, mode)
+
+    device = util.get_devices(gpus)
+    world_size = util.get_world_size()
+    rank = util.get_rank()
+
+    test_triplets = torch.cat([test_data.target_edge_index, test_data.target_edge_type.unsqueeze(0)]).t()
+    sampler = torch_data.DistributedSampler(test_triplets, world_size, rank)
+    test_loader = torch_data.DataLoader(test_triplets, batch_size, sampler=sampler)
+
+    model.eval()
+    rankings = []
+    num_negatives = []
+    tail_rankings, num_tail_negs = [], []  # for explicit tail-only evaluation needed for 5 datasets
+    for batch in test_loader:
+        t_batch, h_batch = tasks.all_negative(test_data, batch)
+        t_pred = model(test_data, t_batch)
+        h_pred = model(test_data, h_batch)
+
+        if filtered_data is None:
+            t_mask, h_mask = tasks.strict_negative_mask(test_data, batch)
+        else:
+            t_mask, h_mask = tasks.strict_negative_mask(filtered_data, batch)
+        pos_h_index, pos_t_index, pos_r_index = batch.t()
+        t_ranking = tasks.compute_ranking(t_pred, pos_t_index, t_mask)
+        h_ranking = tasks.compute_ranking(h_pred, pos_h_index, h_mask)
+        num_t_negative = t_mask.sum(dim=-1)
+        num_h_negative = h_mask.sum(dim=-1)
+
+        rankings += [t_ranking, h_ranking]
+        num_negatives += [num_t_negative, num_h_negative]
+
+        tail_rankings += [t_ranking]
+        num_tail_negs += [num_t_negative]
+
+    ranking = torch.cat(rankings)
+    num_negative = torch.cat(num_negatives)
+    all_size = torch.zeros(world_size, dtype=torch.long, device=device)
+    all_size[rank] = len(ranking)
+
+    # ugly repetitive code for tail-only ranks processing
+    tail_ranking = torch.cat(tail_rankings)
+    num_tail_neg = torch.cat(num_tail_negs)
+    all_size_t = torch.zeros(world_size, dtype=torch.long, device=device)
+    all_size_t[rank] = len(tail_ranking)
+    if world_size > 1:
+        dist.all_reduce(all_size, op=dist.ReduceOp.SUM)
+        dist.all_reduce(all_size_t, op=dist.ReduceOp.SUM)
+
+    # obtaining all ranks 
+    cum_size = all_size.cumsum(0)
+    all_ranking = torch.zeros(all_size.sum(), dtype=torch.long, device=device)
+    all_ranking[cum_size[rank] - all_size[rank]: cum_size[rank]] = ranking
+    all_num_negative = torch.zeros(all_size.sum(), dtype=torch.long, device=device)
+    all_num_negative[cum_size[rank] - all_size[rank]: cum_size[rank]] = num_negative
+
+    # the same for tails-only ranks
+    cum_size_t = all_size_t.cumsum(0)
+    all_ranking_t = torch.zeros(all_size_t.sum(), dtype=torch.long, device=device)
+    all_ranking_t[cum_size_t[rank] - all_size_t[rank]: cum_size_t[rank]] = tail_ranking
+    all_num_negative_t = torch.zeros(all_size_t.sum(), dtype=torch.long, device=device)
+    all_num_negative_t[cum_size_t[rank] - all_size_t[rank]: cum_size_t[rank]] = num_tail_neg
+    if world_size > 1:
+        dist.all_reduce(all_ranking, op=dist.ReduceOp.SUM)
+        dist.all_reduce(all_num_negative, op=dist.ReduceOp.SUM)
+        dist.all_reduce(all_ranking_t, op=dist.ReduceOp.SUM)
+        dist.all_reduce(all_num_negative_t, op=dist.ReduceOp.SUM)
+
+    metrics = {}
+    if rank == 0:
+        for metric in eval_metrics:
+            if "-tail" in metric:
+                _metric_name, direction = metric.split("-")
+                if direction != "tail":
+                    raise ValueError("Only tail metric is supported in this mode")
+                _ranking = all_ranking_t
+                _num_neg = all_num_negative_t
+            else:
+                _ranking = all_ranking 
+                _num_neg = all_num_negative 
+                _metric_name = metric
+            
+            if _metric_name == "mr":
+                score = _ranking.float().mean()
+            elif _metric_name == "mrr":
+                score = (1 / _ranking.float()).mean()
+            elif _metric_name.startswith("hits@"):
+                values = _metric_name[5:].split("_")
+                threshold = int(values[0])
+                if len(values) > 1:
+                    num_sample = int(values[1])
+                    # unbiased estimation
+                    fp_rate = (_ranking - 1).float() / _num_neg
+                    score = 0
+                    for i in range(threshold):
+                        # choose i false positive from num_sample - 1 negatives
+                        num_comb = math.factorial(num_sample - 1) / \
+                                   math.factorial(i) / math.factorial(num_sample - i - 1)
+                        score += num_comb * (fp_rate ** i) * ((1 - fp_rate) ** (num_sample - i - 1))
+                    score = score.mean()
+                else:
+                    score = (_ranking <= threshold).float().mean()
+            logger.warning("%s: %g" % (metric, score))
+            metrics[metric] = score
+    mrr = (1 / all_ranking.float()).mean()
+
+    return mrr if not return_metrics else metrics

ultra/layers.py

@@ -177,7 +177,7 @@ class GeneralizedRelationalConv(MessagePassing):
         # fused computation of message and aggregate steps with the custom rspmm cuda kernel
         # speed up computation by several times
         # reduce memory complexity from O(|E|d) to O(|V|d), so we can apply it to larger graphs
-        from .rspmm import generalized_rspmm
+        from ultra.rspmm.rspmm import generalized_rspmm
 
         batch_size, num_node = input.shape[:2]
         input = input.transpose(0, 1).flatten(1)

ultra/util.py

@@ -109,6 +109,13 @@ def get_device(cfg):
         device = torch.device("cpu")
     return device
 
+def get_devices(gpus):
+    if gpus is not None:
+        device = torch.device(gpus[get_rank()])
+    else:
+        device = torch.device("cpu")
+    return device
+
 
 def create_working_directory(cfg):
     file_name = "working_dir.tmp"