wehub-resource-sync
808 lines
29 KiB
Python
808 lines
29 KiB
Python
"""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
|