Refactor test_transforms.py into classes (#4030)

test/test_transforms.py

@@ -195,100 +195,98 @@ class TestAccImage:
         torch.testing.assert_close(output, expected_output)
 
 
-@pytest.mark.parametrize('channels', [1, 3, 4])
-def test_to_tensor(channels):
-    height, width = 4, 4
-    trans = transforms.ToTensor()
-
-    input_data = torch.ByteTensor(channels, height, width).random_(0, 255).float().div_(255)
-    img = transforms.ToPILImage()(input_data)
-    output = trans(img)
-    torch.testing.assert_close(output, input_data, check_stride=False)
-
-    ndarray = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
-    output = trans(ndarray)
-    expected_output = ndarray.transpose((2, 0, 1)) / 255.0
-    torch.testing.assert_close(output.numpy(), expected_output, check_stride=False, check_dtype=False)
-
-    ndarray = np.random.rand(height, width, channels).astype(np.float32)
-    output = trans(ndarray)
-    expected_output = ndarray.transpose((2, 0, 1))
-    torch.testing.assert_close(output.numpy(), expected_output, check_stride=False, check_dtype=False)
-
-    # separate test for mode '1' PIL images
-    input_data = torch.ByteTensor(1, height, width).bernoulli_()
-    img = transforms.ToPILImage()(input_data.mul(255)).convert('1')
-    output = trans(img)
-    torch.testing.assert_close(input_data, output, check_dtype=False, check_stride=False)
-
-
-def test_to_tensor_errors():
-    height, width = 4, 4
-    trans = transforms.ToTensor()
-
-    with pytest.raises(TypeError):
-        trans(np.random.rand(1, height, width).tolist())
-
-    with pytest.raises(ValueError):
-        trans(np.random.rand(height))
-
-    with pytest.raises(ValueError):
-        trans(np.random.rand(1, 1, height, width))
+class TestToTensor:
 
+    @pytest.mark.parametrize('channels', [1, 3, 4])
+    def test_to_tensor(self, channels):
+        height, width = 4, 4
+        trans = transforms.ToTensor()
 
-@pytest.mark.parametrize('dtype', [torch.float16, torch.float, torch.double])
-def test_to_tensor_with_other_default_dtypes(dtype):
-    current_def_dtype = torch.get_default_dtype()
+        input_data = torch.ByteTensor(channels, height, width).random_(0, 255).float().div_(255)
+        img = transforms.ToPILImage()(input_data)
+        output = trans(img)
+        torch.testing.assert_close(output, input_data, check_stride=False)
+
+        ndarray = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
+        output = trans(ndarray)
+        expected_output = ndarray.transpose((2, 0, 1)) / 255.0
+        torch.testing.assert_close(output.numpy(), expected_output, check_stride=False, check_dtype=False)
+
+        ndarray = np.random.rand(height, width, channels).astype(np.float32)
+        output = trans(ndarray)
+        expected_output = ndarray.transpose((2, 0, 1))
+        torch.testing.assert_close(output.numpy(), expected_output, check_stride=False, check_dtype=False)
+
+        # separate test for mode '1' PIL images
+        input_data = torch.ByteTensor(1, height, width).bernoulli_()
+        img = transforms.ToPILImage()(input_data.mul(255)).convert('1')
+        output = trans(img)
+        torch.testing.assert_close(input_data, output, check_dtype=False, check_stride=False)
+
+    def test_to_tensor_errors(self):
+        height, width = 4, 4
+        trans = transforms.ToTensor()
 
-    t = transforms.ToTensor()
-    np_arr = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)
-    img = Image.fromarray(np_arr)
+        with pytest.raises(TypeError):
+            trans(np.random.rand(1, height, width).tolist())
 
-    torch.set_default_dtype(dtype)
-    res = t(img)
-    assert res.dtype == dtype, f"{res.dtype} vs {dtype}"
+        with pytest.raises(ValueError):
+            trans(np.random.rand(height))
 
-    torch.set_default_dtype(current_def_dtype)
+        with pytest.raises(ValueError):
+            trans(np.random.rand(1, 1, height, width))
 
