pytorch 实现GCN的MPNN框架 gcn代码pytorch讲解

转载

游侠小影 2023-09-27 16:46:03

文章标签 where解析顺序邻接矩阵数组 ci 文章分类 PyTorch 人工智能

本文为gcn的PyTorch版本pygcn代码的注释解析（代码地址），也作为学习PyTorch时的一个实例加深对PyTorch API的理解。

模型代码一般分为下面几个关键步骤：

数据预处理
搭建模型
定义损失函数
训练与测试

其中代码量最大的是前两步，数据预处理包括如何从文件中读取数据，并存储成深度学习框架可处理的tensor类型，构建训练集、测试集和验证集等；搭建模型则是核心，需要对模型内部的运算流程有详细的认识，还包含很多论文中没有提到的细节。由于PyTorch相比Tensorflow更为简便，所以在模型和损失函数的定义方面，代码量要简洁，且易读。

数据预处理

对应于与pygcn/train.py文件中的：

adj, features, labels, idx_train, idx_val, idx_test = load_data()

所以关键在于pygcn/utils.py中的load_data()方法：

def load_data(path="../data/cora/", dataset="cora"):
    """Load citation network dataset (cora only for now)"""
    print('Loading {} dataset...'.format(dataset))

    idx_features_labels = np.genfromtxt("{}{}.content".format(path, dataset),
                                        dtype=np.dtype(str))
    features = sp.csr_matrix(idx_features_labels[:, 1:-1], dtype=np.float32)
    labels = encode_onehot(idx_features_labels[:, -1])

    # build graph
    idx = np.array(idx_features_labels[:, 0], dtype=np.int32)
    idx_map = {j: i for i, j in enumerate(idx)}
    edges_unordered = np.genfromtxt("{}{}.cites".format(path, dataset),
                                    dtype=np.int32)
    edges = np.array(list(map(idx_map.get, edges_unordered.flatten())),
                     dtype=np.int32).reshape(edges_unordered.shape)
    adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])),
                        shape=(labels.shape[0], labels.shape[0]),
                        dtype=np.float32)

    # build symmetric adjacency matrix
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)

    features = normalize(features)
    adj = normalize(adj + sp.eye(adj.shape[0]))

    idx_train = range(140)
    idx_val = range(200, 500)
    idx_test = range(500, 1500)

    features = torch.FloatTensor(np.array(features.todense()))
    labels = torch.LongTensor(np.where(labels)[1])
    adj = sparse_mx_to_torch_sparse_tensor(adj)

    idx_train = torch.LongTensor(idx_train)
    idx_val = torch.LongTensor(idx_val)
    idx_test = torch.LongTensor(idx_test)

    return adj, features, labels, idx_train, idx_val, idx_test

load_data()的参数是用来拼接数据集的字符串。

第5-6行：idx_features_labels用来存储读取data/cora/cora.content文件中的内容。这里涉及numpy中的方法genfromtxt()，关于该方法如何使用可参考[1,2]，可以快速将txt中保存的文本内容转换成numpy中的数组。文件数据的格式为id features labels，因此idx_features_labels[:, 0]、idx_features_labels[:, 1:-1]、idx_features_labels[:, -1]分别代表以上三部分。

第7行：idx_features_labels[:, 1:-1]表示节点本身附带的特征，由于这些特征构成的矩阵是稀疏矩阵，使用csr_matrix()方法将其存储为稀疏矩阵。关于该方法的使用参考[3,4]。

第8行：labels在文件中是用字符串表示的某个论文所属的类别，如“reinforce learning”，要将其表示成one-hot向量的形式，用到了encode_onehot()的方法。具体来看：

def encode_onehot(labels):
    classes = set(labels)
    classes_dict = {c: np.identity(len(classes))[i, :] for i, c in
                    enumerate(classes)}
    labels_onehot = np.array(list(map(classes_dict.get, labels)),
                             dtype=np.int32)
    return labels_onehot

先将所有由字符串表示的标签数组用set保存，set的重要特征就是元素没有重复，因此表示成set后可以直接得到所有标签的总数，随后为每个标签分配一个编号，创建一个单位矩阵，单位矩阵的每一行对应一个one-hot向量，也就是np.identity(len(classes))[i, :]，再将每个数据对应的标签表示成的one-hot向量，类型为numpy数组。

第11行：将所有节点的id表示成numpy数组。

第12行：由于文件中节点并非是按顺序排列的，因此建立一个编号为0-(node_size-1)的哈希表idx_map，哈希表中每一项为id: number，即节点id对应的编号为number。

第13-14行：edges_unordered为直接从边表文件中直接读取的结果，是一个(edge_num, 2)的数组，每一行表示一条边两个端点的idx。

第15-16行：上边的edges_unordered中存储的是端点id，要将每一项的id换成编号。在idx_map中以idx作为键查找得到对应节点的编号，reshape成与edges_unordered形状一样的数组。

第17-19行：首先要明确coo_matrix()的作用，该方法是构建一个矩阵，根据给出的下标、数据和形状，构建一个矩阵，其中下标位置的值对应数据中的值，使用方法见[5]。所以这一段的作用就是，网络有多少条边，邻接矩阵就有多少个1，所以先创建一个长度为edge_num的全1数组，每个1的填充位置就是一条边中两个端点的编号，即edges[:, 0], edges[:, 1]，矩阵的形状为(node_size, node_size)。

第22行：对于无向图，邻接矩阵是对称的。上一步得到的adj是按有向图构建的，转换成无向图的邻接矩阵需要扩充成对称矩阵。

