main_full_batch.py

import logging
import numpy as np
from tqdm import tqdm
import torch
import os
from utils import (
    build_args,
    create_optimizer,
    set_random_seed,
    TBLogger,
    get_current_lr,
    load_best_configs,
)
from datasets.data_proc import load_small_dataset
from models.finetune import linear_probing_full_batch
from models import build_model
import sys
import datetime


class DualLogger(object):
    def __init__(self, filename):
        self.terminal = sys.stdout
        self.log = open(filename, "a")  # Use "a" to append across seeds

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)
        self.terminal.flush()
        self.log.flush()

    def flush(self):
        self.terminal.flush()
        self.log.flush()

    def close(self):
        self.log.close()


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


def pretrain(
    model,
    graph,
    feat,
    optimizer,
    max_epoch,
    device,
    scheduler,
    num_classes,
    lr_f,
    weight_decay_f,
    max_epoch_f,
    linear_prob,
    logger=None,
):
    logging.info("start training..")
    graph = graph.to(device)
    x = feat.to(device)

    target_nodes = torch.arange(x.shape[0], device=x.device, dtype=torch.long)
    epoch_iter = tqdm(range(max_epoch))

    for epoch in epoch_iter:
        model.train()

        loss = model(graph, x, targets=target_nodes)

        loss_dict = {"loss": loss.item()}
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if scheduler is not None:
            scheduler.step()

        epoch_iter.set_description(f"# Epoch {epoch}: train_loss: {loss.item():.4f}")
        if logger is not None:
            loss_dict["lr"] = get_current_lr(optimizer)
            logger.note(loss_dict, step=epoch)

        if (epoch + 1) % 200 == 0:
            linear_probing_full_batch(
                model,
                graph,
                x,
                num_classes,
                lr_f,
                weight_decay_f,
                max_epoch_f,
                device,
                linear_prob,
                mute=True,
            )

    return model


def main(args):
    log_dir = "./logs"
    os.makedirs(log_dir, exist_ok=True)

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    log_filename = os.path.join(log_dir, f"graphmae_{args.dataset}_{timestamp}.txt")

    # --- 2. START REDIRECTION ---
    logger_obj = DualLogger(log_filename)
    sys.stdout = logger_obj

    print("=" * 50)
    print(f"Job Started at {timestamp}")
    print(f"Log file: {log_filename}")
    print("=" * 50)
    device = args.device if args.device >= 0 else "cpu"
    seeds = args.seeds
    dataset_name = args.dataset
    max_epoch = args.max_epoch
    max_epoch_f = args.max_epoch_f
    num_hidden = args.num_hidden
    num_layers = args.num_layers
    encoder_type = args.encoder
    decoder_type = args.decoder
    replace_rate = args.replace_rate

    optim_type = args.optimizer
    loss_fn = args.loss_fn

    lr = args.lr
    weight_decay = args.weight_decay
    lr_f = args.lr_f
    weight_decay_f = args.weight_decay_f
    linear_prob = args.linear_prob
    load_model = args.load_model
    logs = args.logging
    use_scheduler = args.scheduler

    graph, (num_features, num_classes) = load_small_dataset(dataset_name)
    args.num_features = num_features

    acc_list = []
    estp_acc_list = []
    for i, seed in enumerate(seeds):
        print(f"####### Run {i} for seed {seed}")
        set_random_seed(seed)

        if logs:
            logger = TBLogger(
                name=f"{dataset_name}_loss_{loss_fn}_rpr_{replace_rate}_nh_{num_hidden}_nl_{num_layers}_lr_{lr}_mp_{max_epoch}_mpf_{max_epoch_f}_wd_{weight_decay}_wdf_{weight_decay_f}_{encoder_type}_{decoder_type}"
            )
        else:
            logger = None

        model = build_model(args)
        model.to(device)
        optimizer = create_optimizer(optim_type, model, lr, weight_decay)

        if use_scheduler:
            logging.info("Use schedular")
            scheduler = lambda epoch: (1 + np.cos((epoch) * np.pi / max_epoch)) * 0.5
            scheduler = torch.optim.lr_scheduler.LambdaLR(
                optimizer, lr_lambda=scheduler
            )
        else:
            scheduler = None

        x = graph.ndata["feat"]
        if not load_model:
            model = pretrain(
                model,
                graph,
                x,
                optimizer,
                max_epoch,
                device,
                scheduler,
                num_classes,
                lr_f,
                weight_decay_f,
                max_epoch_f,
                linear_prob,
                logger,
            )
            os.makedirs("./checkpoints", exist_ok=True)
            save_path = f"./checkpoints/{dataset_name}_seed{seed}.pt"
            torch.save(model.cpu().state_dict(), save_path)
            logging.info(f"Checkpoint saved to {save_path}")
            model = model.cpu()

        if load_model:
            logging.info("Loading Model ... ")
            model.load_state_dict(torch.load("checkpoint.pt"))

        model = model.to(device)
        model.eval()

        final_acc, estp_acc = linear_probing_full_batch(
            model,
            graph,
            x,
            num_classes,
            lr_f,
            weight_decay_f,
            max_epoch_f,
            device,
            linear_prob,
        )
        acc_list.append(final_acc)
        estp_acc_list.append(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:.4f}±{final_acc_std:.4f}")
    print(f"# early-stopping_acc: {estp_acc:.4f}±{estp_acc_std:.4f}")

    # --- 3. CLEAN UP ---
    print("\nJob Completed.")
    sys.stdout = logger_obj.terminal  # Restore original stdout
    logger_obj.close()


# Press the green button in the gutter to run the script.
if __name__ == "__main__":
    args = build_args()
    if args.use_cfg:
        args = load_best_configs(args)
    print(args)
    main(args)