src/2_lrmc_bilevel.py

# Bi-level Node↔Cluster message passing with fixed LRMC seeds

import argparse, json, os
from pathlib import Path
from typing import Dict, List, Tuple, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

from torch import Tensor
from torch_scatter import scatter_add, scatter_mean
from torch_sparse import coalesce, spspmm

from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
from torch_geometric.datasets import Planetoid, TUDataset
from torch_geometric.nn import GCNConv, global_mean_pool

from rich import print


# ---------------------------
# Utilities: edges and seeds
# ---------------------------

def add_scaled_self_loops(edge_index: Tensor,
                          edge_weight: Optional[Tensor],
                          num_nodes: int,
                          scale: float = 1.0) -> Tuple[Tensor, Tensor]:
    """Add self-loops with a chosen weight (scale). If scale=0, return unchanged."""
    if scale == 0.0:
        if edge_weight is None:
            edge_weight = torch.ones(edge_index.size(1), device=edge_index.device)
        return edge_index, edge_weight
    device = edge_index.device
    self_loops = torch.arange(num_nodes, device=device)
    self_index = torch.stack([self_loops, self_loops], dim=0)
    self_weight = torch.full((num_nodes,), float(scale), device=device)
    if edge_weight is None:
        base_w = torch.ones(edge_index.size(1), device=device)
    else:
        base_w = edge_weight
    ei = torch.cat([edge_index, self_index], dim=1)
    ew = torch.cat([base_w, self_weight], dim=0)
    # coalesce to sum possible duplicates
    ei, ew = coalesce(ei, ew, num_nodes, num_nodes, op='add')
    return ei, ew


def adjacency_power(edge_index: Tensor, num_nodes: int, k: int = 2) -> Tensor:
    """
    Compute (binary) k-th power adjacency using sparse matmul (torch_sparse.spspmm).
    Returns a coalesced edge_index (no weights, duplicates removed).
    """
    # Build A as (row, col) with all weights 1
    device = edge_index.device
    row, col = edge_index
    val = torch.ones(row.numel(), device=device)
    Ai, Av = edge_index, val
    # Repeatedly multiply: A^2, then chain if k>2
    Ri, Rv = spspmm(Ai, Av, Ai, Av, num_nodes, num_nodes, num_nodes)
    # Remove diagonal self-loops in pure power (we can add our own later)
    mask = Ri[0] != Ri[1]
    Ri = Ri[:, mask]
    # (Optional) higher powers: (A^k) – here we keep exactly k=2 for simplicity
    return coalesce(Ri, torch.ones(Ri.size(1), device=device), num_nodes, num_nodes)[0]


def build_cluster_graph(edge_index: Tensor,
                        num_nodes: int,
                        node2cluster: Tensor,
                        weight_per_edge: Optional[Tensor] = None,
                        num_clusters: Optional[int] = None
                        ) -> Tuple[Tensor, Tensor, int]:
    """
    Build cluster graph A_c = S^T A S.
    node2cluster: [N] long tensor with the cluster id for each node (hard assignment).
    Returns (edge_index_c, edge_weight_c, K).
    """
    if num_clusters is None:
        K = int(node2cluster.max().item()) + 1
    else:
        K = num_clusters

    src, dst = edge_index
    csrc = node2cluster[src]
    cdst = node2cluster[dst]
    edge_c = torch.stack([csrc, cdst], dim=0)
    if weight_per_edge is None:
        w = torch.ones(edge_c.size(1), device=edge_c.device)
    else:
        w = weight_per_edge
    edge_c, w = coalesce(edge_c, w, K, K, op='add')  # sum multiplicities

    # # set all weights of cluster edges to 1
    # w = torch.ones_like(w)
    # mask = edge_c[0] != edge_c[1]
    # edge_c, w = edge_c[:, mask], w[mask]

    return edge_c, w, K


# -----
# Seeds
# -----

