p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

limit ssim input img type to fp32 and fp16

This commit is contained in:
zhaozhenlong 2020-06-28 11:11:33 +08:00
parent cc4c40ce2e
commit 15580f8f6e
2 changed files with 7 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mindspore/nn/layer/image.py
View File

@ -110,6 +110,10 @@ def _check_input_filter_size(input_shape, param_name, filter_size, func_name):
validator.check(param_name + " shape[2]", input_shape[2], "filter_size", filter_size, Rel.GE, func_name)
validator.check(param_name + " shape[3]", input_shape[3], "filter_size", filter_size, Rel.GE, func_name)
@constexpr
def _check_input_dtype(input_dtype, param_name, allow_dtypes, cls_name):
validator.check_type_name(param_name, input_dtype, allow_dtypes, cls_name)
class SSIM(Cell):
r"""
Returns SSIM index between img1 and img2.
@ -160,6 +164,7 @@ class SSIM(Cell):
self.mean = P.DepthwiseConv2dNative(channel_multiplier=1, kernel_size=filter_size)
def construct(self, img1, img2):
_check_input_dtype(F.dtype(img1), "img1", [mstype.float32, mstype.float16], self.cls_name)
_check_input_filter_size(F.shape(img1), "img1", self.filter_size, self.cls_name)
P.SameTypeShape()(img1, img2)
max_val = _convert_img_dtype_to_float32(self.max_val, self.max_val)

12
tests/ut/python/nn/test_ssim.py
View File

@ -35,16 +35,8 @@ class SSIMNet(nn.Cell):
def test_compile():
net = SSIMNet()
img1 = Tensor(np.random.random((8, 3, 16, 16)))
img2 = Tensor(np.random.random((8, 3, 16, 16)))
_executor.compile(net, img1, img2)
def test_compile_grayscale():
max_val = 255
net = SSIMNet(max_val=max_val)
img1 = Tensor(np.random.randint(0, 256, (8, 1, 16, 16), np.uint8))
img2 = Tensor(np.random.randint(0, 256, (8, 1, 16, 16), np.uint8))
img1 = Tensor(np.random.random((8, 3, 16, 16)), mstype.float32)
img2 = Tensor(np.random.random((8, 3, 16, 16)), mstype.float32)
_executor.compile(net, img1, img2)