log_dir -> log_file for dllogger

openfold/utils/logger.py

@@ -27,8 +27,8 @@ def is_main_process():
 
 
 class PerformanceLoggingCallback(Callback):
-    def __init__(self, log_dir, global_batch_size, warmup_steps: int = 0, profile: bool = False):
-        logger.init(backends=[JSONStreamBackend(Verbosity.VERBOSE, log_dir), StdOutBackend(Verbosity.VERBOSE)])
+    def __init__(self, log_file, global_batch_size, warmup_steps: int = 0, profile: bool = False):
+        logger.init(backends=[JSONStreamBackend(Verbosity.VERBOSE, log_file), StdOutBackend(Verbosity.VERBOSE)])
         self.warmup_steps = warmup_steps
         self.global_batch_size = global_batch_size
         self.step = 0

train_openfold.py

@@ -152,7 +152,7 @@ def main(args):
     if args.log_performance:
         global_batch_size = args.num_nodes * args.gpus
         perf = PerformanceLoggingCallback(
-            log_dir=args.output_dir,
+            log_file=os.path.join(args.output_dir, "performance_log.json"),
             global_batch_size=global_batch_size,
         )
         callbacks.append(perf)