Predict mnist

run_predict(finetune_dir, test_length)

Run the prediction step.

Parameters:

Name	Type	Description	Default
finetune_dir	str	The directory with the previous step	required
test_length	int	The length of the test step.	required

Returns:

Name	Type	Description
tensor		the outputs of the model.

Source code in 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