Speedups for to_tensorflow

TensorFlow is designed around graphs of ops for a DNN, not eager
channel shuffling and similar operations.  It's much faster to do the
channel shuffles in NumPy, although we still do the dtype conversion
in TensorFlow.

Author: jholveck

src/mss/screenshot.py

@@ -405,39 +405,22 @@ def to_tensorflow(
         # TypeErrors from tf.as_dtype are passed up to the caller.
         tf_dtype = tf.as_dtype(dtype)
 
-        # We don't need to do anything explicit about device management in TensorFlow; it handles that for us.
-        # convert_to_tensor will always copy.  We verified in __init__ that self._raw is a memoryview of unsigned bytes,
-        # so that's what TensorFlow will take as the source format; we just give an explicit dtype for clarity.
-        rv = tf.convert_to_tensor(self._raw, dtype=tf.uint8)
-        rv = tf.reshape(rv, (self.height, self.width, 4))
-
-        if channels == "BGRA":
-            pass
-        elif channels == "BGR":
-            rv = rv[:, :, :3]
-        elif channels == "RGB":
-            rv = tf.gather(rv, [2, 1, 0], axis=2)
-        elif channels == "RGBA":
-            rv = tf.gather(rv, [2, 1, 0, 3], axis=2)
-        else:
-            msg = 'Channels must be "BGRA", "BGR", "RGB", or "RGBA"'
-            raise ValueError(msg)
+        # Whether we're sending to a CPU or GPU, the channels and layout conversion are much faster in NumPy.  I suspect
+        # this is because doing this in TensorFlow eager mode materializes the tensor at each operation, but that's just
+        # a guess, and putting the channel shuffles in a tf.function didn't help.  It's still best (especially on GPU)
+        # to do the dtype conversion in TensorFlow.
+        ndarray = self.to_numpy(channels=channels, layout=layout)
 
-        if layout == "HWC":
-            pass
-        elif layout == "CHW":
-            rv = tf.transpose(rv, perm=(2, 0, 1))
-        else:
-            msg = 'Layout must be "HWC" or "CHW"'
-            raise ValueError(msg)
+        # We don't need to do anything explicit about device management in TensorFlow; it handles that for us.
+        # convert_to_tensor will always copy.
+        rv = tf.convert_to_tensor(ndarray)
 
-        # Do the conversion last to save memory bandwidth during channel shuffles.
         if tf_dtype != tf.uint8:
+            # I have no idea why, but doing the cast here instead of in convert_to_tensor is significantly faster,
+            # especially on the GPU.
             rv = tf.cast(rv, dtype=tf_dtype)
             if tf_dtype.is_floating:
-                # TensorFlow's implicit dtype conversion rules are not trivial.  We use an explicit dtype on both sides
-                # instead, by making a tf.constant.
-                rv = rv / tf.constant(255.0, dtype=tf_dtype)
+                rv /= 255.0
 
         return rv
 

src/tests/third_party/array_frameworks/test_tensorflow_method.py

@@ -45,7 +45,9 @@ def test_to_tensorflow_permutations(framework_test_image: ScreenShot, channels:
     bfloat16_target = uint8_target.astype(np.float32) / 255.0
 
     uint8_result = framework_test_image.to_tensorflow(channels=channels, layout=layout, dtype="uint8")  # type: ignore[arg-type]
+    assert uint8_result.dtype == tf.uint8
     assert np.array_equal(uint8_result.numpy(), uint8_target)
 
     bfloat16_result = framework_test_image.to_tensorflow(channels=channels, layout=layout, dtype="bfloat16")  # type: ignore[arg-type]
+    assert bfloat16_result.dtype == tf.bfloat16
     assert np.allclose(bfloat16_result.numpy(), bfloat16_target, rtol=0, atol=1 / 512.0)

src/tests/third_party/array_frameworks/test_torch_method.py

@@ -41,7 +41,7 @@ def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str,
     if cuda and not torch.cuda.is_available():
         # The CUDA versions won't be run in CI/CD, but it's still worth checking them on developers' machines if they
         # happen to have PyTorch with CUDA support.
-        pytest.skip()
+        pytest.skip("CUDA is not available")
 
     uint8_target = reordered_test_image(channels=channels, layout=layout)
     bfloat16_target = uint8_target.astype(np.float32) / 255.0
@@ -52,6 +52,7 @@ def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str,
         dtype=torch.uint8,
         device="cuda" if cuda else "cpu",
     )
+    assert uint8_result.dtype == torch.uint8
     assert np.array_equal(uint8_result.cpu().numpy(), uint8_target)
 
     bfloat16_result = framework_test_image.to_torch(
@@ -60,7 +61,8 @@ def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str,
         dtype=torch.bfloat16,
         device="cuda" if cuda else "cpu",
     )
+    assert bfloat16_result.dtype == torch.bfloat16
     # We have to explicitly cast back to float32 for the comparison, because PyTorch won't directly convert bfloat16 to
     # NumPy.
-    bfloat16_result = bfloat16_result.to(torch.float32)
-    assert np.allclose(bfloat16_result.cpu().numpy(), bfloat16_target, rtol=0, atol=1 / 512.0)
+    bfloat16_result_f32 = bfloat16_result.to(torch.float32)
+    assert np.allclose(bfloat16_result_f32.cpu().numpy(), bfloat16_target, rtol=0, atol=1 / 512.0)