susan.ml.Noise2Noise

class susan.ml.Noise2Noise(*args: Any, **kwargs: Any)

Bases: Module

Simple Noise2Noise denoiser for 3-D cryo-EM half-map pairs.

Warning

Experimental. This class is part of the machine-learning subpackage and may change or be removed in a future release.

Takes one half-map as input and predicts the other as target. The loss is either L2 (MSE) or L1 (MAE). Backbone is either a sequential ConvNet or a UNet.

Parameters:

• arch (str) – Backbone architecture: 'convnet' or 'unet'. Default: 'convnet'.

• depth (int) – For 'convnet': total number of blocks (≥ 2). For 'unet': number of encoder/decoder levels (≥ 1). Default: 6.

• n_feat (int) – Feature channels (base channels for UNet). Default: 32.

• gpus (list of int) – CUDA device indices. Default: [0].

Forward

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

Forward pass.

Parameters:

x (torch.Tensor) – Shape (N, 1, D, H, W).

Returns:

Predicted denoised volume, same shape, on output_device().

Return type:

torch.Tensor

Persistence

save(path: str)

Save architecture hyperparameters and weights.

classmethod load(path: str, gpus: list = [0]) → Noise2Noise

Reconstruct model from a file saved with save().

Utilities

reset_weights()

Re-initialise all learnable parameters.

n_parameters() → int

Total number of trainable parameters.

input_device() → torch.device

Device expected for network inputs.

output_device() → torch.device

Device on which the prediction is produced.