Resolve #656 by fixing progressbar pickling (#663)

The fix is to ignore the tqdm instance in the returned state of
the progress bar callback.

Author: ottonemo

CHANGES.md

@@ -26,6 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Fixed a bug where `CyclicLR` scheduler would update during both training and validation rather than just during training.
 - Fixed a bug introduced by moving the `optimizer.zero_grad()` call outside of the train step function, making it incompatible with LBFGS and other optimizers that call the train step several times per batch (#636)
+- Fixed pickling of the `ProgressBar` callback (#656)
 
 ## [0.8.0] - 2019-04-11
 

skorch/callbacks/logging.py

@@ -556,8 +556,8 @@ def _get_postfix_dict(self, net):
 
     # pylint: disable=attribute-defined-outside-init
     def on_batch_end(self, net, **kwargs):
-        self.pbar.set_postfix(self._get_postfix_dict(net), refresh=False)
-        self.pbar.update()
+        self.pbar_.set_postfix(self._get_postfix_dict(net), refresh=False)
+        self.pbar_.update()
 
     # pylint: disable=attribute-defined-outside-init, arguments-differ
     def on_epoch_begin(self, net, dataset_train=None, dataset_valid=None, **kwargs):
@@ -576,12 +576,19 @@ def on_epoch_begin(self, net, dataset_train=None, dataset_valid=None, **kwargs):
                 batches_per_epoch = len(net.history[-2, 'batches'])
 
         if self._use_notebook():
-            self.pbar = tqdm.tqdm_notebook(total=batches_per_epoch, leave=False)
+            self.pbar_ = tqdm.tqdm_notebook(total=batches_per_epoch, leave=False)
         else:
-            self.pbar = tqdm.tqdm(total=batches_per_epoch, leave=False)
+            self.pbar_ = tqdm.tqdm(total=batches_per_epoch, leave=False)
 
     def on_epoch_end(self, net, **kwargs):
-        self.pbar.close()
+        self.pbar_.close()
+
+    def __getstate__(self):
+        # don't save away the temporary pbar_ object which gets created on
+        # epoch begin anew anyway. This avoids pickling errors with tqdm.
+        state = self.__dict__.copy()
+        del state['pbar_']
+        return state
 
 
 def rename_tensorboard_key(key):

skorch/tests/callbacks/test_logging.py

@@ -541,6 +541,20 @@ def test_different_count_schemes(
         for i, total in enumerate(expected_total):
             assert tqdm_mock.call_args_list[i][1]['total'] == total
 
+    def test_pickle(self, net_cls, progressbar_cls, data):
+        # pickling was an issue since TQDM progress bar instances cannot
+        # be pickled. Test pickling and restoration.
+        import pickle
+
+        net = net_cls(callbacks=[
+            progressbar_cls(),
+        ])
+        net.fit(*data)
+        dump = pickle.dumps(net)
+
+        net = pickle.loads(dump)
+        net.fit(*data)
+
 
 @pytest.mark.skipif(
     not tensorboard_installed, reason='tensorboard is not installed')
@@ -776,4 +790,4 @@ def forward(self, b, e, c, d, a, **kwargs):
         net.fit(X_dict, y)
 
         # is not empty
-        assert os.listdir(path)
+        assert os.listdir(path)