classifier_model

torchtune.modules.classifier_model(num_classes: int, base_model_path: str, **base_model_kwargs: Dict[str, Any]) → Union[TransformerDecoder, Module]

Create a classifier model from a base model by adapting the output layer.

Note

This builder does not support models which apply PEFT to the output layer.

Parameters:

• num_classes (int) – The number of classes for the classifier.

• base_model_path (str) – The path to the base model builder, which must return an instance of TransformerDecoder, or a model with a decoder attribute that is an instance of TransformerDecoder.

• **base_model_kwargs (Dict[str, Any]) – Keyword arguments for the base model.

Returns:

The base model, with the output layer adapted for the number of classes.

Return type:

Union[TransformerDecoder, nn.Module]

Raises:

ValueError – If the base model does not have a valid output layer to adapt.

Example

>>> from torchtune.modules import classifier_model
>>> model = classifier_model(num_classes=1, base_model_path="torchtune.models.llama3_2.llama3_2_1b")
>>> model.output.weight.shape
torch.Size([1, 4096])