src/imagenet/test.py

import argparse
import os
import jax
import flax
import copy
import torch
import wandb
import optax
import numpy as np
from tqdm import tqdm
import jax.numpy as jnp
from flax import linen as nn
from typing import Any, Dict, List
from flax.jax_utils import replicate, unreplicate
from flax.core.frozen_dict import freeze, unfreeze
from flax.training import checkpoints, train_state
from flax.traverse_util import flatten_dict, unflatten_dict
from flax.training.common_utils import get_metrics, onehot, shard
from lmc_model import  LMCFlaxViTForImageClassification, print_model, print_model_with_prefix
from transformers.models.vit.modeling_flax_vit import ViTConfig
from datasets import build_dataset
import multiprocessing as mp
from pprint import pprint
import json
import shutil

mp.set_start_method("spawn", force=True)
os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7"

# ---------- Dataset Loader ----------
def remove_old_dirs_with_prefix(save_path, prefix, keep_step):
    for fname in os.listdir(save_path):
        if fname.startswith(prefix) and not fname.endswith(str(keep_step)):
            full_path = os.path.join(save_path, fname)
            if os.path.isdir(full_path):
                shutil.rmtree(full_path)
def imagenet_data_loader(args):
    dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
    dataset_val, _ = build_dataset(is_train=False, args=args)
    sampler_train = torch.utils.data.RandomSampler(dataset_train)
    sampler_val = torch.utils.data.SequentialSampler(dataset_val)
    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, sampler=sampler_train, batch_size=args.batch_size,
        num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True
    )
    data_loader_val = torch.utils.data.DataLoader(
        dataset_val, sampler=sampler_val, batch_size=args.batch_size,
        num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False
    )
    return data_loader_train, data_loader_val
def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]:
    images, labels =  jnp.array(images),jnp.array(labels)
    return {'images': shard(images),'labels': shard(labels)}
def accuracy(logits, labels, topk=(1,)):
    maxk = max(topk)
    batch_size = labels.shape[0]
    topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1]  # Top-k predictions
    res = []
    for k in topk:
        correct = (topk_preds[:, :k] == labels[:, None])
        correct = jnp.any(correct, axis=1)
        correct = jnp.sum(correct)
        res.append(100.0 * correct / batch_size)
    return res  # list of [acc@1, acc@5]
# ---------- Training Utilities ----------
def get_trainable_mask(params, config):
    def is_att_param(keys):
        if len(keys) < 5:
            return False
        return (keys[0] == "vit" and keys[1] == "encoder" and keys[2] == "layer" and int(keys[3]) in config.lmc_layer_indices and keys[4] == "attention")
    def label_fn(path, _):
        keys = [str(k.key) for k in path]
        if is_att_param(keys): return "trainable"
        return "frozen"
    return jax.tree_util.tree_map_with_path(label_fn, params)
def pretrained2finetune_params(pretrained_params, finetune_params, config):
    pretrained_params = unfreeze(pretrained_params)
    finetune_params = unfreeze(finetune_params)
    # 1. Copy top-level params (embeddings, layernorm, classifier)
    finetune_params["vit"]["embeddings"] = copy.deepcopy(pretrained_params["vit"]["embeddings"])
    finetune_params["vit"]["layernorm"] = copy.deepcopy(pretrained_params["vit"]["layernorm"])
    finetune_params["classifier"] = copy.deepcopy(pretrained_params["classifier"])
    # 2. Copy encoder layers
    for i in range(config.num_hidden_layers):
        str_i = str(i)
        if i in config.lmc_layer_indices:
            ref_layer = pretrained_params["vit"]["encoder"]["layer"][str_i]  # use layer 0 from pretrained
            target_layer = finetune_params["vit"]["encoder"]["layer"][str_i]
            # Copy shared parts
            target_layer["layernorm_before"] = copy.deepcopy(ref_layer["layernorm_before"])
            target_layer["layernorm_after"] = copy.deepcopy(ref_layer["layernorm_after"])
            target_layer["moe"] = copy.deepcopy(ref_layer["moe"])
            # #Copy Attention weights 
            # target_layer["attention"] = copy.deepcopy(ref_layer["attention"])
        else:
            finetune_params["vit"]["encoder"]["layer"][str_i] = copy.deepcopy(pretrained_params["vit"]["encoder"]["layer"][str_i])
    return freeze(finetune_params)


