def __init__(self, in_channels, out_channels, use_weight=True, use_init=False):

super(GraphConvLayer, self).__init__()

self.use_init = use_init

self.use_weight = use_weight

if self.use_init:

in_channels_ = 2 * in_channels

else:

in_channels_ = in_channels

self.W = nn.Linear(in_channels_, out_channels)

def reset_parameters(self):

self.W.reset_parameters()

def forward(self, x, edge_index, x0):

N = x.shape[0]

# Store original device

device = x.device

row, col = edge_index

# torch_sparse requires indices, values, and input matrix to be on CPU

row_cpu = row.cpu()

col_cpu = col.cpu()

x_cpu = x.cpu()

d = degree(col_cpu, N).float()

d_norm_in = (1. / d[col_cpu]).sqrt()

d_norm_out = (1. / d[row_cpu]).sqrt()

value = torch.ones_like(row_cpu) * d_norm_in * d_norm_out

value = torch.nan_to_num(value, nan=0.0, posinf=0.0, neginf=0.0)

# Keep everything on CPU for torch_sparse operations

adj = SparseTensor(row=col_cpu, col=row_cpu, value=value, sparse_sizes=(N, N))

x = matmul(adj, x_cpu) # [N, D]

# Move result back to original device

x = x.to(device)

if self.use_init:

x = torch.cat([x, x0], 1)

x = self.W(x)

elif self.use_weight:

x = self.W(x)

def __init__(self, in_channels, hidden_channels, num_layers=2, dropout=0.5, use_bn=True, use_residual=True, use_weight=True, use_init=False, use_act=True):

super(GraphConv, self).__init__()

self.convs = nn.ModuleList()

self.fcs = nn.ModuleList()

self.fcs.append(nn.Linear(in_channels, hidden_channels))

self.bns = nn.ModuleList()

self.bns.append(nn.BatchNorm1d(hidden_channels))

for _ in range(num_layers):

self.convs.append(

GraphConvLayer(hidden_channels, hidden_channels, use_weight, use_init))

self.bns.append(nn.BatchNorm1d(hidden_channels))

self.dropout = dropout

self.activation = F.relu

self.use_bn = use_bn

self.use_residual = use_residual

self.use_act = use_act

def reset_parameters(self):

for conv in self.convs:

conv.reset_parameters()

for bn in self.bns:

bn.reset_parameters()

for fc in self.fcs:

fc.reset_parameters()

def forward(self, x, edge_index):

layer_ = []

x = self.fcs[0](x)

if self.use_bn:

x = self.bns[0](x)

x = self.activation(x)

x = F.dropout(x, p=self.dropout, training=self.training)

layer_.append(x)

for i, conv in enumerate(self.convs):

x = conv(x, edge_index, layer_[0])

if self.use_bn:

x = self.bns[i+1](x)

if self.use_act:

x = self.activation(x)

x = F.dropout(x, p=self.dropout, training=self.training)

if self.use_residual:

x = x + layer_[-1]

