import torch
import torch.nn.functional as F
from torch.nn import ModuleList, Embedding
from torch.nn import Sequential, ReLU, Linear
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch_geometric.utils import degree
from torch_geometric.datasets import ZINC
from torch_geometric.loader import DataLoader
from torch_geometric.nn import PNAConv, BatchNorm, global_add_pool
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ZINC')
train_dataset = ZINC(path, subset=True, split='train')
val_dataset = ZINC(path, subset=True, split='val')
test_dataset = ZINC(path, subset=True, split='test')
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=128)
test_loader = DataLoader(test_dataset, batch_size=128)
# Compute in-degree histogram over training data.
deg = torch.zeros(5, dtype=torch.long)
for data in train_dataset:
d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)
deg += torch.bincount(d, minlength=deg.numel())
class Net(torch.nn.Module):
def __init__(self):
super().__init__()
self.node_emb = Embedding(21, 75)
import os.path as osp
import torch
import torch.nn.functional as F
from torch_geometric.datasets import MNISTSuperpixels
import torch_geometric.transforms as T
from torch_geometric.loader import DataLoader
from torch_geometric.nn import SplineConv, voxel_grid, max_pool, max_pool_x
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST')
transform = T.Cartesian(cat=False)
train_dataset = MNISTSuperpixels(path, True, transform=transform)
test_dataset = MNISTSuperpixels(path, False, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64)
d = train_dataset
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = SplineConv(d.num_features, 32, dim=2, kernel_size=5)
self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)
self.conv3 = SplineConv(64, 64, dim=2, kernel_size=5)
self.fc1 = torch.nn.Linear(4 * 64, 128)
self.fc2 = torch.nn.Linear(128, d.num_classes)
def forward(self, data):
data.x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr))
cluster = voxel_grid(data.pos, data.batch, size=5, start=0, end=28)
data.edge_attr = None
for data in loader:
data = data.to(device)
with torch.no_grad():
pred = model(data).max(1)[1]
correct += pred.eq(data.y).sum().item()
return correct / len(loader.dataset)
if __name__ == '__main__':
path = osp.join(osp.dirname(osp.realpath(__file__)), '..',
'data/ModelNet10')
pre_transform, transform = T.NormalizeScale(), T.SamplePoints(1024)
train_dataset = ModelNet(path, '10', True, transform, pre_transform)
test_dataset = ModelNet(path, '10', False, transform, pre_transform)
train_loader = DataLoader(train_dataset,
batch_size=32,
shuffle=True,
num_workers=6)
test_loader = DataLoader(test_dataset,
batch_size=32,
shuffle=False,
num_workers=6)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(1, 201):
train(epoch)
test_acc = test(test_loader)
print(f'Epoch: {epoch:03d}, Test: {test_acc:.4f}')
import os.path as osp
import torch
import torch.nn.functional as F
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GraphConv, TopKPooling
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'PROTEINS')
dataset = TUDataset(path, name='PROTEINS')
dataset = dataset.shuffle()
n = len(dataset) // 10
test_dataset = dataset[:n]
train_dataset = dataset[n:]
test_loader = DataLoader(test_dataset, batch_size=60)
train_loader = DataLoader(train_dataset, batch_size=60)
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = GraphConv(dataset.num_features, 128)
self.pool1 = TopKPooling(128, ratio=0.8)
self.conv2 = GraphConv(128, 128)
self.pool2 = TopKPooling(128, ratio=0.8)
self.conv3 = GraphConv(128, 128)
self.pool3 = TopKPooling(128, ratio=0.8)
self.lin1 = torch.nn.Linear(256, 128)
def cross_validation_with_val_set(dataset,
model,
folds,
epochs,
batch_size,
lr,
lr_decay_factor,
lr_decay_step_size,
weight_decay,
logger=None):
val_losses, accs, durations = [], [], []
for fold, (train_idx, test_idx,
val_idx) in enumerate(zip(*k_fold(dataset, folds))):
train_dataset = dataset[train_idx]
test_dataset = dataset[test_idx]
val_dataset = dataset[val_idx]
if 'adj' in train_dataset[0]:
train_loader = DenseLoader(train_dataset, batch_size, shuffle=True)
val_loader = DenseLoader(val_dataset, batch_size, shuffle=False)
test_loader = DenseLoader(test_dataset, batch_size, shuffle=False)
else:
train_loader = DataLoader(train_dataset, batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size, shuffle=False)
model.to(device).reset_parameters()
optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
if torch.cuda.is_available():
torch.cuda.synchronize()
t_start = time.perf_counter()
for epoch in range(1, epochs + 1):
train_loss = train(model, optimizer, train_loader)
val_losses.append(eval_loss(model, val_loader))
accs.append(eval_acc(model, test_loader))
eval_info = {
'fold': fold,
'epoch': epoch,
'train_loss': train_loss,
'val_loss': val_losses[-1],
'test_acc': accs[-1],
}
if logger is not None:
logger(eval_info)
if epoch % lr_decay_step_size == 0:
for param_group in optimizer.param_groups:
param_group['lr'] = lr_decay_factor * param_group['lr']
if torch.cuda.is_available():
torch.cuda.synchronize()
t_end = time.perf_counter()
durations.append(t_end - t_start)
loss, acc, duration = tensor(val_losses), tensor(accs), tensor(durations)
loss, acc = loss.view(folds, epochs), acc.view(folds, epochs)
loss, argmin = loss.min(dim=1)
acc = acc[torch.arange(folds, dtype=torch.long), argmin]
loss_mean = loss.mean().item()
acc_mean = acc.mean().item()
acc_std = acc.std().item()
duration_mean = duration.mean().item()
print(f'Val Loss: {loss_mean:.4f}, Test Accuracy: {acc_mean:.3f} '
f'± {acc_std:.3f}, Duration: {duration_mean:.3f}')
return loss_mean, acc_mean, acc_std
def train_dataloader(self):
return DataLoader(self.train_dataset, batch_size=64, shuffle=True)
class HandleNodeAttention(object):
def __call__(self, data):
data.attn = torch.softmax(data.x, dim=0).flatten()
data.x = None
return data
transform = T.Compose([HandleNodeAttention(), T.OneHotDegree(max_degree=14)])
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'TRIANGLES')
dataset = TUDataset(path,
name='TRIANGLES',
use_node_attr=True,
transform=transform)
train_loader = DataLoader(dataset[:30000], batch_size=60, shuffle=True)
val_loader = DataLoader(dataset[30000:35000], batch_size=60)
test_loader = DataLoader(dataset[35000:], batch_size=60)
class Net(torch.nn.Module):
def __init__(self, in_channels):
super(Net, self).__init__()
self.conv1 = GINConv(Seq(Lin(in_channels, 64), ReLU(), Lin(64, 64)))
self.pool1 = SAGPooling(64, min_score=0.001, GNN=GCNConv)
self.conv2 = GINConv(Seq(Lin(64, 64), ReLU(), Lin(64, 64)))
self.pool2 = SAGPooling(64, min_score=0.001, GNN=GCNConv)
self.conv3 = GINConv(Seq(Lin(64, 64), ReLU(), Lin(64, 64)))
self.lin = torch.nn.Linear(64, 1)
args = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
init_wandb(name=f'GIN-{args.dataset}',
batch_size=args.batch_size,
lr=args.lr,
epochs=args.epochs,
hidden_channels=args.hidden_channels,
num_layers=args.num_layers,
device=device)
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'TU')
dataset = TUDataset(path, name=args.dataset).shuffle()
train_dataset = dataset[len(dataset) // 10:]
train_loader = DataLoader(train_dataset, args.batch_size, shuffle=True)
test_dataset = dataset[:len(dataset) // 10]
test_loader = DataLoader(test_dataset, args.batch_size)
class Net(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels, num_layers):
super().__init__()
self.convs = torch.nn.ModuleList()
for _ in range(num_layers):
mlp = MLP([in_channels, hidden_channels, hidden_channels])
self.convs.append(GINConv(nn=mlp, train_eps=False))
in_channels = hidden_channels
from torch_geometric.loader import DataLoader
from torch_geometric.nn.models.re_net import RENet
seq_len = 10
# Load the dataset and precompute history objects.
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ICEWS18')
train_dataset = ICEWS18(path, pre_transform=RENet.pre_transform(seq_len))
test_dataset = ICEWS18(path, split='test')
# path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'GDELT')
# train_dataset = GDELT(path, pre_transform=RENet.pre_transform(seq_len))
# test_dataset = ICEWS18(path, split='test')
# Create dataloader for training and test dataset.
train_loader = DataLoader(train_dataset, batch_size=1024, shuffle=True,
follow_batch=['h_sub', 'h_obj'], num_workers=6)
test_loader = DataLoader(test_dataset, batch_size=1024, shuffle=False,
follow_batch=['h_sub', 'h_obj'], num_workers=6)
# Initialize model and optimizer.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = RENet(
train_dataset.num_nodes,
train_dataset.num_rels,
hidden_channels=200,
seq_len=seq_len,
dropout=0.5,
).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001,
weight_decay=0.00001)
from torch_geometric.datasets import QM9
from torch_geometric.loader import DataLoader
from torch_geometric.nn import SchNet
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'QM9')
dataset = QM9(path)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
for target in range(12):
model, datasets = SchNet.from_qm9_pretrained(path, dataset, target)
train_dataset, val_dataset, test_dataset = datasets
model = model.to(device)
loader = DataLoader(test_dataset, batch_size=256)
maes = []
for data in loader:
data = data.to(device)
with torch.no_grad():
pred = model(data.z, data.pos, data.batch)
mae = (pred.view(-1) - data.y[:, target]).abs()
maes.append(mae)
mae = torch.cat(maes, dim=0)
# Report meV instead of eV.
mae = 1000 * mae if target in [2, 3, 4, 6, 7, 8, 9, 10] else mae
print(f'Target: {target:02d}, MAE: {mae.mean():.5f} ± {mae.std():.5f}')
def test_to_superpixels():
import torchvision.transforms as T
from torchvision.datasets.mnist import (
MNIST,
read_image_file,
read_label_file,
)
root = osp.join('/', 'tmp', str(random.randrange(sys.maxsize)))
raw_folder = osp.join(root, 'MNIST', 'raw')
processed_folder = osp.join(root, 'MNIST', 'processed')
makedirs(raw_folder)
makedirs(processed_folder)
for resource in resources:
path = download_url(resource, raw_folder)
extract_gz(path, osp.join(root, raw_folder))
test_set = (
read_image_file(osp.join(raw_folder, 't10k-images-idx3-ubyte')),
read_label_file(osp.join(raw_folder, 't10k-labels-idx1-ubyte')),
)
torch.save(test_set, osp.join(processed_folder, 'training.pt'))
torch.save(test_set, osp.join(processed_folder, 'test.pt'))
dataset = MNIST(root, download=False)
dataset.transform = T.Compose([T.ToTensor(), ToSLIC()])
data, y = dataset[0]
assert len(data) == 2
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.pos.size(0) == data.x.size(0)
assert y == 7
loader = DataLoader(dataset, batch_size=2, shuffle=False)
for data, y in loader:
assert len(data) == 4
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.batch.dim() == 1
assert data.ptr.dim() == 1
assert data.pos.size(0) == data.x.size(0) == data.batch.size(0)
assert y.tolist() == [7, 2]
break
dataset.transform = T.Compose(
[T.ToTensor(), ToSLIC(add_seg=True, add_img=True)])
data, y = dataset[0]
assert len(data) == 4
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.pos.size(0) == data.x.size(0)
assert data.seg.size() == (1, 28, 28)
assert data.img.size() == (1, 1, 28, 28)
assert data.seg.max().item() + 1 == data.x.size(0)
assert y == 7
loader = DataLoader(dataset, batch_size=2, shuffle=False)
for data, y in loader:
assert len(data) == 6
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.batch.dim() == 1
assert data.ptr.dim() == 1
assert data.pos.size(0) == data.x.size(0) == data.batch.size(0)
assert data.seg.size() == (2, 28, 28)
assert data.img.size() == (2, 1, 28, 28)
assert y.tolist() == [7, 2]
break
shutil.rmtree(root)
parser = argparse.ArgumentParser()
parser.add_argument('--use_multi_aggregators', action='store_true',
help='Switch between EGC-S and EGC-M')
args = parser.parse_args()
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'OGB')
dataset = OGBG('ogbg-molhiv', path, pre_transform=T.ToSparseTensor())
evaluator = Evaluator('ogbg-molhiv')
split_idx = dataset.get_idx_split()
train_dataset = dataset[split_idx['train']]
val_dataset = dataset[split_idx['valid']]
test_dataset = dataset[split_idx['test']]
train_loader = DataLoader(train_dataset, batch_size=32, num_workers=4,
shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=256)
test_loader = DataLoader(test_dataset, batch_size=256)
class Net(torch.nn.Module):
def __init__(self, hidden_channels, num_layers, num_heads, num_bases):
super().__init__()
if args.use_multi_aggregators:
aggregators = ['sum', 'mean', 'max']
else:
aggregators = ['symnorm']
self.encoder = AtomEncoder(hidden_channels)
self.convs = torch.nn.ModuleList()
def test_dataloader(self):
return DataLoader(self.test_dataset, batch_size=128)
def val_dataloader(self):
return DataLoader(self.val_dataset, batch_size=128)
train_dataset = PPI(path, split='train')
val_dataset = PPI(path, split='val')
test_dataset = PPI(path, split='test')
# Group all training graphs into a single graph to perform sampling:
train_data = Batch.from_data_list(train_dataset)
loader = LinkNeighborLoader(train_data,
batch_size=2048,
shuffle=True,
neg_sampling_ratio=0.5,
num_neighbors=[10, 10],
num_workers=6,
persistent_workers=True)
# Evaluation loaders (one datapoint corresponds to a graph)
train_loader = DataLoader(train_dataset, batch_size=2)
val_loader = DataLoader(val_dataset, batch_size=2)
test_loader = DataLoader(test_dataset, batch_size=2)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GraphSAGE(
in_channels=train_dataset.num_features,
hidden_channels=64,
num_layers=2,
out_channels=64,
).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005)
def train():
model.train()
def main():
# Training settings
parser = argparse.ArgumentParser(
description='GNN baselines on ogbgmol* data with Pytorch Geometrics')
parser.add_argument('--device',
type=int,
default=0,
help='which gpu to use if any (default: 0)')
parser.add_argument(
'--gnn',
type=str,
default='gin-virtual',
help=
'GNN gin, gin-virtual, or gcn, or gcn-virtual (default: gin-virtual)')
parser.add_argument('--drop_ratio',
type=float,
default=0.5,
help='dropout ratio (default: 0.5)')
parser.add_argument(
'--num_layer',
type=int,
default=5,
help='number of GNN message passing layers (default: 5)')
parser.add_argument(
'--emb_dim',
type=int,
default=300,
help='dimensionality of hidden units in GNNs (default: 300)')
parser.add_argument('--batch_size',
type=int,
default=32,
help='input batch size for training (default: 32)')
parser.add_argument('--epochs',
type=int,
default=100,
help='number of epochs to train (default: 100)')
parser.add_argument('--num_workers',
type=int,
default=0,
help='number of workers (default: 0)')
parser.add_argument('--dataset',
type=str,
default="ogbg-molhiv",
help='dataset name (default: ogbg-molhiv)')
parser.add_argument('--feature',
type=str,
default="full",
help='full feature or simple feature')
parser.add_argument('--filename',
type=str,
default="",
help='filename to output result (default: )')
args = parser.parse_args()
device = torch.device(
"cuda:" +
str(args.device)) if torch.cuda.is_available() else torch.device("cpu")
### automatic dataloading and splitting
dataset = PygGraphPropPredDataset(name=args.dataset)
if args.feature == 'full':
pass
elif args.feature == 'simple':
print('using simple feature')
# only retain the top two node/edge features
dataset.data.x = dataset.data.x[:, :2]
dataset.data.edge_attr = dataset.data.edge_attr[:, :2]
split_idx = dataset.get_idx_split()
### automatic evaluator. takes dataset name as input
evaluator = Evaluator(args.dataset)
train_loader = DataLoader(dataset[split_idx["train"]],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
valid_loader = DataLoader(dataset[split_idx["valid"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
test_loader = DataLoader(dataset[split_idx["test"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
if args.gnn == 'gin':
model = GNN(gnn_type='gin',
num_tasks=dataset.num_tasks,
num_layer=args.num_layer,
emb_dim=args.emb_dim,
drop_ratio=args.drop_ratio,
virtual_node=False).to(device)
elif args.gnn == 'gin-virtual':
model = GNN(gnn_type='gin',
num_tasks=dataset.num_tasks,
num_layer=args.num_layer,
emb_dim=args.emb_dim,
drop_ratio=args.drop_ratio,
virtual_node=True).to(device)
elif args.gnn == 'gcn':
model = GNN(gnn_type='gcn',
num_tasks=dataset.num_tasks,
num_layer=args.num_layer,
emb_dim=args.emb_dim,
drop_ratio=args.drop_ratio,
virtual_node=False).to(device)
elif args.gnn == 'gcn-virtual':
model = GNN(gnn_type='gcn',
num_tasks=dataset.num_tasks,
num_layer=args.num_layer,
emb_dim=args.emb_dim,
drop_ratio=args.drop_ratio,
virtual_node=True).to(device)
else:
raise ValueError('Invalid GNN type')
optimizer = optim.Adam(model.parameters(), lr=0.001)
valid_curve = []
test_curve = []
train_curve = []
for epoch in range(1, args.epochs + 1):
print("=====Epoch {}".format(epoch))
print('Training...')
train(model, device, train_loader, optimizer, dataset.task_type)
print('Evaluating...')
train_perf = eval(model, device, train_loader, evaluator)
valid_perf = eval(model, device, valid_loader, evaluator)
test_perf = eval(model, device, test_loader, evaluator)
print({
'Train': train_perf,
'Validation': valid_perf,
'Test': test_perf
})
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 '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]))
if not args.filename == '':
torch.save(
{
'Val': valid_curve[best_val_epoch],
'Test': test_curve[best_val_epoch],
'Train': train_curve[best_val_epoch],
'BestTrain': best_train
}, args.filename)
import os.path as osp
import torch
import torch.nn.functional as F
from torch_geometric.datasets import PPI
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GATConv
from sklearn.metrics import f1_score
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'PPI')
train_dataset = PPI(path, split='train')
val_dataset = PPI(path, split='val')
test_dataset = PPI(path, split='test')
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False)
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = GATConv(train_dataset.num_features, 256, heads=4)
self.lin1 = torch.nn.Linear(train_dataset.num_features, 4 * 256)
self.conv2 = GATConv(4 * 256, 256, heads=4)
self.lin2 = torch.nn.Linear(4 * 256, 4 * 256)
self.conv3 = GATConv(4 * 256,
train_dataset.num_classes,
heads=6,
concat=False)
self.lin3 = torch.nn.Linear(4 * 256, train_dataset.num_classes)
def main():
# Training settings
parser = argparse.ArgumentParser(
description='GNN baselines on pcqm4m with Pytorch Geometrics')
parser.add_argument('--device',
type=int,
default=0,
help='which gpu to use if any (default: 0)')
parser.add_argument(
'--gnn',
type=str,
default='gin-virtual',
help=
'GNN gin, gin-virtual, or gcn, or gcn-virtual (default: gin-virtual)')
parser.add_argument(
'--graph_pooling',
type=str,
default='sum',
help='graph pooling strategy mean or sum (default: sum)')
parser.add_argument('--drop_ratio',
type=float,
default=0,
help='dropout ratio (default: 0)')
parser.add_argument(
'--num_layers',
type=int,
default=5,
help='number of GNN message passing layers (default: 5)')
parser.add_argument(
'--emb_dim',
type=int,
default=600,
help='dimensionality of hidden units in GNNs (default: 600)')
parser.add_argument('--train_subset', action='store_true')
parser.add_argument('--batch_size',
type=int,
default=256,
help='input batch size for training (default: 256)')
parser.add_argument('--epochs',
type=int,
default=100,
help='number of epochs to train (default: 100)')
parser.add_argument('--num_workers',
type=int,
default=0,
help='number of workers (default: 0)')
parser.add_argument('--log_dir',
type=str,
default="",
help='tensorboard log directory')
parser.add_argument('--checkpoint_dir',
type=str,
default='',
help='directory to save checkpoint')
parser.add_argument('--save_test_dir',
type=str,
default='',
help='directory to save test submission file')
args = parser.parse_args()
print(args)
np.random.seed(42)
torch.manual_seed(42)
torch.cuda.manual_seed(42)
random.seed(42)
device = torch.device(
"cuda:" +
str(args.device)) if torch.cuda.is_available() else torch.device("cpu")
### automatic dataloading and splitting
dataset = PygPCQM4Mv2Dataset(root='dataset/')
split_idx = dataset.get_idx_split()
### automatic evaluator. takes dataset name as input
evaluator = PCQM4Mv2Evaluator()
if args.train_subset:
subset_ratio = 0.1
subset_idx = torch.randperm(len(
split_idx["train"]))[:int(subset_ratio * len(split_idx["train"]))]
train_loader = DataLoader(dataset[split_idx["train"][subset_idx]],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
else:
train_loader = DataLoader(dataset[split_idx["train"]],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
valid_loader = DataLoader(dataset[split_idx["valid"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
if args.save_test_dir != '':
testdev_loader = DataLoader(dataset[split_idx["test-dev"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
testchallenge_loader = DataLoader(dataset[split_idx["test-challenge"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
if args.checkpoint_dir != '':
os.makedirs(args.checkpoint_dir, exist_ok=True)
shared_params = {
'num_layers': args.num_layers,
'emb_dim': args.emb_dim,
'drop_ratio': args.drop_ratio,
'graph_pooling': args.graph_pooling
}
if args.gnn == 'gin':
model = GNN(gnn_type='gin', virtual_node=False,
**shared_params).to(device)
elif args.gnn == 'gin-virtual':
model = GNN(gnn_type='gin', virtual_node=True,
**shared_params).to(device)
elif args.gnn == 'gcn':
model = GNN(gnn_type='gcn', virtual_node=False,
**shared_params).to(device)
elif args.gnn == 'gcn-virtual':
model = GNN(gnn_type='gcn', virtual_node=True,
**shared_params).to(device)
else:
raise ValueError('Invalid GNN type')
num_params = sum(p.numel() for p in model.parameters())
print(f'#Params: {num_params}')
optimizer = optim.Adam(model.parameters(), lr=0.001)
if args.log_dir != '':
writer = SummaryWriter(log_dir=args.log_dir)
best_valid_mae = 1000
if args.train_subset:
scheduler = StepLR(optimizer, step_size=300, gamma=0.25)
args.epochs = 1000
else:
scheduler = StepLR(optimizer, step_size=30, gamma=0.25)
for epoch in range(1, args.epochs + 1):
print("=====Epoch {}".format(epoch))
print('Training...')
train_mae = train(model, device, train_loader, optimizer)
print('Evaluating...')
valid_mae = eval(model, device, valid_loader, evaluator)
print({'Train': train_mae, 'Validation': valid_mae})
if args.log_dir != '':
writer.add_scalar('valid/mae', valid_mae, epoch)
writer.add_scalar('train/mae', train_mae, epoch)
if valid_mae < best_valid_mae:
best_valid_mae = valid_mae
if args.checkpoint_dir != '':
print('Saving checkpoint...')
checkpoint = {
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'best_val_mae': best_valid_mae,
'num_params': num_params
}
torch.save(checkpoint,
os.path.join(args.checkpoint_dir, 'checkpoint.pt'))
if args.save_test_dir != '':
testdev_pred = test(model, device, testdev_loader)
testdev_pred = testdev_pred.cpu().detach().numpy()
testchallenge_pred = test(model, device, testchallenge_loader)
testchallenge_pred = testchallenge_pred.cpu().detach().numpy()
print('Saving test submission file...')
evaluator.save_test_submission({'y_pred': testdev_pred},
args.save_test_dir,
mode='test-dev')
evaluator.save_test_submission({'y_pred': testchallenge_pred},
args.save_test_dir,
mode='test-challenge')
scheduler.step()
print(f'Best validation MAE so far: {best_valid_mae}')
if args.log_dir != '':
writer.close()
