• Stars
    star
    498
  • Rank 88,494 (Top 2 %)
  • Language
    Python
  • License
    MIT License
  • Created over 6 years ago
  • Updated over 1 year ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

Gaussian mixture models in PyTorch.

This repository contains an implementation of a simple Gaussian mixture model (GMM) fitted with Expectation-Maximization in pytorch. The interface closely follows that of sklearn.

Example of a fit via a Gaussian Mixture model.


A new model is instantiated by calling gmm.GaussianMixture(..) and providing as arguments the number of components, as well as the tensor dimension. Note that once instantiated, the model expects tensors in a flattened shape (n, d).

The first step would usually be to fit the model via model.fit(data), then predict with model.predict(data). To reproduce the above figure, just run the provided example.py.

Some sanity checks can be executed by calling python test.py. To fit data on GPUs, ensure that you first call model.cuda().