108 lines
4.5 KiB
Python
108 lines
4.5 KiB
Python
import argparse
|
||
import random
|
||
from collections import defaultdict
|
||
|
||
import dgl
|
||
import torch
|
||
|
||
from gnnrec.hge.hgt.model import HGTFull
|
||
from gnnrec.hge.utils import set_random_seed, get_device, load_data, add_node_feat
|
||
|
||
|
||
def main():
|
||
args = parse_args()
|
||
print(args)
|
||
set_random_seed(args.seed)
|
||
device = get_device(args.device)
|
||
|
||
data, g, _, labels, predict_ntype, train_idx, val_idx, test_idx, _ = load_data(args.dataset)
|
||
g = g.to(device)
|
||
labels = labels.tolist()
|
||
train_idx = torch.cat([train_idx, val_idx])
|
||
add_node_feat(g, 'one-hot')
|
||
|
||
label_neigh = sample_label_neighbors(labels, args.num_samples) # (N, T_pos)
|
||
# List[tensor(N, T_pos)] HGT计算出的注意力权重，M条元路径+一个总体
|
||
attn_pos = calc_attn_pos(g, data.num_classes, predict_ntype, args.num_samples, device, args)
|
||
|
||
# 元路径对应的正样本图
|
||
v = torch.repeat_interleave(g.nodes(predict_ntype), args.num_samples).cpu()
|
||
pos_graphs = []
|
||
for p in attn_pos[:-1]:
|
||
u = p.view(1, -1).squeeze(dim=0) # (N*T_pos,)
|
||
pos_graphs.append(dgl.graph((u, v)))
|
||
|
||
# 整体正样本图
|
||
pos = attn_pos[-1]
|
||
if args.use_label:
|
||
pos[train_idx] = label_neigh[train_idx]
|
||
u = pos.view(1, -1).squeeze(dim=0)
|
||
pos_graphs.append(dgl.graph((u, v)))
|
||
|
||
dgl.save_graphs(args.save_graph_path, pos_graphs)
|
||
print('正样本图已保存到', args.save_graph_path)
|
||
|
||
|
||
def calc_attn_pos(g, num_classes, predict_ntype, num_samples, device, args):
|
||
"""使用预训练的HGT模型计算的注意力权重选择目标顶点的正样本。"""
|
||
# 形如ABA的元路径，其中A是目标顶点类型
|
||
metapaths = []
|
||
for s, e, d in g.canonical_etypes:
|
||
if d == predict_ntype:
|
||
re = next(re for rs, re, rd in g.canonical_etypes if rs == d and rd == s)
|
||
metapaths.append((re, e))
|
||
|
||
model = HGTFull(
|
||
{ntype: g.nodes[ntype].data['feat'].shape[1] for ntype in g.ntypes},
|
||
args.num_hidden, num_classes, args.num_heads, g.ntypes, g.canonical_etypes,
|
||
predict_ntype, 2, args.dropout
|
||
).to(device)
|
||
model.load_state_dict(torch.load(args.hgt_model_path, map_location=device))
|
||
|
||
# 每条元路径ABA对应一个正样本图G_ABA，加一个总体正样本图G_pos
|
||
with torch.no_grad():
|
||
_ = model(g, g.ndata['feat'])
|
||
attn = [calc_attn(mp, model, g, device).t() for mp in metapaths] # List[tensor(N, N)]
|
||
pos = [torch.topk(a, num_samples)[1] for a in attn] # List[tensor(N, T_pos)]
|
||
pos.append(torch.topk(sum(attn), num_samples)[1])
|
||
return [p.cpu() for p in pos]
|
||
|
||
|
||
def calc_attn(metapath, model, g, device):
|
||
"""计算通过指定元路径与目标顶点连接的同类型顶点的注意力权重。"""
|
||
re, e = metapath
|
||
s, _, d = g.to_canonical_etype(re) # s是目标顶点类型, d是中间顶点类型
|
||
a0 = torch.zeros(g.num_nodes(s), g.num_nodes(d), device=device)
|
||
a0[g.edges(etype=re)] = model.layers[0].conv.mods[re].attn.mean(dim=1)
|
||
a1 = torch.zeros(g.num_nodes(d), g.num_nodes(s), device=device)
|
||
a1[g.edges(etype=e)] = model.layers[1].conv.mods[e].attn.mean(dim=1)
|
||
return torch.matmul(a0, a1) # (N, N)
|
||
|
||
|
||
def sample_label_neighbors(labels, num_samples):
|
||
"""为每个顶点采样相同标签的邻居。"""
|
||
label2id = defaultdict(list)
|
||
for i, y in enumerate(labels):
|
||
label2id[y].append(i)
|
||
return torch.tensor([random.sample(label2id[y], num_samples) for y in labels])
|
||
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(description='使用预训练的HGT计算的注意力权重构造正样本图(full-batch)')
|
||
parser.add_argument('--seed', type=int, default=0, help='随机数种子')
|
||
parser.add_argument('--device', type=int, default=0, help='GPU设备')
|
||
parser.add_argument('--dataset', choices=['acm', 'dblp'], default='acm', help='数据集')
|
||
parser.add_argument('--num-hidden', type=int, default=512, help='隐藏层维数')
|
||
parser.add_argument('--num-heads', type=int, default=8, help='注意力头数')
|
||
parser.add_argument('--dropout', type=float, default=0.5, help='Dropout概率')
|
||
parser.add_argument('--num-samples', type=int, default=5, help='每个顶点采样的正样本数量')
|
||
parser.add_argument('--use-label', action='store_true', help='训练集使用真实标签')
|
||
parser.add_argument('hgt_model_path', help='预训练的HGT模型保存路径')
|
||
parser.add_argument('save_graph_path', help='正样本图保存路径')
|
||
args = parser.parse_args()
|
||
return args
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|