cisnlp/simalign

The speed is not improved when using gpu

Closed this issue · 2 comments

Hi,
Thank you for open-source a well-performed word alignment tool. When running simalign, I met a problem about the model speed. It takes about 160 seconds to deal with a batch (100 sentence pairs) of data on cpu. When using the GPU by setting the device as "cuda", the speed is still about 160 seconds a batch. I think the data is loaded on the gpu, cause part of the memory of the gpu is being used. Is there any insights?

hi @Life-0-1 thanks for pointing this issue out. One potential explanation is that the Match (mwmf) method is quite slow. For 100 sentences Match could be the bottleneck. I suggest to just compute the ArgMax or IterMax alignments and then there should be some speed-up (although for 100 sentences it is probably not that big). It took us 15 seconds to align 500 sentence on a GPU. Please let us know whether this resolved your issue.

The default device is CPU on SentenceAligner(). You will need to specify cuda like this:

import torch
myaligner = SentenceAligner(model="xlmr", token_type="bpe", matching_methods="a", device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))

~3 seconds for 100 sentences