CODEWITHSUNDEEP
spacymultiple-gpu
Mohit Sharma
Mohit Sharma

Reputation: 590

running spacy for predicting ner on mulitple GPUs

I am using spacy for predicting ner tags using gpu. I have a bigger machine with 8 GPUs, I want to use all those GPUs.

one way to use all gpus are running 8 different script on each gpu and then passing text to each script using kafka queue.

Is there any other way through which we can use single script to use all the gpus for predicting ner.

Upvotes: 1

Views: 306

Answers (1)

Jinhua Wang
Jinhua Wang

Reputation: 1759

Have you tried running Spacy with MPI? I am experimenting the following code myself, so please let me know if it works!

from mpi4py import MPI
import cupy

comm = MPI.COMM_WORLD
rank = comm.Get_rank()

if rank == 0:
    data = ["His friend Nicolas J. Smith is here with Bart Simpon and Fred."*100]
else:
    data = None

unit = comm.scatter(data, root=0)

with cupy.cuda.Device(rank):
    import spacy
    from thinc.api import set_gpu_allocator, require_gpu
    set_gpu_allocator("pytorch")
    require_gpu(rank)
    nlp = spacy.load('en_core_web_lg')
    nlp.add_pipe("merge_entities")
    tmp_list = []
    for doc in nlp.pipe(unit):
        res = " ".join([t.text if not t.ent_type_ else t.ent_type_ for t in doc])
        tmp_list.append(res)

result = comm.gather(tmp_list, root=0)

if comm.rank == 0:
    print (result)
else:
    result = None

Upvotes: 0

Related Questions