finishing up IT benchmarks

Author: dmoi

config_notebook_1k_epochs.yaml

@@ -11,8 +11,8 @@ overwrite: true
 
 # Training hyperparameters (from notebook)
 epochs: 1000
-batch_size: 8
-gradient_accumulation_steps: 2
+batch_size: 16
+gradient_accumulation_steps: 1
 seed: 0
 
 # Model architecture (from notebook)
@@ -49,12 +49,12 @@ clip_grad: true
 # Loss weights (from notebook cell)
 # Note: These are initial weights - notebook has weight schedulers (currently commented out)
 edgeweight: 0.1
-# logitweight = 0.1
-xweight: 5.0  
+logitweight: 0.01
+xweight: 0.5
 # fft2weight = 0.01
-# vqweight = 0.1
-# angles_weight = 0.1
-# ss_weight = 0.1
+vqweight: 0.1
+angles_weight: 0.1
+ss_weight: 0.1
 
 # Commitment cost scheduling (from notebook encoder config)
 use_commitment_scheduling: true

foldtree2/learn_lightning.py

@@ -1,4 +1,5 @@
 # learn_lightning.py - PyTorch Lightning training with multi-GPU support
+import gc
 import os
 # Set CUDA memory allocator to use expandable segments to reduce fragmentation
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
@@ -407,6 +408,10 @@ def training_step(self, batch, batch_idx):
             current_commitment = self.encoder.vector_quantizer.get_commitment_cost()
             self.log('train/commitment_cost', current_commitment, on_step=False, on_epoch=True, batch_size=batch_size)
 
+            
+        # Clear CUDA cache
+        torch.cuda.empty_cache()
+        gc.collect()
         return loss
 
     def configure_optimizers(self):

foldtree2/learn_monodecoder.py

@@ -194,6 +194,9 @@ def print_about():
 parser.add_argument('--ss-weight', type=float, default=0.25,
                     help='Weight for secondary structure loss (default: 0.25)')
 
+parser.add_argument('--nconv-layers', type=int, default=3,
+                    help='Number of convolutional layers in the geometry decoder (default: 3)')
+
 # Tensor Core precision
 parser.add_argument('--tensor-core-precision', type=str, default='high',
                     choices=['highest', 'high', 'medium'],
@@ -536,7 +539,7 @@ def decode_batch_reconstruction(encoder, decoder, z_batch, device, converter, ve
 				'in_channels': {'res': args.embedding_dim, 'godnode4decoder': ndim_godnode, 'foldx': 23, 'fft2r': ndim_fft2r, 'fft2i': ndim_fft2i},
 				'concat_positions': False,
 				'conv_channels': [2*hidden_size, hidden_size, hidden_size],
-				'kernel_sizes': [3, 3 ,3 ],
+				'kernel_sizes': [3 ]*args.nconv_layers,
 				'FFT2decoder_hidden': [hidden_size//2, hidden_size//2],
 				'contactdecoder_hidden': [hidden_size//2, hidden_size//4],
 				'ssdecoder_hidden': [hidden_size//2, hidden_size//2],