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:
import lightly_train
if __name__ == "__main__":
lightly_train.train(
out="out/my_experiment",
data="my_data_dir",
model="torchvision/resnet50",
method="dino",
)
lightly-train train out="out/my_experiment" data="my_data_dir" model="torchvision/resnet50" method="dino"
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:
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,
method="dino",
)
List Models¶
The list_models command lists all available models. Only models from installed
packages are listed.
import lightly_train
print(lightly_train.list_models())
lightly-train list_models
Supported Libraries¶
The following libraries are supported:
Torchvision¶
Supported models
ResNet
torchvision/resnet18torchvision/resnet34torchvision/resnet50torchvision/resnet101torchvision/resnet152
ConvNext
torchvision/convnext_basetorchvision/convnext_largetorchvision/convnext_smalltorchvision/convnext_tiny
TIMM¶
Important
TIMM must be installed with
pip install lightly-train[timm].
Supported models
timm/<model name>(all models are supported, see timm docs for a full list)
Examples
timm/resnet50timm/convnext_basetimm/vit_base_patch16_224
Ultralytics¶
Important
Ultralytics must be installed with
pip install lightly-train[ultralytics].
Warning
Using Ultralytics models might require a commercial Ultralytics license. See the Ultralytics website for more information.
Note
For YOLOv12, we recommend using the original YOLOv12 package developed by the authors of YOLOv12, since the official Ultralytics implementations are not stable yet.
Models ending with .pt load pre-trained weights by Ultralytics. Models ending with
.yaml are not pre-trained.
Supported models
YOLOv5
ultralytics/yolov5l.yamlultralytics/yolov5l6u.ptultralytics/yolov5lu.ptultralytics/yolov5lu.yamlultralytics/yolov5m.yamlultralytics/yolov5m6u.ptultralytics/yolov5mu.ptultralytics/yolov5mu.yamlultralytics/yolov5n.yamlultralytics/yolov5n6u.ptultralytics/yolov5nu.ptultralytics/yolov5nu.yamlultralytics/yolov5s.yamlultralytics/yolov5s6u.ptultralytics/yolov5su.ptultralytics/yolov5su.yamlultralytics/yolov5x.yamlultralytics/yolov5x6u.ptultralytics/yolov5xu.ptultralytics/yolov5xu.yaml
YOLOv6
ultralytics/yolov6l.yamlultralytics/yolov6m.yamlultralytics/yolov6n.yamlultralytics/yolov6s.yamlultralytics/yolov6x.yaml
YOLOv8
ultralytics/yolov8l-cls.ptultralytics/yolov8l-cls.yamlultralytics/yolov8l-obb.ptultralytics/yolov8l-obb.yamlultralytics/yolov8l-oiv7.ptultralytics/yolov8l-pose.ptultralytics/yolov8l-pose.yamlultralytics/yolov8l-seg.ptultralytics/yolov8l-seg.yamlultralytics/yolov8l-world.ptultralytics/yolov8l-world.yamlultralytics/yolov8l-worldv2.ptultralytics/yolov8l-worldv2.yamlultralytics/yolov8l.ptultralytics/yolov8l.yamlultralytics/yolov8m-cls.ptultralytics/yolov8m-cls.yamlultralytics/yolov8m-obb.ptultralytics/yolov8m-obb.yamlultralytics/yolov8m-oiv7.ptultralytics/yolov8m-pose.ptultralytics/yolov8m-pose.yamlultralytics/yolov8m-seg.ptultralytics/yolov8m-seg.yamlultralytics/yolov8m-world.ptultralytics/yolov8m-world.yamlultralytics/yolov8m-worldv2.ptultralytics/yolov8m-worldv2.yamlultralytics/yolov8m.ptultralytics/yolov8m.yamlultralytics/yolov8n-cls.ptultralytics/yolov8n-cls.yamlultralytics/yolov8n-obb.ptultralytics/yolov8n-obb.yamlultralytics/yolov8n-oiv7.ptultralytics/yolov8n-pose.ptultralytics/yolov8n-pose.yamlultralytics/yolov8n-seg.ptultralytics/yolov8n-seg.yamlultralytics/yolov8n.ptultralytics/yolov8n.yamlultralytics/yolov8s-cls.ptultralytics/yolov8s-cls.yamlultralytics/yolov8s-obb.ptultralytics/yolov8s-obb.yamlultralytics/yolov8s-oiv7.ptultralytics/yolov8s-pose.ptultralytics/yolov8s-pose.yamlultralytics/yolov8s-seg.ptultralytics/yolov8s-seg.yamlultralytics/yolov8s-world.ptultralytics/yolov8s-world.yamlultralytics/yolov8s-worldv2.ptultralytics/yolov8s-worldv2.yamlultralytics/yolov8s.ptultralytics/yolov8s.yamlultralytics/yolov8x-cls.ptultralytics/yolov8x-cls.yamlultralytics/yolov8x-obb.ptultralytics/yolov8x-obb.yamlultralytics/yolov8x-oiv7.ptultralytics/yolov8x-pose.ptultralytics/yolov8x-pose.yamlultralytics/yolov8x-seg.ptultralytics/yolov8x-seg.yamlultralytics/yolov8x-world.ptultralytics/yolov8x-world.yamlultralytics/yolov8x-worldv2.ptultralytics/yolov8x-worldv2.yamlultralytics/yolov8x.ptultralytics/yolov8x.yaml
YOLO11
ultralytics/yolo11n-cls.yamlultralytics/yolo11n-cls.ptultralytics/yolo11n-obb.yamlultralytics/yolo11n-obb.ptultralytics/yolo11n-pose.yamlultralytics/yolo11n-pose.ptultralytics/yolo11n-seg.yamlultralytics/yolo11n-seg.ptultralytics/yolo11n.yamlultralytics/yolo11n.ptultralytics/yolo11s-cls.yamlultralytics/yolo11s-cls.ptultralytics/yolo11s-obb.yamlultralytics/yolo11s-obb.ptultralytics/yolo11s-pose.yamlultralytics/yolo11s-pose.ptultralytics/yolo11s-seg.yamlultralytics/yolo11s-seg.ptultralytics/yolo11s.yamlultralytics/yolo11s.ptultralytics/yolo11m-cls.yamlultralytics/yolo11m-cls.ptultralytics/yolo11m-obb.yamlultralytics/yolo11m-obb.ptultralytics/yolo11m-pose.yamlultralytics/yolo11m-pose.ptultralytics/yolo11m-seg.yamlultralytics/yolo11m-seg.ptultralytics/yolo11m.yamlultralytics/yolo11m.ptultralytics/yolo11l-cls.yamlultralytics/yolo11l-cls.ptultralytics/yolo11l-obb.yamlultralytics/yolo11l-obb.ptultralytics/yolo11l-pose.yamlultralytics/yolo11l-pose.ptultralytics/yolo11l-seg.yamlultralytics/yolo11l-seg.ptultralytics/yolo11l.yamlultralytics/yolo11l.ptultralytics/yolo11x-cls.yamlultralytics/yolo11x-cls.ptultralytics/yolo11x-obb.yamlultralytics/yolo11x-obb.ptultralytics/yolo11x-pose.yamlultralytics/yolo11x-pose.ptultralytics/yolo11x-seg.yamlultralytics/yolo11x-seg.ptultralytics/yolo11x.yamlultralytics/yolo11x.pt
YOLO12
ultralytics/yolo12n-cls.yamlultralytics/yolo12n-cls.ptultralytics/yolo12n-obb.yamlultralytics/yolo12n-obb.ptultralytics/yolo12n-pose.yamlultralytics/yolo12n-pose.ptultralytics/yolo12n-seg.yamlultralytics/yolo12n-seg.ptultralytics/yolo12n.yamlultralytics/yolo12n.ptultralytics/yolo12s-cls.yamlultralytics/yolo12s-cls.ptultralytics/yolo12s-obb.yamlultralytics/yolo12s-obb.ptultralytics/yolo12s-pose.yamlultralytics/yolo12s-pose.ptultralytics/yolo12s-seg.yamlultralytics/yolo12s-seg.ptultralytics/yolo12s.yamlultralytics/yolo12s.ptultralytics/yolo12m-cls.yamlultralytics/yolo12m-cls.ptultralytics/yolo12m-obb.yamlultralytics/yolo12m-obb.ptultralytics/yolo12m-pose.yamlultralytics/yolo12m-pose.ptultralytics/yolo12m-seg.yamlultralytics/yolo12m-seg.ptultralytics/yolo12m.yamlultralytics/yolo12m.ptultralytics/yolo12l-cls.yamlultralytics/yolo12l-cls.ptultralytics/yolo12l-obb.yamlultralytics/yolo12l-obb.ptultralytics/yolo12l-pose.yamlultralytics/yolo12l-pose.ptultralytics/yolo12l-seg.yamlultralytics/yolo12l-seg.ptultralytics/yolo12l.yamlultralytics/yolo12l.ptultralytics/yolo12x-cls.yamlultralytics/yolo12x-cls.ptultralytics/yolo12x-obb.yamlultralytics/yolo12x-obb.ptultralytics/yolo12x-pose.yamlultralytics/yolo12x-pose.ptultralytics/yolo12x-seg.yamlultralytics/yolo12x-seg.ptultralytics/yolo12x.yamlultralytics/yolo12x.pt
SuperGradients¶
Important
SuperGradients must be installed with
pip install lightly-train[super-gradients].
Warning
SuperGradients support is still experimental. There might be unexpected warnings in the logs.
Supported models
PP-LiteSeg
super_gradients/pp_lite_b_segsuper_gradients/pp_lite_b_seg50super_gradients/pp_lite_b_seg75super_gradients/pp_lite_t_segsuper_gradients/pp_lite_t_seg50super_gradients/pp_lite_t_seg75
SSD
super_gradients/ssd_lite_mobilenet_v2super_gradients/ssd_mobilenet_v1
YOLO-NAS
super_gradients/yolo_nas_lsuper_gradients/yolo_nas_msuper_gradients/yolo_nas_pose_lsuper_gradients/yolo_nas_pose_msuper_gradients/yolo_nas_pose_nsuper_gradients/yolo_nas_pose_ssuper_gradients/yolo_nas_s
Custom Models¶
See Custom Models for information on how to train custom models.