| | import logging |
| | import random |
| |
|
| | import numpy as np |
| | import torch |
| | from omegaconf import DictConfig |
| | from torch.utils.data import DataLoader, Dataset |
| | from torch.utils.data.dataloader import default_collate |
| | from torch.utils.data.distributed import DistributedSampler |
| |
|
| | from .eval.audiocaps import AudioCapsData |
| | from .eval.video_dataset import MovieGen, VGGSound |
| | from .extracted_audio import ExtractedAudio |
| | from .extracted_vgg import ExtractedVGG |
| | from .mm_dataset import MultiModalDataset |
| | from ..utils.dist_utils import local_rank |
| |
|
| | log = logging.getLogger() |
| |
|
| |
|
| | |
| | def worker_init_fn(worker_id: int): |
| | worker_seed = torch.initial_seed() % (2**31) + worker_id + local_rank * 1000 |
| | np.random.seed(worker_seed) |
| | random.seed(worker_seed) |
| | log.debug(f'Worker {worker_id} re-seeded with seed {worker_seed} in rank {local_rank}') |
| |
|
| |
|
| | def load_vgg_data(cfg: DictConfig, data_cfg: DictConfig) -> Dataset: |
| | dataset = ExtractedVGG(tsv_path=data_cfg.tsv, |
| | data_dim=cfg.data_dim, |
| | premade_mmap_dir=data_cfg.memmap_dir) |
| |
|
| | return dataset |
| |
|
| |
|
| | def load_audio_data(cfg: DictConfig, data_cfg: DictConfig) -> Dataset: |
| | dataset = ExtractedAudio(tsv_path=data_cfg.tsv, |
| | data_dim=cfg.data_dim, |
| | premade_mmap_dir=data_cfg.memmap_dir) |
| |
|
| | return dataset |
| |
|
| |
|
| | def setup_training_datasets(cfg: DictConfig) -> tuple[Dataset, DistributedSampler, DataLoader]: |
| | if cfg.mini_train: |
| | vgg = load_vgg_data(cfg, cfg.data.ExtractedVGG_val) |
| | audiocaps = load_audio_data(cfg, cfg.data.AudioCaps) |
| | dataset = MultiModalDataset([vgg], [audiocaps]) |
| | if cfg.example_train: |
| | video = load_vgg_data(cfg, cfg.data.Example_video) |
| | audio = load_audio_data(cfg, cfg.data.Example_audio) |
| | dataset = MultiModalDataset([video], [audio]) |
| | else: |
| | |
| | freesound = load_audio_data(cfg, cfg.data.FreeSound) |
| | vgg = load_vgg_data(cfg, cfg.data.ExtractedVGG) |
| | audiocaps = load_audio_data(cfg, cfg.data.AudioCaps) |
| | audioset_sl = load_audio_data(cfg, cfg.data.AudioSetSL) |
| | bbcsound = load_audio_data(cfg, cfg.data.BBCSound) |
| | clotho = load_audio_data(cfg, cfg.data.Clotho) |
| | dataset = MultiModalDataset([vgg] * cfg.vgg_oversample_rate, |
| | [audiocaps, audioset_sl, bbcsound, freesound, clotho]) |
| |
|
| | batch_size = cfg.batch_size |
| | num_workers = cfg.num_workers |
| | pin_memory = cfg.pin_memory |
| | sampler, loader = construct_loader(dataset, |
| | batch_size, |
| | num_workers, |
| | shuffle=True, |
| | drop_last=True, |
| | pin_memory=pin_memory) |
| |
|
| | return dataset, sampler, loader |
| |
|
| |
|
| | def setup_test_datasets(cfg): |
| | dataset = load_vgg_data(cfg, cfg.data.ExtractedVGG_test) |
| |
|
| | batch_size = cfg.batch_size |
| | num_workers = cfg.num_workers |
| | pin_memory = cfg.pin_memory |
| | sampler, loader = construct_loader(dataset, |
| | batch_size, |
| | num_workers, |
| | shuffle=False, |
| | drop_last=False, |
| | pin_memory=pin_memory) |
| |
|
| | return dataset, sampler, loader |
| |
|
| |
|
| | def setup_val_datasets(cfg: DictConfig) -> tuple[Dataset, DataLoader, DataLoader]: |
| | if cfg.example_train: |
| | dataset = load_vgg_data(cfg, cfg.data.Example_video) |
| | else: |
| | dataset = load_vgg_data(cfg, cfg.data.ExtractedVGG_val) |
| |
|
| | val_batch_size = cfg.batch_size |
| | val_eval_batch_size = cfg.eval_batch_size |
| | num_workers = cfg.num_workers |
| | pin_memory = cfg.pin_memory |
| | _, val_loader = construct_loader(dataset, |
| | val_batch_size, |
| | num_workers, |
| | shuffle=False, |
| | drop_last=False, |
| | pin_memory=pin_memory) |
| | _, eval_loader = construct_loader(dataset, |
| | val_eval_batch_size, |
| | num_workers, |
| | shuffle=False, |
| | drop_last=False, |
| | pin_memory=pin_memory) |
| |
|
| | return dataset, val_loader, eval_loader |
| |
|
| |
|
| | def setup_eval_dataset(dataset_name: str, cfg: DictConfig) -> tuple[Dataset, DataLoader]: |
| | if dataset_name.startswith('audiocaps_full'): |
| | dataset = AudioCapsData(cfg.eval_data.AudioCaps_full.audio_path, |
| | cfg.eval_data.AudioCaps_full.csv_path) |
| | elif dataset_name.startswith('audiocaps'): |
| | dataset = AudioCapsData(cfg.eval_data.AudioCaps.audio_path, |
| | cfg.eval_data.AudioCaps.csv_path) |
| | elif dataset_name.startswith('moviegen'): |
| | dataset = MovieGen(cfg.eval_data.MovieGen.video_path, |
| | cfg.eval_data.MovieGen.jsonl_path, |
| | duration_sec=cfg.duration_s) |
| | elif dataset_name.startswith('vggsound'): |
| | dataset = VGGSound(cfg.eval_data.VGGSound.video_path, |
| | cfg.eval_data.VGGSound.csv_path, |
| | duration_sec=cfg.duration_s) |
| | else: |
| | raise ValueError(f'Invalid dataset name: {dataset_name}') |
| |
|
| | batch_size = cfg.batch_size |
| | num_workers = cfg.num_workers |
| | pin_memory = cfg.pin_memory |
| | _, loader = construct_loader(dataset, |
| | batch_size, |
| | num_workers, |
| | shuffle=False, |
| | drop_last=False, |
| | pin_memory=pin_memory, |
| | error_avoidance=True) |
| | return dataset, loader |
| |
|
| |
|
| | def error_avoidance_collate(batch): |
| | batch = list(filter(lambda x: x is not None, batch)) |
| | return default_collate(batch) |
| |
|
| |
|
| | def construct_loader(dataset: Dataset, |
| | batch_size: int, |
| | num_workers: int, |
| | *, |
| | shuffle: bool = True, |
| | drop_last: bool = True, |
| | pin_memory: bool = False, |
| | error_avoidance: bool = False) -> tuple[DistributedSampler, DataLoader]: |
| | train_sampler = DistributedSampler(dataset, rank=local_rank, shuffle=shuffle) |
| | train_loader = DataLoader(dataset, |
| | batch_size, |
| | sampler=train_sampler, |
| | num_workers=num_workers, |
| | worker_init_fn=worker_init_fn, |
| | drop_last=drop_last, |
| | persistent_workers=num_workers > 0, |
| | pin_memory=pin_memory, |
| | collate_fn=error_avoidance_collate if error_avoidance else None) |
| | return train_sampler, train_loader |
| |
|
