Skip to content

Classification training worker

User Configuration parameters should include:

  • class_names: List of class_names to predict, optional. If not provided, we will process all class_names encountered
  • image_size: images will be resized to that size, defaults to 640 (YOLO default)
  • worker_runs: List of strings, filter sources of the classification
  • model_id: ID of the model that will hold the new model version
  • training_kwargs: Additional kwargs passed to model.train, dict
  • num_epochs: int, default to 500

First you need to build the training dataset. The format is described in the doc. Code to do that (using the official Arkindex Export SQLite structure is available in our lib).

You will start from a Cache database generated by https://gitlab.teklia.com/workers/generic-training-dataset. Images are already downloaded, we simply need to move them to the right path. First parse the database and find out where each image should go. Then do the moves.

Up to this point, all is done in process_dataset. The rest is done in run.

Once the database has been generated we can start training.

from ultralytics import YOLO

model = YOLO('yolov8x-cls.pt')  # load a pretrained model (recommended for training)

# Build kwargs from configuration
# Override with self.config["training_kwargs"]

results = model.train(data=path/to/dataset, **kwargs)

You can expose some parameters from the documentation:

  • epochs as num_epochs
  • dropout

Hardcoded parameters:

  • batch = -1 for AutoBatch
  • device = 0 if torch.cuda.is_available() else "cpu"
  • verbose = true
  • plots = true

After training we should evaluate on all splits dataset.

for split in ("train", "val", "test"):
    model.val(split=split)  # no arguments needed, dataset and settings remembered

Then we can publish the model version using the publish_model_version helper.

All these commands generate a great lot of files. Basically we should have a local run/classify/train/weights folder with the model's best.pt checkpoint (should be renamed to model.pt before publication). The validation step should create 3 folders (one for each call) run/classify/val{|2|3}. In each there should be two confusion matrices. We should make a results.zip with train/val/test folders with confusion matrices. Also include run/classify/train/results.png (the recap of the training). This ZIP archive should be exposed as task artifact.

Edited by Yoann Schneider