updated notebook

Author: andreasmardt

Attention_and_coarse_graining.ipynb

@@ -27,16 +27,6 @@
     "from tqdm.notebook import tqdm  # progress bar"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from torch.utils.tensorboard import SummaryWriter\n",
-    "writer = SummaryWriter()"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -280,6 +270,25 @@
     "loader_val = DataLoader(val_data, batch_size=len(val_data), shuffle=False)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Can be used to record the training performance with tensorboard\n",
+    "# it is not necessary for training or using the methods\n",
+    "# if you do not wish to install the additional package just leave the flag to false!\n",
+    "tensorboard_installed = False\n",
+    "if tensorboard_installed:\n",
+    "    from torch.utils.tensorboard import SummaryWriter\n",
+    "    writer = SummaryWriter()\n",
+    "    input_model, _ = next(iter(loader_train))\n",
+    "    writer.add_graph(lobe, input_to_model=input_model.to(device))\n",
+    "else:\n",
+    "    writer=None"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -392,7 +401,7 @@
    "outputs": [],
    "source": [
     "# train only for the matrix S\n",
-    "deepmsm.fit(loader_train, n_epochs=100, validation_loader=loader_val, train_mode='s')\n",
+    "deepmsm.fit(loader_train, n_epochs=100, validation_loader=loader_val, train_mode='s', tb_writer=writer)\n",
     "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
     "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
     "plt.xlabel('step')\n",
@@ -407,7 +416,7 @@
    "outputs": [],
    "source": [
     "# Train for S and u\n",
-    "deepmsm.fit(loader_train, n_epochs=100, validation_loader=loader_val, train_mode='us')\n",
+    "deepmsm.fit(loader_train, n_epochs=1000, validation_loader=loader_val, train_mode='us', tb_writer=writer)\n",
     "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
     "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
     "plt.xlabel('step')\n",
@@ -424,7 +433,7 @@
    "outputs": [],
    "source": [
     "# Train for chi, u, and S in an iterative manner\n",
-    "deepmsm.fit_routine(loader_train, n_epochs=50, validation_loader=loader_val, rel=0.001, reset_u=False, max_iter=1000)\n",
+    "deepmsm.fit_routine(loader_train, n_epochs=50, validation_loader=loader_val, rel=0.001, reset_u=False, max_iter=1000, tb_writer=writer)\n",
     "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
     "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
     "plt.xlabel('step')\n",
@@ -445,7 +454,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "plot_mask(data=data[0], lobe=lobe_msm, mask=mask, mask_const=mask_const, device=device, vmax=0.5)"
+    "plot_mask(data=data[0], lobe=lobe_msm, mask=mask, mask_const=mask_const, device=device, vmax=0.25)"
    ]
   },
   {
@@ -689,9 +698,9 @@
    "outputs": [],
    "source": [
     "for _ in range(5):\n",
-    "    model_msm_final = deepmsm.fit(loader_train, n_epochs=1000, validation_loader=loader_val, train_mode='s').fetch_model()\n",
+    "    model_msm_final = deepmsm.fit(loader_train, n_epochs=1000, validation_loader=loader_val, train_mode='s', tb_writer=writer).fetch_model()\n",
     "    # train for u and S\n",
-    "    model_msm_final = deepmsm.fit(loader_train, n_epochs=100, validation_loader=loader_val, train_mode='us').fetch_model()"
+    "    model_msm_final = deepmsm.fit(loader_train, n_epochs=100, validation_loader=loader_val, train_mode='us', tb_writer=writer).fetch_model()"
    ]
   },
   {

deepmsm.py

@@ -12,6 +12,7 @@
 from deeptime.util.torch import map_data
 from deeptime.markov.tools.analysis import pcca_memberships
 
+CLIP_VALUE = 1.
 
 def symeig_reg(mat, epsilon: float = 1e-6, mode='regularize', eigenvectors=True) \
         -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
@@ -710,6 +711,14 @@ def transform(self, data, **kwargs):
             out.append(net(self.mask(data_tensor)).cpu().numpy())
         return out if len(out) > 1 else out[0]
 
+    def get_mu(self, data_t):
+        self._lobe.eval()
+        net = self._lobe
+        with torch.no_grad():
+            x_t = net(self.mask(torch.Tensor(data_t).to(self._device)))
+            mu = self._ulayer(x_t, x_t, return_mu=True)[-1] # use dummy x_0
+        return mu.detach().to('cpu').numpy()
+    
     def get_transition_matrix(self, data_0, data_t):
         self._lobe.eval()
         net = self._lobe
@@ -776,11 +785,11 @@ def observables(self, data_0, data_t, data_ev=None, data_ac=None, state1=None, s
         with torch.no_grad():
             x_0 = net(self.mask(torch.Tensor(data_0).to(self._device)))
             x_t = net(self.mask(torch.Tensor(data_t).to(self._device)))
-            output_u = self.ulayer(x_0, x_t, return_mu=return_mu)
+            output_u = self._ulayer(x_0, x_t, return_mu=return_mu)
             if return_mu:
                 mu = output_u[5]
             Sigma = output_u[4]
-            output_S = self.slayer(*output_u[:5], return_K=return_K, return_S=return_S)
+            output_S = self._slayer(*output_u[:5], return_K=return_K, return_S=return_S)
             if return_K:
                 K = output_S[1]
             if return_S:
@@ -789,14 +798,14 @@ def observables(self, data_0, data_t, data_ev=None, data_ac=None, state1=None, s
             if data_ev is not None:
                 x_ev = torch.Tensor(data_ev).to(self._device)
                 ev_est = obs_ev(x_ev,mu)
-                ret.append(ev_est)
+                ret.append(ev_est.detach().to('cpu').numpy())
             if data_ac is not None:
                 x_ac = torch.Tensor(data_ac).to(self._device)
                 ac_est = obs_ac(x_ac, mu, x_t, K, Sigma)
-                ret.append(ac_est)
+                ret.append(ac_est.detach().to('cpu').numpy())
             if state1 is not None:
                 its_est = get_process_eigval(S, Sigma, state1, state2, epsilon=self._epsilon, mode=self._mode)
-                ret.append(its_est)
+                ret.append(its_est.detach().to('cpu').numpy())
         return ret
 
 class DeepMSM(DLEstimatorMixin, Transformer):
@@ -887,8 +896,7 @@ def __init__(self, lobe: nn.Module, output_dim: int, coarse_grain: list = None,
         self.optimizer_u = torch.optim.Adam(self.ulayer.parameters(), lr=self.learning_rate*10)
         self.optimizer_s = torch.optim.Adam(self.slayer.parameters(), lr=self.learning_rate*100)
         self.optimizer_lobe = torch.optim.Adam(self.lobe.parameters(), lr=self.learning_rate)
-        self.optimimzer_all = torch.optim.Adam(chain(self.ulayer.parameters(), self.slayer.parameters(), self.lobe.parameters()), 
-                                              lr=self.learning_rate)
+        self.optimimzer_all = torch.optim.Adam(chain(self.ulayer.parameters(), self.slayer.parameters(), self.lobe.parameters()), lr=self.learning_rate)
         self._train_scores = []
         self._validation_scores = []
         self._train_vampe = []
@@ -1075,6 +1083,7 @@ def partial_fit(self, data, mask: bool = False, train_score_callback: Callable[[
 
         loss_value = -vampe_loss_rev(x_0, x_t, self.ulayer, self.slayer)[0]
         loss_value.backward()
+        torch.nn.utils.clip_grad_norm_(chain(self.lobe.parameters(), self.mask.parameters(), self.ulayer.parameters(), self.slayer.parameters()), CLIP_VALUE)
         if self.mask is not None and mask:
             self.optimizer_mask.step()
         self.optimizer_lobe.step()
@@ -1194,6 +1203,7 @@ def fit(self, data_loader: torch.utils.data.DataLoader, n_epochs=1, validation_l
 
                     loss_value = -vampe_loss_rev(x_0, x_t, self.ulayer, self.slayer)[0]
                     loss_value.backward()
+                    torch.nn.utils.clip_grad_norm_(chain(self.ulayer.parameters(), self.slayer.parameters()), CLIP_VALUE)
                     self.optimizer_u.step()
                     if train_mode=='us':
                         self.optimizer_s.step()
@@ -1224,6 +1234,7 @@ def fit(self, data_loader: torch.utils.data.DataLoader, n_epochs=1, validation_l
 
                     loss_value = -vampe_loss_rev_only_S(v, C_00, C_11, C_01, Sigma, self.slayer)[0]
                     loss_value.backward()
+                    torch.nn.utils.clip_grad_norm_(self.slayer.parameters(), CLIP_VALUE)
                     self.optimizer_s.step()
 
                     if train_score_callback is not None:
@@ -1355,6 +1366,7 @@ def fit_routine(self, data_loader: torch.utils.data.DataLoader, n_epochs=1, vali
                 score = vampe_loss_rev(x_0, x_t, self.ulayer, self.slayer)[0]
                 loss_value = -score
                 loss_value.backward()
+                torch.nn.utils.clip_grad_norm_(chain(self.ulayer.parameters(), self.slayer.parameters()), CLIP_VALUE)
                 self.optimizer_u.step()
                 self.optimizer_s.step()
                 if (score-score_value_before) < rel and counter > 0:
@@ -1455,6 +1467,9 @@ def fit_cg(self, data_loader: torch.utils.data.DataLoader, n_epochs=1, validatio
                     loss_value += torch.trace(matrix_cg)
 
                 loss_value.backward()
+                torch.nn.utils.clip_grad_norm_(chain(self.ulayer.parameters(), self.slayer.parameters()), CLIP_VALUE)
+                for lay_cg in self.cg_list:
+                    torch.nn.utils.clip_grad_norm_(lay_cg.parameters(), CLIP_VALUE)
                 self.optimizer_u.step()
                 self.optimizer_s.step()
                 for opt in self.cg_opt_list:
@@ -1503,6 +1518,7 @@ def fit_cg(self, data_loader: torch.utils.data.DataLoader, n_epochs=1, validatio
                 matrix_cg = self.cg_list[idx].get_cg_uS(chi_t, chi_tau, u_n, S_n, return_chi=False)[0]
                 loss_value = torch.trace(matrix_cg)
                 loss_value.backward()
+                torch.nn.utils.clip_grad_norm_(self.cg_list[idx].parameters(), CLIP_VALUE)
                 self.cg_opt_list[idx].step()
 
                 if train_score_callback is not None:
@@ -1634,6 +1650,7 @@ def partial_fit_obs(self, data, data_ev, data_ac, exp_ev=None, exp_ac=None, exp_
             loss_its, est_its = obs_its_loss(S, Sigma, its_state1, its_state2, exp_its, lam_its, epsilon=self.epsilon, mode=self.score_mode)
             loss_value += loss_its
         loss_value.backward()
+        torch.nn.utils.clip_grad_norm_(chain(self.lobe.parameters(), self.mask.parameters(), self.ulayer.parameters(), self.slayer.parameters()), CLIP_VALUE)
         self.optimizer_lobe.step()
         self.optimizer_u.step()
         self.optimizer_s.step()
@@ -1989,6 +2006,7 @@ def fit_obs(self, data_loader: torch.utils.data.DataLoader, n_epochs=1, validati
                             for i in range(est_its.shape[0]):
                                 tb_writer.add_scalars('ITS', {'train_'+str(i+1): est_its[i].item()}, self._step)
                     loss_value.backward()
+                    torch.nn.utils.clip_grad_norm_(chain(self.ulayer.parameters(), self.slayer.parameters()), CLIP_VALUE)
                     self.optimizer_u.step()
                     if train_mode=='us':
                         self.optimizer_s.step()
@@ -2082,6 +2100,7 @@ def fit_obs(self, data_loader: torch.utils.data.DataLoader, n_epochs=1, validati
                             for i in range(est_its.shape[0]):
                                 tb_writer.add_scalars('ITS', {'train_'+str(i+1): est_its[i].item()}, self._step)
                     loss_value.backward()
+                    torch.nn.utils.clip_grad_norm_(self.slayer.parameters(), CLIP_VALUE)
                     self.optimizer_s.step()
 
                     if train_score_callback is not None: