Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Save/load TorchScript object in test #1446

Merged
merged 2 commits into from
Apr 14, 2021
Merged

Save/load TorchScript object in test #1446

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bb886e7
Add dump step to torchscript test
mthrok Apr 9, 2021
00d439b
Fix test
mthrok Apr 14, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions test/torchaudio_unittest/functional/torchscript_consistency_impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,17 +6,21 @@
from parameterized import parameterized

from torchaudio_unittest import common_utils
from torchaudio_unittest.common_utils import TempDirMixin, TestBaseMixin
from torchaudio_unittest.common_utils import (
skipIfRocm,
)


class Functional(common_utils.TestBaseMixin):
class Functional(TempDirMixin, TestBaseMixin):
"""Implements test for `functinoal` modul that are performed for different devices"""
def _assert_consistency(self, func, tensor, shape_only=False):
tensor = tensor.to(device=self.device, dtype=self.dtype)

ts_func = torch.jit.script(func)
Copy link

@anjali411 anjali411 Apr 12, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we shouldn't remove this and continue checking scripted_fn(input) == fn(input) where scripted_fn = torch.jit.script(fn)
Synced with Meghan offline:
As mentioned in the comment below, torch.jit.load returns a scripted module, so the current test is ok, but I think it would be nice to move this code to a helper function that takes in a function/module and returns a scripted module, with a note (see the comment below) indicating what is exactly happening.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with the above. To clarify, torch.jit.load returns a scripted module with a forward method equivalent to the function that was serialized. The caller does not notice the difference because __call__ for the module calls forward.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, I have a mixed feeling about introducing another helper function.

First, this method IS the helper function used by the actual tests.
Secondly, my view is that if it's something that needs explanation with comment, keeping it along side of tests make it easier to follow the logic of tests.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At the moment, you duplicate the same code logic at different places, which I think strongly suggests that we should just have one helper function that can be used in all the tests.

path = self.get_temp_path('func.zip')
torch.jit.script(func).save(path)
ts_func = torch.jit.load(path)
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This returns a scripted module whose forward is the function you saved.


output = func(tensor)
ts_output = ts_func(tensor)
if shape_only:
Expand Down Expand Up @@ -565,15 +569,18 @@ def func(tensor):
self._assert_consistency(func, tensor)


class FunctionalComplex:
class FunctionalComplex(TempDirMixin, TestBaseMixin):
complex_dtype = None
real_dtype = None
device = None

def _assert_consistency(self, func, tensor, test_pseudo_complex=False):
assert tensor.is_complex()
tensor = tensor.to(device=self.device, dtype=self.complex_dtype)
ts_func = torch.jit.script(func)

path = self.get_temp_path('func.zip')
torch.jit.script(func).save(path)
ts_func = torch.jit.load(path)

if test_pseudo_complex:
tensor = torch.view_as_real(tensor)
Expand Down
16 changes: 12 additions & 4 deletions test/torchaudio_unittest/transforms/torchscript_consistency_impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,21 @@
from torchaudio_unittest import common_utils
from torchaudio_unittest.common_utils import (
skipIfRocm,
TempDirMixin,
TestBaseMixin,
)


class Transforms(common_utils.TestBaseMixin):
class Transforms(TempDirMixin, TestBaseMixin):
"""Implements test for Transforms that are performed for different devices"""
def _assert_consistency(self, transform, tensor):
tensor = tensor.to(device=self.device, dtype=self.dtype)
transform = transform.to(device=self.device, dtype=self.dtype)

ts_transform = torch.jit.script(transform)
path = self.get_temp_path('transform.zip')
torch.jit.script(transform).save(path)
ts_transform = torch.jit.load(path)

output = transform(tensor)
ts_output = ts_transform(tensor)
self.assertEqual(ts_output, output)
Expand Down Expand Up @@ -100,7 +105,7 @@ def test_SpectralCentroid(self):
self._assert_consistency(T.SpectralCentroid(sample_rate=sample_rate), waveform)


class TransformsComplex:
class TransformsComplex(TempDirMixin, TestBaseMixin):
complex_dtype = None
real_dtype = None
device = None
Expand All @@ -109,7 +114,10 @@ def _assert_consistency(self, transform, tensor, test_pseudo_complex=False):
assert tensor.is_complex()
tensor = tensor.to(device=self.device, dtype=self.complex_dtype)
transform = transform.to(device=self.device, dtype=self.real_dtype)
ts_transform = torch.jit.script(transform)

path = self.get_temp_path('transform.zip')
torch.jit.script(transform).save(path)
ts_transform = torch.jit.load(path)

if test_pseudo_complex:
tensor = torch.view_as_real(tensor)
Expand Down