Skip to content
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

PoC: Better reordering #129

Draft
wants to merge 9 commits into
base: main
Choose a base branch
from
Draft

PoC: Better reordering #129

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
d06595d
poc: adds a better way to handle rearrangement of symmetric entities
hypnopump Aug 30, 2023
1c94a4b
use ready frames
hypnopump Aug 30, 2023
8c24b80
use precomputed frames, indexed properly
hypnopump Aug 30, 2023
7602039
xx_fape as k order cant be assumed
hypnopump Aug 30, 2023
ccebda4
better type hints
hypnopump Aug 30, 2023
8491d77
add chunking to the mask as well
hypnopump Aug 31, 2023
4a0f28d
typing complains
hypnopump Sep 4, 2023
7a55551
runs with the train_multimer_demo.sh
hypnopump Sep 4, 2023
a62d2b8
pr comments
hypnopump Sep 5, 2023
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
20 changes: 10 additions & 10 deletions unifold/loss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,10 @@ def forward(self, model, batch, reduce=True):
# return config in model.
out, config = model(batch)
num_recycling = batch["msa_feat"].shape[0]

# remove recyling dim
batch = tensor_tree_map(lambda t: t[-1, ...], batch)

loss, sample_size, logging_output = self.loss(out, batch, config)
logging_output["num_recycling"] = num_recycling
return loss, sample_size, logging_output
Expand All @@ -57,7 +57,7 @@ def loss(self, out, batch, config):
if "renamed_atom14_gt_positions" not in out.keys():
batch.update(
compute_renamed_ground_truth(batch, out["sm"]["positions"]))

loss_dict = {}
loss_fns = {
"chain_centre_mass": lambda: chain_centre_mass_loss(
Expand Down Expand Up @@ -143,11 +143,11 @@ def loss(self, out, batch, config):
with torch.no_grad():
seq_len = torch.sum(batch["seq_mask"].float(), dim=-1)
seq_length_weight = seq_len**0.5

assert (
len(seq_length_weight.shape) == 1 and seq_length_weight.shape[0] == bsz
), seq_length_weight.shape

for loss_name, loss_fn in loss_fns.items():
weight = config[loss_name].weight
if weight > 0.:
Expand All @@ -159,7 +159,7 @@ def loss(self, out, batch, config):
if any(torch.isnan(loss)) or any(torch.isinf(loss)):
logging.warning(f"{loss_name} loss is NaN. Skipping...")
loss = loss.new_tensor(0.0, requires_grad=True)

cum_loss = cum_loss + weight * loss

for key in loss_dict:
Expand Down Expand Up @@ -208,10 +208,10 @@ def forward(self, model, batch, reduce=True):
# return config in model.
out, config = model(features)
num_recycling = features["msa_feat"].shape[0]

# remove recycling dim
features = tensor_tree_map(lambda t: t[-1, ...], features)

# perform multi-chain permutation alignment.
if labels:
with torch.no_grad():
Expand All @@ -230,12 +230,12 @@ def forward(self, model, batch, reduce=True):
)
new_labels.append(cur_new_labels)
new_labels = data_utils.collate_dict(new_labels, dim=0)

# check for consistency of label and feature.
assert (new_labels["aatype"] == features["aatype"]).all()
features.update(new_labels)

loss, sample_size, logging_output = self.loss(out, features, config)
logging_output["num_recycling"] = num_recycling

return loss, sample_size, logging_output
Loading