composer.models

Models provided to Trainer must use the basic interface specified by BaseMosaicModel.

Additionally, for convience we provide a number of extensions of BaseMosaicModel as detailed below.

Base Models

BaseMosaicModel

The minimal interface needed to use a model with composer.trainer.Trainer.

MosaicClassifier

Implements the base logic that all classifiers can build on top of.

MosaicTransformer

Implements the base logic that all Transformers can build on top of.

Image Models

CIFAR10_ResNet56

A ResNet-56 model extending MosaicClassifier.

MNIST_Classifier

A simple convolutional neural network.

EfficientNetB0

An EfficientNet-b0 model extending MosaicClassifier.

ResNet18

A ResNet-18 model extending MosaicClassifier.

ResNet50

A ResNet-50 model extending MosaicClassifier.

ResNet101

A ResNet-101 model extending MosaicClassifier.

UNet

A U-Net model extending MosaicClassifier.

Language Models

GPT2Model

Implements a GPT-2 wrapper around a MosaicTransformer.

Metrics and Loss Functions

Evaluation metrics for common tasks are in torchmetrics and are directly compatible with BaseMosaicModel. Additionally, we provide implementations of the following metrics and loss functions.

Dice

The Dice Coefficient for evaluating image segmentation.

CrossEntropyLoss

Torchmetric cross entropy loss implementation.

soft_cross_entropy

Drop-in replacement for torch.CrossEntropy that can handle dense labels.

LanguageCrossEntropyLoss

Hugging Face compatible cross entropy loss.

Perplexity

Subclasses LanguageCrossEntropyLoss to implement perplexity.