Kernel Density Estimation as Loss Function

username: d-diaz
cooked: <p>I am working on a 3D point cloud co-registration problem. I have an observed lidar point cloud and a generative model with several stochastic parameters that simulates point clouds. I am interested in using a Kernel Density Estimate generated from the observed lidar point cloud to compute the log probability of the simulated points and then optimize the generative model’s parameters to minimize the negative log likelihood from the KDE.</p>
<p>The sklearn KDE score_samples method is essentially what I’d like to reproduce in PyTorch as a loss function so that it can capture gradients and not just the log probability. I don’t see a straightforward way to implement KDE in PyTorch beyond building one from scratch (e.g., a list-like object containing as many 3D-Normal distributions as there are observed lidar points, which would be looped through to calculate the log_pdf for every simulated point against every one of these KDE distributions and then summing them across all distributions for all simulated points).  This seems like a lot of looping, and am not sure if there’s an effective way to broadcast this to be efficient.</p>
<p>Any ideas?</p>

username: mmaher
cooked: <p>Did you find any solutions to your question?</p>

username: Arta_A
cooked: <p>If you find any solution please tell us.</p>

username: d-diaz
cooked: <p>Just wanted to provide a quick update that I have not found a good solution yet. This project has been on a back-burner, but will probably be picking up again in later this fall. Would still welcome any suggestions.</p>

username: avloss
cooked: <p>I’m trying to do something similar (still not quite sure, still learning)<br>
Trying to figure out if this is about the same problem:<br>
</p><aside class="onebox whitelistedgeneric">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="16" height="16">
      <a href="https://github.com/ikostrikov/pytorch-flows/" target="_blank" rel="nofollow noopener">GitHub</a>
  </header>
  <article class="onebox-body">
    <img src="https://avatars1.githubusercontent.com/u/1212494?s=400&amp;v=4" class="thumbnail onebox-avatar" width="60" height="60">

<h3><a href="https://github.com/ikostrikov/pytorch-flows/" target="_blank" rel="nofollow noopener">ikostrikov/pytorch-flows</a></h3>

<p>PyTorch implementations of algorithms for density estimation - ikostrikov/pytorch-flows</p>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


username: d-diaz
cooked: <p>I came across this <a href="https://janvdvegt.github.io/2017/06/07/Kernel-Mixture-Networks.html" rel="nofollow noopener">blog post</a> mentioned in a comment to a question that has since been taken down from <a href="https://webcache.googleusercontent.com/search?q=cache:vAEfqljooCIJ:https://stackoverflow.com/questions/61685272/kernel-density-estimation-in-pytorch+&amp;cd=5&amp;hl=en&amp;ct=clnk&amp;gl=us" rel="nofollow noopener">Stack Overflow</a>.</p>
<p>The authors of the blog post seem like they’re attempting much more complicated modeling than I am, and are using TensorFlow, but it appears they are also looping through each point/kernel to calculate an overall density:</p>
<pre><code class="lang-auto">def density_output_layer(w, y, cluster_centers, sigma):
    output_kernels = []
    for center in cluster_centers:
        output_kernels.append(tf.reshape(gaussian_kernel(y, center, sigma), [-1, 1]))
    output_kernels = tf.concat(output_kernels, axis=1)
    output_nodes = tf.multiply(w, output_kernels)
    return tf.reduce_sum(output_nodes, axis=1) / total_weight
</code></pre>

username: d-diaz
cooked: <p>Got something working for vectorized calculation of log probability. I have noticed this can gobble up memory fast if you have large numbers of points. Would welcome any more thoughts. Here is a class to implement it:</p>
<pre><code class="lang-auto">import torch
from torch.distributions import MultivariateNormal, Normal
from torch.distributions.distribution import Distribution

