Add ability to configure model training epochs

Author: PawelPeczek-Roboflow

roboflow/__init__.py

@@ -15,7 +15,7 @@
 from roboflow.models import CLIPModel, GazeModel  # noqa: F401
 from roboflow.util.general import write_line
 
-__version__ = "1.2.10"
+__version__ = "1.2.11"
 
 
 def check_key(api_key, model, notebook, num_retries=0):

roboflow/adapters/rfapi.py

@@ -58,6 +58,7 @@ def start_version_training(
     speed: Optional[str] = None,
     checkpoint: Optional[str] = None,
     model_type: Optional[str] = None,
+    epochs: Optional[int] = None,
 ):
     """
     Start a training job for a specific version.
@@ -74,6 +75,8 @@ def start_version_training(
     if model_type is not None:
         # API expects camelCase
         data["modelType"] = model_type
+    if epochs is not None:
+        data["epochs"] = epochs
 
     response = requests.post(url, json=data)
     if not response.ok:

roboflow/core/version.py

@@ -296,14 +296,15 @@ def export(self, model_format=None) -> bool | None:
         else:
             raise RuntimeError(f"Unexpected export {export_info}")
 
-    def train(self, speed=None, model_type=None, checkpoint=None, plot_in_notebook=False) -> InferenceModel:
+    def train(self, speed=None, model_type=None, checkpoint=None, plot_in_notebook=False, epochs=None) -> InferenceModel:
         """
         Ask the Roboflow API to train a previously exported version's dataset.
 
         Args:
             speed: Whether to train quickly or accurately. Note: accurate training is a paid feature. Default speed is `fast`.
             model_type: The type of model to train. Default depends on kind of project. It takes precedence over speed. You can check the list of model ids by sending an invalid parameter in this argument.
             checkpoint: A string representing the checkpoint to use while training
+            epochs: Number of epochs to train the model
             plot: Whether to plot the training results. Default is `False`.
 
         Returns:
@@ -336,6 +337,7 @@ def train(self, speed=None, model_type=None, checkpoint=None, plot_in_notebook=F
             speed=payload_speed,
             checkpoint=payload_checkpoint,
             model_type=payload_model_type,
+            epochs=epochs,
         )
 
         status = "training"