tests/shell/test_basic_commands.py

import pytest

from tests.helpers.run_command import run_command
from tests.helpers.runif import RunIf

"""
A couple of sanity checks to make sure the model doesn't crash with different running options.
"""


def test_fast_dev_run():
    """Test running for 1 train, val and test batch."""
    command = ["train.py", "++trainer.fast_dev_run=true"]
    run_command(command)


@pytest.mark.slow
def test_cpu():
    """Test running 1 epoch on CPU."""
    command = ["train.py", "++trainer.max_epochs=1", "++trainer.gpus=0"]
    run_command(command)


# use RunIf to skip execution of some tests, e.g. when no gpus are available
@RunIf(min_gpus=1)
@pytest.mark.slow
def test_gpu():
    """Test running 1 epoch on GPU."""
    command = [
        "train.py",
        "++trainer.max_epochs=1",
        "++trainer.gpus=1",
    ]
    run_command(command)


@RunIf(min_gpus=1)
@pytest.mark.slow
def test_mixed_precision():
    """Test running 1 epoch with pytorch native automatic mixed precision (AMP)."""
    command = [
        "train.py",
        "++trainer.max_epochs=1",
        "++trainer.gpus=1",
        "++trainer.precision=16",
    ]
    run_command(command)


@pytest.mark.slow
def test_double_validation_loop():
    """Test running 1 epoch with validation loop twice per epoch."""
    command = [
        "train.py",
        "++trainer.max_epochs=1",
        "++trainer.val_check_interval=0.5",
    ]
    run_command(command)