DGL官方教程——图生成模型

DGL官方教程——图生成模型

        在本教程中，您将学习如何同时训练和生成一个图。 您还将探索图形嵌入过程种进行并行操作，这是必不可少的构建模块。 本教程以一个简单的优化结尾，该优化通过跨图批处理将速度提高了一倍。
        早期的教程显示了嵌入图形或节点是如何使您能够完成诸如半监督节点分类或情感分析之类任务的。 您是否认为预测图形的未来演变并迭代执行分析会很有趣？
        为了解决图的演变，您可以生成各种图样本。 换句话说，您需要图形的生成模型。 除学习节点和边特征外，您还需要对任意图的分布进行建模。 虽然一般的生成模型可以显式和隐式地对密度函数进行建模，并且可以一次或顺序生成样本，但是在此，您仅关注用于顺序生成的显式生成模型。 典型的应用包括药物或材料的发现，化学过程或蛋白质组学。

Introduction

更改Deep Graph Library（DGL）中的图形的基础操作无非就是add_nodes和add_edges。 也就是说，如果您要画一个由三个节点组成的圆。

可以如下编写代码。

import dgl

g = dgl.DGLGraph()
g.add_nodes(1)              # 添加编号为0的节点
g.add_nodes(1)              # 添加编号为1的节点

# 默认情况下，DGLGraph中的边是有向的。
# 对于无方向的边，为两个方向都添加边。
g.add_edges([1, 0], [0, 1]) # 添加由节点1到节点0的边，添加由节点0到节点1的边，
g.add_nodes(1)              # 添加编号为2的节点
g.add_edges([2, 1], [1, 2]) # 添加由节点2到节点1的边，添加由节点1到节点2的边，
g.add_edges([2, 0], [0, 2]) # 添加由节点2到节点0的边，添加由节点0到节点2的边，

现实世界中的图要复杂得多。 他们具有不同的种类，不同种类的图具有不同的大小，拓扑，节点类型，边类型并且可能是多重图。 此外，可以以许多不同的顺序生成相同的图。 无论如何，生成过程需要以下几个步骤。

• 对变换的图形进行编码
• 随机执行操作
• 在训练过程中，收集错误信号并优化模型参数

关于实现，另一个重要方面是速度。 考虑到生成图基本上是一个顺序过程，并行化计算值得我们仔细考虑。

可以肯定的是，这不一定是硬约束。 子图可以并行构建，然后组装。 但是，我们将本教程局限于顺序过程。

DGMG：The main flow

在本教程中，您将使用图的深度生成模型（DGMG）来使用DGL实现图的生成模型。 它的算法框架是通用的，但也很难并行化。

DGMG可以处理带有类型化节点，类型化边缘和多图的复杂图形，但您可以在此处使用其简化版本来生成图形拓扑。

DGMG通过遵循状态机来生成图，该状态机通常上是一个两级循环。 一次生成一个节点，并同时将其连接到现有节点的子集上。 这类似于语言建模。 生成过程是一种迭代过程，它以到目前为止生成的序列为条件，一次延申一个单词，一个字符或一个句子。

在每个时间步骤中，您可以选择一下二者之一：

• 向图中添加一个新节点
• 选择两个现有节点并在它们之间添加一条边

Python代码如下所示。 实际上，这正是在DGL中实现DGMG推理的方式。

def forward_inference(self):
    stop = self.add_node_and_update()
    while (not stop) and (self.g.number_of_nodes() < self.v_max + 1):
        num_trials = 0
        to_add_edge = self.add_edge_or_not()
        while to_add_edge and (num_trials < self.g.number_of_nodes() - 1):
            self.choose_dest_and_update()
            num_trials += 1
            to_add_edge = self.add_edge_or_not()
        stop = self.add_node_and_update()

    return self.g

假设您具有用于生成节点10-20的循环的预训练模型。 它如何在推理过程中动态生成一个周期？ 使用下面的代码使用自己的模型创建动画。

import torch
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import networkx as nx
from copy import deepcopy