+    @pytest.mark.parametrize('dtype', [torch.float16, torch.float, torch.double])
+    def test_to_tensor_with_other_default_dtypes(self, dtype):
+        current_def_dtype = torch.get_default_dtype()
 
-@pytest.mark.parametrize('channels', [1, 3, 4])
-def test_pil_to_tensor(channels):
-    height, width = 4, 4
-    trans = transforms.PILToTensor()
+        t = transforms.ToTensor()
+        np_arr = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)
+        img = Image.fromarray(np_arr)
 
-    input_data = torch.ByteTensor(channels, height, width).random_(0, 255)
-    img = transforms.ToPILImage()(input_data)
-    output = trans(img)
-    torch.testing.assert_close(input_data, output, check_stride=False)
+        torch.set_default_dtype(dtype)
+        res = t(img)
+        assert res.dtype == dtype, f"{res.dtype} vs {dtype}"
 
-    input_data = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
-    img = transforms.ToPILImage()(input_data)
-    output = trans(img)
-    expected_output = input_data.transpose((2, 0, 1))
-    torch.testing.assert_close(output.numpy(), expected_output)
+        torch.set_default_dtype(current_def_dtype)
 
-    input_data = torch.as_tensor(np.random.rand(channels, height, width).astype(np.float32))
-    img = transforms.ToPILImage()(input_data)  # CHW -> HWC and (* 255).byte()
-    output = trans(img)  # HWC -> CHW
-    expected_output = (input_data * 255).byte()
-    torch.testing.assert_close(output, expected_output, check_stride=False)
+    @pytest.mark.parametrize('channels', [1, 3, 4])
+    def test_pil_to_tensor(self, channels):
+        height, width = 4, 4
+        trans = transforms.PILToTensor()
 
-    # separate test for mode '1' PIL images
-    input_data = torch.ByteTensor(1, height, width).bernoulli_()
-    img = transforms.ToPILImage()(input_data.mul(255)).convert('1')
-    output = trans(img).view(torch.uint8).bool().to(torch.uint8)
-    torch.testing.assert_close(input_data, output, check_stride=False)
+        input_data = torch.ByteTensor(channels, height, width).random_(0, 255)
+        img = transforms.ToPILImage()(input_data)
+        output = trans(img)
+        torch.testing.assert_close(input_data, output, check_stride=False)
+
+        input_data = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
+        img = transforms.ToPILImage()(input_data)
+        output = trans(img)
+        expected_output = input_data.transpose((2, 0, 1))
+        torch.testing.assert_close(output.numpy(), expected_output)
+
+        input_data = torch.as_tensor(np.random.rand(channels, height, width).astype(np.float32))
+        img = transforms.ToPILImage()(input_data)  # CHW -> HWC and (* 255).byte()
+        output = trans(img)  # HWC -> CHW
+        expected_output = (input_data * 255).byte()
+        torch.testing.assert_close(output, expected_output, check_stride=False)
 
+        # separate test for mode '1' PIL images
+        input_data = torch.ByteTensor(1, height, width).bernoulli_()
+        img = transforms.ToPILImage()(input_data.mul(255)).convert('1')
+        output = trans(img).view(torch.uint8).bool().to(torch.uint8)
+        torch.testing.assert_close(input_data, output, check_stride=False)
 
-def test_pil_to_tensor_errors():
-    height, width = 4, 4
-    trans = transforms.PILToTensor()
+    def test_pil_to_tensor_errors(self):
+        height, width = 4, 4
+        trans = transforms.PILToTensor()
 
-    with pytest.raises(TypeError):
-        trans(np.random.rand(1, height, width).tolist())
+        with pytest.raises(TypeError):
+            trans(np.random.rand(1, height, width).tolist())
 
-    with pytest.raises(TypeError):
-        trans(np.random.rand(1, height, width))
+        with pytest.raises(TypeError):
+            trans(np.random.rand(1, height, width))
 
 
 def test_randomresized_params():
@@ -537,303 +535,285 @@ def test_randomness(fn, trans, config, p):
     assert p_value > 0.0001
 
 
