finalized flash_examples/face_detection.py for user api (Lightning-Un…

…iverse#508)
borhanMorphy · Jul 16, 2021 · 57226b6 · 57226b6
1 parent 02c14c6
commit 57226b6
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/flash_examples/face_detection.py b/flash_examples/face_detection.py
@@ -30,14 +30,18 @@
 datamodule = FaceDetectionData.from_fddb(
     data_folder="data/",
     train_data=True,
+    val_data=True,
 )
 
 # 2. Build the task
 model = FaceDetector(model="lffd_slim")
 
 # 3. Create the trainer and finetune the model
-trainer = flash.Trainer(max_epochs=3)
-trainer.finetune(model, datamodule=datamodule)
+trainer = flash.Trainer(max_epochs=3,
+    limit_train_batches=0.1,
+    limit_val_batches=0.1)
+
+trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
 # 4. Detect faces in a few images!
 predictions = model.predict([