dmlc--dgl
Mufei Li
492ad9be81
* Update gatedgraphconv.py * Update entity_classify.py * Update data-process.rst * Update reading_data.py * Update data-process.rst * Update utils.py * Update knowledge_graph.py * Update entity_classify.py * Update rdf.py * Update entity_classify_mb.py * Update test_classify.py * Update tensor.py * Update sparse.py * Update entity_classify_mp.py * Update 6_line_graph.py
80 行
2.8 KiB
Python
80 行
2.8 KiB
Python
"""Infering Relational Data with Graph Convolutional Networks
|
|
"""
|
|
import argparse
|
|
import torch as th
|
|
from functools import partial
|
|
import torch.nn.functional as F
|
|
|
|
from dgl.data.rdf import AIFB, MUTAG, BGS, AM
|
|
from entity_classify import EntityClassify
|
|
|
|
def main(args):
|
|
# load graph data
|
|
if args.dataset == 'aifb':
|
|
dataset = AIFBDataset()
|
|
elif args.dataset == 'mutag':
|
|
dataset = MUTAGDataset()
|
|
elif args.dataset == 'bgs':
|
|
dataset = BGSDataset()
|
|
elif args.dataset == 'am':
|
|
dataset = AMDataset()
|
|
else:
|
|
raise ValueError()
|
|
|
|
g = dataset[0]
|
|
category = dataset.predict_category
|
|
num_classes = dataset.num_classes
|
|
test_mask = g.nodes[category].data.pop('test_mask')
|
|
test_idx = th.nonzero(test_mask, as_tuple=False).squeeze()
|
|
labels = g.nodes[category].data.pop('labels')
|
|
|
|
# check cuda
|
|
use_cuda = args.gpu >= 0 and th.cuda.is_available()
|
|
if use_cuda:
|
|
th.cuda.set_device(args.gpu)
|
|
labels = labels.cuda()
|
|
test_idx = test_idx.cuda()
|
|
g = g.to('cuda:%d' % args.gpu)
|
|
|
|
# create model
|
|
model = EntityClassify(g,
|
|
args.n_hidden,
|
|
num_classes,
|
|
num_bases=args.n_bases,
|
|
num_hidden_layers=args.n_layers - 2,
|
|
use_self_loop=args.use_self_loop)
|
|
model.load_state_dict(th.load(args.model_path))
|
|
if use_cuda:
|
|
model.cuda()
|
|
|
|
print("start testing...")
|
|
model.eval()
|
|
logits = model.forward()[category]
|
|
test_loss = F.cross_entropy(logits[test_idx], labels[test_idx])
|
|
test_acc = th.sum(logits[test_idx].argmax(dim=1) == labels[test_idx]).item() / len(test_idx)
|
|
print("Test Acc: {:.4f} | Test loss: {:.4f}".format(test_acc, test_loss.item()))
|
|
print()
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='RGCN')
|
|
parser.add_argument("--n-hidden", type=int, default=16,
|
|
help="number of hidden units")
|
|
parser.add_argument("--gpu", type=int, default=-1,
|
|
help="gpu")
|
|
parser.add_argument("--lr", type=float, default=1e-2,
|
|
help="learning rate")
|
|
parser.add_argument("--n-bases", type=int, default=-1,
|
|
help="number of filter weight matrices, default: -1 [use all]")
|
|
parser.add_argument("--n-layers", type=int, default=2,
|
|
help="number of propagation rounds")
|
|
parser.add_argument("-d", "--dataset", type=str, required=True,
|
|
help="dataset to use")
|
|
parser.add_argument("--model_path", type=str,
|
|
help='path of the model to load from')
|
|
parser.add_argument("--use-self-loop", default=False, action='store_true',
|
|
help="include self feature as a special relation")
|
|
|
|
args = parser.parse_args()
|
|
print(args)
|
|
main(args)
|