Update finetune.py

finetune.py

@@ -61,7 +61,7 @@ def get_dataset(data_path, data_info, test_count=10):
 
 
 class CustomFinetuneModel(ESPnetS2TModel):
-    def __init__(self, model, log_every=500):
+    def __init__(self, model, tempdir_path, log_every=500):
         super().__init__(
             vocab_size=model.vocab_size,
             token_list=model.token_list,
@@ -94,6 +94,7 @@ class CustomFinetuneModel(ESPnetS2TModel):
             'loss': 0.0,
             'acc': 0.0
         }
+        self.tempdir_path = tempdir_path
     
     def forward(self, *args, **kwargs):
         out = super().forward(*args, **kwargs)
@@ -104,7 +105,7 @@ class CustomFinetuneModel(ESPnetS2TModel):
         if self.iter_count % self.log_every == 0:
             loss = self.log_stats['loss'] / self.log_every
             acc = self.log_stats['acc'] / self.log_every
-            print(f"[{self.iter_count}] - loss: {loss:.3f} - acc: {acc:.3f}")
+            log(self.tempdir_path, f"[{self.iter_count}] - loss: {loss:.3f} - acc: {acc:.3f}")
             self.log_stats['loss'] = 0.0
             self.log_stats['acc'] = 0.0
 
@@ -159,8 +160,6 @@ def finetune_model(lang, task, tempdir_path, log_every, max_epoch, scheduler, wa
     finetune_config['scheduler_conf']['warmup_steps'] = warmup_steps
     finetune_config['multiple_iterator'] = False
     finetune_config['num_iters_per_epoch'] = None
-    finetune_config['multiprocessing_distributed'] = False
-    finetune_config['distributed'] = False
 
     def build_model_fn(args):
         model, _ = S2TTask.build_model_from_file(
@@ -169,8 +168,8 @@ def finetune_model(lang, task, tempdir_path, log_every, max_epoch, scheduler, wa
             device="cuda" if torch.cuda.is_available() else "cpu",
         )
         model.train()
-        print(f'Trainable parameters: {count_parameters(model)}')
-        model = CustomFinetuneModel(model, log_every=log_every)
+        log(tempdir_path, f'Trainable parameters: {count_parameters(model)}')
+        model = CustomFinetuneModel(model, tempdir_path, log_every=log_every)
         return model
 
     trainer = ez.Trainer(