def _parse_clusters_single_file(obj: dict, n_nodes: int) -> Tuple[List[List[int]], Tensor]:
    """
    Expect the JSON to have top-level "clusters": [{members:[...], score:...}, ...]
    Unassigned nodes become singleton clusters.
    If a node appears in multiple clusters, we keep the cluster with largest 'score' then by size.
    """
    clusters = obj.get("clusters", [])
    # Collect candidate cluster id per node with priority by cluster score and size:
    per_node = {}  # node_id -> (priority_tuple, cluster_idx)
    out: List[List[int]] = []
    # prepare list of (members, score, size, index)
    cinfo = []
    for idx, c in enumerate(clusters):
        members = c.get("members", [])
        score = float(c.get("score", 0.0))
        cinfo.append((members, score, len(members), idx))

    # make a stable cluster list first
    for members, score, size, idx in cinfo:
        out.append(list(members))

    # assign best cluster per node
    chosen = torch.full((n_nodes,), -1, dtype=torch.long)
    best_key = [(-1e18, -10) for _ in range(n_nodes)]
    for c_idx, (members, score, size, _) in enumerate(cinfo):
        key = (score, size)
        for u in members:
            old = best_key[u]
            if key > old:  # prioritize larger score, then larger cluster
                best_key[u] = key
                chosen[u] = c_idx

    # any unassigned node becomes its own new cluster
    next_c = len(out)
    for u in range(n_nodes):
        if chosen[u] == -1:
            out.append([u])
            chosen[u] = next_c
            next_c += 1

    # Build cluster_scores vector aligned with `out` order, then normalize to [0,1]
    base_scores = [float(s) for (_, s, _, _) in cinfo]
    K = len(out)
    scores = torch.zeros(K, dtype=torch.float32)
    # Fill provided cluster scores first
    for i, sc in enumerate(base_scores):
        scores[i] = sc
    # Singletons (appended) remain 0 by default
    if len(base_scores) > 0:
        smin = min(base_scores)
        smax = max(base_scores)
        if smax > smin:
            # Min-max normalize provided cluster scores; keep singletons at 0
            norm = (scores[:len(base_scores)] - smin) / (smax - smin)
            scores[:len(base_scores)] = norm
        else:
            # All equal: treat as confident -> set to 1 for provided clusters
            scores[:len(base_scores)] = 1.0
    # Shape as (K,1)
    cluster_scores = scores.view(-1, 1)

    # Return clusters and their normalized scores
    return out, cluster_scores


def seeds_to_node2cluster(n_nodes: int, clusters: List[List[int]]) -> Tensor:
    node2cluster = torch.full((n_nodes,), -1, dtype=torch.long)
    for cid, members in enumerate(clusters):
        for u in members:
            node2cluster[u] = cid
    assert int(node2cluster.min()) >= 0, "All nodes must be assigned a cluster."
    return node2cluster


def load_lrmc_seeds_single_graph(seeds_json: str, n_nodes: int) -> Tuple[Tensor, Tensor]:
    """Load seeds for a single big graph (Planetoid)."""
    with open(seeds_json, "r") as f:
        obj = json.load(f)
    clusters, cluster_scores = _parse_clusters_single_file(obj, n_nodes)
    node2cluster = seeds_to_node2cluster(n_nodes, clusters)
    return node2cluster, cluster_scores


# --------------------------
# Bi-level LRMC layer (1x)
# --------------------------

