train_dynamic_ours.py

import argparse
from pathlib import Path
from datetime import datetime
import logging

from omegaconf import OmegaConf

import torch
from torch.utils.data import DataLoader

from torchvision import transforms

from src.utils import get_unet_model
from src.utils import load_dataset_from_config
from src.utils import set_logger
from src.utils import get_num_parameters
from src.utils import seed_everthing
# from src.runners_original import train_dp_promise
# from src.runners import train_dp_promise
# from src.runners_dynamic_phase1 import train_dp_promise
# from src.runners_dynamic_binary import train_dp_promise
from src.runners_dynamic_v0 import train_dp_promise
from src.runners_fedavg import train_fedavg_dp_promise
from src.mechanisms import eps_from_config
from src.tstr import MLP, build_resnet18, train_classifier

def main(config, opt):
    # Prepare
    time = datetime.now().strftime("%m%d-%H:%M:%S")
    running_id = f"dp-promise_{opt.note}_{config.data.name}_eps{config.dp.epsilon}_{time}"
    workdir = Path(opt.workdir) / running_id
    workdir.mkdir(parents=True)

    # Logging configuration
    (workdir / "logs").mkdir()
    set_logger(path=workdir / "logs")

    # Print Configuration
    logging.info(f"Configuration: {config}")

    # Save configuration
    OmegaConf.save(config, workdir / "config.yaml")
    # Setting seed
    seed_everthing(config.train.seed)
    logging.info(f"Using seed: {config.sample.seed}")

    # Load Datasets
    logging.info(f"Load dataset {config.data.name}")
    transform = transforms.Compose([
        transforms.ToTensor(),
        lambda x: x * 2. - 1.,
    ])
    dataset = load_dataset_from_config(config, transform, train=True)

    dataloader1 = DataLoader(
        dataset,
        batch_size=config.train.batch_size1,
        shuffle=True,
        pin_memory=True,
        num_workers=config.data.num_workers,
    )

    dataloader2 = DataLoader(
        dataset,
        batch_size=config.train.batch_size2,
        shuffle=True,
        pin_memory=True,
        num_workers=config.data.num_workers,
    )

    # Calculate privacy budget
    if config.dp.epsilon == "inf":
        logging.info("Epsilon: inf")
    else:
        eps = eps_from_config(config)
        logging.info(opt.note)
        logging.info(f"Satisfy ({eps}, {config.dp.delta})-DP")

    # Load model
    model = get_unet_model(
        config, checkpoint_path=config.train.get("ckpt_path"))
    logging.info(f"Num parameters: {get_num_parameters(model)}")

    # ---------------------------------------------------------
    # NEW: Pre-train classifiers on REAL data for TRTS
    # ---------------------------------------------------------
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logging.info("Pre-training TRTS classifiers on real data...")

    # # Grab dimensions. Adjust img_dim if you are using 32x32 images instead of 28x28
    # num_classes = getattr(config.data, 'num_classes', 10)
    # img_dim = 28 * 28  # 784 for MNIST/FMNIST
    # in_channels = 1    # 1 for grayscale

    # # Instantiate the models
    # mlp_clf = MLP(img_dim=img_dim, num_classes=num_classes).to(device)
    # resnet_clf = build_resnet18(num_classes=num_classes, in_channels=in_channels).to(device)

    # # Train them using dataloader1 (which contains your real dataset)
    # mlp_clf = train_classifier(mlp_clf, dataloader1, device, epochs=10)
    # resnet_clf = train_classifier(resnet_clf, dataloader1, device, epochs=10)

    # logging.info("Classifier pre-training complete.")
    # ---------------------------------------------------------

    logging.info(f"Start training...")

    if opt.is_fed:
        print("Running Federated + DP training")
        train_fedavg_dp_promise(
            model,
            dataloader1,
            dataloader2,
            config=config,
            workdir=workdir,
        )
    else:
        print("Running DP training (single node)")
        train_dp_promise(
            model,
            dataloader1,
            dataloader2,
            config=config,
            workdir=workdir,
            is_notc1=opt.is_notc1,
            sigma_start_arg=opt.sigma_start,
            sigma_end_arg=opt.sigma_end,
            sigma_kind_arg=opt.sigma_kind,
            # pretrained_mlp=mlp_clf,
            # pretrained_resnet=resnet_clf,
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config",
        type=str,
        required=True,
    )
    parser.add_argument(
        "--workdir",
        type=str,
        default="_workdir/dp_promise",
        required=False,
    )
    parser.add_argument(
        "--job-id",
        type=str,
        default="<job-id>",
    )
    parser.add_argument(
        "--is-fed",
        action="store_true",        # ✅ use store_true instead of string
        help="Enable federated training mode",
    )
    parser.add_argument(
        "--is-notc1",
        action="store_true",        # ✅ use store_true instead of string
        help="Enable federated training mode",
    )
    parser.add_argument(
        "--note",
        type=str,
        default="",
    )
    parser.add_argument(
        "--sigma-start",
        type=float,
        default=1.1708,
        help="Starting value for the sigma schedule",
    )
    parser.add_argument(
        "--sigma-end",
        type=float,
        default=1.4340,
        help="Ending value for the sigma schedule",
    )
    parser.add_argument(
        "--sigma-kind",
        type=str,
        default="power2",
        help="Kind of sigma schedule",
    )
    opt, _ = parser.parse_known_args()

    config = OmegaConf.load(opt.config)

    main(config, opt)