forked from BUPT-GAMMA/GammaGL
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathhpn_trainer.py
144 lines (124 loc) · 5.42 KB
/
hpn_trainer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
# !/usr/bin/env python3
# -*- coding:utf-8 -*-
import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# os.environ['TL_BACKEND'] = 'torch'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# 0:Output all; 1:Filter out INFO; 2:Filter out INFO and WARNING; 3:Filter out INFO, WARNING, and ERROR
import argparse
import tensorlayerx as tlx
import gammagl.transforms as T
from gammagl.datasets import IMDB
from gammagl.models import HPN
from gammagl.utils import mask_to_index
from tensorlayerx.model import TrainOneStep, WithLoss
class SemiSpvzLoss(WithLoss):
def __init__(self, net, loss_fn):
super(SemiSpvzLoss, self).__init__(backbone=net, loss_fn=loss_fn)
def forward(self, data, y):
logits = self.backbone_network(data['x_dict'], data['edge_index_dict'], data['num_nodes_dict'])
train_logits = tlx.gather(logits['movie'], data['train_idx'])
train_y = tlx.gather(data['y'], data['train_idx'])
loss = self._loss_fn(train_logits, train_y)
return loss
def calculate_acc(logits, y, metrics):
"""
Args:
logits: node logits
y: node labels
metrics: tensorlayerx.metrics
Returns:
rst
"""
metrics.update(logits, y)
rst = metrics.result()
metrics.reset()
return rst
def main(args):
# NOTE: ONLY IMDB DATASET
# If you want to execute HPN on other dataset (e.g. ACM),
# you will be needed to init `metepaths`
# and set `movie` string with proper values.
# set_device(args.gpu)
# path = osp.join(osp.dirname(osp.realpath(__file__)), '../IMDB')
metapaths = [[('movie', 'actor'), ('actor', 'movie')],
[('movie', 'director'), ('director', 'movie')]]
transform = T.AddMetaPaths(metapaths=metapaths, drop_orig_edges=True,
drop_unconnected_nodes=True)
dataset = IMDB(args.dataset_path, transform=transform)
graph = dataset[0]
y = graph['movie'].y
# for mindspore, it should be passed into node indices
train_idx = mask_to_index(graph['movie'].train_mask,)
test_idx = mask_to_index(graph['movie'].test_mask)
val_idx = mask_to_index(graph['movie'].val_mask)
net = HPN(
in_channels=graph.x_dict['movie'].shape[1],
out_channels=3, # graph.num_classes,
metadata=graph.metadata(),
drop_rate=args.drop_rate,
hidden_channels=args.hidden_dim,
iter_K=args.iter_K,
alpha=args.alpha,
name = 'hpn',
)
optimizer = tlx.optimizers.Adam(lr=args.lr, weight_decay=args.l2_coef)
metrics = tlx.metrics.Accuracy()
train_weights = net.trainable_weights
loss_func = tlx.losses.softmax_cross_entropy_with_logits
semi_spvz_loss = SemiSpvzLoss(net, loss_func)
train_one_step = TrainOneStep(semi_spvz_loss, optimizer, train_weights)
# train test val = 400, 3478, 400
data = {
"x_dict": graph.x_dict,
"y":y,
"edge_index_dict": graph.edge_index_dict,
"train_idx": train_idx,
"test_idx": test_idx,
"val_idx": val_idx,
"num_nodes_dict": {'movie': graph['movie'].num_nodes},
}
best_val_acc = 0
for epoch in range(args.n_epoch):
net.set_train()
train_loss = train_one_step(data, y)
net.set_eval()
logits = net(data['x_dict'], data['edge_index_dict'], data['num_nodes_dict'])
val_logits = tlx.gather(logits['movie'], data['val_idx'])
val_y = tlx.gather(data['y'], data['val_idx'])
val_acc = calculate_acc(val_logits, val_y, metrics)
print("Epoch [{:0>3d}] ".format(epoch + 1)
+ " train_loss: {:.4f}".format(train_loss.item())
# + " train_acc: {:.4f}".format(train_acc)
+ " val_acc: {:.4f}".format(val_acc))
# save best model on evaluation set
if val_acc > best_val_acc:
best_val_acc = val_acc
net.save_weights(args.best_model_path + net.name + ".npz", format='npz_dict')
net.load_weights(args.best_model_path + net.name + ".npz", format='npz_dict')
net.set_eval()
logits = net(data['x_dict'], data['edge_index_dict'], data['num_nodes_dict'])
test_logits = tlx.gather(logits['movie'], data['test_idx'])
test_y = tlx.gather(data['y'], data['test_idx'])
test_acc = calculate_acc(test_logits, test_y, metrics)
print("Test acc: {:.4f}".format(test_acc))
if __name__ == '__main__':
# parameters setting
parser = argparse.ArgumentParser()
parser.add_argument("--lr", type=float, default=0.01, help="learnin rate")
parser.add_argument("--n_epoch", type=int, default=200, help="number of epoch")
parser.add_argument("--hidden_dim", type=int, default=512, help="dimention of hidden layers")
parser.add_argument("--l2_coef", type=float, default=1e-3, help="l2 loss coeficient")
parser.add_argument("--iter_K", type=int, default=1, help="number K of iteration")
parser.add_argument("--drop_rate", type=float, default=0.4, help="drop_rate")
parser.add_argument("--alpha", type=float, default=0.3, help="alpha")
parser.add_argument("--dataset_path", type=str, default=r'', help="path to save dataset")
# parser.add_argument('--dataset', type=str, default='IMDB', help='dataset')
parser.add_argument("--best_model_path", type=str, default=r'./', help="path to save best model")
parser.add_argument("--gpu", type=int, default=0)
args = parser.parse_args()
if args.gpu >= 0:
tlx.set_device("GPU", args.gpu)
else:
tlx.set_device("CPU")
main(args)