圖神經網絡22-DGL實戰：針對邊分類任務的鄰居採樣訓練方法

邊分類/迴歸的訓練與節點分類/迴歸的訓練類似，但還是有一些明顯的區別。

定義鄰居採樣器和數據加載器

用戶可以使用和節點分類一樣的鄰居採樣器 。

    sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)

想要用DGL提供的鄰居採樣器做邊分類，需要將其與
:class:~dgl.dataloading.pytorch.EdgeDataLoader 結合使用。
:class:~dgl.dataloading.pytorch.EdgeDataLoader 以小批次的形式對一組邊進行迭代，
從而產生包含邊小批次的子圖以及供下文中模塊使用的 塊。

例如，以下代碼創建了一個PyTorch數據加載器，該PyTorch數據加載器以批的形式迭代訓練邊ID的數組
train_eids，並將生成的塊列表放到GPU上。

    dataloader = dgl.dataloading.EdgeDataLoader(
        g, train_eid_dict, sampler,
        batch_size=1024,
        shuffle=True,
        drop_last=False,
        num_workers=4)

有關DGL的內置採樣器的完整列表，用戶可以參考neighborhood sampler API reference <api-dataloading-neighbor-sampling>。

如果用戶希望開發自己的鄰居採樣器，或者想要對塊的概念有更詳細的瞭解，請參考guide_cn-minibatch-customizing-neighborhood-sampler。

小批次鄰居採樣訓練時刪邊

用戶在訓練邊分類模型時，有時希望從計算依賴中刪除出現在訓練數據中的邊，就好像這些邊根本不存在一樣。
否則，模型將 "知道" 兩個節點之間存在邊的聯繫，並有可能利用這點 "作弊" 。