-def _get_1_channel_tensor_various_types():
-    img_data_float = torch.Tensor(1, 4, 4).uniform_()
-    expected_output = img_data_float.mul(255).int().float().div(255).numpy()
-    yield img_data_float, expected_output, 'L'
+class TestToPil:
 
-    img_data_byte = torch.ByteTensor(1, 4, 4).random_(0, 255)
-    expected_output = img_data_byte.float().div(255.0).numpy()
-    yield img_data_byte, expected_output, 'L'
+    def _get_1_channel_tensor_various_types():
+        img_data_float = torch.Tensor(1, 4, 4).uniform_()
+        expected_output = img_data_float.mul(255).int().float().div(255).numpy()
+        yield img_data_float, expected_output, 'L'
 
-    img_data_short = torch.ShortTensor(1, 4, 4).random_()
-    expected_output = img_data_short.numpy()
-    yield img_data_short, expected_output, 'I;16'
+        img_data_byte = torch.ByteTensor(1, 4, 4).random_(0, 255)
+        expected_output = img_data_byte.float().div(255.0).numpy()
+        yield img_data_byte, expected_output, 'L'
 
-    img_data_int = torch.IntTensor(1, 4, 4).random_()
-    expected_output = img_data_int.numpy()
-    yield img_data_int, expected_output, 'I'
+        img_data_short = torch.ShortTensor(1, 4, 4).random_()
+        expected_output = img_data_short.numpy()
+        yield img_data_short, expected_output, 'I;16'
 
+        img_data_int = torch.IntTensor(1, 4, 4).random_()
+        expected_output = img_data_int.numpy()
+        yield img_data_int, expected_output, 'I'
 
-@pytest.mark.parametrize('with_mode', [False, True])
-@pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_1_channel_tensor_various_types())
-def test_1_channel_tensor_to_pil_image(with_mode, img_data, expected_output, expected_mode):
-    transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
-    to_tensor = transforms.ToTensor()
+    def _get_2d_tensor_various_types():
+        img_data_float = torch.Tensor(4, 4).uniform_()
+        expected_output = img_data_float.mul(255).int().float().div(255).numpy()
+        yield img_data_float, expected_output, 'L'
 
-    img = transform(img_data)
-    assert img.mode == expected_mode
-    torch.testing.assert_close(expected_output, to_tensor(img).numpy(), check_stride=False)
+        img_data_byte = torch.ByteTensor(4, 4).random_(0, 255)
+        expected_output = img_data_byte.float().div(255.0).numpy()
+        yield img_data_byte, expected_output, 'L'
 
+        img_data_short = torch.ShortTensor(4, 4).random_()
+        expected_output = img_data_short.numpy()
+        yield img_data_short, expected_output, 'I;16'
 
-def test_1_channel_float_tensor_to_pil_image():
-    img_data = torch.Tensor(1, 4, 4).uniform_()
-    # 'F' mode for torch.FloatTensor
-    img_F_mode = transforms.ToPILImage(mode='F')(img_data)
-    assert img_F_mode.mode == 'F'
-    torch.testing.assert_close(
-        np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode='F')), np.array(img_F_mode)
-    )
+        img_data_int = torch.IntTensor(4, 4).random_()
+        expected_output = img_data_int.numpy()
+        yield img_data_int, expected_output, 'I'
 
+    @pytest.mark.parametrize('with_mode', [False, True])
+    @pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_1_channel_tensor_various_types())
+    def test_1_channel_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode):
+        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
+        to_tensor = transforms.ToTensor()
 
