train.py

import os
import time
import logging
import hydra
from omegaconf import DictConfig, OmegaConf
import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import dhg
from dhg.models import HGNNP,HGNN,GCN,UniGCN
from dhg import Graph, Hypergraph
from dhg.random import set_seed
from dhg.utils import split_by_num
from copy import deepcopy
from sklearn.metrics import accuracy_score,f1_score,precision_score
import numpy as np

from config.config import get_config
from datasets.data_helper import load_data
from utils.hypergraph_utils import fix_iso_v
from utils.hyperattack import RandomAttack_hyperedges, load_pickle

from models.hgnn import HGNN as myHGNN
from models.shield import HGNNShield
from models.gcn import GCN
from models.mlp import MLP
from models.gnnrobust import GCNRobust
from models.rgcn import RGCN
from models.elasso import EstimatedH,proximal_op


log = logging.getLogger(__name__)


def get_training_data(cfg):
    seed = cfg.train.seed
    mode = cfg.attack.mode
    rate = cfg.attack.rate
    root = cfg.data.root
    dataset = cfg.data.dataset
    device = cfg.data.device
    num_train = cfg.data.num_train
    num_val = cfg.data.num_val

    set_seed(seed)

    data, edge_list = load_data(cfg)
    X = data['features'].to(device)
    if rate > 0:
        try:
            content = load_pickle(dataset,rate,mode,root)
        except FileNotFoundError:
            print(f"[Data] Cache not found. Generating attacked data...")
            content = RandomAttack_hyperedges(cfg,dataset)

        if mode =='noise':
            X=content.to(device)
        else:
            edge_list=content
    labels = data['labels'].to(device)
    num_classes = data['num_classes']
    train_mask, val_mask, test_mask = split_by_num(X.shape[0],labels,num_train,num_val)

    return edge_list, X, labels, train_mask, val_mask, test_mask, num_classes

def get_model(cfg, dim_features, num_classes, num_vertices=None):
    name = cfg.model.name
    device = cfg.data.device
    hid = cfg.model.hidden_dim
    dropout = cfg.model.dropout

    if name in ['hgnn_shield']:
        backbone = cfg.model.backbone
        if backbone == 'hgnn':
            model = HGNN(dim_features,hid,num_classes,drop_rate=dropout)
        elif backbone == 'hgnnp':
            model = HGNNP(dim_features,hid,num_classes,drop_rate=dropout)
        elif backbone == 'unigcn':
            model = UniGCN(dim_features,hid,num_classes,drop_rate=dropout)
        else:
            raise NotImplementedError(f"Backbone {backbone} not supported.")
        
        model = HGNNShield(device,model)

    elif name in ['hgnn']:
        model = HGNN(dim_features,hid,num_classes,drop_rate=dropout)
        model =HGNNShield(device,model)

    elif name in ['hgnnp']:
        model = HGNNP(dim_features,hid,num_classes,drop_rate=dropout)
        model =HGNNShield(device,model)

    elif name in ['mlp']:
        model = MLP(dim_features,hid,num_classes,dropout)
        model =HGNNShield(device,model)

    elif name in ['unignn']:
        model = UniGCN(dim_features,hid,num_classes,drop_rate=dropout)
        model =HGNNShield(device,model)

    elif name in ['gcn','gcnsvd','gcnj']:
        model=GCN(dim_features,hid,num_classes,dropout)

        model=GCNRobust(model,device)

    elif name in ['rgcn']:
        model=RGCN(nnodes=num_vertices,in_channels=dim_features,hidden_channels=hid,
                   num_classes=num_classes,dropout=dropout,device=device)
        
        model=model.to(device)

    elif name in ['elasso']:
        model=myHGNN(dim_features, hid, num_classes, dropout)
        model=model.to(device)
    else:
        raise NotImplementedError
    
    return model

@torch.no_grad()
def evaluate(model, features, G, labels, mask, num_classes):
    model.eval()
    if isinstance(model, RGCN):
        output = model(features)
    else:
        output = model(features, G)

    pred = output[mask].max(1)[1].cpu().numpy()
    true = labels[mask].cpu().numpy()
    acc = accuracy_score(true, pred)
    f1 = f1_score(true, pred, average='weighted', labels=np.arange(num_classes),zero_division=0)
    pre = precision_score(true, pred, average='weighted', labels=np.arange(num_classes),zero_division=0)
    return acc, f1, pre,true,pred

