DOCSEARCHbyAlgolia
Submit Docs

Squidiff 1.0.8

repository·main·Indexed 41 days ago

https://github.com/siyuh/squidiff

A diffusion model-based generative framework designed to predict transcriptomic changes across diverse cell types in response to environmental changes, drug treatments, cell differentiation, or gene perturbations. It features a specialized MLPModel architecture for single-cell RNA-seq data, supporting semantic latent embeddings and drug structural information.

Tokens
11.9K
Snippets
37
Records
42
Agent score
52%

What's inside Squidiff

  4. How ScheduleSampler works and its subclasses

    main

    A ScheduleSampler defines a distribution over timesteps in the diffusion process to reduce objective variance.

    Core Logic

    By default, samplers use unbiased importance sampling. The sample() method calculates probabilities based on weights(), then returns weights calculated as 1 / (len(p) * p[indices]) to ensure the expected value of the loss remains unchanged.

    Subclasses

    • UniformSampler: Provides a constant weight for all timesteps.
    • LossAwareSampler: An abstract base for samplers that can update their weights based on model loss. It provides the synchronization logic via update_with_local_losses.
    • LossSecondMomentResampler: A concrete implementation of LossAwareSampler. It tracks a history of losses for each timestep and uses the square root of the mean squared loss (second moment) to prioritize timesteps with higher variance/error. It includes a uniform_prob to ensure all timesteps retain a minimum probability of being sampled.

  5. How TimestepBlock and TimestepEmbedSequential work

    main

    The project uses a specialized hierarchy for modules that need to be aware of diffusion timesteps.

    1. TimestepBlock: An abstract base class (interface) for any module whose forward method accepts (x, emb, z_sem).

      • x: The input tensor.
      • emb: Timestep embeddings.
      • z_sem: Semantic latent embeddings.

    2. MLPBlock: A concrete implementation of TimestepBlock. It uses two linear layers with LayerNorm and SiLU activation. It conditionally adds emb and z_sem to the hidden state via learned linear projections (time_dense and zsem_dense).

    3. TimestepEmbedSequential: A specialized version of nn.Sequential. When iterating through its layers, it checks if a layer is an instance of TimestepBlock. If so, it passes x, emb, and z_sem. Otherwise, it passes only x. This allows mixing standard PyTorch layers with timestep-aware layers in a single sequence.

    # Conceptual usage of TimestepEmbedSequential
# It automatically handles the distribution of embeddings to compatible layers
from Squidiff.MLPModel import TimestepEmbedSequential, MLPBlock

layers = TimestepEmbedSequential(
    MLPBlock(128, 128, time_embed_dim=64, latent_dim=32),
    MLPBlock(128, 128, time_embed_dim=64, latent_dim=32)
)

# The call passes emb and z_sem to both MLPBlocks
output = layers(x=input_tensor, emb=time_emb, z_sem=semantic_emb)

  6. Train the Squidiff model via CLI

    main

    The train_squidiff.py script is the primary entrypoint for training the perturbation-conditioned generative diffusion model. It uses a command-line interface to configure model architecture, diffusion parameters, and training hyperparameters.

    To run training, you must provide at least --logger_path and --data_path via the CLI. The script supports distributed training and can resume from a checkpoint using --resume_checkpoint.

    # Example command to start training
python train_squidiff.py \
    --data_path "/path/to/adata/" \
    --logger_path "/path/to/logs/" \
    --batch_size 128 \
    --lr 1e-4 \
    --use_fp16 True

  7. Sample transcriptomics using sample_squidiff.sampler

    main

    Use the sample_squidiff.sampler to generate predicted transcriptomics. The process involves initializing a sampler with a model path, gene size, output dimension, and drug structure preference, then encoding test data through the model's encoder before calling .pred() to get the predictions.

    sampler = sample_squidiff.sampler(
    model_path = 'simu_results/model.pt',
    gene_size = 100,
    output_dim = 100,
    use_drug_structure = False
)

test_adata_scrna = sc.read_h5ad('datasets/sc_simu_test.h5ad')
z_sem_scrna = sampler.model.encoder(torch.tensor(test_adata_scrna.X).to('cuda'))

scrnas_pred = sampler.pred(z_sem_scrna, gene_size = test_adata_scrna.shape[1])

  8. Dump logged metrics to files

    main

    To write the currently buffered key-value metrics to the output files (CSV, JSON, etc.), call dumpkvs(). This operation clears the internal buffer and returns the dictionary of values that were just written.

    from Squidiff import dumpkvs

# Inside a training loop
metrics = dumpkvs()  # Writes to files and clears buffer

  9. Log key-value pairs and running means

    main

    The logger allows tracking metrics using key-value pairs. These metrics are aggregated and can be dumped to the configured output formats (like CSV or JSON).

    • logkv(key, val): Records a single value for a specific key. Subsequent calls to the same key will overwrite the previous value in the current buffer.
    • logkv_mean(key, val): Calculates a running average for the specified key. This is useful for tracking epoch-level losses or accuracies.
    • logkvs(d): A convenience function to log multiple key-value pairs from a dictionary d.
    from Squidiff import logkv, logkv_mean, logkvs

# Log a single metric
logkv("learning_rate", 0.001)

# Log a running average (e.g., loss)
logkv_mean("loss", 0.5)

# Log multiple metrics at once
logkvs({"accuracy": 0.95, "step": 100})

  10. Configure the global logger

    main

    Use configure(dir, format_strs=None, log_suffix="") to initialize the global Logger instance. This sets up where logs are stored and which output formats are used.

    • dir: The directory where log files will be created.
    • format_strs: A list or comma-separated string of formats. Supported formats are stdout, log, json, and csv.
    • log_suffix: An optional string to append to filenames.

    If configure is not called, calling get_current() will trigger a default configuration using the OPENAI_LOG_FORMAT environment variable.

    from Squidiff import configure

configure(dir="./logs", format_strs=["stdout", "csv", "json"], log_suffix="_experiment1")

  11. Use standard logging levels

    main

    The module provides standard logging functions to print messages to the configured output formats (e.g., stdout or a .txt log file).

    Available levels:

    • debug(args)
    • info(args)
    • warn(args)
    • error(args)

    Constants for levels: DEBUG=10, INFO=20, WARN=30, ERROR=40, DISABLED=50.

    from Squidiff import info, warn, error

info("Starting training...")
warn("Loss is diverging!")
error("Failed to load checkpoint.")

  12. Default configuration dictionaries

    main

    The following functions return dictionaries of default parameters for different components of the system:

    • diffusion_defaults(): Returns defaults for diffusion training (e.g., learn_sigma, diffusion_steps, noise_schedule, timestep_respacing, use_kl, predict_xstart, rescale_timesteps, rescale_learned_sigmas).
    • classifier_defaults(): Returns defaults for classifier models (e.g., gene_size, classifier_use_fp16, classifier_width, classifier_depth, classifier_attention_resolutions, classifier_use_scale_shift_norm, classifier_pool).
    • model_and_diffusion_defaults(): Returns a combined dictionary for training a standard model, including output_dim, num_layers, gene_size, num_channels, dropout, class_cond, use_checkpoint, use_scale_shift_norm, use_fp16, use_encoder, use_drug_structure, drug_dimension, and comb_num, merged with diffusion_defaults().
    • classifier_and_diffusion_defaults(): Returns a combined dictionary for classifier training, merging classifier_defaults() with diffusion_defaults().
    • sr_model_and_diffusion_defaults(): Returns defaults specifically for Super Resolution, including large_size and small_size, filtered to match the arguments required by sr_create_model_and_diffusion.