src.score_net

Classes

ScoreNet(*args, **kwargs)

A simple convolutional neural network for score estimation.

class src.score_net.ScoreNet(*args: Any, **kwargs: Any)

Bases: Module

A simple convolutional neural network for score estimation.

This module implements a small encoder-decoder architecture using convolutional and transposed convolutional layers with ELU activations to estimate the score function \(\nabla_x \log p(x)\).

Variables:

main (nn.Sequential) – The sequential model containing the convolutional and transposed convolutional layers.

Initialize the ScoreNet model.

Parameters:

channels (int, optional) – Number of input and output channels. For RGB images is typically 3. Default is 3.

forward(x: torch.Tensor) → torch.Tensor

Compute the forward pass of the ScoreNet.

Parameters:

x (torch.Tensor) – Input image batch of shape (batch_size, channels, height, width).

Returns:

Output tensor with the same shape as the input, representing the estimated score for each pixel.

Return type:

torch.Tensor