From f43ef69f0dc4bcc7dc871b6a1474d0c3d0700f88 Mon Sep 17 00:00:00 2001 From: Jordon Brooks Date: Sun, 13 Aug 2023 14:21:47 +0100 Subject: [PATCH] Added support for continued model training --- train_model.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/train_model.py b/train_model.py index 74aaae6..bddebfd 100644 --- a/train_model.py +++ b/train_model.py @@ -90,7 +90,11 @@ def main(): training_videos = all_videos[:split_index] validation_videos = all_videos[split_index:] - model = VideoCompressionModel() + if args.continue_training: + model = tf.keras.models.load_model(args.continue_training) + else: + model = VideoCompressionModel() + # Set optimizer and compile the model optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)