After I trained a deep neural network (DNN) that takes 6 images as input, I want to evaluate the performance of the model by adding perturbation to the input images. In this case, using gradient of loss function with respect to the input images is an effective method to generate perturbation.

But the trouble is coming.

Simply speaking, the NVIDIA GPU will out of memory if I set the input_images.requires_grad=True. The shape of input_images is torch.Size([1, 6, 3, 928, 1600]), where 1 stands for batch_size, 6 for number of cameras, 3 for image channels, and (928, 1600) for width and height of the image respectively. By the way, the input images is sampled from nuScenes dataset.

The specification of Remote Server below is what I am using.

OS: Ubuntu 22.04.4 LTS
GPU: NVIDIA GeForce RTX 4090 / 24GB Memory
Totally 4 GPU

First, I would paste the output information from the terminal.

  File "/home/server/my_space/src/attack_eval.py", line 184, in main
    predict_labels_vox, predict_labels_pts = my_model(img=imgs,

  ...

  File "/home/server/my_space/modules/image_cross_attention.py", line 389, in forward
    sampling_offsets, attention_weights = self.get_sampling_offsets_and_attention(query)
  File "/home/server/my_space/modules/image_cross_attention.py", line 309, in get_sampling_offsets_and_attention
    attns = torch.cat(attns, dim=1)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 72.00 MiB (GPU 0; 23.65 GiB total capacity; 22.53 GiB already allocated; 13.00 MiB free; 22.71 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

From the terminal’s output, we can see that the cuda had out of memory during forwarding process. In the below, I will also paste the associated codes:

my_model.eval()

for param in my_model.parameters():
    param.requires_grad = False

torch.cuda.empty_cache()

for i_iter_val, (imgs, img_metas, val_vox_label, val_grid, val_pt_labs) in enumerate(val_dataset_loader):
    imgs = imgs.cuda()
    imgs.requires_grad = True

    ### This line caue the error above ###
    predict_labels_vox, predict_labels_pts = my_model(img=imgs,
    ...

    loss = ...
    my_model.zero_grad()
    loss.backward()
    ...

    data_grad = input_images.grad.data 
    
    # Add the grad to the images
    ...

In the above codes, I close the flag of gradient of my_model.parameters() to reduce the ultilize of CUDA memory. The total number of parameters of my_model is about 67, 000, 000. By the way, during training stage, the code can be execuated successfully, and the gradient of parameters can be obtained. And the total pixel of input images is about 26, 700, 000.

Here is the question: If the GPU memory can support training all the parameters (about 0.067B), why can not be able to support calculating the gradient of all input images, that is all pixels (about 0.026B)? Note: During the training stage, I use multi-gpu, so did the eval stage.

There are some solutions which I had tried.

Solution 1: Reduce the resolution of input images. (Successful Run, But not recommoned)

Let’s dive into the code.

def downsample_images(images, scale_factor=0.5):
    # images: [1, 6, 3, 928, 1600]
    batch_size, time_steps, channels, height, width = images.shape
    new_height = int(height * scale_factor)
    new_width = int(width * scale_factor)
    images = images.view(batch_size * time_steps, channels, height, width)
    downsampled_images = F.interpolate(images, size=(new_height, new_width), mode='bilinear', align_corners=False)
    downsampled_images = downsampled_images.view(batch_size, time_steps, channels, new_height, new_width)
    return downsampled_images


imgs = downsample_images(imgs, scale_factor=0.5)

... = my_model(output) 

In this way, the code can executate the model forwarding successfully, but it will also influence the origin performance of neural network. So I am trying to come up with other thoughts.

Solution 2: Splitting the input images. (Failed. Shallow attempt, but the correctness should be verified.)

attack_image = imgs[:, 5:, :, :, :]
attack_image.requires_grad = True

imgs = imgs[:, :5, :, :, :]
imgs.requires_grad = False


class my_model(nn.Module):
    ...

    def forward(imgs, attack_imgs):
        img = torch.cat((img, attack_images), dim=1)
        ...


... = my_model(imgs, attack_imgs)

Solution 3: Partial elements of the tensor are set to True. (Failed.)

imgs[:, 5:, :, :, :].requires_grad = True
... = my_model(imgs)

Based on the above thoughts, I still cannot deal with this problem. So I am looking for help in the community. If you need more details to propose a solution, please let me know. Thanks in advance.