1

我在 colab 上运行以下代码，取自此处的示例：https ://huggingface.co/transformers/model_doc/albert.html#albertformaskedlm

import os
import torch
import torch_xla
import torch_xla.core.xla_model as xm

assert os.environ['COLAB_TPU_ADDR']

dev = xm.xla_device()

from transformers import AlbertTokenizer, AlbertForMaskedLM
import torch

tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
model = AlbertForMaskedLM.from_pretrained('albert-base-v2').to(dev)
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1

data = input_ids.to(dev)

outputs = model(data, masked_lm_labels=data)
loss, prediction_scores = outputs[:2]

我没有对示例代码做任何事情，input_ids除了model使用.to(dev). 似乎一切都移到了 TPU 没有问题，因为当我输入时，data我得到以下输出：tensor([[ 2, 10975, 15, 51, 1952, 25, 10901, 3]], device='xla:1')

但是，当我运行此代码时，出现以下错误：

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-5-f756487db8f7> in <module>()
      1 
----> 2 outputs = model(data, masked_lm_labels=data)
      3 loss, prediction_scores = outputs[:2]

9 frames
/usr/local/lib/python3.6/dist-packages/transformers/modeling_albert.py in forward(self, hidden_states, attention_mask, head_mask)
    277         attention_output = self.attention(hidden_states, attention_mask, head_mask)
    278         ffn_output = self.ffn(attention_output[0])
--> 279         ffn_output = self.activation(ffn_output)
    280         ffn_output = self.ffn_output(ffn_output)
    281         hidden_states = self.full_layer_layer_norm(ffn_output + attention_output[0])

RuntimeError: Unknown device

有谁知道发生了什么？

4

1 回答 1

0

解决方案在这里：https ://github.com/pytorch/xla/issues/1909

在调用之前model.to(dev)，您需要调用xm.send_cpu_data_to_device(model, xm.xla_device())：

model = AlbertForMaskedLM.from_pretrained('albert-base-v2')
model = xm.send_cpu_data_to_device(model, dev)
model = model.to(dev)

让 ALBERT 使用的 gelu 激活函数在 TPU 上工作也存在一些问题，因此在 TPU 上工作时需要使用以下转换器分支：https ://github.com/huggingface/transformers/tree/fix- jitpu

有关完整解决方案，请参阅以下 colab 笔记本（由https://github.com/jysohn23提供）： https ://colab.research.google.com/gist/jysohn23/68d620cda395eab66289115169f43900/getting-started-with-pytorch-on-cloud -tpus.ipynb

于 2020-04-15T10:09:04.767 回答