def main(): # Training settings parser = argparse.ArgumentParser(description='PyTorch implementation of pre-training of graph neural networks') parser.add_argument('--device', type=int, default=0, help='which gpu to use if any (default: 0)') 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('--lr', type=float, default=0.001, help='learning rate (default: 0.001)') parser.add_argument('--decay', type=float, default=0, help='weight decay (default: 0)') 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='embedding dimensions (default: 300)') parser.add_argument('--dropout_ratio', type=float, default=0, help='dropout ratio (default: 0)') parser.add_argument('--JK', type=str, default="last", help='how the node features across layers are combined. last, sum, max or concat') parser.add_argument('--dataset', type=str, default = 'zinc_standard_agent', help='root directory of dataset. For now, only classification.') parser.add_argument('--output_model_file', type = str, default = '', help='filename to output the pre-trained model') parser.add_argument('--gnn_type', type=str, default="gin") parser.add_argument('--num_workers', type=int, default = 8, help='number of workers for dataset loading') args = parser.parse_args() torch.manual_seed(0) np.random.seed(0) device = torch.device("cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu") if torch.cuda.is_available(): torch.cuda.manual_seed_all(0) #set up dataset dataset = MoleculeDataset("dataset/" + args.dataset, dataset=args.dataset, transform = NegativeEdge()) print(dataset[0]) loader = DataLoaderAE(dataset, batch_size=args.batch_size, shuffle=True, num_workers = args.num_workers) #set up model model = GNN(args.num_layer, args.emb_dim, JK = args.JK, drop_ratio = args.dropout_ratio, gnn_type = args.gnn_type) model.to(device) #set up optimizer optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.decay) print(optimizer) for epoch in range(1, args.epochs+1): print("====epoch " + str(epoch)) train_acc, train_loss = train(args, model, device, loader, optimizer) print(train_acc) print(train_loss) if not args.output_model_file == "": torch.save(model.state_dict(), args.output_model_file + ".pth")
model = GNN(*data_train.size, embed_size, dropout).to(device) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) for epoch in range(num_epochs): model.train() y_pred = model(data_train).reshape(-1) y = data_train.edge_type.float() loss = criterion(y_pred, y) optimizer.zero_grad() loss.backward() optimizer.step() rmse_loss_train = loss.item()**0.5 model.eval() with torch.no_grad(): y_pred = model(data_test).reshape(-1) y = data_test.edge_type.float() rmse_loss_test = criterion(y_pred, y).item()**0.5 print(f' [Epoch {epoch + 1:3}/{num_epochs}]', end=' ') print( f'RMSE Loss (Train: {rmse_loss_train:6.4f} | Test: {rmse_loss_test:6.4f})' ) torch.save(model.state_dict(), f'model/model.ckpt')
def main(): # Training settings parser = argparse.ArgumentParser( description= 'PyTorch implementation of pre-training of graph neural networks') parser.add_argument('--device', type=int, default=0, help='which gpu to use if any (default: 0)') parser.add_argument('--batch_size', type=int, default=2, 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('--lr', type=float, default=0.001, help='learning rate (default: 0.001)') parser.add_argument('--decay', type=float, default=0, help='weight decay (default: 0)') parser.add_argument( '--num_layer', type=int, default=5, help='number of GNN message passing layers (default: 5).') parser.add_argument('--l1', type=int, default=1, help='l1 (default: 1).') parser.add_argument('--center', type=int, default=0, help='center (default: 0).') parser.add_argument('--emb_dim', type=int, default=300, help='embedding dimensions (default: 300)') parser.add_argument('--dropout_ratio', type=float, default=0, help='dropout ratio (default: 0)') parser.add_argument( '--neg_samples', type=int, default=1, help='number of negative contexts per positive context (default: 1)') parser.add_argument( '--JK', type=str, default="last", help= 'how the node features are combined across layers. last, sum, max or concat' ) parser.add_argument('--context_pooling', type=str, default="mean", help='how the contexts are pooled (sum, mean, or max)') parser.add_argument('--gnn_type', type=str, default="gat") parser.add_argument('--mode', type=str, default="cbow", help="cbow or skipgram") parser.add_argument('--model_file', type=str, default='', help='filename to output the model') parser.add_argument('--num_workers', type=int, default=4, help='number of workers for dataset loading') args = parser.parse_args() torch.manual_seed(0) np.random.seed(0) device = torch.device( "cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu") if torch.cuda.is_available(): torch.cuda.manual_seed_all(0) print(args.mode) #set up dataset root_unsupervised = 'dataset/unsupervised' dataset = BioDataset(root_unsupervised, data_type='unsupervised', transform=ExtractSubstructureContextPair( l1=args.l1, center=args.center)) print(dataset[0], "\n", dataset[1], "\n", len(dataset)) print("l1: " + str(args.l1)) print("center: " + str(args.center)) loader = DataLoaderSubstructContext(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) #print(dataset[0]) #set up models, one for pre-training and one for context embeddings model_substruct = GNN(args.num_layer, args.emb_dim, JK=args.JK, drop_ratio=args.dropout_ratio, gnn_type=args.gnn_type).to(device) model_context = GNN(3, args.emb_dim, JK=args.JK, drop_ratio=args.dropout_ratio, gnn_type=args.gnn_type).to(device) #set up optimizer for the two GNNs optimizer_substruct = optim.Adam(model_substruct.parameters(), lr=args.lr, weight_decay=args.decay) optimizer_context = optim.Adam(model_context.parameters(), lr=args.lr, weight_decay=args.decay) for epoch in range(1, args.epochs + 1): print("====epoch " + str(epoch)) train_loss, train_acc = train(args, model_substruct, model_context, loader, optimizer_substruct, optimizer_context, device) print(train_loss, train_acc) if not args.model_file == "": torch.save(model_substruct.state_dict(), args.model_file + ".pth")
def main(): # Training settings parser = argparse.ArgumentParser( description= 'PyTorch implementation of pre-training of graph neural networks') parser.add_argument('--device', type=int, default=0, help='which gpu to use if any (default: 0)') 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('--lr', type=float, default=0.001, help='learning rate (default: 0.001)') parser.add_argument('--decay', type=float, default=0, help='weight decay (default: 0)') 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='embedding dimensions (default: 300)') parser.add_argument('--dropout_ratio', type=float, default=0, help='dropout ratio (default: 0)') parser.add_argument('--mask_rate', type=float, default=0.15, help='dropout ratio (default: 0.15)') parser.add_argument( '--mask_edge', type=int, default=0, help='whether to mask edges or not together with atoms') parser.add_argument( '--JK', type=str, default="last", help= 'how the node features are combined across layers. last, sum, max or concat' ) parser.add_argument('--dataset', type=str, default='zinc_standard_agent', help='root directory of dataset for pretraining') parser.add_argument('--output_model_file', type=str, default='', help='filename to output the model') parser.add_argument('--gnn_type', type=str, default="gin") parser.add_argument('--seed', type=int, default=0, help="Seed for splitting dataset.") parser.add_argument('--num_workers', type=int, default=8, help='number of workers for dataset loading') args = parser.parse_args() torch.manual_seed(0) np.random.seed(0) device = torch.device( "cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu") if torch.cuda.is_available(): torch.cuda.manual_seed_all(0) print("num layer: %d mask rate: %f mask edge: %d" % (args.num_layer, args.mask_rate, args.mask_edge)) #set up dataset and transform function. dataset = MoleculeDataset("dataset/" + args.dataset, dataset=args.dataset, transform=MaskAtom(num_atom_type=119, num_edge_type=5, mask_rate=args.mask_rate, mask_edge=args.mask_edge)) loader = DataLoaderMasking(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) #set up models, one for pre-training and one for context embeddings model = GNN(args.num_layer, args.emb_dim, JK=args.JK, drop_ratio=args.dropout_ratio, gnn_type=args.gnn_type).to(device) linear_pred_atoms = torch.nn.Linear(args.emb_dim, 119).to(device) linear_pred_bonds = torch.nn.Linear(args.emb_dim, 4).to(device) model_list = [model, linear_pred_atoms, linear_pred_bonds] #set up optimizers optimizer_model = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.decay) optimizer_linear_pred_atoms = optim.Adam(linear_pred_atoms.parameters(), lr=args.lr, weight_decay=args.decay) optimizer_linear_pred_bonds = optim.Adam(linear_pred_bonds.parameters(), lr=args.lr, weight_decay=args.decay) optimizer_list = [ optimizer_model, optimizer_linear_pred_atoms, optimizer_linear_pred_bonds ] for epoch in range(1, args.epochs + 1): print("====epoch " + str(epoch)) train_loss, train_acc_atom, train_acc_bond = train( args, model_list, loader, optimizer_list, device) print(train_loss, train_acc_atom, train_acc_bond) if not args.output_model_file == "": torch.save(model.state_dict(), args.output_model_file + ".pth")
def main(): # Training settings parser = argparse.ArgumentParser( description= 'PyTorch implementation of pre-training of graph neural networks') parser.add_argument('--device', type=int, default=0, help='which gpu to use if any (default: 0)') parser.add_argument('--batch_size', type=int, default=32, 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('--lr', type=float, default=0.001, help='learning rate (default: 0.001)') parser.add_argument('--decay', type=float, default=0, help='weight decay (default: 0)') 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='embedding dimensions (default: 300)') parser.add_argument('--dropout_ratio', type=float, default=0, help='dropout ratio (default: 0)') parser.add_argument('--mask_rate', type=float, default=0.15, help='dropout ratio (default: 0.15)') parser.add_argument( '--JK', type=str, default="last", help= 'how the node features are combined across layers. last, sum, max or concat' ) parser.add_argument('--gnn_type', type=str, default="gsan") parser.add_argument('--model_file', type=str, default='', help='filename to output the model') parser.add_argument('--seed', type=int, default=0, help="Seed for splitting dataset.") parser.add_argument('--num_workers', type=int, default=8, help='number of workers for dataset loading') args = parser.parse_args() torch.manual_seed(0) np.random.seed(0) device = torch.device( "cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu") if torch.cuda.is_available(): torch.cuda.manual_seed_all(0) print("num layer: %d mask rate: %f" % (args.num_layer, args.mask_rate)) #set up dataset root_unsupervised = 'dataset/unsupervised' dataset = BioDataset(root_unsupervised, data_type='unsupervised', transform=MaskEdge(mask_rate=args.mask_rate)) print(dataset) loader = DataLoaderMasking(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) #set up models, one for pre-training and one for context embeddings model = GNN(args.num_layer, args.emb_dim, JK=args.JK, drop_ratio=args.dropout_ratio, gnn_type=args.gnn_type).to(device) #Linear layer for classifying different edge types linear_pred_edges = torch.nn.Linear(args.emb_dim, 7).to(device) model_list = [model, linear_pred_edges] #set up optimizers optimizer_model = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.decay) optimizer_linear_pred_edges = optim.Adam(linear_pred_edges.parameters(), lr=args.lr, weight_decay=args.decay) optimizer_list = [optimizer_model, optimizer_linear_pred_edges] for epoch in range(1, args.epochs + 1): print("====epoch " + str(epoch)) train_loss, train_acc = train(args, model_list, loader, optimizer_list, device) print("loss :", train_loss, "accuracy :", train_acc) if not args.model_file == "": torch.save(model.state_dict(), args.model_file + ".pth")
def main(): # Training settings parser = argparse.ArgumentParser( description= "PyTorch implementation of pre-training of graph neural networks") parser.add_argument("--device", type=int, default=0, help="which gpu to use if any (default: 0)") 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("--lr", type=float, default=0.001, help="learning rate (default: 0.001)") parser.add_argument("--decay", type=float, default=0, help="weight decay (default: 0)") parser.add_argument( "--num_layer", type=int, default=5, help="number of GNN message passing layers (default: 5).", ) parser.add_argument("--csize", type=int, default=3, help="context size (default: 3).") parser.add_argument("--emb_dim", type=int, default=300, help="embedding dimensions (default: 300)") parser.add_argument("--dropout_ratio", type=float, default=0, help="dropout ratio (default: 0)") parser.add_argument( "--neg_samples", type=int, default=1, help="number of negative contexts per positive context (default: 1)", ) parser.add_argument( "--JK", type=str, default="last", help="how the node features are combined across layers." "last, sum, max or concat", ) parser.add_argument( "--context_pooling", type=str, default="mean", help="how the contexts are pooled (sum, mean, or max)", ) parser.add_argument("--mode", type=str, default="cbow", help="cbow or skipgram") parser.add_argument( "--dataset", type=str, default="contextPred/chem/dataset/zinc_standard_agent", help="root directory of dataset for pretraining", ) parser.add_argument("--output_model_file", type=str, default="", help="filename to output the model") parser.add_argument("--gnn_type", type=str, default="gin") parser.add_argument("--seed", type=int, default=0, help="Seed for splitting dataset.") parser.add_argument( "--num_workers", type=int, default=8, help="number of workers for dataset loading", ) args = parser.parse_args() torch.manual_seed(0) np.random.seed(0) device = (torch.device("cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu")) if torch.cuda.is_available(): torch.cuda.manual_seed_all(0) l1 = args.num_layer - 1 l2 = l1 + args.csize print(args.mode) print("num layer: %d l1: %d l2: %d" % (args.num_layer, l1, l2)) # set up dataset and transform function. dataset = MoleculeDataset( args.dataset, dataset=os.path.basename(args.dataset), transform=ExtractSubstructureContextPair(args.num_layer, l1, l2), ) loader = DataLoaderSubstructContext(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) # set up models, one for pre-training and one for context embeddings model_substruct = GNN( args.num_layer, args.emb_dim, JK=args.JK, drop_ratio=args.dropout_ratio, gnn_type=args.gnn_type, ).to(device) model_context = GNN( int(l2 - l1), args.emb_dim, JK=args.JK, drop_ratio=args.dropout_ratio, gnn_type=args.gnn_type, ).to(device) # set up optimizer for the two GNNs optimizer_substruct = optim.Adam(model_substruct.parameters(), lr=args.lr, weight_decay=args.decay) optimizer_context = optim.Adam(model_context.parameters(), lr=args.lr, weight_decay=args.decay) for epoch in range(1, args.epochs + 1): print("====epoch " + str(epoch)) train_loss, train_acc = train( args, model_substruct, model_context, loader, optimizer_substruct, optimizer_context, device, ) print(train_loss, train_acc) if not args.output_model_file == "": torch.save(model_substruct.state_dict(), args.output_model_file + ".pth")