Efficient metrics evaluation in Training and Validation Loop

nlp
1

I am new to PyTorch and want to efficiently evaluate among others F1 during my Training and my Validation Loop.

So far, my approach was to calculate the predictions on GPU, then push them to CPU and append them to a vector for both Training and Validation. After Training and Validation, I would evaluate both for each epoch using sklearn. However, profiling my code it showed, that pushing to cpu is quite a bottleneck.

for epoch in range(n_epochs):
    model.train()
    avg_loss = 0
    avg_val_loss = 0
    train_pred = np.array([])
    val_pred = np.array([])
    # Training loop (transpose X_batch to fit pretrained (features, samples) style)
    for X_batch, y_batch in train_loader:
        scores = model(X_batch)
        y_pred = F.softmax(scores, dim=1)
        train_pred = np.append(train_pred, self.get_vector(y_pred.detach().cpu().numpy()))

        loss = loss_fn(scores, self.get_vector(y_batch))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        avg_loss += loss.item() / len(train_loader)

    model.eval()
    # Validation loop
    for X_batch, y_batch in val_loader:
        with torch.no_grad():
            scores = model(X_batch)
            y_pred = F.softmax(scores, dim=1)
            val_pred = np.append(val_pred, self.get_vector(y_pred.detach().cpu().numpy()))
            loss = loss_fn(scores, self.get_vector(y_batch))
            avg_val_loss += loss.item() / len(val_loader)

    # Model Checkpoint for best validation f1
    val_f1 = self.calculate_metrics(train_targets[val_index], val_pred, f1_only=True)
    if val_f1 > best_val_f1:
        prev_best_val_f1 = best_val_f1
        best_val_f1 = val_f1
        torch.save(model.state_dict(), self.PATHS['xlm'])
        evaluated_epoch = epoch

    # Calc the metrics
    self.save_metrics(train_targets[train_index], train_pred, avg_loss, 'train')
    self.save_metrics(train_targets[val_index], val_pred, avg_val_loss, 'val')

I am certain there is a more efficient way to a) store the predictions without having to push them to cpu each batch. b) calculate the metrics on GPU directly?

I also posted this question on stackexchange, but thought that this forum is probably even better :slight_smile:
As I am new to PyTorch, I am very grateful for any hints and feedback :slight_smile: