def run(args):
from ogb.graphproppred import DglGraphPropPredDataset, Evaluator, collate_dgl
from torch.utils.data import DataLoader
dataset = DglGraphPropPredDataset(name="ogbg-molhiv")
import os
if not os.path.exists("heterographs.bin"):
dataset.graphs = [hpno.heterograph(graph) for graph in dataset.graphs]
from dgl.data.utils import save_graphs
save_graphs("heterographs.bin", dataset.graphs)
else:
from dgl.data.utils import load_graphs
dataset.graphs = load_graphs("heterographs.bin")[0]
evaluator = Evaluator(name="ogbg-molhiv")
in_features = 9
out_features = 1
split_idx = dataset.get_idx_split()
train_loader = DataLoader(dataset[split_idx["train"]], batch_size=128, drop_last=True, shuffle=True, collate_fn=collate_dgl)
valid_loader = DataLoader(dataset[split_idx["valid"]], batch_size=len(split_idx["valid"]), shuffle=False, collate_fn=collate_dgl)
test_loader = DataLoader(dataset[split_idx["test"]], batch_size=len(split_idx["test"]), shuffle=False, collate_fn=collate_dgl)
model = hpno.HierarchicalPathNetwork(
in_features=in_features,
out_features=args.hidden_features,
hidden_features=args.hidden_features,
depth=args.depth,
readout=hpno.GraphReadout(
in_features=args.hidden_features,
out_features=out_features,
hidden_features=args.hidden_features,
)
)
if torch.cuda.is_available():
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(), args.learning_rate, weight_decay=args.weight_decay)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min", factor=0.5, patience=20)
for idx_epoch in range(args.n_epochs):
print(idx_epoch, flush=True)
model.train()
for g, y in train_loader:
y = y.float()
if torch.cuda.is_available():
g = g.to("cuda:0")
y = y.cuda()
optimizer.zero_grad()
y_hat = model.forward(g, g.nodes['n1'].data["feat"].float())
loss = torch.nn.BCELoss()(
input=y_hat.sigmoid(),
target=y,
)
loss.backward()
optimizer.step()
model.eval()
with torch.no_grad():
g, y = next(iter(valid_loader))
y = y.float()
if torch.cuda.is_available():
g = g.to("cuda:0")
y = y.cuda()
y_hat = model.forward(g, g.nodes['n1'].data["feat"].float())
loss = torch.nn.BCELoss()(
input=y_hat.sigmoid(),
target=y,
)
scheduler.step(loss)
if optimizer.param_groups[0]["lr"] <= 0.01 * args.learning_rate: break
model = model.cpu()
g, y = next(iter(valid_loader))
rocauc_vl = evaluator.eval(
{
"y_true": y.float(),
"y_pred": model.forward(g, g.nodes['n1'].data["feat"].float()).sigmoid()
}
)["rocauc"]
g, y = next(iter(test_loader))
rocauc_te = evaluator.eval(
{
"y_true": y.float(),
"y_pred": model.forward(g, g.nodes['n1'].data["feat"].float()).sigmoid()
}
)["rocauc"]
import pandas as pd
df = pd.DataFrame(
{
args.data: {
"rocauc_te": rocauc_te,
"rocauc_vl": rocauc_vl,
}
}
)
df.to_csv("%s.csv" % args.out)
def main(args):
torch.manual_seed(args.seed)
np.random.seed(args.seed)
device = torch.device(
"cuda:" +
str(args.device)) if torch.cuda.is_available() else torch.device("cpu")
# Load dataset and evaluator
dataset = DglGraphPropPredDataset(name=args.dataset)
split_idx = dataset.get_idx_split()
evaluator = Evaluator(args.dataset)
if args.pos_enc_dim > 0:
# Add graph positional encodings
print("Adding PEs...")
dataset.graphs = [
add_positional_encoding(g, args.pos_enc_dim)
for g in tqdm(dataset.graphs)
]
# Basic pre-processing
if args.dataset == 'ogbg-molpcba':
print("Removing training graphs with 0 edges...")
train_split = []
for idx, g in enumerate(tqdm(dataset.graphs)):
if idx in split_idx["train"] and g.number_of_edges() != 0:
train_split.append(idx)
split_idx["train"] = torch.LongTensor(train_split)
# Prepare dataloaders
train_loader = DataLoader(dataset[split_idx["train"]],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=collate_dgl)
valid_loader = DataLoader(dataset[split_idx["valid"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
collate_fn=collate_dgl)
test_loader = DataLoader(dataset[split_idx["test"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
collate_fn=collate_dgl)
# Initialize model, optimizer and scheduler
if args.gnn in ['gated-gcn', 'gcn', 'mlp']:
model = GNN_mol(gnn_type=args.gnn,
num_tasks=dataset.num_tasks,
num_layer=args.num_layer,
emb_dim=args.emb_dim,
dropout=args.dropout,
batch_norm=True,
residual=True,
pos_enc_dim=args.pos_enc_dim,
graph_pooling=args.pooling,
virtualnode=args.virtualnode)
model.to(device)
print(model)
total_param = 0
for param in model.parameters():
total_param += np.prod(list(param.data.size()))
print(f'Total parameters: {total_param}')
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode='min',
factor=args.lr_reduce_factor,
patience=args.lr_scheduler_patience,
verbose=True)
else:
raise ValueError('Invalid GNN type')
# Define loss function
cls_criterion = torch.nn.BCEWithLogitsLoss()
# Create Tensorboard logger
start_time_str = time.strftime("%Y%m%dT%H%M%S")
log_dir = os.path.join(
"logs", args.dataset,
f"{args.expt_name}-{args.gnn}-L{args.num_layer}-h{args.emb_dim}-d{args.dropout}-LR{args.lr}",
f"{start_time_str}-GPU{args.device}")
tb_logger = SummaryWriter(log_dir)
# Training loop
train_curve = []
valid_curve = []
test_curve = []
for epoch in range(1, args.epochs + 1):
print("=====Epoch {}".format(epoch))
tb_logger.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch)
print('Training...')
train(model, device, train_loader, optimizer, cls_criterion)
print('Evaluating...')
train_loss, train_perf = eval(model, device, train_loader, evaluator,
cls_criterion)
valid_loss, valid_perf = eval(model, device, valid_loader, evaluator,
cls_criterion)
_, test_perf = eval(model, device, test_loader, evaluator,
cls_criterion)
# Log statistics to Tensorboard, etc.
print({
'Train': train_perf,
'Validation': valid_perf,
'Test': test_perf
})
tb_logger.add_scalar('loss/train', train_loss, epoch)
tb_logger.add_scalar(f'{dataset.eval_metric}/train',
train_perf[dataset.eval_metric], epoch)
tb_logger.add_scalar('loss/valid', valid_loss, epoch)
tb_logger.add_scalar(f'{dataset.eval_metric}/valid',
valid_perf[dataset.eval_metric], epoch)
tb_logger.add_scalar(f'{dataset.eval_metric}/test',
test_perf[dataset.eval_metric], epoch)
train_curve.append(train_perf[dataset.eval_metric])
valid_curve.append(valid_perf[dataset.eval_metric])
test_curve.append(test_perf[dataset.eval_metric])
if args.lr_scheduler_patience > 0:
# Reduce LR using scheduler
scheduler.step(valid_loss)
if 'classification' in dataset.task_type:
best_val_epoch = np.argmax(np.array(valid_curve))
best_train = max(train_curve)
else:
best_val_epoch = np.argmin(np.array(valid_curve))
best_train = min(train_curve)
print('Finished training!')
print('Best validation score: {}'.format(valid_curve[best_val_epoch]))
print('Test score: {}'.format(test_curve[best_val_epoch]))
torch.save(
{
'args': args,
'model': model.__repr__,
'total_param': total_param,
'BestEpoch': best_val_epoch,
'Validation': valid_curve[best_val_epoch],
'Test': test_curve[best_val_epoch],
'Train': train_curve[best_val_epoch],
'BestTrain': best_train,
}, os.path.join(log_dir, "results.pt"))