data = HeteData(dataset=dataset, train_percent=train_percent, shuffle=shuffle)
graph_list = data.get_dict_of_list()
homo_graph = nx.to_dict_of_lists(data.homo_graph)
input_dim = data.x.shape[1]
pre_embed_dim = data.type_num * info_section
output_dim = max(data.train_list[:, 1].tolist()) + 1 # 隐藏单元节点数 两层
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # 获取预处理数据
assert select_method in ["end_node", "all_node"]
x = data.x
train_list = data.train_list # 训练节点/数据对应的标签
test_list = data.test_list # 测试节点/数据对应的索引
val_list = data.val_list # 验证节点/数据对应的索引
criterion = nn.CrossEntropyLoss().to(DEVICE)
num_thread = 12
pool = ThreadPool(num_thread)
# 此处枚举metapath并初始化选择模型
metapath_name = mp.enum_metapath_name(data.get_metapath_name(), data.get_metapath_dict(), metapath_length)
metapath_list = mp.enum_longest_metapath_index(data.get_metapath_name(), data.get_metapath_dict(), metapath_length)
select_model = GraphMSE(metapath_list=metapath_name, input_dim=input_dim, pre_embed_dim=pre_embed_dim,
select_dim=output_dim, mlp_settings=mlp_settings).to(DEVICE)
train(model=select_model, epochs=100, method="all_node", ablation=ablation)
# test(model=select_model, test_method="final")
test(model=select_model, test_method="best_val")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # 获取预处理数据
assert select_method in ["end_node", "all_node"]
x = data.x
train_list = data.train_list # 训练节点/数据对应的标签
test_list = data.test_list # 测试节点/数据对应的索引
val_list = data.val_list # 验证节点/数据对应的索引
criterion = nn.CrossEntropyLoss().to(DEVICE)
num_thread = 12
pool = ThreadPool(num_thread)
# 此处枚举metapath并初始化选择模型
metapath_name = mp.enum_metapath_name(data.get_metapath_name(),
data.get_metapath_dict(),
metapath_length)
metapath_list = mp.enum_longest_metapath_index(data.get_metapath_name(),
data.get_metapath_dict(),
metapath_length)
select_model = GraphMSE(metapath_list=metapath_name,
input_dim=input_dim,
pre_embed_dim=pre_embed_dim,
select_dim=output_dim,
mlp_settings=mlp_settings,
endnode_test=True).to(DEVICE)
train(model=select_model, epochs=100, method="end_node", ablation=ablation)
test(model=select_model, test_method="best_val")