PyTorch RuntimeError求助：Input、output和indices需在当前设备

嘿，我看你在PyTorch上用BERT处理3万条多分类情感数据时，明明已经把数据迁移到CUDA了，却还是碰到了RuntimeError: Input, output and indices must be on the current device的报错，咱们来一步步排查解决：

你的训练代码 #

for epoch in tqdm(range(1, epochs+1)):
    model.train()
    loss_train_total = 0
    progress_bar = tqdm(dataloader_train, desc='Epoch {:1d}'.format(epoch), leave=False, disable=False)
    for batch in progress_bar:
        model.zero_grad()
        batch = tuple(b.to(device) for b in batch)
        inputs = {'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[2], }
        outputs = model(**inputs)
        loss = outputs[0]
        loss_train_total += loss.item()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        progress_bar.set_postfix({'training_loss': '{:.3f}'.format(loss.item()/len(batch))})
    torch.save(model.state_dict(), f'finetuned_BERT_epoch_{epoch}.model')
    tqdm.write(f'\nEpoch {epoch}')
    loss_train_avg = loss_train_total/len(dataloader_train)
    tqdm.write(f'Training loss: {loss_train_avg}')
    val_loss, predictions, true_vals = evaluate(dataloader_validation)
    val_f1 = f1_score_func(predictions, true_vals)
    tqdm.write(f'Validation loss: {val_loss}')
    tqdm.write(f'F1 Score (Weighted): {val_f1}')

报错堆栈 #

---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-67-9306225bb55a> in <module>()
     17 }
     18 --> 19 outputs = model(**inputs)
     20 
     21 loss = outputs[0]
8 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
   1850         # remove once script supports set_grad_enabled
   1851         _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
-> 1852         return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
   1853 
   1854 RuntimeError: Input, output and indices must be on the current device

具体解决建议 #

• 确认模型完全迁移到CUDA：很多时候我们以为模型已经在GPU上，但可能初始化后没执行model = model.to(device)，或者执行位置不对。你可以在训练循环前加一行print(next(model.parameters()).device)，如果输出不是cuda:0（或对应GPU编号），那就说明模型还在CPU上，赶紧补上model = model.to(device)。
• 验证batch内所有数据的设备：虽然你写了batch = tuple(b.to(device) for b in batch)，但可以在循环里加个打印验证，比如print(batch[0].device, batch[1].device, batch[2].device)，确保这三个张量都在CUDA上，避免某些特殊情况导致迁移失败。
• 检查embedding层的权重设备：报错来自embedding操作，大概率是BERT的词嵌入权重还在CPU上。如果前面的模型迁移操作没生效，可以手动把embedding层移过去：model.bert.embeddings.word_embeddings = model.bert.embeddings.word_embeddings.to(device)，不过更稳妥的是确保整个模型都完成设备迁移。
• 确认优化器和调度器的初始化时机：优化器需要在模型迁移到CUDA之后再初始化，否则优化器的参数会绑定在CPU上的模型参数。如果你的优化器是在model.to(device)之前定义的，重新调整顺序：先移模型到GPU，再创建优化器和调度器。
• 排查验证集的数据设备：训练集处理了，但验证集的dataloader_validation里的数据可能没迁移到CUDA，记得在evaluate函数里也要把batch数据移到device上，不然验证时也会出现同样的设备不匹配问题。

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