簡介

在之前的文章中，我主要提及了節點級別的任務，這種任務的特點是不斷對節點的特征進行聚合，每次聚合各個節點都會聚集鄰居的信息，但是，圖的結構是不會變化的。圖分類則是一個圖層面的任務，與節點層面的任務不同，它需要的是圖數據的全局信息，既包括圖的結構信息，也包括每個節點的屬性信息。圖分類任務說來也簡單，給定多張圖及其標簽，要求學習一個由圖到相應標簽的圖分類模型，模型的重點是如何學習一個較好的全圖表示向量。

SAGPool算法

圖分類任務和視覺圖像的分類任務類似，都需要對全局的信息進行融合學習。不妨回顧一下CNN中的處理思路，在卷積神經網絡中，我們通常利用層次化池化來逐漸提取全局信息。得益于圖像的柵格結構，池化操作能夠非常簡單高效地的實現并用于高階信息的提取。然而，非規則的圖結構數據使得池化的設計變得比較困難。

SAGPool（Self-Attention Pooling，自注意力池化）是層次化池化的一種實現，它的思路是通過圖卷積從圖中自適應學習到節點的重要性，然后利用TopK機制進行節點丟棄。

具體來看，使用GCN為每個節點賦予重要性得分，如下式所示。其中$\sigma$表示激活函數，$\tilde{A}$表示增加了自連接的鄰接矩陣，$X$表示節點特征，${\Theta }_{att}\in R^{N\times 1}$表示權重參數，這也是自注意力池化層引入的唯一參數。

$$Z=\sigma \left({\tilde{D}}^{?\frac{1}{2}}\tilde{A}{\tilde{D}}^{?\frac{1}{2}}X{\Theta }_{att}\right)$$

根據重要性得分以及圖的拓撲結構可以進行池化操作，下式就是TopK選擇的過程，基于$Z$得分，只會有$?kN?$個節點繼續保留。

$$\mathrm{i}\mathrm{d}\mathrm{x}=\mathrm{t}\mathrm{o}\mathrm{p}-\mathrm{r}\mathrm{a}\mathrm{n}\mathrm{k}(Z,?kN?),Z_{\text{mask}}=Z_{\text{idx?}}$$

就這樣，反復堆疊上述的自注意力池化層就可以進行圖的池化，最后通過全局池化將各個圖降維到同一維度即可。對算法細節感興趣可以查看原論文。

數據集

本項目采用的是D&D數據集，這是一個包含1178個蛋白質結構的數據集，每個蛋白質用圖結構表示，圖中的節點為氨基酸，如果兩個節點之間的距離小于6埃則它們之間有邊相連，每個圖分為酶和非酶兩種類別之一。該數據集源于論文《Distinguishing enzyme structures from non-enzymes without alignments》，下載地址為BaiduNetDisk(code: zczc)。

模型構建

在開始之前我們先來安裝一個包torch-scatter，它是torch的一個拓展包，可以很方便地如下圖般按照索引類來分別進行均值和最大值求解，安裝可以通過pip安裝，也可以通過下面的地址找到合適的版本安裝。

https://s3.eu-central-1.amazonaws.com/pytorch-geometric.com/whl/torch-1.6.0.html

然后我們來構建模型，核心的TopK-rank等方法就不多贅述，可以查看文末的Github鏈接，下面是自注意力池化的具體實現，思路就是我上面說的利用GCN來進行節點評分。

class SelfAttentionPooling(nn.Module):def __init__(self, input_dim, keep_ratio, activation=torch.tanh):super(SelfAttentionPooling, self).__init__()self.input_dim = input_dimself.keep_ratio = keep_ratioself.activation = activationself.attn_gcn = GraphConvolution(input_dim, 1)def forward(self, adjacency, input_feature, graph_indicator):attn_score = self.attn_gcn(adjacency, input_feature).squeeze()attn_score = self.activation(attn_score)mask = top_rank(attn_score, graph_indicator, self.keep_ratio)hidden = input_feature[mask] * attn_score[mask].view(-1, 1)mask_graph_indicator = graph_indicator[mask]mask_adjacency = filter_adjacency(adjacency, mask)return hidden, mask_graph_indicator, mask_adjacency

有了自注意力池化層，就可以定義模型了，論文提出了兩種模型設計如下圖，左邊的模型是global模型，右邊的模型是hierarchical模型，實驗證明前者適合小圖分類，后者適合大圖分類。

