def main():
dataset, learning_rate, epoch, alpha, beta, gamma, pretrain_epoch = parse_argument(
)
random.seed(9001)
dataset_config = {
'feature_file': './Database/' + dataset + '/features.txt',
'graph_file': './Database/' + dataset + '/edges.txt',
'walks_file': './Database/' + dataset + '/walks.txt',
'label_file': './Database/' + dataset + '/group.txt'
}
graph = Dataset(dataset_config)
pretrain_config = {
'net_shape': [6, 1],
'att_shape': [1],
'net_input_dim': graph.num_nodes,
'att_input_dim': graph.num_feas,
'pretrain_params_path': './Log/' + dataset + '/pretrain_params.pkl'
}
model_config = {
'net_shape': [6, 1],
'att_shape': [1],
'net_input_dim': graph.num_nodes,
'att_input_dim': graph.num_feas,
'is_init': True,
'pretrain_params_path': './Log/' + dataset + '/pretrain_params.pkl'
}
trainer_config = {
'net_shape': [6, 1],
'att_shape': [1],
'net_input_dim': graph.num_nodes,
'att_input_dim': graph.num_feas,
'drop_prob': 0.2,
'learning_rate': learning_rate,
'batch_size': 100,
'num_epochs': epoch,
'beta': beta,
'alpha': alpha,
'gamma': gamma,
'model_path': './Log/' + dataset + '/test_model.pkl',
}
print("=" * 30 + "开始进行参数预训练" + "=" * 30)
pretrainer = PreTrainer(pretrain_config)
pretrainer.pretrain(graph.X, 'net', pretrain_epoch) # walks.txt的矩阵
pretrainer.pretrain(graph.Z, 'att', pretrain_epoch)
model = Model(model_config)
trainer = Trainer(model, trainer_config)
trainer.train(graph)
train_emb = trainer.infer(graph)
train_emb = np.array(train_emb)
with open("./Log/" + dataset + "/train_emb.txt", 'w') as f:
np.savetxt(f, train_emb)
'att_shape': [200, 100],
'net_input_dim': graph.num_nodes,
'att_input_dim': graph.num_feas,
'is_init': True,
'pretrain_params_path': './Log/cora/pretrain_params.pkl'
}
trainer_config = {
'net_shape': [200, 100],
'att_shape': [200, 100],
'net_input_dim': graph.num_nodes,
'att_input_dim': graph.num_feas,
'drop_prob': 0.2,
'learning_rate': 1e-5,
'batch_size': 100,
'num_epochs': 500,
'beta': 100,
'alpha': 50,
'gamma': 500,
'model_path': './Log/cora/cora_model.pkl',
}
pretrainer = PreTrainer(pretrain_config)
pretrainer.pretrain(graph.X, 'net')
pretrainer.pretrain(graph.Z, 'att')
model = Model(model_config)
trainer = Trainer(model, trainer_config)
trainer.train(graph)
trainer.infer(graph)
