PyTorch GPU搬迁错误解决方法

一、问题描述

在将原本运行在 CPU 上的 PyTorch 代码移动到 GPU 上时,容易报错：

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

这是由于并未将数据与模型同时加载到 GPU

二、知识背景

在 PyTorch 中，数据 Tensor 与模型都能用 .to()/.cuda()搬运到 GPU上

import torch
from torch.utils.data import DataLoader

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# 搬迁模型到 GPU
model = MyModel().to(device)

# 搬迁数据到 GPU
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch in dataloader:
    inputs, labels = batch[0].to(device), batch[1].to(device)

三、额外的注意点

1. Dataset 实例对象返回数据时，此时无需把返回数据搬移到 GPU
2. 模型的Tokenizer 实例无需搬移到 GPU
3. DataLoader 实例本身无需搬移到 GPU,需要搬移到 GPU的是其返回的数据对象。
4. 定义 DataLoader 时若额外定义了collate_fn,.to(device)应当写在collate_fn()内，举例：

def collate_fn(data):  
    sents = [i[0] for i in data][:1600]  
    labels = [i[1] for i in data][:1600]  
  
    #编码  
    data = token.batch_encode_plus(batch_text_or_text_pairs=sents,  
                    truncation=True,  
                    padding='max_length',  
                    max_length=500,  
                    return_tensors='pt',  
                    return_length=True)  
  
    input_ids = data['input_ids']  
    attention_mask = data['attention_mask']  
    token_type_ids = data['token_type_ids']  
    labels = torch.LongTensor(labels)  
  
    return input_ids.to(device), attention_mask.to(device), token_type_ids.to(device), labels.to(device)
    
#数据加载器
loader = torch.utils.data.DataLoader(dataset=dataset,
                                     batch_size=16,
                                     collate_fn=collate_fn,
                                     shuffle=True,
                                     drop_last=True)

5. model.to(device)会将模型的参数及内嵌子模型的所有参数都递归地搬移到 GPU：

import torch
import torch.nn as nn

# 定义子模型
class SubModel(nn.Module):
    def __init__(self):
        super(SubModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x

# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.submodel = SubModel()
        self.fc1 = nn.Linear(32 * 32 * 32, 10)

    def forward(self, x):
        x = self.submodel(x)
        x = x.view(-1, 32 * 32 * 32)
        x = self.fc1(x)
        return x

# 创建模型实例
model = MyModel()

# 将模型及其子模型的参数全部放到 GPU 上
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 检查模型参数所在的设备
for name, param in model.named_parameters():
    print(name, param.shape, param.device)

在这个示例中，我们定义了一个名为 SubModel 的子模型和一个名为 MyModel 的主模型。 主模型包含一个嵌套的子模型，并且还有一个全连接层。
我们首先创建模型实例 model，然后使用 to() 方法将模型及其子模型的所有参数移动到 GPU 上。
最后，我们使用 named_parameters() 方法来检查模型参数所在的设备，以确保所有参数都已经被移动到 GPU 上。

6.如果下游模型所用到的 sub模型 创建了实例，也需要搬到 GPU:

from transformers import BertModel

# sub 模型创建了实例
sub = BertModel.from_pretrained('bert-base-chinese')
# 必须同样搬运到GPU,否则报错
sub.to(device)

# 下游模型
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 新增一个全连接层
        self.fc = torch.nn.Linear(768, 2)

    def forward(self, arg1,arg2,**kwargs):
        with torch.no_grad():
            bert = sub(arg1,arg2,**kwargs)
        
        out = self.fc(bert.last_hidden_state[:, 0]).to(d2l.try_gpu())

        out = out.softmax(dim=1)

        return out


# 下游模型搬运到 GPU 上
model = Model()
model.to(device)