使用Hugging Face Pipeline批量文本分类遇设备分配问题求助

你的代码中将模型部署到了GPU，但未告知pipeline使用GPU设备，导致tokenizer生成的输入张量留在CPU，触发设备不匹配错误。

原代码 #

tokenizer_filter = AutoTokenizer.from_pretrained("salesken/query_wellformedness_score")
tokenizer_kwargs = {'padding':True,'truncation':True,'max_length':512}
model_filter = AutoModelForSequenceClassification.from_pretrained("salesken/query_wellformedness_score").to(torch.device("cuda"))

filtering = pipeline("text-classification", model=model_filter, tokenizer=tokenizer_filter, batch_size=8)
scores = filtering(df['content'].tolist(), **tokenizer_kwargs)

报错信息 #

Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)

解决方案 #

在创建pipeline时显式指定device参数，让pipeline自动将输入张量转移到GPU：

tokenizer_filter = AutoTokenizer.from_pretrained("salesken/query_wellformedness_score")
tokenizer_kwargs = {'padding':True,'truncation':True,'max_length':512}
model_filter = AutoModelForSequenceClassification.from_pretrained("salesken/query_wellformedness_score").to(torch.device("cuda"))

# 指定device为0（对应第一个GPU）或torch.device("cuda")
filtering = pipeline("text-classification", model=model_filter, tokenizer=tokenizer_filter, batch_size=8, device=0)
scores = filtering(df['content'].tolist(), **tokenizer_kwargs)

说明 #

字符串本身无法直接部署到GPU，需要转移的是tokenizer处理后生成的张量。通过给pipeline指定device参数，pipeline会自动完成输入张量的设备同步，确保其与模型处于同一设备，无需手动处理字符串列表的设备转移。

内容的提问来源于stack exchange，提问作者Lucas Azevedo