def main(args):
cfg = setup(args)
model = build_detection_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
if args.eval_only:
checkpointer = DetectronCheckpointer(cfg,
model,
save_dir=cfg.OUTPUT_DIR)
ckpt = cfg.MODEL.WEIGHT if args.ckpt is None else args.ckpt
_ = checkpointer.load(ckpt, use_latest=args.ckpt is None)
return run_test(cfg, model)
distributed = comm.get_world_size() > 1
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True,
)
train(cfg, model, device, distributed)
def inference(
model,
data_loader,
dataset_name,
eval_types=("detections", ),
device="cuda",
output_folder=None,
):
device = torch.device(device)
num_devices = comm.get_world_size()
logger = logging.getLogger(__name__)
dataset = data_loader.dataset
logger.info("Start evaluation on {} dataset({} images).".format(
dataset_name, len(dataset)))
total_timer = Timer()
inference_timer = Timer()
total_timer.tic()
predictions = compute_on_dataset(model, data_loader, device,
inference_timer)
# predictions = None
comm.synchronize()
total_time = total_timer.toc()
total_time_str = get_time_str(total_time)
logger.info(
"Total run time: {} ({} s / img per device, on {} devices)".format(
total_time_str, total_time * num_devices / len(dataset),
num_devices))
total_infer_time = get_time_str(inference_timer.total_time)
logger.info(
"Model inference time: {} ({} s / img per device, on {} devices)".
format(
total_infer_time,
inference_timer.total_time * num_devices / len(dataset),
num_devices,
))
predictions = _accumulate_predictions_from_multiple_gpus(predictions)
if not comm.is_main_process():
return
return evaluate(
eval_type=eval_types,
dataset=dataset,
predictions=predictions,
output_folder=output_folder,
)
def __init__(self, size: int):
"""
Args:
size (int): the total number of data of the underlying dataset to sample from
"""
self._size = size
assert size > 0
self._rank = comm.get_rank()
self._world_size = comm.get_world_size()
shard_size = (self._size - 1) // self._world_size + 1
begin = shard_size * self._rank
end = min(shard_size * (self._rank + 1), self._size)
self._local_indices = range(begin, end)
def __init__(self,
size: int,
shuffle: bool = True,
seed: Optional[int] = None):
"""
Args:
size (int): the total number of data of the underlying dataset to sample from
shuffle (bool): whether to shuffle the indices or not
seed (int): the initial seed of the shuffle. Must be the same
across all workers. If None, will use a random seed shared
among workers (require synchronization among all workers).
"""
self._size = size
assert size > 0
self._shuffle = shuffle
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()
def reduce_loss_dict(loss_dict):
"""
Reduce the loss dictionary from all processes so that process with rank
0 has the averaged results. Returns a dict with the same fields as
loss_dict, after reduction.
"""
world_size = get_world_size()
if world_size < 2:
return loss_dict
with torch.no_grad():
loss_names = []
all_losses = []
for k in sorted(loss_dict.keys()):
loss_names.append(k)
all_losses.append(loss_dict[k])
all_losses = torch.stack(all_losses, dim=0)
dist.reduce(all_losses, dst=0)
if dist.get_rank() == 0:
# only main process gets accumulated, so only divide by
# world_size in this case
all_losses /= world_size
reduced_losses = {k: v for k, v in zip(loss_names, all_losses)}
return reduced_losses
def __init__(self, dataset_dicts, repeat_thresh, shuffle=True, seed=None):
"""
Args:
dataset_dicts (list[dict]): annotations in Detectron2 dataset format.
repeat_thresh (float): frequency threshold below which data is repeated.
shuffle (bool): whether to shuffle the indices or not
seed (int): the initial seed of the shuffle. Must be the same
across all workers. If None, will use a random seed shared
among workers (require synchronization among all workers).
"""
self._shuffle = shuffle
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()
# Get fractional repeat factors and split into whole number (_int_part)
# and fractional (_frac_part) parts.
rep_factors = self._get_repeat_factors(dataset_dicts, repeat_thresh)
self._int_part = torch.trunc(rep_factors)
self._frac_part = rep_factors - self._int_part
def make_data_loader(cfg, is_train=True):
num_gpus = get_world_size()
if is_train:
images_per_batch = cfg.SOLVER.IMS_PER_BATCH
assert images_per_batch % num_gpus == 0, \
"SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of GPUs ({}) used." \
.format(images_per_batch, num_gpus)
images_per_gpu = images_per_batch // num_gpus
else:
images_per_batch = cfg.TEST.IMS_PER_BATCH
assert images_per_batch % num_gpus == 0, \
"SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of GPUs ({}) used." \
.format(images_per_batch, num_gpus)
images_per_gpu = images_per_batch // num_gpus
# if images_per_gpu > 1:
# logger = logging.getLogger(__name__)
# logger.warning(
# "When using more than one image per GPU you may encounter "
# "an out-of-memory (OOM) error if your GPU does not have "
# "sufficient memory. If this happens, you can reduce "
# "SOLVER.IMS_PER_BATCH (for training) or "
# "TEST.IMS_PER_BATCH (for inference). For training, you must "
# "also adjust the learning rate and schedule length according "
# "to the linear scaling rule. See for example: "
# "https://github.com/facebookresearch/Detectron/blob/master/configs/getting_started/tutorial_1gpu_e2e_faster_rcnn_R-50-FPN.yaml#L14"
# )
# group images which have similar aspect ratio. In this case, we only
# group in two cases: those with width / height > 1, and the other way around,
# but the code supports more general grouping strategy
aspect_grouping = [1] if cfg.DATALOADER.ASPECT_RATIO_GROUPING else []
path_catalog = import_file(
"smoke.config.paths_catalog", cfg.PATHS_CATALOG, True
)
DatasetCatalog = path_catalog.DatasetCatalog
transforms = build_transforms(cfg, is_train)
datasets = build_dataset(cfg, transforms, DatasetCatalog, is_train)
data_loaders = []
for dataset in datasets:
sampler = samplers.TrainingSampler(len(dataset))
batch_sampler = torch.utils.data.sampler.BatchSampler(
sampler, images_per_gpu, drop_last=True
)
collator = BatchCollator(cfg.DATALOADER.SIZE_DIVISIBILITY)
num_workers = cfg.DATALOADER.NUM_WORKERS
data_loader = torch.utils.data.DataLoader(
dataset,
num_workers=num_workers,
batch_sampler=batch_sampler,
collate_fn=collator,
worker_init_fn=worker_init_reset_seed,
)
data_loaders.append(data_loader)
if is_train:
# during training, a single (possibly concatenated) data_loader is returned
assert len(data_loaders) == 1
return data_loaders[0]
return data_loaders