"""Autoencoder ensemble training utilities.
This module provides functions to train a zoo of independent autoencoder
models with early stopping, forming the ensemble used by the GRN inference
stage. Training is hardcoded to CUDA GPU.
"""
from sklearn.model_selection import train_test_split
from netmap.model.nbautoencoder import NegativeBinomialAutoencoder
from netmap.model.zinbautoencoder import ZINBAutoencoder
import torch
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
[docs]
def create_model_zoo(data_tensor, n_models = 10, n_epochs = 10000, model_type = 'ZINBAutoencoder', dropout_rate = 0.1, latent_dim=8, hidden_dim=[64]):
"""Create an ensemble of autoencoders (model zoo) trained on the supplied data.
Each model is initialised with a fresh random seed, trained on an independent
80/20 train/validation split, and added to the zoo only if training converges.
Models that fail to converge (i.e. ``_train_autoencoder_early_stopping()``
returns ``None``) are skipped; training retries until ``n_models`` succeed or
5 consecutive failures occur, after which the loop is aborted.
Args:
data_tensor (torch.Tensor): Raw gene expression data of shape
``(n_cells, n_genes)``. Must be a floating-point tensor on CPU;
it is moved to CUDA internally.
n_models (int, optional): Number of successfully trained models to
collect in the zoo. Defaults to 10.
n_epochs (int, optional): Maximum number of training epochs per model
before early stopping would have been triggered. Defaults to 10000.
model_type (str, optional): Architecture to instantiate. One of
``'ZINBAutoencoder'`` (zero-inflated negative binomial) or
``'NegativeBinomialAutoencoder'``. Any unrecognised value falls
back to ``NegativeBinomialAutoencoder``. Defaults to
``'ZINBAutoencoder'``.
dropout_rate (float, optional): Dropout probability applied during
training in each encoder/decoder layer. Defaults to 0.1.
latent_dim (int, optional): Dimensionality of the bottleneck latent
space. Defaults to 8.
hidden_dim (list of int, optional): Encoder layer sizes, specified as a
list of integers from the input layer towards the bottleneck (e.g.
``[128, 64]`` produces two encoder layers of width 128 and 64). The
decoder mirrors this sequence in reverse. Defaults to ``[64]``.
Returns:
list: A list of trained autoencoder model instances. Length is at most
``n_models``; may be shorter if the 5-consecutive-failure limit is
reached before enough models converge.
"""
model_zoo = []
counter = 0
failures = 0
while (counter < n_models) and (failures < 5):
data_train2, data_test2 = train_test_split(data_tensor, test_size=0.2, shuffle=True)
if model_type == 'ZINBAutoencoder':
trained_model2 = ZINBAutoencoder(input_dim=data_tensor.shape[1], latent_dim=latent_dim, dropout_rate = dropout_rate, hidden_dims = hidden_dim)
elif model_type == 'NegativeBinomialAutoencoder':
trained_model2 = NegativeBinomialAutoencoder(input_dim=data_tensor.shape[1], latent_dim=latent_dim, dropout_rate = dropout_rate, hidden_dims = hidden_dim)
else:
trained_model2 = NegativeBinomialAutoencoder(input_dim=data_tensor.shape[1], latent_dim=latent_dim, dropout_rate = dropout_rate, hidden_dims = hidden_dim)
trained_model2 = trained_model2.cuda()
optimizer2 = torch.optim.Adam(trained_model2.parameters(), lr=1e-4)
trained_model2 = _train_autoencoder_early_stopping(
trained_model2,
data_train2.cuda(),
data_test2.cuda(),
optimizer2,
num_epochs=n_epochs
)
if trained_model2 is not None:
model_zoo.append(trained_model2)
counter +=1
else:
failures+=1
return model_zoo
def _train_autoencoder(
model,
data_train,
optimizer,
batch_size=32, # Minibatch size
num_epochs=100,
):
"""Legacy training loop without early stopping.
Trains the model for a fixed number of epochs using minibatch gradient
descent. Validation and early stopping are not performed. Prefer
``_train_autoencoder_early_stopping()`` for production use.
Args:
model (torch.nn.Module): An initialised autoencoder instance (e.g.
``NegativeBinomialAutoencoder`` or ``ZINBAutoencoder``) that
exposes a ``compute_loss(batch)`` method.
data_train (torch.Tensor): Training data of shape
``(n_cells, n_genes)``.
optimizer (torch.optim.Optimizer): Optimiser to use for parameter
updates (e.g. ``torch.optim.Adam``).
batch_size (int, optional): Minibatch size for the DataLoader.
Defaults to 32.
num_epochs (int, optional): Total number of training epochs.
Defaults to 100.
Returns:
torch.nn.Module: The trained model after ``num_epochs`` epochs.
"""
# Prepare DataLoader for training
train_dataset = TensorDataset(data_train)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
for epoch in range(num_epochs):
model.train()
epoch_loss = 0 # Track loss for the epoch
for batch in train_loader:
data_batch = batch[0] # Unpack the single-element tuple from TensorDataset
# Forward pass
optimizer.zero_grad()
loss = model.compute_loss(data_batch)
# Compute total loss
loss.backward()
optimizer.step()
epoch_loss += loss.item() # Accumulate loss for the epoch
# Print progress
if epoch % 10 == 0:
print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss / len(train_loader):.4f}")
return model
def _train_autoencoder_early_stopping(
model,
data_train,
data_val,
optimizer,
batch_size=32,
num_epochs=10000,
patience=10,
min_delta=0.001,
validation_freq = 10,
):
"""Training loop with early stopping based on validation loss.
Trains the model using minibatch gradient descent and evaluates on the
validation set every ``validation_freq`` epochs. The best model state
(lowest validation loss) is tracked throughout training. Training halts
early when the validation loss fails to improve by more than ``min_delta``
for ``patience`` consecutive validation checks. Before returning, the
best recorded model state is loaded back into the model via
``model.load_state_dict(best_model_state)``, so the returned model
corresponds to the checkpoint with the lowest observed validation loss
rather than the final epoch.
Args:
model (torch.nn.Module): An initialised autoencoder instance (e.g.
``NegativeBinomialAutoencoder`` or ``ZINBAutoencoder``) that
exposes a ``compute_loss(batch)`` method.
data_train (torch.Tensor): Training data of shape
``(n_cells, n_genes)``.
data_val (torch.Tensor): Validation data of shape
``(n_cells, n_genes)`` used exclusively for early stopping
decisions; not used for gradient updates.
optimizer (torch.optim.Optimizer): Optimiser to use for parameter
updates (e.g. ``torch.optim.Adam``).
batch_size (int, optional): Minibatch size for both DataLoaders.
Defaults to 32.
num_epochs (int, optional): Maximum number of training epochs before
the loop terminates regardless of early stopping. Defaults to
10000.
patience (int, optional): Number of validation checks without
sufficient improvement before early stopping is triggered.
Defaults to 10.
min_delta (float, optional): Minimum absolute decrease in validation
loss required to count as an improvement and reset the patience
counter. Defaults to 0.001.
validation_freq (int, optional): Interval in epochs between validation
evaluations. Defaults to 10.
Returns:
torch.nn.Module or None: The trained model loaded with the best
observed ``state_dict``, or ``None`` if a ``TypeError`` is raised
during the early-stopping comparison (indicating a training failure).
"""
# Prepare DataLoaders
train_dataset = TensorDataset(data_train)
val_dataset = TensorDataset(data_val)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
best_val_loss = float('inf')
epochs_no_improve = 0
best_model_state = None
epoch_iterator = tqdm(range(num_epochs), desc="Training Autoencoder")
for epoch in epoch_iterator:
# --- Training loop ---
model.train()
epoch_train_loss = 0
for batch in train_loader:
data_batch = batch[0]
optimizer.zero_grad()
loss = model.compute_loss(data_batch)
loss.backward()
optimizer.step()
current_loss = loss.item()
epoch_train_loss += current_loss
avg_train_loss = epoch_train_loss / len(train_loader)
avg_val_loss = None # Reset for the current epoch
if (epoch + 1) % validation_freq == 0:
model.eval()
epoch_val_loss = 0
with torch.no_grad():
for batch in val_loader:
data_batch = batch[0]
loss = model.compute_loss(data_batch)
epoch_val_loss += loss.item()
avg_val_loss = epoch_val_loss / len(val_loader)
# 1. Update the overall epoch progress bar description
epoch_iterator.set_postfix(
train_loss=f"{avg_train_loss:.4f}",
val_loss=f"{avg_val_loss:.4f}",
best_val=f"{best_val_loss:.4f}"
)
else:
# If no validation was performed, update the epoch iterator with just train loss
epoch_iterator.set_postfix(train_loss=f"{avg_train_loss:.4f}", val_loss = '------', best_val=f"{best_val_loss:.4f}")
if avg_val_loss is not None:
try:
if avg_val_loss < best_val_loss - min_delta:
best_val_loss = avg_val_loss
epochs_no_improve = 0
best_model_state = model.state_dict()
else:
epochs_no_improve += 1
if epochs_no_improve >= patience:
# Use tqdm.write for clean printing of the stopping message
tqdm.write(f"Early stopping triggered after {epoch + 1} epochs due to no improvement in validation loss (Best Loss: {best_val_loss:.4f}).")
model.load_state_dict(best_model_state)
return model
except TypeError:
return None
# Load the best state if the loop finishes
if best_model_state:
model.load_state_dict(best_model_state)
return model