Hey there! Let’s work through this TensorFlow challenge you’re facing with your knowledge graph triples. I get it—switching from NumPy’s easy dictionary lookups to tensor-based operations can feel like hitting a wall at first, but we can build a clean, graph-compatible solution here.

First, Let’s Recap Your Goal #

For each triple in your batch (e.g., <a,y,b>), you need to fetch all other triples from the knowledge graph that share the same subject (so for <a,y,b>, that means <a,y,c> and <a,y,d>). The catch is doing this entirely with TensorFlow tensors, no NumPy workarounds.

Step 1: Preprocess the Knowledge Graph into a Tensor-Friendly Structure #

Instead of a Python dictionary, we’ll use TensorFlow’s RaggedTensor to group triples by their subject. This is perfect for handling variable-length lists of triples per subject.

import tensorflow as tf
import numpy as np

# Your original knowledge graph
X = np.array([['a', 'y', 'b'], ['b', 'y', 'a'], ['a', 'y', 'c'], ['c', 'y', 'a'],
              ['a', 'y', 'd'], ['c', 'y', 'd'], ['b', 'y', 'c'], ['f', 'y', 'e']])

# Convert to a TensorFlow string tensor
X_tf = tf.convert_to_tensor(X, dtype=tf.string)

# Get unique subjects and map each triple to its subject's index
unique_subjects, subject_indices = tf.unique(X_tf[:, 0])

# Group triples by subject using a RaggedTensor
grouped_triples = tf.RaggedTensor.from_value_rowids(
    values=X_tf,
    value_rowids=subject_indices
)

# Now grouped_triples[i] holds all triples for unique_subjects[i]

Step 2: Process Batched Triples to Fetch Same-Subject Triples #

We’ll use tf.map_fn to iterate over each triple in your batch. For each triple, we’ll:

1. Find its subject’s position in the unique subjects list
2. Pull all triples for that subject
3. Filter out the current triple itself

def filter_same_subject_triples(exclude_triple):
    # Extract the subject from the current triple
    current_subject = exclude_triple[0]
    
    # Find the index of this subject in our unique list
    subject_idx = tf.where(tf.equal(unique_subjects, current_subject))[0][0]
    
    # Get all triples for this subject
    all_subject_triples = grouped_triples[subject_idx]
    
    # Create a mask to exclude the current triple
    # Check if each triple matches the one we want to exclude
    is_excluded = tf.reduce_all(tf.equal(all_subject_triples, exclude_triple), axis=1)
    keep_mask = tf.logical_not(is_excluded)
    
    # Filter the triples
    filtered_triples = tf.boolean_mask(all_subject_triples, keep_mask)
    return filtered_triples

# Example batch of triples (replace this with your dataset iterator output)
x_pos_tf = tf.convert_to_tensor([['a','y','b'], ['c','y','a']], dtype=tf.string)

# Process the entire batch
result = tf.map_fn(
    filter_same_subject_triples,
    x_pos_tf,
    fn_output_signature=tf.RaggedTensorSpec(shape=[None, 3], dtype=tf.string)
)

# Print the result to verify
print(result.to_list())
# Output: [[['a', 'y', 'c'], ['a', 'y', 'd']], [['c', 'y', 'd']]]

Key Notes for Your Use Case #

• Integer IDs: If your triples use integer-encoded entities/relations (instead of strings), just change the dtype to tf.int32—the logic stays exactly the same.
• Graph Compatibility: This code runs entirely in TensorFlow’s graph mode, so it works seamlessly with your dataset iterator (dataset_iterator.get_next()). No more mixing NumPy and TensorFlow operations!
• Efficiency: For most knowledge graphs, tf.map_fn will be fast enough. If you’re working with an extremely large batch, you could optimize further with vectorized operations, but this is a solid starting point.

内容的提问来源于stack exchange，提问作者snelzb