Comment savoir si les données transmises au GPU entraîneront ou non CUDA à court de mémoire

Question

J'utilise le GPU pour exécuter de très grands modèles d'apprentissage en profondeur, lorsque je choisis une taille de lot de 8, il peut tenir dans la mémoire, mais si j'utilise une taille de lot de 16, cela provoquera une erreur de mémoire insuffisante CUDA, et je dois tuer le processus.

Ma question est la suivante : avant de transmettre réellement les données au GPU, existe-t-il un moyen de savoir quelle taille les données occuperont dans le GPU ?

Par exemple, le code suivant explique comment je crée un chargeur de données pytorch et transmet chaque lot du chargeur de données au GPU, puis-je savoir quelle est sa taille avant d'appelerbatch.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

Je recommanderais d'utiliser le paquet torchummary ici.

pip install torchsummary

et en cours d'utilisation

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

Cela vous donnera la taille du modèle, la taille de la passe avant et la taille de la passe arrière en Mo pour une taille de lot de 1, et vous pouvez ensuite multiplier par la taille de votre lot.