if __name__ == '__main__':
    # 预训练模型存在“ ./model.pth”路径中
    model = torch.load('./model.pth')
    model.eval()
    g = model()

    src_list = g.edges()[1]
    dest_list = g.edges()[0]

    evolution = []

    nx_g = nx.Graph()
    evolution.append(deepcopy(nx_g))

    for i in range(0, len(src_list), 2):
        src = src_list[i].item()
        dest = dest_list[i].item()
        if src not in nx_g.nodes():
            nx_g.add_node(src)
            evolution.append(deepcopy(nx_g))
        if dest not in nx_g.nodes():
            nx_g.add_node(dest)
            evolution.append(deepcopy(nx_g))
        nx_g.add_edges_from([(src, dest), (dest, src)])
        evolution.append(deepcopy(nx_g))

    def animate(i):
        ax.cla()
        g_t = evolution[i]
        nx.draw_circular(g_t, with_labels=True, ax=ax,
                         node_color=['#FEBD69'] * g_t.number_of_nodes())

    fig, ax = plt.subplots()
    ani = animation.FuncAnimation(fig, animate,
                                  frames=len(evolution),
                                  interval=600)

DGMG: Optimization objective

与语言建模类似，DGMG通过behavior cloning或者teacher forcing来训练模型。 假设每个图都有一系列生成它的预测动作 $a_1$，⋯，$a_T$。 模型要做的是遵循这些动作，计算这些动作序列的联合概率，并使它们最大化。
根据链式规则，取$a_1$，⋯，$a_T$的概率为：

$$p(a_1,\cdots ,a_T)=p(a1)p(a_2|a_1)\cdots p(a_T|a_1,\cdots ,a_{T-1}).$$

那么优化目标就是典型的MLE损失：

$$-logp(a_1,\cdots ,a_T)=-\sum _{t=1}^{T}logp(a_t|a_1,\cdots ,a_{t-1}).$$

def forward_train(self, actions):
    """
    - actions: 列表
        - 包含上述的a_1，...，a_T
    - self.prepare_for_train()
        - 将self.action_step初始化为0，每次调用它将递增1。
        - 初始化记录日志p（a_t | a_1，... a_ {t-1}）的对象

    Returns
    -------
    - self.get_log_prob(): log p(a_1, ..., a_T)
    """
    self.prepare_for_train()

    stop = self.add_node_and_update(a=actions[self.action_step])
    while not stop:
        to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
        while to_add_edge:
            self.choose_dest_and_update(a=actions[self.action_step])
            to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
        stop = self.add_node_and_update(a=actions[self.action_step])

    return self.get_log_prob()

forward_train和forward_inference之间的主要区别在于，训练过程将oracle动作作为输入，并返回用于评估损失的对数概率。

DGMG:The implementation

The DGMG class

您可以在下面找到该模型的框架代码。 您逐步填写每个功能的详细信息。

import torch.nn as nn


class DGMGSkeleton(nn.Module):
    def __init__(self, v_max):
        """
        参数
        ----------
        v_max: 整型
            考虑的最大节点数
        """
        super(DGMGSkeleton, self).__init__()

        # 图形配置
        self.v_max = v_max

    def add_node_and_update(self, a=None):
        """决定是否添加新节点。
        如果应添加新节点，请更新图形。"""
        return NotImplementedError

    def add_edge_or_not(self, a=None):
        """决定是否应添加新边."""
        return NotImplementedError

    def choose_dest_and_update(self, a=None):
        """选择目标并将其连接到最新节点。
           为两个方向添加边并更新图形。"""
        return NotImplementedError

    def forward_train(self, actions):
        """在训练时前向运算。
         它记录了根据这些动作生成基本事实图的可能性。"""
        return NotImplementedError

    def forward_inference(self):
        """在推断时间前向运算。
           它可以即时生成图形。"""
        return NotImplementedError

    def forward(self, actions=None):
        # 您将处理的图形
        self.g = dgl.DGLGraph()

        # 如果节点和边具有某些特征，则将为新节点和边设置零张量。
        self.g.set_n_initializer(dgl.frame.zero_initializer)
        self.g.set_e_initializer(dgl.frame.zero_initializer)

        if self.training:
            return self.forward_train(actions=actions)
        else:
            return self.forward_inference()

Encoding a dynamic graph

所有生成图形的动作都是从概率分布中采样得到的。 为此，可以将结构化数据（即图形）投影到欧几里得空间上。 挑战在于，随着图形的变化，需要重复这种称为嵌入的过程。

Graph embedding

$G=(V,E)$ 表示任意图，对于每一个节点$\nu$都对应一个确定的嵌入向量 ${\text{h}}_v\in {\mathbb{R}}^n$。相似地，图也对应一个嵌入向量${\text{h}}_G\in {\mathbb{R}}^k$。显然，k>n，因为因为图形比单个节点包含更多的信息。

