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.
Hengwei Chen is a new contributor to this site. Take care in asking for clarification, commenting, and answering.
Check out our Code of Conduct.