Efficiency improvement in VecNorm, TensorDict and env tensor casting (p…

…ytorch#168)
ai4co · May 26, 2022 · f0ab441 · f0ab441
1 parent 476ca4e
commit f0ab441
Show file tree

Hide file tree

Showing 12 changed files with 265 additions and 193 deletions.
diff --git a/.circleci/unittest/linux/scripts/environment.yml b/.circleci/unittest/linux/scripts/environment.yml
@@ -4,9 +4,9 @@ channels:
 dependencies:
   - pip
   - cmake >= 3.18
+  - protobuf
   - pip:
     - hypothesis
-    - protobuf
     - future
     - cloudpickle
     - gym_retro

diff --git a/.circleci/unittest/linux_optdeps/scripts/environment.yml b/.circleci/unittest/linux_optdeps/scripts/environment.yml
@@ -6,7 +6,6 @@ dependencies:
   - cmake >= 3.18
   - pip:
     - hypothesis
-    - protobuf
     - future
     - cloudpickle
     - pytest

diff --git a/.circleci/unittest/linux_stable/scripts/environment.yml b/.circleci/unittest/linux_stable/scripts/environment.yml
@@ -5,9 +5,9 @@ dependencies:
   - pip
   - ninja
   - cmake >= 3.18
+  - protobuf
   - pip:
     - hypothesis
-    - protobuf
     - future
     - cloudpickle
     - gym_retro

diff --git a/test/smoke_test_deps.py b/test/smoke_test_deps.py
@@ -1,3 +1,6 @@
+import tempfile
+
+from torch.utils.tensorboard import SummaryWriter
 from torchrl.envs import DMControlEnv, GymEnv
 
 
@@ -20,3 +23,9 @@ def test_gym():
 def test_gym_pixels():
     env = GymEnv("ALE/Pong-v5", from_pixels=True)
     env.reset()
+
+
+def test_tb():
+    with tempfile.TemporaryDirectory() as directory:
+        writer = SummaryWriter(log_dir=directory)
+        writer.add_scalar("a", 1, 1)
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
@@ -600,14 +600,14 @@ def test_masking(self, td_name):
     @pytest.mark.skipif(
         torch.cuda.device_count() == 0, reason="No cuda device detected"
     )
-    @pytest.mark.parametrize("device", [0, "cuda:0", "cuda", torch.device("cuda:0")])
+    @pytest.mark.parametrize("device", [0, "cuda:0", torch.device("cuda:0")])
     def test_pin_memory(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)
         if td_name != "saved_td":
             td.pin_memory()
             td_device = td.to(device)
-            _device = torch.device("cuda:0")
+            _device = torch.device(device)
             assert td_device.device == _device
             assert td_device.clone().device == _device
             assert td_device is not td

diff --git a/test/test_trainer.py b/test/test_trainer.py
@@ -10,8 +10,15 @@
 
 import pytest
 import torch
-from tensorboard.backend.event_processing import event_accumulator
-from torch.utils.tensorboard import SummaryWriter
+
+try:
+    from tensorboard.backend.event_processing import event_accumulator
+    from torch.utils.tensorboard import SummaryWriter
+
+    _has_tb = True
+except ImportError:
+    _has_tb = False
+
 from torchrl.data import (
     TensorDict,
     TensorDictPrioritizedReplayBuffer,
@@ -214,6 +221,7 @@ def test_subsampler():
 
 
 @pytest.mark.skipif(not _has_gym, reason="No gym library")
+@pytest.mark.skipif(not _has_tb, reason="No tensorboard library")
 def test_recorder():
     with tempfile.TemporaryDirectory() as folder:
         writer = SummaryWriter(log_dir=folder)

diff --git a/torchrl/__init__.py b/torchrl/__init__.py
@@ -61,10 +61,14 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self._REG[self.name][2] = count + 1
 
     @staticmethod
-    def print():
+    def print(prefix=None):
         keys = list(timeit._REG)
         keys.sort()
         for name in keys:
-            print(
+            strings = []
+            if prefix:
+                strings.append(prefix)
+            strings.append(
                 f"{name} took {timeit._REG[name][0] * 1000:4.4} msec (total = {timeit._REG[name][1]} sec)"
             )
+            print(" -- ".join(strings))
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -5,7 +5,6 @@
 
 from __future__ import annotations
 
-from copy import deepcopy
 from dataclasses import dataclass
 from textwrap import indent
 from typing import (
@@ -210,24 +209,28 @@ def encode(self, val: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
 
         """
         if not isinstance(val, torch.Tensor):
-            try:
-                val = torch.tensor(val, dtype=self.dtype)
-            except ValueError:
-                val = torch.tensor(deepcopy(val), dtype=self.dtype)
+            if isinstance(val, np.ndarray) and not all(
+                stride > 0 for stride in val.strides
+            ):
+                val = val.copy()
+            val = torch.as_tensor(val, dtype=self.dtype, device=self.device)
         self.assert_is_in(val)
         return val
 
-    def to_numpy(self, val: torch.Tensor) -> np.ndarray:
+    def to_numpy(self, val: torch.Tensor, safe: bool = True) -> np.ndarray:
         """Returns the np.ndarray correspondent of an input tensor.
 
         Args:
             val (torch.Tensor): tensor to be transformed_in to numpy
+            safe (bool): boolean value indicating whether a check should be
+                performed on the value against the domain of the spec.
 
         Returns:
             a np.ndarray
 
         """
-        self.assert_is_in(val)
+        if safe:
+            self.assert_is_in(val)
         return val.detach().cpu().numpy()
 
     def index(self, index: INDEX_TYPING, tensor_to_index: torch.Tensor) -> torch.Tensor:
@@ -490,9 +493,8 @@ def encode(
         space: Optional[DiscreteBox] = None,
     ) -> torch.Tensor:
         if not isinstance(val, torch.Tensor):
-            val = torch.tensor(val)
+            val = torch.as_tensor(val, dtype=self.dtype, device=self.device)
 
-        val = torch.tensor(val, dtype=torch.long)
         if space is None:
             space = self.space
 
@@ -504,10 +506,11 @@ def encode(
         val = torch.nn.functional.one_hot(val, space.n).to(torch.long)
         return val
 
-    def to_numpy(self, val: torch.Tensor) -> np.ndarray:
-        if not isinstance(val, torch.Tensor):
-            raise NotImplementedError
-        self.assert_is_in(val)
+    def to_numpy(self, val: torch.Tensor, safe: bool = True) -> np.ndarray:
+        if safe:
+            if not isinstance(val, torch.Tensor):
+                raise NotImplementedError
+            self.assert_is_in(val)
         val = val.argmax(-1).cpu().numpy()
         if self.use_register:
             inv_reg = self.space.register.inverse()
@@ -794,7 +797,7 @@ def rand(self, shape: torch.Size = torch.Size([])) -> torch.Tensor:
 
     def encode(self, val: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
         if not isinstance(val, torch.Tensor):
-            val = torch.tensor(val)
+            val = torch.tensor(val, device=self.device)
 
         x = []
         for v, space in zip(val.unbind(-1), self.space):
@@ -809,7 +812,9 @@ def _split(self, val: torch.Tensor) -> torch.Tensor:
         vals = val.split([space.n for space in self.space], dim=-1)
         return vals
 
-    def to_numpy(self, val: torch.Tensor) -> np.ndarray:
+    def to_numpy(self, val: torch.Tensor, safe: bool = True) -> np.ndarray:
+        if safe:
+            self.assert_is_in(val)
         vals = self._split(val)
         out = torch.stack([val.argmax(-1) for val in vals], -1).numpy()
         return out