Detic / detic /data /custom_dataset_dataloader.py
AK391
files
159f437
# Copyright (c) Facebook, Inc. and its affiliates.
# Part of the code is from https://github.com/xingyizhou/UniDet/blob/master/projects/UniDet/unidet/data/multi_dataset_dataloader.py (Apache-2.0 License)
import copy
import logging
import numpy as np
import operator
import torch
import torch.utils.data
import json
from detectron2.utils.comm import get_world_size
from detectron2.utils.logger import _log_api_usage, log_first_n
from detectron2.config import configurable
from detectron2.data import samplers
from torch.utils.data.sampler import BatchSampler, Sampler
from detectron2.data.common import DatasetFromList, MapDataset
from detectron2.data.dataset_mapper import DatasetMapper
from detectron2.data.build import get_detection_dataset_dicts, build_batch_data_loader
from detectron2.data.samplers import TrainingSampler, RepeatFactorTrainingSampler
from detectron2.data.build import worker_init_reset_seed, print_instances_class_histogram
from detectron2.data.build import filter_images_with_only_crowd_annotations
from detectron2.data.build import filter_images_with_few_keypoints
from detectron2.data.build import check_metadata_consistency
from detectron2.data.catalog import MetadataCatalog, DatasetCatalog
from detectron2.utils import comm
import itertools
import math
from collections import defaultdict
from typing import Optional
def _custom_train_loader_from_config(cfg, mapper=None, *, dataset=None, sampler=None):
sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
if 'MultiDataset' in sampler_name:
dataset_dicts = get_detection_dataset_dicts_with_source(
cfg.DATASETS.TRAIN,
filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
if cfg.MODEL.KEYPOINT_ON else 0,
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
)
else:
dataset_dicts = get_detection_dataset_dicts(
cfg.DATASETS.TRAIN,
filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
if cfg.MODEL.KEYPOINT_ON else 0,
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
)
if mapper is None:
mapper = DatasetMapper(cfg, True)
if sampler is not None:
pass
elif sampler_name == "TrainingSampler":
sampler = TrainingSampler(len(dataset))
elif sampler_name == "MultiDatasetSampler":
sampler = MultiDatasetSampler(
dataset_dicts,
dataset_ratio = cfg.DATALOADER.DATASET_RATIO,
use_rfs = cfg.DATALOADER.USE_RFS,
dataset_ann = cfg.DATALOADER.DATASET_ANN,
repeat_threshold = cfg.DATALOADER.REPEAT_THRESHOLD,
)
elif sampler_name == "RepeatFactorTrainingSampler":
repeat_factors = RepeatFactorTrainingSampler.repeat_factors_from_category_frequency(
dataset_dicts, cfg.DATALOADER.REPEAT_THRESHOLD
)
sampler = RepeatFactorTrainingSampler(repeat_factors)
else:
raise ValueError("Unknown training sampler: {}".format(sampler_name))
return {
"dataset": dataset_dicts,
"sampler": sampler,
"mapper": mapper,
"total_batch_size": cfg.SOLVER.IMS_PER_BATCH,
"aspect_ratio_grouping": cfg.DATALOADER.ASPECT_RATIO_GROUPING,
"num_workers": cfg.DATALOADER.NUM_WORKERS,
'multi_dataset_grouping': cfg.DATALOADER.MULTI_DATASET_GROUPING,
'use_diff_bs_size': cfg.DATALOADER.USE_DIFF_BS_SIZE,
'dataset_bs': cfg.DATALOADER.DATASET_BS,
'num_datasets': len(cfg.DATASETS.TRAIN)
}
@configurable(from_config=_custom_train_loader_from_config)
def build_custom_train_loader(
dataset, *, mapper, sampler,
total_batch_size=16,
aspect_ratio_grouping=True,
num_workers=0,
num_datasets=1,
multi_dataset_grouping=False,
use_diff_bs_size=False,
dataset_bs=[]
):
"""
Modified from detectron2.data.build.build_custom_train_loader, but supports
different samplers
"""
if isinstance(dataset, list):
dataset = DatasetFromList(dataset, copy=False)
if mapper is not None:
dataset = MapDataset(dataset, mapper)
if sampler is None:
sampler = TrainingSampler(len(dataset))
assert isinstance(sampler, torch.utils.data.sampler.Sampler)
if multi_dataset_grouping:
return build_multi_dataset_batch_data_loader(
use_diff_bs_size,
dataset_bs,
dataset,
sampler,
total_batch_size,
num_datasets=num_datasets,
num_workers=num_workers,
)
else:
return build_batch_data_loader(
dataset,
sampler,
total_batch_size,
aspect_ratio_grouping=aspect_ratio_grouping,
num_workers=num_workers,
)
def build_multi_dataset_batch_data_loader(
use_diff_bs_size, dataset_bs,
dataset, sampler, total_batch_size, num_datasets, num_workers=0
):
"""
"""
world_size = get_world_size()
assert (
total_batch_size > 0 and total_batch_size % world_size == 0
), "Total batch size ({}) must be divisible by the number of gpus ({}).".format(
total_batch_size, world_size
)
batch_size = total_batch_size // world_size
data_loader = torch.utils.data.DataLoader(
dataset,
sampler=sampler,
num_workers=num_workers,
batch_sampler=None,
collate_fn=operator.itemgetter(0), # don't batch, but yield individual elements
worker_init_fn=worker_init_reset_seed,
) # yield individual mapped dict
if use_diff_bs_size:
return DIFFMDAspectRatioGroupedDataset(
data_loader, dataset_bs, num_datasets)
else:
return MDAspectRatioGroupedDataset(
data_loader, batch_size, num_datasets)
def get_detection_dataset_dicts_with_source(
dataset_names, filter_empty=True, min_keypoints=0, proposal_files=None
):
assert len(dataset_names)
dataset_dicts = [DatasetCatalog.get(dataset_name) for dataset_name in dataset_names]
for dataset_name, dicts in zip(dataset_names, dataset_dicts):
assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)
for source_id, (dataset_name, dicts) in \
enumerate(zip(dataset_names, dataset_dicts)):
assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)
for d in dicts:
d['dataset_source'] = source_id
if "annotations" in dicts[0]:
try:
class_names = MetadataCatalog.get(dataset_name).thing_classes
check_metadata_consistency("thing_classes", dataset_name)
print_instances_class_histogram(dicts, class_names)
except AttributeError: # class names are not available for this dataset
pass
assert proposal_files is None
dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))
has_instances = "annotations" in dataset_dicts[0]
if filter_empty and has_instances:
dataset_dicts = filter_images_with_only_crowd_annotations(dataset_dicts)
if min_keypoints > 0 and has_instances:
dataset_dicts = filter_images_with_few_keypoints(dataset_dicts, min_keypoints)
return dataset_dicts
class MultiDatasetSampler(Sampler):
def __init__(
self,
dataset_dicts,
dataset_ratio,
use_rfs,
dataset_ann,
repeat_threshold=0.001,
seed: Optional[int] = None,
):
"""
"""
sizes = [0 for _ in range(len(dataset_ratio))]
for d in dataset_dicts:
sizes[d['dataset_source']] += 1
print('dataset sizes', sizes)
self.sizes = sizes
assert len(dataset_ratio) == len(sizes), \
'length of dataset ratio {} should be equal to number if dataset {}'.format(
len(dataset_ratio), len(sizes)
)
if seed is None:
seed = comm.shared_random_seed()
self._seed = int(seed)
self._rank = comm.get_rank()
self._world_size = comm.get_world_size()
self.dataset_ids = torch.tensor(
[d['dataset_source'] for d in dataset_dicts], dtype=torch.long)
dataset_weight = [torch.ones(s) * max(sizes) / s * r / sum(dataset_ratio) \
for i, (r, s) in enumerate(zip(dataset_ratio, sizes))]
dataset_weight = torch.cat(dataset_weight)
rfs_factors = []
st = 0
for i, s in enumerate(sizes):
if use_rfs[i]:
if dataset_ann[i] == 'box':
rfs_func = RepeatFactorTrainingSampler.repeat_factors_from_category_frequency
else:
rfs_func = repeat_factors_from_tag_frequency
rfs_factor = rfs_func(
dataset_dicts[st: st + s],
repeat_thresh=repeat_threshold)
rfs_factor = rfs_factor * (s / rfs_factor.sum())
else:
rfs_factor = torch.ones(s)
rfs_factors.append(rfs_factor)
st = st + s
rfs_factors = torch.cat(rfs_factors)
self.weights = dataset_weight * rfs_factors
self.sample_epoch_size = len(self.weights)
def __iter__(self):
start = self._rank
yield from itertools.islice(
self._infinite_indices(), start, None, self._world_size)
def _infinite_indices(self):
g = torch.Generator()
g.manual_seed(self._seed)
while True:
ids = torch.multinomial(
self.weights, self.sample_epoch_size, generator=g,
replacement=True)
nums = [(self.dataset_ids[ids] == i).sum().int().item() \
for i in range(len(self.sizes))]
yield from ids
class MDAspectRatioGroupedDataset(torch.utils.data.IterableDataset):
def __init__(self, dataset, batch_size, num_datasets):
"""
"""
self.dataset = dataset
self.batch_size = batch_size
self._buckets = [[] for _ in range(2 * num_datasets)]
def __iter__(self):
for d in self.dataset:
w, h = d["width"], d["height"]
aspect_ratio_bucket_id = 0 if w > h else 1
bucket_id = d['dataset_source'] * 2 + aspect_ratio_bucket_id
bucket = self._buckets[bucket_id]
bucket.append(d)
if len(bucket) == self.batch_size:
yield bucket[:]
del bucket[:]
class DIFFMDAspectRatioGroupedDataset(torch.utils.data.IterableDataset):
def __init__(self, dataset, batch_sizes, num_datasets):
"""
"""
self.dataset = dataset
self.batch_sizes = batch_sizes
self._buckets = [[] for _ in range(2 * num_datasets)]
def __iter__(self):
for d in self.dataset:
w, h = d["width"], d["height"]
aspect_ratio_bucket_id = 0 if w > h else 1
bucket_id = d['dataset_source'] * 2 + aspect_ratio_bucket_id
bucket = self._buckets[bucket_id]
bucket.append(d)
if len(bucket) == self.batch_sizes[d['dataset_source']]:
yield bucket[:]
del bucket[:]
def repeat_factors_from_tag_frequency(dataset_dicts, repeat_thresh):
"""
"""
category_freq = defaultdict(int)
for dataset_dict in dataset_dicts:
cat_ids = dataset_dict['pos_category_ids']
for cat_id in cat_ids:
category_freq[cat_id] += 1
num_images = len(dataset_dicts)
for k, v in category_freq.items():
category_freq[k] = v / num_images
category_rep = {
cat_id: max(1.0, math.sqrt(repeat_thresh / cat_freq))
for cat_id, cat_freq in category_freq.items()
}
rep_factors = []
for dataset_dict in dataset_dicts:
cat_ids = dataset_dict['pos_category_ids']
rep_factor = max({category_rep[cat_id] for cat_id in cat_ids}, default=1.0)
rep_factors.append(rep_factor)
return torch.tensor(rep_factors, dtype=torch.float32)