Trouble with GPU training pipeline

[Angelo92Git]

Hi I'm having trouble with training my model on Google Colab:

I keep getting the error that some tensors are on CPU and some are on GPU despite moving both the model and the HeteroData objects to GPU (I have loaded HeteroData objects into a dataloader)

Here are the relevant sections of my code:

# Create DataLoaders
num_training_instances = len(training_dataset)
num_validation_instances = len(validation_dataset)
num_testing_instances = len(testing_dataset)
print(f"Number of training instances: {num_training_instances}")
print(f"Number of validation instances: {num_validation_instances}")
print(f"Number of testing instances: {num_testing_instances}")
train_loader = DataLoader(training_dataset, batch_size=20, shuffle=True)
val_loader = DataLoader(validation_dataset, batch_size=20, shuffle=True)
test_loader = DataLoader(testing_dataset, batch_size=20, shuffle=True)


hidden_channels = 128
model = m.G4GCN_VCG(hidden_channels=hidden_channels, num_conv_layers=50, include_meta_node=True)
model = model.to(model_device)
decoder = m.MLP([hidden_channels, hidden_channels, hidden_channels, 1])
decoder = decoder.to(model_device)
model_params = list(model.parameters()) + list(decoder.parameters())
optimizer = torch.optim.Adam(model_params, lr=1e-4, weight_decay=1e-8)
train_criterion = BCELoss()

# Train the model
train_loss_record = []
val_loss_record = []
best_val_loss = 999
pbar = tqdm(range(1), desc="Starting", position=0)
for epoch in pbar:
    model.train()
    for batched_graphs in train_loader:
        batched_graphs = ToDevice(device=device)(batched_graphs)
        optimizer.zero_grad()
        out = model(batched_graphs.x_dict, batched_graphs.deg_dict, batched_graphs.edge_index_dict)
        out = scatter_mean(out, batched_graphs['variable']['batch'], dim=0)
        out = decoder(out).sigmoid()
        loss = BCELoss()(out, batched_graphs.y)
        with torch.no_grad():
            pbar.set_description(f"Loss: {loss.item():.2f}")
            train_loss_record.append(loss.item())
        loss.backward()
        optimizer.step()

