wehub/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dgl/nn/pytorch/explain/__init__.py
+6
View File
@@ -0,0 +1,6 @@
"""Torch modules for explanation models."""
# pylint: disable= no-member, arguments-differ, invalid-name
from .gnnexplainer import *
from .subgraphx import *
from .pgexplainer import *
python/dgl/nn/pytorch/explain/gnnexplainer.py
+941
View File
@@ -0,0 +1,941 @@
"""Torch Module for GNNExplainer"""
# pylint: disable= no-member, arguments-differ, invalid-name
from math import sqrt
import torch
from torch import nn
from tqdm.auto import tqdm
from ....base import EID, NID
from ....subgraph import khop_in_subgraph
__all__ = ["GNNExplainer", "HeteroGNNExplainer"]
class GNNExplainer(nn.Module):
r"""GNNExplainer model from `GNNExplainer: Generating Explanations for
Graph Neural Networks <https://arxiv.org/abs/1903.03894>`__
It identifies compact subgraph structures and small subsets of node features that play a
critical role in GNN-based node classification and graph classification.
To generate an explanation, it learns an edge mask :math:`M` and a feature mask :math:`F`
by optimizing the following objective function.
.. math::
l(y, \hat{y}) + \alpha_1 \|M\|_1 + \alpha_2 H(M) + \beta_1 \|F\|_1 + \beta_2 H(F)
where :math:`l` is the loss function, :math:`y` is the original model prediction,
:math:`\hat{y}` is the model prediction with the edge and feature mask applied, :math:`H` is
the entropy function.
Parameters
----------
model : nn.Module
The GNN model to explain.
* The required arguments of its forward function are graph and feat.
The latter one is for input node features.
* It should also optionally take an eweight argument for edge weights
and multiply the messages by it in message passing.
* The output of its forward function is the logits for the predicted
node/graph classes.
See also the example in :func:`explain_node` and :func:`explain_graph`.
num_hops : int
The number of hops for GNN information aggregation.
lr : float, optional
The learning rate to use, default to 0.01.
num_epochs : int, optional
The number of epochs to train.
alpha1 : float, optional
A higher value will make the explanation edge masks more sparse by decreasing
the sum of the edge mask.
alpha2 : float, optional
A higher value will make the explanation edge masks more sparse by decreasing
the entropy of the edge mask.
beta1 : float, optional
A higher value will make the explanation node feature masks more sparse by
decreasing the mean of the node feature mask.
beta2 : float, optional
A higher value will make the explanation node feature masks more sparse by
decreasing the entropy of the node feature mask.
log : bool, optional
If True, it will log the computation process, default to True.
"""
def __init__(
self,
model,
num_hops,
lr=0.01,
num_epochs=100,
*,
alpha1=0.005,
alpha2=1.0,
beta1=1.0,
beta2=0.1,
log=True,
):
super(GNNExplainer, self).__init__()
self.model = model
self.num_hops = num_hops
self.lr = lr
self.num_epochs = num_epochs
self.alpha1 = alpha1
self.alpha2 = alpha2
self.beta1 = beta1
self.beta2 = beta2
self.log = log
def _init_masks(self, graph, feat):
r"""Initialize learnable feature and edge mask.
Parameters
----------
graph : DGLGraph
Input graph.
feat : Tensor
Input node features.
Returns
-------
feat_mask : Tensor
Feature mask of shape :math:`(1, D)`, where :math:`D`
is the feature size.
edge_mask : Tensor
Edge mask of shape :math:`(E)`, where :math:`E` is the
number of edges.
"""
num_nodes, feat_size = feat.size()
num_edges = graph.num_edges()
device = feat.device
std = 0.1
feat_mask = nn.Parameter(torch.randn(1, feat_size, device=device) * std)
std = nn.init.calculate_gain("relu") * sqrt(2.0 / (2 * num_nodes))
edge_mask = nn.Parameter(torch.randn(num_edges, device=device) * std)
return feat_mask, edge_mask
def _loss_regularize(self, loss, feat_mask, edge_mask):
r"""Add regularization terms to the loss.
Parameters
----------
loss : Tensor
Loss value.
feat_mask : Tensor
Feature mask of shape :math:`(1, D)`, where :math:`D`
is the feature size.
edge_mask : Tensor
Edge mask of shape :math:`(E)`, where :math:`E`
is the number of edges.
Returns
-------
Tensor
Loss value with regularization terms added.
"""
# epsilon for numerical stability
eps = 1e-15
edge_mask = edge_mask.sigmoid()
# Edge mask sparsity regularization
loss = loss + self.alpha1 * torch.sum(edge_mask)
# Edge mask entropy regularization
ent = -edge_mask * torch.log(edge_mask + eps) - (
1 - edge_mask
) * torch.log(1 - edge_mask + eps)
loss = loss + self.alpha2 * ent.mean()
feat_mask = feat_mask.sigmoid()
# Feature mask sparsity regularization
loss = loss + self.beta1 * torch.mean(feat_mask)
# Feature mask entropy regularization
ent = -feat_mask * torch.log(feat_mask + eps) - (
1 - feat_mask
) * torch.log(1 - feat_mask + eps)
loss = loss + self.beta2 * ent.mean()
return loss
def explain_node(self, node_id, graph, feat, **kwargs):
r"""Learn and return a node feature mask and subgraph that play a
crucial role to explain the prediction made by the GNN for node
:attr:`node_id`.
Parameters
----------
node_id : int
The node to explain.
graph : DGLGraph
A homogeneous graph.
feat : Tensor
The input feature of shape :math:`(N, D)`. :math:`N` is the
number of nodes, and :math:`D` is the feature size.
kwargs : dict
Additional arguments passed to the GNN model. Tensors whose
first dimension is the number of nodes or edges will be
assumed to be node/edge features.
Returns
-------
new_node_id : Tensor
The new ID of the input center node.
sg : DGLGraph
The subgraph induced on the k-hop in-neighborhood of the input center node.
feat_mask : Tensor
Learned node feature importance mask of shape :math:`(D)`, where :math:`D` is the
feature size. The values are within range :math:`(0, 1)`.
The higher, the more important.
edge_mask : Tensor
Learned importance mask of the edges in the subgraph, which is a tensor
of shape :math:`(E)`, where :math:`E` is the number of edges in the
subgraph. The values are within range :math:`(0, 1)`.
The higher, the more important.
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
>>> import torch.nn as nn
>>> from dgl.data import CoraGraphDataset
>>> from dgl.nn import GNNExplainer
>>> # Load dataset
>>> data = CoraGraphDataset()
>>> g = data[0]
>>> features = g.ndata['feat']
>>> labels = g.ndata['label']
>>> train_mask = g.ndata['train_mask']
>>> # Define a model
>>> class Model(nn.Module):
... def __init__(self, in_feats, out_feats):
... super(Model, self).__init__()
... self.linear = nn.Linear(in_feats, out_feats)
...
... def forward(self, graph, feat, eweight=None):
... with graph.local_scope():
... feat = self.linear(feat)
... graph.ndata['h'] = feat
... if eweight is None:
... graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
... else:
... graph.edata['w'] = eweight
... graph.update_all(fn.u_mul_e('h', 'w', 'm'), fn.sum('m', 'h'))
... return graph.ndata['h']
>>> # Train the model
>>> model = Model(features.shape[1], data.num_classes)
>>> criterion = nn.CrossEntropyLoss()
>>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
>>> for epoch in range(10):
... logits = model(g, features)
... loss = criterion(logits[train_mask], labels[train_mask])
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain the prediction for node 10
>>> explainer = GNNExplainer(model, num_hops=1)
>>> new_center, sg, feat_mask, edge_mask = explainer.explain_node(10, g, features)
>>> new_center
tensor([1])
>>> sg.num_edges()
12
>>> # Old IDs of the nodes in the subgraph
>>> sg.ndata[dgl.NID]
tensor([ 9, 10, 11, 12])
>>> # Old IDs of the edges in the subgraph
>>> sg.edata[dgl.EID]
tensor([51, 53, 56, 48, 52, 57, 47, 50, 55, 46, 49, 54])
>>> feat_mask
tensor([0.2638, 0.2738, 0.3039, ..., 0.2794, 0.2643, 0.2733])
>>> edge_mask
tensor([0.0937, 0.1496, 0.8287, 0.8132, 0.8825, 0.8515, 0.8146, 0.0915, 0.1145,
0.9011, 0.1311, 0.8437])
"""
self.model = self.model.to(graph.device)
self.model.eval()
num_nodes = graph.num_nodes()
num_edges = graph.num_edges()
# Extract node-centered k-hop subgraph and
# its associated node and edge features.
sg, inverse_indices = khop_in_subgraph(graph, node_id, self.num_hops)
sg_nodes = sg.ndata[NID].long()
sg_edges = sg.edata[EID].long()
feat = feat[sg_nodes]
for key, item in kwargs.items():
if torch.is_tensor(item) and item.size(0) == num_nodes:
item = item[sg_nodes]
elif torch.is_tensor(item) and item.size(0) == num_edges:
item = item[sg_edges]
kwargs[key] = item
# Get the initial prediction.
with torch.no_grad():
logits = self.model(graph=sg, feat=feat, **kwargs)
pred_label = logits.argmax(dim=-1)
feat_mask, edge_mask = self._init_masks(sg, feat)
params = [feat_mask, edge_mask]
optimizer = torch.optim.Adam(params, lr=self.lr)
if self.log:
pbar = tqdm(total=self.num_epochs)
pbar.set_description(f"Explain node {node_id}")
for _ in range(self.num_epochs):
optimizer.zero_grad()
h = feat * feat_mask.sigmoid()
logits = self.model(
graph=sg, feat=h, eweight=edge_mask.sigmoid(), **kwargs
)
log_probs = logits.log_softmax(dim=-1)
loss = -log_probs[inverse_indices, pred_label[inverse_indices]]
loss = self._loss_regularize(loss, feat_mask, edge_mask)
loss.backward()
optimizer.step()
if self.log:
pbar.update(1)
if self.log:
pbar.close()
feat_mask = feat_mask.detach().sigmoid().squeeze()
edge_mask = edge_mask.detach().sigmoid()
return inverse_indices, sg, feat_mask, edge_mask
def explain_graph(self, graph, feat, **kwargs):
r"""Learn and return a node feature mask and an edge mask that play a
crucial role to explain the prediction made by the GNN for a graph.
Parameters
----------
graph : DGLGraph
A homogeneous graph.
feat : Tensor
The input feature of shape :math:`(N, D)`. :math:`N` is the
number of nodes, and :math:`D` is the feature size.
kwargs : dict
Additional arguments passed to the GNN model. Tensors whose
first dimension is the number of nodes or edges will be
assumed to be node/edge features.
Returns
-------
feat_mask : Tensor
Learned feature importance mask of shape :math:`(D)`, where :math:`D` is the
feature size. The values are within range :math:`(0, 1)`.
The higher, the more important.
edge_mask : Tensor
Learned importance mask of the edges in the graph, which is a tensor
of shape :math:`(E)`, where :math:`E` is the number of edges in the
graph. The values are within range :math:`(0, 1)`. The higher,
the more important.
Examples
--------
>>> import dgl.function as fn
>>> import torch
>>> import torch.nn as nn
>>> from dgl.data import GINDataset
>>> from dgl.dataloading import GraphDataLoader
>>> from dgl.nn import AvgPooling, GNNExplainer
>>> # Load dataset
>>> data = GINDataset('MUTAG', self_loop=True)
>>> dataloader = GraphDataLoader(data, batch_size=64, shuffle=True)
>>> # Define a model
>>> class Model(nn.Module):
... def __init__(self, in_feats, out_feats):
... super(Model, self).__init__()
... self.linear = nn.Linear(in_feats, out_feats)
... self.pool = AvgPooling()
...
... def forward(self, graph, feat, eweight=None):
... with graph.local_scope():
... feat = self.linear(feat)
... graph.ndata['h'] = feat
... if eweight is None:
... graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
... else:
... graph.edata['w'] = eweight
... graph.update_all(fn.u_mul_e('h', 'w', 'm'), fn.sum('m', 'h'))
... return self.pool(graph, graph.ndata['h'])
>>> # Train the model
>>> feat_size = data[0][0].ndata['attr'].shape[1]
>>> model = Model(feat_size, data.gclasses)
>>> criterion = nn.CrossEntropyLoss()
>>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
>>> for bg, labels in dataloader:
... logits = model(bg, bg.ndata['attr'])
... loss = criterion(logits, labels)
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain the prediction for graph 0
>>> explainer = GNNExplainer(model, num_hops=1)
>>> g, _ = data[0]
>>> features = g.ndata['attr']
>>> feat_mask, edge_mask = explainer.explain_graph(g, features)
>>> feat_mask
tensor([0.2362, 0.2497, 0.2622, 0.2675, 0.2649, 0.2962, 0.2533])
>>> edge_mask
tensor([0.2154, 0.2235, 0.8325, ..., 0.7787, 0.1735, 0.1847])
"""
self.model = self.model.to(graph.device)
self.model.eval()
# Get the initial prediction.
with torch.no_grad():
logits = self.model(graph=graph, feat=feat, **kwargs)
pred_label = logits.argmax(dim=-1)
feat_mask, edge_mask = self._init_masks(graph, feat)
params = [feat_mask, edge_mask]
optimizer = torch.optim.Adam(params, lr=self.lr)
if self.log:
pbar = tqdm(total=self.num_epochs)
pbar.set_description("Explain graph")
for _ in range(self.num_epochs):
optimizer.zero_grad()
h = feat * feat_mask.sigmoid()
logits = self.model(
graph=graph, feat=h, eweight=edge_mask.sigmoid(), **kwargs
)
log_probs = logits.log_softmax(dim=-1)
loss = -log_probs[0, pred_label[0]]
loss = self._loss_regularize(loss, feat_mask, edge_mask)
loss.backward()
optimizer.step()
if self.log:
pbar.update(1)
if self.log:
pbar.close()
feat_mask = feat_mask.detach().sigmoid().squeeze()
edge_mask = edge_mask.detach().sigmoid()
return feat_mask, edge_mask
class HeteroGNNExplainer(nn.Module):
r"""GNNExplainer model from `GNNExplainer: Generating Explanations for
Graph Neural Networks <https://arxiv.org/abs/1903.03894>`__, adapted for heterogeneous graphs
It identifies compact subgraph structures and small subsets of node features that play a
critical role in GNN-based node classification and graph classification.
To generate an explanation, it learns an edge mask :math:`M` and a feature mask :math:`F`
by optimizing the following objective function.
.. math::
l(y, \hat{y}) + \alpha_1 \|M\|_1 + \alpha_2 H(M) + \beta_1 \|F\|_1 + \beta_2 H(F)
where :math:`l` is the loss function, :math:`y` is the original model prediction,
:math:`\hat{y}` is the model prediction with the edge and feature mask applied, :math:`H` is
the entropy function.
Parameters
----------
model : nn.Module
The GNN model to explain.
* The required arguments of its forward function are graph and feat.
The latter one is for input node features.
* It should also optionally take an eweight argument for edge weights
and multiply the messages by it in message passing.
* The output of its forward function is the logits for the predicted
node/graph classes.
See also the example in :func:`explain_node` and :func:`explain_graph`.
num_hops : int
The number of hops for GNN information aggregation.
lr : float, optional
The learning rate to use, default to 0.01.
num_epochs : int, optional
The number of epochs to train.
alpha1 : float, optional
A higher value will make the explanation edge masks more sparse by decreasing
the sum of the edge mask.
alpha2 : float, optional
A higher value will make the explanation edge masks more sparse by decreasing
the entropy of the edge mask.
beta1 : float, optional
A higher value will make the explanation node feature masks more sparse by
decreasing the mean of the node feature mask.
beta2 : float, optional
A higher value will make the explanation node feature masks more sparse by
decreasing the entropy of the node feature mask.
log : bool, optional
If True, it will log the computation process, default to True.
"""
def __init__(
self,
model,
num_hops,
lr=0.01,
num_epochs=100,
*,
alpha1=0.005,
alpha2=1.0,
beta1=1.0,
beta2=0.1,
log=True,
):
super(HeteroGNNExplainer, self).__init__()
self.model = model
self.num_hops = num_hops
self.lr = lr
self.num_epochs = num_epochs
self.alpha1 = alpha1
self.alpha2 = alpha2
self.beta1 = beta1
self.beta2 = beta2
self.log = log
def _init_masks(self, graph, feat):
r"""Initialize learnable feature and edge mask.
Parameters
----------
graph : DGLGraph
Input graph.
feat : dict[str, Tensor]
The dictionary that associates input node features (values) with
the respective node types (keys) present in the graph.
Returns
-------
feat_masks : dict[str, Tensor]
The dictionary that associates the node feature masks (values) with
the respective node types (keys). The feature masks are of shape :math:`(1, D_t)`,
where :math:`D_t` is the feature size for node type :math:`t`.
edge_masks : dict[tuple[str], Tensor]
The dictionary that associates the edge masks (values) with
the respective canonical edge types (keys). The edge masks are of shape :math:`(E_t)`,
where :math:`E_t` is the number of edges for canonical edge type :math:`t`.
"""
device = graph.device
feat_masks = {}
std = 0.1
for node_type, feature in feat.items():
_, feat_size = feature.size()
feat_masks[node_type] = nn.Parameter(
torch.randn(1, feat_size, device=device) * std
)
edge_masks = {}
for canonical_etype in graph.canonical_etypes:
src_num_nodes = graph.num_nodes(canonical_etype[0])
dst_num_nodes = graph.num_nodes(canonical_etype[-1])
num_nodes_sum = src_num_nodes + dst_num_nodes
num_edges = graph.num_edges(canonical_etype)
std = nn.init.calculate_gain("relu")
if num_nodes_sum > 0:
std *= sqrt(2.0 / num_nodes_sum)
edge_masks[canonical_etype] = nn.Parameter(
torch.randn(num_edges, device=device) * std
)
return feat_masks, edge_masks
def _loss_regularize(self, loss, feat_masks, edge_masks):
r"""Add regularization terms to the loss.
Parameters
----------
loss : Tensor
Loss value.
feat_masks : dict[str, Tensor]
The dictionary that associates the node feature masks (values) with
the respective node types (keys). The feature masks are of shape :math:`(1, D_t)`,
where :math:`D_t` is the feature size for node type :math:`t`.
edge_masks : dict[tuple[str], Tensor]
The dictionary that associates the edge masks (values) with
the respective canonical edge types (keys). The edge masks are of shape :math:`(E_t)`,
where :math:`E_t` is the number of edges for canonical edge type :math:`t`.
Returns
-------
Tensor
Loss value with regularization terms added.
"""
# epsilon for numerical stability
eps = 1e-15
for edge_mask in edge_masks.values():
edge_mask = edge_mask.sigmoid()
# Edge mask sparsity regularization
loss = loss + self.alpha1 * torch.sum(edge_mask)
# Edge mask entropy regularization
ent = -edge_mask * torch.log(edge_mask + eps) - (
1 - edge_mask
) * torch.log(1 - edge_mask + eps)
loss = loss + self.alpha2 * ent.mean()
for feat_mask in feat_masks.values():
feat_mask = feat_mask.sigmoid()
# Feature mask sparsity regularization
loss = loss + self.beta1 * torch.mean(feat_mask)
# Feature mask entropy regularization
ent = -feat_mask * torch.log(feat_mask + eps) - (
1 - feat_mask
) * torch.log(1 - feat_mask + eps)
loss = loss + self.beta2 * ent.mean()
return loss
def explain_node(self, ntype, node_id, graph, feat, **kwargs):
r"""Learn and return node feature masks and a subgraph that play a
crucial role to explain the prediction made by the GNN for node
:attr:`node_id` of type :attr:`ntype`.
It requires :attr:`model` to return a dictionary mapping node types to type-specific
predictions.
Parameters
----------
ntype : str
The type of the node to explain. :attr:`model` must be trained to
make predictions for this particular node type.
node_id : int
The ID of the node to explain.
graph : DGLGraph
A heterogeneous graph.
feat : dict[str, Tensor]
The dictionary that associates input node features (values) with
the respective node types (keys) present in the graph.
The input features are of shape :math:`(N_t, D_t)`. :math:`N_t` is the
number of nodes for node type :math:`t`, and :math:`D_t` is the feature size for
node type :math:`t`
kwargs : dict
Additional arguments passed to the GNN model.
Returns
-------
new_node_id : Tensor
The new ID of the input center node.
sg : DGLGraph
The subgraph induced on the k-hop in-neighborhood of the input center node.
feat_mask : dict[str, Tensor]
The dictionary that associates the learned node feature importance masks (values) with
the respective node types (keys). The masks are of shape :math:`(D_t)`, where
:math:`D_t` is the node feature size for node type :attr:`t`. The values are within
range :math:`(0, 1)`. The higher, the more important.
edge_mask : dict[Tuple[str], Tensor]
The dictionary that associates the learned edge importance masks (values) with
the respective canonical edge types (keys). The masks are of shape :math:`(E_t)`,
where :math:`E_t` is the number of edges for canonical edge type :math:`t` in the
subgraph. The values are within range :math:`(0, 1)`.
The higher, the more important.
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch as th
>>> import torch.nn as nn
>>> import torch.nn.functional as F
>>> from dgl.nn import HeteroGNNExplainer
>>> class Model(nn.Module):
... def __init__(self, in_dim, num_classes, canonical_etypes):
... super(Model, self).__init__()
... self.etype_weights = nn.ModuleDict({
... '_'.join(c_etype): nn.Linear(in_dim, num_classes)
... for c_etype in canonical_etypes
... })
...
... def forward(self, graph, feat, eweight=None):
... with graph.local_scope():
... c_etype_func_dict = {}
... for c_etype in graph.canonical_etypes:
... src_type, etype, dst_type = c_etype
... wh = self.etype_weights['_'.join(c_etype)](feat[src_type])
... graph.nodes[src_type].data[f'h_{c_etype}'] = wh
... if eweight is None:
... c_etype_func_dict[c_etype] = (fn.copy_u(f'h_{c_etype}', 'm'),
... fn.mean('m', 'h'))
... else:
... graph.edges[c_etype].data['w'] = eweight[c_etype]
... c_etype_func_dict[c_etype] = (
... fn.u_mul_e(f'h_{c_etype}', 'w', 'm'), fn.mean('m', 'h'))
... graph.multi_update_all(c_etype_func_dict, 'sum')
... return graph.ndata['h']
>>> input_dim = 5
>>> num_classes = 2
>>> g = dgl.heterograph({
... ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1])})
>>> g.nodes['user'].data['h'] = th.randn(g.num_nodes('user'), input_dim)
>>> g.nodes['game'].data['h'] = th.randn(g.num_nodes('game'), input_dim)
>>> transform = dgl.transforms.AddReverse()
>>> g = transform(g)
>>> # define and train the model
>>> model = Model(input_dim, num_classes, g.canonical_etypes)
>>> feat = g.ndata['h']
>>> optimizer = th.optim.Adam(model.parameters())
>>> for epoch in range(10):
... logits = model(g, feat)['user']
... loss = F.cross_entropy(logits, th.tensor([1, 1, 1]))
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain the prediction for node 0 of type 'user'
>>> explainer = HeteroGNNExplainer(model, num_hops=1)
>>> new_center, sg, feat_mask, edge_mask = explainer.explain_node('user', 0, g, feat)
>>> new_center
tensor([0])
>>> sg
Graph(num_nodes={'game': 1, 'user': 1},
num_edges={('game', 'rev_plays', 'user'): 1, ('user', 'plays', 'game'): 1,
('user', 'rev_rev_plays', 'game'): 1},
metagraph=[('game', 'user', 'rev_plays'), ('user', 'game', 'plays'),
('user', 'game', 'rev_rev_plays')])
>>> feat_mask
{'game': tensor([0.2348, 0.2780, 0.2611, 0.2513, 0.2823]),
'user': tensor([0.2716, 0.2450, 0.2658, 0.2876, 0.2738])}
>>> edge_mask
{('game', 'rev_plays', 'user'): tensor([0.0630]),
('user', 'plays', 'game'): tensor([0.1939]),
('user', 'rev_rev_plays', 'game'): tensor([0.9166])}
"""
self.model = self.model.to(graph.device)
self.model.eval()
# Extract node-centered k-hop subgraph and
# its associated node and edge features.
sg, inverse_indices = khop_in_subgraph(
graph, {ntype: node_id}, self.num_hops
)
inverse_indices = inverse_indices[ntype]
sg_nodes = sg.ndata[NID]
sg_feat = {}
for node_type in sg_nodes.keys():
sg_feat[node_type] = feat[node_type][sg_nodes[node_type].long()]
# Get the initial prediction.
with torch.no_grad():
logits = self.model(graph=sg, feat=sg_feat, **kwargs)[ntype]
pred_label = logits.argmax(dim=-1)
feat_mask, edge_mask = self._init_masks(sg, sg_feat)
params = [*feat_mask.values(), *edge_mask.values()]
optimizer = torch.optim.Adam(params, lr=self.lr)
if self.log:
pbar = tqdm(total=self.num_epochs)
pbar.set_description(f"Explain node {node_id} with type {ntype}")
for _ in range(self.num_epochs):
optimizer.zero_grad()
h = {}
for node_type, sg_node_feat in sg_feat.items():
h[node_type] = sg_node_feat * feat_mask[node_type].sigmoid()
eweight = {}
for canonical_etype, canonical_etype_mask in edge_mask.items():
eweight[canonical_etype] = canonical_etype_mask.sigmoid()
logits = self.model(graph=sg, feat=h, eweight=eweight, **kwargs)[
ntype
]
log_probs = logits.log_softmax(dim=-1)
loss = -log_probs[inverse_indices, pred_label[inverse_indices]]
loss = self._loss_regularize(loss, feat_mask, edge_mask)
loss.backward()
optimizer.step()
if self.log:
pbar.update(1)
if self.log:
pbar.close()
for node_type in feat_mask:
feat_mask[node_type] = (
feat_mask[node_type].detach().sigmoid().squeeze()
)
for canonical_etype in edge_mask:
edge_mask[canonical_etype] = (
edge_mask[canonical_etype].detach().sigmoid()
)
return inverse_indices, sg, feat_mask, edge_mask
def explain_graph(self, graph, feat, **kwargs):
r"""Learn and return node feature masks and edge masks that play a
crucial role to explain the prediction made by the GNN for a graph.
Parameters
----------
graph : DGLGraph
A heterogeneous graph that will be explained.
feat : dict[str, Tensor]
The dictionary that associates input node features (values) with
the respective node types (keys) present in the graph.
The input features are of shape :math:`(N_t, D_t)`. :math:`N_t` is the
number of nodes for node type :math:`t`, and :math:`D_t` is the feature size for
node type :math:`t`
kwargs : dict
Additional arguments passed to the GNN model.
Returns
-------
feat_mask : dict[str, Tensor]
The dictionary that associates the learned node feature importance masks (values) with
the respective node types (keys). The masks are of shape :math:`(D_t)`, where
:math:`D_t` is the node feature size for node type :attr:`t`. The values are within
range :math:`(0, 1)`. The higher, the more important.
edge_mask : dict[Tuple[str], Tensor]
The dictionary that associates the learned edge importance masks (values) with
the respective canonical edge types (keys). The masks are of shape :math:`(E_t)`,
where :math:`E_t` is the number of edges for canonical edge type :math:`t` in the
graph. The values are within range :math:`(0, 1)`. The higher, the more important.
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch as th
>>> import torch.nn as nn
>>> import torch.nn.functional as F
>>> from dgl.nn import HeteroGNNExplainer
>>> class Model(nn.Module):
... def __init__(self, in_dim, num_classes, canonical_etypes):
... super(Model, self).__init__()
... self.etype_weights = nn.ModuleDict({
... '_'.join(c_etype): nn.Linear(in_dim, num_classes)
... for c_etype in canonical_etypes
... })
...
... def forward(self, graph, feat, eweight=None):
... with graph.local_scope():
... c_etype_func_dict = {}
... for c_etype in graph.canonical_etypes:
... src_type, etype, dst_type = c_etype
... wh = self.etype_weights['_'.join(c_etype)](feat[src_type])
... graph.nodes[src_type].data[f'h_{c_etype}'] = wh
... if eweight is None:
... c_etype_func_dict[c_etype] = (fn.copy_u(f'h_{c_etype}', 'm'),
... fn.mean('m', 'h'))
... else:
... graph.edges[c_etype].data['w'] = eweight[c_etype]
... c_etype_func_dict[c_etype] = (
... fn.u_mul_e(f'h_{c_etype}', 'w', 'm'), fn.mean('m', 'h'))
... graph.multi_update_all(c_etype_func_dict, 'sum')
... hg = 0
... for ntype in graph.ntypes:
... if graph.num_nodes(ntype):
... hg = hg + dgl.mean_nodes(graph, 'h', ntype=ntype)
... return hg
>>> input_dim = 5
>>> num_classes = 2
>>> g = dgl.heterograph({
... ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1])})
>>> g.nodes['user'].data['h'] = th.randn(g.num_nodes('user'), input_dim)
>>> g.nodes['game'].data['h'] = th.randn(g.num_nodes('game'), input_dim)
>>> transform = dgl.transforms.AddReverse()
>>> g = transform(g)
>>> # define and train the model
>>> model = Model(input_dim, num_classes, g.canonical_etypes)
>>> feat = g.ndata['h']
>>> optimizer = th.optim.Adam(model.parameters())
>>> for epoch in range(10):
... logits = model(g, feat)
... loss = F.cross_entropy(logits, th.tensor([1]))
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain for the graph
>>> explainer = HeteroGNNExplainer(model, num_hops=1)
>>> feat_mask, edge_mask = explainer.explain_graph(g, feat)
>>> feat_mask
{'game': tensor([0.2684, 0.2597, 0.3135, 0.2976, 0.2607]),
'user': tensor([0.2216, 0.2908, 0.2644, 0.2738, 0.2663])}
>>> edge_mask
{('game', 'rev_plays', 'user'): tensor([0.8922, 0.1966, 0.8371, 0.1330]),
('user', 'plays', 'game'): tensor([0.1785, 0.1696, 0.8065, 0.2167])}
"""
self.model = self.model.to(graph.device)
self.model.eval()
# Get the initial prediction.
with torch.no_grad():
logits = self.model(graph=graph, feat=feat, **kwargs)
pred_label = logits.argmax(dim=-1)
feat_mask, edge_mask = self._init_masks(graph, feat)
params = [*feat_mask.values(), *edge_mask.values()]
optimizer = torch.optim.Adam(params, lr=self.lr)
if self.log:
pbar = tqdm(total=self.num_epochs)
pbar.set_description("Explain graph")
for _ in range(self.num_epochs):
optimizer.zero_grad()
h = {}
for node_type, node_feat in feat.items():
h[node_type] = node_feat * feat_mask[node_type].sigmoid()
eweight = {}
for canonical_etype, canonical_etype_mask in edge_mask.items():
eweight[canonical_etype] = canonical_etype_mask.sigmoid()
logits = self.model(graph=graph, feat=h, eweight=eweight, **kwargs)
log_probs = logits.log_softmax(dim=-1)
loss = -log_probs[0, pred_label[0]]
loss = self._loss_regularize(loss, feat_mask, edge_mask)
loss.backward()
optimizer.step()
if self.log:
pbar.update(1)
if self.log:
pbar.close()
for node_type in feat_mask:
feat_mask[node_type] = (
feat_mask[node_type].detach().sigmoid().squeeze()
)
for canonical_etype in edge_mask:
edge_mask[canonical_etype] = (
edge_mask[canonical_etype].detach().sigmoid()
)
return feat_mask, edge_mask
python/dgl/nn/pytorch/explain/pgexplainer.py
+1139
View File
File diff suppressed because it is too large Load Diff
python/dgl/nn/pytorch/explain/subgraphx.py
+807
View File
@@ -0,0 +1,807 @@
"""Torch Module for SubgraphX"""
import math
import networkx as nx
import numpy as np
import torch
import torch.nn as nn
from .... import to_heterogeneous, to_homogeneous
from ....base import NID
from ....convert import to_networkx
from ....subgraph import node_subgraph
from ....transforms.functional import remove_nodes
__all__ = ["SubgraphX", "HeteroSubgraphX"]
class MCTSNode:
r"""Monte Carlo Tree Search Node
Parameters
----------
nodes : Tensor
The node IDs of the graph that are associated with this tree node
"""
def __init__(self, nodes):
self.nodes = nodes
self.num_visit = 0
self.total_reward = 0.0
self.immediate_reward = 0.0
self.children = []
def __repr__(self):
r"""Get the string representation of the node.
Returns
-------
str
The string representation of the node
"""
return str(self.nodes)
class SubgraphX(nn.Module):
r"""SubgraphX from `On Explainability of Graph Neural Networks via Subgraph
Explorations <https://arxiv.org/abs/2102.05152>`
It identifies the most important subgraph from the original graph that
plays a critical role in GNN-based graph classification.
It employs Monte Carlo tree search (MCTS) in efficiently exploring
different subgraphs for explanation and uses Shapley values as the measure
of subgraph importance.
Parameters
----------
model : nn.Module
The GNN model to explain that tackles multiclass graph classification
* Its forward function must have the form
:attr:`forward(self, graph, nfeat)`.
* The output of its forward function is the logits.
num_hops : int
Number of message passing layers in the model
coef : float, optional
This hyperparameter controls the trade-off between exploration and
exploitation. A higher value encourages the algorithm to explore
relatively unvisited nodes. Default: 10.0
high2low : bool, optional
If True, it will use the "High2low" strategy for pruning actions,
expanding children nodes from high degree to low degree when extending
the children nodes in the search tree. Otherwise, it will use the
"Low2high" strategy. Default: True
num_child : int, optional
This is the number of children nodes to expand when extending the
children nodes in the search tree. Default: 12
num_rollouts : int, optional
This is the number of rollouts for MCTS. Default: 20
node_min : int, optional
This is the threshold to define a leaf node based on the number of
nodes in a subgraph. Default: 3
shapley_steps : int, optional
This is the number of steps for Monte Carlo sampling in estimating
Shapley values. Default: 100
log : bool, optional
If True, it will log the progress. Default: False
"""
def __init__(
self,
model,
num_hops,
coef=10.0,
high2low=True,
num_child=12,
num_rollouts=20,
node_min=3,
shapley_steps=100,
log=False,
):
super().__init__()
self.num_hops = num_hops
self.coef = coef
self.high2low = high2low
self.num_child = num_child
self.num_rollouts = num_rollouts
self.node_min = node_min
self.shapley_steps = shapley_steps
self.log = log
self.model = model
def shapley(self, subgraph_nodes):
r"""Compute Shapley value with Monte Carlo approximation.
Parameters
----------
subgraph_nodes : tensor
The tensor node ids of the subgraph that are associated with this
tree node
Returns
-------
float
Shapley value
"""
num_nodes = self.graph.num_nodes()
subgraph_nodes = subgraph_nodes.tolist()
# Obtain neighboring nodes of the subgraph g_i, P'.
local_region = subgraph_nodes
for _ in range(self.num_hops - 1):
in_neighbors, _ = self.graph.in_edges(local_region)
_, out_neighbors = self.graph.out_edges(local_region)
neighbors = torch.cat([in_neighbors, out_neighbors]).tolist()
local_region = list(set(local_region + neighbors))
split_point = num_nodes
coalition_space = list(set(local_region) - set(subgraph_nodes)) + [
split_point
]
marginal_contributions = []
device = self.feat.device
for _ in range(self.shapley_steps):
permuted_space = np.random.permutation(coalition_space)
split_idx = int(np.where(permuted_space == split_point)[0])
selected_nodes = permuted_space[:split_idx]
# Mask for coalition set S_i
exclude_mask = torch.ones(num_nodes)
exclude_mask[local_region] = 0.0
exclude_mask[selected_nodes] = 1.0
# Mask for set S_i and g_i
include_mask = exclude_mask.clone()
include_mask[subgraph_nodes] = 1.0
exclude_feat = self.feat * exclude_mask.unsqueeze(1).to(device)
include_feat = self.feat * include_mask.unsqueeze(1).to(device)
with torch.no_grad():
exclude_probs = self.model(
self.graph, exclude_feat, **self.kwargs
).softmax(dim=-1)
exclude_value = exclude_probs[:, self.target_class]
include_probs = self.model(
self.graph, include_feat, **self.kwargs
).softmax(dim=-1)
include_value = include_probs[:, self.target_class]
marginal_contributions.append(include_value - exclude_value)
return torch.cat(marginal_contributions).mean().item()
def get_mcts_children(self, mcts_node):
r"""Get the children of the MCTS node for the search.
Parameters
----------
mcts_node : MCTSNode
Node in MCTS
Returns
-------
list
Children nodes after pruning
"""
if len(mcts_node.children) > 0:
return mcts_node.children
subg = node_subgraph(self.graph, mcts_node.nodes)
node_degrees = subg.out_degrees() + subg.in_degrees()
k = min(subg.num_nodes(), self.num_child)
chosen_nodes = torch.topk(
node_degrees, k, largest=self.high2low
).indices
mcts_children_maps = dict()
for node in chosen_nodes:
new_subg = remove_nodes(subg, node.to(subg.idtype), store_ids=True)
# Get the largest weakly connected component in the subgraph.
nx_graph = to_networkx(new_subg.cpu())
largest_cc_nids = list(
max(nx.weakly_connected_components(nx_graph), key=len)
)
# Map to the original node IDs.
largest_cc_nids = new_subg.ndata[NID][largest_cc_nids].long()
largest_cc_nids = subg.ndata[NID][largest_cc_nids].sort().values
if str(largest_cc_nids) not in self.mcts_node_maps:
child_mcts_node = MCTSNode(largest_cc_nids)
self.mcts_node_maps[str(child_mcts_node)] = child_mcts_node
else:
child_mcts_node = self.mcts_node_maps[str(largest_cc_nids)]
if str(child_mcts_node) not in mcts_children_maps:
mcts_children_maps[str(child_mcts_node)] = child_mcts_node
mcts_node.children = list(mcts_children_maps.values())
for child_mcts_node in mcts_node.children:
if child_mcts_node.immediate_reward == 0:
child_mcts_node.immediate_reward = self.shapley(
child_mcts_node.nodes
)
return mcts_node.children
def mcts_rollout(self, mcts_node):
r"""Perform a MCTS rollout.
Parameters
----------
mcts_node : MCTSNode
Starting node for MCTS
Returns
-------
float
Reward for visiting the node this time
"""
if len(mcts_node.nodes) <= self.node_min:
return mcts_node.immediate_reward
children_nodes = self.get_mcts_children(mcts_node)
children_visit_sum = sum([child.num_visit for child in children_nodes])
children_visit_sum_sqrt = math.sqrt(children_visit_sum)
chosen_child = max(
children_nodes,
key=lambda c: c.total_reward / max(c.num_visit, 1)
+ self.coef
* c.immediate_reward
* children_visit_sum_sqrt
/ (1 + c.num_visit),
)
reward = self.mcts_rollout(chosen_child)
chosen_child.num_visit += 1
chosen_child.total_reward += reward
return reward
def explain_graph(self, graph, feat, target_class, **kwargs):
r"""Find the most important subgraph from the original graph for the
model to classify the graph into the target class.
Parameters
----------
graph : DGLGraph
A homogeneous graph
feat : Tensor
The input node feature of shape :math:`(N, D)`, :math:`N` is the
number of nodes, and :math:`D` is the feature size
target_class : int
The target class to explain
kwargs : dict
Additional arguments passed to the GNN model
Returns
-------
Tensor
Nodes that represent the most important subgraph
Examples
--------
>>> import torch
>>> import torch.nn as nn
>>> import torch.nn.functional as F
>>> from dgl.data import GINDataset
>>> from dgl.dataloading import GraphDataLoader
>>> from dgl.nn import GraphConv, AvgPooling, SubgraphX
>>> # Define the model
>>> class Model(nn.Module):
... def __init__(self, in_dim, n_classes, hidden_dim=128):
... super().__init__()
... self.conv1 = GraphConv(in_dim, hidden_dim)
... self.conv2 = GraphConv(hidden_dim, n_classes)
... self.pool = AvgPooling()
...
... def forward(self, g, h):
... h = F.relu(self.conv1(g, h))
... h = self.conv2(g, h)
... return self.pool(g, h)
>>> # Load dataset
>>> data = GINDataset('MUTAG', self_loop=True)
>>> dataloader = GraphDataLoader(data, batch_size=64, shuffle=True)
>>> # Train the model
>>> feat_size = data[0][0].ndata['attr'].shape[1]
>>> model = Model(feat_size, data.gclasses)
>>> criterion = nn.CrossEntropyLoss()
>>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
>>> for bg, labels in dataloader:
... logits = model(bg, bg.ndata['attr'])
... loss = criterion(logits, labels)
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Initialize the explainer
>>> explainer = SubgraphX(model, num_hops=2)
>>> # Explain the prediction for graph 0
>>> graph, l = data[0]
>>> graph_feat = graph.ndata.pop("attr")
>>> g_nodes_explain = explainer.explain_graph(graph, graph_feat,
... target_class=l)
"""
self.model.eval()
assert (
graph.num_nodes() > self.node_min
), f"The number of nodes in the\
graph {graph.num_nodes()} should be bigger than {self.node_min}."
self.graph = graph
self.feat = feat
self.target_class = target_class
self.kwargs = kwargs
# book all nodes in MCTS
self.mcts_node_maps = dict()
root = MCTSNode(graph.nodes())
self.mcts_node_maps[str(root)] = root
for i in range(self.num_rollouts):
if self.log:
print(
f"Rollout {i}/{self.num_rollouts}, \
{len(self.mcts_node_maps)} subgraphs have been explored."
)
self.mcts_rollout(root)
best_leaf = None
best_immediate_reward = float("-inf")
for mcts_node in self.mcts_node_maps.values():
if len(mcts_node.nodes) > self.node_min:
continue
if mcts_node.immediate_reward > best_immediate_reward:
best_leaf = mcts_node
best_immediate_reward = best_leaf.immediate_reward
return best_leaf.nodes
class HeteroSubgraphX(nn.Module):
r"""SubgraphX from `On Explainability of Graph Neural Networks via Subgraph
Explorations <https://arxiv.org/abs/2102.05152>`__, adapted for heterogeneous graphs
It identifies the most important subgraph from the original graph that
plays a critical role in GNN-based graph classification.
It employs Monte Carlo tree search (MCTS) in efficiently exploring
different subgraphs for explanation and uses Shapley values as the measure
of subgraph importance.
Parameters
----------
model : nn.Module
The GNN model to explain that tackles multiclass graph classification
* Its forward function must have the form
:attr:`forward(self, graph, nfeat)`.
* The output of its forward function is the logits.
num_hops : int
Number of message passing layers in the model
coef : float, optional
This hyperparameter controls the trade-off between exploration and
exploitation. A higher value encourages the algorithm to explore
relatively unvisited nodes. Default: 10.0
high2low : bool, optional
If True, it will use the "High2low" strategy for pruning actions,
expanding children nodes from high degree to low degree when extending
the children nodes in the search tree. Otherwise, it will use the
"Low2high" strategy. Default: True
num_child : int, optional
This is the number of children nodes to expand when extending the
children nodes in the search tree. Default: 12
num_rollouts : int, optional
This is the number of rollouts for MCTS. Default: 20
node_min : int, optional
This is the threshold to define a leaf node based on the number of
nodes in a subgraph. Default: 3
shapley_steps : int, optional
This is the number of steps for Monte Carlo sampling in estimating
Shapley values. Default: 100
log : bool, optional
If True, it will log the progress. Default: False
"""
def __init__(
self,
model,
num_hops,
coef=10.0,
high2low=True,
num_child=12,
num_rollouts=20,
node_min=3,
shapley_steps=100,
log=False,
):
super().__init__()
self.num_hops = num_hops
self.coef = coef
self.high2low = high2low
self.num_child = num_child
self.num_rollouts = num_rollouts
self.node_min = node_min
self.shapley_steps = shapley_steps
self.log = log
self.model = model
def shapley(self, subgraph_nodes):
r"""Compute Shapley value with Monte Carlo approximation.
Parameters
----------
subgraph_nodes : dict[str, Tensor]
subgraph_nodes[nty] gives the tensor node IDs of node type nty
in the subgraph, which are associated with this tree node
Returns
-------
float
Shapley value
"""
# Obtain neighboring nodes of the subgraph g_i, P'.
local_regions = {
ntype: nodes.tolist() for ntype, nodes in subgraph_nodes.items()
}
for _ in range(self.num_hops - 1):
for c_etype in self.graph.canonical_etypes:
src_ntype, _, dst_ntype = c_etype
if (
src_ntype not in local_regions
or dst_ntype not in local_regions
):
continue
in_neighbors, _ = self.graph.in_edges(
local_regions[dst_ntype], etype=c_etype
)
_, out_neighbors = self.graph.out_edges(
local_regions[src_ntype], etype=c_etype
)
local_regions[src_ntype] = list(
set(local_regions[src_ntype] + in_neighbors.tolist())
)
local_regions[dst_ntype] = list(
set(local_regions[dst_ntype] + out_neighbors.tolist())
)
split_point = self.graph.num_nodes()
coalition_space = {
ntype: list(
set(local_regions[ntype]) - set(subgraph_nodes[ntype].tolist())
)
+ [split_point]
for ntype in subgraph_nodes.keys()
}
marginal_contributions = []
for _ in range(self.shapley_steps):
selected_node_map = dict()
for ntype, nodes in coalition_space.items():
permuted_space = np.random.permutation(nodes)
split_idx = int(np.where(permuted_space == split_point)[0])
selected_node_map[ntype] = permuted_space[:split_idx]
# Mask for coalition set S_i
exclude_mask = {
ntype: torch.ones(self.graph.num_nodes(ntype))
for ntype in self.graph.ntypes
}
for ntype, region in local_regions.items():
exclude_mask[ntype][region] = 0.0
for ntype, selected_nodes in selected_node_map.items():
exclude_mask[ntype][selected_nodes] = 1.0
# Mask for set S_i and g_i
include_mask = {
ntype: exclude_mask[ntype].clone()
for ntype in self.graph.ntypes
}
for ntype, subgn in subgraph_nodes.items():
exclude_mask[ntype][subgn] = 1.0
exclude_feat = {
ntype: self.feat[ntype]
* exclude_mask[ntype].unsqueeze(1).to(self.feat[ntype].device)
for ntype in self.graph.ntypes
}
include_feat = {
ntype: self.feat[ntype]
* include_mask[ntype].unsqueeze(1).to(self.feat[ntype].device)
for ntype in self.graph.ntypes
}
with torch.no_grad():
exclude_probs = self.model(
self.graph, exclude_feat, **self.kwargs
).softmax(dim=-1)
exclude_value = exclude_probs[:, self.target_class]
include_probs = self.model(
self.graph, include_feat, **self.kwargs
).softmax(dim=-1)
include_value = include_probs[:, self.target_class]
marginal_contributions.append(include_value - exclude_value)
return torch.cat(marginal_contributions).mean().item()
def get_mcts_children(self, mcts_node):
r"""Get the children of the MCTS node for the search.
Parameters
----------
mcts_node : MCTSNode
Node in MCTS
Returns
-------
list
Children nodes after pruning
"""
if len(mcts_node.children) > 0:
return mcts_node.children
subg = node_subgraph(self.graph, mcts_node.nodes)
# Choose k nodes based on the highest degree in the subgraph
node_degrees_map = {
ntype: torch.zeros(
subg.num_nodes(ntype), device=subg.nodes(ntype).device
)
for ntype in subg.ntypes
}
for c_etype in subg.canonical_etypes:
src_ntype, _, dst_ntype = c_etype
node_degrees_map[src_ntype] += subg.out_degrees(etype=c_etype)
node_degrees_map[dst_ntype] += subg.in_degrees(etype=c_etype)
node_degrees_list = [
((ntype, i), degree)
for ntype, node_degrees in node_degrees_map.items()
for i, degree in enumerate(node_degrees)
]
node_degrees = torch.stack([v for _, v in node_degrees_list])
k = min(subg.num_nodes(), self.num_child)
chosen_node_indicies = torch.topk(
node_degrees, k, largest=self.high2low
).indices
chosen_nodes = [node_degrees_list[i][0] for i in chosen_node_indicies]
mcts_children_maps = dict()
for ntype, node in chosen_nodes:
new_subg = remove_nodes(subg, node, ntype, store_ids=True)
if new_subg.num_edges() > 0:
new_subg_homo = to_homogeneous(new_subg)
# Get the largest weakly connected component in the subgraph.
nx_graph = to_networkx(new_subg_homo.cpu())
largest_cc_nids = list(
max(nx.weakly_connected_components(nx_graph), key=len)
)
largest_cc_homo = node_subgraph(new_subg_homo, largest_cc_nids)
largest_cc_hetero = to_heterogeneous(
largest_cc_homo, new_subg.ntypes, new_subg.etypes
)
# Follow steps for backtracking to original graph node ids
# 1. retrieve instanced homograph from connected-component homograph
# 2. retrieve instanced heterograph from instanced homograph
# 3. retrieve hetero-subgraph from instanced heterograph
# 4. retrieve orignal graph ids from subgraph node ids
cc_nodes = {
ntype: subg.ndata[NID][ntype][
new_subg.ndata[NID][ntype][
new_subg_homo.ndata[NID][
largest_cc_homo.ndata[NID][indicies]
]
]
]
for ntype, indicies in largest_cc_hetero.ndata[NID].items()
}
else:
available_ntypes = [
ntype
for ntype in new_subg.ntypes
if new_subg.num_nodes(ntype) > 0
]
chosen_ntype = np.random.choice(available_ntypes)
# backtrack from subgraph node ids to entire graph
chosen_node = subg.ndata[NID][chosen_ntype][
np.random.choice(new_subg.nodes[chosen_ntype].data[NID])
]
cc_nodes = {
chosen_ntype: torch.tensor(
[chosen_node],
device=subg.device,
)
}
if str(cc_nodes) not in self.mcts_node_maps:
child_mcts_node = MCTSNode(cc_nodes)
self.mcts_node_maps[str(child_mcts_node)] = child_mcts_node
else:
child_mcts_node = self.mcts_node_maps[str(cc_nodes)]
if str(child_mcts_node) not in mcts_children_maps:
mcts_children_maps[str(child_mcts_node)] = child_mcts_node
mcts_node.children = list(mcts_children_maps.values())
for child_mcts_node in mcts_node.children:
if child_mcts_node.immediate_reward == 0:
child_mcts_node.immediate_reward = self.shapley(
child_mcts_node.nodes
)
return mcts_node.children
def mcts_rollout(self, mcts_node):
r"""Perform a MCTS rollout.
Parameters
----------
mcts_node : MCTSNode
Starting node for MCTS
Returns
-------
float
Reward for visiting the node this time
"""
if (
sum(len(nodes) for nodes in mcts_node.nodes.values())
<= self.node_min
):
return mcts_node.immediate_reward
children_nodes = self.get_mcts_children(mcts_node)
children_visit_sum = sum([child.num_visit for child in children_nodes])
children_visit_sum_sqrt = math.sqrt(children_visit_sum)
chosen_child = max(
children_nodes,
key=lambda c: c.total_reward / max(c.num_visit, 1)
+ self.coef
* c.immediate_reward
* children_visit_sum_sqrt
/ (1 + c.num_visit),
)
reward = self.mcts_rollout(chosen_child)
chosen_child.num_visit += 1
chosen_child.total_reward += reward
return reward
def explain_graph(self, graph, feat, target_class, **kwargs):
r"""Find the most important subgraph from the original graph for the
model to classify the graph into the target class.
Parameters
----------
graph : DGLGraph
A heterogeneous graph
feat : dict[str, Tensor]
The dictionary that associates input node features (values) with
the respective node types (keys) present in the graph.
The input features are of shape :math:`(N_t, D_t)`. :math:`N_t` is the
number of nodes for node type :math:`t`, and :math:`D_t` is the feature size for
node type :math:`t`
target_class : int
The target class to explain
kwargs : dict
Additional arguments passed to the GNN model
Returns
-------
dict[str, Tensor]
The dictionary associating tensor node ids (values) to
node types (keys) that represents the most important subgraph
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch as th
>>> import torch.nn as nn
>>> import torch.nn.functional as F
>>> from dgl.nn import HeteroSubgraphX
>>> class Model(nn.Module):
... def __init__(self, in_dim, num_classes, canonical_etypes):
... super(Model, self).__init__()
... self.etype_weights = nn.ModuleDict(
... {
... "_".join(c_etype): nn.Linear(in_dim, num_classes)
... for c_etype in canonical_etypes
... }
... )
...
... def forward(self, graph, feat):
... with graph.local_scope():
... c_etype_func_dict = {}
... for c_etype in graph.canonical_etypes:
... src_type, etype, dst_type = c_etype
... wh = self.etype_weights["_".join(c_etype)](feat[src_type])
... graph.nodes[src_type].data[f"h_{c_etype}"] = wh
... c_etype_func_dict[c_etype] = (
... fn.copy_u(f"h_{c_etype}", "m"),
... fn.mean("m", "h"),
... )
... graph.multi_update_all(c_etype_func_dict, "sum")
... hg = 0
... for ntype in graph.ntypes:
... if graph.num_nodes(ntype):
... hg = hg + dgl.mean_nodes(graph, "h", ntype=ntype)
... return hg
>>> input_dim = 5
>>> num_classes = 2
>>> g = dgl.heterograph({("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 1, 1])})
>>> g.nodes["user"].data["h"] = th.randn(g.num_nodes("user"), input_dim)
>>> g.nodes["game"].data["h"] = th.randn(g.num_nodes("game"), input_dim)
>>> transform = dgl.transforms.AddReverse()
>>> g = transform(g)
>>> # define and train the model
>>> model = Model(input_dim, num_classes, g.canonical_etypes)
>>> feat = g.ndata["h"]
>>> optimizer = th.optim.Adam(model.parameters())
>>> for epoch in range(10):
... logits = model(g, feat)
... loss = F.cross_entropy(logits, th.tensor([1]))
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain for the graph
>>> explainer = HeteroSubgraphX(model, num_hops=1)
>>> explainer.explain_graph(g, feat, target_class=1)
{'game': tensor([0, 1]), 'user': tensor([1, 2])}
"""
self.model.eval()
assert (
graph.num_nodes() > self.node_min
), f"The number of nodes in the\
graph {graph.num_nodes()} should be bigger than {self.node_min}."
self.graph = graph
self.feat = feat
self.target_class = target_class
self.kwargs = kwargs
# book all nodes in MCTS
self.mcts_node_maps = dict()
root_dict = {ntype: graph.nodes(ntype) for ntype in graph.ntypes}
root = MCTSNode(root_dict)
self.mcts_node_maps[str(root)] = root
for i in range(self.num_rollouts):
if self.log:
print(
f"Rollout {i}/{self.num_rollouts}, \
{len(self.mcts_node_maps)} subgraphs have been explored."
)
self.mcts_rollout(root)
best_leaf = None
best_immediate_reward = float("-inf")
for mcts_node in self.mcts_node_maps.values():
if len(mcts_node.nodes) > self.node_min:
continue
if mcts_node.immediate_reward > best_immediate_reward:
best_leaf = mcts_node
best_immediate_reward = best_leaf.immediate_reward
return best_leaf.nodes