Spaces:
Running
Running
| 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 mmaudio.data.eval.audiocaps import AudioCapsData | |
| from mmaudio.data.eval.video_dataset import MovieGen, VGGSound | |
| from mmaudio.data.extracted_audio import ExtractedAudio | |
| from mmaudio.data.extracted_vgg import ExtractedVGG | |
| from mmaudio.data.mm_dataset import MultiModalDataset | |
| from mmaudio.utils.dist_utils import local_rank | |
| log = logging.getLogger() | |
| # Re-seed randomness every time we start a worker | |
| 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: | |
| # load the largest one first | |
| 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 | |