Skip to content

Implemented batch_size parameter #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JoshuaMHardy wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Implemented batch_size parameter #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/boltz/data/module/inferencev2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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]
Expand All @@ -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
Expand Down
62 changes: 37 additions & 25 deletions src/boltz/data/write/writer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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]}"
Expand Down Expand Up @@ -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]
}
Expand All @@ -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"
Expand All @@ -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"
Expand All @@ -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
Expand Down
17 changes: 13 additions & 4 deletions src/boltz/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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")

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
33 changes: 26 additions & 7 deletions src/boltz/model/modules/diffusionv2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
),
)
Expand Down