基于PyTorch如何實現EdgeCNN

這篇文章主要講解了“基于PyTorch如何實現EdgeCNN”，文中的講解內容簡單清晰，易于學習與理解，下面請大家跟著小編的思路慢慢深入，一起來研究和學習“基于PyTorch如何實現EdgeCNN”吧！

一、導入相關庫

本項目是采用自己實現的EdgeCNN，并沒有使用 PyG 庫，原因是為了幫助新手朋友們能夠對EdgeConv的原理有個更深刻的理解，如果熟悉之后可以嘗試使用PyG庫直接調用 EdgeConv 這個圖層即可。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.utils import scatter
from torch_geometric.datasets import Planetoid

二、加載Cora數據集

本文使用的數據集是比較經典的Cora數據集，它是一個根據科學論文之間相互引用關系而構建的Graph數據集合，論文分為7類，共2708篇。

• Genetic_Algorithms

• Neural_Networks

• Probabilistic_Methods

• Reinforcement_Learning

• Rule_Learning

• Theory

這個數據集是一個用于圖節點分類的任務，數據集中只有一張圖，這張圖中含有2708個節點，10556條邊，每個節點的特征維度為1433。

# 1.加載Cora數據集
dataset = Planetoid(root='./data/Cora', name='Cora')

三、定義EdgeCNN網絡

3.1 定義EdgeConv層

這里我們就不重點介紹EdgeCNN網絡了，相信大家能夠掌握基本原理，本文我們使用的是PyTorch定義網絡層。

對于EdgeConv的常用參數：

• nn：進行節點特征轉換使用的 MLP網絡，需要自己定義傳入

• aggr：聚合鄰居節點特征時采用的方式，默認為 max

我們在實現時也是考慮這幾個常見參數

對于EdgeConv的傳播公式為：

上式子中的 x i x_i xi 代表中心節點特征信息， x j x_j xj 代表鄰居節點的特征信息，對于 h θ h_{\theta} hθ 代表每個 EdgeConv 層的可學習參數，也就是對應傳入的MLP層中的可學習參數。

所以我們的任務無非就是獲取這幾個變量，然后進行傳播計算即可

3.1.1 特征拼接

該環節實現的公式為： x i ∣ ∣ x j − x i x_i||x_j-x_i xi∣∣xj−xi，對于這個公式來說，我們要獲得兩個變量，一個是中心節點 x i x_i xi（target）的特征信息，一個是鄰居節點 x j x_j xj（source）的特征信息。

對于這兩個變量的獲取很容易，利用 edge_index 就可以提取出來，edge_index 中保存的是每一條邊的一對起始節點與終止節點，對于起始節點可以認為就是 i，對于終止節點就可以認為是 j，然后我們就會獲得兩個向量，分別為 row 和 col ，這兩個向量就是起始頂點和終止頂點的集合。

然后我們在根據索引進行提取特征，利用 x_i = x[row] 和 x_j = x[col] 就可以將中心節點和終止節點對應的特征獲取，維度為【E，feature_size】。

然后就可以按照公式實現做差然后與中心節點的特征進行拼接，獲得拼接后的特征維度為原來的2倍。

row, col = edge_index # 獲取target、source節點索引 [E]
x_i = x[row] # 獲取target節點信息，中心節點 [E, feature_size]
x_j = x[col] # 獲取source節點信息，鄰居節點 [E, feature_size]

x_cat = torch.cat([x_i, x_j - x_i], dim=1) # 拼接特征 [E, 2 * feature_size]

對于這里 x_i 和 x_j 以及起始節點的索引初學可能混淆，所以多多打印中間結果一步一步調試進行理解。

3.1.2 max聚合

對于 EdgeConv 的默認聚合方式為 max，其實還可以使用 mean 、sum 等排列不變函數進行聚合。

對于聚合操作就是公式中求和符號那里，只不過框架給的公式是 sum ，對于聚合我們希望做的是將中心節點的鄰居特征按照指定的聚合方式進行聚合。

