Merge branch 'feature/deeptfa_gradient_flow' into develop

Author: Eli Sennesh

htfa_torch/dtfa.py

@@ -38,6 +38,8 @@
 from . import tfa_models
 from . import utils
 
+EPOCH_MSG = '[Epoch %d] (%dms) Posterior free-energy %.8e = KL from prior %.8e - log-likelihood %.8e'
+
 class DeepTFA:
     """Overall container for a run of Deep TFA"""
     def __init__(self, query, mask, num_factors=tfa_models.NUM_FACTORS,
@@ -123,7 +125,7 @@ def train(self, num_steps=10, learning_rate=tfa.LEARNING_RATE,
         optimizer = torch.optim.Adam(list(variational.parameters()),
                                      lr=learning_rate, weight_decay=1e-2)
         scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-            optimizer, factor=1e-1, min_lr=5e-5, patience=patience,
+            optimizer, factor=0.5, min_lr=1e-5, patience=patience,
             verbose=True
         )
         variational.train()
@@ -136,9 +138,13 @@ def train(self, num_steps=10, learning_rate=tfa.LEARNING_RATE,
         for epoch in range(num_steps):
             start = time.time()
             epoch_free_energies = list(range(len(activations_loader)))
+            epoch_lls = list(range(len(activations_loader)))
+            epoch_prior_kls = list(range(len(activations_loader)))
 
             for (batch, data) in enumerate(activations_loader):
                 epoch_free_energies[batch] = 0.0
+                epoch_lls[batch] = 0.0
+                epoch_prior_kls[batch] = 0.0
                 block_batches = utils.chunks(list(range(self.num_blocks)),
                                              n=blocks_batch_size)
                 for block_batch in block_batches:
@@ -161,13 +167,13 @@ def train(self, num_steps=10, learning_rate=tfa.LEARNING_RATE,
                     generative(p, times=trs, guide=q, observations=activations,
                                blocks=block_batch)
 
-                    def block_rv_weight(node):
+                    def block_rv_weight(node, prior=True):
                         result = 1.0
                         if measure_occurrences:
                             rv_occurrences[node] += 1
                         result /= rv_occurrences[node]
                         return result
-                    free_energy = tfa.hierarchical_free_energy(
+                    free_energy, ll, prior_kl = tfa.hierarchical_free_energy(
                         q, p,
                         rv_weight=block_rv_weight,
                         num_particles=num_particles
@@ -176,6 +182,8 @@ def block_rv_weight(node):
                     free_energy.backward()
                     optimizer.step()
                     epoch_free_energies[batch] += free_energy
+                    epoch_lls[batch] += ll
+                    epoch_prior_kls[batch] += prior_kl
 
                     if tfa.CUDA and use_cuda:
                         del activations
@@ -185,8 +193,12 @@ def block_rv_weight(node):
                         torch.cuda.empty_cache()
                 if tfa.CUDA and use_cuda:
                     epoch_free_energies[batch] = epoch_free_energies[batch].cpu().data.numpy()
+                    epoch_lls[batch] = epoch_lls[batch].cpu().data.numpy()
+                    epoch_prior_kls[batch] = epoch_prior_kls[batch].cpu().data.numpy()
                 else:
                     epoch_free_energies[batch] = epoch_free_energies[batch].data.numpy()
+                    epoch_lls[batch] = epoch_lls[batch].data.numpy()
+                    epoch_prior_kls[batch] = epoch_prior_kls[batch].data.numpy()
 
             free_energies[epoch] = np.array(epoch_free_energies).sum(0)
             free_energies[epoch] = free_energies[epoch].sum(0)
@@ -195,7 +207,9 @@ def block_rv_weight(node):
             measure_occurrences = False
 
             end = time.time()
-            msg = tfa.EPOCH_MSG % (epoch + 1, (end - start) * 1000, free_energies[epoch])
+            msg = EPOCH_MSG % (epoch + 1, (end - start) * 1000,
+                               free_energies[epoch], sum(epoch_prior_kls),
+                               sum(epoch_lls))
             logging.info(msg)
             if checkpoint_steps is not None and epoch % checkpoint_steps == 0:
                 now = datetime.datetime.now()
@@ -214,22 +228,33 @@ def block_rv_weight(node):
 
         return np.vstack([free_energies])
 
-    def results(self, block):
+    def results(self, block, hist_weights=False):
         hyperparams = self.variational.hyperparams.state_vardict()
         subject = self.generative.block_subjects[block]
+        task = self.generative.block_tasks[block]
 
         factors_embed = hyperparams['factors']['mu'][subject]
 
-        weights = hyperparams['block']['weights']['mu'][block]\
-                             [self._blocks[block].start_time:
-                              self._blocks[block].end_time]
-        factor_params = self.variational.factors_embedding(factors_embed).view(
-            self.num_factors, 8
+        factor_params = self.variational.factors_embedding(factors_embed)
+        factor_centers = self.variational.centers_embedding(factor_params).view(
+            self.num_factors, 3
         )
-        factor_centers = factor_params[:, :3]
-        factor_log_widths = factor_params[:, 6].contiguous().view(
-            self.num_factors
+        factor_log_widths = self.variational.log_widths_embedding(factor_params)
+
+        weight_deltas = hyperparams['block']['weights']['mu'][block]\
+                                   [self._blocks[block].start_time:
+                                    self._blocks[block].end_time]
+        subject_embed = hyperparams['subject']['mu'][subject]
+        task_embed = hyperparams['task']['mu'][task]
+        weights_embed = torch.cat((subject_embed, task_embed), dim=-1)
+        weight_params = self.variational.weights_embedding(weights_embed).view(
+            self.num_factors, 2
         )
+        weights = weight_params[:, 0] + weight_deltas
+
+        if hist_weights:
+            plt.hist(weights.view(weights.numel()).data.numpy())
+            plt.show()
 
         return {
             'weights': weights.data,
@@ -356,6 +381,46 @@ def plot_reconstruction(self, block=None, filename=None, show=True,
 
         return plot
 
+    def visualize_factor_embedding(self, filename=None, show=True,
+                                   num_samples=100, hist_log_widths=True,
+                                   **kwargs):
+        hyperprior = self.generative.hyperparams.state_vardict()
+
+        factor_prior = utils.unsqueeze_and_expand_vardict({
+            'mu': hyperprior['factors']['mu'][0],
+            'sigma': hyperprior['factors']['sigma'][0]
+        }, 0, num_samples, True)
+
+        embedding = torch.normal(factor_prior['mu'], factor_prior['sigma'] * 2)
+        factor_params = self.variational.factors_embedding(embedding)
+        centers = self.variational.centers_embedding(factor_params).view(
+            -1, self.num_factors, 3
+        ).data
+        widths = torch.exp(self.variational.log_widths_embedding(factor_params))
+        widths = widths.view(-1, self.num_factors).data
+
+        plot = niplot.plot_connectome(
+            np.eye(num_samples * self.num_factors),
+            centers.view(num_samples * self.num_factors, 3).numpy(),
+            node_size=widths.view(num_samples * self.num_factors).numpy(),
+            title="$z^F$ std-dev %.8e, $x^F$ std-dev %.8e, $\\rho^F$ std-dev %.8e" %
+            (embedding.std(0).norm(), centers.std(0).norm(),
+             torch.log(widths).std(0).norm()),
+            **kwargs
+        )
+
+        if filename is not None:
+            plot.savefig(filename)
+        if show:
+            niplot.show()
+
+        if hist_log_widths:
+            log_widths = torch.log(widths)
+            plt.hist(log_widths.view(log_widths.numel()).numpy())
+            plt.show()
+
+        return plot, centers, torch.log(widths)
+
     def scatter_factor_embedding(self, labeler=None, filename=None, show=True,
                                  xlims=None, ylims=None, figsize=(3.75, 2.75),
                                  colormap='Set1'):

htfa_torch/dtfa_models.py

@@ -76,16 +76,19 @@ def __init__(self, num_blocks, num_times, num_factors, num_subjects,
 
         params = utils.vardict({
             'factors': {
-                'mu': torch.zeros(self.num_blocks, self.embedding_dim),
-                'sigma': torch.sqrt(torch.rand(self.num_blocks, self.embedding_dim)),
+                'mu': torch.zeros(self.num_subjects, self.embedding_dim),
+                'sigma': torch.ones(self.num_subjects, self.embedding_dim) *\
+                         tfa_models.SOURCE_LOG_WIDTH_STD_DEV,
             },
             'subject': {
                 'mu': torch.zeros(self.num_subjects, self.embedding_dim),
-                'sigma': torch.sqrt(torch.rand(self.num_blocks, self.embedding_dim)),
+                'sigma': torch.ones(self.num_subjects, self.embedding_dim) *\
+                         tfa_models.SOURCE_WEIGHT_STD_DEV,
             },
             'task': {
                 'mu': torch.zeros(self.num_tasks, self.embedding_dim),
-                'sigma': torch.sqrt(torch.rand(self.num_blocks, self.embedding_dim)),
+                'sigma': torch.ones(self.num_tasks, self.embedding_dim) *\
+                         tfa_models.SOURCE_WEIGHT_STD_DEV,
             },
             'template': {
                 'factor_centers': {
@@ -95,7 +98,8 @@ def __init__(self, num_blocks, num_times, num_factors, num_subjects,
                 'factor_log_widths': {
                     'mu': hyper_means['factor_log_widths'] *\
                           torch.ones(self._num_factors),
-                    'sigma': torch.sqrt(torch.rand(self._num_factors))
+                    'sigma': torch.ones(self._num_factors) *
+                             tfa_models.SOURCE_LOG_WIDTH_STD_DEV,
                 }
             },
             'block': {
@@ -136,32 +140,28 @@ def __init__(self, num_factors, block_subjects, block_tasks, num_blocks=1,
                                                    embedding_dim)
         self.factors_embedding = nn.Sequential(
             nn.Linear(self._embedding_dim, self._num_factors),
-            nn.Tanhshrink(),
-            nn.Linear(self._num_factors, self._num_factors * 8),
+            nn.Softsign(),
         )
+        self.centers_embedding = nn.Linear(self._num_factors,
+                                           self._num_factors * 3)
+        self.log_widths_embedding = nn.Linear(self._num_factors,
+                                              self._num_factors)
         self.weights_embedding = nn.Sequential(
             nn.Linear(self._embedding_dim * 2, self._num_factors),
-            nn.Tanhshrink(),
+            nn.Softsign(),
             nn.Linear(self._num_factors, self._num_factors * 2),
         )
         self.softplus = nn.Softplus()
 
         self.epsilon = nn.Parameter(torch.Tensor([tfa_models.VOXEL_NOISE]))
 
         if hyper_means is not None:
-            self.weights_embedding[-1].bias = nn.Parameter(torch.cat(
-                (hyper_means['weights'].mean(0),
-                 torch.sqrt(torch.rand(self._num_factors))),
-                dim=0
-            ))
-            self.factors_embedding[-1].bias = nn.Parameter(torch.cat(
-                (hyper_means['factor_centers'],
-                 torch.ones(self._num_factors, 3),
-                 torch.ones(self._num_factors, 1) *
-                 hyper_means['factor_log_widths'],
-                 torch.sqrt(torch.rand(self._num_factors, 1))),
-                dim=1,
-            ).view(self._num_factors * 8))
+            self.centers_embedding.bias = nn.Parameter(
+                hyper_means['factor_centers'].view(self._num_factors * 3)
+            )
+            self.log_widths_embedding.bias = nn.Parameter(
+                torch.ones(self._num_factors) * hyper_means['factor_log_widths']
+            )
 
     def forward(self, trace, times=None, blocks=None,
                 num_particles=tfa_models.NUM_PARTICLES):
@@ -187,48 +187,53 @@ def forward(self, trace, times=None, blocks=None,
             if ('z^F_%d' % subject) not in trace:
                 factors_embed = trace.normal(
                     params['factors']['mu'][:, subject, :],
-                    params['factors']['sigma'][:, subject, :],
+                    self.softplus(params['factors']['sigma'][:, subject, :]),
                     name='z^F_%d' % subject
                 )
             if ('z^P_%d' % subject) not in trace:
                 subject_embed = trace.normal(
                     params['subject']['mu'][:, subject, :],
-                    params['subject']['sigma'][:, subject, :],
+                    self.softplus(params['subject']['sigma'][:, subject, :]),
                     name='z^P_%d' % subject
                 )
             if ('z^S_%d' % task) not in trace:
-                task_embed = trace.normal(params['task']['mu'][:, task],
-                                          params['task']['sigma'][:, task],
-                                          name='z^S_%d' % task)
+                task_embed = trace.normal(
+                    params['task']['mu'][:, task],
+                    self.softplus(params['task']['sigma'][:, task]),
+                    name='z^S_%d' % task
+                )
 
             factor_params = self.factors_embedding(factors_embed)
-            factor_params = factor_params.view(-1, self._num_factors, 8)
+            centers_predictions = self.centers_embedding(factor_params).view(
+                -1, self._num_factors, 3
+            )
+            log_widths_predictions = self.log_widths_embedding(factor_params).\
+                                     view(-1, self._num_factors)
             weights_embed = torch.cat((subject_embed, task_embed), dim=-1)
-            weight_params = self.weights_embedding(weights_embed).view(
+            weight_predictions = self.weights_embedding(weights_embed).view(
                 -1, self._num_factors, 2
             )
 
-            trace.normal(weight_params[:, :, 0], self.epsilon[0],
-                         name='mu^W_%d' % b)
-            trace.normal(self.softplus(weight_params[:, :, 1]), self.epsilon[0],
-                         name='sigma^W_%d' % b)
+            weights_mu = trace.normal(weight_predictions[:, :, 0],
+                                      self.epsilon[0], name='mu^W_%d' % b)
+            weights_sigma = trace.normal(weight_predictions[:, :, 1],
+                                         self.epsilon[0], name='sigma^W_%d' % b)
+            weights_params = params['block']['weights']
             weights[i] = trace.normal(
-                params['block']['weights']['mu'][:, b, ts[0]:ts[1], :],
-                params['block']['weights']['sigma'][:, b, ts[0]:ts[1], :],
+                weights_params['mu'][:, b, ts[0]:ts[1], :] +
+                weights_mu.unsqueeze(1),
+                self.softplus(weights_params['sigma'][:, b, ts[0]:ts[1], :] +
+                              weights_sigma.unsqueeze(1)),
                 name='Weights%dt%d-%d' % (b, ts[0], ts[1])
             )
             factor_centers[i] = trace.normal(
-                factor_params[:, :, 0:3],
-                self.softplus(factor_params[:, :, 3:6]),
+                centers_predictions,
+                self.epsilon[0],
                 name='FactorCenters%d' % b
             )
             factor_log_widths[i] = trace.normal(
-                factor_params[:, :, 6].contiguous().view(
-                    -1, self._num_factors
-                ),
-                self.softplus(factor_params[:, :, 7].contiguous().view(
-                    -1, self._num_factors
-                )), name='FactorLogWidths%d' % b
+                log_widths_predictions,
+                self.epsilon[0], name='FactorLogWidths%d' % b
             )
 
         return weights, factor_centers, factor_log_widths
@@ -297,5 +302,4 @@ def forward(self, trace, times=None, guide=probtorch.Trace(),
             }
 
         return self.htfa_model(trace, times, guide, blocks=blocks,
-                               observations=observations,
-                               weights_params=weight_params)
+                               observations=observations)

htfa_torch/htfa.py

@@ -130,13 +130,13 @@ def train(self, num_steps=10, learning_rate=tfa.LEARNING_RATE,
                     dec(p, times=trs, guide=q, observations=activations,
                         blocks=block_batch)
 
-                    def block_rv_weight(node):
+                    def block_rv_weight(node, prior=True):
                         result = 1.0
                         if measure_occurrences:
                             rv_occurrences[node] += 1
                         result /= rv_occurrences[node]
                         return result
-                    free_energy = tfa.hierarchical_free_energy(
+                    free_energy, _, _ = tfa.hierarchical_free_energy(
                         q, p,
                         rv_weight=block_rv_weight,
                         num_particles=num_particles

htfa_torch/htfa_models.py

@@ -179,7 +179,7 @@ def __init__(self, brain_center, brain_center_std_dev, num_blocks,
             'factor_center_noise': torch.ones(self._num_blocks),
             'factor_log_width_noise': torch.ones(self._num_blocks),
             'weights': {
-                'mu': torch.randn(self._num_blocks, self._num_factors),
+                'mu': torch.zeros(self._num_blocks, self._num_factors),
                 'sigma': tfa_models.SOURCE_WEIGHT_STD_DEV *\
                          torch.ones(self._num_blocks, self._num_factors)
             },