def train(model,features,hg,labels,idx_train,idx_val,idx_test,cfg,num_classes):
    optimizer = optim.Adam(model.parameters(), lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
    best_val_acc = 0.0
    best_state = None

    t_start= time.perf_counter()

    for epoch in range(cfg.train.epochs):
        model.train()
        optimizer.zero_grad()

        output = model(features, hg)
        if isinstance(model.model, HGNN):
            output = F.log_softmax(output, dim=1)

        loss = F.nll_loss(output[idx_train], labels[idx_train])
        loss.backward()
        optimizer.step()

        val_acc,_,_,_,_=evaluate(model,features,hg,labels,idx_val, num_classes)

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_state = deepcopy(model.state_dict())

        print(f"Epoch {epoch:03d} | Loss: {loss:.4f} | Val: {val_acc:.4f}")

    t_train = time.perf_counter() - t_start

    if best_state is not None:
        model.load_state_dict(best_state)

    t_infer_start = time.perf_counter()
    acc,f1,pre,y_true,y_pred = evaluate(model,features,hg,labels,idx_test, num_classes)
    t_infer = time.perf_counter() - t_infer_start

    return {'acc': acc, 'f1': f1, 'pre': pre, 'train_time': t_train, 'infer_time': t_infer, 'y_true': y_true, 'y_pred': y_pred}
    
def train_rgcn(model,features,adj,labels,idx_train,idx_val,idx_test,cfg,num_classes):
    A = adj + torch.eye(model.nnodes).to(model.device)
    D=A.sum(1)
    D_inv_sqrt = torch.pow(D, -0.5)
    D_inv_sqrt[torch.isinf(D_inv_sqrt)] = 0.0
    D_inv_sqrt = torch.diag(D_inv_sqrt)

    D_inv = torch.pow(D, -1)
    D_inv[torch.isinf(D_inv)] = 0.0
    D_inv = torch.diag(D_inv)

    model.adj_norm1 = D_inv_sqrt @ A @ D_inv_sqrt
    model.adj_norm2 = D_inv @ A @ D_inv

    optimizer = optim.Adam(model.parameters(),lr=cfg.train.lr,weight_decay=cfg.train.weight_decay)
    beta1,beta2=5e-4,5e-4

    best_val_acc = 0.0
    best_output = None

    t_start= time.perf_counter()

    for epoch in range(cfg.train.epochs):
        model.train()
        optimizer.zero_grad()
        output = model(features)

        loss_cls = F.nll_loss(output[idx_train], labels[idx_train])
        miu1,sigma1=model.gc1.miu,model.gc1.sigma
        kl_loss = 0.5*(miu1.pow(2)+sigma1-torch.log(1e-8+sigma1)).mean(1).sum()
        norm2=torch.norm(model.gc1.weight_miu,2).pow(2)+torch.norm(model.gc1.weight_sigma,2).pow(2)
        loss = loss_cls+beta1*kl_loss+beta2*norm2

        loss.backward()
        optimizer.step()
        val_acc,_,_,_,_=evaluate(model,features,None,labels,idx_val, num_classes)
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_output = output.detach()

        print(f"Epoch {epoch:03d} (RGCN) | Loss: {loss:.4f} | Val: {val_acc:.4f}")

    t_train = time.perf_counter() - t_start

    t_infer_start = time.perf_counter()
    pred = torch.argmax(best_output, dim=1)
    y_test=labels[idx_test].detach().cpu().numpy()
    p_test=pred[idx_test].detach().cpu().numpy()
    acc = accuracy_score(y_test, p_test)
    f1 = f1_score(y_true=y_test,y_pred=p_test,
                  average='weighted', labels=np.arange(num_classes))
    pre = precision_score(y_true=y_test,y_pred=p_test,
                          average='weighted', labels=np.arange(num_classes))
    t_infer = time.perf_counter() - t_infer_start

    return {'acc': acc, 'f1': f1, 'pre': pre, 'train_time': t_train, 'infer_time': t_infer, 'y_true': y_test, 'y_pred': p_test}

def train_elasso(model, features, H_init, labels, idx_train, idx_val, idx_test, cfg,num_classes):
    estimator = EstimatedH(H_init).to(cfg.data.device)

    optimizer_model = optim.Adam(model.parameters(), lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
    optimizer_H = optim.SGD(estimator.parameters(), momentum=0.9,lr=5e-4)

    best_val_acc = 0.
    best_graph = None
    best_state = None

    outer_steps,inner_steps,beta=0,2,0.00005
    t_start = time.perf_counter()

    for epoch in range(cfg.train.epochs):
        for _ in range(outer_steps):
            estimator.train()
            optimizer_H.zero_grad()
            H_hat = estimator.normalized()
            output= model(features, H_hat)
            loss_H = F.nll_loss(output[idx_train], labels[idx_train])
            loss_H.backward()
            optimizer_H.step()
            with torch.no_grad():
                new_H = proximal_op(H_init,estimator.estimated_H,beta)
                estimator.estimated_H.data.copy_(new_H)

            normalized_H = estimator.normalized()

            val_acc,_,_,_,_=evaluate(model,features,normalized_H,labels,idx_val, num_classes)
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                best_graph = normalized_H.detach()
                best_state = deepcopy(model.state_dict())

        for _ in range(inner_steps):
            model.train()
            optimizer_model.zero_grad()
            H_hat = estimator.normalized()
            output = model(features, H_hat)
            loss = F.nll_loss(output[idx_train], labels[idx_train])
            loss.backward()
            optimizer_model.step()

            val_acc,_,_,_,_=evaluate(model,features,H_hat,labels,idx_val, num_classes)
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                best_graph = H_hat.detach()
                best_state = deepcopy(model.state_dict())
                
        print(f"Epoch {epoch:03d} (Elasso) | Loss: {loss:.4f} | Val: {val_acc:.4f}")

    t_train = time.perf_counter() - t_start

    if best_state is not None:
        model.load_state_dict(best_state)

    t_infer_start = time.perf_counter()
    acc,f1,pre,y_true,y_pred = evaluate(model,features,best_graph,labels,idx_test,num_classes)
    t_infer = time.perf_counter() - t_infer_start

    return {'acc': acc, 'f1': f1, 'pre': pre, 'train_time': t_train, 'infer_time': t_infer, 'y_true': y_true, 'y_pred': y_pred}

def run_experiment(cfg):
    device = torch.device(cfg.data.device)

    edge_list, X, labels, idx_train, idx_val, idx_test, num_classes = get_training_data(cfg)
    num_vertices = X.shape[0]

    G = Hypergraph(num_vertices,edge_list,device=device)
    G = fix_iso_v(G)

    model_name = cfg.model.name

    model = get_model(cfg, X.shape[1], num_classes, num_vertices)

    if model_name == 'elasso':
        H_init = G.H.to_dense()
        return train_elasso(model, X, H_init, labels, idx_train, idx_val, idx_test, cfg,num_classes)
    elif isinstance(model,GCNRobust):
        g=Graph.from_hypergraph_clique(G).to(device)
        if model_name == 'gcn':
            model.prepare_gcn(g)
        elif model_name == 'gcnj':
            model.prepare_jaccard(X,g,threshold=0.02)
        elif model_name == 'gcnsvd':
            model.prepare_svd(X,g)
        return train(model,X,None,labels,idx_train,idx_val,idx_test,cfg,num_classes)
    elif isinstance(model,RGCN):
        g=Graph.from_hypergraph_clique(G).to(device)
        return train_rgcn(model,X,g.A.to_dense(),labels,idx_train,idx_val,idx_test,cfg,num_classes)
    elif isinstance(model,HGNNShield):
        if model_name in ['hgnn_shield']:
            G = model.purify_structure(X, G, cfg)
            G = fix_iso_v(G)
        
        return train(model,X,G,labels,idx_train,idx_val,idx_test,cfg,num_classes)
    else:
        raise NotImplementedError(f"Training strategy for {model_name} not found.")
    
@hydra.main(config_path="config", config_name="config", version_base="1.2")
def main(cfg: DictConfig):
    log.info(OmegaConf.to_yaml(cfg))
    res = run_experiment(cfg)
    log.info(f"result: {res}")

if __name__ == '__main__':
    main()