我們可以利用 PyG 工具庫中提供的 scatter 函數進行操作，該函數可以指定聚合方式以及聚合維度等參數，使用方法就是需要傳入需要聚合的 Tensor ，此外還需要傳入一個 index ，指明哪些向量為同一個鄰居的節點，舉個例子，我們傳入的 index=[0,0,0,1,1] ，這就代表第一個、第二個、第三個為同一節點的鄰居，所以就會將待聚合的 Tensor 的第一個向量、第二個向量、第三個向量按照指定聚合方式進行聚合。

這里說的有點抽象，自己嘗試一個簡單示例就明白了。

out = scatter(src=x_cat, index=row, dim=0, reduce='max') # max聚合操作 [num_nodes, feature_size]

3.1.3 特征映射

在公式中有個 h θ h_{\theta} hθ，這個就代表 MLP 做特征映射做的，對于官方給的 EdgeConv 需要我們手動傳入 MLP 模型，所以本項目自實現也是按照這種方式，MLP 的操作在 EdgeConv 中并沒有實現，而是利用傳入的模型進行操作。

這里注意一點就是定義的 MLP 模型的輸入維度應該為原始維度的2倍，因為我們在這之前進行了特征拼接操作，所以特征維度進行了加倍。

out = self.mlp(out) # 特征映射 [num_nodes, out_channels]

3.1.4 EdgeConv層

接下來就可以定義EdgeConv層了，該層實現了1個函數，為 forward()

forward()：這個函數定義模型的傳播過程，也就是上面公式的 x i ′ = ∑ j ∈ N ( i ) h θ ( x i ∣ ∣ x j − x i ) x_i'=\sum_{j\in N(i)}h_{\theta}(x_i||x_j-x_i) xi′=∑j∈N(i)hθ(xi∣∣xj−xi)

# 2.定義EdgeConv層
class EdgeConv(nn.Module):
    def __init__(self, nn, aggr='max'):
        super(EdgeConv, self).__init__()
        self.mlp = nn # MLP網絡
    
    def forward(self, x, edge_index):
        row, col = edge_index # 獲取target、source節點索引 [E]
        x_i = x[row] # 獲取target節點信息，中心節點 [E, feature_size]
        x_j = x[col] # 獲取source節點信息，鄰居節點 [E, feature_size]
        
        x_cat = torch.cat([x_i, x_j - x_i], dim=1) # 拼接特征 [E, 2 * feature_size]
        out = scatter(src=x_cat, index=row, dim=0, reduce='max') # max聚合操作 [num_nodes, feature_size]
        
        out = self.mlp(out) # 特征映射 [num_nodes, out_channels]
        
        return out

對于我們實現這個網絡的實現效率上來講比PyG框架內置的 EdgeConv 層稍差一點，因為我們是按照公式來一步一步利用矩陣計算得到，沒有對矩陣計算以及算法進行優化，不然初學者可能看不太懂，不利于理解EdgeConv公式的傳播過程，有能力的小伙伴可以看下官方源碼學習一下，框架內是按照消息傳遞方式實現的。

3.2 定義EdgeCNN網絡

上面我們已經實現好了 EdgeConv 的網絡層，之后就可以調用這個層來搭建 EdgeCNN 網絡。

# 3.定義EdgeConv網絡
class EdgeCNN(nn.Module):
    def __init__(self, num_node_features, num_classes):
        super(EdgeCNN, self).__init__()
        self.conv1 = EdgeConv(nn=nn.Linear(2 * num_node_features, 16), aggr='max')
        self.conv2 = EdgeConv(nn=nn.Linear(2 * 16, num_classes), aggr='max')
        
    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        
        return F.log_softmax(x, dim=1)

上面網絡我們定義了兩個EdgeConv層，第一層的參數的輸入維度就是初始每個節點的特征維度 * 2，輸出維度是16。

第二個層的輸入維度為16 * 2，輸出維度為分類個數，因為我們需要對每個節點進行分類，最終加上softmax操作。

這里說明一下為什么要將輸入乘以2，原因是在使用MLP進行特征轉換之前，會將中心節點的特征與中心節點和鄰居節點的差向量做拼接，所以得到的輸出維度為節點的特征維度 * 2。

