To enable the export stage with a ViT model, configure the configuration files:
In the
defaults
section ofconf/config.yaml
, update theexport
field to point to the desired ViT configuration file. For example, if you want to use thevit/export_vit
configuration, change theexport
field tovit/export_vit
.defaults: - export: vit/export_vit ...
In the
stages
field ofconf/config.yaml
, make sure theexport
stage is included. For example,stages: - export ...
Configure
infer.max_batch_size
of theconf/export/vit/export_vit.yaml
file to set the max_batch_size to use for the ONNX and NVIDIA TensorRT model.Set the resolution of the model with
max_dim
in theinfer
field. This will be used to generate the ONNX and NVIDIA TensorRT formats.
Remarks:
To load a pretrained checkpoint for inference, set the
restore_from_path
field in themodel
section to the path of the pretrained checkpoint in.nemo
format inconf/export/vit/export_vit.yaml
. By default, this field links to the.nemo
format checkpoint located in the ImageNet 1K fine-tuning checkpoints folder.