score-models 0.5.11

ScoreModel

The ScoreModel class is the main interface for training and using score models, defined as

where \(\sigma(t)\) is the standard deviation of the perturbation kernel \(p_t(\mathbf{x} \mid \mathbf{x}_0)\) of an SDE.

The ScoreModel class extends the torch.nn.Module class. Example usage:

from score_models import ScoreModel, EnergyModel, NCSNpp, MLP, DDPM

# B is the batch size
# C is the input channels
# dimensions are the spatial dimensions (e.g. [28, 28] for MNIST)

# Create a ScoreModel instance with Yang Song's NCSN++ architecture and the VESDE
net = NCSNpp(channels=C, dimensions=len(dimensions), nf=128, ch_mult=[2, 2, 2, 2])
model = ScoreModel(model=net, sigma_min=1e-2, sigma_max=50, device="cuda")
# ... or the VPSDE
model = ScoreModel(model=net, beta_min=1e-2, beta_max=20, device="cuda")

# NN Architectures support a Unet with 1D convolutions for time series input data
net = NCSNpp(channels=1, nf=128, ch_mult=[2, 2, 2, 2], dimensions=1)
# ... or 3D convolutions for videos/voxels
net = NCSNpp(channels=1, nf=128, ch_mult=[2, 2, 2, 2], attention=False, dimensions=3)
# You can also use a simpler MLP architecture (dimensions=0)
net = MLP(dimensions=C, layers=4, units=100)
# ... or Jonathan Ho's DDPM architecture
net = DDPM(channels=1, dimensions=2, nf=128, ch_mult=[2, 2, 2, 2])

# Train the score model, and save its weight in checkpoints_directory
model.fit(dataset=your_dataset, epochs=100, batch_size=B, learning_rate=1e-4, checkpoints_directory=checkpoints_directory)

# Generate samples from the trained model (steps is the number of Euler-Maruyama steps)
samples = model.sample(shape=[B, *dimensions], steps=1000)

# Compute the score for a given input
score = model.score(t, x)

# Initialize the score model and its neural network from a path to a checkpoint directory
score = ScoreModel(checkpoints_directory=checkpoints_directory)

EnergyModel

The EnergyModel class works in pretty much the same way as ScoreModel, but implements the score via the automatic differentiation of an energy model

This is to say that the score is defined as

Note: When using the MLP architecture, the energy model can be constructed more efficiently as the output of the neural network by specifying nn_is_energy in the hyperparameters of the MLP, which will modify the neural network architecture to be a function \(f_\theta: \mathbb{R}^d \to \mathbb{R}\) instead of \(f_\theta: \mathbb{R}^d \to \mathbb{R}^d\). An output_activation like relu can also be specified to make the energy positive or bounded from below.

Training Parameters

When training the model using the fit method, you can provide various parameters to customize the training process. Some important parameters include:

• dataset: The training dataset (torch.utils.data.Dataset).

• epochs: The number of training epochs.

• learning_rate: The learning rate for the ADAM optimizer.

• batch_size: The batch size for training.

• checkpoints_directory: The directory to save model checkpoints (default: None).

• seed: The random seed for numpy and torch.

Refer to the method’s docstring or the class definition for more details on available parameters.