NLP工具——Stanza设置GPU device

1. 简介

这篇博客介绍如何在stanza工具中修改设置device。由于stanza模型代码中只预留了设置cpu还是cuda，但是没有给出设置device的选项，这导致我们在多卡的情况下调用模型时不够灵活。所以本文对这一内容进行介绍。

原理很简单，把所有的.cuda()修改为.to(device)即可。此方法同样适用于其他开源项目。

2. 修改

pipeline/core.py中，修改：

class Pipeline的__init__中增加一个参数，device=None：

# self.use_gpu = torch.cuda.is_available() and use_gpu
# 修改为：
self.use_gpu = device
1
2
3

models/depparse/trainer.py中，修改：

# inputs = [b.cuda() if b is not None else None for b in batch[:11]]
# 修改为：
inputs = [b.to(torch.device(use_cuda)) if b is not None else None for b in batch[:11]]
1
2
3

models/lemma/trainer.py中，类似的修改：

# inputs = [b.cuda() if b is not None else None for b in batch[:6]]
# 修改为：
inputs = [b.to(torch.device(use_cuda)) if b is not None else None for b in batch[:6]]

# self.model.cuda()
# self.crit.cuda()
# 修改为：
self.model.to(torch.device(use_cuda))
self.crit.to(torch.device(use_cuda))
1
2
3
4
5
6
7
8
9

models/mwt/trainer.py中，也是类似的修改：

# inputs = [b.cuda() if b is not None else None for b in batch[:4]]
# 修改为
inputs = [b.to(torch.device(use_cuda)) if b is not None else None for b in batch[:4]]

# self.model.cuda()
# self.crit.cuda()
# 修改为：
self.model.to(torch.device(use_cuda))
self.crit.to(torch.device(use_cuda))
1
2
3
4
5
6
7
8
9

pipeline/sentiment_processor.py：

# self._model.cuda()
# 修改为：
self._model.to(torch.device(use_gpu))
1
2
3

models/tokenization/trainer.py同理：
所有.cuda()替换为.to(torch.device(self.use_cuda))

models/common/seq2seq_model.py：

# self.SOS_tensor = self.SOS_tensor.cuda() if self.use_cuda else self.SOS_tensor
# 修改为
if self.use_cuda.startswith('cuda'):
	self.SOS_tensor = self.SOS_tensor.to(torch.device(self.use_cuda))

# return h0.cuda(), c0.cuda()
# 修改为：
h0 = h0.to(torch.device(use_cuda))
c0 = c0.to(torch.device(use_cuda))
1
2
3
4
5
6
7
8
9

以上内容可能仍有遗漏，如果修改之后还是不行，则找到报错的py中，搜搜cuda，然后做出同样的修改即可。

如有疑问，欢迎留言。