def forward(self, embeddings: FloatTensorType) -> FloatTensorType:
match_shape(embeddings, ..., self.dim)
# We add a dimension so that matmul performs a matrix-vector product.
return (torch.matmul(self.linear_transformation,
embeddings.unsqueeze(-1)).squeeze(-1)
+ self.translation)
def prepare_negatives(
self,
pos_input: EntityList,
pos_embs: FloatTensorType,
module: AbstractEmbedding,
type_: Negatives,
num_uniform_neg: int,
rel: Union[int, LongTensorType],
entity_type: str,
operator: Union[None, AbstractOperator, AbstractDynamicOperator],
) -> Tuple[FloatTensorType, Mask]:
"""Given some chunked positives, set up chunks of negatives.
This function operates on one side (left-hand or right-hand) at a time.
It takes all the information about the positives on that side (the
original input value, the corresponding embeddings, and the module used
to convert one to the other). It then produces negatives for that side
according to the specified mode. The positive embeddings come in in
chunked form and the negatives are produced within each of these chunks.
The negatives can be either none, or the positives from the same chunk,
or all the possible entities. In the second mode, uniformly-sampled
entities can also be appended to the per-chunk negatives (each chunk
having a different sample). This function returns both the chunked
embeddings of the negatives and a mask of the same size as the chunked
positives-vs-negatives scores, whose non-zero elements correspond to the
scores that must be ignored.
"""
num_pos = len(pos_input)
num_chunks, chunk_size, dim = match_shape(pos_embs, -1, -1, -1)
last_chunk_size = num_pos - (num_chunks - 1) * chunk_size
ignore_mask: Mask = []
if type_ is Negatives.NONE:
neg_embs = pos_embs.new_empty((num_chunks, 0, dim))
elif type_ is Negatives.UNIFORM:
uniform_neg_embs = module.sample_entities(num_chunks,
num_uniform_neg)
neg_embs = self.adjust_embs(uniform_neg_embs, rel, entity_type,
operator)
elif type_ is Negatives.BATCH_UNIFORM:
neg_embs = pos_embs
if num_uniform_neg > 0:
try:
uniform_neg_embs = module.sample_entities(
num_chunks, num_uniform_neg)
except NotImplementedError:
pass # only use pos_embs i.e. batch negatives
else:
neg_embs = torch.cat(
[
pos_embs,
self.adjust_embs(uniform_neg_embs, rel,
entity_type, operator),
],
dim=1,
)
chunk_indices = torch.arange(chunk_size,
dtype=torch.long,
device=pos_embs.device)
last_chunk_indices = chunk_indices[:last_chunk_size]
# Ignore scores between positive pairs.
ignore_mask.append(
(slice(num_chunks - 1), chunk_indices, chunk_indices))
ignore_mask.append((-1, last_chunk_indices, last_chunk_indices))
# In the last chunk, ignore the scores between the positives that
# are not padding (i.e., the first last_chunk_size ones) and the
# negatives that are padding (i.e., all of them except the first
# last_chunk_size ones). Stop the last slice at chunk_size so that
# it doesn't also affect the uniformly-sampled negatives.
ignore_mask.append(
(-1, slice(last_chunk_size), slice(last_chunk_size,
chunk_size)))
elif type_ is Negatives.ALL:
pos_input_ten = pos_input.to_tensor()
neg_embs = self.adjust_embs(
module.get_all_entities().expand(num_chunks, -1, dim),
rel,
entity_type,
operator,
)
if num_uniform_neg > 0:
logger.warning("Adding uniform negatives makes no sense "
"when already using all negatives")
chunk_indices = torch.arange(chunk_size,
dtype=torch.long,
device=pos_embs.device)
last_chunk_indices = chunk_indices[:last_chunk_size]
# Ignore scores between positive pairs: since the i-th such pair has
# the pos_input[i] entity on this side, ignore_mask[i, pos_input[i]]
# must be set to 1 for every i. This becomes slightly more tricky as
# the rows may be wrapped into multiple chunks (the last of which
# may be smaller).
ignore_mask.append((
torch.arange(num_chunks - 1,
dtype=torch.long,
device=pos_embs.device).unsqueeze(1),
chunk_indices.unsqueeze(0),
pos_input_ten[:-last_chunk_size].view(num_chunks - 1,
chunk_size),
))
ignore_mask.append(
(-1, last_chunk_indices, pos_input_ten[-last_chunk_size:]))
else:
raise NotImplementedError("Unknown negative type %s" % type_)
return neg_embs, ignore_mask
def forward_direction_agnostic(
self,
src: EntityList,
dst: EntityList,
rel: Union[int, LongTensorType],
src_entity_type: str,
dst_entity_type: str,
src_operator: Union[None, AbstractOperator, AbstractDynamicOperator],
dst_operator: Union[None, AbstractOperator, AbstractDynamicOperator],
src_module: AbstractEmbedding,
dst_module: AbstractEmbedding,
src_pos: FloatTensorType,
dst_pos: FloatTensorType,
chunk_size: int,
src_negative_sampling_method: Negatives,
dst_negative_sampling_method: Negatives,
):
num_pos = len(src)
assert len(dst) == num_pos
src_pos = self.adjust_embs(src_pos, rel, src_entity_type, src_operator)
dst_pos = self.adjust_embs(dst_pos, rel, dst_entity_type, dst_operator)
num_chunks = ceil_of_ratio(num_pos, chunk_size)
src_dim = src_pos.size(-1)
dst_dim = dst_pos.size(-1)
if num_pos < num_chunks * chunk_size:
src_padding = src_pos.new_zeros(()).expand(
(num_chunks * chunk_size - num_pos, src_dim))
src_pos = torch.cat((src_pos, src_padding), dim=0)
dst_padding = dst_pos.new_zeros(()).expand(
(num_chunks * chunk_size - num_pos, dst_dim))
dst_pos = torch.cat((dst_pos, dst_padding), dim=0)
src_pos = src_pos.view((num_chunks, chunk_size, src_dim))
dst_pos = dst_pos.view((num_chunks, chunk_size, dst_dim))
src_neg, src_ignore_mask = self.prepare_negatives(
src,
src_pos,
src_module,
src_negative_sampling_method,
self.num_uniform_negs,
rel,
src_entity_type,
src_operator,
)
dst_neg, dst_ignore_mask = self.prepare_negatives(
dst,
dst_pos,
dst_module,
dst_negative_sampling_method,
self.num_uniform_negs,
rel,
dst_entity_type,
dst_operator,
)
pos_scores, src_neg_scores, dst_neg_scores = self.comparator(
src_pos, dst_pos, src_neg, dst_neg)
pos_scores = pos_scores.float()
src_neg_scores = src_neg_scores.float()
dst_neg_scores = dst_neg_scores.float()
# The masks tell us which negative scores (i.e., scores for non-existing
# edges) must be ignored because they come from pairs we don't actually
# intend to compare (say, positive pairs or interactions with padding).
# We do it by replacing them with a "very negative" value so that they
# are considered spot-on predictions with minimal impact on the loss.
for ignore_mask in src_ignore_mask:
src_neg_scores[ignore_mask] = -1e9
for ignore_mask in dst_ignore_mask:
dst_neg_scores[ignore_mask] = -1e9
# De-chunk the scores and ignore the ones whose positives were padding.
pos_scores = pos_scores.flatten(0, 1)[:num_pos]
src_neg_scores = src_neg_scores.flatten(0, 1)[:num_pos]
dst_neg_scores = dst_neg_scores.flatten(0, 1)[:num_pos]
reg = None
if self.regularizer is not None:
assert (src_operator is None) != (
dst_operator is
None), "Exactly one of src or dst operator should be None"
operator = src_operator if src_operator is not None else dst_operator
if self.num_dynamic_rels > 0:
reg = self.regularizer.forward_dynamic(src_pos, dst_pos,
operator, rel)
else:
reg = self.regularizer.forward(src_pos, dst_pos, operator)
return pos_scores, src_neg_scores, dst_neg_scores, reg
def prepare(self, embs: FloatTensorType) -> FloatTensorType:
# Dividing by the norm costs N * dim divisions, multiplying by the
# reciprocal of the norm costs N divisions and N * dim multiplications.
# The latter one is faster.
norm = embs.norm(2, dim=-1)
return embs * norm.reciprocal().unsqueeze(-1)
def save_embeddings(hf: h5py.File, embeddings: FloatTensorType) -> None:
hf.create_dataset(EMBEDDING_DATASET, data=embeddings.numpy())
def prepare_embs_for_reg(self, embs: FloatTensorType) -> FloatTensorType:
return embs.abs()
def forward(
self, embeddings: FloatTensorType, operator_idxs: LongTensorType
) -> FloatTensorType:
match_shape(embeddings, ..., self.dim)
match_shape(operator_idxs, *embeddings.size()[:-1])
return self.diagonals.to(device=embeddings.device)[operator_idxs] * embeddings
def forward(
self, embeddings: FloatTensorType, operator_idxs: LongTensorType
) -> FloatTensorType:
match_shape(embeddings, ..., self.dim)
match_shape(operator_idxs, *embeddings.size()[:-1])
return embeddings