Can I create a model with some layers on CPU, some on GPU?

1

I want to validate the output of each layers running on two different platforms. In order to do so, I have registered forward/backward hooks, use same input and model weights, and dumped the data. But I noticed that there is dropout or even LSTM with dropout layers in my models. The dropout algorithm is different on the two platforms, so the output tensor after dropout is different. Thus any layer’s output after the first dropout cannot be compared directly. One way to address this issue, afaik, is maybe move those dropout layers onto my CPU, and I just checked that the two CPUs dropout can produce same results. Thus, I’m wondering if it is possible to put those dropout layers on CPU and other “more deterministic” layers still on gpus. Or do you have any other suggestions to compare the layers’ outputs?

2

Yes, that’s possible and you could move the layers directly to the desired device via model.layer.to(device) and would need to move the data in the forward method as well:

def forward(self, x):
    x = self.gpu_layer1(x)
    x = x.cpu()
    x = self.cpu_layer1(x)
    x = self.cpu_layer2(x)
    x = x.to('cuda')
    x = self.gpu_layer2(x)
    return x
3

Thanks, I will give it a try.