[Python] SpaCy GPU memory utilization for NER training

My training code:

spacy.require_gpu()
nlp = spacy.blank('en')

if 'ner' not in nlp.pipe_names:
ner = nlp.add_pipe('ner')
else:
ner = nlp.get_pipe('ner')

docs = load_data(ANNOTATED_DATA_FILENAME_BIN)
train_data, test_data = split_data(docs, DATA_SPLIT)

unique_labels = set(ent.label_ for doc in train_data for ent in doc.ents)
for label in unique_labels:
ner.add_label(label)

optimizer = nlp.initialize()

for i in range(EPOCHS):
print(f"Starting Epoch {i+1}...")
losses = {}
batches = minibatch(train_data, size=compounding(4., 4096, 1.001))
for batch in batches:
for doc in batch:
example = Example.from_dict(doc, {'entities': [(ent.start_char, ent.end_char, ent.label_) for ent in doc.ents]})
nlp.update([example], drop=0.5, losses=losses, sgd=optimizer)
print(f"Losses at iteration {i}: {losses}")


This code almost completely does not utilize GPU memory. Utilization is about 11-13% during training, which is almost the same as idle.



I did allocation test with torch, and all 8Gigs are allocated, so server works fine. The problem is with SpaCy or my code.

Could you please help?

Continue reading...