class BiLevelLRMC(nn.Module):
    """
    One round:
      1) Node GCN: H1 = GCN_node(X, A_node)
      2) Up: Z = mean_{i in c} H1[i]         (cluster means via scatter)
         Cluster graph: A_c = S^T A_node S
      3) Cluster GCN: Z2 = GCN_cluster(Z, A_c)
      4) Down: H2 = H1 + W (S Z2)
    """

    def __init__(self,
                 in_dim: int,
                 hidden_dim: int,
                 node2cluster: Tensor,
                 cluster_scores: Tensor,
                 edge_index_node: Tensor,
                 num_nodes: int,
                 self_loop_scale: float = 0.0,
                 use_a2: bool = False):
        super().__init__()
        self.num_nodes = num_nodes
        self.node2cluster = node2cluster.clone().long()
        self.register_buffer("node2cluster_buf", self.node2cluster)
        # cluster_scores: (K,1) in [0,1]
        self.register_buffer("cluster_scores", cluster_scores.clone().float())
        # # Node graph (optionally with A^2 and/or scaled self-loops)
        # ei = edge_index_node
        # if use_a2:
        #     ei = adjacency_power(ei, num_nodes, k=2)
        # ei, ew = add_scaled_self_loops(ei, None, num_nodes, scale=self_loop_scale)
        # self.register_buffer("edge_index_node", ei)
        # self.register_buffer("edge_weight_node", ew)

        # 1) Node graph: keep raw A (no A^2), but use A+2I by default
        ei_node = edge_index_node
        ei_node, ew_node = add_scaled_self_loops(ei_node, None, num_nodes, scale=self_loop_scale)
        self.register_buffer("edge_index_node", ei_node)
        self.register_buffer("edge_weight_node", ew_node)

        # 2) Cluster graph: build from A^2 to keep coarsened graph well connected
        ei_base_for_clusters = edge_index_node
        if use_a2:
            ei_base_for_clusters = adjacency_power(edge_index_node, num_nodes, k=2)

        edge_index_c, edge_weight_c, K = build_cluster_graph(
            ei_base_for_clusters, num_nodes, self.node2cluster
        )
        self.register_buffer("edge_index_c", edge_index_c)
        self.register_buffer("edge_weight_c", edge_weight_c)
        self.num_clusters = K

        # GCNs
        self.gcn_node = GCNConv(in_dim, hidden_dim, add_self_loops=False, normalize=True)
        # self.gcn_cluster = GCNConv(hidden_dim, hidden_dim, add_self_loops=True, normalize=True)
        # self.down = nn.Linear(hidden_dim, hidden_dim)
        # self.gate = nn.Sequential(
        #     nn.Linear(2 * hidden_dim, hidden_dim // 2),
        #     nn.ReLU(),
        #     nn.Linear(hidden_dim // 2, 1)
        # )
        # self.lambda_logit = nn.Parameter(torch.tensor(0.0))

    def forward(self, x: Tensor) -> Tensor:
        # Node GCN
        h1 = self.gcn_node(x, self.edge_index_node, self.edge_weight_node)
        h1 = F.relu(h1)

        # # Up: cluster means
        # counts = torch.bincount(self.node2cluster_buf, minlength=self.num_clusters).clamp(min=1).unsqueeze(-1)
        # z = scatter_add(h1, self.node2cluster_buf, dim=0, dim_size=self.num_clusters) / counts

        # # Cluster GCN
        # z2 = self.gcn_cluster(z, self.edge_index_c, self.edge_weight_c)
        # z2 = F.relu(z2)

        # # Down: broadcast to nodes + residual, scaled by cluster_scores
        # z2_nodes = z2[self.node2cluster_buf]
        # inj = self.down(z2_nodes)
        # gate_in = torch.cat([h1, inj], dim=-1)
        # gate_dyn = torch.sigmoid(self.gate(gate_in))
        # alpha_seed = 0.25 + 0.75 * self.cluster_scores[self.node2cluster_buf]
        # # print(alpha_seed)
        # lam = torch.sigmoid(self.lambda_logit)
        # # print(lam)
        # alpha = lam * alpha_seed + (1 - lam) * gate_dyn
        # h2 = h1 + alpha * inj
        return h1

# -----------------------------------
# Node classification model (Planetoid)
# -----------------------------------

class NodeLRMCGCN(nn.Module):
    def __init__(self, in_dim: int, hidden: int, num_classes: int,
                 node2cluster: Tensor, cluster_scores: Tensor, edge_index: Tensor, num_nodes: int,
                 layers: int = 1, self_loop_scale: float = 0.0, use_a2: bool = False, dropout: float = 0.5):
        super().__init__()
        self.layer = BiLevelLRMC(in_dim, hidden, node2cluster, cluster_scores, edge_index, num_nodes, self_loop_scale, use_a2))
        self.cls = nn.Linear(hidden, num_classes)
        self.dropout = dropout

    def forward(self, x: Tensor) -> Tensor:
        h = x
        h = layer(h)
        h = F.dropout(h, p=self.dropout, training=self.training)
        out = self.cls(h)
        return out


# ---------------------------------------
# Graph classification with batching (TU)
# ---------------------------------------

