一、问题描述
使用transformers包调用pytorch框架的Bert预训练模型时,使用正常的bert-base-cased在其他数据集上正常运行,但是运用Roberta的时候却一直报错:RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling `cublasCreate(handle)`
忙活了好几天也没查出是什么错误,网上一直提醒batch_size是否太大导致,修改为16->8->4->2都没有用。
通过与其他数据集的对比,发现我在tokenizer中加入了新的special_token,导致可能报错了!
二、问题解决
在原始的tokenizer中加入special_tokens时,忘记将model的tokenizer的词表进行更新导致!
完整更新方式为:
from transformers import BertTokenizer, BertModel
tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
# 添加特殊词
tokenizer.add_special_tokens({'additional_special_tokens':["<S>"]})
model = BertModel.from_pretrained("bert-base-cased")
# 在模型中更新词表的大小!
# 重要!
model.resize_token_embeddings(len(tokenizer))
三、问题解决
可以通过,开始训练!