class ModelA(nn.Module):def __init__(self, input_dim, hidden_dim, num_classes=2):"""SAGPool Global Model:param input_dim::param hidden_dim::param num_classes:"""super(ModelA, self).__init__()self.input_dim = input_dimself.hidden_dim = hidden_dimself.num_classes = num_classesself.gcn1 = GraphConvolution(input_dim, hidden_dim)self.gcn2 = GraphConvolution(hidden_dim, hidden_dim)self.gcn3 = GraphConvolution(hidden_dim, hidden_dim)self.pool = SelfAttentionPooling(hidden_dim * 3, 0.5)self.fc1 = nn.Linear(hidden_dim * 3 * 2, hidden_dim)self.fc2 = nn.Linear(hidden_dim, hidden_dim // 2)self.fc3 = nn.Linear(hidden_dim // 2, num_classes)def forward(self, adjacency, input_feature, graph_indicator):gcn1 = F.relu(self.gcn1(adjacency, input_feature))gcn2 = F.relu(self.gcn2(adjacency, gcn1))gcn3 = F.relu(self.gcn3(adjacency, gcn2))gcn_feature = torch.cat((gcn1, gcn2, gcn3), dim=1)pool, pool_graph_indicator, pool_adjacency = self.pool(adjacency, gcn_feature,graph_indicator)readout = torch.cat((global_avg_pool(pool, pool_graph_indicator),global_max_pool(pool, pool_graph_indicator)), dim=1)fc1 = F.relu(self.fc1(readout))fc2 = F.relu(self.fc2(fc1))logits = self.fc3(fc2)return logitsclass ModelB(nn.Module):def __init__(self, input_dim, hidden_dim, num_classes=2):"""SAGPool Hierarchical Model:param input_dim::param hidden_dim::param num_classes:"""super(ModelB, self).__init__()self.input_dim = input_dimself.hidden_dim = hidden_dimself.num_classes = num_classesself.gcn1 = GraphConvolution(input_dim, hidden_dim)self.pool1 = SelfAttentionPooling(hidden_dim, 0.5)self.gcn2 = GraphConvolution(hidden_dim, hidden_dim)self.pool2 = SelfAttentionPooling(hidden_dim, 0.5)self.gcn3 = GraphConvolution(hidden_dim, hidden_dim)self.pool3 = SelfAttentionPooling(hidden_dim, 0.5)self.mlp = nn.Sequential(nn.Linear(hidden_dim * 2, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim // 2),nn.ReLU(),nn.Linear(hidden_dim // 2, num_classes))def forward(self, adjacency, input_feature, graph_indicator):gcn1 = F.relu(self.gcn1(adjacency, input_feature))pool1, pool1_graph_indicator, pool1_adjacency = \self.pool1(adjacency, gcn1, graph_indicator)global_pool1 = torch.cat([global_avg_pool(pool1, pool1_graph_indicator),global_max_pool(pool1, pool1_graph_indicator)],dim=1)gcn2 = F.relu(self.gcn2(pool1_adjacency, pool1))pool2, pool2_graph_indicator, pool2_adjacency = \self.pool2(pool1_adjacency, gcn2, pool1_graph_indicator)global_pool2 = torch.cat([global_avg_pool(pool2, pool2_graph_indicator),global_max_pool(pool2, pool2_graph_indicator)],dim=1)gcn3 = F.relu(self.gcn3(pool2_adjacency, pool2))pool3, pool3_graph_indicator, pool3_adjacency = \self.pool3(pool2_adjacency, gcn3, pool2_graph_indicator)global_pool3 = torch.cat([global_avg_pool(pool3, pool3_graph_indicator),global_max_pool(pool3, pool3_graph_indicator)],dim=1)readout = global_pool1 + global_pool2 + global_pool3logits = self.mlp(readout)return logits

模型訓練

下面使用兩種模型分別在數據集上進行訓練和推理，global模型測試集準確率0.75，hierarchical模型測試集準確率0.72。下圖是兩個模型的實驗對比圖，設計的代碼見文末Github。

補充說明

本文簡單實現D&D數據集上實驗自注意力池化進行圖分類任務，內容參考《深入淺出圖神經網絡》以及SAGPool的論文，代碼開放于Github，歡迎star和fork。

總結

以上是生活随笔為你收集整理的SAGPool图分类的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。