图嵌入是线性变换下的节点嵌入的加权和：
$${\text{h}}_G=\sum _{v\in V}\text{Sigmoid}(g_m({\text{h}}_v))f_m({\text{h}}_v),$$
第一阶段，$Sigmoid(g_m(h_v))$,计算一个选通函数，并且可以考虑整个图形嵌入在每个节点上的参与量。第二阶段，$f_m:R_n\to R^k$，将节点嵌入映射到图嵌入的空间。
将图形嵌入实现为GraphEmbed类。

import torch


class GraphEmbed(nn.Module):
    def __init__(self, node_hidden_size):
        super(GraphEmbed, self).__init__()

        # 从文章中进行设置
        self.graph_hidden_size = 2 * node_hidden_size

        # 图嵌入
        self.node_gating = nn.Sequential(
            nn.Linear(node_hidden_size, 1),
            nn.Sigmoid()
        )
        self.node_to_graph = nn.Linear(node_hidden_size,
                                       self.graph_hidden_size)

    def forward(self, g):
        if g.number_of_nodes() == 0:
            return torch.zeros(1, self.graph_hidden_size)
        else:
            # 节点特征以hv的形式存储在ndata中。
            hvs = g.ndata['hv']
            return (self.node_gating(hvs) *
                    self.node_to_graph(hvs)).sum(0, keepdim=True)

Update node embeddings via graph propagation

DGMG中节点嵌入的更新机制与图卷积网络相似。 对于图中的节点$v$，其邻居$u$向其发送一条消息。
$${\text{m}}_{u\to v}={\text{W}}_m\text{concat}([{\text{h}}_v,{\text{h}}_u,{\text{x}}_{u,v}])+{\text{b}}_m,$$
其中${\text{x}}_{u,v}$是$h_v$和$h_u$之间的边缘的嵌入。
从所有邻居收到消息后，$v$用节点激活向量对其进行汇总
$${\text{a}}_v=\sum _{u:(u,v)\in E}{\text{m}}_{u\to v}$$
并使用此信息来更新其自身的功能：
$${\text{h}}_v^{\prime }=\text{GRU}({\text{h}}_v,{\text{a}}_v).$$
同步地对所有节点一次执行以上所有操作称为一轮图传播。 您执行的图形传播次数越多，消息在整个图形中传播的距离就越长。
使用DGL，您可以使用g.update_all实现图传播。 这里的消息符号可能会有些混乱。 研究人员可以将${\text{m}}_{u\to v}$称为消息，但是下面的消息功能仅通过$\text{concat}([{\text{h}}_u,{\text{x}}_{u,v}])$。 然后出于效率考虑，一次跨所有边缘执行操作${\text{W}}_m\text{concat}([{\text{h}}_v,{\text{h}}_u,{\text{x}}_{u,v}])+{\text{b}}_m$。

from functools import partial

class GraphProp(nn.Module):
    def __init__(self, num_prop_rounds, node_hidden_size):
        super(GraphProp, self).__init__()

        self.num_prop_rounds = num_prop_rounds

        # Setting from the paper
        self.node_activation_hidden_size = 2 * node_hidden_size

        message_funcs = []
        node_update_funcs = []
        self.reduce_funcs = []

        for t in range(num_prop_rounds):
            # input being [hv, hu, xuv]
            message_funcs.append(nn.Linear(2 * node_hidden_size + 1,
                                           self.node_activation_hidden_size))

            self.reduce_funcs.append(partial(self.dgmg_reduce, round=t))
            node_update_funcs.append(
                nn.GRUCell(self.node_activation_hidden_size,
                           node_hidden_size))

        self.message_funcs = nn.ModuleList(message_funcs)
        self.node_update_funcs = nn.ModuleList(node_update_funcs)

    def dgmg_msg(self, edges):
        """For an edge u->v, return concat([h_u, x_uv])"""
        return {'m': torch.cat([edges.src['hv'],
                                edges.data['he']],
                               dim=1)}

    def dgmg_reduce(self, nodes, round):
        hv_old = nodes.data['hv']
        m = nodes.mailbox['m']
        message = torch.cat([
            hv_old.unsqueeze(1).expand(-1, m.size(1), -1), m], dim=2)
        node_activation = (self.message_funcs[round](message)).sum(1)

        return {'a': node_activation}

    def forward(self, g):
        if g.number_of_edges() > 0:
            for t in range(self.num_prop_rounds):
                g.update_all(message_func=self.dgmg_msg,
                             reduce_func=self.reduce_funcs[t])
                g.ndata['hv'] = self.node_update_funcs[t](
                     g.ndata['a'], g.ndata['hv'])

