测试工具

compute_biobert_loss_singlegpu(trainer, pl_module)

计算单 GPU 上 BioBert 模型的损失。

这在多 GPU 设置或不符合 BioBert 模型的模型中不起作用。

参数

名称	类型	描述	默认值
trainer	Trainer	Lightning Trainer 对象。	必需
pl_module	LightningModule	正在训练的 LightningModule。	必需

返回

名称	类型	描述
float		平均损失。

另请参阅： - :class: BioBertModel

源代码位于 bionemo/llm/model/biobert/testing_utils.py

def compute_biobert_loss_singlegpu(trainer: pl.Trainer, pl_module: pl.LightningModule):
    """Computes the loss for BioBert models on a single GPU.

    This will not function in multi-gpu settings nor with models that do not conform to BioBert.

    Args:
        trainer (pl.Trainer): The Lightning Trainer object.
        pl_module (pl.LightningModule): The LightningModule being trained.

    Returns:
        float: The mean loss.

    See Also:
    - :class: BioBertModel
    """
    model = pl_module
    dl = trainer.datamodule.val_dataloader()

    n, loss = -1, 0.0
    model.eval()
    # batch = next(iter(dl))
    batch = model.data_step(iter(dl))
    result = model(
        input_ids=batch["text"].cuda(),  # 'tokens' also a valid input for MockGPTDataModule
        attention_mask=batch["attention_mask"].cuda(),
    )
    loss_mask = batch["loss_mask"].cuda()
    # Not guaranteed i guess?
    logits = result["token_logits"]
    target = batch["labels"].cuda()
    loss += F.cross_entropy(logits[loss_mask].float(), target[loss_mask], reduction="sum")
    n += loss_mask.sum()

    mean_loss: float = (loss / n).detach().cpu().numpy().item()
    model.train()
    return mean_loss