File size: 6,636 Bytes
436b829 | 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 176 177 | from copy import deepcopy
import pytorch_lightning as pl
from hydra.utils import instantiate
from omegaconf import DictConfig, ListConfig
from pytorch_lightning.utilities.combined_loader import CombinedLoader
from torch.utils.data import ConcatDataset, DataLoader, Subset
import numpy as np
def mix_datasets(datasets, names, ratios, total=None):
if total is None:
total = min(int(len(ds) / ratios[n]) for ds, n in zip(datasets, names))
subsets = []
for ds, n in zip(datasets, names):
want = int(ratios[n] * total)
# Allow oversampling when the requested per-dataset count exceeds the
# dataset size — required for very small training subsets.
replace = want > len(ds)
idx = np.random.choice(len(ds), want, replace=replace)
subsets.append(Subset(ds, idx))
return ConcatDataset(subsets)
class GeneralDataModule(pl.LightningDataModule):
default_train_loader_opts = DictConfig(
{
"batch_size": 4,
"num_workers": 4,
"shuffle": True,
"pin_memory": True,
"drop_last": True,
# "persistent_workers": True,
}
)
default_val_loader_opts = DictConfig(
{
"batch_size": 1,
"num_workers": 4,
"shuffle": False,
"pin_memory": True,
"drop_last": False,
# "persistent_workers": True,
}
)
def __init__(
self,
train_dataset: DictConfig = None,
val_dataset: DictConfig = None,
test_dataset: DictConfig = None,
train_loader_opts: DictConfig = None,
val_loader_opts: DictConfig = None,
**kwargs,
):
"""
Initialize the GeneralDataModule with datasets and loader options.
This is a general datamodule that can be used for any dataset.
Train uses ConcatDataset. Val and Test use CombinedLoader, sequentially
consuming each iterable and returning a triplet (data, idx, iterable_idx).
Args:
train_dataset (DictConfig): Configuration for the training dataset.
val_dataset (DictConfig): Configuration for the validation dataset.
train_loader_opts (DictConfig): Options for the training data loader.
val_loader_opts (DictConfig): Options for the validation data loader.
**kwargs: Additional keyword arguments.
"""
super().__init__()
self.train_dataset = train_dataset
self.val_dataset = val_dataset
self.train_loader_opts = self.default_train_loader_opts
self.val_loader_opts = self.default_val_loader_opts
if train_loader_opts is not None:
self.train_loader_opts.update(train_loader_opts)
if val_loader_opts is not None:
self.val_loader_opts.update(val_loader_opts)
def val_dataloader(self):
"""
Create and return the validation data loader.
Returns:
CombinedLoader or DataLoader: The validation data loader.
"""
loaders = GeneralDataModule._parse_loaders(self.val_dataset, self.val_loader_opts)
if isinstance(loaders, list):
return CombinedLoader(loaders, mode="sequential")
else:
return loaders
def train_dataloader(self):
"""
Create and return the training data loader.
Returns:
DataLoader: The training data loader.
"""
return GeneralDataModule._parse_train_dataloader(self.train_dataset, self.train_loader_opts)
@staticmethod
def _parse_train_dataloader(config, loader_opts):
"""
Parse and create the training data loader from the configuration.
Args:
config (DictConfig): Configuration for the dataset.
loader_opts (DictConfig): Options for the data loader.
Returns:
DataLoader or CombinedLoader: The training data loader.
"""
if isinstance(config.dataset_opts, ListConfig):
datasets = GeneralDataModule._parse_datasets(config)
if config.pretrain:
names = ["hypersim"]
ratios = {"hypersim": 1.0}
else:
names = ["hypersim", "urbansyn", "unrealstereo4k", "vkitti", "tartanair"]
ratios = {"hypersim": 0.5, "urbansyn": 0.15, "unrealstereo4k": 0.15, "vkitti": 0.1, "tartanair": 0.1}
dataset = mix_datasets(datasets, names, ratios, total=48000)
return DataLoader(dataset, **loader_opts)
else:
return GeneralDataModule._parse_loaders(config, loader_opts)
@staticmethod
def _parse_datasets(config):
"""
Parse and instantiate datasets from the configuration.
Args:
config (DictConfig): Configuration for the datasets.
Returns:
list: A list of instantiated datasets.
"""
datasets = []
for idx, dataset_opt in enumerate(config.dataset_opts):
dataset = instantiate(dataset_opt)
datasets.append(dataset)
return datasets
@staticmethod
def _parse_loaders(config, loader_opts):
"""
Parse and create data loaders from the configuration.
Args:
config (DictConfig): Configuration for the datasets.
loader_opts (DictConfig): Options for the data loaders.
Returns:
DataLoader or list: A single DataLoader or a list of DataLoaders.
"""
if not isinstance(config.dataset_opts, ListConfig):
dataset = instantiate(config.dataset_opts)
if "loader_opts" in config:
loader_opts = deepcopy(loader_opts)
loader_opts.update(config.loader_opts)
return DataLoader(dataset, **loader_opts)
else:
dataloaders = []
for idx, dataset_opt in enumerate(config.dataset_opts):
if isinstance(dataset_opt, ListConfig):
datasets = [instantiate(opt) for opt in dataset_opt]
dataset = ConcatDataset(datasets)
else:
dataset = instantiate(dataset_opt)
if "loader_opts" in config:
loader_opt = deepcopy(loader_opts)
if isinstance(config.loader_opts, ListConfig):
loader_opt.update(config.loader_opts[idx])
else:
loader_opt.update(config.loader_opts)
dataloaders.append(DataLoader(dataset, **loader_opts))
return dataloaders |