He pasado un tiempo considerable tratando de depurar algún código de Pytorch del que he creado un ejemplo mínimo con el fin de ayudar a comprender mejor cuál podría ser el problema.

Eliminé todas las partes necesarias del código que no están relacionadas con el problema, por lo que el fragmento de código restante no tendrá mucho sentido desde un punto de vista funcional, pero aún muestra el error al que me enfrento.

La tarea general en la que estoy trabajando es en un bucle y cada paso del bucle calcula la incrustación de la imagen y la agrega a una variable que la almacena. Lo está agregando efectivamente (no concatenando, por lo que el tamaño sigue siendo el mismo). No espero que el número de iteraciones fuerce el desbordamiento del tipo de datos, no veo que esto suceda aquí ni en mi código.

  • He agregado varias métricas para evaluar el tamaño de los tensores con los que estoy trabajando para asegurarme de que no estén creciendo en la huella de memoria.
  • Estoy comprobando el uso general de la memoria de la GPU para verificar el problema que condujo al RuntimeError: CUDA out of memory. final.

Mi entorno es el siguiente:

 - python 3.6.2
 - Pytorch 1.4.0
 - Cudatoolkit 10.0
 - Driver version 410.78
 - GPU: Nvidia GeForce GT 1030  (2GB VRAM)   
(though I've replicated this experiment with the same result on a Titan RTX with 24GB,
same pytorch version and cuda toolkit and driver, it only goes out of memory further in the loop).

Complete el código a continuación. He marcado 2 líneas como culpables, ya que eliminarlas elimina el problema, aunque obviamente necesito encontrar una manera de ejecutarlas sin tener problemas de memoria. ¡Cualquier ayuda será muy apreciada! Puede intentar con cualquier imagen llamada "source_image.bmp" para replicar el problema.

import torch
from PIL import Image
import torchvision
from torchvision import transforms
from pynvml import nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo, nvmlInit
import sys
import os
os.environ["CUDA_VISIBLE_DEVICES"]='0'      # this is necessary on my system to allow the environment to recognize my nvidia GPU for some reason
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'    # to debug by having all CUDA functions executed in place
torch.set_default_tensor_type('torch.cuda.FloatTensor')

# Preprocess image
tfms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224), 
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),])
img = tfms(Image.open('source_image.bmp')).unsqueeze(0).cuda()

model = torchvision.models.resnet50(pretrained=True).cuda()
model.eval()    # we put the model in evaluation mode, to prevent storage of gradient which might accumulate

nvmlInit()
h = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(h)
print(f'Total available memory   : {info.total / 1000000000}')

feature_extractor = torch.nn.Sequential(*list(model.children())[:-1])
orig_embedding = feature_extractor(img)

embedding_depth = 2048

mem0 = 0

embedding = torch.zeros(2048, img.shape[2], img.shape[3]) #, dtype=torch.float)

patch_size=[4,4]
patch_stride=[2,2]
patch_value=0.0

# Here, we iterate over the patch placement, defined at the top left location
for row in range(img.shape[2]-1):
    for col in range(img.shape[3]-1):
        print("######################################################")        
                
        ######################################################
        # Isolated line, culprit 1 of the GPU memory leak
        ######################################################
        patched_embedding = feature_extractor(img)
        
        delta_embedding = (patched_embedding - orig_embedding).view(-1, 1, 1)
        
        ######################################################
        # Isolated line, culprit 2 of the GPU memory leak
        ######################################################
        embedding[:,row:row+1,col:col+1] = torch.add(embedding[:,row:row+1,col:col+1], delta_embedding)

        print("img size:\t\t", img.element_size() * img.nelement())
        print("patched_embedding size:\t", patched_embedding.element_size() * patched_embedding.nelement())
        print("delta_embedding size:\t", delta_embedding.element_size() * delta_embedding.nelement())
        print("Embedding size:\t\t", embedding.element_size() * embedding.nelement())

        del patched_embedding, delta_embedding
        torch.cuda.empty_cache()
        
        info = nvmlDeviceGetMemoryInfo(h)
        print("\nMem usage increase:\t", info.used / 1000000000 - mem0)
        mem0 = info.used / 1000000000
        print(f'Free:\t\t\t {(info.total - info.used) / 1000000000}')

print("Done.")
1
Alex U. 26 ago. 2020 a las 01:39

1 respuesta

La mejor respuesta

Agregue esto a su código tan pronto como cargue el modelo

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

De https://pytorch.org/docs/stable /notes/autograd.html#excluyendo-subgraphs-from-backward

0
183.amir 27 ago. 2020 a las 12:31