[Feature]: faster meta-tensor API for TensorDict (pytorch#272)

ai4co · Jul 14, 2022 · e5bea04 · e5bea04
1 parent 4c81c6f
commit e5bea04
Show file tree

Hide file tree

Showing 5 changed files with 181 additions and 119 deletions.
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
@@ -63,7 +63,7 @@ def test_tensordict_set(device):
         torch.randn(4, 5, 1, 2, dtype=torch.double, device=device),
         inplace=False,
     )
-    assert td._tensordict_meta["key1"].shape == td._tensordict["key1"].shape
+    assert td._dict_meta["key1"].shape == td._tensordict["key1"].shape
 
 
 def test_pad():

diff --git a/torchrl/__init__.py b/torchrl/__init__.py
@@ -114,6 +114,7 @@ def __init__(self, fun=lambda x: x):
 
     def __missing__(self, key):
         value = self.fun(key)
+        self[key] = value
         return value
 
 

diff --git a/torchrl/data/tensordict/memmap.py b/torchrl/data/tensordict/memmap.py
@@ -106,6 +106,8 @@ class MemmapTensor(object):
 
     """
 
+    requires_grad = False
+
     def __init__(
         self,
         elem: Union[torch.Tensor, MemmapTensor],

diff --git a/torchrl/data/tensordict/metatensor.py b/torchrl/data/tensordict/metatensor.py
@@ -65,7 +65,7 @@ def __init__(
         self,
         *shape: Union[int, torch.Tensor, "MemmapTensor"],
         device: Optional[DEVICE_TYPING] = "cpu",
-        dtype: torch.dtype = torch.get_default_dtype(),
+        dtype: torch.dtype = None,
         requires_grad: bool = False,
         _is_shared: Optional[bool] = None,
         _is_memmap: Optional[bool] = None,
@@ -81,8 +81,8 @@ def __init__(
                 _is_memmap = isinstance(tensor, MemmapTensor)
             device = tensor.device if not tensor.is_meta else device
             if _is_tensordict is None:
-                _is_tensordict = not isinstance(tensor, (MemmapTensor, torch.Tensor))
-            if isinstance(tensor, (MemmapTensor, torch.Tensor)):
+                _is_tensordict = not _is_memmap and not isinstance(tensor, torch.Tensor)
+            if not _is_tensordict:
                 dtype = tensor.dtype
             else:
                 dtype = None
@@ -97,11 +97,11 @@ def __init__(
         if not isinstance(shape, torch.Size):
             shape = torch.Size(shape)
         self.shape = shape
-        self.device = torch.device(device)
-        self.dtype = dtype
+        self.device = device
+        self.dtype = dtype if dtype is not None else torch.get_default_dtype()
         self.requires_grad = requires_grad
         self._ndim = len(shape)
-        self._numel = np.prod(shape)
+        self._numel = None
         self._is_shared = bool(_is_shared)
         self._is_memmap = bool(_is_memmap)
         self._is_tensordict = bool(_is_tensordict)
@@ -155,6 +155,8 @@ def is_tensordict(self) -> bool:
         return self._is_tensordict
 
     def numel(self) -> int:
+        if self._numel is None:
+            self._numel = np.prod(self.shape)
         return self._numel
 
     def ndimension(self) -> int:
-Original file line number
+Diff line change
@@ Expand Up / @@ -106,6 +106,8 @@ class MemmapTensor(object): @@
         """
+        requires_grad = False
         def __init__(
             self,
             elem: Union[torch.Tensor, MemmapTensor],
@@ Expand Down @@