Как рассчитать потребность Берта в памяти?

Мне интересно узнать об использовании памяти трансформаторами. Я хотел бы использовать предварительно обученную модель для преобразования текста и сохранения вывода токена [CLS]. Никакого обучения, только умозаключение.

Мой вклад в bert составляет 511 жетонов. При размере пакета 16 моему коду не хватает памяти. Графический процессор имеет 32 ГБ памяти. У меня вопрос, как оценить использование памяти Бертом.

Как ни странно, другая работа с размером пакета 32 успешно завершилась с той же настройкой. Мой код указан ниже.

# Create dataloader
bs = 16
train_comb = ConcatDataset([train_data, valid_data])
train_dl = DataLoader(train_comb, sampler=RandomSampler(train_data), batch_size=bs)

model = BertModel.from_pretrained('/my_dir/bert_base_uncased/',
                                  output_attentions=False,
                                  output_hidden_states=False)
model.cuda()
out_list = []
model.eval()
with torch.no_grad():
    for d in train_dl:
        d = [i.cuda() for i in d].      # d = [input_ids, attention_mask, token_type_ids, labels] 
        inputs, labels = d[:3], d[3]    # input_ids has shape 16 x 511
        output = model(*inputs)[0][:, 0, :]
        out_list.append(output)

outputs = torch.cat(out_list)

Позже я изменил цикл for на ниже

with torch.no_grad():
    for d in train_dl:
        d = [i.cuda() for i in d[:3]]          # don't care about the labels
        out_list.append(model(*d)[0][:, 0, :]) # remove the intermediary variables
    del d

Подводя итог, мои вопросы:

  1. Как оценить использование памяти Бертом? Я хочу использовать его, чтобы оценить размер партии.
  2. Моя вторая работа с размером партии 32 завершилась успешно. Это потому, что у него больше набивок?
  3. Есть ли предложения по повышению эффективности использования памяти в моем коде?

person Lei Hao    schedule 24.07.2020    source источник


Ответы (1)


После некоторого поиска выясняется, что ошибка была вызвана добавлением вывода в список в GPU. Со следующим кодом ошибка исчезла.

with torch.no_grad():
    for d in train_dl:
        d = [i.cuda() for i in d[:3]]          
        out_list.append(model(*d)[0][:, 0, :].cpu()) 
    del d

Без .cpu () память продолжает увеличиваться

Tensor size: torch.Size([4, 511]), Memory allocated: 418.7685546875MB
Tensor size: torch.Size([4, 768]), Memory allocated: 424.7568359375MB


Tensor size: torch.Size([4, 511]), Memory allocated: 424.7568359375MB
Tensor size: torch.Size([4, 768]), Memory allocated: 430.7451171875MB


Tensor size: torch.Size([4, 511]), Memory allocated: 430.7451171875MB
Tensor size: torch.Size([4, 768]), Memory allocated: 436.7333984375MB

С .cpu () память не меняется.

Tensor size: torch.Size([128, 511]), Memory allocated: 420.21875MB
Tensor size: torch.Size([128, 768]), Memory allocated: 420.21875MB


Tensor size: torch.Size([128, 511]), Memory allocated: 420.21875MB
Tensor size: torch.Size([128, 768]), Memory allocated: 420.21875MB


Tensor size: torch.Size([128, 511]), Memory allocated: 420.21875MB
Tensor size: torch.Size([128, 768]), Memory allocated: 420.21875MB
person Lei Hao    schedule 25.07.2020