torchcell.models.gpu_edge_mask_generator.GPUEdgeMaskGenerator
- class GPUEdgeMaskGenerator(cell_graph: HeteroData, device: device)[source]
Bases:
ModuleGenerates edge masks on GPU from perturbation indices.
Stores incidence cache and base masks on GPU to avoid CPU→GPU transfer of large mask tensors during training.
- Architecture:
One-time GPU setup (during __init__): - Build incidence cache on GPU (~19 MB) - Create base edge masks (all True)
Per-batch (in training_step): - Transfer perturbation indices (~448 bytes) - Generate masks on GPU using incidence cache - Use masks in message passing
- Expected Performance:
Data transfer: 65 MB → 0.0004 MB per batch (162,500x reduction)
Training speed: ~0.3 it/s → 8-10 it/s (25-33x speedup)
- generate_batch_masks(batch_perturbation_indices: list[Tensor], batch_size: int) dict[tuple[str, str, str], Tensor][source]
Generate edge masks for a batch of perturbations.
For batched graphs in PyG, we need to concatenate masks across samples with proper offsets matching the concatenated edge_index.
- Parameters:
batch_perturbation_indices – List of perturbation index tensors, one per sample in batch. Each tensor contains indices of perturbed genes.
batch_size – Number of samples in batch
- Returns:
- {edge_type: concatenated_mask} where mask length
matches the batched edge_index length
- Return type:
edge_mask_dict
Example
>>> # Sample 1: perturb genes [42, 100] >>> # Sample 2: perturb genes [7] >>> pert_indices = [torch.tensor([42, 100]), torch.tensor([7])] >>> masks = generator.generate_batch_masks(pert_indices, batch_size=2) >>> # masks[edge_type] has length = 2 * num_edges (concatenated)
- generate_batch_masks_vectorized(batch_perturbation_indices: list[Tensor], batch_size: int) dict[tuple[str, str, str], Tensor][source]
VECTORIZED: Generate edge masks with zero Python loops and zero .item() calls.
This is the optimized version that eliminates GPU→CPU synchronizations. Expected speedup: 10-20x compared to the loop-based version.
- Parameters:
batch_perturbation_indices – List of perturbation index tensors, one per sample in batch
batch_size – Number of samples in batch
- Returns:
{edge_type: concatenated_mask}
- Return type:
edge_mask_dict
- generate_single_mask(perturbation_indices: Tensor) dict[tuple[str, str, str], Tensor][source]
Generate edge masks for a single sample.
Useful for validation or single-sample inference.
- Parameters:
perturbation_indices – Tensor of perturbed gene indices
- Returns:
{edge_type: mask_tensor} for one sample
- Return type:
edge_mask_dict