四、定義模型

下面就是定義了一些模型需要的參數，像學習率、迭代次數這些超參數，然后是模型的定義以及優化器及損失函數的定義，和pytorch定義網絡是一樣的。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 設備
epochs = 10 # 學習輪數
lr = 0.003 # 學習率
num_node_features = dataset.num_node_features # 每個節點的特征數
num_classes = dataset.num_classes # 每個節點的類別數
data = dataset[0].to(device) # Cora的一張圖

# 3.定義模型
model = EdgeCNN(num_node_features, num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr) # 優化器
loss_function = nn.NLLLoss() # 損失函數

五、模型訓練

模型訓練部分也是和pytorch定義網絡一樣，因為都是需要經過前向傳播、反向傳播這些過程，對于損失、精度這些指標可以自己添加。

# 訓練模式
model.train()

for epoch in range(epochs):
    optimizer.zero_grad()
    pred = model(data)
    
    loss = loss_function(pred[data.train_mask], data.y[data.train_mask]) # 損失
    correct_count_train = pred.argmax(axis=1)[data.train_mask].eq(data.y[data.train_mask]).sum().item() # epoch正確分類數目
    acc_train = correct_count_train / data.train_mask.sum().item() # epoch訓練精度
    
    loss.backward()
    optimizer.step()
    
    if epoch % 20 == 0:
        print("【EPOCH: 】%s" % str(epoch + 1))
        print('訓練損失為：{:.4f}'.format(loss.item()), '訓練精度為：{:.4f}'.format(acc_train))

print('【Finished Training！】')

六、模型驗證

下面就是模型驗證階段，在訓練時我們是只使用了訓練集，測試的時候我們使用的是測試集，注意這和傳統網絡測試不太一樣，在圖像分類一些經典任務中，我們是把數據集分成了兩份，分別是訓練集、測試集，但是在Cora這個數據集中并沒有這樣，它區分訓練集還是測試集使用的是掩碼機制，就是定義了一個和節點長度相同緯度的數組，該數組的每個位置為True或者False，標記著是否使用該節點的數據進行訓練。

# 模型驗證
model.eval()
pred = model(data)

# 訓練集（使用了掩碼）
correct_count_train = pred.argmax(axis=1)[data.train_mask].eq(data.y[data.train_mask]).sum().item()
acc_train = correct_count_train / data.train_mask.sum().item()
loss_train = loss_function(pred[data.train_mask], data.y[data.train_mask]).item()

# 測試集
correct_count_test = pred.argmax(axis=1)[data.test_mask].eq(data.y[data.test_mask]).sum().item()
acc_test = correct_count_test / data.test_mask.sum().item()
loss_test = loss_function(pred[data.test_mask], data.y[data.test_mask]).item()

print('Train Accuracy: {:.4f}'.format(acc_train), 'Train Loss: {:.4f}'.format(loss_train))
print('Test  Accuracy: {:.4f}'.format(acc_test), 'Test  Loss: {:.4f}'.format(loss_test))

七、結果

【EPOCH: 】1
訓練損失為：1.9629 訓練精度為：0.1214
【EPOCH: 】21
訓練損失為：1.6709 訓練精度為：0.5714
【EPOCH: 】41
訓練損失為：1.3965 訓練精度為：0.7571
【EPOCH: 】61
訓練損失為：1.1095 訓練精度為：0.8643
【EPOCH: 】81
訓練損失為：0.9088 訓練精度為：0.9286
【EPOCH: 】101
訓練損失為：0.7454 訓練精度為：0.9643
【EPOCH: 】121
訓練損失為：0.5841 訓練精度為：0.9643
【EPOCH: 】141
訓練損失為：0.4985 訓練精度為：0.9714
【EPOCH: 】161
訓練損失為：0.3954 訓練精度為：0.9714
【EPOCH: 】181
訓練損失為：0.3339 訓練精度為：0.9857
【Finished Training！】

