Snapshotter for torch

src/garage/experiment/snapshotter.py

@@ -3,8 +3,12 @@
 import errno
 import os
 import pathlib
+import sys
 
 import cloudpickle
+from dowel import logger
+
+# pylint: disable=no-name-in-module
 
 SnapshotConfig = collections.namedtuple(
     'SnapshotConfig', ['snapshot_dir', 'snapshot_mode', 'snapshot_gap'])
@@ -82,6 +86,7 @@ def snapshot_gap(self):
         """
         return self._snapshot_gap
 
+    # pylint: disable=too-many-branches
     def save_itr_params(self, itr, params):
         """Save the parameters if at the right iteration.
 
@@ -94,8 +99,13 @@ def save_itr_params(self, itr, params):
                 "gap_overwrite", "gap_and_last", or "none".
 
         """
+        # pylint: disable=import-outside-toplevel
+        torch = False
+        if 'torch' in sys.modules:
+            import torch
+            from garage.torch import global_device
         file_name = None
-
+        # pylint: enable=import-outside-toplevel
         if self._snapshot_mode == 'all':
             file_name = os.path.join(self._snapshot_dir, 'itr_%d.pkl' % itr)
         elif self._snapshot_mode == 'gap_overwrite':
@@ -113,17 +123,32 @@ def save_itr_params(self, itr, params):
                 file_name = os.path.join(self._snapshot_dir,
                                          'itr_%d.pkl' % itr)
             file_name_last = os.path.join(self._snapshot_dir, 'params.pkl')
-            with open(file_name_last, 'wb') as file:
-                cloudpickle.dump(params, file)
+            if torch:
+                torch.save(params, file_name_last, pickle_module=cloudpickle)
+            else:
+                with open(file_name_last, 'wb') as file:
+                    cloudpickle.dump(params, file)
         elif self._snapshot_mode == 'none':
             pass
         else:
             raise ValueError('Invalid snapshot mode {}'.format(
                 self._snapshot_mode))
 
         if file_name:
-            with open(file_name, 'wb') as file:
-                cloudpickle.dump(params, file)
+            if torch:
+
+                class _pickle_module:
+                    dump = cloudpickle.dump
+                    Pickler = cloudpickle.CloudPickler
+
+                params['global_device'] = global_device()
+                torch.save(params,
+                           file_name,
+                           pickle_module=_pickle_module,
+                           _use_new_zipfile_serialization=False)
+            else:
+                with open(file_name, 'wb') as file:
+                    cloudpickle.dump(params, file)
 
     def load(self, load_dir, itr='last'):
         # pylint: disable=no-self-use
@@ -145,6 +170,12 @@ def load(self, load_dir, itr='last'):
             NotAFileError: If the snapshot exists but is not a file.
 
         """
+        torch = False
+        # pylint: disable=import-outside-toplevel
+        if 'torch' in sys.modules:
+            import torch
+            from garage.torch import global_device
+        # pylint: enable=import-outside-toplevel
         if isinstance(itr, int) or itr.isdigit():
             load_from_file = os.path.join(load_dir, 'itr_{}.pkl'.format(itr))
         else:
@@ -165,7 +196,13 @@ def load(self, load_dir, itr='last'):
 
         if not os.path.isfile(load_from_file):
             raise NotAFileError('File not existing: ', load_from_file)
-
+        if torch:
+            device = global_device()
+            params = torch.load(load_from_file, map_location=device)
+            origin_device = params['global_device']
+            del params['global_device']
+            logger.log(f'Resuming experiment from {origin_device} on {device}')
+            return params
         with open(load_from_file, 'rb') as file:
             return cloudpickle.load(file)
 

tests/fixtures/__init__.py

@@ -1,9 +1,10 @@
 """Test fixtures."""
 # yapf: disable
-from tests.fixtures.fixtures import (snapshot_config,
-                                     TfGraphTestCase,
-                                     TfTestCase)
+from tests.fixtures.fixtures import (reset_gpu_mode, snapshot_config,
+                                     TfGraphTestCase, TfTestCase)
 
 # yapf: enable
 
-__all__ = ['snapshot_config', 'TfGraphTestCase', 'TfTestCase']
+__all__ = [
+    'reset_gpu_mode', 'snapshot_config', 'TfGraphTestCase', 'TfTestCase'
+]

tests/fixtures/fixtures.py

@@ -9,6 +9,7 @@
 
 from garage.experiment import deterministic
 from garage.experiment.snapshotter import SnapshotConfig
+from garage.torch import set_gpu_mode
 
 from tests.fixtures.logger import NullOutput
 
@@ -64,3 +65,8 @@ def teardown_method(self):
         del self.graph
         del self.sess
         gc.collect()
+
+
+def reset_gpu_mode():
+    """Reset mode to CPU after test."""
+    set_gpu_mode(False)

tests/garage/experiment/test_snapshotter.py

@@ -22,6 +22,7 @@
 class TestSnapshotter:
 
     def setup_method(self):
+        # pylint: disable=consider-using-with
         self.temp_dir = tempfile.TemporaryDirectory()
 
     def teardown_method(self):
@@ -44,6 +45,7 @@ def test_snapshotter(self, mode, files):
             assert osp.exists(filename)
             with open(filename, 'rb') as pkl_file:
                 data = pickle.load(pkl_file)
+                snapshot_data[num]['global_device'] = None
                 assert data == snapshot_data[num]
 
     def test_gap_overwrite(self):
@@ -60,7 +62,7 @@ def test_gap_overwrite(self):
         assert osp.exists(filename)
         with open(filename, 'rb') as pkl_file:
             data = pickle.load(pkl_file)
-            assert data == snapshot_data[1]
+            assert data == {'global_device': None, 'testparam': 4}
 
     def test_invalid_snapshot_mode(self):
         with pytest.raises(ValueError):