wehub-resource-sync
418 lines
14 KiB
Python
418 lines
14 KiB
Python
from typing import List, Tuple, Union
|
|
|
|
from layers import *
|
|
import dgl.function as fn
|
|
import torch
|
|
import torch.nn
|
|
import torch.nn.functional as F
|
|
from dgl.nn.pytorch.glob import SortPooling
|
|
|
|
|
|
class GraphCrossModule(torch.nn.Module):
|
|
"""
|
|
Description
|
|
-----------
|
|
The Graph Cross Module used by Graph Cross Networks.
|
|
This module only contains graph cross layers.
|
|
|
|
Parameters
|
|
----------
|
|
pool_ratios : Union[float, List[float]]
|
|
The pooling ratios (for keeping nodes) for each layer.
|
|
For example, if `pool_ratio=0.8`, 80\% nodes will be preserved.
|
|
If a single float number is given, all pooling layers will have the
|
|
same pooling ratio.
|
|
in_dim : int
|
|
The number of input node feature channels.
|
|
out_dim : int
|
|
The number of output node feature channels.
|
|
hidden_dim : int
|
|
The number of hidden node feature channels.
|
|
cross_weight : float, optional
|
|
The weight parameter used in graph cross layers
|
|
Default: :obj:`1.0`
|
|
fuse_weight : float, optional
|
|
The weight parameter used at the end of GXN for channel fusion.
|
|
Default: :obj:`1.0`
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pool_ratios: Union[float, List[float]],
|
|
in_dim: int,
|
|
out_dim: int,
|
|
hidden_dim: int,
|
|
cross_weight: float = 1.0,
|
|
fuse_weight: float = 1.0,
|
|
dist: int = 1,
|
|
num_cross_layers: int = 2,
|
|
):
|
|
super(GraphCrossModule, self).__init__()
|
|
if isinstance(pool_ratios, float):
|
|
pool_ratios = (pool_ratios, pool_ratios)
|
|
self.cross_weight = cross_weight
|
|
self.fuse_weight = fuse_weight
|
|
self.num_cross_layers = num_cross_layers
|
|
|
|
# build network
|
|
self.start_gcn_scale1 = GraphConvWithDropout(in_dim, hidden_dim)
|
|
self.start_gcn_scale2 = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
self.end_gcn = GraphConvWithDropout(2 * hidden_dim, out_dim)
|
|
|
|
self.index_select_scale1 = IndexSelect(
|
|
pool_ratios[0], hidden_dim, act="prelu", dist=dist
|
|
)
|
|
self.index_select_scale2 = IndexSelect(
|
|
pool_ratios[1], hidden_dim, act="prelu", dist=dist
|
|
)
|
|
self.start_pool_s12 = GraphPool(hidden_dim)
|
|
self.start_pool_s23 = GraphPool(hidden_dim)
|
|
self.end_unpool_s21 = GraphUnpool(hidden_dim)
|
|
self.end_unpool_s32 = GraphUnpool(hidden_dim)
|
|
|
|
self.s1_l1_gcn = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
self.s1_l2_gcn = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
self.s1_l3_gcn = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
|
|
self.s2_l1_gcn = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
self.s2_l2_gcn = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
self.s2_l3_gcn = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
|
|
self.s3_l1_gcn = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
self.s3_l2_gcn = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
self.s3_l3_gcn = GraphConvWithDropout(hidden_dim, hidden_dim)
|
|
|
|
if num_cross_layers >= 1:
|
|
self.pool_s12_1 = GraphPool(hidden_dim, use_gcn=True)
|
|
self.unpool_s21_1 = GraphUnpool(hidden_dim)
|
|
self.pool_s23_1 = GraphPool(hidden_dim, use_gcn=True)
|
|
self.unpool_s32_1 = GraphUnpool(hidden_dim)
|
|
if num_cross_layers >= 2:
|
|
self.pool_s12_2 = GraphPool(hidden_dim, use_gcn=True)
|
|
self.unpool_s21_2 = GraphUnpool(hidden_dim)
|
|
self.pool_s23_2 = GraphPool(hidden_dim, use_gcn=True)
|
|
self.unpool_s32_2 = GraphUnpool(hidden_dim)
|
|
|
|
def forward(self, graph, feat):
|
|
# start of scale-1
|
|
graph_scale1 = graph
|
|
feat_scale1 = self.start_gcn_scale1(graph_scale1, feat)
|
|
feat_origin = feat_scale1
|
|
feat_scale1_neg = feat_scale1[
|
|
torch.randperm(feat_scale1.size(0))
|
|
] # negative samples
|
|
(
|
|
logit_s1,
|
|
scores_s1,
|
|
select_idx_s1,
|
|
non_select_idx_s1,
|
|
feat_down_s1,
|
|
) = self.index_select_scale1(graph_scale1, feat_scale1, feat_scale1_neg)
|
|
feat_scale2, graph_scale2 = self.start_pool_s12(
|
|
graph_scale1,
|
|
feat_scale1,
|
|
select_idx_s1,
|
|
non_select_idx_s1,
|
|
scores_s1,
|
|
pool_graph=True,
|
|
)
|
|
|
|
# start of scale-2
|
|
feat_scale2 = self.start_gcn_scale2(graph_scale2, feat_scale2)
|
|
feat_scale2_neg = feat_scale2[
|
|
torch.randperm(feat_scale2.size(0))
|
|
] # negative samples
|
|
(
|
|
logit_s2,
|
|
scores_s2,
|
|
select_idx_s2,
|
|
non_select_idx_s2,
|
|
feat_down_s2,
|
|
) = self.index_select_scale2(graph_scale2, feat_scale2, feat_scale2_neg)
|
|
feat_scale3, graph_scale3 = self.start_pool_s23(
|
|
graph_scale2,
|
|
feat_scale2,
|
|
select_idx_s2,
|
|
non_select_idx_s2,
|
|
scores_s2,
|
|
pool_graph=True,
|
|
)
|
|
|
|
# layer-1
|
|
res_s1_0, res_s2_0, res_s3_0 = feat_scale1, feat_scale2, feat_scale3
|
|
|
|
feat_scale1 = F.relu(self.s1_l1_gcn(graph_scale1, feat_scale1))
|
|
feat_scale2 = F.relu(self.s2_l1_gcn(graph_scale2, feat_scale2))
|
|
feat_scale3 = F.relu(self.s3_l1_gcn(graph_scale3, feat_scale3))
|
|
|
|
if self.num_cross_layers >= 1:
|
|
feat_s12_fu = self.pool_s12_1(
|
|
graph_scale1,
|
|
feat_scale1,
|
|
select_idx_s1,
|
|
non_select_idx_s1,
|
|
scores_s1,
|
|
)
|
|
feat_s21_fu = self.unpool_s21_1(
|
|
graph_scale1, feat_scale2, select_idx_s1
|
|
)
|
|
feat_s23_fu = self.pool_s23_1(
|
|
graph_scale2,
|
|
feat_scale2,
|
|
select_idx_s2,
|
|
non_select_idx_s2,
|
|
scores_s2,
|
|
)
|
|
feat_s32_fu = self.unpool_s32_1(
|
|
graph_scale2, feat_scale3, select_idx_s2
|
|
)
|
|
|
|
feat_scale1 = (
|
|
feat_scale1 + self.cross_weight * feat_s21_fu + res_s1_0
|
|
)
|
|
feat_scale2 = (
|
|
feat_scale2
|
|
+ self.cross_weight * (feat_s12_fu + feat_s32_fu) / 2
|
|
+ res_s2_0
|
|
)
|
|
feat_scale3 = (
|
|
feat_scale3 + self.cross_weight * feat_s23_fu + res_s3_0
|
|
)
|
|
|
|
# layer-2
|
|
feat_scale1 = F.relu(self.s1_l2_gcn(graph_scale1, feat_scale1))
|
|
feat_scale2 = F.relu(self.s2_l2_gcn(graph_scale2, feat_scale2))
|
|
feat_scale3 = F.relu(self.s3_l2_gcn(graph_scale3, feat_scale3))
|
|
|
|
if self.num_cross_layers >= 2:
|
|
feat_s12_fu = self.pool_s12_2(
|
|
graph_scale1,
|
|
feat_scale1,
|
|
select_idx_s1,
|
|
non_select_idx_s1,
|
|
scores_s1,
|
|
)
|
|
feat_s21_fu = self.unpool_s21_2(
|
|
graph_scale1, feat_scale2, select_idx_s1
|
|
)
|
|
feat_s23_fu = self.pool_s23_2(
|
|
graph_scale2,
|
|
feat_scale2,
|
|
select_idx_s2,
|
|
non_select_idx_s2,
|
|
scores_s2,
|
|
)
|
|
feat_s32_fu = self.unpool_s32_2(
|
|
graph_scale2, feat_scale3, select_idx_s2
|
|
)
|
|
|
|
cross_weight = self.cross_weight * 0.05
|
|
feat_scale1 = feat_scale1 + cross_weight * feat_s21_fu
|
|
feat_scale2 = (
|
|
feat_scale2 + cross_weight * (feat_s12_fu + feat_s32_fu) / 2
|
|
)
|
|
feat_scale3 = feat_scale3 + cross_weight * feat_s23_fu
|
|
|
|
# layer-3
|
|
feat_scale1 = F.relu(self.s1_l3_gcn(graph_scale1, feat_scale1))
|
|
feat_scale2 = F.relu(self.s2_l3_gcn(graph_scale2, feat_scale2))
|
|
feat_scale3 = F.relu(self.s3_l3_gcn(graph_scale3, feat_scale3))
|
|
|
|
# final layers
|
|
feat_s3_out = (
|
|
self.end_unpool_s32(graph_scale2, feat_scale3, select_idx_s2)
|
|
+ feat_down_s2
|
|
)
|
|
feat_s2_out = self.end_unpool_s21(
|
|
graph_scale1, feat_scale2 + feat_s3_out, select_idx_s1
|
|
)
|
|
feat_agg = (
|
|
feat_scale1
|
|
+ self.fuse_weight * feat_s2_out
|
|
+ self.fuse_weight * feat_down_s1
|
|
)
|
|
feat_agg = torch.cat((feat_agg, feat_origin), dim=1)
|
|
feat_agg = self.end_gcn(graph_scale1, feat_agg)
|
|
|
|
return feat_agg, logit_s1, logit_s2
|
|
|
|
|
|
class GraphCrossNet(torch.nn.Module):
|
|
"""
|
|
Description
|
|
-----------
|
|
The Graph Cross Network.
|
|
|
|
Parameters
|
|
----------
|
|
in_dim : int
|
|
The number of input node feature channels.
|
|
out_dim : int
|
|
The number of output node feature channels.
|
|
edge_feat_dim : int, optional
|
|
The number of input edge feature channels. Edge feature
|
|
will be passed to a Linear layer and concatenated to
|
|
input node features. Default: :obj:`0`
|
|
hidden_dim : int, optional
|
|
The number of hidden node feature channels.
|
|
Default: :obj:`96`
|
|
pool_ratios : Union[float, List[float]], optional
|
|
The pooling ratios (for keeping nodes) for each layer.
|
|
For example, if `pool_ratio=0.8`, 80\% nodes will be preserved.
|
|
If a single float number is given, all pooling layers will have the
|
|
same pooling ratio.
|
|
Default: :obj:`[0.9, 0.7]`
|
|
readout_nodes : int, optional
|
|
Number of nodes perserved in the final sort pool operation.
|
|
Default: :obj:`30`
|
|
conv1d_dims : List[int], optional
|
|
The number of kernels of Conv1d operations.
|
|
Default: :obj:`[16, 32]`
|
|
conv1d_kws : List[int], optional
|
|
The kernel size of Conv1d.
|
|
Default: :obj:`[5]`
|
|
cross_weight : float, optional
|
|
The weight parameter used in graph cross layers
|
|
Default: :obj:`1.0`
|
|
fuse_weight : float, optional
|
|
The weight parameter used at the end of GXN for channel fusion.
|
|
Default: :obj:`1.0`
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_dim: int,
|
|
out_dim: int,
|
|
edge_feat_dim: int = 0,
|
|
hidden_dim: int = 96,
|
|
pool_ratios: Union[List[float], float] = [0.9, 0.7],
|
|
readout_nodes: int = 30,
|
|
conv1d_dims: List[int] = [16, 32],
|
|
conv1d_kws: List[int] = [5],
|
|
cross_weight: float = 1.0,
|
|
fuse_weight: float = 1.0,
|
|
dist: int = 1,
|
|
):
|
|
super(GraphCrossNet, self).__init__()
|
|
self.in_dim = in_dim
|
|
self.out_dim = out_dim
|
|
self.hidden_dim = hidden_dim
|
|
self.edge_feat_dim = edge_feat_dim
|
|
self.readout_nodes = readout_nodes
|
|
conv1d_kws = [hidden_dim] + conv1d_kws
|
|
|
|
if edge_feat_dim > 0:
|
|
self.in_dim += hidden_dim
|
|
self.e2l_lin = torch.nn.Linear(edge_feat_dim, hidden_dim)
|
|
else:
|
|
self.e2l_lin = None
|
|
|
|
self.gxn = GraphCrossModule(
|
|
pool_ratios,
|
|
in_dim=self.in_dim,
|
|
out_dim=hidden_dim,
|
|
hidden_dim=hidden_dim // 2,
|
|
cross_weight=cross_weight,
|
|
fuse_weight=fuse_weight,
|
|
dist=dist,
|
|
)
|
|
self.sortpool = SortPooling(readout_nodes)
|
|
|
|
# final updates
|
|
self.final_conv1 = torch.nn.Conv1d(
|
|
1, conv1d_dims[0], kernel_size=conv1d_kws[0], stride=conv1d_kws[0]
|
|
)
|
|
self.final_maxpool = torch.nn.MaxPool1d(2, 2)
|
|
self.final_conv2 = torch.nn.Conv1d(
|
|
conv1d_dims[0], conv1d_dims[1], kernel_size=conv1d_kws[1], stride=1
|
|
)
|
|
self.final_dense_dim = int((readout_nodes - 2) / 2 + 1)
|
|
self.final_dense_dim = (
|
|
self.final_dense_dim - conv1d_kws[1] + 1
|
|
) * conv1d_dims[1]
|
|
|
|
if self.out_dim > 0:
|
|
self.out_lin = torch.nn.Linear(self.final_dense_dim, out_dim)
|
|
|
|
self.init_weights()
|
|
|
|
def init_weights(self):
|
|
if self.e2l_lin is not None:
|
|
torch.nn.init.xavier_normal_(self.e2l_lin.weight)
|
|
torch.nn.init.xavier_normal_(self.final_conv1.weight)
|
|
torch.nn.init.xavier_normal_(self.final_conv2.weight)
|
|
if self.out_dim > 0:
|
|
torch.nn.init.xavier_normal_(self.out_lin.weight)
|
|
|
|
def forward(
|
|
self,
|
|
graph: DGLGraph,
|
|
node_feat: Tensor,
|
|
edge_feat: Optional[Tensor] = None,
|
|
):
|
|
num_batch = graph.batch_size
|
|
if edge_feat is not None:
|
|
edge_feat = self.e2l_lin(edge_feat)
|
|
with graph.local_scope():
|
|
graph.edata["he"] = edge_feat
|
|
graph.update_all(fn.copy_e("he", "m"), fn.sum("m", "hn"))
|
|
edge2node_feat = graph.ndata.pop("hn")
|
|
node_feat = torch.cat((node_feat, edge2node_feat), dim=1)
|
|
|
|
node_feat, logits1, logits2 = self.gxn(graph, node_feat)
|
|
batch_sortpool_feats = self.sortpool(graph, node_feat)
|
|
|
|
# final updates
|
|
to_conv1d = batch_sortpool_feats.unsqueeze(1)
|
|
conv1d_result = F.relu(self.final_conv1(to_conv1d))
|
|
conv1d_result = self.final_maxpool(conv1d_result)
|
|
conv1d_result = F.relu(self.final_conv2(conv1d_result))
|
|
|
|
to_dense = conv1d_result.view(num_batch, -1)
|
|
if self.out_dim > 0:
|
|
out = F.relu(self.out_lin(to_dense))
|
|
else:
|
|
out = to_dense
|
|
|
|
return out, logits1, logits2
|
|
|
|
|
|
class GraphClassifier(torch.nn.Module):
|
|
"""
|
|
Description
|
|
-----------
|
|
Graph Classifier for graph classification.
|
|
GXN + MLP
|
|
"""
|
|
|
|
def __init__(self, args):
|
|
super(GraphClassifier, self).__init__()
|
|
self.gxn = GraphCrossNet(
|
|
in_dim=args.in_dim,
|
|
out_dim=args.embed_dim,
|
|
edge_feat_dim=args.edge_feat_dim,
|
|
hidden_dim=args.hidden_dim,
|
|
pool_ratios=args.pool_ratios,
|
|
readout_nodes=args.readout_nodes,
|
|
conv1d_dims=args.conv1d_dims,
|
|
conv1d_kws=args.conv1d_kws,
|
|
cross_weight=args.cross_weight,
|
|
fuse_weight=args.fuse_weight,
|
|
)
|
|
self.lin1 = torch.nn.Linear(args.embed_dim, args.final_dense_hidden_dim)
|
|
self.lin2 = torch.nn.Linear(args.final_dense_hidden_dim, args.out_dim)
|
|
self.dropout = args.dropout
|
|
|
|
def forward(
|
|
self,
|
|
graph: DGLGraph,
|
|
node_feat: Tensor,
|
|
edge_feat: Optional[Tensor] = None,
|
|
):
|
|
embed, logits1, logits2 = self.gxn(graph, node_feat, edge_feat)
|
|
logits = F.relu(self.lin1(embed))
|
|
if self.dropout > 0:
|
|
logits = F.dropout(logits, p=self.dropout, training=self.training)
|
|
logits = self.lin2(logits)
|
|
return F.log_softmax(logits, dim=1), logits1, logits2
|