diff --git a/dreamer.py b/dreamer.py index 0ffc536..bbc4d8f 100644 --- a/dreamer.py +++ b/dreamer.py @@ -311,7 +311,7 @@ def main(config): logger.warning(f"Loaded model from {logdir / 'latest.pt'}") # make sure eval will be executed once after config.steps - with tqdm(total=config.steps + config.eval_every, unit='step') as pbar: + with tqdm(total=config.steps + config.eval_every, unit='step', mininterval=60) as pbar: while agent._step < config.steps + config.eval_every: tlogger.write() if config.eval_episode_num > 0: diff --git a/justfile b/justfile index e441114..125e3de 100644 --- a/justfile +++ b/justfile @@ -6,7 +6,7 @@ export TQDM_MININTERVAL := "30" main: . ./.venv/bin/activate - python dreamer.py --configs craftax_smaller --logdir ./logdir/crafterer + python dreamer.py --configs craftax_smaller --logdir ./logdir/craftax_smaller logs: tensorboard --logdir logdir/craftax diff --git a/models.py b/models.py index 84b3f52..de3ca75 100644 --- a/models.py +++ b/models.py @@ -99,7 +99,7 @@ class WorldModel(nn.Module): opt=config.opt, use_amp=self._use_amp, ) - logger.info(f"World Model\n{summary(self, row_settings=['var_names'],)}") + logger.info(f"World Model\n{summary(self, row_settings=['depth', 'var_names'], verbose=False)}") # other losses are scaled by 1.0. self._scales = dict( reward=config.reward_head["loss_scale"], diff --git a/networks.py b/networks.py index 915c82a..b109cc3 100644 --- a/networks.py +++ b/networks.py @@ -12,7 +12,7 @@ from einops import rearrange from torchinfo import summary def my_summary(model, input_data): - return summary(model, input_data, col_names=('input_size', 'output_size', 'num_params', 'mult_adds'), verbose=0, row_settings=['depth', 'var_names', 'ascii_only']) + return summary(model, input_data, col_names=('input_size', 'output_size', 'num_params', 'mult_adds'), verbose=0, row_settings=['depth', 'var_names']) class RSSM(nn.Module): diff --git a/tools.py b/tools.py index ef05a82..6d48135 100644 --- a/tools.py +++ b/tools.py @@ -973,7 +973,7 @@ def set_seed_everywhere(seed): def enable_deterministic_run(): os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - torch.backends.cudnn.benchmark = False + # torch.backends.cudnn.benchmark = False torch.use_deterministic_algorithms(True)