tests/tools/cli/test_eval_cli.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0
#

import os
import pytest
import shutil
import tempfile

from click.testing import CliRunner
from pathlib import Path
from unittest.mock import MagicMock, patch
from tlt.tools.cli.commands.eval import eval
from tlt.utils.types import FrameworkType


@pytest.mark.common
@pytest.mark.parametrize('model_name,framework',
                         [['efficientnet_b0', FrameworkType.TENSORFLOW],
                          ['resnet50', FrameworkType.PYTORCH]])
@patch("tlt.models.model_factory.get_model")
@patch("tlt.datasets.dataset_factory.load_dataset")
@patch("inspect.getfullargspec")
def test_eval_preprocess_with_image_size(mock_inspect, mock_load_dataset, mock_get_model, model_name, framework):
    """
    Tests the eval command with a dataset preprocessing method that has an image_size. Actual calls for the model and
    dataset are mocked out. The test verifies that the proper args are used for calling preprocess()
    """
    runner = CliRunner()

    tmp_dir = tempfile.mkdtemp()
    dataset_dir = os.path.join(tmp_dir, 'data')
    model_dir = os.path.join(tmp_dir, 'model')
    dummy_image_size = 100

    try:
        for new_dir in [model_dir, dataset_dir]:
            os.makedirs(new_dir)

        # Create dummy model file
        if framework == FrameworkType.TENSORFLOW:
            Path(os.path.join(model_dir, 'saved_model.pb')).touch()
        elif framework == FrameworkType.PYTORCH:
            Path(os.path.join(model_dir, 'model.pt')).touch()

        model_mock = MagicMock()
        model_mock.image_size = dummy_image_size
        data_mock = MagicMock()

        # Test where the preprocessing command will have an image size
        inspect_mock = MagicMock()
        inspect_mock.args = ['image_size', 'batch_size']
        mock_inspect.return_value = inspect_mock
        mock_get_model.return_value = model_mock
        mock_load_dataset.return_value = data_mock

        # Call the eval command
        result = runner.invoke(eval, ["--model-dir", model_dir, "--dataset_dir", dataset_dir])

        # Verify that the expected calls were made
        mock_load_dataset.assert_called_once_with(dataset_dir, model_mock.use_case, model_mock.framework)
        assert mock_get_model.called
        assert data_mock.shuffle_split.called
        assert model_mock.evaluate.called

        # Verify that preprocess was called with an image size
        data_mock.preprocess.assert_called_once_with(image_size=dummy_image_size, batch_size=32)

        # Verify that the eval command exit code is successful
        assert result.exit_code == 0
    finally:
        if os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)


@pytest.mark.common
@pytest.mark.parametrize('model_name,framework',
                         [['google/bert_uncased_L-10_H-128_A-2', FrameworkType.TENSORFLOW],
                          ['bert_en_uncased_L-12_H-768_A-12', FrameworkType.PYTORCH]])
@patch("tlt.models.model_factory.get_model")
@patch("tlt.datasets.dataset_factory.load_dataset")
@patch("inspect.getfullargspec")
def test_eval_preprocess_without_image_size(mock_inspect, mock_load_dataset, mock_get_model, model_name, framework):
    """
    Tests the eval command with a dataset preprocessing method that just has a batch size arg. Actual calls for the
    model and dataset are mocked out. The test verifies that the proper args are used for calling preprocess()
    """
    runner = CliRunner()

    tmp_dir = tempfile.mkdtemp()
    dataset_dir = os.path.join(tmp_dir, 'data')
    model_dir = os.path.join(tmp_dir, 'model')
    dummy_image_size = 100

    try:
        for new_dir in [model_dir, dataset_dir]:
            os.makedirs(new_dir)

        # Create dummy model file
        if framework == FrameworkType.TENSORFLOW:
            Path(os.path.join(model_dir, 'saved_model.pb')).touch()
        elif framework == FrameworkType.PYTORCH:
            Path(os.path.join(model_dir, 'model.pt')).touch()

        model_mock = MagicMock()
        model_mock.image_size = dummy_image_size
        data_mock = MagicMock()

        # Test where the preprocessing command just has a batch_size arg
        inspect_mock = MagicMock()
        inspect_mock.args = ['batch_size']

        mock_inspect.return_value = inspect_mock
        mock_get_model.return_value = model_mock
        mock_load_dataset.return_value = data_mock

        # Call the eval command
        result = runner.invoke(eval, ["--model-dir", model_dir, "--dataset_dir", dataset_dir])

        # Verify that the eval command exit code is successful
        assert result.exit_code == 0

        # Verify that the expected calls were made
        mock_load_dataset.assert_called_once_with(dataset_dir, model_mock.use_case, model_mock.framework)
        assert mock_get_model.called
        assert data_mock.shuffle_split.called
        assert model_mock.evaluate.called

        # Verify that preprocess was called with just batch size
        data_mock.preprocess.assert_called_once_with(batch_size=32)
    finally:
        if os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)


@pytest.mark.common
@pytest.mark.parametrize('provided_model_name,model_dir,expected_model_name,framework',
                         [['mymodel/name', 'model/abc', 'mymodel/name', FrameworkType.TENSORFLOW],
                          ['', 'bert_en_uncased_L-12_H-768_A-12/3', 'bert_en_uncased_L-12_H-768_A-12',
                           FrameworkType.PYTORCH],
                          ['test', 'bert_en_uncased_L-12_H-768_A-12/3', 'test',
                           FrameworkType.PYTORCH]
                          ])
