PyTorch Lightning & Jupyter Compatibility

ivon/_ivon.py

@@ -74,7 +74,7 @@ def __init__(
         self.mc_samples = mc_samples
         self.hess_approx = hess_approx
         self.sync = sync
-        self._numel, self._device, self._dtype = self._get_param_configs()
+        self._numel, self._dtype = self._get_param_configs()
         self.current_step = 0
         self.debias = debias
         self.rescale_lr = rescale_lr
@@ -90,14 +90,8 @@ def _get_param_configs(self):
             pg["numel"] = sum(p.numel() for p in pg["params"] if p is not None)
             all_params += [p for p in pg["params"] if p is not None]
         if len(all_params) == 0:
-            return 0, torch.device("cpu"), torch.get_default_dtype()
-        devices = {p.device for p in all_params}
-        if len(devices) > 1:
-            raise ValueError(
-                "Parameters are on different devices: "
-                f"{[str(d) for d in devices]}"
-            )
-        device = next(iter(devices))
+            return 0, torch.get_default_dtype()
+
         dtypes = {p.dtype for p in all_params}
         if len(dtypes) > 1:
             raise ValueError(
@@ -106,7 +100,7 @@ def _get_param_configs(self):
             )
         dtype = next(iter(dtypes))
         total = sum(pg["numel"] for pg in self.param_groups)
-        return total, device, dtype
+        return total, dtype
 
     def _reset_samples(self):
         self.state['count'] = 0
@@ -118,18 +112,22 @@ def _init_buffers(self):
         for group in self.param_groups:
             hess_init, numel = group["hess_init"], group["numel"]
 
+            group_device = IVON._find_group_device(group, check_same_device=True)
+
             group["momentum"] = torch.zeros(
-                numel, device=self._device, dtype=self._dtype
+                numel, device=group_device, dtype=self._dtype
             )
             group["hess"] = torch.zeros(
-                numel, device=self._device, dtype=self._dtype
+                numel, device=group_device, dtype=self._dtype
             ).add(torch.as_tensor(hess_init))
 
     @contextmanager
     def sampled_params(self, train: bool = False):
         param_avg, noise = self._sample_params()
-        yield
-        self._restore_param_average(train, param_avg, noise)
+        try:
+            yield
+        finally:
+            self._restore_param_average(train, param_avg, noise)
 
     def _restore_param_average(
         self, train: bool, param_avg: Tensor, noise: Tensor
@@ -174,7 +172,8 @@ def step(self, closure: ClosureType = None) -> Optional[Tensor]:
             losses = []
             for _ in range(self.mc_samples):
                 with torch.enable_grad():
-                    loss = closure()
+                    with self.sampled_params(train=True):
+                        loss = closure()
                 losses.append(loss)
             loss = sum(losses) / self.mc_samples
         if self.sync and dist.is_initialized():  # explicit sync
@@ -197,8 +196,13 @@ def _sample_params(self) -> Tuple[Tensor, Tensor]:
         offset = 0
         for group in self.param_groups:
             gnumel = group["numel"]
+
+            group_device = IVON._find_group_device(group, check_same_device=False)
+
+            group["hess"] = group["hess"].to(group_device)
+
             noise_sample = (
-                torch.randn(gnumel, device=self._device, dtype=self._dtype)
+                torch.randn(gnumel, device=group_device, dtype=self._dtype)
                 / (
                     group["ess"] * (group["hess"] + group["weight_decay"])
                 ).sqrt()
@@ -228,6 +232,8 @@ def _update(self):
 
         offset = 0
         for group in self.param_groups:
+            group_device = IVON._find_group_device(group, check_same_device=True)
+
             lr = group["lr"]
             b1 = group["beta1"]
             b2 = group["beta2"]
@@ -237,6 +243,8 @@ def _update(self):
                 [p.flatten() for p in group["params"] if p is not None], 0
             )
 
+            group["momentum"] = group["momentum"].to(group_device)
+
             group["momentum"] = self._new_momentum(
                 self.state["avg_grad"][pg_slice], group["momentum"], b1
             )
@@ -307,3 +315,26 @@ def _new_param_averages(
             min=-clip_radius,
             max=clip_radius,
         )
+
+    @staticmethod
+    def _find_group_device(group, check_same_device=True):
+        group_device = None
+        for p in group["params"]:
+            if p is None:
+                continue
+
+            if group_device is None:
+                group_device = p.device
+            elif group_device != p.device:
+                raise ValueError(
+                    "Parameters are on different devices: "
+                    f"{group_device} and {p.device}"
+                )
+
+            if not check_same_device:
+                return group_device
+
+        if group_device is None:
+            return torch.device("cpu")
+
+        return group_device