And here is the stack trace:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-45-0f4843b8d92b>](https://localhost:8080/#) in <cell line: 82>()
     88         # batched_graphs.edge_index_dict = batched_graphs.edge_index_dict.to(device)
     89         optimizer.zero_grad()
---> 90         out = model(batched_graphs.x_dict, batched_graphs.deg_dict, batched_graphs.edge_index_dict)
     91         out = scatter_mean(out, batched_graphs['variable']['batch'], dim=0)
     92         out = decoder(out).sigmoid()

5 frames
[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1509             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1510         else:
-> 1511             return self._call_impl(*args, **kwargs)
   1512 
   1513     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1518                 or _global_backward_pre_hooks or _global_backward_hooks
   1519                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1520             return forward_call(*args, **kwargs)
   1521 
   1522         try:

[/tmp/GNN_SAT_proj/models.py](https://localhost:8080/#) in forward(self, x_dict, deg_dict, edge_index_dict)
    110             x_dict['variable'] = torch.cat([x_dict['variable'], x_dict_prev['variable']], dim=1)
    111             x_dict['clause'] = self.lins_c[conv_layer](x_dict['clause'])
--> 112             x_dict['variable'] = self.lins_v[conv_layer](x_dict['variable'])
    113         out = x_dict['variable']
    114         return out

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1509             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1510         else:
-> 1511             return self._call_impl(*args, **kwargs)
   1512 
   1513     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1518                 or _global_backward_pre_hooks or _global_backward_hooks
   1519                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1520             return forward_call(*args, **kwargs)
   1521 
   1522         try:

[/usr/local/lib/python3.10/dist-packages/torch_geometric/nn/dense/linear.py](https://localhost:8080/#) in forward(self, x)
    145             x (torch.Tensor): The input features.
    146         """
--> 147         return F.linear(x, self.weight, self.bias)
    148 
    149     @torch.no_grad()

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 mat1 in method wrapper_CUDA_addmm)

[akihironitta]

Is model_device the same device as device in your code?

[Angelo92Git]

Hi, here is how I set device and model_device, I was thinking of experimenting with DataListLoader but hadn't gotten around to it

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

[Angelo92Git]

I'm trying to implement my own version of the paper G4SATBench, in particular the GCN model on the Variable-Clause Graph (VCG) representation of Boolean SAT instances

In case it's relevant here is the metadata for the batched_graphs object:

batched_graphs.metadata()[0]
['clause', 'variable']

batched_graphs.metadata()[1]
[('clause', 'contains_pos', 'variable'), ('clause', 'contains_neg', 'variable'), ('variable', 'rev_contains_pos', 'clause'), ('variable', 'rev_contains_neg', 'clause')]

and my model defintion:

class G4GCN_VCG(nn.Module):
    def __init__(self, hidden_channels=8, num_conv_layers=4, include_meta_node=False):  # meta node not used
        super().__init__()
        self.include_meta_node = include_meta_node
        if self.include_meta_node:
            self.lin0_m = Linear(1, hidden_channels, bias=False)
            self.lins_m = []
        self.lin0_c = Linear(1, hidden_channels, bias=False)
        self.lin0_v = Linear(1, hidden_channels, bias=False)
        self.convs = torch.nn.ModuleList()
        self.lins_c = torch.nn.ModuleList()
        self.lins_v = torch.nn.ModuleList()

        for _ in range(num_conv_layers): # set up convolution layers
            conv_dict = {
                ('clause', 'contains_pos', 'variable'): G4GCNConv(*[hidden_channels]*3),
                ('clause', 'contains_neg', 'variable'): G4GCNConv(*[hidden_channels]*3),
                ('variable', 'rev_contains_pos', 'clause'): G4GCNConv(*[hidden_channels]*3),
                ('variable', 'rev_contains_neg', 'clause'): G4GCNConv(*[hidden_channels]*3)
            }
            if self.include_meta_node:
                conv_dict[('meta', 'connects', 'clause')] = G4GCNConv(*[hidden_channels]*3)
                conv_dict[('clause', 'rev_connects', 'meta')] = G4GCNConv(*[hidden_channels]*3)
                self.lins_m.append(Linear(hidden_channels*2, hidden_channels))
                self.lins_c.append(Linear(hidden_channels*4, hidden_channels))
                self.lins_v.append(Linear(hidden_channels*3, hidden_channels))
            else:
                self.lins_c.append(Linear(hidden_channels*3, hidden_channels)) # set up linear layers to compute updates to node features
                self.lins_v.append(Linear(hidden_channels*3, hidden_channels))

            conv = HeteroConv(conv_dict, aggr='cat')
            self.convs.append(conv)

    def forward(self, x_dict, deg_dict, edge_index_dict):
        if self.include_meta_node:
            x_dict['meta'] = self.lin0_m(x_dict['meta'])
        x_dict['clause'] = self.lin0_c(x_dict['clause'])   # Lift variable and clause features to latend representations space
        x_dict['variable'] = self.lin0_v(x_dict['variable'])
        for conv_layer, conv in enumerate(self.convs): # apply convolutions to compute aggregated messages from neighbours
            x_dict_prev = x_dict # store previous node represntation for update
            x_dict = conv(x_dict, deg_dict, edge_index_dict)
            if self.include_meta_node:
                x_dict['meta'] = torch.cat([x_dict['meta'], x_dict_prev['meta']], dim=1)
                x_dict['meta'] = self.lins_m[conv_layer](x_dict['meta'])
            x_dict['clause'] = torch.cat([x_dict['clause'], x_dict_prev['clause']], dim=1) # concatenate aggregated messges and previous node feature
            x_dict['variable'] = torch.cat([x_dict['variable'], x_dict_prev['variable']], dim=1)
            x_dict['clause'] = self.lins_c[conv_layer](x_dict['clause']) # Update node features based on messaged and previous node feature
            x_dict['variable'] = self.lins_v[conv_layer](x_dict['variable'])
        out = x_dict['variable']
        return out

class G4GCNConv(MessagePassing):
    # Must manually ensure that models have the same parameters when initialized
    def __init__(self, src_channels, tgt_channels, out_channels):
        super().__init__(aggr='add')  # "Add" aggregation (Step 5).
        self.lin_src = Linear(src_channels, out_channels, bias=True)
        self.lin_src.reset_parameters()
        self.MLP = MLP([out_channels, out_channels, out_channels, out_channels])

    def forward(self, x, deg, edge_index):
        # Linearly transform node feature matrix.
        x_src, x_trg = x
        x_src = self.lin_src(x_src).relu()

        # Compute normalization.
        src, trg = edge_index
        deg_src, deg_trg = deg
        deg_src_inv_sqrt = deg_src.pow(-0.5)
        deg_src_inv_sqrt[deg_src_inv_sqrt == float('inf')] = 0
        deg_trg_inv_sqrt = deg_trg.pow(-0.5)
        deg_trg_inv_sqrt[deg_trg_inv_sqrt == float('inf')] = 0
        norm = deg_src_inv_sqrt[src] * deg_trg_inv_sqrt[trg]

        # Start propagating messages.
        out = self.propagate(edge_index, x=x, norm=norm)
        return out

    def message(self, x_i, x_j, norm):
        # x_j has shape [E, out_channels]

        # Normalize node features.
        return norm.view(-1, 1) * self.MLP(x_j).relu()


#endregion
class MLP(nn.Module):
    def __init__(self, layer_sizes):
        """
        layer_sizes is a list of integers with the input size as the 
        first element and the output size as the last element
        """
        super().__init__()
        self.layer_sizes = layer_sizes
        layers = nn.ModuleList()
        for layer, layer_size in enumerate(layer_sizes[:-1]):
            layers.append(Linear(layer_sizes[layer], layer_sizes[layer+1], bias=True))
            if layer != len(layer_sizes[:-1]) - 1:
                layers.append(nn.ReLU())
        self.linear_layers = nn.Sequential(*layers)
    
    def forward(self, x):
        x = self.linear_layers(x)
        return x

[akihironitta]

Have you had a chance to check which device model and batched_graphs are on?

        out = model(batched_graphs.x_dict, batched_graphs.deg_dict, batched_graphs.edge_index_dict)

You could try to debug it by setting a breakpoint between these lines and seeing whether x_dict['variable'] is on the correct device:

            x_dict['clause'] = self.lins_c[conv_layer](x_dict['clause']) # Update node features based on messaged and previous node feature
+           breakpoint()
            x_dict['variable'] = self.lins_v[conv_layer](x_dict['variable'])

If you could provide me with a complete script that I can simply run, I may be able to point out what the issue is :)

[Angelo92Git]

Thank you very much, it's resolved now!