Train module

Model training with YOLOv8.

• Mlflow: Model tracking. Can access in localhost:5000.

• boto3: Connect to S3 bucket

User can easily modify hiperparameters used for training. Logs are saved in logs/model_train.log. Model is saved in .onnx format. This format is lighter and for predictions, it is not necessary to install pytorch and ultralytics packages. Very useful to make the ECR Docker lighter. The only difference is that it is necessary to do some preprocessing, that is described in the lambda function.

train.main(hp: dict)

Train a YOLO model using specified hyperparameters and save the trained model to an S3 bucket.

This function initializes a YOLO model, trains it on the Stanford car dataset using the provided hyperparameters, and then exports the best model in ONNX format. The trained model is uploaded to an S3 bucket for storage.

Set mlflow to save tracking inside folder model/runs/mlflow. Logs are saved in logs/model_train.log.

Parameters:

hpdict

A dictionary containing hyperparameters for training, including:

• ‘epochs’ (int): Number of training epochs.

• ‘imgsz’ (int): Image size for training.

• ‘scale’ (float): Scaling factor for augmentation.

• ‘batch’ (int): Batch size for training.

• ‘optimizer’ (str): Optimizer to use for training.

Returns:

bool

True if training was successful and the model was uploaded to the S3 bucket, False otherwise.

train.train_with_YOLO(hp: dict)