main.py

import logging
import numpy as np
import torch
import math
import os
from collections import defaultdict
from config import load_config
from dhg import Hypergraph
from data import load_dataset, rand_train_test_idx, permute_edges
from models import cot_numpy, PreModel
from utils import TBLogger, set_seed
from train import create_optimizer, pretrain, node_classification_evaluation

logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO)

if __name__ == "__main__":

    cfg = load_config()
    # ------- parameters -------
    seed = cfg['seed']
    train_rate = cfg['train_rate']
    valid_rate = cfg['valid_rate']
    loss_fn = cfg['loss_fn']
    replace_rate = cfg['replace_rate']
    mask_rate = cfg['mask_rate']
    num_hidden = cfg['num_hidden']
    lr = cfg['lr']
    max_epoch = cfg['max_epoch']
    weight_decay = cfg['weight_decay']
    lr_f = cfg['lr_f']
    max_epoch_f = cfg['max_epoch_f']
    weight_decay_f = cfg['weight_decay_f']
    encoder_type = cfg['encoder_type']
    decoder_type = cfg['decoder_type']
    dropout = cfg['dropout']
    optim_type = cfg['optim_type']
    runs = cfg['runs']
    cl = cfg['cl']
    attr = cfg['attr']
    lamda = cfg['lamda']
    aug_ratio = cfg['aug_ratio']
    data_name = cfg['data_name']
    result_file = os.path.join(cfg['result_dir'], "parameter_and_result.txt")
    acc_list = []
    estp_acc_list = []

    file = open(result_file, "a")
    set_seed(seed)
    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
    print(f"CUDA available: {torch.cuda.is_available()}")

    all_dataset = [
        'Cora', 'Citeseer', 'CA-Cora', 'CC-Cora', 'CC-Citeseer',
        'DBLP-Paper', 'DBLP-Conf', 'DBLP-Term', 'IMDB-Actor', 'IMDB-Director'
    ]

    # ------- load data -------
    print(f"Loading dataset:{data_name}")
    data = load_dataset(dataset=data_name)
    data = data.to(device)
    labels = data.y
    num_classes = len(np.unique(labels.cpu().numpy()))
    num_nodes = data.x.shape[0]
    feature_dim = data.x.shape[1]

    # ------- data split -------
    split_idx_lst = []
    for run in range(runs):
        split_idx = rand_train_test_idx(
            data.y, train_prop=train_rate, valid_prop=valid_rate)
        split_idx_lst.append(split_idx)
    he_index = defaultdict(list)
    edge_index = data.edge_index.cpu().numpy()
    for i, he in enumerate(edge_index[1]):
        he_index[he].append(edge_index[0][i])
    HG = Hypergraph(num_nodes, list(he_index.values())).to(device)

    # ------- data augmentation -------
    sorted_index1 = permute_edges(HG.H.coalesce().indices(), aug_ratio)
    HG_aug1 = Hypergraph(num_nodes, list(sorted_index1.values())).to(device)
    sorted_index2 = permute_edges(HG.H.coalesce().indices(), aug_ratio)
    HG_aug2 = Hypergraph(num_nodes, list(sorted_index2.values())).to(device)

    op_sim = []
    for i, e in enumerate(HG_aug1.e[0]):
        E1 = np.zeros((len(e), 1))
        for j, v in enumerate(e):
            E1[j] = HG_aug1.deg_v[v]
        E2 = np.zeros((len(HG_aug2.e[0][i]), 1))
        for j, v in enumerate(HG_aug2.e[0][i]):
            E2[j] = HG_aug2.deg_v[v]
        _, _, cost = cot_numpy(E1, E2)
        op_sim.append(math.exp(-lamda*cost)+1e-15)
    topo_sim = torch.tensor(op_sim).to(device)

    # ------- pretraining -------
    for run in range(runs):

        print(f"############################ Run {run} for seed {seed} #####################")
        split_idx = split_idx_lst[run]
        train_mask = split_idx['train'].to(device)
        val_mask = split_idx['valid'].to(device)
        test_mask = split_idx['test'].to(device)
        logger = TBLogger(log_path=cfg['log_dir'],
                          name=f"{data_name}_loss_{loss_fn}_rpr_{replace_rate}_nh_\
                          {num_hidden}_lr_{lr}_mp_{max_epoch}_mpf_{max_epoch_f}_wd_\
                          {weight_decay}_wdf_{weight_decay_f}_{encoder_type}_{decoder_type}")
        
        model = PreModel(
            in_dim=feature_dim,
            hid_dim=num_hidden,
            edge_dim=len(HG.e[0]),
            feat_drop=dropout,
            use_bn=True,
            mask_rate=mask_rate,
            encoder_type=encoder_type,
            decoder_type=decoder_type,
            loss_fn=loss_fn,
            replace_rate=replace_rate,
        )
        model.to(device)
        optimizer = create_optimizer(optim_type, model, lr, weight_decay)
        scheduler = None

        model = pretrain(model, data.x, HG, HG_aug1, HG_aug2,
                         optimizer, max_epoch, device, topo_sim, cl, attr, use_sim=True)
        
        # ------- linear evaluation -------
        model = model.to(device)
        X, lbl = data.x.to(device), labels.to(device)
        HG = HG.to(X.device)
        final_acc, estp_acc = node_classification_evaluation(run, model, X, HG, num_classes,
                                                             lr_f, weight_decay_f, max_epoch_f, device,
                                                             train_mask, val_mask, test_mask, lbl, linear_prob=True)
        acc_list.append(100*final_acc)
        estp_acc_list.append(100*estp_acc)

        if logger is not None:
            logger.finish()

    final_acc, final_acc_std = np.mean(acc_list), np.std(acc_list)
    estp_acc, estp_acc_std = np.mean(estp_acc_list), np.std(estp_acc_list)
    print(f"# final_acc: {final_acc:.2f} ± {final_acc_std:.2f}")
    print(f"# early-stopping_acc: {estp_acc:.2f} ± {estp_acc_std:.2f}")
    line = ("dataset = " + data_name + ", cl = " + str(cl) + ", attr = " + str(attr) +
            ", acc = " + "{:.2f}".format(estp_acc) + " ± " + "{:.2f}".format(estp_acc_std) + "\n")
    print(line)
    file.write(line)
    file.close()