class GraphLRMCProvider:
    """
    Holds per-graph LRMC assignments and cluster graphs.
    Expects a directory with one JSON per graph OR a single JSON with {"graphs":[{"graph_id":int,"clusters":[...]},...]}.
    Node indices are local per-graph [0..n_i-1].
    """

    def __init__(self, dataset, seeds_path: str, use_a2: bool = True):
        """
        dataset: any iterable/sequence of torch_geometric.data.Data
        """
        self.dataset = dataset
        self.root = Path(seeds_path)
        self.per_graph: Dict[int, Dict[str, Tensor]] = {}
        # Try single JSON with all graphs
        single_json = None
        if self.root.is_file() and self.root.suffix.lower() == ".json":
            single_json = json.loads(Path(self.root).read_text())
        for gid, data in enumerate(dataset):
            n = data.num_nodes
            if single_json is not None and "graphs" in single_json:
                # Structure: {"graphs":[{"graph_id":int,"clusters":[...]}]}
                entry = None
                for g in single_json["graphs"]:
                    if int(g.get("graph_id", -1)) == gid:
                        entry = g
                        break
                if entry is None:
                    # fallback: singleton clusters
                    node2cluster = torch.arange(n, dtype=torch.long)
                    cluster_scores = torch.ones(n, 1, dtype=torch.float32)  # singletons -> treat as 1
                else:
                    clusters, cluster_scores = _parse_clusters_single_file(entry, n)
                    node2cluster = seeds_to_node2cluster(n, clusters)
            else:
                # One JSON per graph e.g. seeds_dir/graph_000123.json
                guess = self.root / f"graph_{gid:06d}.json"
                if guess.exists():
                    obj = json.loads(guess.read_text())
                    clusters, cluster_scores = _parse_clusters_single_file(obj, n)
                    node2cluster = seeds_to_node2cluster(n, clusters)
                else:
                    node2cluster = torch.arange(n, dtype=torch.long)  # singleton fallback
                    cluster_scores = torch.ones(n, 1, dtype=torch.float32)

            ei = data.edge_index
            if use_a2:
                ei = adjacency_power(ei, n, k=2)
            ei_c, ew_c, K = build_cluster_graph(ei, n, node2cluster)
            self.per_graph[gid] = {
                "node2cluster": node2cluster,
                "cluster_scores": cluster_scores,
                "edge_index_c": ei_c,
                "edge_weight_c": ew_c,
                "num_clusters": torch.tensor([K]),
            }

    def get(self, graph_id: int):
        rec = self.per_graph[graph_id]
        return (rec["node2cluster"], rec["cluster_scores"], rec["edge_index_c"], rec["edge_weight_c"],
                int(rec["num_clusters"][0].item()))


