netmap.model.train_model.create_model_zoo

netmap.model.train_model.create_model_zoo(data_tensor, n_models=10, n_epochs=10000, model_type='ZINBAutoencoder', dropout_rate=0.1, latent_dim=8, hidden_dim=[64])[source]

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.

Parameters:

  • 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:

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.

Return type:

list