Как узнать, вызовут ли данные, переданные на GPU, CUDA из памяти или нет

Aug 17 2020

Я использую графический процессор для запуска некоторых очень больших моделей глубокого обучения, когда я выбираю размер пакета 8, он может поместиться в памяти, но если я использую размер пакета 16, это вызовет ошибку CUDA из-за нехватки памяти, и я должен убить процесс.

У меня вопрос, прежде чем передавать данные в GPU, могу ли я узнать, насколько большие данные будут занимать в GPU?

Например, следующий код рассказывает о том, как я создаю загрузчик данных pytorch и передаю каждый пакет загрузчика данных в графический процессор, могу ли я узнать, насколько он велик, прежде чем я позвоню 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)

Ответы

3 David Aug 17 2020 at 15:36

Я бы порекомендовал использовать здесь пакет torchsummary.

pip install torchsummary

и в использовании

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

Это даст вам размер модели, размер прямого прохода и размер обратного прохода в МБ для размера пакета 1, и затем вы можете умножить его на размер пакета.