chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,168 @@
+import argparse
+
+import dgl
+
+import torch as th
+from dgl.data import GINDataset
+from dgl.dataloading import GraphDataLoader
+from evaluate_embedding import evaluate_embedding
+from model import InfoGraph
+
+
+def argument():
+    parser = argparse.ArgumentParser(description="InfoGraph")
+    # data source params
+    parser.add_argument(
+        "--dataname", type=str, default="MUTAG", help="Name of dataset."
+    )
+
+    # training params
+    parser.add_argument(
+        "--gpu", type=int, default=-1, help="GPU index, default:-1, using CPU."
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=20, help="Training epochs."
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=128, help="Training batch size."
+    )
+    parser.add_argument("--lr", type=float, default=0.01, help="Learning rate.")
+    parser.add_argument(
+        "--log_interval",
+        type=int,
+        default=1,
+        help="Interval between two evaluations.",
+    )
+
+    # model params
+    parser.add_argument(
+        "--n_layers",
+        type=int,
+        default=3,
+        help="Number of graph convolution layers before each pooling.",
+    )
+    parser.add_argument(
+        "--hid_dim", type=int, default=32, help="Hidden layer dimensionalities."
+    )
+
+    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
+
+
+def collate(samples):
+    """collate function for building graph dataloader"""
+
+    graphs, labels = map(list, zip(*samples))
+
+    # generate batched graphs and labels
+    batched_graph = dgl.batch(graphs)
+    batched_labels = th.tensor(labels)
+
+    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_labels
+
+
+if __name__ == "__main__":
+    # Step 1: Prepare graph data   ===================================== #
+    args = argument()
+    print(args)
+
+    # load dataset from dgl.data.GINDataset
+    dataset = GINDataset(args.dataname, self_loop=False)
+
+    # get graphs and labels
+    graphs, labels = map(list, zip(*dataset))
+
+    # generate a full-graph with all examples for evaluation
+    wholegraph = dgl.batch(graphs)
+    wholegraph.ndata["attr"] = wholegraph.ndata["attr"].to(th.float32)
+
+    # create dataloader for batch training
+    dataloader = GraphDataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        collate_fn=collate,
+        drop_last=False,
+        shuffle=True,
+    )
+
+    in_dim = wholegraph.ndata["attr"].shape[1]
+
+    # Step 2: Create model =================================================================== #
+    model = InfoGraph(in_dim, args.hid_dim, args.n_layers)
+    model = model.to(args.device)
+
+    # Step 3: Create training components ===================================================== #
+    optimizer = th.optim.Adam(model.parameters(), lr=args.lr)
+
+    print("===== Before training ======")
+
+    wholegraph = wholegraph.to(args.device)
+    wholefeat = wholegraph.ndata["attr"]
+
+    emb = model.get_embedding(wholegraph, wholefeat).cpu()
+    res = evaluate_embedding(emb, labels, args.device)
+
+    """ Evaluate the initialized embeddings """
+    """ using logistic regression and SVM(non-linear) """
+    print("logreg {:4f}, svc {:4f}".format(res[0], res[1]))
+
+    best_logreg = 0
+    best_logreg_epoch = 0
+    best_svc = 0
+    best_svc_epoch = 0
+
+    # Step 4: training epochs =============================================================== #
+    for epoch in range(args.epochs):
+        loss_all = 0
+        model.train()
+
+        for graph, label in dataloader:
+            graph = graph.to(args.device)
+            feat = graph.ndata["attr"]
+            graph_id = graph.ndata["graph_id"]
+
+            n_graph = label.shape[0]
+
+            optimizer.zero_grad()
+            loss = model(graph, feat, graph_id)
+            loss.backward()
+            optimizer.step()
+            loss_all += loss.item()
+
+        print("Epoch {}, Loss {:.4f}".format(epoch, loss_all))
+
+        if epoch % args.log_interval == 0:
+            # evaluate embeddings
+            model.eval()
+            emb = model.get_embedding(wholegraph, wholefeat).cpu()
+            res = evaluate_embedding(emb, labels, args.device)
+
+            if res[0] > best_logreg:
+                best_logreg = res[0]
+                best_logreg_epoch = epoch
+
+            if res[1] > best_svc:
+                best_svc = res[1]
+                best_svc_epoch = epoch
+
+            print(
+                "best logreg {:4f}, epoch {} | best svc: {:4f}, epoch {}".format(
+                    best_logreg, best_logreg_epoch, best_svc, best_svc_epoch
+                )
+            )
+
+    print("Training End")
+    print("best logreg {:4f} ,best svc {:4f}".format(best_logreg, best_svc))