class GaussianKDE(Distribution):
    def __init__(self, X, bw):
        """
        X : tensor (n, d)
          `n` points with `d` dimensions to which KDE will be fit
        bw : numeric
          bandwidth for Gaussian kernel
        """
        self.X = X
        self.bw = bw
        self.dims = X.shape[-1]
        self.n = X.shape[0]
        self.mvn = MultivariateNormal(loc=torch.zeros(self.dims),
                                      covariance_matrix=torch.eye(self.dims))

    def sample(self, num_samples):
        idxs = (np.random.uniform(0, 1, num_samples) * self.n).astype(int)
        norm = Normal(loc=self.X[idxs], scale=self.bw)
        return norm.sample()

    def score_samples(self, Y, X=None):
        """Returns the kernel density estimates of each point in `Y`.

        Parameters
        ----------
        Y : tensor (m, d)
          `m` points with `d` dimensions for which the probability density will
          be calculated
        X : tensor (n, d), optional
          `n` points with `d` dimensions to which KDE will be fit. Provided to
          allow batch calculations in `log_prob`. By default, `X` is None and
          all points used to initialize KernelDensityEstimator are included.


        Returns
        -------
        log_probs : tensor (m)
          log probability densities for each of the queried points in `Y`
        """
        if X == None:
            X = self.X
        log_probs = torch.log(
            (self.bw**(-self.dims) *
             torch.exp(self.mvn.log_prob(
                 (X.unsqueeze(1) - Y) / self.bw))).sum(dim=0) / self.n)

        return log_probs

    def log_prob(self, Y):
        """Returns the total log probability of one or more points, `Y`, using
        a Multivariate Normal kernel fit to `X` and scaled using `bw`.

        Parameters
        ----------
        Y : tensor (m, d)
          `m` points with `d` dimensions for which the probability density will
          be calculated

        Returns
        -------
        log_prob : numeric
          total log probability density for the queried points, `Y`
        """

        X_chunks = self.X.split(1000)
        Y_chunks = Y.split(1000)

        log_prob = 0

        for x in X_chunks:
            for y in Y_chunks:
                log_prob += self.score_samples(y, x).sum(dim=0)

        return log_prob
</code></pre>
<p>Here’s a toy example working with gradients:</p>
<pre><code class="lang-auto">from sklearn.datasets import make_blobs
from torch.optim import LBFGS

# centers of blobs we'll try to find through optimization 
# are at (0.1, 0.3) and (-0.2, -0.1)
X, y = make_blobs(5000, centers=[[0.1,0.3],[-0.2,-0.1]], cluster_std=0.1) 
kde = GaussianKDE(X=torch.tensor(X, dtype=torch.float32), bw=0.03)
test_pts = torch.tensor([[-0.75,-0.25],[0.6,0.4]], requires_grad=True)   
optimizer = LBFGS([test_pts])

for _ in range(10):  # take 10 optimization steps
    def closure():
        optimizer.zero_grad()
        loss = -kde.log_prob(test_pts)
        loss.backward()
        return loss
    optimizer.step(closure)

test_pts
&gt;&gt;&gt; tensor([[-0.2470, -0.1488],
            [ 0.1720,  0.3189]], requires_grad=True)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/e/6ea3f71ea856a79b23f5809acaa46ee82c1b514a.png" data-download-href="https://discuss.pytorch.org/uploads/default/6ea3f71ea856a79b23f5809acaa46ee82c1b514a" title="finding_blob_centers"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/e/6ea3f71ea856a79b23f5809acaa46ee82c1b514a_2_250x250.png" alt="finding_blob_centers" data-base62-sha1="fMLOzYbnabt0iXXlFiPokqGV42C" width="250" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/e/6ea3f71ea856a79b23f5809acaa46ee82c1b514a_2_250x250.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/e/6ea3f71ea856a79b23f5809acaa46ee82c1b514a_2_375x375.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/e/6ea3f71ea856a79b23f5809acaa46ee82c1b514a_2_500x500.png 2x" data-dominant-color="FDFDFD"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">finding_blob_centers</span><span class="informations">1200×1200 77.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: Bohan_Li
cooked: <p>Hi Diaz, I am interested in this work. How can I cite it in my research paper?</p>

username: Fort23
cooked: <p>Did you find the solution<a href="https://dgcustomerfirst.fun/" rel="noopener nofollow ugc">.</a>?</p>

username: kkladny
cooked: <p>I created a proper package for kernel density estimation in PyTorch called <a href="https://github.com/rudolfwilliam/torch-kde" rel="noopener nofollow ugc">torch-kde</a> that provides a <code>KernelDensity</code> class like in <a href="https://scikit-learn.org/dev/modules/generated/sklearn.neighbors.KernelDensity.html" rel="noopener nofollow ugc">scikit-learn</a>, but is fully differentiable. Enjoy!</p>