在跑Pytorch官方的Bilstm-CRF模型代码的时候,希望可以利用实验室的服务器的GPU加速训练。
第一次尝试
下面是我第一次尝试使用的部分代码
# GPU
import torch
dev_id = 1
dev = torch.device(f'cuda:{dev_id}')
# 将模型转入到GPU
model = model.to(dev)
# 将相应的数据和loss函数也要转移到GPU中
# sentence就是模型训练输入的数据,tags是对应的标签
sentence = sentence.to(dev)
tags = tags.to(dev)
loss = model.neg_log_likelihood(sentence, tags)
loss.to(dev)
但是loss函数在计算的时候一直出现
cuDNN error: CUDNN_STATUS_INTERNAL_ERROR
第二次尝试
将代码里面的.to(dev)
全部改成.cuda(dev)
就成功了,没有再出现
cuDNN error: CUDNN_STATUS_INTERNAL_ERROR。
目前我也不太清楚原因
# GPU
import torch
dev_id = 1
dev = torch.device(f'cuda:{dev_id}')
# 将模型转入到GPU
model = model.cuda(dev)
# 将相应的数据和loss函数也要转移到GPU中
# sentence就是模型训练输入的数据,tags是对应的标签
sentence = sentence.cuda(dev)
tags = tags.cuda(dev)
loss = model.neg_log_likelihood(sentence, tags)
loss.cuda(dev)