def main(args):
# check cuda
device = 'cpu'
use_cuda = args.gpu >= 0 and th.cuda.is_available()
if use_cuda:
th.cuda.set_device(args.gpu)
device = 'cuda:%d' % args.gpu
# 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
train_mask = g.nodes[category].data.pop('train_mask')
test_mask = g.nodes[category].data.pop('test_mask')
train_idx = th.nonzero(train_mask, as_tuple=False).squeeze()
test_idx = th.nonzero(test_mask, as_tuple=False).squeeze()
labels = g.nodes[category].data.pop('labels')
# split dataset into train, validate, test
if args.validation:
val_idx = train_idx[:len(train_idx) // 5]
train_idx = train_idx[len(train_idx) // 5:]
else:
val_idx = train_idx
# create embeddings
embed_layer = RelGraphEmbed(g, args.n_hidden)
if not args.data_cpu:
labels = labels.to(device)
embed_layer = embed_layer.to(device)
node_embed = embed_layer()
# create model
model = EntityClassify(g,
args.n_hidden,
num_classes,
num_bases=args.n_bases,
num_hidden_layers=args.n_layers - 2,
dropout=args.dropout,
use_self_loop=args.use_self_loop)
if use_cuda:
model.cuda()
# train sampler
sampler = dgl.dataloading.MultiLayerNeighborSampler([args.fanout] *
args.n_layers)
loader = dgl.dataloading.DataLoader(g, {category: train_idx},
sampler,
batch_size=args.batch_size,
shuffle=True,
num_workers=0)
# validation sampler
# we do not use full neighbor to save computation resources
val_sampler = dgl.dataloading.MultiLayerNeighborSampler([args.fanout] *
args.n_layers)
val_loader = dgl.dataloading.DataLoader(g, {category: val_idx},
val_sampler,
batch_size=args.batch_size,
shuffle=True,
num_workers=0)
# optimizer
all_params = itertools.chain(model.parameters(), embed_layer.parameters())
optimizer = th.optim.Adam(all_params, lr=args.lr, weight_decay=args.l2norm)
# training loop
print("start training...")
dur = []
for epoch in range(args.n_epochs):
model.train()
optimizer.zero_grad()
if epoch > 3:
t0 = time.time()
for i, (input_nodes, seeds, blocks) in enumerate(loader):
blocks = [blk.to(device) for blk in blocks]
seeds = seeds[
category] # we only predict the nodes with type "category"
batch_tic = time.time()
emb = extract_embed(node_embed, input_nodes)
lbl = labels[seeds]
if use_cuda:
emb = {k: e.cuda() for k, e in emb.items()}
lbl = lbl.cuda()
logits = model(emb, blocks)[category]
loss = F.cross_entropy(logits, lbl)
loss.backward()
optimizer.step()
train_acc = th.sum(logits.argmax(dim=1) == lbl).item() / len(seeds)
print(
"Epoch {:05d} | Batch {:03d} | Train Acc: {:.4f} | Train Loss: {:.4f} | Time: {:.4f}"
.format(epoch, i, train_acc, loss.item(),
time.time() - batch_tic))
if epoch > 3:
dur.append(time.time() - t0)
val_loss, val_acc = evaluate(model, val_loader, node_embed, labels,
category, device)
print(
"Epoch {:05d} | Valid Acc: {:.4f} | Valid loss: {:.4f} | Time: {:.4f}"
.format(epoch, val_acc, val_loss, np.average(dur)))
print()
if args.model_path is not None:
th.save(model.state_dict(), args.model_path)
output = model.inference(g, args.batch_size, 'cuda' if use_cuda else 'cpu',
0, node_embed)
test_pred = output[category][test_idx]
test_labels = labels[test_idx].to(test_pred.device)
test_acc = (test_pred.argmax(1) == test_labels).float().mean()
print("Test Acc: {:.4f}".format(test_acc))
print()
def main(args):
# load graph data
if args.dataset == 'aifb':
dataset = AIFB()
elif args.dataset == 'mutag':
dataset = MUTAG()
elif args.dataset == 'bgs':
dataset = BGS()
elif args.dataset == 'am':
dataset = AM()
else:
raise ValueError()
g = dataset.graph
category = dataset.predict_category
num_classes = dataset.num_classes
train_idx = dataset.train_idx
test_idx = dataset.test_idx
labels = dataset.labels
# split dataset into train, validate, test
if args.validation:
val_idx = train_idx[:len(train_idx) // 5]
train_idx = train_idx[len(train_idx) // 5:]
else:
val_idx = train_idx
# check cuda
use_cuda = args.gpu >= 0 and th.cuda.is_available()
if use_cuda:
th.cuda.set_device(args.gpu)
train_label = labels[train_idx]
val_label = labels[val_idx]
test_label = labels[test_idx]
# create embeddings
embed_layer = RelGraphEmbed(g, args.n_hidden)
node_embed = embed_layer()
# create model
model = EntityClassify(g,
args.n_hidden,
num_classes,
num_bases=args.n_bases,
num_hidden_layers=args.n_layers - 2,
dropout=args.dropout,
use_self_loop=args.use_self_loop)
if use_cuda:
model.cuda()
# train sampler
sampler = HeteroNeighborSampler(g, category, [args.fanout] * args.n_layers)
loader = DataLoader(dataset=train_idx.numpy(),
batch_size=args.batch_size,
collate_fn=sampler.sample_blocks,
shuffle=True,
num_workers=0)
# validation sampler
val_sampler = HeteroNeighborSampler(g, category, [None] * args.n_layers)
_, val_blocks = val_sampler.sample_blocks(val_idx)
# test sampler
test_sampler = HeteroNeighborSampler(g, category, [None] * args.n_layers)
_, test_blocks = test_sampler.sample_blocks(test_idx)
# optimizer
all_params = itertools.chain(model.parameters(), embed_layer.parameters())
optimizer = th.optim.Adam(all_params, lr=args.lr, weight_decay=args.l2norm)
# training loop
print("start training...")
dur = []
for epoch in range(args.n_epochs):
model.train()
optimizer.zero_grad()
if epoch > 3:
t0 = time.time()
for i, (seeds, blocks) in enumerate(loader):
batch_tic = time.time()
emb = extract_embed(node_embed, blocks[0])
lbl = labels[seeds[category]]
if use_cuda:
emb = {k: e.cuda() for k, e in emb.items()}
lbl = lbl.cuda()
logits = model(emb, blocks)[category]
loss = F.cross_entropy(logits, lbl)
loss.backward()
optimizer.step()
train_acc = th.sum(logits.argmax(dim=1) == lbl).item() / len(
seeds[category])
print(
"Epoch {:05d} | Batch {:03d} | Train Acc: {:.4f} | Train Loss: {:.4f} | Time: {:.4f}"
.format(epoch, i, train_acc, loss.item(),
time.time() - batch_tic))
if epoch > 3:
dur.append(time.time() - t0)
val_loss, val_acc = evaluate(model, val_idx, val_blocks, node_embed,
labels, category, use_cuda)
print(
"Epoch {:05d} | Valid Acc: {:.4f} | Valid loss: {:.4f} | Time: {:.4f}"
.format(epoch, val_acc, val_loss.item(), np.average(dur)))
print()
if args.model_path is not None:
th.save(model.state_dict(), args.model_path)
test_loss, test_acc = evaluate(model, test_idx, test_blocks, node_embed,
labels, category, use_cuda)
print("Test Acc: {:.4f} | Test loss: {:.4f}".format(
test_acc, test_loss.item()))
print()
def main(args):
# load graph data
if args.dataset == 'aifb':
dataset = AIFB()
elif args.dataset == 'mutag':
dataset = MUTAG()
elif args.dataset == 'bgs':
dataset = BGS()
elif args.dataset == 'am':
dataset = AM()
else:
raise ValueError()
g = dataset.graph
category = dataset.predict_category
num_classes = dataset.num_classes
train_idx = dataset.train_idx
test_idx = dataset.test_idx
labels = dataset.labels
category_id = len(g.ntypes)
for i, ntype in enumerate(g.ntypes):
if ntype == category:
category_id = i
# split dataset into train, validate, test
if args.validation:
val_idx = train_idx[:len(train_idx) // 5]
train_idx = train_idx[len(train_idx) // 5:]
else:
val_idx = train_idx
# check cuda
use_cuda = args.gpu >= 0 and th.cuda.is_available()
if use_cuda:
th.cuda.set_device(args.gpu)
g = g.to('cuda:%d' % args.gpu)
labels = labels.cuda()
train_idx = train_idx.cuda()
test_idx = test_idx.cuda()
# create model
model = EntityClassify(g,
args.n_hidden,
num_classes,
num_bases=args.n_bases,
num_hidden_layers=args.n_layers - 2,
dropout=args.dropout,
use_self_loop=args.use_self_loop)
if use_cuda:
model.cuda()
# optimizer
optimizer = th.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.l2norm)
# training loop
print("start training...")
dur = []
model.train()
for epoch in range(args.n_epochs):
optimizer.zero_grad()
if epoch > 5:
t0 = time.time()
logits = model()[category]
loss = F.cross_entropy(logits[train_idx], labels[train_idx])
loss.backward()
optimizer.step()
t1 = time.time()
if epoch > 5:
dur.append(t1 - t0)
train_acc = th.sum(logits[train_idx].argmax(
dim=1) == labels[train_idx]).item() / len(train_idx)
val_loss = F.cross_entropy(logits[val_idx], labels[val_idx])
val_acc = th.sum(logits[val_idx].argmax(
dim=1) == labels[val_idx]).item() / len(val_idx)
print(
"Epoch {:05d} | Train Acc: {:.4f} | Train Loss: {:.4f} | Valid Acc: {:.4f} | Valid loss: {:.4f} | Time: {:.4f}"
.format(epoch, train_acc, loss.item(), val_acc, val_loss.item(),
np.average(dur)))
print()
if args.model_path is not None:
th.save(model.state_dict(), args.model_path)
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()
def run(args: argparse.ArgumentParser) -> None:
torch.manual_seed(args.seed)
dataset, hg, train_idx, valid_idx, test_idx = utils.process_dataset(
args.dataset,
root=args.dataset_root,
)
predict_category = dataset.predict_category
labels = hg.nodes[predict_category].data['labels']
training_device = torch.device('cuda' if args.gpu_training else 'cpu')
inference_device = torch.device('cuda' if args.gpu_inference else 'cpu')
inferfence_mode = args.inference_mode
fanouts = [int(fanout) for fanout in args.fanouts.split(',')]
train_sampler = dgl.dataloading.MultiLayerNeighborSampler(fanouts)
train_dataloader = dgl.dataloading.NodeDataLoader(
hg,
{predict_category: train_idx},
train_sampler,
batch_size=args.batch_size,
shuffle=True,
drop_last=False,
num_workers=args.num_workers,
)
if inferfence_mode == 'neighbor_sampler':
valid_sampler = dgl.dataloading.MultiLayerNeighborSampler(fanouts)
valid_dataloader = dgl.dataloading.NodeDataLoader(
hg,
{predict_category: valid_idx},
valid_sampler,
batch_size=args.eval_batch_size,
shuffle=False,
drop_last=False,
num_workers=args.eval_num_workers,
)
if args.test_validation:
test_sampler = dgl.dataloading.MultiLayerNeighborSampler(fanouts)
test_dataloader = dgl.dataloading.NodeDataLoader(
hg,
{predict_category: test_idx},
test_sampler,
batch_size=args.eval_batch_size,
shuffle=False,
drop_last=False,
num_workers=args.eval_num_workers,
)
else:
valid_dataloader = None
if args.test_validation:
test_dataloader = None
in_feats = hg.nodes[predict_category].data['feat'].shape[-1]
out_feats = dataset.num_classes
num_nodes = {}
node_feats = {}
for ntype in hg.ntypes:
num_nodes[ntype] = hg.num_nodes(ntype)
node_feats[ntype] = hg.nodes[ntype].data.get('feat')
activations = {'leaky_relu': F.leaky_relu, 'relu': F.relu}
embedding_layer = RelGraphEmbedding(hg, in_feats, num_nodes, node_feats)
model = EntityClassify(
hg,
in_feats,
args.hidden_feats,
out_feats,
args.num_bases,
args.num_layers,
norm=args.norm,
layer_norm=args.layer_norm,
input_dropout=args.input_dropout,
dropout=args.dropout,
activation=activations[args.activation],
self_loop=args.self_loop,
)
loss_function = nn.CrossEntropyLoss()
embedding_optimizer = torch.optim.SparseAdam(
embedding_layer.node_embeddings.parameters(), lr=args.embedding_lr)
if args.node_feats_projection:
all_parameters = chain(model.parameters(),
embedding_layer.embeddings.parameters())
model_optimizer = torch.optim.Adam(all_parameters, lr=args.model_lr)
else:
model_optimizer = torch.optim.Adam(model.parameters(),
lr=args.model_lr)
checkpoint = utils.Callback(args.early_stopping_patience,
args.early_stopping_monitor)
print('## Training started ##')
for epoch in range(args.num_epochs):
train_time, train_loss, train_accuracy = train(
embedding_layer,
model,
training_device,
embedding_optimizer,
model_optimizer,
loss_function,
labels,
predict_category,
train_dataloader,
)
valid_time, valid_loss, valid_accuracy = validate(
embedding_layer,
model,
inference_device,
inferfence_mode,
loss_function,
hg,
labels,
predict_category=predict_category,
dataloader=valid_dataloader,
eval_batch_size=args.eval_batch_size,
eval_num_workers=args.eval_num_workers,
mask=valid_idx,
)
checkpoint.create(
epoch,
train_time,
valid_time,
train_loss,
valid_loss,
train_accuracy,
valid_accuracy,
{
'embedding_layer': embedding_layer,
'model': model
},
)
print(f'Epoch: {epoch + 1:03} '
f'Train Loss: {train_loss:.2f} '
f'Valid Loss: {valid_loss:.2f} '
f'Train Accuracy: {train_accuracy:.4f} '
f'Valid Accuracy: {valid_accuracy:.4f} '
f'Train Epoch Time: {train_time:.2f} '
f'Valid Epoch Time: {valid_time:.2f}')
if checkpoint.should_stop:
print('## Training finished: early stopping ##')
break
elif epoch >= args.num_epochs - 1:
print('## Training finished ##')
print(f'Best Epoch: {checkpoint.best_epoch} '
f'Train Loss: {checkpoint.best_epoch_train_loss:.2f} '
f'Valid Loss: {checkpoint.best_epoch_valid_loss:.2f} '
f'Train Accuracy: {checkpoint.best_epoch_train_accuracy:.4f} '
f'Valid Accuracy: {checkpoint.best_epoch_valid_accuracy:.4f}')
if args.test_validation:
print('## Test data validation ##')
embedding_layer.load_state_dict(
checkpoint.best_epoch_model_parameters['embedding_layer'])
model.load_state_dict(checkpoint.best_epoch_model_parameters['model'])
test_time, test_loss, test_accuracy = validate(
embedding_layer,
model,
inference_device,
inferfence_mode,
loss_function,
hg,
labels,
predict_category=predict_category,
dataloader=test_dataloader,
eval_batch_size=args.eval_batch_size,
eval_num_workers=args.eval_num_workers,
mask=test_idx,
)
print(f'Test Loss: {test_loss:.2f} '
f'Test Accuracy: {test_accuracy:.4f} '
f'Test Epoch Time: {test_time:.2f}')
def main(args):
#labels? category?
g, category, num_classes, train_idx, test_idx, labels = load_dataset(
args.dataset)
print('done loading dataset')
for ntype in g.ntypes:
print(ntype, type(ntype))
for i, ntype in enumerate(g.ntypes):
if ntype == category:
category_id = i
# split dataset into train, validate, test
if args.validation:
val_idx = train_idx[:len(train_idx) // 5]
train_idx = train_idx[len(train_idx) // 5:]
else:
val_idx = train_idx
# check cuda
use_cuda = args.gpu >= 0 and th.cuda.is_available()
if use_cuda:
th.cuda.set_device(args.gpu)
labels = labels.cuda()
train_idx = train_idx.cuda()
test_idx = test_idx.cuda()
# create model
model = EntityClassify(g,
args.hidden_size,
num_classes,
num_bases=args.num_bases,
num_hidden_layers=args.n_layers - 2,
dropout=args.dropout,
use_self_loop=args.use_self_loop)
if use_cuda:
model.cuda()
# optimizer
optimizer = th.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.l2norm)
# training loop
print("start training...")
dur = []
model.train()
for epoch in range(args.num_epochs):
optimizer.zero_grad()
th.cuda.synchronize()
if epoch > 3:
t0 = time.time()
logits = model()[category]
loss = F.cross_entropy(logits[train_idx], labels[train_idx])
loss.backward()
optimizer.step()
th.cuda.synchronize()
t1 = time.time()
if epoch > 5:
dur.append(t1 - t0)
train_acc = th.sum(logits[train_idx].argmax(
dim=1) == labels[train_idx]).item() / len(train_idx)
val_loss = F.cross_entropy(logits[val_idx], labels[val_idx])
val_acc = th.sum(logits[val_idx].argmax(
dim=1) == labels[val_idx]).item() / len(val_idx)
print(
"Epoch {:05d} | Train Acc: {:.4f} | Train Loss: {:.4f} | Valid Acc: {:.4f} | Valid loss: {:.4f} | Time: {:.4f}"
.format(epoch, train_acc, loss.item(), val_acc, val_loss.item(),
np.average(dur)))
#print()
if args.model_path is not None:
th.save(model.state_dict(), args.model_path)
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()))
#output we need
Used_memory = th.cuda.max_memory_allocated(0) / (1024**3)
print('^^^{:6f}^^^{:6f}'.format(Used_memory, np.average(dur)))