scattering_run.py

import random
import time
from parser import get_parser

import numpy as np
import seaborn as sns
import torch
import torch.nn as nn
from sklearn.metrics import roc_auc_score
from torch_geometric.datasets import HeterophilousGraphDataset

from scattering_method import LogReg, scatter_transform_diffusion


def index_to_mask(index, size):
    mask = torch.zeros(size, dtype=torch.bool)
    mask[index] = 1
    return mask


parser = get_parser()
args = parser.parse_args()

random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)

if __name__ == "__main__":
    dataset = HeterophilousGraphDataset(
        root=f"./data/{args.dataset}", name=args.dataset
    )
    args.graph_size = dataset[0].x.size(0)
    args.input_dim = dataset.num_features
    args.nclass = dataset.num_classes

    args.device = torch.device(f"cuda:0" if torch.cuda.is_available() else "cpu")

    print(args)
    args = vars(args)

    dname = args["dataset"]

    data = dataset[0]

    feat = data.x
    num_nodes = feat.shape[0]

    label = data.y

    edge_index = data.edge_index

    start_time = time.time()
    features = scatter_transform_diffusion(
        feat,
        edge_index,
        num_nodes,
        depth=args["depth"],
        K=args["K"],
        mode=args["mode"],
        pruning=args["prune"],
        th=args["th"],
    )
    end_time = time.time()
    print(f"Total time taken {end_time-start_time}")
    print(features.shape)

    results = []

    for i in range(10):
        train_mask, val_mask, test_mask = (
            data.train_mask[:, i],
            data.val_mask[:, i],
            data.test_mask[:, i],
        )
        train_mask, val_mask, test_mask = (
            train_mask.to(args["device"]),
            val_mask.to(args["device"]),
            test_mask.to(args["device"]),
        )
        features = features.to(args["device"])
        label = label.to(args["device"])

        train_embs = features[test_mask]
        val_embs = features[val_mask]
        test_embs = features[train_mask]

        train_labels = label[test_mask]
        val_labels = label[val_mask]
        test_labels = label[train_mask]

        best_val_acc = 0
        eval_acc = 0
        bad_counter = 0

        n_classes = (
            1
            if args["dataset"] in ["Minesweeper", "Tolokers", "Questions"]
            else dataset.num_classes
        )

        logreg = LogReg(features.shape[1], n_classes)
        opt = torch.optim.Adam(
            logreg.parameters(), lr=args["lr"], weight_decay=args["wd"]
        )
        logreg = logreg.to(args["device"])

        loss_fn = (
            nn.BCEWithLogitsLoss()
            if args["dataset"] in ["Minesweeper", "Tolokers", "Questions"]
            else nn.CrossEntropyLoss()
        )

        for epoch in range(args["epochs"]):
            logreg.train()
            opt.zero_grad()
            logits = logreg(train_embs)
            logits = (
                logits.squeeze(-1)
                if args["dataset"] in ["Minesweeper", "Tolokers", "Questions"]
                else logits
            )
            loss = (
                loss_fn(logits, train_labels.float())
                if args["dataset"] in ["Minesweeper", "Tolokers", "Questions"]
                else loss_fn(logits, train_labels)
            )
            loss.backward()
            opt.step()

            logreg.eval()
            with torch.no_grad():
                val_logits = logreg(val_embs)
                test_logits = logreg(test_embs)

                if args["dataset"] in ["Minesweeper", "Tolokers", "Questions"]:
                    val_acc = roc_auc_score(
                        y_true=val_labels.cpu().numpy(),
                        y_score=val_logits.squeeze(-1).cpu().numpy(),
                    )
                    test_acc = roc_auc_score(
                        y_true=test_labels.cpu().numpy(),
                        y_score=test_logits.squeeze(-1).cpu().numpy(),
                    )

                else:
                    val_preds = torch.argmax(val_logits, dim=1)
                    test_preds = torch.argmax(test_logits, dim=1)

                    val_acc = (
                        torch.sum(val_preds == val_labels).float() / val_labels.shape[0]
                    )
                    test_acc = (
                        torch.sum(test_preds == test_labels).float()
                        / test_labels.shape[0]
                    )

                if val_acc >= best_val_acc:
                    bad_counter = 0
                    best_val_acc = val_acc
                    if test_acc > eval_acc:
                        eval_acc = test_acc
                else:
                    bad_counter += 1

        print("Linear evaluation acc: {:.4f}".format(eval_acc))
        if torch.is_tensor(eval_acc):
            results.append(eval_acc.item())
        else:
            results.append(float(eval_acc))

    test_acc_mean = np.mean(results, axis=0) * 100
    values = np.asarray(results, dtype=object)
    uncertainty = np.max(
        np.abs(
            sns.utils.ci(
                sns.algorithms.bootstrap(values, func=np.mean, n_boot=1000), 95
            )
            - values.mean()
        )
    )
    print(f"test acc mean = {test_acc_mean:.4f} ± {uncertainty * 100:.4f}")