Pytorch 几何图批处理不使用 DataLoader 进行强化学习

Question

我对使用图神经网络（GNN）还很陌生。我正在使用 PyTorch 几何。我正在创建一个强化学习算法，因此我希望在生成数据/观察时避免使用内置的 DataLoader。但是，我在传递一批 PyTorch 几何图时遇到了问题。我有一个带有 PyG 图的 numpy 内存数组。我从这段记忆中挑选并尝试将其通过神经网络（NN）。

通过神经网络推送单个图似乎效果很好。我得到每个节点的表示。然而，当使用批处理时，就会出现问题。通常，我可以创建 numpy_array 批次的张量。但是，PyTorch 无法执行此操作，因为它无法处理 PyG 数据类型。因此，我使用 PyTorch Geometric 的内置功能创建一个批次。它通过神经网络；然而，输出维度似乎很奇怪。看起来这些图被组合成单个对象，然后作为单个图传递。然而，我期望的输出是

[batch_size, n_nodes]

而不是

[batch_size * n_nodes]

。我想知道我这样做是否正确。有没有更好的方法来处理这个问题以避免维度问题？我不相信我可以将输出数组拆分为数组中的每个 n_nodes。

一种选择是使用 for 循环将每个单独的图推入前向传递，但这非常低效。也许我缺少一个简单的设置？我已经提供了一个工作示例。

提前致谢。

import torch as T
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.nn import GCNConv
from torch_geometric.data import Batch
from torch_geometric.data import Data
import numpy as np


class DeepNetworkGCN(nn.Module):
    def __init__(self, lr=0.001, input_dims=[1], fc1_dims=128, fc2_dims=128, out_dims=[1]):
        super(DeepNetworkGCN, self).__init__()

        # CNN part of network
        self.GCNconv1 = GCNConv(*input_dims, fc1_dims)
        self.GCNconv2 = GCNConv(fc1_dims, fc2_dims)

        # conform to output dimension
        self.fc1 = nn.Linear(fc2_dims, *out_dims)

        self.optimizer = optim.Adam(self.parameters(), lr=lr)
        self.loss = nn.MSELoss()
        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        # Process graph data using GCN layers
        x = self.GCNconv1(state.x, state.edge_index)
        x = F.relu(x)
        x = self.GCNconv2(x, state.edge_index)

        # Final fully connected layer
        out = self.fc1(x)

        return out


def random_pyg_graph(num_nodes=3):  
    # random node features
    node_features = T.randint(0, 5, (num_nodes, 1), dtype=T.float)

    # random edge features
    edge_features = T.randn(num_nodes, num_nodes)

    # random edge indices
    edge_index = T.randint(0, num_nodes, (2, num_nodes * 2))

    # Remove self-loops
    edge_index = edge_index[:, edge_index[0] != edge_index[1]]

    # graph
    graph_data = Data(x=node_features, edge_index=edge_index, edge_attr=edge_features)

    return graph_data


# setup example
batch_size = 3
memory = np.zeros(batch_size, dtype=object)

# fill memory
for i in range(batch_size):
    memory[i] = random_pyg_graph()

# define model
CNN = DeepNetworkGCN()

# test for single PyG
output = CNN.forward(memory[0])
print(output)
# output 1 for each node e.g.
# tensor([[0.3770],
#        [0.6119],
#        [0.2014]], grad_fn=<AddmmBackward0>)

# test for numpy.ndarray
# FAILS! # FAILS! # FAILS!
# output = CNN.forward(memory[:]) # FAILS!
# FAILS! # FAILS! # FAILS!

# Create batch and do forward pass.
output = CNN.forward(Batch.from_data_list(memory[:]))
print(output)
# output dimension is weird. ( n_nodes*batch_size).
# tensor([[ 0.0173],
#         [ 0.0316],
#         [ 0.0282],
#         [ 0.0147],
#         [-0.0201],
#         [-0.0264],
#         [ 0.0147],
#         [-0.0084],
#         [ 0.0021]], grad_fn=<AddmmBackward0>)

Answer 1

在 PyTorch 中，如果您要动态生成数据集，通常的做法是使用 IterableDataset。它只需要一个

__iter__

而不是

__len__

和

__getitem__

。

示例取自此处：

class DataStream1(IterableDataset):

    def __init__(self) -> None:
        super().__init__()
        self.size_input = 4
        self.size_output = 2

    def generate(self):
        while True:
            x = torch.rand(self.size_input)
            y = torch.rand(self.size_output)
            yield x, y

    def __iter__(self):
        return iter(self.generate())

dataset = DataStream1()

train_loader = DataLoader(dataset=dataset)

for i, data in enumerate(train_loader):
    print (i, data)