INNER CODE UNIT · Python
pocket_cls_loss
QizhiPei/FABind · FABind/fabind/main_fabind.py:398
pocket_cls_loss = args.pocket_cls_loss_weight * pocket_cls_criterion(pocket_cls_pred, pocket_cls.float()) * (protein_out_mask_whole.numel() / protein_out_mask_whole.sum())
pocket_coord_loss = args.pocket_distance_loss_weight * pocket_coord_criterion(pred_pocket_center, data.coords_center)
contact_loss = args.pair_distance_loss_weight * criterion(y_pred, dis_map) if len(dis_map) > 0 else torch.tensor([0])
contact_by_pred_loss = args.pair_distance_loss_weight * criterion(y_pred_by_coord, dis_map) if len(dis_map) > 0 else torch.tensor([0])
contact_distill_loss = args.pair_distance_distill_loss_weight * criterion(y_pred_by_coord, y_pred) if len(y_pred) > 0 else torch.tensor([0])
com_coord_loss = args.coord_loss_weight * com_coord_criterion(com_coord_pred, com_coord) if len(com_coord) > 0 else torch.tensor([0])
sd = ((com_coord_pred.detach() - com_coord) ** 2).sum(dim=-1)
rmsd = scatter_mean(sd, index=compound_batch, dim=0).sqrt().detach()
centroid_pred = scatter_mean(src=com_coord_pred, index=compound_batch, dim=0)
centroid_true = scatter_mean(src=com_coord, index=compound_batch, dim=0)
centroid_dis = (centroid_pred - centroid_true).norm(dim=-1)
loss = com_coord_loss + \
contact_loss + contact_by_pred_loss + contact_distill_loss + \