Update train_trocr.py

train_trocr.py

@@ -7,23 +7,22 @@ from evaluate import load
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator, AdamW
 import torchvision.transforms as transforms
 from augments import RandAug, RandRotate
-#import torch_optimizer as optim
 
 parser = argparse.ArgumentParser('arguments for the code')
 
 parser.add_argument('--root_path', type=str, default="",
                     help='Root path to data files.')
-parser.add_argument('--tr_data_path', type=str, default="/data/htr/taulukkosolut/train/trocr/train_new_anno.csv",
+parser.add_argument('--tr_data_path', type=str, default="/path/to/train_data.csv",
                     help='Path to .csv file containing the training data.')
-parser.add_argument('--val_data_path', type=str, default="/data/htr/taulukkosolut/val/trocr/val_new_anno.csv",
+parser.add_argument('--val_data_path', type=str, default="/path/to/val_data.csv",
                     help='Path to .csv file containing the validation data.')
-parser.add_argument('--output_path', type=str, default="./output/no_aug_1beam_07092024/",
+parser.add_argument('--output_path', type=str, default="./output/path/",
                     help='Path for saving training results.')
-parser.add_argument('--model_path', type=str, default="/4tb_01/models/htr/supermalli/202405_fp16/",
+parser.add_argument('--model_path', type=str, default="/model/path/",
                     help='Path to trocr model')
-parser.add_argument('--processor_path', type=str, default="/4tb_01/models/htr/supermalli/202405_fp16/processor",
+parser.add_argument('--processor_path', type=str, default="/processor/path/",
                     help='Path to trocr processor')
-parser.add_argument('--epochs', type=int, default=20,
+parser.add_argument('--epochs', type=int, default=15,
                     help='Training epochs.')
 parser.add_argument('--batch_size', type=int, default=16,
                     help='Training epochs.')
@@ -34,13 +33,6 @@ parser.add_argument('--augment', type=int, default=0,
 
 args = parser.parse_args()
 
-# nohup python train_trocr.py > logs/taulukkosolut_no_aug_1beam_07092024.txt 2>&1 &
-# echo $! > logs/save_pid.txt
-
-#image_size = (224,224)
-#resized_images = []
-# run using 2 GPUs: torchrun --nproc_per_node=2 train_trocr.py
-
 # Initialize processor and model
 processor = TrOCRProcessor.from_pretrained(args.processor_path)
 model = VisionEncoderDecoderModel.from_pretrained(args.model_path)
@@ -53,8 +45,6 @@ wer_metric = load("wer")
 # Load train and validation data to dataframes
 train_df = pd.read_csv(args.tr_data_path)
 val_df = pd.read_csv(args.val_data_path)
-#train_df = train_df.iloc[:10]
-#val_df = val_df.iloc[:5]
 
 # Reset the indices to start from zero
 train_df.reset_index(drop=True, inplace=True)
@@ -129,7 +119,6 @@ model.config.num_beams = 1
 # Set arguments for model training
 # For all argumenst see https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments
 training_args = Seq2SeqTrainingArguments(
-    learning_rate=2.779367469510554e-05,
     predict_with_generate=True,
     eval_strategy="epoch",
     save_strategy="epoch",
@@ -162,7 +151,7 @@ def compute_metrics(pred):
     return {"cer": cer, "wer": wer}
 
 
-# instantiate trainer
+# Instantiate trainer
 # For all parameters see: https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainer
 trainer = Seq2SeqTrainer(
     model=model,
@@ -178,5 +167,4 @@ trainer = Seq2SeqTrainer(
 trainer.train()
 #trainer.train(resume_from_checkpoint = True)
 model.save_pretrained(args.output_path)
-processor.save_pretrained(args.output_path + "/processor")
-
+processor.save_pretrained(args.output_path + "/processor")