from transform import JunctionTree
from model import Net
parser = argparse.ArgumentParser()
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--hidden_channels', type=int, default=128)
parser.add_argument('--num_layers', type=int, default=3)
parser.add_argument('--dropout', type=float, default=0.0)
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--no_inter_message_passing', action='store_true')
args = parser.parse_args()
print(args)
root = 'data/ZINC'
transform = JunctionTree()
train_dataset = ZINC(root, subset=True, split='train', pre_transform=transform)
val_dataset = ZINC(root, subset=True, split='val', pre_transform=transform)
test_dataset = ZINC(root, subset=True, split='test', pre_transform=transform)
train_loader = DataLoader(train_dataset, 128, shuffle=True, num_workers=12)
val_loader = DataLoader(val_dataset, 1000, shuffle=False, num_workers=12)
test_loader = DataLoader(test_dataset, 1000, shuffle=False, num_workers=12)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
model = Net(hidden_channels=args.hidden_channels,
out_channels=1,
num_layers=args.num_layers,
dropout=args.dropout,
inter_message_passing=not args.no_inter_message_passing).to(device)
parser.add_argument('--dropout', type=float, default=0.5)
parser.add_argument('--epochs', type=int, default=100)
parser.add_argument('--no_inter_message_passing', action='store_true')
args = parser.parse_args()
print(args)
class OGBTransform(object):
# OGB saves atom and bond types zero-index based. We need to revert that.
def __call__(self, data):
data.x[:, 0] += 1
data.edge_attr[:, 0] += 1
return data
transform = Compose([OGBTransform(), JunctionTree()])
name = 'ogbg-molhiv'
dataset = PygGraphPropPredDataset(name, 'data', pre_transform=transform)
dataset = dataset.shuffle()
num_train, num_trainval = round(0.8 * len(dataset)), round(0.9 * len(dataset))
train_dataset = dataset[:num_train]
val_dataset = dataset[num_train:num_trainval]
test_dataset = dataset[num_trainval:]
train_loader = DataLoader(train_dataset, 128, shuffle=True, num_workers=12)
val_loader = DataLoader(val_dataset, 1000, shuffle=False, num_workers=12)
test_loader = DataLoader(test_dataset, 1000, shuffle=False, num_workers=12)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'