Spaces:
Running
Running
File size: 7,205 Bytes
1fd4e9c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
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
|