import argparse import os import jax import flax import copy import torch import wandb import optax import itertools import numpy as np from tqdm import tqdm import jax.numpy as jnp from jax import random, vmap from flax import linen as nn from typing import Any, Dict, List from jax import random, vmap 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 utils import flatten_params, lerp, unflatten_params from lmc_model import LMCFlaxViTForImageClassification, print_model, print_model_with_prefix from transformers.models.vit.modeling_flax_vit import ViTConfig from matching_utils import weight_matching_attn from datasets import build_dataset from flax.serialization import from_bytes import multiprocessing as mp import matplotlib.pyplot as plt import numpy as np import json mp.set_start_method("spawn", force=True) def load_flax_params(checkpoint_dir, target): msgpack_path = os.path.join(checkpoint_dir, "flax_model.msgpack") with open(msgpack_path, "rb") as f: packed_bytes = f.read() params = from_bytes(target, packed_bytes) return params def imagenet_data_loader(args,is_generalization): dataset_val, args.nb_classes = build_dataset(is_train=False, args=args,is_generalization=is_generalization) sampler_val = torch.utils.data.SequentialSampler(dataset_val) 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_val def sample_random_batch(dataloader, seed=0): """Return (images, labels) from a random batch without materializing the loader.""" n = len(dataloader) # DataLoader should define __len__ idx = int(np.random.default_rng(seed).integers(n)) return next(itertools.islice(iter(dataloader), idx, None)) 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] def prepare_image_batch(images: torch.Tensor, labels: torch.Tensor): # torch -> numpy images = images.detach().cpu().numpy() # (B, C, H, W) labels = labels.detach().cpu().numpy() images = np.transpose(images, (0, 2, 3, 1)) # (B, H, W, C) images = images.astype(np.float32) labels = labels.astype(np.int32) return {"images": shard(images), "labels": shard(labels)} def make_stuff(model): apply_fn = model.__call__ def batch_eval(params, batch): outputs = apply_fn(params=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_batch_eval = jax.pmap(batch_eval, axis_name="batch") def dataset_loss_and_acc(params, dataloader): """Evaluate over the dataloader on all devices""" eval_results = [] pbar = tqdm(dataloader, desc="Evaluating", leave=False) rep_params = replicate(params) for step_idx, (images, labels) in enumerate(pbar): batch = prepare_image_batch(images, labels) eval_metric = parallel_batch_eval(rep_params, batch) eval_results.append(copy.deepcopy(eval_metric)) eval_metrics = get_metrics(eval_results) eval_metrics = unreplicate(eval_metrics) eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics) loss, acc1 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]) return loss, acc1 def get_attention_inputs(params, dataloader): activation = {} for images, labels in dataloader: images_host = jnp.array(images) _ = apply_fn(params=params,pixel_values=images_host,train=False,attention_input=activation,) break return activation return {"batch_eval": batch_eval, "dataset_loss_and_acc": dataset_loss_and_acc, "get_attention_inputs": get_attention_inputs} def compute_interpolation(params_a, params_b_target, lambdas, stuff, test_loader, desc="Interpolation"): train_loss_interp, test_loss_interp = [], [] train_acc_interp, test_acc_interp = [], [] for lam in tqdm(lambdas, desc=desc): p_interp = freeze(lerp(lam, unfreeze(params_a), unfreeze(params_b_target))) test_loss, test_acc = stuff["dataset_loss_and_acc"](p_interp, test_loader) test_loss_interp.append(test_loss) test_acc_interp.append(test_acc) return {"Test Loss": [float(f"{x:.4f}") for x in test_loss_interp],"Test Acc": [float(f"{x:.4f}") for x in test_acc_interp]} def main(): parser = argparse.ArgumentParser() parser.add_argument("--model-a", type=str, required=True, help="Path to first fine-tuned GPT2 model checkpoint") parser.add_argument("--model-b", type=str, required=True, help="Path to second fine-tuned GPT@ model checkpoint") parser.add_argument("--batch-size", type=int, default=256) parser.add_argument("--data-path", type=str, default=None) parser.add_argument("--data-original", type=str, required=True) parser.add_argument("--data-generalization", 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("--save-path", type=str, default="/", help="Path to plot directory") args = parser.parse_args() config = ViTConfig.from_json_file(os.path.join(os.path.dirname(args.model_a).rstrip("/"),'config.json')) lmc_config = ViTConfig(**config.lmc_config) config.lmc_config = lmc_config model = LMCFlaxViTForImageClassification(config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype),) print_model(model.params) if os.path.exists(args.model_a) and os.path.exists(args.model_b): params_a = load_flax_params(args.model_a,copy.deepcopy(model.params)) params_b = load_flax_params(args.model_b,copy.deepcopy(model.params)) else: raise FileNotFoundError(f"Checkpoint path does not exist") stuff = make_stuff(model = model) args.data_path = args.data_original original_val_loader = imagenet_data_loader(args,False) args.data_path = args.data_generalization generalization_val_loader = imagenet_data_loader(args,True) lambdas = jnp.linspace(0, 1, num=3) rng = random.PRNGKey(args.seed) # Compute naive interpolation # all_results = {"Naive": naive_results} naive_original_results = compute_interpolation(params_a, params_b, lambdas, stuff, original_val_loader, desc="Naive Original Interpolation") print(json.dumps({"Naive-Original": naive_original_results}, indent=2)) naive_original_results = compute_interpolation(params_a, params_b, lambdas, stuff, generalization_val_loader, desc="Naive Generalization Interpolation") print(json.dumps({"Naive-Generalization": naive_original_results}, indent=2)) exit() # Compute weight matching interpolations for each method activation = None #stuff["get_attention_inputs"](params_b, val_loader) aligned_models = weight_matching_attn(rng, params_a, params_b, activation, config) for method, params_b_aligned in aligned_models.items(): original_method_results = compute_interpolation(params_a, params_b_aligned, lambdas, stuff, original_val_loader, desc=f"{method} Interpolation") print(json.dumps({method+"original": original_method_results}, indent=2)) generalization_method_results = compute_interpolation(params_a, params_b_aligned, lambdas, stuff, generalization_val_loader, desc=f"{method} Interpolation") print(json.dumps({method+"generalization": generalization_method_results}, indent=2)) exit() # Save directories os.makedirs(f"./plots/imagenet", exist_ok=True) os.makedirs(f"./results/imagenet", exist_ok=True) # Save results JSON print("Save List of Values...") name_a = os.path.basename(os.path.dirname(args.model_a).rstrip("/")) name_b = os.path.basename(os.path.dirname(args.model_b).rstrip("/")) result_path = f'results/imagenet/{name_a}+{name_b}.json' with open(result_path, 'w') as f: json.dump(all_results, f, indent=2) # Plot print("Generating plots...") plot_path = f"./plots/imagenet/{name_a}+{name_b}.pdf" plt.rcParams.update({ "font.family": "serif", 'legend.frameon': False, 'lines.linewidth': 2, 'font.size': 13, 'axes.labelsize': 16, 'xtick.labelsize': 11, 'ytick.labelsize': 11, 'legend.fontsize': 11, }) plt.style.use('tableau-colorblind10') num_points = len(all_results["Naive"]["Test Loss"]) lambda_values = np.linspace(0, 1, num_points) fig, axs = plt.subplots(2, 2, figsize=(12, 10)) metrics = ["Test Loss", "Test Acc", "Test Loss", "Test Acc",] positions = [(0, 0), (0, 1), (1, 0), (1, 1)] for metric, pos in zip(metrics, positions): row, col = pos ax = axs[row, col] for method in all_results: ax.plot(lambda_values, all_results[method][metric], label=method) ax.set_xticks([0, 0.5, 1]) ax.set_xticklabels(["Model 1", r"$\lambda$", "Model 2"]) ax.set_ylabel(metric) ax.legend(loc='best') plt.tight_layout() plt.savefig(plot_path) plt.close() if __name__ == "__main__": main()