Repository: libffcv/ffcv (language: python)

Issue:
Support `TorchTensorField`
There should be a counterpart to `NDArrayField` for torch Tensors. This would be useful to convert existing `torch.utils.data.Dataset`s to `.beton`.

Following the API described in this guide: https://docs.ffcv.io/writing_datasets.html#writing-a-dataset-to-ffcv-format

```python
dataset = Dataset(...)
fields = {"covariate": TorchTensorField(shape=..., dtype=torch.float32)}
writer = DatasetWriter(dataset_path, fields)
writer.from_indexed_dataset(dataset)
```

Using `NDArrayField` currently fails with:

```python
  File "/home/carlos/miniconda3/envs/ffcv/lib/python3.8/site-packages/ffcv/fields/ndarray.py", line 94, in encode
    data_region[:] = field.reshape(-1).view('<u1')
TypeError: view() received an invalid combination of arguments - got (str), but expected one of:
 * (tuple of ints size)
      didn't match because some of the arguments have invalid types: (str)
 * (torch.dtype dtype)
      didn't match because some of the arguments have invalid types: (str)
```


Discussion / hints:
I think it would just be easier to just check the type and convert tensors as part of the encode procedure of the NDarray field. Thoughts ?

I could also add what you suggest and make it extend NDarray field. 

Any preference ? (Both are super easy and I'll add that in the next release for sure)
I think the added class (even if it's a minimal extension) is the simplest option and most natural to users
Trying to balance ICML submissions and adding features to FFCV but I'll try to have this today. 

Fix patch:
diff --git a/ffcv/fields/__init__.py b/ffcv/fields/__init__.py
index 90987fe7..79b76d52 100644
--- a/ffcv/fields/__init__.py
+++ b/ffcv/fields/__init__.py
@@ -2,8 +2,8 @@
 from .basics import FloatField, IntField
 from .rgb_image import RGBImageField
 from .bytes import BytesField
-from .ndarray import NDArrayField
+from .ndarray import NDArrayField, TorchTensorField
 from .json import JSONField
 
 __all__ = ['Field', 'BytesField', 'IntField', 'FloatField', 'RGBImageField',
-           'NDArrayField', 'JSONField']
\ No newline at end of file
+           'NDArrayField', 'JSONField', 'TorchTensorField']
\ No newline at end of file
diff --git a/ffcv/fields/ndarray.py b/ffcv/fields/ndarray.py
index 740f8bf5..2e2003fc 100644
--- a/ffcv/fields/ndarray.py
+++ b/ffcv/fields/ndarray.py
@@ -3,6 +3,7 @@
 from dataclasses import replace
 
 import numpy as np
+import torch as ch
 
 from .base import Field, ARG_TYPE
 from ..pipeline.operation import Operation
@@ -93,4 +94,21 @@ def encode(self, destination, field, malloc):
         data_region[:] = field.reshape(-1).view('<u1')
 
     def get_decoder_class(self) -> Type[Operation]:
-        return NDArrayDecoder
\ No newline at end of file
+        return NDArrayDecoder
+
+
+class TorchTensorField(NDArrayField):
+    """A subclass of :class:`~ffcv.fields.Field` supporting
+    multi-dimensional fixed size matrices of any torch type.
+    """
+    def __init__(self, dtype:ch.dtype, shape:Tuple[int, ...]):
+        self.dtype = dtype
+        self.shape = shape
+        dtype = ch.zeros(0, dtype=dtype).numpy().dtype
+
+        super().__init__(dtype, shape)
+
+
+    def encode(self, destination, field, malloc):
+        field = field.numpy()
+        return super().encode(destination, field, malloc)
diff --git a/ffcv/types.py b/ffcv/types.py
index 2b669ebc..3b2123f0 100644
--- a/ffcv/types.py
+++ b/ffcv/types.py
@@ -6,7 +6,8 @@
 from .fields.base import Field
 from .fields import (
     FloatField, IntField, RGBImageField,
-    BytesField, NDArrayField, JSONField
+    BytesField, NDArrayField, JSONField,
+    TorchTensorField
 )
 
 CURRENT_VERSION = 2
@@ -49,7 +50,8 @@
     2   : RGBImageField,
     3   : BytesField,
     4   : NDArrayField,
-    5   : JSONField
+    5   : JSONField,
+    6   : TorchTensorField
 }
 
 # Parse the fields descriptors from the header of the dataset


Test patch:
diff --git a/tests/test_array_field.py b/tests/test_array_field.py
index 3f6692dd..5ae359d3 100644
--- a/tests/test_array_field.py
+++ b/tests/test_array_field.py
@@ -3,19 +3,21 @@
 from collections import defaultdict
 from assertpy.assertpy import assert_that
 
+import torch as ch
 from assertpy import assert_that
 import numpy as np
 from torch.utils.data import Dataset
 from ffcv import DatasetWriter
-from ffcv.fields import IntField, NDArrayField
+from ffcv.fields import IntField, NDArrayField, TorchTensorField
 from ffcv import Loader
 
 class DummyActivationsDataset(Dataset):
 
-    def __init__(self, n_samples, shape):
+    def __init__(self, n_samples, shape, is_ch=False):
         self.n_samples = n_samples
         self.shape = shape
-        
+        self.is_ch = is_ch
+
     def __len__(self):
         return self.n_samples
 
@@ -23,7 +25,11 @@ def __getitem__(self, index):
         if index >= self.n_samples:
             raise IndexError()
         np.random.seed(index)
-        return index, np.random.randn(*self.shape).astype('<f4')
+        to_return = np.random.randn(*self.shape).astype('<f4')
+        if self.is_ch:
+            to_return = ch.from_numpy(to_return)
+
+        return index, to_return
 
 class TripleDummyActivationsDataset(Dataset):
 
@@ -43,13 +49,19 @@ def __getitem__(self, index):
         d3 = np.random.randn(*self.shape).astype('<f4')
         return index, d1, d2, d3
 
-def run_test(n_samples, shape):
+def run_test(n_samples, shape, is_ch=False):
     with NamedTemporaryFile() as handle:
         name = handle.name
-        dataset = DummyActivationsDataset(n_samples, shape)
+        dataset = DummyActivationsDataset(n_samples, shape, is_ch)
+
+        if is_ch:
+            field = TorchTensorField(ch.float32, shape)
+        else:
+            field = NDArrayField(np.dtype('<f4'), shape)
+
         writer = DatasetWriter(name, {
             'index': IntField(),
-            'activations': NDArrayField(np.dtype('<f4'), shape)
+            'activations': field
         }, num_workers=3)
 
         writer.from_indexed_dataset(dataset)
@@ -57,12 +69,18 @@ def run_test(n_samples, shape):
         loader = Loader(name, batch_size=3, num_workers=5)
         for ixes, activations in loader:
             for ix, activation in zip(ixes, activations):
-                assert_that(np.all(dataset[ix][1] == activation.numpy())).is_true()
+                d = dataset[ix][1]
+                if is_ch:
+                    d = d.numpy()
+                assert_that(np.all(d == activation.numpy())).is_true()
 
 
 def test_simple_activations():
     run_test(4096, (2048,))
 
+def test_simple_activations_ch():
+    run_test(4096, (2048,), True)
+
 def test_multi_fields():
     n_samples = 4096
     shape = (2048,)