# ---------- Main ----------
def main(args: argparse.Namespace):
    if os.path.exists(args.model_path):
        config = ViTConfig.from_pretrained(os.path.dirname(args.model_path))        
    else:
        raise FileNotFoundError(f"Config directory does not exist: {args.model_path}")
    os.makedirs(args.wandb_run_dir, exist_ok=True)
    # --- Seeds & RNG ---
    save_path = os.path.join(
        args.save_dir,f"finetune-{config.position_embeddings}-indice{str(args.lmc_layer_indices)}-heads{args.num_attention_heads}"
        f"-shared{config.num_shared_experts}-routed{config.num_routed_experts}-topk{config.topk}-seed{args.seed}"
    )
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    rng = jax.random.PRNGKey(args.seed)
    train_loader, val_loader = imagenet_data_loader(args)
    # --- Load pretrained model ---
    pretrained_model = LMCFlaxViTForImageClassification(config, dtype=jnp.dtype(args.dtype))
    lmc_config = copy.deepcopy(config)
    lmc_config.num_attention_heads = args.num_attention_heads
    lmc_config.routed_scaling_factor = args.routed_scaling_factor
    config.lmc_config = lmc_config
    config.lmc_layer_indices = args.lmc_layer_indices
    pretrained_params = checkpoints.restore_checkpoint(ckpt_dir=args.model_path, target={"params": pretrained_model.params})["params"]
    pretrained_model.params = pretrained_params
    # --- Initialize fine-tuning model ---
    model = LMCFlaxViTForImageClassification(config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype),)
    model.config.save_pretrained(save_path)
    print_model(model.params)
    model.params = pretrained2finetune_params(pretrained_model.params,model.params,config)
    model = pretrained_model
    label_mask = get_trainable_mask(model.params,config)
    print(json.dumps(label_mask, indent=2))
    num_train_steps = len(train_loader)*args.epochs 
    num_warmup_steps = len(train_loader)*args.warmup_epochs
    lr_schedule =optax.warmup_cosine_decay_schedule(
        init_value=args.warmup_lr,
        peak_value=args.lr,
        warmup_steps=num_warmup_steps,
        decay_steps=num_train_steps,
        end_value=args.min_lr,
    )    
    tx = optax.multi_transform(
        transforms={
            'trainable': optax.adamw(
                learning_rate=lr_schedule,
                b1=args.adamw_beta1,
                b2=args.adamw_beta2,
                eps=args.adamw_eps,
                weight_decay=args.weight_decay
            ),
            'frozen': optax.set_to_zero()
        },
        param_labels=label_mask  
    )
    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx)
    if args.restore_checkpoint_path:
        state = checkpoints.restore_checkpoint(ckpt_dir=args.restore_checkpoint_path, target=state)
        print(f"train state restored from {args.restore_checkpoint_path}")
        print(f"skip trian step to {state.step}")
    latest_global_step = state.step
    curr_epoch = latest_global_step//len(train_loader)
    state = replicate(state)
    def train_step(state, batch, rng):
        dropout_rng, new_dropout_rng = jax.random.split(rng)
        def loss_fn(params):
            outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,)
            logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
            loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean()
            return loss, logits
        (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params)
        grads = jax.lax.pmean(grads, axis_name="batch")
        state = state.apply_gradients(grads=grads)
        acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5))
        metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),}
        metrics = jax.lax.pmean(metrics, axis_name="batch")
        return state, metrics, new_dropout_rng
    def eval_step(state, batch):
        outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,)
        logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
        loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean()
        acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5))
        metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,}
        metrics = jax.lax.pmean(metrics, axis_name="batch")
        return metrics
    parallel_train_step = jax.pmap(train_step, "batch")
    parallel_eval_step = jax.pmap(eval_step, "batch")
    rng = jax.random.PRNGKey(args.seed)
    train_metrics_stack = []
    train_loss = 0.0
    global_step = latest_global_step
    best_val_acc1 = 0.0
    ###JUST FOR TESTING####
    eval_results = []
    pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False)
    for batch_idx, (images, labels) in pbar:
        batch = prepare_image_batch(images, labels)
        eval_metric = parallel_eval_step(state, batch)
        eval_results.append(eval_metric)            
    # Compute mean metrics across all eval batches
    eval_metrics = get_metrics(eval_results)
    eval_metrics = unreplicate(eval_metrics)
    eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
    val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"])
    # val_loss, val_acc1, val_acc5 = float(np.mean(eval_metrics["loss"])), float(np.mean(eval_metrics["acc1"])), float(np.mean(eval_metrics["acc5"]))
    print("-" * 100)
    print(f"valid loss {val_loss:5.4f} | valid acc@1 {val_acc1:6.2f}% | valid acc@5 {val_acc5:6.2f}%")
    print("-" * 100)
    exit()
    #### START FINETUNING ####
    print("Starting Finetuning...")
    print(f"JAX devices: {jax.devices()}")
    print(f"Using {jax.local_device_count()} devices")
    for epoch in range(curr_epoch,args.epochs):
        print(f"Epoch {epoch}")
        dropout_rngs = jax.random.split(rng, jax.local_device_count())
        train_metrics_stack = []
        pbar = tqdm(enumerate(train_loader), desc="Training", leave=False)
        for batch_idx, (images, labels) in pbar:
            # Prepare and shard batch
            batch = prepare_image_batch(images,labels)
            # Run train step
            state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
            train_metrics_stack.append(train_metric)
            train_metrics = get_metrics(train_metrics_stack)
            train_metrics = unreplicate(train_metrics)
            train_metrics = jax.tree_util.tree_map(lambda x: x.mean(), train_metrics)
            train_metrics_stack = []
            loss, acc1, acc5 = float(train_metrics["loss"]), float(train_metrics["acc1"]), float(train_metrics["acc5"])
            curr_lr = float(lr_schedule(global_step))
            if global_step % args.wandb_logging_frequency == 0:
                wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step)
            pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",})
            global_step += 1
        eval_results = []
        pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False)
        for batch_idx, (images, labels) in pbar:
            batch = prepare_image_batch(images, labels)
            # Run eval step
            eval_metric = parallel_eval_step(state, batch)
            eval_results.append(eval_metric)            # Run train step
        # Compute mean metrics across all eval batches
        eval_metrics = get_metrics(eval_results)
        eval_metrics = unreplicate(eval_metrics)
        eval_metrics = jax.tree_util.tree_map(lambda x: x.mean(), eval_metrics)
        val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"])
        wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step)
        # Save best model
        if best_val_acc1 < val_acc1:
            best_val_acc1 = val_acc1
            model.params = unreplicate(state).params
            best_dir = os.path.join(save_path, f"best_{global_step}")
            model.save_pretrained(best_dir)
            print(f"Best model saved at step {global_step}")
            remove_old_dirs_with_prefix(save_path, "best_", global_step)
            # checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1,overwrite=True,orbax_checkpointer=None)                    
        # Save last model
        model.params = unreplicate(state).params
        last_dir = os.path.join(save_path, f"last_{global_step}")
        model.save_pretrained(last_dir)
        print(f"Last model saved at step {global_step}")
        remove_old_dirs_with_prefix(save_path, "last_", global_step)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet")
    # --- Model & Training Config ---
    parser.add_argument("--model-path", type=str, default="", help="Path of Pretrained Model")
    parser.add_argument("--num-attention-heads", type=int, default = 1)
    parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='')
    parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)")
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER')
    parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)')
    parser.add_argument('--warmup-epochs', type=int, default=3, metavar='N',help='epochs to warmup LR, if scheduler supports')
    parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)')
    parser.add_argument("--adamw-beta1", type=float, default=0.9)
    parser.add_argument("--adamw-beta2", type=float, default=0.999)    
    parser.add_argument("--adamw-eps", type=float, default=1e-8)
    parser.add_argument("--patience", type=int, default=10, help="Early stopping patience")
    parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path")
    # --- Data Config ---
    parser.add_argument("--data-path", type=str, required=True)
    parser.add_argument("--save-dir", type=str, required=True)
    parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19'])
    parser.add_argument("--input-size", type=int, default=224)
    parser.add_argument('--num_workers', type=int, default=8)
    parser.add_argument('--pin-mem', action='store_true')
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--color-jitter', type=float, default=0.4)
    parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1')
    parser.add_argument('--train-interpolation', type=str, default='bicubic')
    parser.add_argument('--reprob', type=float, default=0.25)
    parser.add_argument('--remode', type=str, default='pixel')
    parser.add_argument('--recount', type=int, default=1)
    parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype")
    parser.add_argument("--wandb-entity", default="", help="wandb entity for logging")
    parser.add_argument("--wandb-group", default="", help="wandb group for logging")
    parser.add_argument("--wandb-project", default="", help="wandb project name for logging")
    parser.add_argument("--wandb-run-dir", default=".wandb", help="wandb run dir")
    parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step")
    args = parser.parse_args()
    main(args)