Source code for composer.algorithms.no_op_model.no_op_model

# Copyright 2022 MosaicML Composer authors
# SPDX-License-Identifier: Apache-2.0

"""NoOpModel algorithm and class."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, Dict, Optional

import torch
import torch.nn.functional as F
from torchmetrics import Metric
from torchmetrics.classification import BinaryAccuracy

from composer.core import Algorithm, Event, State
from composer.loggers import Logger
from composer.models.base import ComposerModel
from composer.utils import module_surgery

    from composer.core.types import Batch

log = logging.getLogger(__name__)

__all__ = ['NoOpModelClass', 'NoOpModel']

class NoOpModelClass(ComposerModel):
    """Dummy model used for performance measurements.

    The :class:`.NoOpModel` algorithm uses this class to replace a :class:`torch.nn.Module`.

        original_model (torch.nn.Module): Model to replace.

    def __init__(self, original_model: torch.nn.Module):
        original_device = next(original_model.parameters()).device
        self.weights = torch.nn.Parameter(torch.Tensor([1.5]).to(original_device))
            # For classification
            self.num_classes = original_model.num_classes
        except AttributeError:

    def loss(self, outputs: torch.Tensor, batch: Batch):
        x, y = batch
        assert isinstance(y, torch.Tensor)
        del x  # unused
        return F.mse_loss(outputs,

    def forward(self, batch: Batch):
        x, y = batch
        del x  # unused
        assert isinstance(y, torch.Tensor)
        return y * self.weights

    def get_metrics(self, is_train: bool) -> Dict[str, Metric]:
        return {'BinaryAccuracy': BinaryAccuracy()}

    def eval_forward(self, batch: Batch, outputs: Optional[Any] = None):
        x, y = batch
        del x  # unused
        return y

    def update_metric(self, batch: Any, outputs: Any, metric: Metric) -> None:
        _, targets = batch
        metric.update(outputs, targets)

[docs]class NoOpModel(Algorithm): """Runs on :attr:`Event.INIT` and replaces the model with a dummy :class:`.NoOpModelClass` instance.""" def __init__(self): # No arguments pass def __repr__(self) -> str: return f'{self.__class__.__name__}()' @staticmethod def required_on_load() -> bool: return True def match(self, event: Event, state: State) -> bool: return event == Event.INIT def apply(self, event: Event, state: State, logger: Logger) -> Optional[int]: new_model = NoOpModelClass(state.model) module_surgery.update_params_in_optimizer(old_params=state.model.parameters(), new_params=new_model.parameters(), optimizers=state.optimizers) state.model = new_model'Replaced model with a NoOpModel')