-@pytest.mark.parametrize('with_mode', [False, True])
-@pytest.mark.parametrize('img_data, expected_mode', [
-    (torch.Tensor(4, 4, 1).uniform_().numpy(), 'F'),
-    (torch.ByteTensor(4, 4, 1).random_(0, 255).numpy(), 'L'),
-    (torch.ShortTensor(4, 4, 1).random_().numpy(), 'I;16'),
-    (torch.IntTensor(4, 4, 1).random_().numpy(), 'I'),
-])
-def test_1_channel_ndarray_to_pil_image(with_mode, img_data, expected_mode):
-    transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
-    img = transform(img_data)
-    assert img.mode == expected_mode
-    # note: we explicitly convert img's dtype because pytorch doesn't support uint16
-    # and otherwise assert_close wouldn't be able to construct a tensor from the uint16 array
-    torch.testing.assert_close(img_data[:, :, 0], np.asarray(img).astype(img_data.dtype))
-
-
-@pytest.mark.parametrize('expected_mode', [None, 'LA'])
-def test_2_channel_ndarray_to_pil_image(expected_mode):
-    img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy()
-
-    if expected_mode is None:
-        img = transforms.ToPILImage()(img_data)
-        assert img.mode == 'LA'  # default should assume LA
-    else:
-        img = transforms.ToPILImage(mode=expected_mode)(img_data)
-        assert img.mode == expected_mode
-    split = img.split()
-    for i in range(2):
-        torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False)
-
-
-def test_2_channel_ndarray_to_pil_image_error():
-    img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy()
-    transforms.ToPILImage().__repr__()
-
-    # should raise if we try a mode for 4 or 1 or 3 channel images
-    with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
-        transforms.ToPILImage(mode='RGBA')(img_data)
-    with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
-        transforms.ToPILImage(mode='P')(img_data)
-    with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
-        transforms.ToPILImage(mode='RGB')(img_data)
-
-
-@pytest.mark.parametrize('expected_mode', [None, 'LA'])
-def test_2_channel_tensor_to_pil_image(expected_mode):
-    img_data = torch.Tensor(2, 4, 4).uniform_()
-    expected_output = img_data.mul(255).int().float().div(255)
-    if expected_mode is None:
-        img = transforms.ToPILImage()(img_data)
-        assert img.mode == 'LA'  # default should assume LA
-    else:
-        img = transforms.ToPILImage(mode=expected_mode)(img_data)
+        img = transform(img_data)
         assert img.mode == expected_mode
+        torch.testing.assert_close(expected_output, to_tensor(img).numpy(), check_stride=False)
 
-    split = img.split()
-    for i in range(2):
-        torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
-
-
-def test_2_channel_tensor_to_pil_image_error():
-    img_data = torch.Tensor(2, 4, 4).uniform_()
-
-    # should raise if we try a mode for 4 or 1 or 3 channel images
-    with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
-        transforms.ToPILImage(mode='RGBA')(img_data)
-    with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
-        transforms.ToPILImage(mode='P')(img_data)
-    with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
-        transforms.ToPILImage(mode='RGB')(img_data)
-
-
-def _get_2d_tensor_various_types():
-    img_data_float = torch.Tensor(4, 4).uniform_()
-    expected_output = img_data_float.mul(255).int().float().div(255).numpy()
-    yield img_data_float, expected_output, 'L'
-
-    img_data_byte = torch.ByteTensor(4, 4).random_(0, 255)
-    expected_output = img_data_byte.float().div(255.0).numpy()
-    yield img_data_byte, expected_output, 'L'
-
-    img_data_short = torch.ShortTensor(4, 4).random_()
-    expected_output = img_data_short.numpy()
-    yield img_data_short, expected_output, 'I;16'
-
-    img_data_int = torch.IntTensor(4, 4).random_()
-    expected_output = img_data_int.numpy()
-    yield img_data_int, expected_output, 'I'
-
-
-@pytest.mark.parametrize('with_mode', [False, True])
-@pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_2d_tensor_various_types())
-def test_2d_tensor_to_pil_image(with_mode, img_data, expected_output, expected_mode):
-    transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
-    to_tensor = transforms.ToTensor()
-
-    img = transform(img_data)
-    assert img.mode == expected_mode
-    torch.testing.assert_close(expected_output, to_tensor(img).numpy()[0])
-
-
-@pytest.mark.parametrize('with_mode', [False, True])
-@pytest.mark.parametrize('img_data, expected_mode', [
-    (torch.Tensor(4, 4).uniform_().numpy(), 'F'),
-    (torch.ByteTensor(4, 4).random_(0, 255).numpy(), 'L'),
-    (torch.ShortTensor(4, 4).random_().numpy(), 'I;16'),
-    (torch.IntTensor(4, 4).random_().numpy(), 'I'),
-])
-def test_2d_ndarray_to_pil_image(with_mode, img_data, expected_mode):
-    transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
-    img = transform(img_data)
-    assert img.mode == expected_mode
-    np.testing.assert_allclose(img_data, img)
-
-
-@pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr'])
-def test_3_channel_tensor_to_pil_image(expected_mode):
-    img_data = torch.Tensor(3, 4, 4).uniform_()
-    expected_output = img_data.mul(255).int().float().div(255)
+    def test_1_channel_float_tensor_to_pil_image(self):
+        img_data = torch.Tensor(1, 4, 4).uniform_()
+        # 'F' mode for torch.FloatTensor
+        img_F_mode = transforms.ToPILImage(mode='F')(img_data)
+        assert img_F_mode.mode == 'F'
+        torch.testing.assert_close(
+            np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode='F')), np.array(img_F_mode)
+        )
 
