def package_devices_model(self):
with TemporaryDirectory() as test_dir:
neuropod_path = os.path.join(test_dir, "test_neuropod")
# `create_torchscript_neuropod` runs inference with the test data immediately
# after creating the neuropod. Raises a ValueError if the model output
# does not match the expected output.
create_torchscript_neuropod(
neuropod_path=neuropod_path,
model_name="devices_model",
module=DevicesModel(),
input_spec=[
{
"name": "x",
"dtype": "float32",
"shape": (None, )
},
{
"name": "y",
"dtype": "float32",
"shape": (None, )
},
],
output_spec=[
{
"name": "x",
"dtype": "int64",
"shape": (None, )
},
{
"name": "y",
"dtype": "int64",
"shape": (None, )
},
],
test_input_data={
"x": np.arange(5).astype(np.float32),
"y": np.arange(5).astype(np.float32),
},
test_expected_out={
"x": np.array([0], dtype=np.int64),
"y": np.array([1], dtype=np.int64),
},
input_tensor_device={"x": "CPU"},
default_input_tensor_device="GPU",
)
# Ensure all inputs are moved to CPU if we run with no visible GPUs
load_and_test_neuropod(
neuropod_path,
test_input_data={
"x": np.arange(5).astype(np.float32),
"y": np.arange(5).astype(np.float32),
},
test_expected_out={
"x": np.array([0], dtype=np.int64),
"y": np.array([0], dtype=np.int64),
},
neuropod_load_args={"visible_gpu": None},
)
def package_named_tuple_model(self, do_fail=False):
with TemporaryDirectory() as test_dir:
neuropod_path = os.path.join(test_dir, "test_neuropod")
# `create_torchscript_neuropod` runs inference with the test data immediately
# after creating the neuropod. Raises a ValueError if the model output
# does not match the expected output.
create_torchscript_neuropod(
neuropod_path=neuropod_path,
model_name="named_tuple_model",
module=NamedTupleModel(),
input_spec=[
{
"name": "x",
"dtype": "float32",
"shape": ("batch_size", )
},
{
"name": "y",
"dtype": "float32",
"shape": ("batch_size", )
},
],
output_spec=[
{
"name": "sum",
"dtype": "float32",
"shape": ("batch_size", )
},
{
"name": "difference",
"dtype": "float32",
"shape": ("batch_size", ),
},
{
"name": "product",
"dtype": "float32",
"shape": ("batch_size", )
},
],
test_input_data={
"x": np.arange(5, dtype=np.float32),
"y": np.arange(5, dtype=np.float32),
},
test_expected_out={
"sum":
np.zeros(5) if do_fail else np.arange(5) + np.arange(5),
"difference":
np.zeros(5) if do_fail else np.arange(5) - np.arange(5),
"product":
np.zeros(5) if do_fail else np.arange(5) * np.arange(5),
},
)
def package_mixed_types_model(self, do_fail=False):
with TemporaryDirectory() as test_dir:
neuropod_path = os.path.join(test_dir, "test_neuropod")
# `create_torchscript_neuropod` runs inference with the test data immediately
# after creating the neuropod. Raises a ValueError if the model output
# does not match the expected output.
create_torchscript_neuropod(
neuropod_path=neuropod_path,
model_name="mixed_types_model",
module=MixedReturnTypesModel(),
# Get the input/output spec along with test data
**get_mixed_model_spec(do_fail=do_fail))
def package_simple_addition_model(self, do_fail=False):
for model in [CustomOpModel]:
with TemporaryDirectory() as test_dir:
neuropod_path = os.path.join(test_dir, "test_neuropod")
# `create_torchscript_neuropod` runs inference with the test data immediately
# after creating the neuropod. Raises a ValueError if the model output
# does not match the expected output.
create_torchscript_neuropod(
neuropod_path=neuropod_path,
model_name="addition_model",
module=model(),
custom_ops=[self.custom_op_path, self.second_custom_op],
# Get the input/output spec along with test data
**get_addition_model_spec(do_fail=do_fail))
def package_strings_model(out_dir, model=StringsModel, do_fail=False):
neuropod_path = os.path.join(out_dir, "test_neuropod")
# `create_torchscript_neuropod` runs inference with the test data immediately
# after creating the neuropod. Raises a ValueError if the model output
# does not match the expected output.
create_torchscript_neuropod(
neuropod_path=neuropod_path,
model_name="strings_model",
module=model(),
# Get the input/output spec along with test data
**get_string_concat_model_spec(do_fail=do_fail))
# Run some additional checks
check_strings_model(neuropod_path)
def test_mixed_types_model_failure_duplicate_item(self):
# Tests a model that returns duplicate items across multiple dictionaries
# This is either a CalledProcessError or a RuntimeError depending on whether
# we're using the native bindings or not
with self.assertRaises((CalledProcessError, RuntimeError)):
with TemporaryDirectory() as test_dir:
neuropod_path = os.path.join(test_dir, "test_neuropod")
# `create_torchscript_neuropod` runs inference with the test data immediately
# after creating the neuropod. Raises a ValueError if the model output
# does not match the expected output.
create_torchscript_neuropod(
neuropod_path=neuropod_path,
model_name="mixed_types_model",
module=MixedReturnTypesModelDuplicateItem(),
# Get the input/output spec along with test data
**get_mixed_model_spec())
def package_simple_addition_model(self, do_fail=False):
for model in [
AdditionModel, AdditionModelDictInput,
AdditionModelTensorOutput
]:
with TemporaryDirectory() as test_dir:
neuropod_path = os.path.join(test_dir, "test_neuropod")
# `create_torchscript_neuropod` runs inference with the test data immediately
# after creating the neuropod. Raises a ValueError if the model output
# does not match the expected output.
create_torchscript_neuropod(
neuropod_path=neuropod_path,
model_name="addition_model",
module=model(),
# Get the input/output spec along with test data
**get_addition_model_spec(do_fail=do_fail))
# Run some additional checks
check_addition_model(neuropod_path)