add num_timesteps

checkpoints/ROMA_UNSB_001/loss_log.txt

@@ -1,2 +1,3 @@
 ================ Training Loss (Sun Feb 23 15:46:44 2025) ================
 ================ Training Loss (Sun Feb 23 15:52:29 2025) ================
+================ Training Loss (Sun Feb 23 16:00:07 2025) ================

checkpoints/ROMA_UNSB_001/train_opt.txt

@@ -76,6 +76,7 @@ nce_includes_all_negatives_from_minibatch: False
            serial_batches: False                         
 stylegan2_G_num_downsampling: 1                             
                    suffix:                               
+                      tau: 0.1                           
          update_html_freq: 1000                          
                   use_idt: False                         
                   verbose: False                         

models/roma_unsb_model.py

@@ -222,6 +222,7 @@ class RomaUnsbModel(BaseModel):
         parser.add_argument('--atten_layers', type=str, default='1,3,5', help='compute Cross-Similarity on which layers')
         
         parser.add_argument('--tau', type=float, default=0.1, help='used in unsb')
+        parser.add_argument('--num_timesteps', type=int, default=10, help='used in unsb')
 
         parser.set_defaults(pool_size=0)  # no image pooling
 

scripts/train.sh

@@ -28,4 +28,5 @@ python train.py \
     --num_patches 256 \
     --flip_equivariance False \
     --eta_ratio 0.1 \
-    --tau 0.1
+    --tau 0.1 \
+    --num_timesteps 10