File size: 12,021 Bytes
a01ef8c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
#!/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