第24-25行：分别对特征矩阵features和邻接矩阵adj做标准化，用到了normalize()方法：

def normalize(mx):
    """Row-normalize sparse matrix"""
    rowsum = np.array(mx.sum(1))
    r_inv = np.power(rowsum, -1).flatten()
    r_inv[np.isinf(r_inv)] = 0.
    r_mat_inv = sp.diags(r_inv)
    mx = r_mat_inv.dot(mx)
    return mx

首先对每一行求和得到rowsum；求倒数得到r_inv；如果某一行全为0，则r_inv算出来会等于无穷大，将这些行的r_inv置为0；构建对角元素为r_inv的对角矩阵；用对角矩阵与原始矩阵的点积起到标准化的作用，原始矩阵中每一行元素都会与对应的r_inv相乘。

第27-37行：分别构建训练集、验证集、测试集，并创建特征矩阵、标签向量和邻接矩阵的tensor，用来做模型的输入。

搭建模型

模型部分对应于pygcn/train.py中的：

model = GCN(nfeat=features.shape[1],
            nhid=args.hidden,
            nclass=labels.max().item() + 1,
            dropout=args.dropout)

模型定义在了pygcn/model.py中，下面看看模型是怎么定义的：

class GCN(nn.Module):
    def __init__(self, nfeat, nhid, nclass, dropout):
        super(GCN, self).__init__()

        self.gc1 = GraphConvolution(nfeat, nhid)
        self.gc2 = GraphConvolution(nhid, nclass)
        self.dropout = dropout

    def forward(self, x, adj):
        x = F.relu(self.gc1(x, adj))
        x = F.dropout(x, self.dropout, training=self.training)
        x = self.gc2(x, adj)
        return F.log_softmax(x, dim=1)

模型的定义看起来十分简洁，如果熟悉GCN的公式的话都不要注释。GCN由两个GraphConvolution层构成。GCN的forward(self, x, adj)方法对应输入分别是特征和邻接矩阵。最后输出为输出层做log_softmax变换的结果。GraphConvolution层定义如下：

class GraphConvolution(Module):
    """
    Simple GCN layer, similar to https://arxiv.org/abs/1609.02907
    """

    def __init__(self, in_features, out_features, bias=True):
        super(GraphConvolution, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
        if bias:
            self.bias = Parameter(torch.FloatTensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, input, adj):
        support = torch.mm(input, self.weight)
        output = torch.spmm(adj, support)
        if self.bias is not None:
            return output + self.bias
        else:
            return output

    def __repr__(self):
        return self.__class__.__name__ + ' (' 
               + str(self.in_features) + ' -> ' 
               + str(self.out_features) + ')'

__init__部分初始化了一些用到的参数，包括输入和输出的维度，并初始化了每一层的权重。

forward()方法说明了每一层对数据的操作，其实就是先将输入特征矩阵与权重矩阵相乘，在左乘标准化的邻接矩阵，由于邻接矩阵的存储时用的是稀疏矩阵，所以有别于上一行，torch.spmm表示sparse_tensor与dense_tensor相乘。

定义损失函数

GCN的损失函数分为两部分，一部分是分类损失，一部分是权重的正则化。

定义优化器（pygcn/train.py）：

optimizer = optim.Adam(model.parameters(),
                       lr=args.lr, weight_decay=args.weight_decay)

weight_decay表示正则化的系数，默认取的是5e-4。所以在定义优化器的时候就已将将正则化的损失定义了，不要在后续看代码的过程中发出找不到定义权重正则化的代码。

loss_train = F.nll_loss(output[idx_train], labels[idx_train])

另一部分是分类损失，用的是交叉熵，由于在算output时已经使用了log_softmax，这里使用的损失函数就是NLLloss，如果前面没有加log运算，这里就要使用CrossEntropyLoss`了，有关两者之间的差别，见参考[6]。

训练与测试

# train
def train(epoch):
    t = time.time()
    model.train()
    optimizer.zero_grad()
    output = model(features, adj)
    loss_train = F.nll_loss(output[idx_train], labels[idx_train])
    acc_train = accuracy(output[idx_train], labels[idx_train])
    loss_train.backward()
    optimizer.step()

    # val
    loss_val = F.nll_loss(output[idx_val], labels[idx_val])
    acc_val = accuracy(output[idx_val], labels[idx_val])
    print('Epoch: {:04d}'.format(epoch+1),
          'loss_train: {:.4f}'.format(loss_train.item()),
          'acc_train: {:.4f}'.format(acc_train.item()),
          'loss_val: {:.4f}'.format(loss_val.item()),
          'acc_val: {:.4f}'.format(acc_val.item()),
          'time: {:.4f}s'.format(time.time() - t))
# test    
def test():
    model.eval（)
    output = model(features, adj)
    loss_test = F.nll_loss(output[idx_test], labels[idx_test])
    acc_test = accuracy(output[idx_test], labels[idx_test])
    print("Test set results:",
          "loss= {:.4f}".format(loss_test.item()),
          "accuracy= {:.4f}".format(acc_test.item()))

训练与测试则比较常规，注意一下PyTorch的顺序，先将model置为训练状态；梯度清零；将输入送到模型得到输出结果；计算损失与准确率；反向传播求梯度更新参数。不得不说，见识过了Tensorflow后真的对PyTorch这种极其简介的定义颇为好感，api简单易懂。

本文章为转载内容，我们尊重原作者对文章享有的著作权。如有内容错误或侵权问题，欢迎原作者联系我们进行内容更正或删除文章。