class GraphLRMCGCN(nn.Module):
    """
    Batched version:
      - Run node-level GCN over batch graph (standard).
      - Up: per-graph scatter to cluster means; build a batched cluster-graph by offsetting cluster ids.
      - Cluster GCN over the batched cluster graph.
      - Down: broadcast cluster features back to nodes and residual.
      - Graph head: global mean pooling -> MLP.
    """

    def __init__(self, in_dim: int, hidden: int, num_classes: int,
                 self_loop_scale: float = 0.0, use_a2: bool = False, dropout: float = 0.5):
        super().__init__()
        self.gcn_node = GCNConv(in_dim, hidden, add_self_loops=False, normalize=True)
        self.gcn_cluster = GCNConv(hidden, hidden, add_self_loops=True, normalize=True)
        self.down = nn.Linear(hidden, hidden)
        # Classifier takes concatenated node and cluster embeddings (2 * hidden)
        self.cls = nn.Linear(2 * hidden, num_classes)
        self.self_loop_scale = self_loop_scale
        self.use_a2 = use_a2
        self.dropout = dropout
        self.gate = nn.Sequential(
            nn.Linear(2 * hidden, hidden // 2),
            nn.ReLU(),
            nn.Linear(hidden // 2, 1)
        )
        self.lambda_logit = nn.Parameter(torch.tensor(0.0))

    def forward(self, data: Data, provider: GraphLRMCProvider) -> Tensor:
        # Single-graph only: no batching.
        x, edge_index = data.x, data.edge_index
        num_nodes = x.size(0)

        # Node graph prep
        ei = edge_index
        if self.use_a2:
            ei = adjacency_power(ei, num_nodes, k=2)
        ei, ew = add_scaled_self_loops(ei, None, num_nodes, scale=self.self_loop_scale)

        # Node GCN
        h1 = self.gcn_node(x, ei, ew)
        h1 = F.relu(h1)

        # Fetch LRMC seeds/cluster-graph for this graph
        assert hasattr(data, 'gid'), "Each graph must carry a 'gid' attribute for provider lookup."
        gid = int(data.gid.view(-1)[0].item())
        node2cluster_g, cluster_scores_g, edge_index_c, edge_weight_c, K = provider.get(gid)
        node2cluster_g = node2cluster_g.to(x.device)
        edge_index_c = edge_index_c.to(x.device)
        edge_weight_c = edge_weight_c.to(x.device)
        cluster_scores_g = cluster_scores_g.to(x.device)

        # Up: cluster means
        counts = torch.bincount(node2cluster_g, minlength=K).clamp(min=1).unsqueeze(-1)
        z = scatter_add(h1, node2cluster_g, dim=0, dim_size=K) / counts

        # Cluster GCN
        z2 = self.gcn_cluster(z, edge_index_c, edge_weight_c)
        z2 = F.relu(z2)

        # Down: broadcast to nodes and residual
        z2_nodes = z2[node2cluster_g]
        inj = self.down(z2_nodes)
        gate_in = torch.cat([h1, inj], dim=-1)  # (N, 2H)
        gate_dyn = torch.sigmoid(self.gate(gate_in))  # (N, 1)
        # normalize cluster_scores to [0.25,1] so singletons still pass some signal
        alpha_seed = 0.25 + 0.75 * cluster_scores_g[node2cluster_g]
        lam = torch.sigmoid(self.lambda_logit)
        alpha = lam * alpha_seed + (1 - lam) * gate_dyn
        print(lam)
        h2 = h1 + alpha * inj

        # Graph head: simple mean over nodes
        h2 = F.dropout(h2, p=self.dropout, training=self.training)
        g_nodes = h2.mean(dim=0, keepdim=True)
        g_clust = z2.mean(dim=0, keepdim=True)
        g = torch.cat([g_nodes, g_clust], dim=-1)
        out = self.cls(g)
        return out


# -------------
# Training glue
# -------------

def train_node(task_ds: str, seeds_json: str, hidden=64, layers=1, epochs=300,
               lr=0.01, weight_decay=5e-4, dropout=0.5, self_loop_scale=0.0, use_a2=False, seed=0):
    torch.manual_seed(seed)
    ds = Planetoid(root=f"./data/{task_ds}", name=task_ds)
    data = ds[0]
    n, c_in, n_cls = data.num_nodes, ds.num_node_features, ds.num_classes

    node2cluster, cluster_scores = load_lrmc_seeds_single_graph(seeds_json, n)
    model = NodeLRMCGCN(c_in, hidden, n_cls, node2cluster, cluster_scores, data.edge_index, n,
                        layers=layers, self_loop_scale=self_loop_scale, use_a2=use_a2, dropout=dropout).to('cpu')
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

    def step():
        # Train step: compute loss with dropout on, then evaluate metrics with dropout off.
        model.train()
        opt.zero_grad(set_to_none=True)
        out_train = model(data.x)
        loss = F.cross_entropy(out_train[data.train_mask], data.y[data.train_mask])
        loss.backward()
        opt.step()

        # Evaluation pass in eval mode to report metrics without dropout.
        with torch.no_grad():
            model.eval()
            out_eval = model(data.x)

            def acc(mask):
                pred = out_eval[mask].argmax(dim=1)
                pred_t = torch.as_tensor(pred)
                y_t = torch.as_tensor(data.y)
                return (pred_t == y_t[mask]).float().mean().item()

            return loss.item(), acc(data.train_mask), acc(data.val_mask), acc(data.test_mask)

    best_val, best_test = 0.0, 0.0
    for ep in range(1, epochs + 1):
        loss, tr, va, te = step()
        if va > best_val:
            best_val, best_test = va, te
        if ep % 20 == 0:
            print(f"[{ep:04d}] loss={loss:.4f}  train={tr:.3f}  val={va:.3f}  test={te:.3f}  best_test={best_test:.3f}")
    print(f"Best val={best_val:.3f}  test@best={best_test:.3f}")


def train_graph(dataset_name: str, seeds_path: str, hidden=64, epochs=100,
                lr=0.001, weight_decay=1e-4, dropout=0.5, self_loop_scale=0.0, use_a2=False, seed=0):
    torch.manual_seed(seed)
    ds = TUDataset(root=f"./data/{dataset_name}", name=dataset_name)
    num_classes = ds.num_classes
    c_in = ds.num_node_features if ds.num_node_features > 0 else 1

    # Materialize dataset into a list of Data objects to make mutations persistent.
    graphs: List[Data] = []
    for i, g in enumerate(ds):
        gc = g.clone()
        # Attach persistent global id for provider lookup across splits/batches
        gc.gid = torch.tensor([i], dtype=torch.long)
        graphs.append(gc)

    # If dataset has no node features, use degree as a 1-D feature for each graph.
    if ds.num_node_features == 0:
        for g in graphs:
            deg = torch.bincount(g.edge_index[0], minlength=g.num_nodes).float().view(-1, 1)
            g.x = deg

    provider = GraphLRMCProvider(graphs, seeds_path)
    idx = torch.randperm(len(graphs))
    ntrain = int(0.8 * len(ds))
    nval = int(0.1 * len(ds))
    # Build splits from the materialized list
    train_ds = [graphs[i] for i in idx[:ntrain]]
    val_ds = [graphs[i] for i in idx[ntrain:ntrain + nval]]
    test_ds = [graphs[i] for i in idx[ntrain + nval:]]

    device = 'cpu'
    model = GraphLRMCGCN(c_in, hidden, num_classes,
                         self_loop_scale=self_loop_scale, use_a2=use_a2, dropout=dropout).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

    @torch.no_grad()
    def evaluate(graph_list: List[Data]):
        model.eval()
        tot, correct = 0, 0
        for g in graph_list:
            g = g.to(device)
            logits = model(g, provider)
            pred = logits.argmax(dim=1)
            pred_t = torch.as_tensor(pred)
            y_t = torch.as_tensor(g.y)
            correct += (pred_t == y_t).sum().item()
            tot += g.y.size(0)
        return correct / tot

    best_val, best_test = 0.0, 0.0
    for ep in range(1, epochs + 1):
        model.train()
        for g in train_ds:
            g = g.to(device)
            opt.zero_grad(set_to_none=True)
            logits = model(g, provider)
            loss = F.cross_entropy(logits, g.y)
            loss.backward()
            opt.step()
        if ep % 5 == 0:
            va = evaluate(val_ds)
            te = evaluate(test_ds)
            if va > best_val:
                best_val, best_test = va, te
            print(f"[{ep:03d}] val={va:.3f} test={te:.3f} best_test@val={best_test:.3f}")
    print(f"Best val={best_val:.3f} test@best={best_test:.3f}")


# -----------
# Entrypoint
# -----------

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--task", choices=["node", "graph"], required=True)
    p.add_argument("--dataset", required=True, help="Cora/Citeseer/Pubmed or DD/PROTEINS/COLLAB/ENZYMES")
    p.add_argument("--seeds", required=True, help="Path to seeds JSON (node task) or dir/single JSON (graph task)")
    p.add_argument("--hidden", type=int, default=64)
    p.add_argument("--layers", type=int, default=1)
    p.add_argument("--epochs", type=int, default=300)
    p.add_argument("--batch_size", type=int, default=64)
    p.add_argument("--lr", type=float, default=0.01)
    p.add_argument("--wd", type=float, default=5e-4)
    p.add_argument("--dropout", type=float, default=0.5)
    p.add_argument("--self_loop_scale", type=float, default=0.0, help="use 2.0 to mimic A+2I")
    p.add_argument("--use_a2", action="store_true", help="use A^2 connectivity augmentation")
    p.add_argument("--seed", type=int, default=0)
    args = p.parse_args()

    if args.task == "node":
        for i in range(42, 60):
            train_node(args.dataset, args.seeds, hidden=args.hidden, layers=args.layers, epochs=args.epochs, lr=args.lr,
                       weight_decay=args.wd, dropout=args.dropout, self_loop_scale=args.self_loop_scale,
                       use_a2=args.use_a2, seed=i)
    else:
        for i in range(42, 60):
            train_graph(args.dataset, args.seeds, hidden=args.hidden, epochs=max(100, args.epochs),
                        lr=min(args.lr, 0.001), weight_decay=args.wd, dropout=args.dropout,
                        self_loop_scale=args.self_loop_scale, use_a2=args.use_a2, seed=i)


if __name__ == "__main__":
    main()