wehub-resource-sync
314 lines
8.8 KiB
Python
314 lines
8.8 KiB
Python
import argparse
|
|
|
|
import dgl
|
|
|
|
import numpy as np
|
|
import torch as th
|
|
import torch.nn.functional as F
|
|
from dgl.data import QM9EdgeDataset
|
|
from dgl.data.utils import Subset
|
|
from dgl.dataloading import GraphDataLoader
|
|
from model import InfoGraphS
|
|
|
|
|
|
def argument():
|
|
parser = argparse.ArgumentParser(description="InfoGraphS")
|
|
|
|
# data source params
|
|
parser.add_argument(
|
|
"--target", type=str, default="mu", help="Choose regression task"
|
|
)
|
|
parser.add_argument(
|
|
"--train_num", type=int, default=5000, help="Size of training set"
|
|
)
|
|
|
|
# training params
|
|
parser.add_argument(
|
|
"--gpu", type=int, default=-1, help="GPU index, default:-1, using CPU."
|
|
)
|
|
parser.add_argument(
|
|
"--epochs", type=int, default=200, help="Training epochs."
|
|
)
|
|
parser.add_argument(
|
|
"--batch_size", type=int, default=20, help="Training batch size."
|
|
)
|
|
parser.add_argument(
|
|
"--val_batch_size", type=int, default=100, help="Validation batch size."
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--lr", type=float, default=0.001, help="Learning rate."
|
|
)
|
|
parser.add_argument("--wd", type=float, default=0, help="Weight decay.")
|
|
|
|
# model params
|
|
parser.add_argument(
|
|
"--hid_dim", type=int, default=64, help="Hidden layer dimensionality"
|
|
)
|
|
parser.add_argument(
|
|
"--reg", type=float, default=0.001, help="Regularization coefficient"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# check cuda
|
|
if args.gpu != -1 and th.cuda.is_available():
|
|
args.device = "cuda:{}".format(args.gpu)
|
|
else:
|
|
args.device = "cpu"
|
|
|
|
return args
|
|
|
|
|
|
class DenseQM9EdgeDataset(QM9EdgeDataset):
|
|
def __getitem__(self, idx):
|
|
r"""Get graph and label by index
|
|
|
|
Parameters
|
|
----------
|
|
idx : int
|
|
Item index
|
|
|
|
Returns
|
|
-------
|
|
dgl.DGLGraph
|
|
The graph contains:
|
|
|
|
- ``ndata['pos']``: the coordinates of each atom
|
|
- ``ndata['attr']``: the features of each atom
|
|
- ``edata['edge_attr']``: the features of each bond
|
|
|
|
Tensor
|
|
Property values of molecular graphs
|
|
"""
|
|
|
|
pos = self.node_pos[self.n_cumsum[idx] : self.n_cumsum[idx + 1]]
|
|
src = self.src[self.ne_cumsum[idx] : self.ne_cumsum[idx + 1]]
|
|
dst = self.dst[self.ne_cumsum[idx] : self.ne_cumsum[idx + 1]]
|
|
|
|
g = dgl.graph((src, dst))
|
|
|
|
g.ndata["pos"] = th.tensor(pos).float()
|
|
g.ndata["attr"] = th.tensor(
|
|
self.node_attr[self.n_cumsum[idx] : self.n_cumsum[idx + 1]]
|
|
).float()
|
|
g.edata["edge_attr"] = th.tensor(
|
|
self.edge_attr[self.ne_cumsum[idx] : self.ne_cumsum[idx + 1]]
|
|
).float()
|
|
|
|
label = th.tensor(self.targets[idx][self.label_keys]).float()
|
|
|
|
n_nodes = g.num_nodes()
|
|
row = th.arange(n_nodes)
|
|
col = th.arange(n_nodes)
|
|
|
|
row = row.view(-1, 1).repeat(1, n_nodes).view(-1)
|
|
col = col.repeat(n_nodes)
|
|
|
|
src = g.edges()[0]
|
|
dst = g.edges()[1]
|
|
|
|
idx = src * n_nodes + dst
|
|
size = list(g.edata["edge_attr"].size())
|
|
size[0] = n_nodes * n_nodes
|
|
edge_attr = g.edata["edge_attr"].new_zeros(size)
|
|
|
|
edge_attr[idx] = g.edata["edge_attr"]
|
|
|
|
pos = g.ndata["pos"]
|
|
dist = th.norm(pos[col] - pos[row], p=2, dim=-1).view(-1, 1)
|
|
|
|
new_edge_attr = th.cat([edge_attr, dist.type_as(edge_attr)], dim=-1)
|
|
|
|
graph = dgl.graph((row, col))
|
|
graph.ndata["attr"] = g.ndata["attr"]
|
|
graph.edata["edge_attr"] = new_edge_attr
|
|
graph = graph.remove_self_loop()
|
|
|
|
return graph, label
|
|
|
|
|
|
def collate(samples):
|
|
"""collate function for building graph dataloader"""
|
|
|
|
# generate batched graphs and labels
|
|
graphs, targets = map(list, zip(*samples))
|
|
batched_graph = dgl.batch(graphs)
|
|
batched_targets = th.Tensor(targets)
|
|
|
|
n_graphs = len(graphs)
|
|
graph_id = th.arange(n_graphs)
|
|
graph_id = dgl.broadcast_nodes(batched_graph, graph_id)
|
|
|
|
batched_graph.ndata["graph_id"] = graph_id
|
|
|
|
return batched_graph, batched_targets
|
|
|
|
|
|
def evaluate(model, loader, num, device):
|
|
error = 0
|
|
for graphs, targets in loader:
|
|
graphs = graphs.to(device)
|
|
|
|
nfeat, efeat = graphs.ndata["attr"], graphs.edata["edge_attr"]
|
|
targets = targets.to(device)
|
|
error += (model(graphs, nfeat, efeat) - targets).abs().sum().item()
|
|
|
|
error = error / num
|
|
|
|
return error
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Step 1: Prepare graph data ===================================== #
|
|
args = argument()
|
|
label_keys = [args.target]
|
|
print(args)
|
|
|
|
dataset = DenseQM9EdgeDataset(label_keys=label_keys)
|
|
|
|
# Train/Val/Test Splitting
|
|
N = dataset.targets.shape[0]
|
|
all_idx = np.arange(N)
|
|
np.random.shuffle(all_idx)
|
|
|
|
val_num = 10000
|
|
test_num = 10000
|
|
|
|
val_idx = all_idx[:val_num]
|
|
test_idx = all_idx[val_num : val_num + test_num]
|
|
train_idx = all_idx[
|
|
val_num + test_num : val_num + test_num + args.train_num
|
|
]
|
|
|
|
train_data = Subset(dataset, train_idx)
|
|
val_data = Subset(dataset, val_idx)
|
|
test_data = Subset(dataset, test_idx)
|
|
|
|
unsup_idx = all_idx[val_num + test_num :]
|
|
unsup_data = Subset(dataset, unsup_idx)
|
|
|
|
# generate supervised training dataloader and unsupervised training dataloader
|
|
train_loader = GraphDataLoader(
|
|
train_data,
|
|
batch_size=args.batch_size,
|
|
collate_fn=collate,
|
|
drop_last=False,
|
|
shuffle=True,
|
|
)
|
|
|
|
unsup_loader = GraphDataLoader(
|
|
unsup_data,
|
|
batch_size=args.batch_size,
|
|
collate_fn=collate,
|
|
drop_last=False,
|
|
shuffle=True,
|
|
)
|
|
|
|
# generate validation & testing dataloader
|
|
val_loader = GraphDataLoader(
|
|
val_data,
|
|
batch_size=args.val_batch_size,
|
|
collate_fn=collate,
|
|
drop_last=False,
|
|
shuffle=True,
|
|
)
|
|
|
|
test_loader = GraphDataLoader(
|
|
test_data,
|
|
batch_size=args.val_batch_size,
|
|
collate_fn=collate,
|
|
drop_last=False,
|
|
shuffle=True,
|
|
)
|
|
|
|
print("======== target = {} ========".format(args.target))
|
|
|
|
in_dim = dataset[0][0].ndata["attr"].shape[1]
|
|
|
|
# Step 2: Create model =================================================================== #
|
|
model = InfoGraphS(in_dim, args.hid_dim)
|
|
model = model.to(args.device)
|
|
|
|
# Step 3: Create training components ===================================================== #
|
|
optimizer = th.optim.Adam(
|
|
model.parameters(), lr=args.lr, weight_decay=args.wd
|
|
)
|
|
scheduler = th.optim.lr_scheduler.ReduceLROnPlateau(
|
|
optimizer, mode="min", factor=0.7, patience=5, min_lr=0.000001
|
|
)
|
|
|
|
# Step 4: training epochs =============================================================== #
|
|
best_val_error = float("inf")
|
|
test_error = float("inf")
|
|
|
|
for epoch in range(args.epochs):
|
|
"""Training"""
|
|
model.train()
|
|
lr = scheduler.optimizer.param_groups[0]["lr"]
|
|
|
|
iteration = 0
|
|
sup_loss_all = 0
|
|
unsup_loss_all = 0
|
|
consis_loss_all = 0
|
|
|
|
for sup_data, unsup_data in zip(train_loader, unsup_loader):
|
|
sup_graph, sup_target = sup_data
|
|
unsup_graph, _ = unsup_data
|
|
|
|
sup_graph = sup_graph.to(args.device)
|
|
unsup_graph = unsup_graph.to(args.device)
|
|
|
|
sup_nfeat, sup_efeat = (
|
|
sup_graph.ndata["attr"],
|
|
sup_graph.edata["edge_attr"],
|
|
)
|
|
unsup_nfeat, unsup_efeat, unsup_graph_id = (
|
|
unsup_graph.ndata["attr"],
|
|
unsup_graph.edata["edge_attr"],
|
|
unsup_graph.ndata["graph_id"],
|
|
)
|
|
|
|
sup_target = sup_target
|
|
sup_target = sup_target.to(args.device)
|
|
|
|
optimizer.zero_grad()
|
|
|
|
sup_loss = F.mse_loss(
|
|
model(sup_graph, sup_nfeat, sup_efeat), sup_target
|
|
)
|
|
unsup_loss, consis_loss = model.unsup_forward(
|
|
unsup_graph, unsup_nfeat, unsup_efeat, unsup_graph_id
|
|
)
|
|
|
|
loss = sup_loss + unsup_loss + args.reg * consis_loss
|
|
|
|
loss.backward()
|
|
|
|
sup_loss_all += sup_loss.item()
|
|
unsup_loss_all += unsup_loss.item()
|
|
consis_loss_all += consis_loss.item()
|
|
|
|
optimizer.step()
|
|
|
|
print(
|
|
"Epoch: {}, Sup_Loss: {:4f}, Unsup_loss: {:.4f}, Consis_loss: {:.4f}".format(
|
|
epoch, sup_loss_all, unsup_loss_all, consis_loss_all
|
|
)
|
|
)
|
|
|
|
model.eval()
|
|
|
|
val_error = evaluate(model, val_loader, val_num, args.device)
|
|
scheduler.step(val_error)
|
|
|
|
if val_error < best_val_error:
|
|
best_val_error = val_error
|
|
test_error = evaluate(model, test_loader, test_num, args.device)
|
|
|
|
print(
|
|
"Epoch: {}, LR: {}, val_error: {:.4f}, best_test_error: {:.4f}".format(
|
|
epoch, lr, val_error, test_error
|
|
)
|
|
)
|