python - 如何测试 stanfordnlp 是否在 GPU 上运行？

Question

如何测试是否 stanfordnlp正在 GPU 上运行吗？

Here是一些示例代码:

import stanfordnlp
stanfordnlp.download('en')   # This downloads the English models for the neural pipeline
nlp = stanfordnlp.Pipeline() # This sets up a default neural pipeline in English
doc = nlp("Barack Obama was born in Hawaii.  He was elected president in 2008.")
doc.sentences[0].print_dependencies()

我稍微改变了这一点，指向我的模型目录:

config = {'models_dir': '/scratch/lklein/models'}
nlp = stanfordnlp.Pipeline(**config)

我的机器有 CUDA，当我运行代码时，我得到以下输出:

Use device: gpu
---
Loading: tokenize
With settings: 
{'model_path': '/scratch/lklein/models/en_ewt_models/en_ewt_tokenizer.pt', 'lang': 'en', 'shorthand': 'en_ewt', 'mode': 'predict'}

...

所以设置是正确的，它确实检测到了 GPU。我如何通过计算来查询这个？理想情况下，我正在寻找类似 spacy 的 require_gpu 的东西。

Answer 1

您可以尝试在代码执行时在 bash 中运行 nvidia-smi。

话虽这么说，默认行为是在可用时使用 GPU https://github.com/stanfordnlp/stanfordnlp/blob/d8061501ff14c73734e834a08fa33c58c4a6d917/stanfordnlp/pipeline/core.py#L102