Skip to main content

A pytorch package for Non-negative Matrix Factorization

Project description

Non-negative Matrix Fatorization in PyTorch

build codecov Documentation Status PyPI version

PyTorch is not only a good deep learning framework, but also a fast tool when it comes to matrix operations and convolutions on large data. A great example is PyTorchWavelets.

In this package I implement NMF, PLCA and their deconvolutional variations in PyTorch based on torch.nn.Module, so the models can be moved freely among CPU/GPU devices and utilize parallel computation of cuda. We also utilize the computational graph from torch.autograd to derive updated coefficients so the amount of codes is reduced and easy to maintain.



Basic NMF and NMFD module minimizing beta-divergence using multiplicative update rules.

The interface is similar to sklearn.decomposition.NMF with some extra options.

  • NMF: Original NMF algorithm.
  • NMFD: 1-D deconvolutional NMF algorithm.
  • NMF2D: 2-D deconvolutional NMF algorithm.
  • NMF3D: 3-D deconvolutional NMF algorithm.


Basic PLCA and SIPLCA module using EM algorithm to minimize KL-divergence between the target distribution and the estimated distribution.

  • PLCA: Original PLCA (Probabilistic Latent Component Analysis) algorithm.
  • SIPLCA: Shift-Invariant PLCA algorithm (similar to NMFD).
  • SIPLCA2: 2-D deconvolutional SIPLCA algorithm.
  • SIPLCA3: 3-D deconvolutional SIPLCA algorithm.


Here is a short example of decompose a spectrogram using deconvolutional NMF:

import torch
import librosa
from torchnmf.nmf import NMFD
from torchnmf.metrics import kl_div

y, sr = librosa.load(librosa.util.example_audio_file())
y = torch.from_numpy(y)
windowsize = 2048
S = torch.stft(y, windowsize, 
S = S.unsqueeze(0)

R = 8   # number of components
T = 400 # size of convolution window

net = NMFD(S.shape, rank=R, T=T).cuda()
# run extremely fast on gpu      # fit to target matrix S
V = net()
print(kl_div(V, S))        # KL divergence to S

A more detailed version can be found here. See our documentation to find out more usage of this package.

Compare to sklearn

The barchart shows the time cost per iteration with different beta-divergence. It shows that pytorch-based NMF has a much more constant process time across different beta values, which can take advantage when beta is not 0, 1, or 2. This is because our implementation use the same computational graph regardless which beta-divergence are we minimizing. It runs even faster when computation is done on GPU. The test is conducted on a Acer E5 laptop with i5-7200U CPU and GTX 950M GPU.


pip install torchnmf


  • PyTorch
  • tqdm


  • If you notice significant slow down when operating on CPU, please flush denormal numbers by torch.set_flush_denormal(True).


  • Support sparse matrix target (only on NMF module).
  • Regularization.
  • NNDSVD initialization.
  • 2/3-D deconvolutional module.
  • PLCA.
  • Documentation.
  • ipynb examples.
  • Refactor PLCA module.

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

torchnmf-0.3.4.tar.gz (21.4 kB view hashes)

Uploaded source

Built Distribution

torchnmf-0.3.4-py3-none-any.whl (24.5 kB view hashes)

Uploaded py3

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Huawei Huawei PSF Sponsor Microsoft Microsoft PSF Sponsor NVIDIA NVIDIA PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page