fix discrete cat grad and losses

Author: cactuskid

foldtree2/notebooks/experiments/test_monodecoders.ipynb

@@ -529,8 +529,8 @@ def aa_reconstruction_loss(x, recon_x , normalize = False):
 		The function expects one-hot encoded targets (not class indices).
 		Cross-entropy will internally convert these to class indices.
 	"""
-
-	return F.cross_entropy(recon_x, x)
+	target = x.argmax(dim=-1)   # integers in [0..19], no gradients needed
+	return F.cross_entropy(recon_x, target)
 
 
 def ss_reconstruction_loss(ss, recon_ss, mask_plddt=False, plddt_threshold=0.3 , plddt_mask = None , normalize = False):
@@ -570,18 +570,20 @@ def ss_reconstruction_loss(ss, recon_ss, mask_plddt=False, plddt_threshold=0.3 ,
 		Masking by pLDDT is recommended for AlphaFold structures where low-confidence
 		regions may have unreliable secondary structure assignments.
 	"""
+	
+	target = ss.argmax(dim=-1)
 	if mask_plddt:
 		# Create boolean mask for high-confidence residues
 		mask = (plddt_mask > plddt_threshold).squeeze()
 		if mask.sum() > 0:
 			# Compute loss only on masked residues
-			ss_loss = F.cross_entropy(recon_ss[mask], ss[mask])
+			ss_loss = F.cross_entropy(recon_ss[mask], target[mask])
 		else:
 			# No residues pass threshold - return zero loss to prevent NaN
 			ss_loss = torch.tensor(0.0, device=recon_ss.device)
 	else:
 		# Compute loss on all residues
-		ss_loss = F.cross_entropy(recon_ss, ss)
+		ss_loss = F.cross_entropy(recon_ss, target)
 	return ss_loss
 
 

foldtree2/src/losses/losses.py

@@ -486,7 +486,7 @@ def __init__(self, in_channels={'res': 10, 'godnode4decoder': 5, 'foldx': 23},
 				nn.Linear(anglesdecoder_hidden[0], anglesdecoder_hidden[1]),
 				nn.GELU(),
 				nn.Linear(anglesdecoder_hidden[1], 3),
-				nn.Softmax(dim=-1)
+				#nn.Softmax(dim=-1)
 			)
 
 		# Bond angles prediction
@@ -758,7 +758,7 @@ def __init__(
 			nn.Conv1d(AAdecoder_hidden[1], AAdecoder_hidden[2] if len(AAdecoder_hidden) > 2 else AAdecoder_hidden[1], kernel_size=3, padding=1),
 			nn.GELU(),
 			nn.Conv1d(AAdecoder_hidden[2] if len(AAdecoder_hidden) > 2 else AAdecoder_hidden[1], 20, kernel_size=1),  # 20 amino acids
-			nn.Softmax(dim=1)  # Probabilities for amino acid classes
+			#nn.Softmax(dim=1)  # Probabilities for amino acid classes
 		)
 
 		# Optional secondary structure prediction head
@@ -772,7 +772,7 @@ def __init__(
 				nn.Conv1d(ssdecoder_hidden[1], ssdecoder_hidden[2] if len(ssdecoder_hidden) > 2 else ssdecoder_hidden[1], kernel_size=3, padding=1),
 				nn.GELU(),
 				nn.Conv1d(ssdecoder_hidden[2] if len(ssdecoder_hidden) > 2 else ssdecoder_hidden[1], 3, kernel_size=1),  # 3 SS classes
-				nn.Softmax(dim=1)  # Probabilities for SS classes
+				#nn.Softmax(dim=1)  # Probabilities for SS classes
 			)
 
 	def forward(self, data, **kwargs):
@@ -862,7 +862,7 @@ def forward(self, data, **kwargs):
 				# AA prediction
 				aa_out = self.head['aa_cnn'](x_cnn)  # (1, 20, seq_len)
 				aa_out = aa_out.permute(2, 0, 1).squeeze(1)  # (seq_len, 20)
-				aa_list.append(F.log_softmax(aa_out, dim=-1))
+				aa_list.append(aa_out)
 
 				# SS prediction if enabled
 				if self.output_ss and 'ss_cnn' in self.head:
