`Trainer.predict` takes twice as long as progress bar shows

Am using Trainer.predict but have noticed that it’s actually taking twice as long as displayed by the progress bar.

print('Predicting on test samples...')
t0 = time.time()
predictions, label_ids, _ = trainer.predict(tokenized_datasets['test'])
print(f'completed in {(time.time() - t0) / 60:.2f} mins.')
print('Argmaxing...')
t0 = time.time()
predictions = predictions.argmax(axis=2)
print(f'completed in {(time.time() - t0) / 60:.2f} mins.')

During the period shown by the progress bar, the GPU is used, but it’s not used after the progress bar has completed. Different sizes of tokenized_datasets['test'] has been tried, and it appears that the trend is that it takes twice as long regardless of the size.

Is this normal, or expected?

It depends on the Trainer you are using, which you did not share. The progress bar only represents the evaluation loop (goign through all the batches of your evaluation set and getting the predictions). After this is done, the Trainer computes metrics if you gave it a compute metrics function, which could take some time, and there can also be some additional post-processing, depending on your task.

1 Like