PyJuice

PyJuice is a library for Probabilistic Circuits (PCs) written in PyTorch. It has code for inference (e.g., marginals, sampling) and learning (e.g., EM, pruning) in PCs, which can be either defined by hand or generated directly from pre-specified structures (e.g., PD, RAT-SPN, HCLT).

Why PyJuice?

The biggest advantage of PyJuice is its speed and scalability. We benchmark PyJuice against prior PC packages SPFlow, EiNet, and Juice.jl on the PD and HCLT structures with various sizes by variating their width. We report the average ($\pm$ standard deviation of 5 runs) runtime (in seconds) per training epoch of 60K samples. All experiments were carried out on an RTX 4090 GPU with 24GB memory. To maximize parallelism, we always use the maximum possible batch size. "OOM" denotes out-of-memory with batch size 2.

As indicated by the tables, PyJuice is not only much faster than existing implementations, but it is also more scalable -- it can train much larger PCs in a reasonable amount of time without suffering from OOM issues.

Technical details of PyJuice can be found here.

Installation

The easiest way to install PyJuice is via pip:

pip install pyjuice

Alternatively, you can install PyJuice from source by running the following commands:

git clone [email protected]:Tractables/pyjuice.git
cd pyjuice
pip install -e .

Example Usage

The best way to get familiar with PyJuice is by referring to our documentation and the tutorials, which provide a gentle introduction to the main functionalities of PyJuice. In the following, we provide some quick start guides to get your first PyJuice program running.

Train a PC with pre-specified structures

pyjuice.structures contains a collection of widely used PC structures. In the following, we use HCLT as an example to demonstrate how to define an HCLT and learn its parameters with the mini-batch EM algorithm.

Assume that we have a training dataset train_data of size [num_samples, num_vars]. Assume all variables are categorical with 256 categories. We start by importing the necessary functions:

import torch
import pyjuice as juice
import pyjuice.nodes.distributions as juice_dists
from pyjuice.structures import HCLT

An HCLT with latent size 32 can be defined by:

root_ns = HCLT(train_data, num_latents = 32, input_dist = juice_dists.Categorical(num_cats = 256))

Here, the input train_data is used to generate the backbone Chow-Liu Tree structure of the HCLT. It is possible to supply a subset of the training set (e.g., train_data[1:100,:]). input_dist specifies the distributions of the input nodes (here every input is a categorical distribution with 256 categories).

The returned PC is stored as a Directed Acyclic Graph (DAG). Specifically, every node in the PC is an instance of CircuitNodes (defined here), which stores the properties of the current node (e.g., parameters) as well as its children (which are also instances of CircuitNodes). The returned object root_ns denotes the root node of the DAG.

As hinted by the name CircuitNodes, rather than representing a single PC node (i.e., input node, product node, or sum node), every CircuitNodes encodes a vector of nodes with the same type. For example, a SumNodes (a subclass of CircuitNodes) represents a vector of sum nodes, and a ProdNodes models a vector of product nodes. In principle, we can still define every CircuitNodes to be just 1 node, but often the vectorized representation will greatly simplify the code (and also speed it up quite a bit).

For now, we will move on to training the HCLT, and leave the detailed description of the API for defining PCs to the next example. Although the DAG-based representation is straightforward to understand and manipulate, it is not a good structure for efficient computation. PyJuice uses a compilation stage to convert the PC into an equivalent GPU-friendly representation:

pc = juice.compile(root_ns)

Note that one can equivalently use pc = juice.TensorCircuit(root_ns). From this point, we can treat pc as an instance of torch.nn.Module, and the training procedure is very similar to that of a neural network defined with PyTorch. We first move the PC to a GPU:

device = torch.device("cuda:0")
pc.to(device)

The training loop of a single epoch can be written as:

여기까지가 README 앞부분입니다. 나머지는 원문에서 읽을 수 있습니다. README 전문 보기 ↗