import numpy as np
from torch.utils.data import DataLoader
from torch.utils.data.sampler import WeightedRandomSampler, SubsetRandomSampler
from milliontrees.common.utils import get_counts, split_into_groups
from milliontrees.datasets.milliontrees_dataset import MillionTreesDataset, MillionTreesSubset
[docs]
def get_train_loader(loader,
dataset,
batch_size,
uniform_over_groups=None,
grouper=None,
distinct_groups=True,
n_groups_per_batch=None,
**loader_kwargs):
"""Constructs and returns the data loader for training.
Args:
- loader (str): Loader type. 'standard' for standard loaders and 'group' for group loaders,
which first samples groups and then samples a fixed number of examples belonging
to each group.
- dataset (milliontreesDataset or milliontreesSubset): Data
- batch_size (int): Batch size
- uniform_over_groups (None or bool): Whether to sample the groups uniformly or according
to the natural data distribution.
Setting to None applies the defaults for each type of loaders.
For standard loaders, the default is False. For group loaders,
the default is True.
- grouper (Grouper): Grouper used for group loaders or for uniform_over_groups=True
- distinct_groups (bool): Whether to sample distinct_groups within each minibatch for group loaders.
- n_groups_per_batch (int): Number of groups to sample in each minibatch for group loaders.
- loader_kwargs: kwargs passed into torch DataLoader initialization.
Output:
- data loader (DataLoader): Data loader.
"""
if isinstance(dataset, MillionTreesDataset) and not isinstance(
dataset, MillionTreesSubset):
print(
"Warning: You are loading the entire dataset. Consider using dataset.get_subset('train') for a portion of the dataset if intended."
)
if loader == 'standard':
if uniform_over_groups is None or not uniform_over_groups:
return DataLoader(
dataset,
shuffle=True, # Shuffle training dataset
sampler=None,
collate_fn=dataset.collate,
batch_size=batch_size,
**loader_kwargs)
else:
assert grouper is not None
groups, group_counts = grouper.metadata_to_group(
dataset.metadata_array, return_counts=True)
group_weights = 1 / group_counts
weights = group_weights[groups]
# Replacement needs to be set to True, otherwise we'll run out of minority samples
sampler = WeightedRandomSampler(weights,
len(dataset),
replacement=True)
return DataLoader(
dataset,
shuffle=False, # The WeightedRandomSampler already shuffles
sampler=sampler,
collate_fn=dataset.collate,
batch_size=batch_size,
**loader_kwargs)
elif loader == 'group':
if uniform_over_groups is None:
uniform_over_groups = True
assert grouper is not None
assert n_groups_per_batch is not None
if n_groups_per_batch > grouper.n_groups:
raise ValueError(
f'n_groups_per_batch was set to {n_groups_per_batch} but there are only {grouper.n_groups} groups specified.'
)
group_ids = grouper.metadata_to_group(dataset.metadata_array)
batch_sampler = GroupSampler(group_ids=group_ids,
batch_size=batch_size,
n_groups_per_batch=n_groups_per_batch,
uniform_over_groups=uniform_over_groups,
distinct_groups=distinct_groups)
return DataLoader(dataset,
shuffle=None,
sampler=None,
collate_fn=dataset.collate,
batch_sampler=batch_sampler,
drop_last=False,
**loader_kwargs)
[docs]
def get_eval_loader(loader, dataset, batch_size, grouper=None, **loader_kwargs):
"""Constructs and returns the data loader for evaluation.
Args:
- loader (str): Loader type. 'standard' for standard loaders.
- dataset (milliontreesDataset or milliontreesSubset): Data
- batch_size (int): Batch size
- loader_kwargs: kwargs passed into torch DataLoader initialization.
Output:
- data loader (DataLoader): Data loader.
"""
if loader == 'standard':
return DataLoader(
dataset,
shuffle=False, # Do not shuffle eval datasets
sampler=None,
collate_fn=dataset.collate,
batch_size=batch_size,
**loader_kwargs)
[docs]
class GroupSampler:
"""Constructs batches by first sampling groups, then sampling data from those groups.
It drops the last batch if it's incomplete.
"""
def __init__(self, group_ids, batch_size, n_groups_per_batch,
uniform_over_groups, distinct_groups):
if batch_size % n_groups_per_batch != 0:
raise ValueError(
f'batch_size ({batch_size}) must be evenly divisible by n_groups_per_batch ({n_groups_per_batch}).'
)
if len(group_ids) < batch_size:
raise ValueError(
f'The dataset has only {len(group_ids)} examples but the batch size is {batch_size}. There must be enough examples to form at least one complete batch.'
)
self.group_ids = group_ids
self.unique_groups, self.group_indices, unique_counts = split_into_groups(
group_ids)
self.distinct_groups = distinct_groups
self.n_groups_per_batch = n_groups_per_batch
self.n_points_per_group = batch_size // n_groups_per_batch
self.dataset_size = len(group_ids)
self.num_batches = self.dataset_size // batch_size
if uniform_over_groups: # Sample uniformly over groups
self.group_prob = None
else: # Sample a group proportionately to its size
self.group_prob = unique_counts.numpy() / unique_counts.numpy().sum(
)
def __iter__(self):
for batch_id in range(self.num_batches):
# Note that we are selecting group indices rather than groups
groups_for_batch = np.random.choice(
len(self.unique_groups),
size=self.n_groups_per_batch,
replace=(not self.distinct_groups),
p=self.group_prob)
sampled_ids = [
np.random.choice(
self.group_indices[group],
size=self.n_points_per_group,
replace=len(self.group_indices[group]) <= self.
n_points_per_group, # False if the group is larger than the sample size
p=None) for group in groups_for_batch
]
# Flatten
sampled_ids = np.concatenate(sampled_ids)
yield sampled_ids
def __len__(self):
return self.num_batches