Cómo saber si los datos pasados ​​a la GPU harán que CUDA se quede sin memoria o no

Question

Estoy usando GPU para ejecutar algunos modelos de aprendizaje profundo muy grandes, cuando elijo un tamaño de lote de 8, puede caber en la memoria, pero si uso un tamaño de lote de 16, causará un error de memoria insuficiente de CUDA, y tengo que matar el proceso.

Mi pregunta es, antes de pasar los datos a la GPU, ¿hay alguna manera de saber qué tamaño ocuparán los datos en la GPU?

Por ejemplo, el siguiente código trata sobre cómo creo un cargador de datos pytorch y paso cada lote del cargador de datos a la GPU, ¿podría saber qué tan grande es antes de llamar?batch.to(device)

train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)
for step, batch in enumerate(train_dataloader):
    b_input_ids = batch[0].to(device)
    b_input_mask = batch[1].to(device)
    b_labels = batch[2].to(device)

Answer 1

Recomendaría usar el paquete torchsummary aquí.

pip install torchsummary

y en uso

from torchsummary import summary
myModel.cuda()
summary(myModel, (shapeOfInput)) # where shapeOfInput is a tuple of the sample's dimensions

Esto le dará el tamaño del modelo, el tamaño del pase hacia adelante y el tamaño del pase hacia atrás en MB para un tamaño de lote de 1, y luego puede multiplicar por el tamaño de su lote.