PyTorch Mask 50% image during training process with medium boxes

I would like to mask 50% image randomly with medium boxes during training a network. Not more than two boxes are at same position. I have an idea but it is little bit hard how to implement during training loop.

Code Idea

  1. you have image_width, image_height, and image_channels and image size is 224
  2. So we define a range: boxes=np.arange(49)
  3. batch.size() (100, 224, 224, 3)
  4. So we resize it to batchview=batch.view(-1, 224*7, 7, 3)
  5. batch.size() (100, 224, 224, 3) change to batchview=batch.view(batch_size, 7*7, -1, 3)

Input

enter image description here

Output

enter image description here

Training loop

def train(train_loader, model, criterion, optimizer, epoch, args):
batch_time = AverageMeter('Time', ':6.3f')
data_time = AverageMeter('Data', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
top5 = AverageMeter('Acc@5', ':6.2f')
progress = ProgressMeter(
    len(train_loader),
    [batch_time, data_time, losses, top1, top5],
    prefix="Epoch: [{}]".format(epoch))

# switch to train mode
model.train()

end = time.time()
for i, data in enumerate(train_loader):
    # measure data loading time
    data_time.update(time.time() - end)
    images = data['image']
    target = data['label']

    if torch.cuda.is_available():
        images = images.cuda(args.gpu, non_blocking=True)
        target = target.cuda(args.gpu, non_blocking=True)