multigeom.encoders.NeuralEncoder#

class multigeom.encoders.NeuralEncoder(*args, **kwargs)[source]#

A backbone followed by a readout.

Args: backbone: an nn.Module mapping images [B, 1, 64, 64] to features [B, C, S, S], with attributes out_channels and out_spatial. Pretrained backbones can be frozen already; a backbone that you trained is frozen by fit_readout. n_neurons: number of neurons to predict. readout: the name of a registered readout (“gaussian”, “linear”, “factorized”, or one you registered) or a readout module (features [B, C, S, S] -> [B, n_neurons]). dropout: dropout on the feature map while the readout is fitted (used when readout is a name).

Parameters:

  • backbone (nn.Module)

  • n_neurons (int)

  • readout (str | nn.Module)

  • dropout (float)

__init__(backbone, n_neurons, readout='gaussian', dropout=0.25)[source]#
Parameters:

  • backbone (nn.Module)

  • n_neurons (int)

  • readout (str | nn.Module)

  • dropout (float)

Methods

__init__(backbone, n_neurons[, readout, dropout])

features(x)

The backbone's feature map only (no readout, no gradient): what fit_readout caches.

fit_readout(images, responses, *[, epochs, ...])

Fit the readout to neural data and keep the best epoch.

forward(x)

x: [B, 1, 64, 64] in [-1, 1] -> responses [B, n_neurons].

freeze_backbone()

Stop training the backbone: its parameters get no gradient and its BatchNorm statistics stay fixed.

from_config(config)

Build an encoder (with fresh weights) from config, as returned by encoder.config.

load(path[, device])

Rebuild an encoder from a checkpoint written by save (or one of the paper's checkpoints) and load its weights.

load_checkpoint(path[, strict])

Load the weights of a checkpoint into this encoder (its architecture has to match).

predict

save(path, **metadata)

Save the weights, the architecture and metadata (anything picklable, e.g. the area or the test correlation) to one file.

summary()

One line per part: the backbone, the readout and the number of trainable parameters.

train([mode])

Set the mode; a frozen backbone stays in evaluation mode.

Attributes

config

enough to rebuild the encoder (without its weights).