|
|
|
|
|
|
|
|
import logging |
|
|
from dataclasses import dataclass |
|
|
from functools import partial |
|
|
from typing import Any, Dict, Optional |
|
|
|
|
|
from core.data.data_collators import MllmPaddingCollator |
|
|
from core.data.data_mixer import DatasetMixer, PersistentDataLoader |
|
|
from core.data.preprocessor import VisionPreprocessor |
|
|
from core.transforms.image_transform import get_image_transform |
|
|
from core.transforms.region_transform import get_region_transform |
|
|
from core.transforms.video_transform import get_video_transform |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class DataloadArgs: |
|
|
datamix: str = ( |
|
|
"dummy_image:1,dummy_multi_image:1,dummy_image_region:1,dummy_video:1,dummy_text:1,dummy_stc_RDCap:1,dummy_stc_RCap:1,dummy_stc_RTLoc:1" |
|
|
) |
|
|
batch_size: int = 2 |
|
|
seed: int = 42 |
|
|
image_res: Optional[int] = None |
|
|
max_num_tiles: Optional[int] = None |
|
|
vision_input_type: Optional[str] = None |
|
|
num_workers: Optional[int] = None |
|
|
tokenizer_path: Optional[str] = None |
|
|
tokenizer_name: Optional[str] = None |
|
|
conversation_format: Optional[str] = None |
|
|
patch_size: Optional[int] = None |
|
|
seq_len: Optional[int] = None |
|
|
max_video_frames: Optional[int] = None |
|
|
show_first_batch: Optional[bool] = False |
|
|
|
|
|
|
|
|
def get_rank_position(positions, rank, workers, world_size): |
|
|
if positions is not None and rank in positions: |
|
|
if positions["num_workers"] != workers or positions["world_size"] != world_size: |
|
|
logger.warning( |
|
|
f"Checkpoint resumed with different number of total dataloader workers. Dataloaders have been reset. " |
|
|
f"num_workers: {positions['num_workers']} -> {workers}, " |
|
|
f"world_size: {positions['world_size']} -> {world_size}" |
|
|
) |
|
|
return None |
|
|
return positions[rank] |
|
|
return None |
|
|
|
|
|
|
|
|
def get_dataloader( |
|
|
args, |
|
|
dp_rank, |
|
|
dp_world_size, |
|
|
dataset_configs: Dict[str, Any], |
|
|
tokenizer=None, |
|
|
positions=None, |
|
|
): |
|
|
vision_input_type = args.vision_input_type |
|
|
image_res = args.image_res |
|
|
max_num_tiles = args.max_num_tiles |
|
|
max_video_frames = args.max_video_frames |
|
|
|
|
|
preprocessor = partial( |
|
|
VisionPreprocessor, |
|
|
transform={ |
|
|
"image": get_image_transform( |
|
|
vision_input_type=vision_input_type, |
|
|
image_res=image_res, |
|
|
max_num_tiles=max_num_tiles, |
|
|
), |
|
|
"video": get_video_transform(image_res=image_res), |
|
|
"region": get_region_transform(), |
|
|
}, |
|
|
tokenizer=tokenizer, |
|
|
max_video_frames=max_video_frames, |
|
|
) |
|
|
|
|
|
dataset = DatasetMixer( |
|
|
args.datamix, |
|
|
global_rank=dp_rank, |
|
|
world_size=dp_world_size, |
|
|
seed=args.seed, |
|
|
preprocessors=[preprocessor], |
|
|
dataset_configs=dataset_configs, |
|
|
) |
|
|
|
|
|
|
|
|
dataloader = PersistentDataLoader( |
|
|
dataset, |
|
|
args.batch_size, |
|
|
args.num_workers, |
|
|
collate_fn=MllmPaddingCollator( |
|
|
tokenizer, |
|
|
show_first_batch=args.show_first_batch, |
|
|
), |
|
|
positions=get_rank_position( |
|
|
positions, dp_rank, args.num_workers, dp_world_size |
|
|
), |
|
|
) |
|
|
|
|
|
return dataloader |
|
|
|
