composer.algorithms.channels_last.channels_last

ChannelsLast algorithm.

Functions

apply_channels_last

Changes the memory format of the model to torch.channels_last.

Classes

ChannelsLast

Changes the memory format of the model to torch.channels_last.

class composer.algorithms.channels_last.channels_last.ChannelsLast

Bases: composer.core.algorithm.Algorithm

Changes the memory format of the model to torch.channels_last. This usually improves GPU utilization.

Runs on Event.INIT, so it can set the memory format before the model is DDP wrapped. Has no hyperparameters.

Example

from composer.algorithms import ChannelsLast
algorithm = ChannelsLast()
trainer = Trainer(
    model=model,
    train_dataloader=train_dataloader,
    eval_dataloader=eval_dataloader,
    max_duration="1ep",
    algorithms=[algorithm],
    optimizers=[optimizer]
)

composer.algorithms.channels_last.channels_last.apply_channels_last(model)

Changes the memory format of the model to torch.channels_last.

This usually yields improved GPU utilization.

Parameters

model (torch.nn.Module) – model or module to modify