使用GPU,哪些内容需要to(device)

目录

模型.to(device)

数据张量.to(device)

1、模型to(device)

2、但是数据那么多，哪些数据需要to(device)呢

网络接收数据的to(device) 

 其他数据to(device)

3、损失函数是否要to(device)呢

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

模型.to(device)

数据张量.to(device)

模型要么在CPU上跑要么在GPU上跑

如果没有统一，就会出现下面的报错

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

1、模型to(device)

我们通常遇到的形式如下

class Model(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4, 64),
            nn.Linear(64, 10)
        )
 
    def forward(self, x):
        x = self.model(x)
        return x
 
 
model = Model()

 那么只需要将模型的实例化对象to（device）就可以了

model=model.to(device)

model.to(device)的本质是将model网络里面的参数放到device中

2、但是数据那么多，哪些数据需要to(device)呢

我们将数据和模型to(device)的原因在于要将数据和模型都放在GPU上运算，换句话说，这里的数据是模型中运行的数据。

我们要确保的是模型中用到数据要和模型在一个device（GPU或者CPU）上面，其他的数据不需要管。

网络接收数据的to(device) 

我们一般见到的情况如下：（这里的模型是net）

        input_q = torch.from_numpy(input_q).long().to(device)
        input_qa = torch.from_numpy(input_qa).long().to(device)
        target = torch.from_numpy(target_1).float().to(device)
        if pid_flag:
            pid_one_seq = pid_data[:, idx * batch_size: (idx + 1) * batch_size]
            input_pid = np.transpose(pid_one_seq[:, :])
            input_pid = torch.from_numpy(input_pid).long().to(device)

            loss, pred, true_ct = net(input_q, input_qa, target, input_pid)
        else:
            loss, pred, true_ct = net(input_q, input_qa, target)

在将数据放在模型里之前就已经将其放到GPU上了

那可不可以在模型里面接受之后再to(device)

可以的 

    def forward(self, q_data, qa_data, target, pid_data=None):
        q_data= q_data.to(device)
        qa_data=qa_data.to(device)
        target=target.to(device)
        pid_data=pid_data.to(device)

 其他数据to(device)

一般情况对网络接收数据to(device)就可以了，网络处理过程中的数据一般都是从网络接收数据演变出来的，不需要再单独to(device)

凡事都有例外

x1 = torch.arange(seqlen).expand(seqlen, -1).to(device)

 比如说这条代码

他不是从网络接受数据中演变出来的，并且它还要参与网络后面的计算，那么他就需要单独to（device） 

3、损失函数是否要to(device)呢

我在查找哪些内容需要放到GPU上的时候，查到两种答案：

1、模型、数据张量

2、模型、数据张量和损失函数

所以现在的问题就变成了损失函数是否需要to(device)

看到了这样一个答案，觉得非常有道理

通常不需要。

简单地说，损失函数接收一个或多个输入tensor，如果输入tensor本身就是在gpu上，则输出tensor自然就在gpu上。

构建网络时，需要主动to或cuda的包括模型输入tensor，以及对应的Module(网络模型)对象。模型输入tensor是数据不用多说，Module对象使用.cuda是为了将模型内部的数据(即每个子叶Module的Parameter或Buffer成员变量)递归地传入gpu。

再进一步，如CrossEntropyCriterion等pytorch自带的loss类(或称loss layer)虽然也是Module的子类，但是这些类的成员变量不含Parameter对象或Buffer对象，所以无需to或cuda处理。Module对象是否持有数据，在设计上也是区分network layer和loss layer的关键。

链接：https://www.zhihu.com/question/314169571/answer/618862874