Actions

从使用神经网络参数化的分布中采样所有动作，然后依次进行。

Action 1: Add nodes

给定图嵌入向量${\text{h}}_G$，
$$\text{Sigmoid}({\text{W}}_{\text{add node}}{\text{h}}_G+b_{\text{add node}})$$
然后将其用于参数化Bernoulli分布，以决定是否添加新节点。

如果要添加新节点，请使用以下命令初始化其功能:
$${\text{W}}_{\text{init}}\text{concat}([{\text{h}}_{\text{init}},{\text{h}}_G])+{\text{b}}_{\text{init}},$$
其中，${\text{h}}_{\text{init}}$是可学习的用于无类型节点的嵌入模块。

import torch.nn.functional as F
from torch.distributions import Bernoulli

def bernoulli_action_log_prob(logit, action):
    """计算有关伯努利分布的操作的对数p。 使用logit而不是prob以获得数值稳定性。"""
    if action == 0:
        return F.logsigmoid(-logit)
    else:
        return F.logsigmoid(logit)

class AddNode(nn.Module):
    def __init__(self, graph_embed_func, node_hidden_size):
        super(AddNode, self).__init__()

        self.graph_op = {'embed': graph_embed_func}

        self.stop = 1
        self.add_node = nn.Linear(graph_embed_func.graph_hidden_size, 1)

        # If to add a node, initialize its hv
        self.node_type_embed = nn.Embedding(1, node_hidden_size)
        self.initialize_hv = nn.Linear(node_hidden_size + \
                                       graph_embed_func.graph_hidden_size,
                                       node_hidden_size)

        self.init_node_activation = torch.zeros(1, 2 * node_hidden_size)

    def _initialize_node_repr(self, g, node_type, graph_embed):
        """Whenver a node is added, initialize its representation."""
        num_nodes = g.number_of_nodes()
        hv_init = self.initialize_hv(
            torch.cat([
                self.node_type_embed(torch.LongTensor([node_type])),
                graph_embed], dim=1))
        g.nodes[num_nodes - 1].data['hv'] = hv_init
        g.nodes[num_nodes - 1].data['a'] = self.init_node_activation

    def prepare_training(self):
        self.log_prob = []

    def forward(self, g, action=None):
        graph_embed = self.graph_op['embed'](g)

        logit = self.add_node(graph_embed)
        prob = torch.sigmoid(logit)

        if not self.training:
            action = Bernoulli(prob).sample().item()
        stop = bool(action == self.stop)

        if not stop:
            g.add_nodes(1)
            self._initialize_node_repr(g, action, graph_embed)

        if self.training:
            sample_log_prob = bernoulli_action_log_prob(logit, action)

            self.log_prob.append(sample_log_prob)

        return stop

Action 2: Add edges

给定图嵌入向量${\text{h}}_G$和最新节点$v$的节点嵌入向量${\text{h}}_v$，您可以评估:
$$\text{Sigmoid}({\text{W}}_{\text{add edge}}\text{concat}([{\text{h}}_G,{\text{h}}_v])+b_{\text{add edge}}),$$
然后将其用于参数化Bernoulli分布，以决定是否从$v$开始添加新边。

class AddEdge(nn.Module):
    def __init__(self, graph_embed_func, node_hidden_size):
        super(AddEdge, self).__init__()

        self.graph_op = {'embed': graph_embed_func}
        self.add_edge = nn.Linear(graph_embed_func.graph_hidden_size + \
                                  node_hidden_size, 1)

    def prepare_training(self):
        self.log_prob = []

    def forward(self, g, action=None):
        graph_embed = self.graph_op['embed'](g)
        src_embed = g.nodes[g.number_of_nodes() - 1].data['hv']

        logit = self.add_edge(torch.cat(
            [graph_embed, src_embed], dim=1))
        prob = torch.sigmoid(logit)

        if self.training:
            sample_log_prob = bernoulli_action_log_prob(logit, action)
            self.log_prob.append(sample_log_prob)
        else:
            action = Bernoulli(prob).sample().item()

        to_add_edge = bool(action == 0)
        return to_add_edge