因此，在基於鄰居採樣的邊分類中，用戶有時會希望從採樣得到的小批次圖中刪去部分邊及其對應的反向邊。
用戶可以在實例化
:class:~dgl.dataloading.pytorch.EdgeDataLoader
時設置 exclude='reverse_id'，同時將邊ID映射到其反向邊ID。
通常這樣做會導致採樣過程變慢很多，這是因爲DGL要定位並刪除包含在小批次中的反向邊。

    n_edges = g.number_of_edges()
    dataloader = dgl.dataloading.EdgeDataLoader(
        g, train_eid_dict, sampler,

        # 下面的兩個參數專門用於在鄰居採樣時刪除小批次的一些邊和它們的反向邊
        exclude='reverse_id',
        reverse_eids=torch.cat([
            torch.arange(n_edges // 2, n_edges), torch.arange(0, n_edges // 2)]),
    
        batch_size=1024,
        shuffle=True,
        drop_last=False,
        num_workers=4)

調整模型以適用小批次訓練

邊分類模型通常由兩部分組成：

• 獲取邊兩端節點的表示。
• 用邊兩端節點表示爲每個類別打分。

第一部分與
:ref:隨機批次訓練節點分類 <guide_cn-minibatch-node-classification-model>
完全相同，用戶可以簡單地複用它。輸入仍然是DGL的數據加載器生成的塊列表和輸入特徵。

    class StochasticTwoLayerGCN(nn.Module):
        def __init__(self, in_features, hidden_features, out_features):
            super().__init__()
            self.conv1 = dglnn.GraphConv(in_features, hidden_features)
            self.conv2 = dglnn.GraphConv(hidden_features, out_features)
    
        def forward(self, blocks, x):
            x = F.relu(self.conv1(blocks[0], x))
            x = F.relu(self.conv2(blocks[1], x))
            return x

第二部分的輸入通常是前一部分的輸出，以及由小批次邊導出的原始圖的子圖。
子圖是從相同的數據加載器產生的。用戶可以調用 :meth:dgl.DGLHeteroGraph.apply_edges 計算邊子圖中邊的得分。

以下代碼片段實現了通過合併邊兩端節點的特徵並將其映射到全連接層來預測邊的得分。

    class ScorePredictor(nn.Module):
        def __init__(self, num_classes, in_features):
            super().__init__()
            self.W = nn.Linear(2 * in_features, num_classes)
    
        def apply_edges(self, edges):
            data = torch.cat([edges.src['x'], edges.dst['x']])
            return {'score': self.W(data)}
    
        def forward(self, edge_subgraph, x):
            with edge_subgraph.local_scope():
                edge_subgraph.ndata['x'] = x
                edge_subgraph.apply_edges(self.apply_edges)
                return edge_subgraph.edata['score']

模型接受數據加載器生成的塊列表、邊子圖以及輸入節點特徵進行前向傳播，如下所示：

    class Model(nn.Module):
        def __init__(self, in_features, hidden_features, out_features, num_classes):
            super().__init__()
            self.gcn = StochasticTwoLayerGCN(
                in_features, hidden_features, out_features)
            self.predictor = ScorePredictor(num_classes, out_features)
    
        def forward(self, edge_subgraph, blocks, x):
            x = self.gcn(blocks, x)
            return self.predictor(edge_subgraph, x)

DGL保證邊子圖中的節點與生成的塊列表中最後一個塊的輸出節點相同。

模型的訓練

模型的訓練與節點分類的隨機批次訓練的情況非常相似。用戶可以遍歷數據加載器以獲得由小批次邊組成的子圖，
以及計算其兩端節點表示所需的塊列表。

    model = Model(in_features, hidden_features, out_features, num_classes)
    model = model.cuda()
    opt = torch.optim.Adam(model.parameters())
    
    for input_nodes, edge_subgraph, blocks in dataloader:
        blocks = [b.to(torch.device('cuda')) for b in blocks]
        edge_subgraph = edge_subgraph.to(torch.device('cuda'))
        input_features = blocks[0].srcdata['features']
        edge_labels = edge_subgraph.edata['labels']
        edge_predictions = model(edge_subgraph, blocks, input_features)
        loss = compute_loss(edge_labels, edge_predictions)
        opt.zero_grad()
        loss.backward()
        opt.step()

異構圖上的模型訓練

在異構圖上，計算節點表示的模型也可以用於計算邊分類/迴歸所需的兩端節點的表示。

    class StochasticTwoLayerRGCN(nn.Module):
        def __init__(self, in_feat, hidden_feat, out_feat, rel_names):
            super().__init__()
            self.conv1 = dglnn.HeteroGraphConv({
                    rel : dglnn.GraphConv(in_feat, hidden_feat, norm='right')
                    for rel in rel_names
                })
            self.conv2 = dglnn.HeteroGraphConv({
                    rel : dglnn.GraphConv(hidden_feat, out_feat, norm='right')
                    for rel in rel_names
                })
    
        def forward(self, blocks, x):
            x = self.conv1(blocks[0], x)
            x = self.conv2(blocks[1], x)
            return x

在同構圖和異構圖上做評分預測時，代碼實現的唯一不同在於調用
:meth:~dgl.DGLHeteroGraph.apply_edges
時需要在特定類型的邊上進行迭代。

    class ScorePredictor(nn.Module):
        def __init__(self, num_classes, in_features):
            super().__init__()
            self.W = nn.Linear(2 * in_features, num_classes)
    
        def apply_edges(self, edges):
            data = torch.cat([edges.src['x'], edges.dst['x']])
            return {'score': self.W(data)}
    
        def forward(self, edge_subgraph, x):
            with edge_subgraph.local_scope():
                edge_subgraph.ndata['x'] = x
                for etype in edge_subgraph.canonical_etypes:
                    edge_subgraph.apply_edges(self.apply_edges, etype=etype)
                return edge_subgraph.edata['score']

    class Model(nn.Module):
        def __init__(self, in_features, hidden_features, out_features, num_classes,
                     etypes):
            super().__init__()
            self.rgcn = StochasticTwoLayerRGCN(
                in_features, hidden_features, out_features, etypes)
            self.pred = ScorePredictor(num_classes, out_features)

        def forward(self, edge_subgraph, blocks, x):
            x = self.rgcn(blocks, x)
            return self.pred(edge_subgraph, x)

數據加載器的定義也與節點分類的非常相似。唯一的區別是用戶需要使用
~dgl.dataloading.pytorch.EdgeDataLoader而不是~dgl.dataloading.pytorch.NodeDataLoader，
並且提供邊類型和邊ID張量的字典，而不是節點類型和節點ID張量的字典。

    sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)
    dataloader = dgl.dataloading.EdgeDataLoader(
        g, train_eid_dict, sampler,
        batch_size=1024,
        shuffle=True,
        drop_last=False,
        num_workers=4)

如果用戶希望刪除異構圖中的反向邊，情況會有所不同。在異構圖上，
反向邊通常具有與正向邊本身不同的邊類型，以便區分 向前 和 向後 關係。
例如，關注 和 被關注 是一對相反的關係， 購買 和 被買下 也是一對相反的關係。

如果一個類型中的每個邊都有一個與之對應的ID相同、屬於另一類型的反向邊，
則用戶可以指定邊類型及其反向邊類型之間的映射。刪除小批次中的邊及其反向邊的方法如下。

    dataloader = dgl.dataloading.EdgeDataLoader(
        g, train_eid_dict, sampler,
    
        # 下面的兩個參數專門用於在鄰居採樣時刪除小批次的一些邊和它們的反向邊
        exclude='reverse_types',
        reverse_etypes={'follow': 'followed by', 'followed by': 'follow',
                        'purchase': 'purchased by', 'purchased by': 'purchase'}
    
        batch_size=1024,
        shuffle=True,
        drop_last=False,
        num_workers=4)

除了 compute_loss 的代碼實現有所不同，異構圖的訓練循環與同構圖中的訓練循環幾乎相同，
計算損失函數接受節點類型和預測的兩個字典。

    model = Model(in_features, hidden_features, out_features, num_classes, etypes)
    model = model.cuda()
    opt = torch.optim.Adam(model.parameters())
    
    for input_nodes, edge_subgraph, blocks in dataloader:
        blocks = [b.to(torch.device('cuda')) for b in blocks]
        edge_subgraph = edge_subgraph.to(torch.device('cuda'))
        input_features = blocks[0].srcdata['features']
        edge_labels = edge_subgraph.edata['labels']
        edge_predictions = model(edge_subgraph, blocks, input_features)
        loss = compute_loss(edge_labels, edge_predictions)
        opt.zero_grad()
        loss.backward()
        opt.step()

完整代碼可以查看：GCMC <https://github.com/dmlc/dgl/tree/master/examples/pytorch/gcmc>