-    if expected_mode is None:
-        img = transforms.ToPILImage()(img_data)
-        assert img.mode == 'RGB'  # default should assume RGB
-    else:
-        img = transforms.ToPILImage(mode=expected_mode)(img_data)
+    @pytest.mark.parametrize('with_mode', [False, True])
+    @pytest.mark.parametrize('img_data, expected_mode', [
+        (torch.Tensor(4, 4, 1).uniform_().numpy(), 'F'),
+        (torch.ByteTensor(4, 4, 1).random_(0, 255).numpy(), 'L'),
+        (torch.ShortTensor(4, 4, 1).random_().numpy(), 'I;16'),
+        (torch.IntTensor(4, 4, 1).random_().numpy(), 'I'),
+    ])
+    def test_1_channel_ndarray_to_pil_image(self, with_mode, img_data, expected_mode):
+        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
+        img = transform(img_data)
         assert img.mode == expected_mode
-    split = img.split()
-    for i in range(3):
-        torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
+        # note: we explicitly convert img's dtype because pytorch doesn't support uint16
+        # and otherwise assert_close wouldn't be able to construct a tensor from the uint16 array
+        torch.testing.assert_close(img_data[:, :, 0], np.asarray(img).astype(img_data.dtype))
 
+    @pytest.mark.parametrize('expected_mode', [None, 'LA'])
+    def test_2_channel_ndarray_to_pil_image(self, expected_mode):
+        img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy()
 
