From 2cee3f56e346354f0e8c90f454f9c34620eddc91 Mon Sep 17 00:00:00 2001 From: JoshuaMHardy <173318912+JoshuaMHardy@users.noreply.github.com> Date: Thu, 26 Feb 2026 02:35:59 +0000 Subject: [PATCH] Implemented batch_size parameter --- src/boltz/data/module/inferencev2.py | 4 ++ src/boltz/data/write/writer.py | 62 +++++++++++++++----------- src/boltz/main.py | 17 +++++-- src/boltz/model/modules/diffusionv2.py | 33 +++++++++++--- 4 files changed, 80 insertions(+), 36 deletions(-) diff --git a/src/boltz/data/module/inferencev2.py b/src/boltz/data/module/inferencev2.py index 590297d26..3b7df3eff 100644 --- a/src/boltz/data/module/inferencev2.py +++ b/src/boltz/data/module/inferencev2.py @@ -324,6 +324,7 @@ def __init__( msa_dir: Path, mol_dir: Path, num_workers: int, + batch_size: int = 1, constraints_dir: Optional[Path] = None, template_dir: Optional[Path] = None, extra_mols_dir: Optional[Path] = None, @@ -344,6 +345,8 @@ def __init__( The path to the moldir. num_workers : int The number of workers to use. + batch_size : int + The number of inputs to process in parallel per batch. Default is 1. constraints_dir : Optional[Path] The path to the constraints directory. template_dir : Optional[Path] @@ -356,6 +359,7 @@ def __init__( """ super().__init__() self.num_workers = num_workers + self.batch_size = batch_size self.manifest = manifest self.target_dir = target_dir self.msa_dir = msa_dir diff --git a/src/boltz/data/write/writer.py b/src/boltz/data/write/writer.py index 984be2ae5..05a5195f3 100644 --- a/src/boltz/data/write/writer.py +++ b/src/boltz/data/write/writer.py @@ -65,21 +65,30 @@ def write_on_batch_end( records: list[Record] = batch["record"] # Get the predictions + # coords has shape [B*M, L, 3] where B=batch_size, M=diffusion_samples. coords = prediction["coords"] - coords = coords.unsqueeze(0) - - pad_masks = prediction["masks"] - - # Get ranking - if "confidence_score" in prediction: - argsort = torch.argsort(prediction["confidence_score"], descending=True) - idx_to_rank = {idx.item(): rank for rank, idx in enumerate(argsort)} - # Handles cases where confidence summary is False - else: - idx_to_rank = {i: i for i in range(len(records))} + pad_masks = prediction["masks"] # [B, L] + + B = len(records) + total_samples = coords.shape[0] + M = total_samples // B # diffusion_samples per batch item + + # Iterate over the batch items + for b, (record, pad_mask) in enumerate(zip(records, pad_masks)): + # Slice the coords and all per-sample prediction tensors for this item. + start = b * M + end = start + M + coord = coords[start:end] # [M, L, 3] + + # Per-item ranking over the M diffusion samples + if "confidence_score" in prediction: + conf_scores = prediction["confidence_score"][start:end] # [M] + argsort = torch.argsort(conf_scores, descending=True) + idx_to_rank = {idx.item(): rank for rank, idx in enumerate(argsort)} + # Handles cases where confidence summary is False + else: + idx_to_rank = {i: i for i in range(M)} - # Iterate over the records - for record, coord, pad_mask in zip(records, coords, pad_masks): # Load the structure path = self.data_dir / f"{record.id}.npz" if self.boltz2: @@ -96,7 +105,10 @@ def write_on_batch_end( # Remove masked chains completely structure = structure.remove_invalid_chains() - for model_idx in range(coord.shape[0]): + for model_idx in range(M): + # global_idx indexes into the flat [B*M, ...] prediction tensors + global_idx = start + model_idx + # Get model coord model_coord = coord[model_idx] # Unpad @@ -153,7 +165,7 @@ def write_on_batch_end( # Get plddt's plddts = None if "plddt" in prediction: - plddts = prediction["plddt"][model_idx] + plddts = prediction["plddt"][global_idx] # Create path name outname = f"{record.id}_model_{idx_to_rank[model_idx]}" @@ -198,15 +210,15 @@ def write_on_batch_end( "complex_pde", "complex_ipde", ]: - confidence_summary_dict[key] = prediction[key][model_idx].item() + confidence_summary_dict[key] = prediction[key][global_idx].item() confidence_summary_dict["chains_ptm"] = { - idx: prediction["pair_chains_iptm"][idx][idx][model_idx].item() + idx: prediction["pair_chains_iptm"][idx][idx][global_idx].item() for idx in prediction["pair_chains_iptm"] } confidence_summary_dict["pair_chains_iptm"] = { idx1: { idx2: prediction["pair_chains_iptm"][idx1][idx2][ - model_idx + global_idx ].item() for idx2 in prediction["pair_chains_iptm"][idx1] } @@ -221,7 +233,7 @@ def write_on_batch_end( ) # Save plddt - plddt = prediction["plddt"][model_idx] + plddt = prediction["plddt"][global_idx] path = ( struct_dir / f"plddt_{record.id}_model_{idx_to_rank[model_idx]}.npz" @@ -230,7 +242,7 @@ def write_on_batch_end( # Save pae if "pae" in prediction: - pae = prediction["pae"][model_idx] + pae = prediction["pae"][global_idx] path = ( struct_dir / f"pae_{record.id}_model_{idx_to_rank[model_idx]}.npz" @@ -239,17 +251,17 @@ def write_on_batch_end( # Save pde if "pde" in prediction: - pde = prediction["pde"][model_idx] + pde = prediction["pde"][global_idx] path = ( struct_dir / f"pde_{record.id}_model_{idx_to_rank[model_idx]}.npz" ) np.savez_compressed(path, pde=pde.cpu().numpy()) - - # Save embeddings + + # Save embeddings (s/z are [B, ...]; slice for this batch item) if self.write_embeddings and "s" in prediction and "z" in prediction: - s = prediction["s"].cpu().numpy() - z = prediction["z"].cpu().numpy() + s = prediction["s"][b].cpu().numpy() + z = prediction["z"][b].cpu().numpy() path = ( struct_dir diff --git a/src/boltz/main.py b/src/boltz/main.py index 4a3750fec..92287836c 100644 --- a/src/boltz/main.py +++ b/src/boltz/main.py @@ -296,7 +296,7 @@ def check_inputs(data: Path) -> list[Path]: # Check if data is a directory if data.is_dir(): - data: list[Path] = list(data.glob("*")) + data: list[Path] = sorted(data.glob("*")) # Filter out non .fasta or .yaml files, raise # an error on directory and other file types @@ -725,7 +725,7 @@ def process_inputs( records_dir = out_dir / "processed" / "records" if records_dir.exists(): # Load existing records - existing = [Record.load(p) for p in records_dir.glob("*.json")] + existing = [Record.load(p) for p in sorted(records_dir.glob("*.json"))] processed_ids = {record.id for record in existing} # Filter to missing only @@ -802,8 +802,8 @@ def process_inputs( for path in tqdm(data): process_input_partial(path) - # Load all records and write manifest - records = [Record.load(p) for p in records_dir.glob("*.json")] + # Load all records and write manifest (sorted for deterministic order) + records = [Record.load(p) for p in sorted(records_dir.glob("*.json"))] manifest = Manifest(records) manifest.dump(out_dir / "processed" / "manifest.json") @@ -910,6 +910,12 @@ def cli() -> None: help="The number of dataloader workers to use for prediction. Default is 2.", default=2, ) +@click.option( + "--batch_size", + type=int, + help="The number of inputs to process in parallel per batch. Default is 1.", + default=1, +) @click.option( "--override", is_flag=True, @@ -1058,6 +1064,7 @@ def predict( # noqa: C901, PLR0915, PLR0912 write_full_pde: bool = False, output_format: Literal["pdb", "mmcif"] = "mmcif", num_workers: int = 2, + batch_size: int = 1, override: bool = False, seed: Optional[int] = None, use_msa_server: bool = False, @@ -1275,6 +1282,7 @@ def predict( # noqa: C901, PLR0915, PLR0912 msa_dir=processed.msa_dir, mol_dir=mol_dir, num_workers=num_workers, + batch_size=batch_size, constraints_dir=processed.constraints_dir, template_dir=processed.template_dir, extra_mols_dir=processed.extra_mols_dir, @@ -1362,6 +1370,7 @@ def predict( # noqa: C901, PLR0915, PLR0912 msa_dir=processed.msa_dir, mol_dir=mol_dir, num_workers=num_workers, + batch_size=batch_size, constraints_dir=processed.constraints_dir, template_dir=processed.template_dir, extra_mols_dir=processed.extra_mols_dir, diff --git a/src/boltz/model/modules/diffusionv2.py b/src/boltz/model/modules/diffusionv2.py index fd1af56d3..5579e0254 100644 --- a/src/boltz/model/modules/diffusionv2.py +++ b/src/boltz/model/modules/diffusionv2.py @@ -379,17 +379,36 @@ def sample( with torch.no_grad(): atom_coords_denoised = torch.zeros_like(atom_coords_noisy) - sample_ids = torch.arange(multiplicity).to(atom_coords_noisy.device) - sample_ids_chunks = sample_ids.chunk( - multiplicity % max_parallel_samples + 1 - ) - - for sample_ids_chunk in sample_ids_chunks: + # atom_mask has shape [batch_size * multiplicity, L] after repeat_interleave. + # Build a [multiplicity, batch_size] index matrix so that each chunk of + # `max_parallel_samples` diffusion steps contains one sample from every + # batch item, matching the repeat_interleave ordering the network expects. + batch_size = atom_mask.shape[0] // multiplicity + sample_ids_2d = ( + torch.arange( + batch_size * multiplicity, device=atom_coords_noisy.device + ) + .view(batch_size, multiplicity) + .t() + ) # [multiplicity, batch_size] + + for chunk_start in range(0, multiplicity, max_parallel_samples): + chunk_end = min(chunk_start + max_parallel_samples, multiplicity) + chunk_m = chunk_end - chunk_start + # Transpose [chunk_m, batch_size] -> [batch_size, chunk_m] then + # flatten: gives [B0_sc..B0_{sc+k}, B1_sc..B1_{sc+k}, ...] which + # aligns with the network's repeat_interleave(chunk_m) on feats. + sample_ids_chunk = ( + sample_ids_2d[chunk_start:chunk_end] + .t() + .contiguous() + .view(-1) + ) atom_coords_denoised_chunk = self.preconditioned_network_forward( atom_coords_noisy[sample_ids_chunk], t_hat, network_condition_kwargs=dict( - multiplicity=sample_ids_chunk.numel(), + multiplicity=chunk_m, **network_condition_kwargs, ), )