All in One: Multi-Task Prompting for Graph Neural Networks学习笔记

今天阅读的是All in One: Multi-Task Prompting for Graph Neural Networks这一篇论文的代码，讨论提示图（Prompt Graph）的生成过程，以及它如何与原始图结合。

提示图的生成主要通过LightPrompt和HeavyPrompt类来实现。在代码中的体现是：

1. LightPrompt类
   LightPrompt类负责生成提示图的节点（tokens）及其内部连接:

  class LightPrompt(torch.nn.Module):

      def __init__(self, token_dim, token_num_per_group, group_num=1, inner_prune=None):

          super(LightPrompt, self).__init__()

          self.inner_prune = inner_prune

          self.token_list = torch.nn.ParameterList(
              [torch.nn.Parameter(torch.empty(token_num_per_group, token_dim)) for i in range(group_num)])

          self.token_init(init_method="kaiming_uniform")

      def token_init(self, init_method="kaiming_uniform"):
          if init_method == "kaiming_uniform":
              for token in self.token_list:
                  torch.nn.init.kaiming_uniform_(token, nonlinearity='leaky_relu', mode='fan_in', a=0.01)
          else:
              raise ValueError("only support kaiming_uniform init, more init methods will be included soon")

      def inner_structure_update(self):
          return self.token_view()

      def token_view(self):
          pg_list = []
          for i, tokens in enumerate(self.token_list):
              token_dot = torch.mm(tokens, torch.transpose(tokens, 0, 1))
              token_sim = torch.sigmoid(token_dot)
              inner_adj = torch.where(token_sim < self.inner_prune, 0, token_sim)
              edge_index = inner_adj.nonzero().t().contiguous()
              pg_list.append(Data(x=tokens, edge_index=edge_index, y=torch.tensor([i]).long()))

          pg_batch = Batch.from_data_list(pg_list)
          return pg_batch

初始化：在__init__方法中，创建一个torch.nn.ParameterList，其中包含多个torch.nn.Parameter对象，每个对象代表一组tokens。
参数初始化：在token_init方法中，对这些tokens进行参数初始化（这里使用了Kaiming Uniform初始化）。
内部结构更新：在inner_structure_update方法中，调用token_view方法生成提示图的内部结构。

2. HeavyPrompt类
   HeavyPrompt类继承自LightPrompt，并增加了跨图连接的功能。

	class HeavyPrompt(LightPrompt):

		def __init__(self, token_dim, token_num, cross_prune=0.1, inner_prune=0.01):

			super(HeavyPrompt, self).__init__(token_dim, token_num, 1, inner_prune)

			self.cross_prune = cross_prune

		def forward(self, graph_batch: Batch):
			device = torch.device("cpu")

			pg = self.inner_structure_update()

			inner_edge_index = pg.edge_index
			token_num = pg.x.shape[0]

			re_graph_list = []
			for g in Batch.to_data_list(graph_batch):
				g_edge_index = g.edge_index + token_num
				pg_x = pg.x.to(device)
				g_x = g.x.to(device)

				cross_dot = torch.mm(pg.x, torch.transpose(g.x, 0, 1))
				cross_sim = torch.sigmoid(cross_dot)
				cross_adj = torch.where(cross_sim < self.cross_prune, 0, cross_sim)

				cross_edge_index = cross_adj.nonzero().t().contiguous()
				cross_edge_index[1] = cross_edge_index[1] + token_num

				x = torch.cat([pg.x, g.x], dim=0)
				y = g.y

				edge_index = torch.cat([inner_edge_index, g_edge_index, cross_edge_index], dim=1)
				data = Data(x=x, edge_index=edge_index, y=y)
				re_graph_list.append(data)

			graphp_batch = Batch.from_data_list(re_graph_list)
			return graphp_batch

跨图连接：在forward方法中，生成提示图与输入图之间的跨图连接。

提示图与原始图的结合过程:
提示图与原始图的结合过程主要在HeavyPrompt类的forward方法中进行。
首先，通过inner_structure_update方法生成提示图的内部结构。这一步会调用LightPrompt类中的token_view方法:

inner_edge_index = pg.edge_index
token_num = pg.x.shape[0]

pg 是生成的提示图的批次。
inner_edge_index 是提示图内部的边索引。
token_num 是提示图中节点的数量。

2. 遍历原始图的批次
   接下来，遍历原始图的批次，将每个原始图与提示图结合生成新的图。

re_graph_list = []
for g in Batch.to_data_list(graph_batch):
    g_edge_index = g.edge_index + token_num
    pg_x = pg.x.to(device)
    g_x = g.x.to(device)

graph_batch 是原始图的批次。
Batch.to_data_list(graph_batch) 将批次转换为单个图的列表。
g 是原始图中的一个图。
g_edge_index 是原始图的边索引，所有边索引都加上token_num，以避免与提示图的节点索引冲突。
pg_x 和 g_x 分别是提示图和原始图的节点特征矩阵，转换到相同的设备上。

3. 计算提示图与原始图的跨图相似度
   计算提示图与原始图之间的相似度（内积），并生成跨图的边。

cross_dot = torch.mm(pg.x, torch.transpose(g.x, 0, 1))
cross_sim = torch.sigmoid(cross_dot)
cross_adj = torch.where(cross_sim < self.cross_prune, 0, cross_sim)
cross_edge_index = cross_adj.nonzero().t().contiguous()
cross_edge_index[1] = cross_edge_index[1] + token_num

cross_dot 是提示图节点特征与原始图节点特征的内积矩阵。
cross_sim 是通过sigmoid函数将内积矩阵映射到0到1之间的相似度矩阵。
cross_adj 是根据cross_prune阈值进行剪枝后的相似度矩阵，低于阈值的相似度被置为0。
cross_edge_index 是跨图的边索引。注意，这里将跨图边的目标节点索引加上token_num，以匹配原始图节点的索引。

4. 拼接节点特征和边索引
   将提示图和原始图的节点特征和边索引拼接在一起，形成新的图。

x = torch.cat([pg.x, g.x], dim=0)
y = g.y
edge_index = torch.cat([inner_edge_index, g_edge_index, cross_edge_index], dim=1)
data = Data(x=x, edge_index=edge_index, y=y)
re_graph_list.append(data)

x 是提示图和原始图的节点特征矩阵拼接在一起。
y 是原始图的标签。
edge_index 是提示图的内部边、原始图的边和跨图的边拼接在一起的边索引矩阵。
data 是新的图数据对象，包含节点特征、边索引和标签。

5. 返回新的图批次

graphp_batch = Batch.from_data_list(re_graph_list)