-def test_3_channel_tensor_to_pil_image_error():
-    img_data = torch.Tensor(3, 4, 4).uniform_()
-    error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs"
-    # should raise if we try a mode for 4 or 1 or 2 channel images
-    with pytest.raises(ValueError, match=error_message_3d):
-        transforms.ToPILImage(mode='RGBA')(img_data)
-    with pytest.raises(ValueError, match=error_message_3d):
-        transforms.ToPILImage(mode='P')(img_data)
-    with pytest.raises(ValueError, match=error_message_3d):
-        transforms.ToPILImage(mode='LA')(img_data)
-
-    with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'):
-        transforms.ToPILImage()(torch.Tensor(1, 3, 4, 4).uniform_())
-
-
-@pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr'])
-def test_3_channel_ndarray_to_pil_image(expected_mode):
-    img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy()
-
-    if expected_mode is None:
-        img = transforms.ToPILImage()(img_data)
-        assert img.mode == 'RGB'  # default should assume RGB
-    else:
-        img = transforms.ToPILImage(mode=expected_mode)(img_data)
+        if expected_mode is None:
+            img = transforms.ToPILImage()(img_data)
+            assert img.mode == 'LA'  # default should assume LA
+        else:
+            img = transforms.ToPILImage(mode=expected_mode)(img_data)
+            assert img.mode == expected_mode
+        split = img.split()
+        for i in range(2):
+            torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False)
+
+    def test_2_channel_ndarray_to_pil_image_error(self):
+        img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy()
+        transforms.ToPILImage().__repr__()
+
+        # should raise if we try a mode for 4 or 1 or 3 channel images
+        with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
+            transforms.ToPILImage(mode='RGBA')(img_data)
+        with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
+            transforms.ToPILImage(mode='P')(img_data)
+        with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
+            transforms.ToPILImage(mode='RGB')(img_data)
+
+    @pytest.mark.parametrize('expected_mode', [None, 'LA'])
+    def test_2_channel_tensor_to_pil_image(self, expected_mode):
+        img_data = torch.Tensor(2, 4, 4).uniform_()
+        expected_output = img_data.mul(255).int().float().div(255)
+        if expected_mode is None:
+            img = transforms.ToPILImage()(img_data)
+            assert img.mode == 'LA'  # default should assume LA
+        else:
+            img = transforms.ToPILImage(mode=expected_mode)(img_data)
+            assert img.mode == expected_mode
+
+        split = img.split()
+        for i in range(2):
+            torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
+
+    def test_2_channel_tensor_to_pil_image_error(self):
+        img_data = torch.Tensor(2, 4, 4).uniform_()
+
+        # should raise if we try a mode for 4 or 1 or 3 channel images
+        with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
+            transforms.ToPILImage(mode='RGBA')(img_data)
+        with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
+            transforms.ToPILImage(mode='P')(img_data)
+        with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"):
+            transforms.ToPILImage(mode='RGB')(img_data)
+
+    @pytest.mark.parametrize('with_mode', [False, True])
+    @pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_2d_tensor_various_types())
+    def test_2d_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode):
+        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
+        to_tensor = transforms.ToTensor()
+
+        img = transform(img_data)
         assert img.mode == expected_mode
-    split = img.split()
-    for i in range(3):
-        torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False)
-
-
-def test_3_channel_ndarray_to_pil_image_error():
-    img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy()
-
-    # Checking if ToPILImage can be printed as string
-    transforms.ToPILImage().__repr__()
-
-    error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs"
-    # should raise if we try a mode for 4 or 1 or 2 channel images
-    with pytest.raises(ValueError, match=error_message_3d):
-        transforms.ToPILImage(mode='RGBA')(img_data)
-    with pytest.raises(ValueError, match=error_message_3d):
-        transforms.ToPILImage(mode='P')(img_data)
-    with pytest.raises(ValueError, match=error_message_3d):
-        transforms.ToPILImage(mode='LA')(img_data)
-
-
-@pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX'])
-def test_4_channel_tensor_to_pil_image(expected_mode):
-    img_data = torch.Tensor(4, 4, 4).uniform_()
-    expected_output = img_data.mul(255).int().float().div(255)
-
-    if expected_mode is None:
-        img = transforms.ToPILImage()(img_data)
-        assert img.mode == 'RGBA'  # default should assume RGBA
-    else:
-        img = transforms.ToPILImage(mode=expected_mode)(img_data)
+        torch.testing.assert_close(expected_output, to_tensor(img).numpy()[0])
+
+    @pytest.mark.parametrize('with_mode', [False, True])
+    @pytest.mark.parametrize('img_data, expected_mode', [
+        (torch.Tensor(4, 4).uniform_().numpy(), 'F'),
+        (torch.ByteTensor(4, 4).random_(0, 255).numpy(), 'L'),
+        (torch.ShortTensor(4, 4).random_().numpy(), 'I;16'),
+        (torch.IntTensor(4, 4).random_().numpy(), 'I'),
+    ])
+    def test_2d_ndarray_to_pil_image(self, with_mode, img_data, expected_mode):
+        transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage()
+        img = transform(img_data)
         assert img.mode == expected_mode
+        np.testing.assert_allclose(img_data, img)
 