Action 3: Choose a destination

当Action 2返回True时，请选择最新节点$v$的目标。
对于每个可能的目标$u\in \{0,\cdots ,v-1\}$，选择它的概率为:
$$\frac{\text{exp}({\text{W}}_{\text{dest}}\text{concat}([{\text{h}}_u,{\text{h}}_v])+{\text{b}}_{\text{dest}})}{{\sum }_{i=0}^{v-1}\text{exp}({\text{W}}_{\text{dest}}\text{concat}([{\text{h}}_i,{\text{h}}_v])+{\text{b}}_{\text{dest}})}$$

from torch.distributions import Categorical

class ChooseDestAndUpdate(nn.Module):
    def __init__(self, graph_prop_func, node_hidden_size):
        super(ChooseDestAndUpdate, self).__init__()

        self.graph_op = {'prop': graph_prop_func}
        self.choose_dest = nn.Linear(2 * node_hidden_size, 1)

    def _initialize_edge_repr(self, g, src_list, dest_list):
        # 对于无类型的边，仅加1表示其存在。
        # 对于多种边缘类型，请使用one-hot表示或嵌入模块。
        edge_repr = torch.ones(len(src_list), 1)
        g.edges[src_list, dest_list].data['he'] = edge_repr

    def prepare_training(self):
        self.log_prob = []

    def forward(self, g, dest):
        src = g.number_of_nodes() - 1
        possible_dests = range(src)

        src_embed_expand = g.nodes[src].data['hv'].expand(src, -1)
        possible_dests_embed = g.nodes[possible_dests].data['hv']

        dests_scores = self.choose_dest(
            torch.cat([possible_dests_embed,
                       src_embed_expand], dim=1)).view(1, -1)
        dests_probs = F.softmax(dests_scores, dim=1)

        if not self.training:
            dest = Categorical(dests_probs).sample().item()

        if not g.has_edge_between(src, dest):
            # 对于无向图，请同时添加两个方向的边，以便执行图传播。
            src_list = [src, dest]
            dest_list = [dest, src]

            g.add_edges(src_list, dest_list)
            self._initialize_edge_repr(g, src_list, dest_list)

            self.graph_op['prop'](g)

        if self.training:
            if dests_probs.nelement() > 1:
                self.log_prob.append(
                    F.log_softmax(dests_scores, dim=1)[:, dest: dest + 1])

Putting it together

现在，您可以准备完整的模型类实现了。

class DGMG(DGMGSkeleton):
    def __init__(self, v_max, node_hidden_size,
                 num_prop_rounds):
        super(DGMG, self).__init__(v_max)

        # 图形嵌入模型
        self.graph_embed = GraphEmbed(node_hidden_size)

        # 图传播模块
        self.graph_prop = GraphProp(num_prop_rounds,
                                    node_hidden_size)

        # Actions
        self.add_node_agent = AddNode(
            self.graph_embed, node_hidden_size)
        self.add_edge_agent = AddEdge(
            self.graph_embed, node_hidden_size)
        self.choose_dest_agent = ChooseDestAndUpdate(
            self.graph_prop, node_hidden_size)

        # 前向函数
        self.forward_train = partial(forward_train, self=self)
        self.forward_inference = partial(forward_inference, self=self)

    @property
    def action_step(self):
        old_step_count = self.step_count
        self.step_count += 1

        return old_step_count

    def prepare_for_train(self):
        self.step_count = 0

        self.add_node_agent.prepare_training()
        self.add_edge_agent.prepare_training()
        self.choose_dest_agent.prepare_training()

    def add_node_and_update(self, a=None):
        """判断是否添加一个新的节点
        如果应该添加一个新的节点,那么更新整个图"""

        return self.add_node_agent(self.g, a)

    def add_edge_or_not(self, a=None):
        """判断是否应该添加一个边."""

        return self.add_edge_agent(self.g, a)

    def choose_dest_and_update(self, a=None):
        """选择目标并将其连接到最新节点。
        为两个方向添加边并更新图形。"""

        self.choose_dest_agent(self.g, a)

    def get_log_prob(self):
        add_node_log_p = torch.cat(self.add_node_agent.log_prob).sum()
        add_edge_log_p = torch.cat(self.add_edge_agent.log_prob).sum()
        choose_dest_log_p = torch.cat(self.choose_dest_agent.log_prob).sum()
        return add_node_log_p + add_edge_log_p + choose_dest_log_p

