def create_train_dataloader(configs):
"""Create dataloader for training"""
train_lidar_transforms = OneOf([
Random_Rotation(limit_angle=20., p=1.0),
Random_Scaling(scaling_range=(0.95, 1.05), p=1.0)
], p=0.66)
train_aug_transforms = Compose([
Horizontal_Flip(p=configs.hflip_prob),
Cutout(n_holes=configs.cutout_nholes, ratio=configs.cutout_ratio, fill_value=configs.cutout_fill_value,
p=configs.cutout_prob)
], p=1.)
train_dataset = KittiDataset(configs.dataset_dir, mode='train', lidar_transforms=train_lidar_transforms,
aug_transforms=train_aug_transforms, multiscale=configs.multiscale_training,
num_samples=configs.num_samples, mosaic=configs.mosaic,
random_padding=configs.random_padding)
train_sampler = None
if configs.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_dataloader = DataLoader(train_dataset, batch_size=configs.batch_size, shuffle=(train_sampler is None),
pin_memory=configs.pin_memory, num_workers=configs.num_workers, sampler=train_sampler,
collate_fn=train_dataset.collate_fn)
return train_dataloader, train_sampler
def create_val_dataloader(configs):
"""Create dataloader for validation"""
val_aug_transforms = Compose([
Horizontal_Flip(p=configs.hflip_prob),
Cutout(n_holes=configs.cutout_nholes, ratio=configs.cutout_ratio, fill_value=configs.cutout_fill_value,
p=configs.cutout_prob)
], p=1.)
val_sampler = None
val_dataset = KittiDataset(configs.dataset_dir, mode='val', lidar_transforms=None,
aug_transforms=val_aug_transforms, multiscale=False, num_samples=configs.num_samples,
mosaic=False, random_padding=False)
if configs.distributed:
val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False)
val_dataloader = DataLoader(val_dataset, batch_size=configs.batch_size, shuffle=False,
pin_memory=configs.pin_memory, num_workers=configs.num_workers, sampler=val_sampler,
collate_fn=val_dataset.collate_fn)
return val_dataloader