预测 mnist

run_predict(finetune_dir, test_length)

运行预测步骤。

参数

名称	类型	描述	默认值
finetune_dir	str	包含上一步的目录	必需
test_length	int	测试步骤的长度。	必需

返回值

名称	类型	描述
tensor		模型的输出。

源代码位于 bionemo/example_model/training_scripts/predict_mnist.py

def run_predict(finetune_dir: str, test_length: int):
    """Run the prediction step.

    Args:
        finetune_dir: The directory with the previous step
        test_length: The length of the test step.

    Returns:
        tensor: the outputs of the model.
    """
    strategy = nl.MegatronStrategy(
        tensor_model_parallel_size=1,
        pipeline_model_parallel_size=1,
        ddp="megatron",
        find_unused_parameters=True,
        always_save_context=True,
    )

    test_run_trainer = nl.Trainer(
        accelerator="gpu",
        devices=1,
        strategy=strategy,
        num_nodes=1,
        plugins=nl.MegatronMixedPrecision(precision="bf16-mixed"),
    )

    lightning_module3 = BionemoLightningModule(config=ExampleFineTuneConfig(initial_ckpt_path=finetune_dir))
    new_data_module = MNISTDataModule(data_dir=str(BIONEMO_CACHE_DIR), batch_size=test_length, output_log=False)

    results = test_run_trainer.predict(lightning_module3, datamodule=new_data_module)
    return results