Skip to content

Commit

Permalink
Clear predict_progress_bar in ProgressBar.__getstate__ (#7608)
Browse files Browse the repository at this point in the history
Co-authored-by: Yifu Wang <yifuwang@2012@gmail.com>
Co-authored-by: Justus Schock <12886177+justusschock@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Carlos Mocholi <carlossmocholi@gmail.com>
  • Loading branch information
5 people committed May 25, 2021
1 parent 6b81cd8 commit 2573470
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 0 deletions.
8 changes: 8 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).





- Fixed `ProgressBar` pickling after calling `trainer.predict` ([#7608](https://github.com/PyTorchLightning/pytorch-lightning/pull/7608))




## [1.3.2] - 2021-05-18

### Changed
Expand Down
2 changes: 2 additions & 0 deletions pytorch_lightning/callbacks/progress.py
Original file line number Diff line number Diff line change
Expand Up @@ -283,13 +283,15 @@ def __init__(self, refresh_rate: int = 1, process_position: int = 0):
self.main_progress_bar = None
self.val_progress_bar = None
self.test_progress_bar = None
self.predict_progress_bar = None

def __getstate__(self):
# can't pickle the tqdm objects
state = self.__dict__.copy()
state['main_progress_bar'] = None
state['val_progress_bar'] = None
state['test_progress_bar'] = None
state['predict_progress_bar'] = None
return state

@property
Expand Down
15 changes: 15 additions & 0 deletions tests/callbacks/test_progress_bar.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import pickle
import sys
from typing import Optional, Union
from unittest import mock
Expand Down Expand Up @@ -482,3 +483,17 @@ def test_progress_bar_print_disabled(tqdm_write, mock_print, tmpdir):
call("test_step"),
])
tqdm_write.assert_not_called()


def test_progress_bar_can_be_pickled():
bar = ProgressBar()
trainer = Trainer(fast_dev_run=True, callbacks=[bar], max_steps=1)
model = BoringModel()

pickle.dumps(bar)
trainer.fit(model)
pickle.dumps(bar)
trainer.test(model)
pickle.dumps(bar)
trainer.predict(model)
pickle.dumps(bar)

0 comments on commit 2573470

Please sign in to comment.