Как правильно реализовать NNConv в Pytorch?

Вопрос или проблема

x torch.Size([39, 5])
edge_attr torch.Size([44, 3])
edge_index torch.Size([2, 88])
batch_index torch.Size([39])

model = GNN(test[0].x.shape[1], test[0].edge_attr.shape[1], test[0].edge_index) 
model = model.to(device)
pred = model(batch.x.float(),  batch.edge_attr.float(), batch.edge_index,  batch.batch)

44 x = self.conv1(x, edge_index, edge_attr)
1501 return forward_call(*args, **kwargs)
102 out = self.propagate(edge_index, x=x, edge_attr=edge_attr, size=size)
469 out = self.message(**msg_kwargs)
115 weight = weight.view(-1, self.in_channels_l, self.out_channels)

1. RuntimeError: shape ‘[-1, 5, 32]’ is invalid for input of size 1408

x torch.Size([24, 5])
edge_attr torch.Size([25, 3])
edge_index torch.Size([2, 50])
batch_index torch.Size([24])

model = GNN(test[0].x.shape[1], test[0].edge_attr.shape[1], test[0].edge_index) 
model = model.to(device)
pred = model(batch.x.float(),  batch.edge_attr.float(), batch.edge_index,  batch.batch)

RuntimeError: The size of tensor a (50) must match the size of tensor b (5) at non-singleton dimension 0

class GNN(torch.nn.Module):
    def __init__(self, feature_size,edge_feat,edge_index):
        super(GNN, self).__init__()
        num_classes = 2
        embedding_size = 32
        n_out = 16

        nn1 = nn.Sequential(nn.Linear(edge_feat,32), nn.ReLU())
        nn2 = nn.Sequential(nn.Linear(edge_feat,32), nn.ReLU())

        # Уровни GNN

        self.conv1 = NNConv(feature_size, 32, nn1, aggr="mean")
        self.pool1 = TopKPooling(32, ratio=0.8)

        self.conv2 = NNConv(32, 64, nn2, aggr="mean")
        self.pool2 = TopKPooling(32, ratio=0.8)

        # Линейные уровни
        self.linear1 = Linear(32, 16)
        self.linear2 = Linear(16, num_classes)  

    def forward(self, x, edge_attr, edge_index, batch_index):

        # Первый блок
        x = self.conv1(x, edge_index, edge_attr)
        x, edge_index, edge_attr, batch_index, _, _ = self.pool1(x, 
                                                        edge_index, 
                                                        None, 
                                                        batch_index)
        x1 = torch.cat([gmp(x, batch_index), gap(x, batch_index)], dim=1)

        # Второй блок
        x = self.conv2(x, edge_index)
        x, edge_index, edge_attr, batch_index, _, _ = self.pool2(x, 
                                                        edge_index, 
                                                        None, 
                                                        batch_index)
        x2 = torch.cat([gmp(x, batch_index), gap(x, batch_index)], dim=1)

        # Конкатенация объединенных векторов
        x = x1 + x2

        # Блок вывода
        x = self.linear1(x).relu()
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.linear2(x)

        return x

Ответ или решение

Чтобы корректно реализовать NNConv в PyTorch, важно учитывать несколько ключевых аспектов, связанных с архитектурой графовой нейронной сети (GNN) и спецификой библиотеки PyTorch Geometric. Ниже будет представлена детализированная инструкция по исправлению проблем в вашем коде и правильной реализации NNConv.

Основные ошибки и их решение

1. Неверная форма выходных данных nn1 и nn2

В процессе использования NNConv в PyTorch Geometric критически важно, чтобы выходные данные ваших последовательностей nn1 и nn2 соответствовали ожиданиям NNConv. Согласно документации от PyTorch Geometric, выходные данные nn1 и nn2 должны иметь форму (num_edges, in_channels * out_channels).

В вашем коде вы определили сети nn1 и nn2 следующим образом:

nn1 = nn.Sequential(nn.Linear(edge_feat, 32), nn.ReLU())
nn2 = nn.Sequential(nn.Linear(edge_feat, 32), nn.ReLU())

Однако их выходы не соответствуют требованиям NNConv. Чтобы исправить это, измените последний слой каждой из последовательностей. Например, если вы хотите использовать 32 входных канала и 32 выходных канала, вы можете сделать следующее:

nn1 = nn.Sequential(nn.Linear(edge_feat, 32), nn.ReLU(), nn.Linear(32, 32*32))
nn2 = nn.Sequential(nn.Linear(edge_feat, 32), nn.ReLU(), nn.Linear(32, 32*64))

Это обеспечит, что выходные данные будут соответствовать размерности (num_edges, in_channels * out_channels).

2. Ошибки при передаче размеров тензоров

У вас возникают ошибки, связанные с несовпадением размеров тензоров. Например, ошибка RuntimeError: shape ‘[-1, 5, 32]’ is invalid for input of size 1408 говорит о том, что в процессе выполнения операции reshaping ожидаются данные другой формы.

Убедитесь, что во всех шагах передачи данных в нейронные сети используется правильная размерность. Проверьте, что входы x и edge_attr правильно передаются через все слои и при необходимости пересчитывайте их размерности.

3. Корректная обработка пакетов графов

Работа с батчами графов требует аккуратной обработки индексов. Ваша функция forward должна правильно обрабатывать batch_index, чтобы делающие обработку значения из разных графов не перепутались.

Пример исправленного кода

Ниже приведён пример исправленного кода для класса GNN с учетом вышеизложенных рекомендаций:

class GNN(torch.nn.Module):
    def __init__(self, feature_size, edge_feat):
        super(GNN, self).__init__()
        num_classes = 2
        self.embedding_size1 = 32
        self.embedding_size2 = 64

        nn1 = nn.Sequential(nn.Linear(edge_feat, 32), nn.ReLU(), nn.Linear(32, self.embedding_size1 * self.embedding_size2))
        nn2 = nn.Sequential(nn.Linear(edge_feat, 32), nn.ReLU(), nn.Linear(32, self.embedding_size2 * num_classes))

        self.conv1 = NNConv(feature_size, self.embedding_size1, nn1, aggr="mean")
        self.pool1 = TopKPooling(self.embedding_size1, ratio=0.8)

        self.conv2 = NNConv(self.embedding_size1, self.embedding_size2, nn2, aggr="mean")
        self.pool2 = TopKPooling(self.embedding_size2, ratio=0.8)

        self.linear1 = Linear(self.embedding_size1 + self.embedding_size2, 16)
        self.linear2 = Linear(16, num_classes)

    def forward(self, x, edge_attr, edge_index, batch_index):
        x = self.conv1(x, edge_index, edge_attr)
        x, edge_index, edge_attr, batch_index, _, _ = self.pool1(x, edge_index, None, batch_index)

        x1 = torch.cat([gmp(x, batch_index), gap(x, batch_index)], dim=1)

        x = self.conv2(x, edge_index)
        x, edge_index, edge_attr, batch_index, _, _ = self.pool2(x, edge_index, None, batch_index)

        x2 = torch.cat([gmp(x, batch_index), gap(x, batch_index)], dim=1)

        x = torch.cat([x1, x2], dim=1)  # Конкатенация pooled векторов
        x = self.linear1(x).relu()
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.linear2(x)

        return x

Заключение

Эти корректировки должны помочь устранить ошибки при выполнении и позволить вам успешно использовать NNConv в построении графовых нейронных сетей. Следуйте указаниям выше, чтобы адаптировать вашу архитектуру и учитывать размерности тензоров в PyTorch. Успехов в разработке вашего проекта!