Some tensors are not moved to the device I want when doing .to(device)

user_9793316 0
Viewed 163

The Problem

I have a model that requires different linear models for different time steps:

class CPC(nn.Module):
    def __init__(self, batch_size, timestep):

        super().__init__()
        self.batch_size = batch_size

        self.timestep = timestep  
        self.Wk  = nn.ModuleList([nn.Linear(77, 55) for i in range(timestep)])

        def _weights_init(m):
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        pred = torch.empty((self.batch_size,self.timestep,55)).float()
        for i in np.arange(0, self.timestep):
            linear = self.Wk[i]
            pred[:,i] = linear(x[:,i]) # Wk*c_t e.g. size 8*512

        return pred

When doing model=CPC(batch_size, timestep).to(device), pred is still in cpu instead of the device I want.

My attempt to solve it

Then I modified my code to be something like this:

class CPC(nn.Module):
    def __init__(self, batch_size, timestep):

        super().__init__()

        self.timestep = timestep  
        pred = torch.nn.Parameter(torch.empty((batch_size,self.timestep,55)).float())
        self.register_parameter('pred', pred)
        self.Wk  = nn.ModuleList([nn.Linear(77, 55) for i in range(timestep)])

        def _weights_init(m):
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        for i in np.arange(0, self.timestep):
            linear = self.Wk[i]
            self.pred[:,i] = linear(x[:,i]) # Wk*c_t e.g. size 8*512

        return self.pred

Now, model.to(device) moves everything to device, but I get the following error when training.

RuntimeError: leaf variable has been moved into the graph interior

What should I do such that model.to(device) moves everything to the device I want, and still make the model trainable?

Dummy code for debugging

Here's my full code for training:

import torch
import torch.nn as nn
import numpy as np
from torchvision import transforms, datasets

device='cuda:0'

class CPC(nn.Module):
    def __init__(self, batch_size, timestep):

        super().__init__()

        self.timestep = timestep  
        pred = torch.nn.Parameter(torch.empty((batch_size,self.timestep,55)).float())
        self.register_parameter('pred', pred)
        self.Wk  = nn.ModuleList([nn.Linear(77, 55) for i in range(timestep)])

        def _weights_init(m):
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        for i in np.arange(0, self.timestep):
            linear = self.Wk[i]
            self.pred[:,i] = linear(x[:,i]) # Wk*c_t e.g. size 8*512

        return self.pred   

x = torch.randn(16, 12, 77)
A = torch.randn(77,55)
y = torch.matmul(x, A)

model = CPC(x.size(0), x.size(1))
model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
loss_fn = torch.nn.MSELoss()

x = x.to(device)
y = y.to(device)
for i in range(10):
    optimizer.zero_grad()      
    pred = model(x)
#     print(f'model.pred={model.pred}')
    loss = loss_fn(pred, y)
    loss.backward()
    optimizer.step()
pytorch python
1 Answers

The problem is torch.empty by default gives you a cpu tensor and it doesn't look like you've moved it elsewhere.

This works, although not very elegant.

class CPC(nn.Module):
    def __init__(self, batch_size, timestep):

        super().__init__()
        self.batch_size = batch_size

        self.timestep = timestep  
        self.Wk  = nn.ModuleList([nn.Linear(77, 55) for i in range(timestep)])

        def _weights_init(m):
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        pred = torch.empty((self.batch_size, self.timestep, 55)).float().to(self.Wk[0].weight.device)
        for i in np.arange(0, self.timestep):
            linear = self.Wk[i]
            pred[:,i] = linear(x[:,i]) # Wk*c_t e.g. size 8*512

        return pred

This looks better (probably) but I usually too lazy to write 4 more lines.

    def __init__

        self.device = 'cpu'

    def to(self, device):
        super().to(device)
        self.device = device

    def forward(self, x):
        pred = torch.empty((self.batch_size, self.timestep, 55)).float().to(self.device)
user_94773380
Related