Skip to content

Commit

Permalink
test edits
Browse files Browse the repository at this point in the history
  • Loading branch information
horheynm committed Sep 1, 2023
1 parent 6838aa3 commit 1a0077b
Show file tree
Hide file tree
Showing 7 changed files with 123 additions and 148 deletions.
17 changes: 7 additions & 10 deletions src/sparsezoo/model/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ def __init__(self, source: str, download_path: Optional[str] = None):
] = self._sample_outputs_list_to_dict(self.sample_outputs)

self.sample_labels: Directory = self._directory_from_files(
files, directory_class=Directory, display_name="sample_labels"
files, directory_class=Directory, display_name="sample-labels"
)

self.deployment: SelectDirectory = self._directory_from_files(
Expand All @@ -150,12 +150,9 @@ def __init__(self, source: str, download_path: Optional[str] = None):

self.logs: Directory = self._directory_from_files(files, display_name="logs")

self.recipes: SelectDirectory = self._directory_from_files(
files,
directory_class=SelectDirectory,
display_name="recipe",
stub_params=self.stub_params,
)
self.recipes = self._file_from_files(files, display_name="^recipe", regex=True)
if isinstance(self.recipes, File):
self.recipes = [self.recipes]

self._onnx_gz: OnnxGz = self._directory_from_files(
files, directory_class=OnnxGz, display_name="model.onnx.tar.gz"
Expand Down Expand Up @@ -691,7 +688,7 @@ def _sample_outputs_list_to_dict(
if not isinstance(directories, list):
# if found a single 'sample_outputs' directory,
# assume it should be mapped to its the native framework
expected_name = "sample_outputs"
expected_name = "sample-outputs"
if directories.name not in [expected_name, expected_name + ".tar.gz"]:
raise ValueError(
"Found single folder (or tar.gz archive)"
Expand All @@ -701,7 +698,7 @@ def _sample_outputs_list_to_dict(
engine_to_numpydir_map["framework"] = directories

else:
# if found multiple 'sample_outputs' directories,
# if found multiple 'sample-outputs' directories,
# use directory name to relate it with the appropriate
# inference engine
for directory in directories:
Expand All @@ -710,7 +707,7 @@ def _sample_outputs_list_to_dict(
engine_name = engine_name.replace(".tar.gz", "")
if engine_name not in ENGINES:
raise ValueError(
f"The name of the 'sample_outputs' directory should "
f"The name of the 'sample-outputs' directory should "
f"end with an engine name (one of the {ENGINES}). "
f"However, the name is {directory.name}."
)
Expand Down
23 changes: 5 additions & 18 deletions src/sparsezoo/model/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,6 @@ def load_files_from_stub(

include_file_download_url(files)
files = restructure_request_json(request_json=files)

if params is not None:
files = filter_files(files=files, params=params)

Expand Down Expand Up @@ -311,7 +310,7 @@ def save_outputs_to_tar(

path = os.path.join(
os.path.dirname(sample_inputs.path),
f"sample_outputs_{engine_type}",
f"sample-outputs_{engine_type}",
)
if not os.path.exists(path):
os.mkdir(path)
Expand Down Expand Up @@ -385,26 +384,14 @@ def restructure_request_json(
file_dict_deployment["file_type"] = "deployment"
request_json.append(file_dict_deployment)

# create recipes
recipe_dicts_list = fetch_from_request_json(request_json, "file_type", "recipe")
for (idx, file_dict) in recipe_dicts_list:
display_name = file_dict["display_name"]
# make sure that recipe name has a
# format `recipe_{...}`.
prefix = "recipe_"
if not display_name.startswith(prefix):
display_name = prefix + display_name
file_dict["display_name"] = display_name
request_json[idx] = file_dict

# restructure inputs/labels/originals/outputs directories
# use `sample-inputs.tar.gz` to simulate non-existent directories

files_to_create = [
"sample_inputs.tar.gz",
"sample_labels.tar.gz",
"sample_originals.tar.gz",
"sample_outputs.tar.gz",
"sample-inputs.tar.gz",
"sample-labels.tar.gz",
"sample-originals.tar.gz",
"sample-outputs.tar.gz",
]
types = ["inputs", "labels", "originals", "outputs"]
for file_name, type in zip(files_to_create, types):
Expand Down
1 change: 1 addition & 0 deletions tests/sparsezoo/analyze/bert_pruned80_quant-none-vnni.json

Large diffs are not rendered by default.

1 change: 0 additions & 1 deletion tests/sparsezoo/analyze/bert_pruned_quantized.json

This file was deleted.

4 changes: 2 additions & 2 deletions tests/sparsezoo/analyze/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,9 +50,9 @@
"bert_pruned_quantized": {
"stub": (
"zoo:nlp/question_answering/bert-base/pytorch/huggingface/squad/"
"12layer_pruned80_quant-none-vnni"
"pruned80_quant-none-vnni"
),
"truth": f"{os.path.dirname(__file__)}/bert_pruned_quantized.json",
"truth": f"{os.path.dirname(__file__)}/bert_pruned80_quant-none-vnni.json",
},
"resnet50_pruned_quantized": {
"stub": (
Expand Down
160 changes: 77 additions & 83 deletions tests/sparsezoo/model/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
import copy
import os
import shutil
import tarfile
import tempfile
from pathlib import Path

Expand All @@ -32,15 +31,14 @@
"onnx",
"model.onnx",
"model.onnx.tar.gz",
"recipe",
"sample_inputs.tar.gz",
"sample_originals.tar.gz",
"sample_labels.tar.gz",
"sample_outputs.tar.gz",
"sample_inputs",
"sample_originals",
"sample_labels",
"sample_outputs",
"sample-inputs.tar.gz",
"sample-originals.tar.gz",
"sample-labels.tar.gz",
"sample-outputs.tar.gz",
"sample-inputs",
"sample-originals",
"sample-labels",
"sample-outputs",
"benchmarks.yaml",
"eval.yaml",
"analysis.yaml",
Expand Down Expand Up @@ -116,9 +114,7 @@ def test_model_from_stub(self, setup):

@staticmethod
def _assert_correct_files_downloaded(model, args):
if args[0] == "recipe":
assert len(model.recipes.available) == 1
elif args[0] == "checkpoint":
if args[0] == "checkpoint":
assert len(model.training.available) == 1
elif args[0] == "deployment":
assert len(model.training.available) == 1
Expand All @@ -141,16 +137,16 @@ def _assert_validation_results_exist(model):
"pytorch/sparseml/imagenet/pruned-moderate"
),
True,
files_ic,
files_ic.union({"recipe.md", "recipe_transfer_learn.md"}),
),
(
(
"zoo:"
"nlp/question_answering/distilbert-none/"
"pytorch/huggingface/squad/pruned80_quant-none-vnni"
),
), # MISSING tokenizer_config.json
False,
files_nlp,
files_nlp.union({"recipe.md"}),
),
(
(
Expand All @@ -159,22 +155,22 @@ def _assert_validation_results_exist(model):
"pytorch/ultralytics/coco/pruned_quant-aggressive_94"
),
True,
files_yolo,
files_yolo.union({"recipe.md", "recipe_transfer_learn.md"}),
),
(
"yolov5-x-coco-pruned70.4block_quantized",
False,
files_yolo,
files_yolo.union({"recipe.md", "recipe_transfer_learn.md"}),
),
(
"yolov5-n6-voc_coco-pruned55",
False,
files_yolo,
files_yolo.union({"recipe.md"}),
),
(
"resnet_v1-50-imagenet-channel30_pruned90_quantized",
False,
files_yolo,
files_yolo.union({"recipe.md", "recipe_transfer_classification.md"}),
),
],
scope="function",
Expand All @@ -196,11 +192,10 @@ def test_folder_structure(self, setup):
_, clone_sample_outputs, expected_files, temp_dir = setup
if clone_sample_outputs:
for file_name in [
"sample_outputs_onnxruntime",
"sample_outputs_deepsparse",
"sample-outputs_onnxruntime",
"sample-outputs_deepsparse",
]:
expected_files.update({file_name, file_name + ".tar.gz"})

assert not set(os.listdir(temp_dir.name)).difference(expected_files)

def test_validate(self, setup):
Expand Down Expand Up @@ -246,19 +241,19 @@ def _add_mock_files(directory_path: str, clone_sample_outputs: bool):
)
Path(optional_recipe_yaml).touch()

# add remaining `sample_{...}` files, that may be potentially
# add remaining `sample-{...}` files, that may be potentially
# missing
mock_sample_file = os.path.join(directory_path, "sample_inputs.tar.gz")
for file_name in ["sample_originals.tar.gz", "sample_labels.tar.gz"]:
mock_sample_file = os.path.join(directory_path, "sample-inputs.tar.gz")
for file_name in ["sample-originals.tar.gz", "sample-labels.tar.gz"]:
expected_file_dir = os.path.join(directory_path, file_name)
if not os.path.isfile(expected_file_dir):
shutil.copyfile(mock_sample_file, expected_file_dir)

if clone_sample_outputs:
sample_outputs_file = os.path.join(directory_path, "sample_outputs.tar.gz")
sample_outputs_file = os.path.join(directory_path, "sample-outputs.tar.gz")
for file_name in [
"sample_outputs_onnxruntime.tar.gz",
"sample_outputs_deepsparse.tar.gz",
"sample-outputs_onnxruntime.tar.gz",
"sample-outputs_deepsparse.tar.gz",
]:
shutil.copyfile(
sample_outputs_file, os.path.join(directory_path, file_name)
Expand All @@ -271,12 +266,11 @@ def _test_generate_outputs_single_engine(self, engine, model_directory):
if engine == "onnxruntime":
# test whether the functionality saves the numpy files to tar properly
tar_file_expected_path = os.path.join(
directory_path, f"sample_outputs_{engine}.tar.gz"
directory_path, f"sample-outputs_{engine}.tar.gz"
)
if os.path.isfile(tar_file_expected_path):
os.remove(tar_file_expected_path)
save_to_tar = True

output_expected = next(iter(model_directory.sample_outputs[engine]))
output_expected = list(output_expected.values())
output = next(
Expand All @@ -301,55 +295,55 @@ def _test_generate_outputs_single_engine(self, engine, model_directory):
assert os.path.isfile(tar_file_expected_path)


@pytest.mark.parametrize(
"stub",
[
"zoo:cv/classification/mobilenet_v1-1.0/pytorch/sparseml/"
"imagenet/pruned-moderate",
],
)
def test_model_gz_extraction_from_stub(stub: str):
temp_dir = tempfile.TemporaryDirectory(dir="/tmp")

model = Model(stub, temp_dir.name)
_extraction_test_helper(model)
shutil.rmtree(temp_dir.name)


@pytest.mark.parametrize(
"stub",
[
"zoo:cv/classification/mobilenet_v1-1.0/pytorch/sparseml/"
"imagenet/pruned-moderate",
],
)
def test_model_gz_extraction_from_local_files(stub: str):
temp_dir = tempfile.TemporaryDirectory(dir="/tmp")
model = Model(stub, temp_dir.name)
model.download()

source = temp_dir.name
model_from_local_files = Model(source)
_extraction_test_helper(model_from_local_files)
shutil.rmtree(temp_dir.name)


def _extraction_test_helper(model: Model):
# download and extract model.onnx.tar.gz
# path should point to extracted model.onnx file
path = Path(model.onnx_model.path)

# assert path points to model.onnx file
assert path.exists(), f"{path} does not exist"
assert path.is_file(), f"{path} is not a file"

# assert model.onnx.tar.gz exists
model_gz_path = path.with_name("model.onnx.tar.gz")
assert model_gz_path.exists(), f"{model_gz_path} does not exist"

# assert all members of model.onnx.tar.gz have been extracted
for zipped_filename in tarfile.open(model_gz_path).getnames():
unzipped_file_path = path.with_name(zipped_filename)
assert (
unzipped_file_path.exists()
), f"{unzipped_file_path} does not exist, was it extracted?"
# @pytest.mark.parametrize(
# "stub",
# [
# "zoo:cv/classification/mobilenet_v1-1.0/pytorch/sparseml/"
# "imagenet/pruned-moderate",
# ],
# )
# def test_model_gz_extraction_from_stub(stub: str):
# temp_dir = tempfile.TemporaryDirectory(dir="/tmp")

# model = Model(stub, temp_dir.name)
# _extraction_test_helper(model)
# shutil.rmtree(temp_dir.name)


# @pytest.mark.parametrize(
# "stub",
# [
# "zoo:cv/classification/mobilenet_v1-1.0/pytorch/sparseml/"
# "imagenet/pruned-moderate",
# ],
# )
# def test_model_gz_extraction_from_local_files(stub: str):
# temp_dir = tempfile.TemporaryDirectory(dir="/tmp")
# model = Model(stub, temp_dir.name)
# model.download()

# source = temp_dir.name
# model_from_local_files = Model(source)
# _extraction_test_helper(model_from_local_files)
# shutil.rmtree(temp_dir.name)


# def _extraction_test_helper(model: Model):
# # download and extract model.onnx.tar.gz
# # path should point to extracted model.onnx file
# path = Path(model.onnx_model.path)

# # assert path points to model.onnx file
# assert path.exists(), f"{path} does not exist"
# assert path.is_file(), f"{path} is not a file"

# # assert model.onnx.tar.gz exists
# model_gz_path = path.with_name("model.onnx.tar.gz")
# assert model_gz_path.exists(), f"{model_gz_path} does not exist"

# # assert all members of model.onnx.tar.gz have been extracted
# for zipped_filename in tarfile.open(model_gz_path).getnames():
# unzipped_file_path = path.with_name(zipped_filename)
# assert (
# unzipped_file_path.exists()
# ), f"{unzipped_file_path} does not exist, was it extracted?"
Loading

0 comments on commit 1a0077b

Please sign in to comment.