support tensorboard

finetune.py

@@ -1,3 +1,5 @@
+from transformers.integrations import TensorBoardCallback
+from torch.utils.tensorboard import SummaryWriter
 from transformers import TrainingArguments
 from transformers import Trainer, HfArgumentParser
 from transformers import AutoTokenizer
@@ -114,6 +116,7 @@ def save_model(self, output_dir=None, _internal_call=False):
 
 
 def main():
+    writer = SummaryWriter()
     finetune_args, training_args = HfArgumentParser(
         (FinetuneArguments, TrainingArguments)
     ).parse_args_into_dataclasses()
@@ -149,10 +152,11 @@ def main():
         model=model,
         train_dataset=dataset,
         args=training_args,
+        callbacks=[TensorBoardCallback(writer)],
         data_collator=data_collator,
     )
     trainer.train()
-
+    writer.close()
     # save model
     model.save_pretrained(training_args.output_dir)
 

requirements.txt

@@ -8,7 +8,8 @@ transformers==4.27.1
 icetk
 cpm_kernels==1.0.11
 torch>=1.13.1
+tensorboard
 
 # 
 datasets==2.10.1
-git+https://github.com/huggingface/peft.git  # 最新版本 >=0.3.0.dev0
+git+https://github.com/huggingface/peft.git  # 最新版本 >=0.3.0.dev0