import inspect import json import math import os import time from datetime import datetime from pathlib import Path from typing import Any, Optional, Union, List, Iterator, TypeVar import pandas as pd import torch import torch.distributed as dist from PIL import Image from torch.utils.data import DataLoader, Dataset, Sampler from transformers.generation.utils import GenerateOutput from evaluation.entry import create_sampler_for_pipeline from rosetta.utils import get_args, get_logger, get_parallel_state from rosetta.utils import ParallelState from evaluation.sampling_dataset import MessageListDataset from evaluation.metrics import load_metric from rosetta.modeling import build_model from rosetta.utils import default, readable_time from rosetta.utils import safe_save_file, save_to_csv, save_to_json from rosetta.utils import Timer T_co = TypeVar('T_co', covariant=True) def _make_json_serializable(obj): if isinstance(obj, dict): return {k: _make_json_serializable(v) for k, v in obj.items()} if isinstance(obj, list): return [_make_json_serializable(v) for v in obj] if isinstance(obj, Image.Image): return f"" return obj def _serialize_output_data(data): if isinstance(data, torch.Tensor): return data.cpu().tolist() if isinstance(data, pd.DataFrame): return None if isinstance(data, list): return [ json.dumps(_make_json_serializable(d), ensure_ascii=False) if isinstance(d, (list, dict)) else d for d in data ] raise ValueError(f"Unsupported data type: {type(data)}") def _postprocess_images(outputs, save_images: bool): if not save_images: outputs.images = None if outputs.images is not None and isinstance(outputs.images, torch.Tensor): images_np = outputs.images.cpu().permute(0, 2, 3, 1).float().numpy() images_np = (images_np * 255).round().astype("uint8") outputs.images = [Image.fromarray(img) for img in images_np] def _save_generation_outputs(outputs, save_base: Path, summary_file_name: str): if outputs.is_empty(): return response = [dict(role="assistant", content=[]) for _ in outputs.batch["index"]] if outputs.texts is not None: texts = outputs.texts if isinstance(texts, GenerateOutput) or hasattr(texts, "sequences"): texts = texts.sequences if isinstance(texts, (list, tuple)) and all(isinstance(text, str) for text in texts): for i, text in enumerate(texts): response[i]["content"].append(dict(type="text", text=text)) if outputs.images is not None: save_image_base = save_base / "images" save_image_base.mkdir(parents=True, exist_ok=True) outputs.batch["gen_images"] = [] for i, image in enumerate(outputs.images): image_path = save_image_base / f"{outputs.batch['index'][i]}_0.png" image.save(image_path) outputs.batch["gen_images"].append(str(image_path)) response[i]["content"].append(dict(type="image", image_path=str(image_path))) outputs.batch["response"] = response serialized = {k: _serialize_output_data(v) for k, v in outputs.batch.items()} serialized = {k: v for k, v in serialized.items() if v is not None} save_to_csv(pd.DataFrame(serialized), save_base / summary_file_name, append=True) def _load_pretrained_model(model, dtype, ckpt_path): import re from accelerate import dispatch_model from transformers.modeling_utils import ( PreTrainedModel, _get_device_map, _get_resolved_checkpoint_files, # noqa ) from transformers.quantizers.quantizers_utils import get_module_from_name from rosetta.utils import is_package_version keep_in_fp32_regex = re.compile(r"\.gate\.wg") get_device_map_kwargs = dict(keep_in_fp32_regex=keep_in_fp32_regex) if is_package_version("transformers", ">=", "4.56"): get_device_map_kwargs["dtype"] = dtype else: get_device_map_kwargs["torch_dtype"] = dtype model.device_map = _get_device_map( model, device_map="auto" if torch.cuda.device_count() > 1 else "sequential", max_memory=None, hf_quantizer=None, **get_device_map_kwargs, ) print(f"Device map: \n{json.dumps(model.device_map, indent=4)}", flush=True) kwargs = {} if is_package_version("transformers", ">=", "4.53"): kwargs["is_remote_code"] = False checkpoint_files, sharded_metadata = _get_resolved_checkpoint_files( pretrained_model_name_or_path=ckpt_path, subfolder='', variant=None, gguf_file=None, from_tf=False, from_flax=False, use_safetensors=None, # noqa cache_dir=None, # noqa force_download=False, proxies=None, local_files_only=False, token=False, user_agent={'file_type': 'model', 'framework': 'pytorch', 'from_auto_class': False}, revision='main', commit_hash=None, **kwargs, ) ( model, missing_keys, unexpected_keys, mismatched_keys, offload_index, error_msgs, ) = PreTrainedModel._load_pretrained_model( # noqa model, None, checkpoint_files, ckpt_path, sharded_metadata=sharded_metadata, device_map=model.device_map, dtype=dtype, keep_in_fp32_regex=keep_in_fp32_regex, key_mapping=(model.get_key_mapping() if hasattr(model, "get_key_mapping") else None), weights_only=True, ) if len(missing_keys) > 0: for key in missing_keys: module, _ = get_module_from_name(model, key) if hasattr(module, "reset_parameters"): module.reset_parameters() model.tie_weights() print(f"Missing keys: {missing_keys}\nUnexpected keys: {unexpected_keys}", flush=True) if model.device_map is not None: dispatch_model(model, device_map=model.device_map) def build_tkwrapper(): from rosetta.tokenizer import load_tokenizer args = get_args() tkwrapper = load_tokenizer(args.tokenizer_name, args.tokenizer_class) return tkwrapper def build_vae(dp_rank=None): from rosetta.autoencoder import load_vae args = get_args() logger = get_logger() local_rank = int(os.environ.get("LOCAL_RANK", torch.cuda.current_device())) device = torch.device("cuda", local_rank) logger.info("Building VAE...") vae = load_vae( args.vae_type, args.vae_precision, device=device, logger=logger, args=args, ) if dp_rank is None: dp_rank = get_parallel_state().dp_rank generator = torch.Generator(device).manual_seed(args.seed + dp_rank) vae.generator = generator return vae class DistributedSamplerFix(Sampler[T_co]): def __init__(self, dataset: Dataset, num_replicas: Optional[int] = None, rank: Optional[int] = None, shuffle: bool = True, seed: int = 0, drop_last: bool = False, add_extra_samples: bool | str = False, ) -> None: if num_replicas is None: if not dist.is_available(): raise RuntimeError("Requires distributed package to be available") raise RuntimeError('Using `dist.get_world_size()` is dangerous.') num_replicas = dist.get_world_size() if rank is None: if not dist.is_available(): raise RuntimeError("Requires distributed package to be available") raise RuntimeError('Using `dist.get_rank()` is dangerous.') rank = dist.get_rank() if rank >= num_replicas or rank < 0: raise ValueError( f"Invalid rank {rank}, rank should be in the interval [0, {num_replicas - 1}]") self.dataset = dataset self.num_replicas = num_replicas self.rank = rank self.drop_last = drop_last self.add_extra_samples = add_extra_samples if self.drop_last and len(self.dataset) % self.num_replicas != 0: # type: ignore[arg-type] self.num_samples = math.ceil( (len(self.dataset) - self.num_replicas) / self.num_replicas # type: ignore[arg-type] ) self.total_size = self.num_samples * self.num_replicas elif self.add_extra_samples: self.num_samples = math.ceil(len(self.dataset) / self.num_replicas) # type: ignore[arg-type] self.total_size = self.num_samples * self.num_replicas else: total_size = len(self.dataset) # type: ignore[arg-type] self.num_samples = total_size // self.num_replicas + int( rank < total_size % self.num_replicas ) self.total_size = total_size self.shuffle = shuffle self.seed = seed def __iter__(self) -> Iterator[T_co]: dataset_length = len(self.dataset) # type: ignore[arg-type] if self.shuffle: g = torch.Generator() g.manual_seed(self.seed) indices = torch.randperm(dataset_length, generator=g).tolist() else: indices = list(range(dataset_length)) if not self.drop_last: if not self.add_extra_samples: pass elif self.add_extra_samples == "extend": padding_size = self.total_size - len(indices) indices += [index + dataset_length for index in range(padding_size)] else: padding_size = self.total_size - len(indices) if padding_size <= len(indices): indices += indices[:padding_size] else: indices += (indices * math.ceil(padding_size / len(indices)))[ :padding_size ] else: indices = indices[:self.total_size] assert len(indices) == self.total_size indices = indices[self.rank:self.total_size:self.num_replicas] assert len(indices) == self.num_samples return iter(indices) def __len__(self) -> int: return self.num_samples class MultimodalSampler(object): @classmethod def from_pretrained( cls, ckpt_path: Union[str, Path], config_path: Optional[Union[str, Path]] = None, device: int = 0, logger=None, extra_args: Optional[List[str]] = None, ): """Load sampler for pipeline/inference without distributed. Reuses evaluation.entry parsing (parse_argv_from_yaml + add_core_args) so that Gradio and run_sample.sh share the same config/arg path. extra_args: optional CLI args (e.g. ["--framework", "hf", "--sequence-template", "pretrain"]). """ ckpt_path = Path(ckpt_path) config_path = Path(config_path) if not ckpt_path.exists(): raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}.") if not config_path.exists(): raise FileNotFoundError(f"Config not found at {config_path}.") if logger is None: from loguru import logger as _logger logger = _logger return create_sampler_for_pipeline( config_path=str(config_path), ckpt_path=str(ckpt_path), extra_args=extra_args, sampler_name="multimodal_sampler.MultimodalSampler", framework="hf", device=device, logger_instance=logger, ) def __init__(self, init_args, rank, world_size): super().__init__() self.init_args = init_args self.device = rank % torch.cuda.device_count() self.rank = rank self.world_size = world_size self.parallel_state: ParallelState = get_parallel_state() self.logger = get_logger() self.setup_models(init_args.framework) self.setup_extra_models() self.pure_text_tasks = ["auto"] def setup_models(self, framework): args = get_args() # Initialize model. # For inference we always honour --bf16 so ViT flash_attention_2 gets bf16/fp16. dtype = torch.bfloat16 if args.bf16 else torch.float32 # Use args.init_device as-is (may be "meta" for FSDP training configs). # For the fsdp inference path we follow the standard PyTorch pattern: # build on meta → to_empty(cuda) → load_state_dict # This avoids OOM from keeping a full CPU copy before moving to CUDA. self.model, self.model_config = build_model( args, dtype=dtype, device=args.init_device, initialize_weights=False, ) assert args.ckpt is not None, "Checkpoint path `--ckpt` must be provided for sampling." if framework == "hf": _load_pretrained_model(self.model, dtype=dtype, ckpt_path=args.ckpt) self.model.requires_grad_(False) self.model.eval() self._cast_model_for_inference(dtype) elif framework == "fsdp": from pathlib import Path as _Path _ckpt = _Path(args.ckpt) if (_ckpt / "model.safetensors.index.json").exists(): self._load_hf_checkpoint(str(args.ckpt)) elif (_ckpt / "model.safetensors").exists(): self._load_single_safetensors_checkpoint(str(args.ckpt)) elif (_ckpt / ".metadata").exists() or (_ckpt / "weights" / ".metadata").exists(): self._load_dcp_checkpoint(str(args.ckpt)) else: raise ValueError(f"Unsupported checkpoint format for FSDP inference: {args.ckpt}") cuda_device = f"cuda:{self.device}" try: first_param = next(self.model.parameters()) if first_param.device.type == 'meta': self.model = self.model.to_empty(device=cuda_device) else: self.model = self.model.cuda() except StopIteration: self.model = self.model.cuda() self.model.requires_grad_(False) self.model.eval() self.model_engine = self.model self._cast_model_for_inference(dtype) else: raise NotImplementedError(f"Framework {framework} not supported.") self.model.load_generation_config(default(args.generation_config, args.ckpt)) def _cast_model_for_inference(self, dtype: torch.dtype): if dtype == torch.float32: return self.model = self.model.to(dtype=dtype) self._keep_moe_router_fp32() if hasattr(self.model, "_dtype"): self.model._dtype = dtype if hasattr(self.model, "vit_precision"): self.model.vit_precision = dtype def _keep_moe_router_fp32(self): for module in self.model.modules(): for attr in ("wg", "wg_text", "wg_vit", "wg_vae"): linear = getattr(module, attr, None) if isinstance(linear, torch.nn.Linear): linear.to(dtype=torch.float32) def _validate_loaded_checkpoint(self, missing, unexpected, strict: bool, label: str): allowed_missing_prefixes = ("timestep_emb.",) disallowed_missing = [ key for key in missing if not key.startswith(allowed_missing_prefixes) ] if strict and (disallowed_missing or unexpected): raise RuntimeError( f"Error(s) in loading state_dict for {label} checkpoint:\n" f"\tMissing key(s): {disallowed_missing}\n" f"\tUnexpected key(s): {unexpected}" ) meta_params = [name for name, param in self.model.named_parameters() if param.is_meta] if meta_params: raise RuntimeError( f"Model still has meta parameters after loading {label} checkpoint: " f"{meta_params[:8]}" ) def _load_dcp_checkpoint(self, load_dir: str, strict: bool = True): import torch import torch.distributed.checkpoint as DCP from torch.distributed.checkpoint import FileSystemReader from pathlib import Path load_path = Path(load_dir) if not (load_path / ".metadata").exists() and (load_path / "weights" / ".metadata").exists(): load_path = load_path / "weights" load_dir = str(load_path) if self.rank == 0: print(f"[DCP] Loading checkpoint from {load_dir} ...", flush=True) cuda_device = f"cuda:{self.device}" try: first_param = next(self.model.parameters()) if first_param.device.type == 'meta': if self.rank == 0: print(f"[DCP] Materialising meta tensors to {cuda_device} ...", flush=True) self.model = self.model.to_empty(device=cuda_device) except StopIteration: pass model_state_dict = dict(self.model.state_dict()) wrapped_state_dict = {"model": model_state_dict} try: from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner planner = DefaultLoadPlanner(flatten_sharded_tensors=True) except (ImportError, TypeError): planner = None DCP.load( state_dict=wrapped_state_dict, storage_reader=FileSystemReader(load_dir), planner=planner, ) missing, unexpected = self.model.load_state_dict(model_state_dict, strict=False) self._validate_loaded_checkpoint(missing, unexpected, strict, "DCP") if self.rank == 0: if missing: print(f"[DCP] Missing keys ({len(missing)}): {missing[:3]}{'...' if len(missing)>3 else ''}", flush=True) if unexpected: print(f"[DCP] Unexpected keys ({len(unexpected)}): {unexpected[:3]}{'...' if len(unexpected)>3 else ''}", flush=True) print(f"[DCP] Done.", flush=True) def _load_hf_checkpoint(self, load_dir: str, strict: bool = True): import json from pathlib import Path from safetensors.torch import load_file load_dir = Path(load_dir) index_file = load_dir / "model.safetensors.index.json" if self.rank == 0: print(f"[HF] Loading checkpoint from {load_dir} ...", flush=True) # Materialise meta tensors to CUDA before loading weights. cuda_device = f"cuda:{self.device}" try: first_param = next(self.model.parameters()) if first_param.device.type == "meta": if self.rank == 0: print(f"[HF] Materialising meta tensors to {cuda_device} ...", flush=True) self.model = self.model.to_empty(device=cuda_device) except StopIteration: pass # All ranks load from the shared filesystem in parallel. with index_file.open() as f: index_data = json.load(f) weight_map = index_data["weight_map"] shard_files = sorted(set(weight_map.values())) full_state_dict = {} for shard_fname in shard_files: if self.rank == 0: print(f"[HF] Loading shard {shard_fname} ...", flush=True) shard = load_file(str(load_dir / shard_fname), device=cuda_device) full_state_dict.update(shard) missing, unexpected = self.model.load_state_dict(full_state_dict, strict=False) self._validate_loaded_checkpoint(missing, unexpected, strict, "HF") if self.rank == 0: if missing: print(f"[HF] Missing keys ({len(missing)}): {missing[:3]}{'...' if len(missing)>3 else ''}", flush=True) if unexpected: print(f"[HF] Unexpected keys ({len(unexpected)}): {unexpected[:3]}{'...' if len(unexpected)>3 else ''}", flush=True) print(f"[HF] Done.", flush=True) def _load_single_safetensors_checkpoint(self, load_dir: str, strict: bool = True): from pathlib import Path from safetensors.torch import load_file load_dir = Path(load_dir) checkpoint_file = load_dir / "model.safetensors" if self.rank == 0: print(f"[safetensors] Loading checkpoint from {checkpoint_file} ...", flush=True) cuda_device = f"cuda:{self.device}" try: first_param = next(self.model.parameters()) if first_param.device.type == "meta": if self.rank == 0: print(f"[safetensors] Materialising meta tensors to {cuda_device} ...", flush=True) self.model = self.model.to_empty(device=cuda_device) except StopIteration: pass state_dict = load_file(str(checkpoint_file), device=cuda_device) missing, unexpected = self.model.load_state_dict(state_dict, strict=False) self._validate_loaded_checkpoint(missing, unexpected, strict, "single-file safetensors") if self.rank == 0: if missing: print( f"[safetensors] Missing keys ({len(missing)}): " f"{missing[:3]}{'...' if len(missing)>3 else ''}", flush=True, ) if unexpected: print( f"[safetensors] Unexpected keys ({len(unexpected)}): " f"{unexpected[:3]}{'...' if len(unexpected)>3 else ''}", flush=True, ) print(f"[safetensors] Done.", flush=True) def setup_extra_models(self): args = get_args() # Initialize vae, tokenizer self.model.tokenizer = build_tkwrapper() if args.use_vae: self.model.model_dict["vae"] = build_vae(dp_rank=self.rank) def run(self): args = get_args() dp_rank = self.parallel_state.dp_rank assert args.prompt is not None, "'prompt' must be provided for generation." bot_task = self.model.generation_config.bot_task if bot_task in self.pure_text_tasks: # Pure text generation inputs = self.model.prepare_model_inputs(prompt=args.prompt, image=args.image, bot_task=bot_task) self.model.generate(**inputs, verbose=2) else: # Hybrid text-image generation generation_outputs = self.model.generate_image( prompt=args.prompt[self.rank % len(args.prompt)], image=args.image, seed=args.seed, image_size=args.image_size, bot_task=bot_task, verbose=2, ) texts, images = generation_outputs.texts, generation_outputs.images if texts is not None: print(f"[rank {dp_rank}] Generated Text: {texts}") for i, image in enumerate(images): image.save(f"image_{dp_rank}_{i}.png") print(f"Image saved to image_{dp_rank}_{i}.png") @staticmethod def postprocess_results(save_base): results_dir = save_base / "results" if not results_dir.exists(): return if torch.distributed.get_rank() == 0: all_dfs = [] for csv_file in sorted(results_dir.glob("results_*.csv"), key=lambda x: int(x.stem.split("_")[-1])): df = pd.read_csv(csv_file) all_dfs.append(df) merged_df = pd.concat(all_dfs, ignore_index=True) if 'index' in merged_df.columns: merged_df = merged_df.sort_values(by='index') merged_save_path = save_base / "results/all_results.csv" merged_df.to_csv(merged_save_path, index=False) def runtime_media_generation_config(self, run_task_kwargs): args = get_args() runtime_config = dict( bot_task=run_task_kwargs.get("bot_task", self.model.generation_config.bot_task), image_size=run_task_kwargs.get("image_size", args.image_size), ) return runtime_config @staticmethod def per_batch_runtime_config(batch, runtime_config): if "height" in batch and "width" in batch: image_size = [(int(h), int(w)) for h, w in zip(batch["height"], batch["width"])] else: image_size = runtime_config["image_size"] per_batch_runtime_config = {**runtime_config, "image_size": image_size} return per_batch_runtime_config def build_dataset(self, testset): args = get_args() return MessageListDataset(testset, args.sample_save_base, tokenizer=self.model.tokenizer) def run_testsets(self): args = get_args() assert args.testsets is not None or args.eval_metrics is not None, \ "'testsets' or 'eval_metrics' must be provided for run_testsets." assert args.sample_save_base is not None, "'sample_save_base' must be provided for run_testsets." run_tasks = [("sample", testset) for testset in (args.testsets or [])] + [ ("eval_metric", testset) for testset in (args.eval_metrics or []) ] timer = Timer(enabled=True) timer.start("global") for task_idx, (run_task_type, testset) in enumerate(run_tasks): generate_task_specific_kwargs = dict() dataset = self.build_dataset(testset) run_task_kwargs = dataset.task_kwargs sampler = DistributedSamplerFix(dataset, num_replicas=self.parallel_state.dp_size, rank=self.parallel_state.dp_rank, shuffle=False, drop_last=False, add_extra_samples="extend") dataloader = DataLoader(dataset, batch_size=args.sample_batch_size, shuffle=False, sampler=sampler, drop_last=False, collate_fn=getattr(dataset, "collate_fn", None)) save_base = dataset.save_dir self.logger.info(f"=" * 80) self.logger.info(f"Running task {testset}({task_idx + 1} / {len(run_tasks)})") self.logger.info(f"Save directory: {save_base}") self.logger.info(f"=" * 80) metric_instances = [] if run_task_type == "eval_metric": kwargs = {} LOGIT_BASED_TESTSETS = {"mmlu_bench", "arc_challenge", "mmmlu"} if dataset.testset in LOGIT_BASED_TESTSETS: kwargs = {"tokenizer": self.model.tokenizer} generate_task_specific_kwargs["return_dict_in_generate"] = True generate_task_specific_kwargs["output_logits"] = True assert "metric" in run_task_kwargs, f"'metric' must be specified in testset for eval_metric task." metric_types = run_task_kwargs["metric"].split("+") for metric_type in metric_types: metric = load_metric(f"{metric_type}@{dataset.testset}", **kwargs) metric.load_model(self.logger) metric_instances.append((metric, metric_type)) bot_task = run_task_kwargs.get("bot_task", self.model.generation_config.bot_task) generate_task_specific_kwargs["bot_task"] = bot_task verbose = args.verbose if args.sample_batch_size == 1 else min(1, args.verbose) if bot_task in self.pure_text_tasks: max_new_tokens = int(run_task_kwargs.get("max_new_tokens", self.model.generation_config.max_new_tokens)) generate_task_specific_kwargs["max_new_tokens"] = max_new_tokens model_type = run_task_kwargs.get("model_type", None) skip_special_tokens = False if model_type else True runtime_config = {**self.model.generation_config.to_dict(), **generate_task_specific_kwargs} timer.start(f"Task {task_idx}") timer.start(f"Batch") for batch_idx, batch in enumerate(dataloader): batch: dict[str, Any] self.logger.info(f"Generating batch {batch_idx + 1} / {len(dataloader)} ...") inputs = self.model.prepare_model_inputs( message_list=batch[dataset.name_mapper("message_list")], mode="gen_text", **generate_task_specific_kwargs, ) outputs = self.model.generate( **inputs, decode_text=True, verbose=verbose, skip_special_tokens=skip_special_tokens ) outputs = outputs.postprocess_outputs(batch) if run_task_type == "eval_metric" and outputs is not None: start_time = time.time() for metric, _ in metric_instances: proc_inputs = {k: v for k, v in batch.items()} proc_inputs["answers"] = outputs.texts metric.process(**proc_inputs, model_type=model_type) if isinstance(outputs.texts, GenerateOutput) or hasattr(outputs.texts, "sequences"): outputs.texts = outputs.texts.sequences gen_time = time.time() - start_time self.logger.info(f"Metric process time: {gen_time}") _save_generation_outputs( outputs, save_base=save_base, summary_file_name=f"results/results_{self.parallel_state.dp_rank}.csv", ) timer.stop(f"Batch") self.logger.info(f"Task {testset}({task_idx + 1}/{len(run_tasks)})" f"[{batch_idx + 1} / {len(dataloader)}] " f"| {readable_time(timer, 'Batch', len(dataloader) - batch_idx - 1)}") timer.start(f"Batch") timer.stop("Batch") timer.stop(f"Task {task_idx}") self.logger.info(f"Save directory: {save_base}") else: runtime_config = self.runtime_media_generation_config(run_task_kwargs) _image_output = dict(sample="pil", eval_metric="pt") output_type = _image_output[run_task_type] timer.start(f"Task {task_idx}") timer.start(f"Batch") for batch_idx, batch in enumerate(dataloader): batch: dict[str, Any] self.logger.info(f"Generating batch {batch_idx + 1} / {len(dataloader)} ...") batch_config = self.per_batch_runtime_config(batch, runtime_config) outputs = self.model.generate_image( message_list=batch[dataset.name_mapper("message_list")], seed=batch["seed"], **batch_config, image_output_type=output_type, verbose=verbose, ) outputs = outputs.postprocess_outputs(batch) if run_task_type == "eval_metric" and not outputs.is_empty(): start_time = time.time() for metric, _ in metric_instances: proc_inputs = {k: v for k, v in outputs.batch.items()} proc_inputs["images"] = outputs.images.float() proc_inputs.setdefault("ids", outputs.batch["index"]) metric.process(**proc_inputs) gen_time = time.time() - start_time _postprocess_images(outputs, args.eval_save_images) self.logger.info(f"Metric process time: {gen_time}") _save_generation_outputs( outputs, save_base=save_base, summary_file_name=f"results/results_{self.parallel_state.dp_rank}.csv", ) timer.stop(f"Batch") self.logger.info(f"[Task {testset}({task_idx + 1}/{len(run_tasks)})] " f"[{batch_idx + 1} / {len(dataloader)}] " f"| {readable_time(timer, 'Batch', len(dataloader) - batch_idx - 1)}") timer.start(f"Batch") timer.stop("Batch") timer.stop(f"Task {task_idx}") self.logger.info(f"Save directory: {save_base}") torch.distributed.barrier() self.postprocess_results(save_base) if run_task_type == "eval_metric": dist.barrier() self.logger.info(f"All processes have finished the evaluation. Start all gathering results...") metric_gather_results = [metric.all_gather_results() for metric, _ in metric_instances] self.logger.info(f"All gather results for all metrics finished") if self.rank == 0: timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") for (metric, metric_type), gathered_results in zip(metric_instances, metric_gather_results): metric_kwargs = {} accept_save_file = 'save_file' in list(inspect.signature(metric.compute_metrics).parameters.keys()) if accept_save_file and hasattr(metric, "save_file_template"): metric_kwargs['save_file'] = save_base / f"metric_temp/{metric_type}" \ / metric.save_file_template.format(timestamp.replace('-', '_')) self.logger.info(f"[{metric_type}] Temp results will be saved to {metric_kwargs['save_file']}") output = metric.compute_metrics(gathered_results, **metric_kwargs) if isinstance(output, tuple): output = {"_": output} results = [] for key, (value, count, *extra_outputs) in output.items(): suffix = "" if key == "_" else f"_{key}" results.append({ "timestamp": timestamp, "metric": f"{metric_type}{suffix}", "testset": dataset.testset, "value": value, "count": count, "runtime_config": { **self.model.generation_config.to_dict(), **runtime_config, }, }) if len(extra_outputs) > 0: # Metrics like VQAv2 return a dict to save some extra information if 'extra_info_dict' in extra_outputs[0]: results[-1]["extra_metric_stats"] = extra_outputs[0]['extra_info_dict'] self.logger.info(results) accumulated_results = results[:] save_path = save_base / f"metric_results/{metric_type}.json" if save_path.exists(): with open(save_path, "r") as f: ori_results = json.load(f) accumulated_results = ori_results + results save_to = safe_save_file(save_path, accumulated_results, save_fn=save_to_json) self.logger.info(f"Evaluation results saved to {save_to}") dist.barrier() timer.stop("global") self.logger.info(f"Total time cost: {readable_time(timer.elapsed('global'))}.") def exit(self): torch.cuda.empty_cache() if torch.distributed.is_initialized(): dist.barrier() print(f"[rank {self.rank}] Sampling is complete. Exiting now.", flush=True) dist.destroy_process_group()