Ability to Disable One Side's Negative Sampling

Summary: As titled.

Reviewed By: lerks

Differential Revision: D17337389

fbshipit-source-id: fdf9884b51d08dd4bb9ffde45d0e1926d0d7c390

Author: Xin Jin

torchbiggraph/config.py

@@ -259,6 +259,14 @@ class ConfigSchema(Schema):
         metadata={'help': "The number of negatives uniformly sampled from the "
                           "currently active partition, per positive edge."},
     )
+    disable_lhs_negs : bool = attr.ib(
+        default=False,
+        metadata={'help': "Disable negative sampling on the left-hand side."},
+    )
+    disable_rhs_negs : bool = attr.ib(
+        default=False,
+        metadata={'help': "Disable negative sampling on the right-hand side."},
+    )
     lr: float = attr.ib(
         default=1e-2,
         validator=non_negative,
@@ -373,6 +381,9 @@ def __attrs_post_init__(self):
         if self.loss_fn == "logistic" and self.comparator == "cos":
             logger.warning("You have logistic loss and cosine distance. Are you sure?")
 
+        if self.disable_lhs_negs and self.disable_rhs_negs:
+            raise ValueError("Cannot disable negative sampling on both sides.")
+
 
 # TODO make this a non-inplace operation
 def override_config_dict(config_dict: Any, overrides: List[str]) -> Any:

torchbiggraph/eval.py

@@ -65,21 +65,28 @@ def eval(
     ) -> Stats:
         batch_size = len(batch_edges)
 
-        lhs_rank = (scores.lhs_neg >= scores.lhs_pos.unsqueeze(1)).sum(1) + 1
-        rhs_rank = (scores.rhs_neg >= scores.rhs_pos.unsqueeze(1)).sum(1) + 1
-
-        lhs_auc = compute_randomized_auc(scores.lhs_pos, scores.lhs_neg, batch_size)
-        rhs_auc = compute_randomized_auc(scores.rhs_pos, scores.rhs_neg, batch_size)
+        ranks = []
+        aucs = []
+        if scores.lhs_neg.nelement() > 0:
+            lhs_rank = (scores.lhs_neg >= scores.lhs_pos.unsqueeze(1)).sum(1) + 1
+            lhs_auc = compute_randomized_auc(scores.lhs_pos, scores.lhs_neg, batch_size)
+            ranks.append(lhs_rank)
+            aucs.append(lhs_auc)
+
+        if scores.rhs_neg.nelement() > 0:
+            rhs_rank = (scores.rhs_neg >= scores.rhs_pos.unsqueeze(1)).sum(1) + 1
+            rhs_auc = compute_randomized_auc(scores.rhs_pos, scores.rhs_neg, batch_size)
+            ranks.append(rhs_rank)
+            aucs.append(rhs_auc)
 
         return Stats(
-            pos_rank=average_of_sums(lhs_rank, rhs_rank),
-            mrr=average_of_sums(lhs_rank.float().reciprocal(),
-                                rhs_rank.float().reciprocal()),
-            r1=average_of_sums(lhs_rank.le(1), rhs_rank.le(1)),
-            r10=average_of_sums(lhs_rank.le(10), rhs_rank.le(10)),
-            r50=average_of_sums(lhs_rank.le(50), rhs_rank.le(50)),
+            pos_rank=average_of_sums(*ranks),
+            mrr=average_of_sums(*(rank.float().reciprocal() for rank in ranks)),
+            r1=average_of_sums(*(rank.le(1) for rank in ranks)),
+            r10=average_of_sums(*(rank.le(10) for rank in ranks)),
+            r50=average_of_sums(*(rank.le(50) for rank in ranks)),
             # At the end the AUC will be averaged over count.
-            auc=batch_size * (lhs_auc + rhs_auc) / 2,
+            auc=batch_size * sum(aucs) / len(aucs),
             count=batch_size)
 
 

torchbiggraph/model.py

@@ -772,6 +772,8 @@ def __init__(
         entities: Dict[str, EntitySchema],
         num_batch_negs: int,
         num_uniform_negs: int,
+        disable_lhs_negs: bool,
+        disable_rhs_negs: bool,
         lhs_operators: Sequence[Optional[Union[AbstractOperator, AbstractDynamicOperator]]],
         rhs_operators: Sequence[Optional[Union[AbstractOperator, AbstractDynamicOperator]]],
         comparator: AbstractComparator,
@@ -795,6 +797,9 @@ def __init__(
         self.num_batch_negs: int = num_batch_negs
         self.num_uniform_negs: int = num_uniform_negs
 
+        self.disable_lhs_negs = disable_lhs_negs
+        self.disable_rhs_negs = disable_rhs_negs
+
         self.comparator = comparator
 
         self.lhs_embs: nn.ParameterDict = nn.ModuleDict()
@@ -1000,6 +1005,14 @@ def forward(
             chunk_size = self.num_batch_negs
             negative_sampling_method = Negatives.BATCH_UNIFORM
 
+        lhs_negative_sampling_method = negative_sampling_method
+        rhs_negative_sampling_method = negative_sampling_method
+
+        if self.disable_lhs_negs:
+            lhs_negative_sampling_method = Negatives.NONE
+        if self.disable_rhs_negs:
+            rhs_negative_sampling_method = Negatives.NONE
+
         if self.num_dynamic_rels == 0:
             # In this case the operator is only applied to the RHS. This means
             # that an edge (u, r, v) is scored with c(u, f_r(v)), whereas the
@@ -1012,7 +1025,8 @@ def forward(
                 raise RuntimeError("In non-dynamic relation mode there should "
                                    "be only a right-hand side operator")
 
-            # Apply operator to right-hand side, sample negatives on both sides.
+            # Apply operator to right-hand side, sample negatives on both sides unless
+            # one side is disabled.
             pos_scores, lhs_neg_scores, rhs_neg_scores = self.forward_direction_agnostic(
                 edges.lhs,
                 edges.rhs,
@@ -1026,8 +1040,8 @@ def forward(
                 lhs_pos,
                 rhs_pos,
                 chunk_size,
-                negative_sampling_method,
-                negative_sampling_method,
+                lhs_negative_sampling_method,
+                rhs_negative_sampling_method,
             )
             lhs_pos_scores = rhs_pos_scores = pos_scores
 
@@ -1061,7 +1075,7 @@ def forward(
                 lhs_pos,
                 rhs_pos,
                 chunk_size,
-                negative_sampling_method,
+                lhs_negative_sampling_method,
                 Negatives.NONE,
             )
             # "Reverse" edges: apply operator to lhs, sample negatives on rhs.
@@ -1078,7 +1092,7 @@ def forward(
                 rhs_pos,
                 lhs_pos,
                 chunk_size,
-                negative_sampling_method,
+                rhs_negative_sampling_method,
                 Negatives.NONE,
             )
 
@@ -1187,6 +1201,8 @@ def make_model(config: ConfigSchema) -> MultiRelationEmbedder:
         config.entities,
         num_uniform_negs=config.num_uniform_negs,
         num_batch_negs=config.num_batch_negs,
+        disable_lhs_negs=config.disable_lhs_negs,
+        disable_rhs_negs=config.disable_rhs_negs,
         lhs_operators=lhs_operators,
         rhs_operators=rhs_operators,
         comparator=comparator,