Edge Classification

[sciencewu]

Hi everyone! I wanted to do an edge classification model as I saw nothing online. I have some problems with it though, and Im not sure if the idea is correct. Any help would be appreciated.

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class Encoder(torch.nn.Module):
    def __init__(self, in_channel, hidden_dim):
        super(Encoder, self).__init__()
        self.conv1 = GCNConv(in_channel, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, hidden_dim)

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

class Decoder(torch.nn.Module):
    def __init__(self, in_channels, hidden_dim, out_channels):
        super(Decoder, self).__init__()
        self.fc1 = torch.nn.Linear(in_channels, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, out_channels)

    def forward(self, x, edge_index):
        x = torch.cat([x[edge_index[0]], x[edge_index[1]]], dim=-1)
        print(x.shape)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

class EdgeClassifier(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(EdgeClassifier, self).__init__()
        self.encoder = Encoder(in_channels, 32)
        self.decoder = Decoder(32 * 2, 64, out_channels)

    def forward(self, x, edge_index):
        x = self.encoder(x, edge_index)
        x = self.decoder(x, edge_index)
        return x

and this is the training process and a simple example

import torch_geometric.data as data

# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load your graph data into a torch_geometric.data object
graph_data = data.Data(
                    x=torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]], dtype=torch.float), 
                    edge_index=torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long), 
                    edge_labels=torch.tensor([0, 1, 0, 1], dtype=torch.float)
                    )

model = EdgeClassifier(in_channels=3, out_channels=2)#.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# # Define the loss function
criterion = torch.nn.CrossEntropyLoss()

graph_data = graph_data#.to(device)

# # Train the model
for epoch in range(200):
    optimizer.zero_grad()
    output = model(graph_data.x, graph_data.edge_index)
    print(output)
    loss = criterion(output, graph_data.edge_labels)
    loss.backward()
    optimizer.step()

# # Test the model
# model.eval()
# with torch.no_grad():
#     output = model(graph_data.x, graph_data.edge_index)
#     pred = output.argmax(dim=1)
#     correct = pred.eq(graph_data.edge_labels).sum().item()
#     accuracy = correct / graph_data.edge_labels.size(0)
#     print('Accuracy: {:.4f}'.format(accuracy))

I do get the RuntimeError: expected scalar type Long but found Float error, but I think the dtype is fine? Also is this the correct approach to solve this? Thank you!

[rusty1s]

edge_labels need to be of dtype=torch.long.

[sciencewu]

Thank you, and I'm sorry for this trivial question. I thought I tried also that combination.