Update train_dtat.py

train_dtat.py

@@ -82,11 +82,12 @@ def get_lr(it, config):
     if it < config.warmup_iters:
         return config.learning_rate * it / config.warmup_iters
     
-    # Cosine decay
+    # Cosine decay with minimum learning rate
     if config.decay_lr:
-        decay_ratio = (it - config.warmup_iters) / (config.max_iters - config.warmup_iters)
-        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # Cosine decay
-        return config.learning_rate * coeff
+        decay_ratio = (it - config.warmup_iters) / (config.lr_decay_iters - config.warmup_iters)
+        decay_ratio = min(decay_ratio, 1.0)  # Cap at 1.0
+        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+        return config.min_lr + coeff * (config.learning_rate - config.min_lr)
     
     return config.learning_rate
 
@@ -280,7 +281,7 @@ def main():
             })
             
             # Save regular checkpoint every 5000 iterations
-            if iter_num % 5000 == 0:
+            if iter_num % 1000 == 0:
                 checkpoint = {
                     'model_state_dict': model.state_dict(),
                     'optimizer_state_dict': optimizer.state_dict(),