-    split = img.split()
-    for i in range(4):
-        torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
-
-
-def test_4_channel_tensor_to_pil_image_error():
-    img_data = torch.Tensor(4, 4, 4).uniform_()
-
-    error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs"
-    # should raise if we try a mode for 3 or 1 or 2 channel images
-    with pytest.raises(ValueError, match=error_message_4d):
-        transforms.ToPILImage(mode='RGB')(img_data)
-    with pytest.raises(ValueError, match=error_message_4d):
-        transforms.ToPILImage(mode='P')(img_data)
-    with pytest.raises(ValueError, match=error_message_4d):
-        transforms.ToPILImage(mode='LA')(img_data)
+    @pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr'])
+    def test_3_channel_tensor_to_pil_image(self, expected_mode):
+        img_data = torch.Tensor(3, 4, 4).uniform_()
+        expected_output = img_data.mul(255).int().float().div(255)
 
-
-@pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX'])
-def test_4_channel_ndarray_to_pil_image(expected_mode):
-    img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy()
-
-    if expected_mode is None:
-        img = transforms.ToPILImage()(img_data)
-        assert img.mode == 'RGBA'  # default should assume RGBA
-    else:
-        img = transforms.ToPILImage(mode=expected_mode)(img_data)
-        assert img.mode == expected_mode
-    split = img.split()
-    for i in range(4):
-        torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False)
-
-
-def test_4_channel_ndarray_to_pil_image_error():
-    img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy()
-
-    error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs"
-    # should raise if we try a mode for 3 or 1 or 2 channel images
-    with pytest.raises(ValueError, match=error_message_4d):
-        transforms.ToPILImage(mode='RGB')(img_data)
-    with pytest.raises(ValueError, match=error_message_4d):
-        transforms.ToPILImage(mode='P')(img_data)
-    with pytest.raises(ValueError, match=error_message_4d):
-        transforms.ToPILImage(mode='LA')(img_data)
-
-
-def test_ndarray_bad_types_to_pil_image():
-    trans = transforms.ToPILImage()
-    reg_msg = r'Input type \w+ is not supported'
-    with pytest.raises(TypeError, match=reg_msg):
-        trans(np.ones([4, 4, 1], np.int64))
-    with pytest.raises(TypeError, match=reg_msg):
-        trans(np.ones([4, 4, 1], np.uint16))
-    with pytest.raises(TypeError, match=reg_msg):
-        trans(np.ones([4, 4, 1], np.uint32))
-    with pytest.raises(TypeError, match=reg_msg):
-        trans(np.ones([4, 4, 1], np.float64))
-
-    with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'):
-        transforms.ToPILImage()(np.ones([1, 4, 4, 3]))
-    with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'):
-        transforms.ToPILImage()(np.ones([4, 4, 6]))
-
-
-def test_tensor_bad_types_to_pil_image():
-    with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'):
-        transforms.ToPILImage()(torch.ones(1, 3, 4, 4))
-    with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'):
-        transforms.ToPILImage()(torch.ones(6, 4, 4))
+        if expected_mode is None:
+            img = transforms.ToPILImage()(img_data)
+            assert img.mode == 'RGB'  # default should assume RGB
+        else:
+            img = transforms.ToPILImage(mode=expected_mode)(img_data)
+            assert img.mode == expected_mode
+        split = img.split()
+        for i in range(3):
+            torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
+
+    def test_3_channel_tensor_to_pil_image_error(self):
+        img_data = torch.Tensor(3, 4, 4).uniform_()
+        error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs"
+        # should raise if we try a mode for 4 or 1 or 2 channel images
+        with pytest.raises(ValueError, match=error_message_3d):
+            transforms.ToPILImage(mode='RGBA')(img_data)
+        with pytest.raises(ValueError, match=error_message_3d):
+            transforms.ToPILImage(mode='P')(img_data)
+        with pytest.raises(ValueError, match=error_message_3d):
+            transforms.ToPILImage(mode='LA')(img_data)
+
+        with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'):
+            transforms.ToPILImage()(torch.Tensor(1, 3, 4, 4).uniform_())
+
+    @pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr'])
+    def test_3_channel_ndarray_to_pil_image(self, expected_mode):
+        img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy()
+
+        if expected_mode is None:
+            img = transforms.ToPILImage()(img_data)
+            assert img.mode == 'RGB'  # default should assume RGB
+        else:
+            img = transforms.ToPILImage(mode=expected_mode)(img_data)
+            assert img.mode == expected_mode
+        split = img.split()
+        for i in range(3):
+            torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False)
+
+    def test_3_channel_ndarray_to_pil_image_error(self):
+        img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy()
+
+        # Checking if ToPILImage can be printed as string
+        transforms.ToPILImage().__repr__()
+
+        error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs"
+        # should raise if we try a mode for 4 or 1 or 2 channel images
+        with pytest.raises(ValueError, match=error_message_3d):
+            transforms.ToPILImage(mode='RGBA')(img_data)
+        with pytest.raises(ValueError, match=error_message_3d):
+            transforms.ToPILImage(mode='P')(img_data)
+        with pytest.raises(ValueError, match=error_message_3d):
+            transforms.ToPILImage(mode='LA')(img_data)
+
+    @pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX'])
+    def test_4_channel_tensor_to_pil_image(self, expected_mode):
+        img_data = torch.Tensor(4, 4, 4).uniform_()
+        expected_output = img_data.mul(255).int().float().div(255)
+
+        if expected_mode is None:
+            img = transforms.ToPILImage()(img_data)
+            assert img.mode == 'RGBA'  # default should assume RGBA
+        else:
+            img = transforms.ToPILImage(mode=expected_mode)(img_data)
+            assert img.mode == expected_mode
+
+        split = img.split()
+        for i in range(4):
+            torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
+
+    def test_4_channel_tensor_to_pil_image_error(self):
+        img_data = torch.Tensor(4, 4, 4).uniform_()
+
+        error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs"
+        # should raise if we try a mode for 3 or 1 or 2 channel images
+        with pytest.raises(ValueError, match=error_message_4d):
+            transforms.ToPILImage(mode='RGB')(img_data)
+        with pytest.raises(ValueError, match=error_message_4d):
+            transforms.ToPILImage(mode='P')(img_data)
+        with pytest.raises(ValueError, match=error_message_4d):
+            transforms.ToPILImage(mode='LA')(img_data)
+
+    @pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX'])
+    def test_4_channel_ndarray_to_pil_image(self, expected_mode):
+        img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy()
+
+        if expected_mode is None:
+            img = transforms.ToPILImage()(img_data)
+            assert img.mode == 'RGBA'  # default should assume RGBA
+        else:
+            img = transforms.ToPILImage(mode=expected_mode)(img_data)
+            assert img.mode == expected_mode
+        split = img.split()
+        for i in range(4):
+            torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False)
+
+    def test_4_channel_ndarray_to_pil_image_error(self):
+        img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy()
+
+        error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs"
+        # should raise if we try a mode for 3 or 1 or 2 channel images
+        with pytest.raises(ValueError, match=error_message_4d):
+            transforms.ToPILImage(mode='RGB')(img_data)
+        with pytest.raises(ValueError, match=error_message_4d):
+            transforms.ToPILImage(mode='P')(img_data)
+        with pytest.raises(ValueError, match=error_message_4d):
+            transforms.ToPILImage(mode='LA')(img_data)
+
+    def test_ndarray_bad_types_to_pil_image(self):
+        trans = transforms.ToPILImage()
+        reg_msg = r'Input type \w+ is not supported'
+        with pytest.raises(TypeError, match=reg_msg):
+            trans(np.ones([4, 4, 1], np.int64))
+        with pytest.raises(TypeError, match=reg_msg):
+            trans(np.ones([4, 4, 1], np.uint16))
+        with pytest.raises(TypeError, match=reg_msg):
+            trans(np.ones([4, 4, 1], np.uint32))
+        with pytest.raises(TypeError, match=reg_msg):
+            trans(np.ones([4, 4, 1], np.float64))
+
+        with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'):
+            transforms.ToPILImage()(np.ones([1, 4, 4, 3]))
+        with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'):
+            transforms.ToPILImage()(np.ones([4, 4, 6]))
+
+    def test_tensor_bad_types_to_pil_image(self):
+        with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'):
+            transforms.ToPILImage()(torch.ones(1, 3, 4, 4))
+        with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'):
+            transforms.ToPILImage()(torch.ones(6, 4, 4))
 
 
 def test_adjust_brightness():