def main(args):
g, num_rels, num_classes, labels, train_idx, test_idx, target_idx, inv_target = load_data(
args.dataset, inv_target=True)
if args.gpu >= 0 and th.cuda.is_available():
device = th.device(args.gpu)
else:
device = th.device('cpu')
train_loader, val_loader, test_loader = init_dataloaders(
args, g, train_idx, test_idx, target_idx, args.gpu)
model = RGCN(g.num_nodes(),
args.n_hidden,
num_classes,
num_rels,
num_bases=args.n_bases,
dropout=args.dropout,
self_loop=args.use_self_loop,
ns_mode=True)
labels = labels.to(device)
model = model.to(device)
optimizer = th.optim.Adam(model.parameters(),
lr=1e-2,
weight_decay=args.wd)
for epoch in range(args.n_epochs):
train_acc, loss = train(model, train_loader, inv_target, labels,
optimizer)
print(
"Epoch {:05d}/{:05d} | Train Accuracy: {:.4f} | Train Loss: {:.4f}"
.format(epoch, args.n_epochs, train_acc, loss))
val_logits, val_seeds = evaluate(model, val_loader, inv_target)
val_acc = accuracy(val_logits.argmax(dim=1),
labels[val_seeds].cpu()).item()
print("Validation Accuracy: {:.4f}".format(val_acc))
test_logits, test_seeds = evaluate(model, test_loader, inv_target)
test_acc = accuracy(test_logits.argmax(dim=1),
labels[test_seeds].cpu()).item()
print("Final Test Accuracy: {:.4f}".format(test_acc))
def main(args):
g, num_rels, num_classes, labels, train_idx, test_idx, target_idx = load_data(
args.dataset, get_norm=True)
num_nodes = g.num_nodes()
# Since the nodes are featureless, learn node embeddings from scratch
# This requires passing the node IDs to the model.
feats = th.arange(num_nodes)
model = RGCN(num_nodes,
args.n_hidden,
num_classes,
num_rels,
num_bases=args.n_bases)
if args.gpu >= 0 and th.cuda.is_available():
device = th.device(args.gpu)
else:
device = th.device('cpu')
feats = feats.to(device)
labels = labels.to(device)
model = model.to(device)
g = g.to(device)
optimizer = th.optim.Adam(model.parameters(),
lr=1e-2,
weight_decay=args.l2norm)
model.train()
for epoch in range(50):
logits = model(g, feats)
logits = logits[target_idx]
loss = F.cross_entropy(logits[train_idx], labels[train_idx])
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_acc = accuracy(logits[train_idx].argmax(dim=1),
labels[train_idx]).item()
print("Epoch {:05d} | Train Accuracy: {:.4f} | Train Loss: {:.4f}".
format(epoch, train_acc, loss.item()))
print()
model.eval()
with th.no_grad():
logits = model(g, feats)
logits = logits[target_idx]
test_acc = accuracy(logits[test_idx].argmax(dim=1),
labels[test_idx]).item()
print("Test Accuracy: {:.4f}".format(test_acc))
def init_models(args, device, num_nodes, num_classes, num_rels):
embed_layer = RelGraphEmbedLayer(device, num_nodes, args.n_hidden)
model = RGCN(args.n_hidden,
args.n_hidden,
num_classes,
num_rels,
num_bases=args.n_bases,
dropout=args.dropout,
self_loop=args.use_self_loop)
return embed_layer, model
def main(args):
g, num_rels, num_classes, labels, train_idx, test_idx, target_idx = load_data(
args.dataset, get_norm=True)
model = RGCN(g.num_nodes(),
args.n_hidden,
num_classes,
num_rels,
num_bases=args.n_bases)
if args.gpu >= 0 and th.cuda.is_available():
device = th.device(args.gpu)
else:
device = th.device('cpu')
labels = labels.to(device)
model = model.to(device)
g = g.int().to(device)
optimizer = th.optim.Adam(model.parameters(),
lr=1e-2,
weight_decay=args.wd)
model.train()
for epoch in range(100):
logits = model(g)
logits = logits[target_idx]
loss = F.cross_entropy(logits[train_idx], labels[train_idx])
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_acc = accuracy(logits[train_idx].argmax(dim=1),
labels[train_idx]).item()
print("Epoch {:05d} | Train Accuracy: {:.4f} | Train Loss: {:.4f}".
format(epoch, train_acc, loss.item()))
print()
model.eval()
with th.no_grad():
logits = model(g)
logits = logits[target_idx]
test_acc = accuracy(logits[test_idx].argmax(dim=1),
labels[test_idx]).item()
print("Test Accuracy: {:.4f}".format(test_acc))
def __init__(self,
in_dim,
num_rels,
h_dim=500,
num_bases=100,
dropout=0.2,
reg_param=0.01):
super(LinkPredict, self).__init__()
self.rgcn = RGCN(in_dim,
h_dim,
h_dim,
num_rels * 2,
regularizer="bdd",
num_bases=num_bases,
dropout=dropout,
self_loop=True)
self.dropout = nn.Dropout(dropout)
self.reg_param = reg_param
self.w_relation = nn.Parameter(th.Tensor(num_rels, h_dim))
nn.init.xavier_uniform_(self.w_relation,
gain=nn.init.calculate_gain('relu'))
def main(args):
paddle.seed(args.seed)
np.random.seed(args.seed)
if paddle.distributed.get_world_size() > 1:
paddle.distributed.init_parallel_env()
node_labels = pd.read_csv(
os.path.join(args.data_path, 'node_labels.txt'), header=None,
sep="\t").values.astype("int64")
node_labels = node_labels[:, 1:2]
print(node_labels.shape)
num_nodes = len(node_labels)
train_idx = pd.read_csv(
os.path.join(args.data_path, 'train_idx.txt'), header=None,
sep="\t").values.astype("int64").reshape(-1, ).tolist()
test_idx = pd.read_csv(
os.path.join(args.data_path, 'test_idx.txt'), header=None,
sep="\t").values.astype("int64").reshape(-1, ).tolist()
g = build_heter_graph(os.path.join(args.data_path, "edges"), num_nodes)
model = RGCN(
num_nodes=num_nodes,
input_size=args.input_size,
hidden_size=args.hidden_size,
num_class=args.num_class,
num_layers=args.num_layers,
etypes=g.edge_types,
num_bases=args.num_bases, )
model = paddle.DataParallel(model)
criterion = paddle.nn.loss.CrossEntropyLoss()
optim = paddle.optimizer.Adam(
learning_rate=args.lr,
parameters=model.parameters(),
weight_decay=0.001)
test_acc_list = []
g.tensor()
train_labels = paddle.to_tensor(node_labels[train_idx])
test_labels = paddle.to_tensor(node_labels[test_idx])
train_idx = paddle.to_tensor(train_idx)
test_idx = paddle.to_tensor(test_idx)
for epoch in range(args.epochs):
logits = model(g)
train_logits = paddle.gather(logits, train_idx)
train_loss = criterion(train_logits, train_labels)
train_loss.backward()
train_acc = paddle.metric.accuracy(
train_logits, label=train_labels, k=1)
optim.step()
optim.clear_grad()
test_logits = paddle.gather(logits, test_idx)
test_loss = criterion(test_logits, test_labels)
test_acc = paddle.metric.accuracy(test_logits, label=test_labels, k=1)
msg = "epoch: %s" % epoch
msg += " | train_loss: %.4f | train_acc: %.4f" \
% (train_loss.numpy()[0], train_acc.numpy()[0])
msg += " | test_loss: %.4f | test_acc: %.4f" \
% (test_loss.numpy()[0], test_acc.numpy()[0])
log.info(msg)
test_acc_list.append(test_acc.numpy()[0])
log.info("best test acc result: %.4f" % (np.max(test_acc_list)))
def run(proc_id, n_gpus, n_cpus, args, devices, dataset, queue=None):
dev_id = devices[proc_id]
th.cuda.set_device(dev_id)
g, num_rels, num_classes, labels, train_idx, test_idx,\
target_idx, inv_target = dataset
dist_init_method = 'tcp://{master_ip}:{master_port}'.format(
master_ip='127.0.0.1', master_port='12345')
backend = 'nccl'
if proc_id == 0:
print("backend using {}".format(backend))
th.distributed.init_process_group(backend=backend,
init_method=dist_init_method,
world_size=n_gpus,
rank=proc_id)
device = th.device(dev_id)
use_ddp = True if n_gpus > 1 else False
train_loader, val_loader, test_loader = init_dataloaders(args,
g,
train_idx,
test_idx,
target_idx,
dev_id,
use_ddp=use_ddp)
model = RGCN(g.num_nodes(),
args.n_hidden,
num_classes,
num_rels,
num_bases=args.n_bases,
dropout=args.dropout,
self_loop=args.use_self_loop,
ns_mode=True)
labels = labels.to(device)
model = model.to(device)
inv_target = inv_target.to(device)
model = DistributedDataParallel(model,
device_ids=[dev_id],
output_device=dev_id)
optimizer = th.optim.Adam(model.parameters(),
lr=1e-2,
weight_decay=args.wd)
th.set_num_threads(n_cpus)
for epoch in range(args.n_epochs):
train_acc, loss = train(model, train_loader, inv_target, labels,
optimizer)
if proc_id == 0:
print(
"Epoch {:05d}/{:05d} | Train Accuracy: {:.4f} | Train Loss: {:.4f}"
.format(epoch, args.n_epochs, train_acc, loss))
# garbage collection that empties the queue
gc.collect()
val_logits, val_seeds = evaluate(model, val_loader, inv_target)
queue.put((val_logits, val_seeds))
# gather evaluation result from multiple processes
if proc_id == 0:
val_acc = collect_eval(n_gpus, queue, labels)
print("Validation Accuracy: {:.4f}".format(val_acc))
# garbage collection that empties the queue
gc.collect()
test_logits, test_seeds = evaluate(model, test_loader, inv_target)
queue.put((test_logits, test_seeds))
if proc_id == 0:
test_acc = collect_eval(n_gpus, queue, labels)
print("Final Test Accuracy: {:.4f}".format(test_acc))
th.distributed.barrier()
def run_simulation():
############## Train graph construction #########################
train_graph = DGLGraph()
train_graph.add_nodes(len(N))
train_edges = np.array([e for e in edges if e[4] == 'Train'])
train_click_data = torch.Tensor(np.stack(train_edges[:, 3])).to(device)
train_edge_src = torch.from_numpy(
train_edges[:, 0].flatten().astype(int)).to(device)
train_edge_dst = torch.from_numpy(
train_edges[:, 1].flatten().astype(int)).to(device)
train_edge_type = train_edges[:, 2].flatten().astype(int)
train_edge_norm = np.ones(len(train_edge_dst),
dtype=np.float32) #/ degrees.astype(np.float32)
train_edge_type = torch.from_numpy(train_edge_type).to(device)
train_edge_norm = torch.from_numpy(train_edge_norm).unsqueeze(1).to(device)
train_graph.add_edges(train_edge_src, train_edge_dst, {
"norm": train_edge_norm,
"type": train_edge_type
})
train_graph.add_edges(train_edge_dst, train_edge_src, {
"norm": train_edge_norm,
"type": train_edge_type
})
train_graph.set_n_initializer(dgl.init.zero_initializer)
train_y = torch.LongTensor(train_edges[:,
2].flatten().astype(int)).to(device)
############## Test graph construction #########################
test_graph = DGLGraph()
test_graph.add_nodes(len(N))
test_edges = np.array([e for e in edges if e[4] == 'Test'])
test_click_data = torch.Tensor(np.stack(test_edges[:, 3])).to(device)
test_edge_src = torch.from_numpy(
test_edges[:, 0].flatten().astype(int)).to(device)
test_edge_dst = torch.from_numpy(
test_edges[:, 1].flatten().astype(int)).to(device)
test_edge_type = test_edges[:, 2].flatten().astype(int)
test_edge_norm = np.ones(len(test_edge_dst),
dtype=np.float32) #/ degrees.astype(np.float32)
test_edge_type = torch.from_numpy(test_edge_type).to(device)
test_edge_norm = torch.from_numpy(test_edge_norm).unsqueeze(1).to(device)
test_graph.add_edges(test_edge_src, test_edge_dst, {
"norm": test_edge_norm,
"type": test_edge_type
})
test_graph.add_edges(test_edge_dst, test_edge_src, {
"norm": test_edge_norm,
"type": test_edge_type
})
test_graph.set_n_initializer(dgl.init.zero_initializer)
test_y = torch.LongTensor(test_edges[:,
2].flatten().astype(int)).to(device)
#################################################################
all_node_features = torch.Tensor(node_features).to(device)
model = RGCN(graph=train_graph,
features=all_node_features,
n_hidden_feats=args.num_hidden_features_graph,
n_hidden_layers=args.num_hidden_layers_graph,
n_classes=args.num_classes,
activation=F.relu,
dropout=args.dropout,
n_rels=args.num_classes,
n_bases=-1,
self_loop=True,
click_lstm_hidden_size=args.click_lstm_hidden_size,
click_sequence_length=40,
click_lstm_num_layers=args.click_lstm_num_layers).to(device)
def test():
model.load_state_dict(torch.load(save_dir + 'model.m'))
model.eval()
with torch.no_grad():
output = model(all_node_features,
test_click_data,
test_edge_src,
test_edge_dst,
truncate_sequence_lstm=args.weeks_test - 1)
predictions = torch.argmax(F.softmax(output, dim=1),
dim=1).cpu().numpy()
ground_truths = test_y.cpu().numpy()
f1 = f1_score(y_true=ground_truths,
y_pred=predictions,
average='weighted')
performance_file = open(save_dir + 'test_performance_log.txt', 'w')
report = classification_report(y_true=ground_truths,
y_pred=predictions)
conf_matrix = confusion_matrix(y_true=ground_truths,
y_pred=predictions)
performance_file.write(
"classification_report:\n{}\n".format(report))
performance_file.write("conf_matrix:\n{}\n".format(conf_matrix))
print("Test : {}".format(f1))
return f1
def train():
model.train()
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
scheduler = StepLR(optimizer, step_size=100, gamma=0.99)
criterion = nn.CrossEntropyLoss()
for step in tqdm(range(args.steps)):
output = model(all_node_features,
train_click_data,
train_edge_src,
train_edge_dst,
truncate_sequence_lstm=39)
loss = criterion(output, train_y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
if step % 20 == 0 and step:
predictions = torch.argmax(F.softmax(output, dim=1),
dim=1).cpu().numpy()
ground_truths = train_y.cpu().numpy()
train_f1 = f1_score(y_true=ground_truths,
y_pred=predictions,
average='weighted')
print('Step {} Loss {} Train F1-score {}'.format(
step, loss.data, train_f1))
model.train()
torch.save(model.state_dict(), save_dir + 'model.m')
train()
test()
def train(args, device, data):
# Unpack data
train_nid, val_nid, test_nid, input_dim, left_pad_size, right_pad_size, labels, n_classes, entity_features, g = data
# Create PyTorch DataLoader for constructing blocks
sampler = dgl.dataloading.MultiLayerNeighborSampler(
[int(fanout) for fanout in args.fan_out.split(',')])
dataloader = dgl.dataloading.NodeDataLoader(
g,
{args.label_entity: train_nid},
sampler,
batch_size=args.batch_size,
shuffle=True,
drop_last=False,
num_workers=args.num_workers) # may change device to gpu? default device='cpu'
# Define model and optimizer
model = RGCN(input_dim, args.hidden_dim, n_classes, args.num_layers, F.relu, args.dropout, g.etypes)
model = model.to(device)
loss_fcn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
# Training loop
avg = 0
iter_tput = []
best_eval_acc = 0
best_test_acc = 0
for epoch in range(args.num_epochs):
tic = time.time()
# Loop over the dataloader to sample the computation dependency graph as a list of
# blocks.
for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
tic_step = time.time()
# copy block to gpu
blocks = [blk.int().to(device) for blk in blocks]
# Load the input features as well as output labels
batch_inputs, batch_labels = load_subtensor(entity_features, labels, seeds, input_nodes, args.label_entity, args.is_pad, device)
# Compute loss and prediction
batch_pred = model(blocks, batch_inputs)[args.label_entity]
loss = loss_fcn(batch_pred, batch_labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
iter_tput.append(len(seeds) / (time.time() - tic_step))
if step % args.log_every == 0:
acc = compute_acc(batch_pred, batch_labels)
gpu_mem_alloc = torch.cuda.max_memory_allocated() / 1000000 if torch.cuda.is_available() else 0
print('Epoch {:05d} | Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | Speed (samples/sec) {:.4f} | GPU {:.1f} MB'.format(
epoch, step, loss.item(), acc.item(), np.mean(iter_tput[3:]), gpu_mem_alloc))
toc = time.time()
print('Epoch Time(s): {:.4f}'.format(toc - tic))
if epoch >= 5:
avg += toc - tic
if epoch % args.eval_every == 0 and epoch != 0:
eval_acc, test_acc = evaluate(model, g, entity_features, labels, val_nid, test_nid, device,
args.batch_size, args.num_workers, args.label_entity, args.is_pad)
# if args.save_pred:
# np.savetxt(args.save_pred + '%02d' % epoch, pred.argmax(1).cpu().numpy(), '%d')
print('Eval Acc {:.4f}'.format(eval_acc))
if eval_acc > best_eval_acc:
best_eval_acc = eval_acc
best_test_acc = test_acc
print('Best Eval Acc {:.4f} Test Acc {:.4f}'.format(best_eval_acc, best_test_acc))
print('Avg epoch time: {}'.format(avg / (epoch - 4)))
return best_test_acc