'''

def __init__(self, in_channels,

out_channels,

num_heads,

use_weight=True):

super().__init__()

self.Wk = nn.Linear(in_channels, out_channels * num_heads)

self.Wq = nn.Linear(in_channels, out_channels * num_heads)

if use_weight:

self.Wv = nn.Linear(in_channels, out_channels * num_heads)

self.out_channels = out_channels

self.num_heads = num_heads

self.use_weight = use_weight

def reset_parameters(self):

self.Wk.reset_parameters()

self.Wq.reset_parameters()

if self.use_weight:

self.Wv.reset_parameters()

def forward(self, query_input, source_input, output_attn=False):

# feature transformation

qs = self.Wq(query_input).reshape(-1, self.num_heads, self.out_channels)

ks = self.Wk(source_input).reshape(-1, self.num_heads, self.out_channels)

if self.use_weight:

vs = self.Wv(source_input).reshape(-1, self.num_heads, self.out_channels)

else:

vs = source_input.reshape(-1, 1, self.out_channels)

# normalize input

qs = qs / torch.norm(qs, p=2) # [N, H, M]

ks = ks / torch.norm(ks, p=2) # [L, H, M]

N = qs.shape[0]

# # numerator

# kvs = torch.einsum("lhm,lhd->hmd", ks, vs)

# attention_num = torch.einsum("nhm,hmd->nhd", qs, kvs) # [N, H, D]

# attention_num += N * vs

#

# # denominator

# all_ones = torch.ones([ks.shape[0]]).to(ks.device)

# ks_sum = torch.einsum("lhm,l->hm", ks, all_ones)

# attention_normalizer = torch.einsum("nhm,hm->nh", qs, ks_sum) # [N, H]

#

# # attentive aggregated results

# attention_normalizer = torch.unsqueeze(

# attention_normalizer, len(attention_normalizer.shape)) # [N, H, 1]

# attention_normalizer += torch.ones_like(attention_normalizer) * N

# attn_output = attention_num / attention_normalizer # [N, H, D]

# numerator

attention_num = torch.sigmoid(torch.einsum("nhm,lhm->nlh", qs, ks)) # [N, L, H]

# denominator

all_ones = torch.ones([ks.shape[0]]).to(ks.device)

attention_normalizer = torch.einsum("nlh,l->nh", attention_num, all_ones)

attention_normalizer = attention_normalizer.unsqueeze(1).repeat(1, ks.shape[0], 1) # [N, L, H]

# compute attention and attentive aggregated results

attention = attention_num / attention_normalizer

attn_output = torch.einsum("nlh,lhd->nhd", attention, vs) # [N, H, D]

# compute attention for visualization if needed

if output_attn:

attention = torch.einsum("nhm,lhm->nlh", qs, ks).mean(dim=-1) # [N, N]

normalizer = attention_normalizer.squeeze(dim=-1).mean(dim=-1, keepdims=True) # [N,1]

attention = attention / normalizer

final_output = attn_output.mean(dim=1)

if output_attn:

return final_output, attention

else:

def __init__(self, in_channels, hidden_channels, num_layers=2, num_heads=1,

dropout=0.5, use_bn=True, use_residual=True, use_weight=True, use_act=True):

super().__init__()

self.convs = nn.ModuleList()

self.fcs = nn.ModuleList()

self.fcs.append(nn.Linear(in_channels, hidden_channels))

self.bns = nn.ModuleList()

self.bns.append(nn.LayerNorm(hidden_channels))

for i in range(num_layers):

self.convs.append(

TransConvLayer(hidden_channels, hidden_channels, num_heads=num_heads, use_weight=use_weight))

self.bns.append(nn.LayerNorm(hidden_channels))

self.dropout = dropout

self.activation = F.relu

self.use_bn = use_bn

self.use_residual = use_residual

self.use_act = use_act

def reset_parameters(self):

for conv in self.convs:

conv.reset_parameters()

for bn in self.bns:

bn.reset_parameters()

for fc in self.fcs:

fc.reset_parameters()

def forward(self, x):

layer_ = []

# input MLP layer

x = self.fcs[0](x)

if self.use_bn:

x = self.bns[0](x)

x = self.activation(x)

x = F.dropout(x, p=self.dropout, training=self.training)

# store as residual link

layer_.append(x)

for i, conv in enumerate(self.convs):

# graph convolution with full attention aggregation

x = conv(x, x)

if self.use_residual:

x = (x + layer_[i]) / 2.

if self.use_bn:

x = self.bns[i + 1](x)

if self.use_act:

x = self.activation(x)

x = F.dropout(x, p=self.dropout, training=self.training)

layer_.append(x)

return x

def get_attentions(self, x):

layer_, attentions = [], []

x = self.fcs[0](x)

if self.use_bn:

x = self.bns[0](x)

x = self.activation(x)

layer_.append(x)

for i, conv in enumerate(self.convs):

x, attn = conv(x, x, output_attn=True)

attentions.append(attn)

if self.use_residual:

x = (x + layer_[i]) / 2.

if self.use_bn:

x = self.bns[i + 1](x)

if self.use_act:

x = self.activation(x)

layer_.append(x)

def __init__(self, in_channels, hidden_channels, out_channels,

trans_num_layers=1, trans_num_heads=1, trans_dropout=0.5, trans_use_bn=True, trans_use_residual=True, trans_use_weight=True, trans_use_act=True,

gnn_num_layers=1, gnn_dropout=0.5, gnn_use_weight=True, gnn_use_init=False, gnn_use_bn=True, gnn_use_residual=True, gnn_use_act=True,

use_graph=True, graph_weight=0.8, aggregate='add'):

super().__init__()

self.trans_conv = TransConv(in_channels, hidden_channels, trans_num_layers, trans_num_heads, trans_dropout, trans_use_bn, trans_use_residual, trans_use_weight, trans_use_act)

self.graph_conv = GraphConv(in_channels, hidden_channels, gnn_num_layers, gnn_dropout, gnn_use_bn, gnn_use_residual, gnn_use_weight, gnn_use_init, gnn_use_act)

self.use_graph = use_graph

self.graph_weight = graph_weight

self.aggregate = aggregate

if aggregate == 'add':

self.fc = nn.Linear(hidden_channels, out_channels)

elif aggregate == 'cat':

self.fc = nn.Linear(2 * hidden_channels, out_channels)

else:

raise ValueError(f'Invalid aggregate type:{aggregate}')

self.params1 = list(self.trans_conv.parameters())

self.params2 = list(self.graph_conv.parameters()) if self.graph_conv is not None else []

self.params2.extend(list(self.fc.parameters()))

def forward(self, x, edge_index,batch):

x1 = self.trans_conv(x)

if self.use_graph:

x2 = self.graph_conv(x, edge_index)

if self.aggregate == 'add':

x = self.graph_weight * x2 + (1 - self.graph_weight) * x1

else:

x = torch.cat((x1, x2), dim=1)

else:

x = x1

x = self.fc(x)

# 汇聚步骤

if self.use_graph:

x = torch.cat((x1, x2), dim=1) if self.aggregate == 'cat' else self.graph_weight * x2 + (1 - self.graph_weight) * x1

else:

x = x1

# 使用全局平均汇聚

x = global_mean_pool(x, batch)

# 全连接层

x = self.fc(x)

return x

# def forward(self, x, edge_index, batch):

# x = self.trans_conv(x) # 使用 Transformer 部分

# x = global_mean_pool(x, batch) # 汇聚节点特征到图特征

# x = self.fc(x) # 全连接层

# return x

# def forward(self, x, edge_index, batch):

# x = self.graph_conv(x, edge_index) # 使用 GNN 部分

# x = global_mean_pool(x, batch) # 汇聚节点特征到图特征

# x = self.fc(x) # 全连接层

# return x

def get_attentions(self, x):

attns = self.trans_conv.get_attentions(x) # [layer num, N, N]

return attns

def reset_parameters(self):

self.trans_conv.reset_parameters()

if self.use_graph:

self.graph_conv.reset_parameters()