Update train_trocr.py

train_trocr.py

@@ -5,7 +5,6 @@ import argparse
 from evaluate import load
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator, AdamW
 import torchvision.transforms as transforms
-#import torch_optimizer as optim
 
 from dataset import TextlineDataset
 
@@ -13,14 +12,12 @@ 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/trocr_data/trocr_tuomiokirjat/train/trocr/data.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/trocr_data/trocr_tuomiokirjat/val/trocr/data.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="/koodit/htr/text_recognition/trocr/tuomiokirjat/models/22112023/",
+parser.add_argument('--output_path', type=str, default="/output/path/",
                     help='Path for saving training results.')
-parser.add_argument('--resume_path', type=str, default="/koodit/htr/text_recognition/trocr/tuomiokirjat/models/22112023",
-                    help='Path to the previous model')
 parser.add_argument('--batch_size', type=int, default=24,
                     help='Batch size per device.')
 parser.add_argument('--epochs', type=int, default=13,
@@ -28,20 +25,12 @@ parser.add_argument('--epochs', type=int, default=13,
 
 args = parser.parse_args()
 
-# nohup python train_trocr.py > logs/tuomiokirjat_resume_23112023.txt 2>&1 &
-# echo $! > logs/save_pid.txt
-
-# run using 2 GPUs: torchrun --nproc_per_node=2 train_trocr.py > logs/tuomiokirjat_22112023.txt 2>&1 &
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print('Device: ', device)
 
 # Initialize processor and model
-#processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
-#model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten") 
-processor =TrOCRProcessor.from_pretrained(args.resume_path + "/processor")
-model = VisionEncoderDecoderModel.from_pretrained(args.resume_path + "/checkpoint-13094")
-
+processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten") 
 model.to(device)
 
 # Initialize metrics
@@ -51,8 +40,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[:50]
-#val_df = val_df.iloc[:10]
 
 # Reset the indices to start from zero
 train_df.reset_index(drop=True, inplace=True)
@@ -88,7 +75,7 @@ model.config.length_penalty = 2.0
 model.config.num_beams = 4
 
 # Set arguments for model training
-# For all argumenst see https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments
+# For all arguments see https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments
 training_args = Seq2SeqTrainingArguments(
     predict_with_generate=True,
     evaluation_strategy="epoch",
@@ -122,7 +109,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,
@@ -138,5 +125,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")