Функция потери в обучении нескольких графических процессоров (PyTorch)

Question

Я использую Pytorch и BERT для отслеживания модели. Everithing прекрасно работает на одном GPU, но когда я пытаюсь использовать несколько GPU, я получаю сообщение об ошибке:

ValueError                                Traceback (most recent call last)
<ipython-input-168-507223f9879c> in <module>()
     92         # single value; the `.item()` function just returns the Python value
     93         # from the tensor.
---> 94         total_loss += loss.item()
     95 
     96         # Perform a backward pass to calculate the gradients.

ValueError: only one element tensors can be converted to Python scalars

Может ли кто-нибудь помочь мне в том, что мне не хватает и как мне это исправить?

Вот мой код для обучения:

import random
seed_val = 42
random.seed(seed_val)
np.random.seed(seed_val)
torch.manual_seed(seed_val)
torch.cuda.manual_seed_all(seed_val)
loss_values = []
for epoch_i in range(0, epochs):

    t0 = time.time()
    total_loss = 0
    for step, batch in enumerate(train_dataloader):
        if step % 40 == 0 and not step == 0:
            elapsed = format_time(time.time() - t0)

        b_input_ids = batch[0].to(device).long() 
        b_input_mask = batch[1].to(device).long()
        b_labels = batch[2].to(device).long()
        model.zero_grad()        
        outputs = model(b_input_ids, 
                    token_type_ids=None, 
                    attention_mask=b_input_mask, 
                    labels=b_labels)

        loss = outputs[0]
        total_loss += loss.item()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
    avg_train_loss = total_loss / len(train_dataloader)            
    loss_values.append(avg_train_loss)
    print("")
    print("  Average training loss: {0:.2f}".format(avg_train_loss))
    print("  Training epcoh took: {:}".format(format_time(time.time() - t0)))

А вот мой код для модели:

from transformers import BertForSequenceClassification, AdamW, BertConfig
model_to_parallel = BertForSequenceClassification.from_pretrained(
    "./bert_cache.zip", 
    num_labels = 2, 
    output_attentions = False,
    output_hidden_states = False,
)
model = nn.DataParallel(model_to_parallel,  device_ids=[0,1,2,3]) 
model.to(device) 

Answer 1

После потери loss = outputs[0] loss является многоэлементным тензором, размер - количество графических процессоров.

Вместо него можно использовать loss = outputs[0].mean().