Skip to content

Ikomia-hub/train_timm_image_classification

Repository files navigation

Algorithm icon

train_timm_image_classification


Stars Website GitHub
Discord community

Train timm image classification models.

Rock paper scissors

🚀 Use with Ikomia API

1. Install Ikomia API

We strongly recommend using a virtual environment. If you're not sure where to start, we offer a tutorial here.

pip install ikomia

2. Create your workflow

from ikomia.dataprocess.workflow import Workflow

# Init your workflow
wf = Workflow()    

# Add dataset loader
data_loader = wf.add_task(name="dataset_classification")

data_loader.set_parameters({
    "dataset_folder": "path/to/dataset/folder",
}) 

train = wf.add_task(name="train_timm_image_classification", auto_connect=True)

# Launch your training on your data
wf.run()

☀️ Use with Ikomia Studio

Ikomia Studio offers a friendly UI with the same features as the API.

  • If you haven't started using Ikomia Studio yet, download and install it from this page.

  • For additional guidance on getting started with Ikomia Studio, check out this blog post.

📝 Set algorithm parameters

  • model_name (str) - default 'resnet18': Name of the pre-trained model.
    • There are over 700 timm models. You can list them using: timm.list_models()
  • input_size (list) - default '[224, 224]': Size of the input image.
  • epochs (int) - default '100': Number of complete passes through the training dataset.
  • batch_size (int) - default '16': Number of samples processed before the model is updated.
  • learning_rate (float) - default '0.0050': Step size at which the model's parameters are updated during training.
  • output_folder (str, optional): path to where the model will be saved.

Parameters should be in strings format when added to the dictionary.

from ikomia.dataprocess.workflow import Workflow

# Init your workflow
wf = Workflow()    

# Add dataset loader
data_loader = wf.add_task(name="dataset_classification")

data_loader.set_parameters({
    "dataset_folder": "C:/Users/allan/OneDrive/Desktop/ik-desktop/Images/datasets/Fruit",
}) 

# Add train algorithm 
train = wf.add_task(name="train_timm_image_classification", auto_connect=True)
train.set_parameters({
    "model_name": "resnet34",
    "batch_size": "8",
    "epochs": "5",
    "learning_rate": "0.0050",
}) 
# Launch your training on your data
wf.run()