Arbitrary batch sizes after passing dataloader through privacy engine

In my code, I have defined privacy engine as follows: -

unet, optimizer, trainloader = privacy_engine.make_private_with_epsilon(
        module = unet,
        data_loader = trainloader,
        optimizer = optimizer,
        epochs = global_epochs * local_epochs,
        target_epsilon = 1.0,
        batch_first = True,
        target_delta = PRIVACY_PARAMS['target_delta'],
        max_grad_norm = PRIVACY_PARAMS['max_grad_norm'])

I sampled a batch from the dataloader before and after passing through the privacy engine and every time I execute it, the wrapped dataloader shows some arbitrary batch size. For example, if my original dataloader was made with batch size =6, the wrapped dataloader shows 7, 10, etc, which changes with every execution. Why does this happen?

Hello @Anirban_Nath,
This is a consequence of using Poisson sampling that is needed for a Differentially Private Data Loader. This has been called out in the doc strings here (Opacus · Train PyTorch models with Differential Privacy). I’ll make sure this also captured in our FAQs at FAQ · Opacus

[Tracking the update to documentation at Add an FAQ about variable batch size · Issue #514 · pytorch/opacus · GitHub]

Hello @karthikprasad

So I figured that out, thanks to you and I also found out that there is a module named BatchMemoryManager that helps to keep maximum batch size in check. However, with the two in place, when I am trying to execute my code, I keep getting a “ValueError: Per sample gradient is not initialized. Not updated in backward pass?” error at loss.backward().

Prior to that, there is a warning when I pass a batch through my model which says “Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.”

Maybe the two are correlated?

Here is a snippet of my code for your reference: -

privacy_engine = PrivacyEngine()
model , optimizer, trainloader = privacy_engine.make_private_with_epsilon(
        module = unet,
        data_loader = trainloader,
        optimizer = optimizer,
        epochs = 200,
        target_epsilon = 1.0,
        batch_first = True,
        target_delta = PRIVACY_PARAMS['target_delta'],
        max_grad_norm = PRIVACY_PARAMS['max_grad_norm'])

print(f"Training local epochs")
for l in range(local_epochs):
    with BatchMemoryManager(data_loader=trainloader, max_physical_batch_size=4, optimizer=optimizer) as new_data_loader:   
    for i_batch, sampled_batch in enumerate(new_data_loader):

        image_batch, segmask = sampled_batch['image'], sampled_batch['mask']
        image_batch, segmask  = image_batch.to(device), segmask.to(device)

        #Segmentation
        outs_seg = unet(image_batch, task = 'segment') #I get warning here
        softmax = torch.nn.functional.log_softmax(outs_seg, dim=1)
        loss_seg = torch.nn.functional.nll_loss(softmax, segmask[:].long())

        optimizer.zero_grad()
        loss.backward() #I get the error here
        optimizer.step()