Merge pull request #132 from anwai98/main

Update Teacher Initialization for MT Setup
constantinpape · May 26, 2023 · 9bd75ae · 9bd75ae
2 parents a93ff57 + c2eaed8
commit 9bd75ae
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/experiments/probabilistic_domain_adaptation/livecell/punet_adamt.py b/experiments/probabilistic_domain_adaptation/livecell/punet_adamt.py
@@ -78,6 +78,7 @@ def _train_source_target(args, source_cell_type, target_cell_type):
         device=device,
         log_image_interval=100,
         save_root=args.save_root,
+        reinit_teacher=True,
         compile_model=False
     )
     trainer.fit(args.n_iterations)
@@ -119,7 +120,7 @@ def run_evaluation(args):
 
 
 def main():
-    parser = common.get_parser(default_iterations=100000, default_batch_size=4)
+    parser = common.get_parser(default_iterations=10000, default_batch_size=2)
     parser.add_argument("--confidence_threshold", default=None, type=float)
     parser.add_argument("--consensus_masking", action='store_true')
     args = parser.parse_args()

diff --git a/experiments/probabilistic_domain_adaptation/livecell/punet_mean_teacher.py b/experiments/probabilistic_domain_adaptation/livecell/punet_mean_teacher.py
@@ -82,6 +82,7 @@ def _train_source_target(args, source_cell_type, target_cell_type):
         device=device,
         log_image_interval=100,
         save_root=args.save_root,
+        reinit_teacher=False,
         compile_model=False
     )
     trainer.fit(args.n_iterations)