Getting model predictions in parallel

[oliver-lloyd]

Hi

I have a model from which I am trying to get scores from a large number of triples. To speed up the process I would like to get these scores in parallel, but I am having an issue where the models will only return predictions in serial. When running in parallel with multiprocessing the program runs indefinitely without returning anything. Please see code below.

Any help would be much appreciated!

def score_triples(triple_list, model):
    s = Tensor([edge[0] for edge in triple_list])
    p = Tensor([edge[1]for edge in triple_list])
    o = Tensor([edge[2] for edge in triple_list])
    preds = model.score_spo(s, p, o).tolist()
    triples_scored = pd.DataFrame(triple_list, columns=['head', 'relation', 'tail'])
    triples_scored['score'] = preds
    return triples_scored

# Load model
checkpoint = load_checkpoint('checkpoint.pt')
model = KgeModel.create_from(checkpoint)

# Group triples by relation
test_drugs = [3565, 5022, 8174, 11516, 8361, 4264]
test_rels = [1151, 1223, 1193, 1786]
triple_sets = [
    list(product(
        test_drugs,
        [rel],
        test_drugs
    ))
    for rel in test_rels
]

# Group args for multiprocessing
mp_args = [[trip_set, model] for trip_set in triple_sets]

# Runs fine when predicting for one set of triples
score_triples(*mp_args[0])

# Process runs forever when predicting for several sets in parallel
with mp.Pool(mp.cpu_count()) as pool:
    results = pool.starmap(score_triples, mp_args)

[rgemulla]

This seems to be more of a torch than a libkge issues. Do you use import torch.multiprocessing as mp?

Also note that this code would probably more efficient (and in no need for multiprocessing) if you did not use score_spo but score_sp_po (e.g., call it once per relation to score all so pairs at once). Quoting from the docs:

    def score_sp_po(
        self, s: Tensor, p: Tensor, o: Tensor, entity_subset: Tensor = None
    ) -> Tensor:
        r"""Combine `score_sp` and `score_po`.

        `s`, `p` and `o` are vectors of common size :math:`n`, holding the indexes of
        the subjects, relations, and objects to score.

        Each sp-pair and each po-pair is scored against the entities in `entity_subset`
        (also holds indexes). If set to `entity_subset` is `None`, scores against all
        entities.

        The result is the horizontal concatenation of the outputs of
        :code:`score_sp(s,p,entity_subset)` and :code:`score_po(p,o,entity_subset)`.
        I.e., returns an :math:`n\times 2E` tensor, where :math:`E` is the size of
        `entity_subset`. For :math:`j<E`, the :math:`(i,j)`-entry holds the score for
        triple :math:`(s_i, p_i, e_j)`. For :math:`j\ge E`, the :math:`(i,j)`-entry
        holds the score for triple :math:`(e_{j-E}, p_i, o_i)`.

        """