使用HuggingFace AutoModel搭配TensorFlow做迁移学习报错如何解决

报错原因 #

• 核心问题是混用了TensorFlow和PyTorch两个框架的API：
  • 调用tokenizer.encode时指定return_tensors="tf"，得到的是TensorFlow格式的张量
  • AutoModel.from_pretrained默认加载的是PyTorch版本的模型，.to("cuda")也是PyTorch专属的模型/张量移动语法，PyTorch模型无法直接接收TensorFlow张量作为输入

解决方案 #

方案1：全程使用TensorFlow栈（匹配你原本搭配TensorFlow API的需求） #

修改模型加载相关代码，使用Transformers库专门适配TensorFlow的模型类即可，TensorFlow会自动调度可用计算设备，无需显式调用to方法移动模型：

from transformers import AutoTokenizer, TFAutoModel # 把AutoModel替换为TFAutoModel

model_name = "distilbert-base-uncased"
text = "this is a test"
tokenizer = AutoTokenizer.from_pretrained(model_name)    
text_tensor = tokenizer.encode(text, return_tensors="tf")

model = TFAutoModel.from_pretrained(model_name) # 去掉.to("cuda")语句
output = model(text_tensor)

方案2：全程使用PyTorch栈 #

如果可以切换为PyTorch实现，仅需修改张量返回格式为PyTorch版本，同时把输入张量同步移动到对应设备即可：

from transformers import AutoTokenizer, AutoModel

model_name = "distilbert-base-uncased"
text = "this is a test"
tokenizer = AutoTokenizer.from_pretrained(model_name)    
text_tensor = tokenizer.encode(text, return_tensors="pt") # 把tf替换为pt

model = AutoModel.from_pretrained(model_name).to("cuda")
output = model(text_tensor.to("cuda")) # 输入张量同步移动到cuda设备，避免设备不匹配报错

补充说明 #

如果你的设备没有可用的NVIDIA GPU，调用.to("cuda")也会触发报错，这种情况下把参数"cuda"改为"cpu"即可。

内容的提问来源于stack exchange，提问作者hans glick