@@ -880,7 +880,6 @@ def forward(self, data, **kwargs):
 			# AA prediction
 			aa = self.head['aa_cnn'](x_cnn)  # (1, 20, seq_len)
 			aa = aa.permute(2, 0, 1).squeeze(1)  # (seq_len, 20)
-			aa = F.log_softmax(aa, dim=-1)
 
 			# SS prediction if enabled
 			if self.output_ss and 'ss_cnn' in self.head:
@@ -995,7 +994,7 @@ def __init__(
 				nn.Conv1d(AAdecoder_hidden[1], AAdecoder_hidden[2], kernel_size=3, padding=1),
 				nn.GELU(),
 				nn.Conv1d(AAdecoder_hidden[2], 20, kernel_size=1),
-				nn.Softmax(dim=-1)
+				#nn.Softmax(dim=-1)
 			)
 		else:
 			# DNN decoder for amino acid prediction
@@ -1007,7 +1006,7 @@ def __init__(
 				nn.Linear(AAdecoder_hidden[1], AAdecoder_hidden[2]),
 				nn.GELU(),
 				nn.Linear(AAdecoder_hidden[2], 20),
-				nn.Softmax(dim=-1)
+				#nn.Softmax(dim=-1)
 			)
 
 		# Optional secondary structure prediction head
@@ -1020,7 +1019,7 @@ def __init__(
 				nn.Linear(AAdecoder_hidden[1], AAdecoder_hidden[2]),
 				nn.GELU(),
 				nn.Linear(AAdecoder_hidden[2], 3),
-				nn.Softmax(dim=-1)
+				#nn.Softmax(dim=-1)
 			)
 
 	def forward(self, data, **kwargs):
@@ -1089,7 +1088,7 @@ def forward(self, data, **kwargs):
 					xi_cnn = xi.permute(1, 0).unsqueeze(0)  # (1, d_model, seq_len)
 					xi_cnn = self.head['cnn_decoder'](xi_cnn)  # (1, 20, seq_len)
 					xi_cnn = xi_cnn.permute(2, 0, 1).squeeze(1)  # (seq_len, 20)
-					aa_list.append(F.log_softmax(xi_cnn[:seq_len, :], dim=-1))
+					aa_list.append(xi_cnn[:seq_len, :])
 				else:
 					aa_list.append(self.head['dnn_decoder'](xi[:seq_len, 0]))
 					if 'ss_head' in self.head:
@@ -1108,7 +1107,7 @@ def forward(self, data, **kwargs):
 				x_cnn = x.permute(1, 2, 0)  # (batch, d_model, seq_len)
 				x_cnn = self.head['cnn_decoder'](x_cnn)  # (batch, 20, seq_len)
 				x_cnn = x_cnn.permute(2, 0, 1)  # (seq_len, batch, 20)
-				aa = F.log_softmax(x_cnn, dim=-1)
+				aa = x_cnn.squeeze(1)  # (seq_len, 20)
 			else:
 				aa = self.head['dnn_decoder'](x)
 				if 'ss_head' in self.head:
@@ -1261,7 +1260,7 @@ def __init__(
 					nn.Conv1d(ssdecoder_hidden[1], ssdecoder_hidden[2] if len(ssdecoder_hidden) > 2 else ssdecoder_hidden[1], kernel_size=3, padding=1),
 					nn.GELU(),
 					nn.Conv1d(ssdecoder_hidden[2] if len(ssdecoder_hidden) > 2 else ssdecoder_hidden[1], 3, kernel_size=1),
-					nn.LogSoftmax(dim=-1)
+					#nn.LogSoftmax(dim=-1)
 
 				)
 			else:
@@ -1274,7 +1273,7 @@ def __init__(
 					nn.Linear(ssdecoder_hidden[1], ssdecoder_hidden[2] if len(ssdecoder_hidden) > 2 else ssdecoder_hidden[1]),
 					nn.GELU(),
 					nn.Linear(ssdecoder_hidden[2] if len(ssdecoder_hidden) > 2 else ssdecoder_hidden[1], 3),
-					nn.LogSoftmax(dim=-1)
+					#nn.LogSoftmax(dim=-1)
 				)
 
 		# Bond angles prediction head (phi, psi, omega)