TwT-6's picture
Upload 2667 files
256a159 verified
import argparse
import json
import os
import os.path as osp
import random
import time
from typing import List, Sequence
import mmengine
import torch
import torch.distributed as dist
from mmengine.config import Config, ConfigDict
from mmengine.device import get_device
from mmengine.dist import init_dist
from mmengine.evaluator import Evaluator
from mmengine.logging import print_log
from mmengine.model.wrappers import MMDistributedDataParallel
from mmengine.utils import track_iter_progress
from opencompass.registry import MM_MODELS, TASKS
from opencompass.utils import get_logger
def build_model(cfg):
model = MM_MODELS.build(cfg['model'])
load_from = cfg.get('load_from', None)
if load_from is not None:
state_dict = torch.load(cfg['load_from'], map_location='cpu')
if 'model' in state_dict:
state_dict = state_dict['model']
elif 'state_dict' in state_dict:
state_dict = state_dict['state_dict']
msg = model.load_state_dict(state_dict, strict=False)
print_log(msg)
model.to(get_device())
if dist.is_initialized():
model = MMDistributedDataParallel(
model,
device_ids=[int(os.environ['LOCAL_RANK'])],
broadcast_buffers=False)
return model
@TASKS.register_module(force=(__name__ == '__main__')) # A hack for script run
class MultimodalInferTask:
"""Multimodal Inference Task.
This task is used to run the inference process.
"""
def __init__(self, cfg: ConfigDict):
self.num_gpus = cfg.get('num_gpus', 0)
self.num_procs = cfg.get('num_procs', 1)
self.dataloader = cfg.get('dataset')
self.model = cfg.get('model')
self.evaluator = cfg.get('evaluator')
self.cfg = cfg
self.logger = get_logger()
@property
def name(self) -> str:
model_name = self.model['type']
dataset_name = self.dataloader['dataset']['type']
evaluator_name = self.evaluator[0]['type']
return f'{model_name}-{dataset_name}-{evaluator_name}'
def get_log_path(self, file_extension: str = 'json') -> str:
"""Get the path to the log file.
Args:
file_extension (str): The file extension of the log file.
Default: 'json'.
"""
model_name = self.model['type']
dataset_name = self.dataloader['dataset']['type']
evaluator_name = self.evaluator[0]['type']
return osp.join(self.cfg.work_dir, model_name, dataset_name,
f'{evaluator_name}.{file_extension}')
def get_output_paths(self, file_extension: str = 'json') -> List[str]:
"""Get the path to the output file.
Args:
file_extension (str): The file extension of the log file.
Default: 'json'.
"""
model_name = self.model['type']
dataset_name = self.dataloader['dataset']['type']
evaluator_name = self.evaluator[0]['type']
return [
osp.join(self.cfg.work_dir, model_name, dataset_name,
f'{evaluator_name}.{file_extension}')
]
def get_command(self, cfg_path, template):
"""Get the command template for the task.
Args:
cfg_path (str): The path to the config file of the task.
template (str): The template which have '{task_cmd}' to format
the command.
"""
script_path = __file__
if self.num_gpus > 0:
port = random.randint(12000, 32000)
command = (f'torchrun --master_port={port} '
f'--nproc_per_node {self.num_procs} '
f'{script_path} {cfg_path}')
else:
command = f'python {script_path} {cfg_path}'
return template.format(task_cmd=command)
def run(self):
from mmengine.runner import Runner
# only support slurm, pytorch, mpi
init_dist(self.cfg.launcher)
self.logger.info(f'Task {self.name}')
# build dataloader
dataloader = Runner.build_dataloader(self.dataloader)
# build model
model = build_model(self.cfg)
model.eval()
# build evaluator
evaluator = Evaluator(self.evaluator)
for batch in track_iter_progress(dataloader):
if dist.is_initialized():
data_samples = model.module.forward(batch)
else:
data_samples = model.forward(batch)
if not isinstance(data_samples, Sequence):
data_samples = [data_samples]
evaluator.process(data_samples)
metrics = evaluator.evaluate(len(dataloader.dataset))
metrics_file = self.get_output_paths()[0]
mmengine.mkdir_or_exist(osp.split(metrics_file)[0])
with open(metrics_file, 'w') as f:
json.dump(metrics, f)
def parse_args():
parser = argparse.ArgumentParser(description='Model Inferencer')
parser.add_argument('config', help='Config file path')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
cfg = Config.fromfile(args.config)
start_time = time.time()
inferencer = MultimodalInferTask(cfg)
inferencer.run()
end_time = time.time()
get_logger().info(f'time elapsed: {end_time - start_time:.2f}s')