CNN Model

A Convolutional Neural Network (CNN) model implementation for protein sequence prediction. This model uses convolutional layers to learn representations from protein sequences and predict fitness values.

class alf_tools.models.cnn.CNNModel(name='cnn_model', model_config=None, train_config=None, alphabet='ARNDCQEGHILKMFPSTWYV', device=None)

Bases: BaseModel

Minimal CNN model for sequence fitness prediction.

One-hot encodes sequences, trains a simple 1D CNN with MSE loss.

featurise(inputs)

Convert inputs to one-hot encoded tensors.

Parameters:

inputs (Union[LabelledCandidates, list[Candidate]]) – Either LabelledCandidates or list of Candidates to featurise.

Return type:

Float[Tensor, 'batch_size alphabet_size seq_length']

Returns:

A one-hot encoded tensor of shape (batch_size, alphabet_size, seq_length).

Raises:

ValueError – If the input is not LabelledCandidates or list of Candidates.

get_epoch_metrics()

Return per-epoch metrics from the most recent train() call.

Return type:

list[SurrogateEpochMetrics]

Returns:

List of SurrogateEpochMetrics, one per epoch trained.

get_training_summary_metrics()

Return training metrics.

Return type:

dict[str, Union[float, int, number]]

Returns:

Dictionary of training metrics including losses and Spearman correlations.

predict(candidate_points)

Make predictions for candidates.

For regression, returns raw scalar predictions. For binary classification, returns class probabilities of shape (n_samples, 2) via sigmoid on the single output neuron. For multiclass, returns softmax probabilities of shape (n_samples, num_classes).

Parameters:

candidate_points (list[Candidate]) – List of candidates to predict for.

Return type:

Predictions

Returns:

Predictions containing means of shape (n_samples,) for regression or (n_samples, num_classes) for classification.

Raises:

RuntimeError – If the model has not been trained yet.

sample(*args, **kwargs)

Sample candidate points from the model.

Return type:

list[Candidate]

setup(dataset)

Configure the model from the full dataset before training begins.

Determines the number of output neurons from the complete label space and validates consistency with the declared problem type.

Parameters:

dataset (BaseDataset) – The dataset this model will be trained on.

Raises:

ValueError – If the model’s problem_type disagrees with the dataset’s.

Return type:

None

train(train_data, val_data=None)

Train the CNN model.

Parameters:

• train_data (LabelledCandidates) – Training data containing sequences and oracle values.

• val_data (LabelledCandidates | None) – Optional validation data.

Raises:

• ValueError – If the model’s problem_type disagrees with the dataset’s.

• RuntimeError – If the model is not properly initialized.

Return type:

None

class alf_tools.models.cnn.CNNModelConfig(num_filters=128, kernel_size=3, num_conv_layers=3, fc_hidden_dim=256, dropout=0.3)

Bases: object

Configuration for CNN model architecture.

Parameters:

• num_filters (int) – Number of filters in convolutional layers.

• kernel_size (int) – Size of convolutional kernels.

• num_conv_layers (int) – Number of convolutional layers.

• fc_hidden_dim (int) – Dimension of fully connected hidden layers.

• dropout (float) – Dropout rate.

dropout: float = 0.3

fc_hidden_dim: int = 256

kernel_size: int = 3

num_conv_layers: int = 3

num_filters: int = 128

class alf_tools.models.cnn.CNNTrainConfig(learning_rate=0.001, log_frequency=10, normalise_inputs_strategy=None, standardise_outputs=False, label_dtype=None, batch_size=32, num_epochs=50)

Bases: BaseTrainConfig

Configuration for CNN training.

Parameters:

• batch_size (int) – Batch size for training.

• num_epochs (int) – Number of epochs to train for.

• learning_rate (float) – Inherited from BaseTrainConfig. Default: 1e-3.

• log_frequency (int) – Inherited from BaseTrainConfig. Default: 10.

• normalise_inputs_strategy (Literal[‘minmax’, ‘zscore’] | None) – Inherited from BaseTrainConfig. Default: None (no input normalisation). Z-score standardisation is a poor fit for the one-hot sequence inputs CNNModel uses, so it is left disabled by default; set explicitly to opt in for continuous-feature inputs.

• standardise_outputs (bool) – Inherited from BaseTrainConfig. Default: False.

• label_dtype (TorchDtype | None) – Inherited from BaseTrainConfig. None uses the model default (float32 for CNN regression). Override to force a dtype.

batch_size: int = 32

num_epochs: int = 50

class alf_tools.models.cnn.SequenceCNN(seq_length, alphabet_size=20, num_filters=128, kernel_size=3, num_conv_layers=3, fc_hidden_dim=256, dropout=0.3, output_neurons=1)

Bases: Module

Simple 1D CNN for sequence prediction.

Architecture: - One-hot encoding → Conv1D layers → Fully connected → output

forward(x)

Forward pass through the network.

Parameters:

x (Float[Tensor, 'batch_size alphabet_size seq_length']) – One-hot encoded sequences (batch_size, alphabet_size, seq_length).

Return type:

Tensor

Returns:

Logits of shape (batch_size,) when output_neurons==1, else (batch_size, output_neurons).