add torch accessor

Author: Joseph Hamman

xbatcher/accessors.py

@@ -24,3 +24,21 @@ def generator(self, *args, **kwargs):
             Keyword arguments to pass to the `BatchGenerator` constructor.
         '''
         return BatchGenerator(self._obj, *args, **kwargs)
+
+
+@xr.register_dataarray_accessor('torch')
+class TorchAccessor:
+    def __init__(self, xarray_obj):
+        self._obj = xarray_obj
+
+    def to_tensor(self):
+        """Convert this DataArray to a torch.Tensor"""
+        import torch
+
+        return torch.tensor(self._obj.data)
+
+    def to_named_tensor(self):
+        """Convert this DataArray to a torch.Tensor with named dimensions"""
+        import torch
+
+        return torch.tensor(self._obj.data, names=self._obj.dims)

xbatcher/loaders/torch.py

@@ -49,8 +49,8 @@ def __getitem__(self, idx) -> Tuple[Any, Any]:
 
         # TODO: figure out the dataset -> array workflow
         # currently hardcoding a variable name
-        X_batch = self.X_generator[idx]['x'].data
-        y_batch = self.y_generator[idx]['y'].data
+        X_batch = self.X_generator[idx]['x'].torch.to_tensor()
+        y_batch = self.y_generator[idx]['y'].torch.to_tensor()
 
         if self.transform:
             X_batch = self.transform(X_batch)
@@ -80,4 +80,4 @@ def __init__(
 
     def __iter__(self):
         for xb, yb in zip(self.X_generator, self.y_generator):
-            yield (xb['x'].data, yb['y'].data)
+            yield (xb['x'].torch.to_tensor(), yb['y'].torch.to_tensor())

xbatcher/tests/test_accessors.py

@@ -38,3 +38,25 @@ def test_batch_accessor_da(sample_ds_3d):
     assert isinstance(bg_acc, BatchGenerator)
     for batch_class, batch_acc in zip(bg_class, bg_acc):
         assert batch_class.equals(batch_acc)
+
+
+def test_torch_to_tensor(sample_ds_3d):
+    torch = pytest.importorskip('torch')
+
+    da = sample_ds_3d['foo']
+    t = da.torch.to_tensor()
+    assert isinstance(t, torch.Tensor)
+    assert t.names == (None, None, None)
+    assert t.shape == da.shape
+    np.testing.assert_array_equal(t, da.values)
+
+
+def test_torch_to_named_tensor(sample_ds_3d):
+    torch = pytest.importorskip('torch')
+
+    da = sample_ds_3d['foo']
+    t = da.torch.to_named_tensor()
+    assert isinstance(t, torch.Tensor)
+    assert t.names == da.dims
+    assert t.shape == da.shape
+    np.testing.assert_array_equal(t, da.values)

xbatcher/tests/test_torch_loaders.py

@@ -37,7 +37,7 @@ def test_map_dataset(ds_xy):
     # test __getitem__
     x_batch, y_batch = dataset[0]
     assert len(x_batch) == len(y_batch)
-    assert isinstance(x_batch, np.ndarray)
+    assert isinstance(x_batch, torch.Tensor)
 
     # test __len__
     assert len(dataset) == len(x_gen)