下面是一个动画，其中在前400批次的每10批次训练之后，动态生成了图形。 您可以看到模型随着时间的推移如何改进并开始生成周期。

对于生成模型，您可以通过检查动态生成的图形中有效图形的百分比来评估性能。

import torch.utils.model_zoo as model_zoo

# 下载预训练的模型状态字典，以生成具有10-20个节点的循环。
state_dict = model_zoo.load_url('https://s3.us-east-2.amazonaws.com/dgl.ai/model/dgmg_cycles-5a0c40be.pth')
model = DGMG(v_max=20, node_hidden_size=16, num_prop_rounds=2)
model.load_state_dict(state_dict)
model.eval()

def is_valid(g):
    # 检查g是否是具有10-20个节点的循环。
    def _get_previous(i, v_max):
        if i == 0:
            return v_max
        else:
            return i - 1

    def _get_next(i, v_max):
        if i == v_max:
            return 0
        else:
            return i + 1

    size = g.number_of_nodes()

    if size < 10 or size > 20:
        return False

    for node in range(size):
        neighbors = g.successors(node)

        if len(neighbors) != 2:
            return False

        if _get_previous(node, size - 1) not in neighbors:
            return False

        if _get_next(node, size - 1) not in neighbors:
            return False

    return True

num_valid = 0
for i in range(100):
    g = model()
    num_valid += is_valid(g)

del model
print('Among 100 graphs generated, {}% are valid.'.format(num_valid))

Out:

Among 100 graphs generated, 94% are valid.
有关完整的实现，请参见DGL DGMG示例。

Batched graph generation

加快DGMG的速度非常困难，因为仅可以通过唯一的操作序列来生成每个图形。 探索并行性的一种方法是采用具有多个过程的异步梯度下降。 它们每一次只能处理一张图，并且过程由参数服务器松散地协调。

DGL在框架提供的张量操作之上探索消息传递框架中的并行性。 较早的教程已经在消息传播和图嵌入阶段中做到了这一点，但仅在一个图内。 对于一批图形，则需要一个for循环：

for g in g_list:
	self.graph_prop(g)

修改代码以一次处理一批图形，方法是将这些行替换为以下内容。 在具有macOS的CPU上，您可以立即将图形传播部分减少6到7倍。

bg = dgl.batch(g_list)
self.graph_prop(bg)
g_list = dgl.unbatch(bg)

您已经在Tree-LSTM教程中使用了调用dgl.batch的技巧，并且值得再解释一次。
通过批处理许多小图，DGL在内部维护了一个大容器图（BatchedDGLGraph），在该图上update_all推动消息在所有边缘和节点上传递。
使用dgl.batch，可以将**g_ {1}，…，g_ {N}**合并为一个包含N个孤立的小图的单个巨型图。 例如，如果我们有两个带有邻接矩阵的图

[0, 1]
[1, 0]

[0, 1, 0]
[1, 0, 0]
[0, 1, 0]

dgl.batch只是给出了一个邻接矩阵为

[0, 1, 0, 0, 0]
[1, 0, 0, 0, 0]
[0, 1, 0, 0, 0]
[1, 0, 0, 0, 0]
[0, 1, 0, 0, 0]

在DGL中，消息功能是在边缘上定义的，因此批处理线性地扩展了边缘用户定义功能（UDF）的处理。
reduce UDF或dgmg_reduce在节点上工作。 它们中的每一个可以具有不同数量的进入边缘。 使用degree bucketing，DGL在内部以相同的度对节点进行分组，并为每个组调用一次减少UDF。 因此，批处理还减少了对这些UDF的调用次数。

BatchedDGLGraph对象的节点/边缘特征的修改不会影响原始小图的特征，因此我们需要用新图列表**g_list = dgl.unbatch（bg）**替换旧图列表。

批处理版本的完整代码也可以在示例中找到。 在测试平台上，与之前的实现相比，您的速度大约提高了一倍。

脚本的总运行时间：（0分钟6.276秒）

$$[下载脚本：5dgmg.py](https://docs.dgl.ai{/}_{d}ownloads/5_{d}gmg.py)$$

$$[下载Jupyter脚本：5dgmg.ipynb](https://docs.dgl.ai{/}_{d}ownloads/5_{d}gmg.ipynb)$$