Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Refactoring] Minor tweaks to recorder and logger #489

Merged
merged 2 commits into from
Sep 26, 2022
Merged

[Refactoring] Minor tweaks to recorder and logger #489

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
94d8a07
Minor tweaks to recorder and logger
nicolas-dufour Sep 26, 2022
524943e
Removed cuda empty cache
nicolas-dufour Sep 26, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion torchrl/record/recorder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,7 @@ def _call(self, td: TensorDictBase) -> TensorDictBase:
if self.count % self.skip == 0:
_td = td
if self.keys_in:
_td = td.select(*self.keys_in).clone()
_td = td.select(*self.keys_in).to_tensordict()
self.td.append(_td)
return td

Expand Down
2 changes: 2 additions & 0 deletions torchrl/trainers/helpers/logger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,3 +24,5 @@ class LoggerConfig:
# number of steps in validation rollouts. " "Default=1000.
recorder_log_keys: Any = field(default_factory=lambda: ["reward"])
# Keys to log in the recorder
offline_logging: bool = True
# If True, Wandb will do the logging offline
1 change: 1 addition & 0 deletions torchrl/trainers/loggers/wandb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,7 @@ def log_video(self, name: str, video: Tensor, **kwargs) -> None:
f"be silenced from now on but the values will keep being incremented."
)
step = self._prev_video_step + 1
self._prev_video_step = step if step is not None else self._prev_video_step + 1
self.experiment.log(
{name: wandb.Video(video, fps=fps, format=format)}, step=step, **kwargs
)
Expand Down
6 changes: 2 additions & 4 deletions torchrl/trainers/trainers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@
from torchrl.data.tensordict.tensordict import TensorDictBase, pad
from torchrl.data.utils import expand_right, DEVICE_TYPING
from torchrl.envs.common import EnvBase
from torchrl.envs.transforms import TransformedEnv
from torchrl.envs.utils import set_exploration_mode
from torchrl.modules import TensorDictModule
from torchrl.objectives.costs.common import LossModule
Expand Down Expand Up @@ -905,7 +904,7 @@ def __init__(
frame_skip: int,
policy_exploration: TensorDictModule,
recorder: EnvBase,
exploration_mode: str = "mean",
exploration_mode: str = "random",
log_keys: Optional[List[str]] = None,
out_keys: Optional[Dict[str, str]] = None,
suffix: Optional[str] = None,
Expand Down Expand Up @@ -933,12 +932,11 @@ def __init__(
def __call__(self, batch: TensorDictBase) -> Dict:
out = None
if self._count % self.record_interval == 0:
torch.cuda.empty_cache()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we added this because of weird cuda mem issues which I believe are solved now. Can we remove this and try to train something like dreamer while looking at the mem on devices?
Asking this because this operation takes a considerable amount of time :)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes i agree let's remove this

with set_exploration_mode(self.exploration_mode):
if isinstance(self.policy_exploration, torch.nn.Module):
self.policy_exploration.eval()
self.recorder.eval()
if isinstance(self.recorder, TransformedEnv):
self.recorder.transform.eval()
td_record = self.recorder.rollout(
policy=self.policy_exploration,
max_steps=self.record_frames,
Expand Down