Fix pytorch loader (#812)

* Throw error if can't apply preprocessing. Fix pytorch label conversion

* black formatting

* Disabling preprocessing durings tests

* Raise error if value is not a tf tensor

* Fixing broken tf1 framework dataset test

Author: kmerchant31892

armory/data/datasets.py

@@ -268,6 +268,13 @@ def _generator_from_tfds(
         ds = ds.batch(batch_size, drop_remainder=False)
     ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
 
+    if framework != "numpy" and (
+        preprocessing_fn is not None or label_preprocessing_fn is not None
+    ):
+        raise ValueError(
+            f"Data/label preprocessing functions only supported for numpy framework.  Selected {framework} framework"
+        )
+
     if framework == "numpy":
         ds = tfds.as_numpy(ds, graph=default_graph)
         generator = ArmoryDataGenerator(

armory/data/pytorch_loader.py

@@ -1,4 +1,5 @@
 import torch
+import tensorflow as tf
 
 
 class TFToTorchGenerator(torch.utils.data.IterableDataset):
@@ -9,12 +10,26 @@ def __init__(self, tf_dataset):
     def __iter__(self):
         for ex in self.tf_dataset.take(-1):
             x, y = ex
-            # manually handle adverarial dataset
+            # separately handle benign/adversarial data formats
             if isinstance(x, tuple):
-                x = (torch.from_numpy(x[0].numpy()), torch.from_numpy(x[1].numpy()))
-                y = torch.from_numpy(y.numpy())
-            # non-adversarial dataset
+                x_torch = (
+                    torch.from_numpy(x[0].numpy()),
+                    torch.from_numpy(x[1].numpy()),
+                )
             else:
-                x = torch.from_numpy(x.numpy())
-                y = torch.from_numpy(y.numpy())
-            yield x, y
+                x_torch = torch.from_numpy(x.numpy())
+
+            # separately handle tensor/object detection label formats
+            if isinstance(y, dict):
+                y_torch = {}
+                for k, v in y.items():
+                    if isinstance(v, tf.Tensor):
+                        y_torch[k] = torch.from_numpy(v.numpy())
+                    else:
+                        raise ValueError(
+                            f"Expected all values to be of type tf.Tensor, but value at key {k} is of type {type(v)}"
+                        )
+            else:
+                y_torch = torch.from_numpy(y.numpy())
+
+            yield x_torch, y_torch

tests/test_pytorch/test_framework_dataset.py

@@ -19,6 +19,8 @@ def test_pytorch_generator_cifar10():
         batch_size=batch_size,
         dataset_dir=DATASET_DIR,
         framework="pytorch",
+        preprocessing_fn=None,
+        fit_preprocessing_fn=None,
     )
 
     assert isinstance(dataset, torch.utils.data.DataLoader)
@@ -38,6 +40,8 @@ def test_pytorch_generator_mnist():
         batch_size=batch_size,
         dataset_dir=DATASET_DIR,
         framework="pytorch",
+        preprocessing_fn=None,
+        fit_preprocessing_fn=None,
     )
 
     assert isinstance(dataset, torch.utils.data.DataLoader)
@@ -57,6 +61,8 @@ def test_pytorch_generator_resisc():
         batch_size=batch_size,
         dataset_dir=DATASET_DIR,
         framework="pytorch",
+        preprocessing_fn=None,
+        fit_preprocessing_fn=None,
     )
 
     assert isinstance(dataset, torch.utils.data.DataLoader)
@@ -76,6 +82,8 @@ def test_pytorch_generator_epochs():
         batch_size=batch_size,
         dataset_dir=DATASET_DIR,
         framework="pytorch",
+        preprocessing_fn=None,
+        fit_preprocessing_fn=None,
     )
 
     cnt = 0
@@ -100,6 +108,8 @@ def test_tf_pytorch_equality():
         dataset_dir=DATASET_DIR,
         framework="tf",
         shuffle_files=False,
+        preprocessing_fn=None,
+        fit_preprocessing_fn=None,
     )
 
     ds_pytorch = iter(
@@ -109,6 +119,8 @@ def test_tf_pytorch_equality():
             dataset_dir=DATASET_DIR,
             framework="pytorch",
             shuffle_files=False,
+            preprocessing_fn=None,
+            fit_preprocessing_fn=None,
         )
     )
 

tests/test_tf1/test_framework_dataset.py

@@ -17,5 +17,7 @@ def test_tf_generator():
         batch_size=16,
         dataset_dir=DATASET_DIR,
         framework="tf",
+        preprocessing_fn=None,
+        fit_preprocessing_fn=None,
     )
     assert isinstance(dataset, (tf.compat.v2.data.Dataset, tf.compat.v1.data.Dataset))