def make_data_sampler(dataset, shuffle, distributed):
if distributed:
return samplers.DistributedSampler(dataset, shuffle=shuffle)
if shuffle:
sampler = torch.utils.data.sampler.RandomSampler(dataset)
else:
sampler = torch.utils.data.sampler.SequentialSampler(dataset)
return sampler
def make_train_data_loader(is_distributed=False,
start_iter=0,
image_size_divisibility=32,
num_workers=1,
batch_size=2,
num_iters=90000,
shuffle=True,
single_block=False):
num_gpus = get_world_size()
assert batch_size % num_gpus == 0, "batch_size ({}) must be divisible by the number "
"of GPUs ({}) used.".format(batch_size, num_gpus)
images_per_gpu = batch_size // 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"
)
transforms = build_transforms(is_train=True)
dataset = Scarlet300MaskDataset('train', transforms=transforms)
if is_distributed:
sampler = samplers.DistributedSampler(dataset, shuffle=shuffle)
else:
sampler = torch.utils.data.sampler.RandomSampler(dataset)
batch_sampler = samplers.IterationBasedBatchSampler(
torch.utils.data.sampler.BatchSampler(sampler,
images_per_gpu,
drop_last=False), num_iters,
start_iter)
return torch.utils.data.DataLoader(
dataset,
num_workers=num_workers,
batch_sampler=batch_sampler,
collate_fn=TTT.BatchCollator(),
)
return data_loader