>>>Train Accuracy: 1.0000 Train Loss: 0.3133
>>>Test  Accuracy: 0.4230 Test  Loss: 1.6562

完整代碼：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.utils import scatter
from torch_geometric.datasets import Planetoid

# 1.加載Cora數據集
dataset = Planetoid(root='./data/Cora', name='Cora')

# 2.定義EdgeConv層
class EdgeConv(nn.Module):
    def __init__(self, nn, aggr='max'):
        super(EdgeConv, self).__init__()
        self.mlp = nn # MLP網絡
    
    def forward(self, x, edge_index):
        row, col = edge_index # 獲取target、source節點索引 [E]
        x_i = x[row] # 獲取target節點信息，中心節點 [E, feature_size]
        x_j = x[col] # 獲取source節點信息，鄰居節點 [E, feature_size]
        
        x_cat = torch.cat([x_i, x_j - x_i], dim=1) # 拼接特征 [E, 2 * feature_size]
        out = scatter(src=x_cat, index=row, dim=0, reduce='max') # max聚合操作 [num_nodes, feature_size]
        
        out = self.mlp(out) # 特征映射 [num_nodes, out_channels]
        
        return out

# 3.定義EdgeConv網絡
class EdgeCNN(nn.Module):
    def __init__(self, num_node_features, num_classes):
        super(EdgeCNN, self).__init__()
        self.conv1 = EdgeConv(nn=nn.Linear(2 * num_node_features, 16), aggr='max')
        self.conv2 = EdgeConv(nn=nn.Linear(2 * 16, num_classes), aggr='max')
        
    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        
        return F.log_softmax(x, dim=1)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 設備
epochs = 200 # 學習輪數
lr = 0.0003 # 學習率
num_node_features = dataset.num_node_features # 每個節點的特征數
num_classes = dataset.num_classes # 每個節點的類別數
data = dataset[0].to(device) # Cora的一張圖

# 4.定義模型
model = EdgeCNN(num_node_features, num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr) # 優化器
loss_function = nn.NLLLoss() # 損失函數

# 訓練模式
model.train()

for epoch in range(epochs):
    optimizer.zero_grad()
    pred = model(data)
    
    loss = loss_function(pred[data.train_mask], data.y[data.train_mask]) # 損失
    correct_count_train = pred.argmax(axis=1)[data.train_mask].eq(data.y[data.train_mask]).sum().item() # epoch正確分類數目
    acc_train = correct_count_train / data.train_mask.sum().item() # epoch訓練精度
    
    loss.backward()
    optimizer.step()
    
    if epoch % 20 == 0:
        print("【EPOCH: 】%s" % str(epoch + 1))
        print('訓練損失為：{:.4f}'.format(loss.item()), '訓練精度為：{:.4f}'.format(acc_train))

print('【Finished Training！】')

# 模型驗證
model.eval()
pred = model(data)

# 訓練集（使用了掩碼）
correct_count_train = pred.argmax(axis=1)[data.train_mask].eq(data.y[data.train_mask]).sum().item()
acc_train = correct_count_train / data.train_mask.sum().item()
loss_train = loss_function(pred[data.train_mask], data.y[data.train_mask]).item()

# 測試集
correct_count_test = pred.argmax(axis=1)[data.test_mask].eq(data.y[data.test_mask]).sum().item()
acc_test = correct_count_test / data.test_mask.sum().item()
loss_test = loss_function(pred[data.test_mask], data.y[data.test_mask]).item()

print('Train Accuracy: {:.4f}'.format(acc_train), 'Train Loss: {:.4f}'.format(loss_train))
print('Test  Accuracy: {:.4f}'.format(acc_test), 'Test  Loss: {:.4f}'.format(loss_test))

感謝各位的閱讀，以上就是“基于PyTorch如何實現EdgeCNN”的內容了，經過本文的學習后，相信大家對基于PyTorch如何實現EdgeCNN這一問題有了更深刻的體會，具體使用情況還需要大家實踐驗證。這里是億速云，小編將為大家推送更多相關知識點的文章，歡迎關注！