@patch("tlt.models.model_factory.get_model")
@patch("tlt.datasets.dataset_factory.load_dataset")
def test_eval_model_name(mock_load_dataset, mock_get_model, provided_model_name, model_dir,
                         expected_model_name, framework):
    """
    Tests the eval command with and without providing a model name to verify that when a model name is provided, that
    is what's used, and when a model name is not provided, we use the model_dir folder as the model name.
    """
    runner = CliRunner()

    tmp_dir = tempfile.mkdtemp()
    dataset_dir = os.path.join(tmp_dir, 'data')
    model_dir = os.path.join(tmp_dir, model_dir)

    try:
        for new_dir in [model_dir, dataset_dir]:
            os.makedirs(new_dir)

        # Create dummy model file
        if framework == FrameworkType.TENSORFLOW:
            Path(os.path.join(model_dir, 'saved_model.pb')).touch()
        elif framework == FrameworkType.PYTORCH:
            Path(os.path.join(model_dir, 'model.pt')).touch()

        model_mock = MagicMock()
        data_mock = MagicMock()

        mock_get_model.return_value = model_mock
        mock_load_dataset.return_value = data_mock

        # Call the eval command
        eval_params = ["--model-dir", model_dir, "--dataset_dir", dataset_dir]

        if provided_model_name:
            eval_params += ["--model-name", provided_model_name]
        result = runner.invoke(eval, eval_params)

        # Verify that the expected calls were made
        mock_get_model.assert_called_once_with(expected_model_name, framework)
        mock_load_dataset.assert_called_once_with(dataset_dir, model_mock.use_case, model_mock.framework)
        assert model_mock.evaluate.called

        # Verify that the eval command exit code is successful
        assert result.exit_code == 0
    finally:
        if os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)


@pytest.mark.common
@pytest.mark.parametrize('model_name,framework,dataset_name,dataset_catalog',
                         [['efficientnet_b0', FrameworkType.TENSORFLOW, 'tf_flowers', 'tf_datasets'],
                          ['resnet50', FrameworkType.PYTORCH, 'cifar10', 'torchvision']])
@patch("tlt.models.model_factory.get_model")
@patch("tlt.datasets.dataset_factory.get_dataset")
def test_eval_dataset_catalog(mock_get_dataset, mock_get_model, model_name, framework, dataset_name, dataset_catalog):
    """
    Tests the eval command a named dataset and verifies that get_dataset is called (vs load_dataset, which is used
    for custom dataset directories in other tests).
    """
    runner = CliRunner()

    tmp_dir = tempfile.mkdtemp()
    dataset_dir = os.path.join(tmp_dir, 'data')
    model_dir = os.path.join(tmp_dir, 'model')

    try:
        for new_dir in [model_dir, dataset_dir]:
            os.makedirs(new_dir)

        # Create dummy model file
        if framework == FrameworkType.TENSORFLOW:
            Path(os.path.join(model_dir, 'saved_model.pb')).touch()
        elif framework == FrameworkType.PYTORCH:
            Path(os.path.join(model_dir, 'model.pt')).touch()

        # Setup mocks
        model_mock = MagicMock()
        data_mock = MagicMock()
        mock_get_model.return_value = model_mock
        mock_get_dataset.return_value = data_mock

        # Call the eval command
        result = runner.invoke(eval,
                               ["--model-dir", str(model_dir), "--model-name", model_name, "--dataset_dir", dataset_dir,
                                "--dataset-name", dataset_name, "--dataset-catalog", dataset_catalog])

        # Verify that the expected calls were made
        mock_get_model.assert_called_once_with(model_name, framework)
        mock_get_dataset.assert_called_once_with(dataset_dir, model_mock.use_case, model_mock.framework,
                                                 dataset_name, dataset_catalog)

        # Verify that the evaluate command exit code is successful
        assert model_mock.evaluate.called
        assert result.exit_code == 0
    finally:
        if os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)


class TestEvalArgs:
    """
    Class for tests that are testing bad inputs for evaluation
    """

    def setup_class(self):
        self._runner = CliRunner()

        self._tmp_dir = tempfile.mkdtemp()
        self._dataset_dir = os.path.join(self._tmp_dir, 'data')
        self._model_dir = os.path.join(self._tmp_dir, 'model')

    def setup_method(self):
        for new_dir in [self._model_dir, self._dataset_dir]:
            if not os.path.exists(new_dir):
                os.makedirs(new_dir)

        Path(os.path.join(self._model_dir, 'saved_model.pb')).touch()

    def teardown_method(self):
        if os.path.exists(self._tmp_dir):
            shutil.rmtree(self._tmp_dir)

    def teardown_class(self):
        if os.path.exists(self._tmp_dir):
            shutil.rmtree(self._tmp_dir)

    @pytest.mark.common
    @pytest.mark.parametrize('dataset_catalog', ['foo', 'benchmark', '0'])
    def test_eval_invalid_dataset_catalog(self, dataset_catalog):
        """
        Verifies that eval command fails if the dataset catalog value is invalid
        """

        result = self._runner.invoke(eval,
                                     ["--model-dir", self._model_dir,
                                      "--dataset-dir", self._dataset_dir,
                                      "--dataset-name", "foo",
                                      "--dataset-catalog", dataset_catalog])

        assert result.exit_code == 2
        assert "Invalid value for '--dataset-catalog'" in result.output
        assert "'{}' is not one of 'tf_datasets', 'torchvision', 'huggingface'".format(dataset_catalog) in result.output