faster ranking at prediction time
@@ -9,6 +9,7 @@ import os
import torch
from torch import nn
+from itertools import islice
from parlai.core.torch_agent import TorchAgent, Output
from parlai.core.thread_utils import SharedTable
from parlai.core.utils import round_sigfigs, padded_3d, warn_once
@@ -47,6 +48,9 @@ class TorchRankerAgent(TorchAgent):
help='Get predictions and calculate mean rank during the train '
'step. Turning this on may slow down training.'
)
+ agent.add_argument(
+ '--cap-num-predictions', type=int, default=100,
+ help='Limit to the number of predictions in output.text_candidates')
def __init__(self, opt, shared=None):
# Must call _get_model_file() first so that paths are updated if necessary
@@ -180,7 +184,7 @@ class TorchRankerAgent(TorchAgent):
scores = self.score_candidates(batch, cand_vecs)
_, ranks = scores.sort(1, descending=True)
-
+ ranks = ranks.cpu()
# Update metrics
if label_inds is not None:
loss = self.rank_loss(scores, label_inds)
@@ -190,6 +194,7 @@ class TorchRankerAgent(TorchAgent):
rank = (ranks[b] == label_inds[b]).nonzero().item()
self.metrics['rank'] += 1 + rank
+ max_preds = self.opt["cap_num_predictions"]
cand_preds = []
for i, ordering in enumerate(ranks):
if cand_vecs.dim() == 2:
@@ -200,7 +205,10 @@ class TorchRankerAgent(TorchAgent):
# ignore padding
true_ordering = [x for x in ordering if x < len(cand_list)]
ordering = true_ordering
- cand_preds.append([cand_list[rank] for rank in ordering])
+ # using a generator instead of a list comprehension allows
+ # to cap the number of elements.
+ cand_preds_generator = (cand_list[rank] for rank in ordering)
+ cand_preds.append(list(islice(cand_preds_generator, max_preds)))
preds = [cand_preds[i][0] for i in range(batchsize)]
return Output(preds, cand_preds)
