Models¶

LightlyTrain supports training models from various libraries. See Supported Libraries for a list of supported libraries and models.

The model is specified in the train command with the model argument:

Python

import lightly_train

if __name__ == "__main__":
    lightly_train.train(
        out="out/my_experiment",
        data="my_data_dir",
        model="torchvision/resnet50",
    )

Command Line

lightly-train train out="out/my_experiment" data="my_data_dir" model="torchvision/resnet50"

Model names always follow the pattern <library name>/<model name>.

Instead of passing a model name, it is also possible to pass a model instance directly to the train function:

Python

import lightly_train
from torchvision.models import resnet50

if __name__ == "__main__":
    model = resnet50()

    lightly_train.train(
        out="out/my_experiment",
        data="my_data_dir",
        model=model,
    )

List Models¶

The list_models command lists all available models. Only models from installed packages are listed.

Python

import lightly_train

print(lightly_train.list_models())

Command Line

lightly-train list_models

Supported Libraries¶

The following libraries are supported (follow the links to get to the respective docs pages):

Custom Models¶

See Custom Models for information on how to train custom models.