wehub-resource-sync
73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# 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.
|
|
#
|
|
|
|
|
|
"""Tests of the classification flow"""
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from os import path
|
|
import glob
|
|
import pytest
|
|
# pylint:disable=missing-docstring, no-self-use
|
|
|
|
class TestClassificationFlow():
|
|
|
|
def test_resnet18(self, request, pytestconfig):
|
|
dir_path = os.path.dirname(os.path.realpath(__file__))
|
|
dataset_dir = pytestconfig.getoption('--data-dir')
|
|
|
|
# skip if the data dir flag was not set
|
|
if not dataset_dir:
|
|
pytest.skip("Prepare required dataset and use --data-dir option to enable")
|
|
|
|
# Verify data dir exists
|
|
if not path.exists(dataset_dir):
|
|
print("Dataset path %s doesn't exist"%(dataset_dir), file=sys.stderr)
|
|
assert path.exists(dataset_dir)
|
|
|
|
# Append required paths to PYTHONPATH
|
|
test_env = os.environ.copy()
|
|
if 'PYTHONPATH' not in test_env:
|
|
test_env['PYTHONPATH'] = ""
|
|
|
|
# Add project root and torchvision to the path (assuming running in nvcr.io/nvidia/pytorch:20.08-py3)
|
|
test_env['PYTHONPATH'] += ":/opt/pytorch/vision/references/classification/:%s/../"%(dir_path)
|
|
|
|
# Add requirement egg files manually to path since we're spawning a new process (downloaded by setuptools)
|
|
for egg in glob.glob(dir_path + "/../.eggs/*.egg"):
|
|
test_env['PYTHONPATH'] += ":%s"%(egg)
|
|
|
|
# Run in a subprocess to avoid contaminating the module state for other test cases
|
|
ret = subprocess.run(
|
|
[
|
|
'python3', dir_path + '/../examples/torchvision/classification_flow.py',
|
|
'--data-dir', dataset_dir,
|
|
'--model', 'resnet18', '--pretrained',
|
|
'-t', '0.5',
|
|
'--num-finetune-epochs', '2',
|
|
'--evaluate-onnx',
|
|
],
|
|
env=test_env,
|
|
check=False, stdout=subprocess.PIPE)
|
|
|
|
# If the test failed dump the output to stderr for better logging
|
|
if ret.returncode != 0:
|
|
print(ret.stdout, file=sys.stderr)
|
|
|
|
assert ret.returncode == 0
|