From 6c00565778e5815e4589afc7499aafbd020535ae Mon Sep 17 00:00:00 2001 From: Ole Petersen <56505957+peteole@users.noreply.github.com> Date: Wed, 15 May 2024 15:59:32 +0200 Subject: [PATCH] Fix memory leak in base_class.py (#1908) * Fix memory leak in base_class.py Loading the data return value is not necessary since it is unused. Loading the data causes a memory leak through the ep_info_buffer variable. I found this while loading a PPO learner from storage on a multi-GPU system since the ep_info_buffer is loaded to the memory location it was on while it was saved to disk, instead of the target loading location, and is then not cleaned up. * Update changelog.rst * Update changelog --------- Co-authored-by: Antonin RAFFIN --- docs/misc/changelog.rst | 8 +++++--- stable_baselines3/common/base_class.py | 2 +- stable_baselines3/version.txt | 2 +- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index a4c333e08..d77388534 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,7 +3,7 @@ Changelog ========== -Release 2.4.0a0 (WIP) +Release 2.4.0a1 (WIP) -------------------------- Breaking Changes: @@ -14,6 +14,8 @@ New Features: Bug Fixes: ^^^^^^^^^^ +- Fixed memory leak when loading learner from storage, ``set_parameters()`` does not try to load the object data anymore + and only loads the PyTorch parameters (@peteole) - Cast type in compute gae method to avoid error when using torch compile (@amjames) `SB3-Contrib`_ @@ -30,7 +32,7 @@ Deprecations: Others: ^^^^^^^ -- Fix various typos (@cschindlbeck) +- Fixed various typos (@cschindlbeck) Bug Fixes: ^^^^^^^^^^ @@ -1659,4 +1661,4 @@ And all the contributors: @anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong @ReHoss @DavyMorgan @luizapozzobon @Bonifatius94 @theSquaredError @harveybellini @DavyMorgan @FieteO @jonasreiher @npit @WeberSamuel @troiganto @lutogniew @lbergmann1 @lukashass @BertrandDecoster @pseudo-rnd-thoughts @stefanbschneider @kyle-he @PatrickHelm @corentinlger -@marekm4 @stagoverflow @rushitnshah @markscsmith @NickLucche @cschindlbeck +@marekm4 @stagoverflow @rushitnshah @markscsmith @NickLucche @cschindlbeck @peteole diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 054e58a1f..4be61d65e 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -592,7 +592,7 @@ def set_parameters( if isinstance(load_path_or_dict, dict): params = load_path_or_dict else: - _, params, _ = load_from_zip_file(load_path_or_dict, device=device) + _, params, _ = load_from_zip_file(load_path_or_dict, device=device, load_data=False) # Keep track which objects were updated. # `_get_torch_save_params` returns [params, other_pytorch_variables]. diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index e96f44fb3..48adc0106 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -2.4.0a0 +2.4.0a1