Export - LightlyTrain documentation (original) (raw)

View this page

Toggle table of contents sidebar

The export command is used to prepare a model for fine-tuning or inference. It allows exporting the model from training checkpoints which contain additional information such as optimizer states that are not needed for fine-tuning or inference.

Tip

After training the model is automatically exported in the default format of the used library to out/my_experiment/exported_models/exported_last.pt.

Python

import lightly_train

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

lightly_train.export(
    out="my_exported_model.pt",
    checkpoint="out/my_experiment/checkpoints/last.ckpt",
    part="model",
    format="torch_state_dict",
)

Command Line

lightly-train train out="out/my_experiment" data="my_data_dir" model="torchvision/resnet50" lightly-train export out="my_exported_model.pt" checkpoint="out/my_experiment/checkpoints/last.ckpt" part="model" format="torch_state_dict"

The above code example trains a model and exports the last training checkpoint as a torch state dictionary.

Out¶

The out argument specifies the output file where the exported model is saved.

Checkpoint¶

The checkpoint argument specifies the LightlyTrain checkpoint to use for exporting the model. This is the checkpoint saved to out/my_experiment/checkpoints/<some>.ckpt after training.

Format¶

The optional format argument specifies the format in which the model is exported. The following formats are supported.

• package_default (default)
  This format option can be used to automatically determine the required export format. It will fall back to torch_state_dict format if the model cannot be associated with a supported library.
  The model is saved in the native format of the used package:

  Model	Format
  custom	state dict
  super_gradients	super_gradients
  timm	timm
  torchvision	state dict
  ultralytics	ultralytics
  Usage examples:
  import torch
  from torchvision.models import resnet50
  model = resnet50()
  model.load_state_dict(torch.load("my_exported_model.pt"))
  from ultralytics import YOLO
  model = YOLO("my_exported_model.pt")
  from super_gradients.training import models
  model = models.get(
  model_name="yolo_nas_s",
  num_classes=3,
  checkpoint_path="my_exported_model.pt",
  )
  import timm
  model = timm.create_model(
  "resnet18", pretrained=False, checkpoint_path="my_exported_model.pt",
  )

• torch_state_dict
  Only the model’s state dict is saved which can be loaded with:
  import torch
  from torchvision.models import resnet50
  model = resnet50()
  model.load_state_dict(torch.load("my_exported_model.pt"))
• torch_model
  The model is saved as a torch module which can be loaded with:
  import torch
  model = torch.load("my_exported_model.pt")
  This requires that the same LightlyTrain version is installed when the model is exported and when it is loaded again.

Part¶

The optional part argument specifies which part of the model to export. The following parts are supported.

• model (default)
  Exports the model as passed with the model argument in the train function.
• embedding_model
  Exports the embedding model. This includes the model passed with the model argument in the train function and an extra embedding layer if the embed_dim argument was set during training. This is useful if you want to use the model for embedding images.