jawahar-konathala/Tryon2

1# Copyright (c) Facebook, Inc. and its affiliates.2# -*- coding: utf-8 -*-3 4import typing5from typing import Any, List6import fvcore7from fvcore.nn import activation_count, flop_count, parameter_count, parameter_count_table8from torch import nn9 10from detectron2.export import TracingAdapter11 12__all__ = [13    "activation_count_operators",14    "flop_count_operators",15    "parameter_count_table",16    "parameter_count",17    "FlopCountAnalysis",18]19 20FLOPS_MODE = "flops"21ACTIVATIONS_MODE = "activations"22 23 24# Some extra ops to ignore from counting, including elementwise and reduction ops25_IGNORED_OPS = {26    "aten::add",27    "aten::add_",28    "aten::argmax",29    "aten::argsort",30    "aten::batch_norm",31    "aten::constant_pad_nd",32    "aten::div",33    "aten::div_",34    "aten::exp",35    "aten::log2",36    "aten::max_pool2d",37    "aten::meshgrid",38    "aten::mul",39    "aten::mul_",40    "aten::neg",41    "aten::nonzero_numpy",42    "aten::reciprocal",43    "aten::repeat_interleave",44    "aten::rsub",45    "aten::sigmoid",46    "aten::sigmoid_",47    "aten::softmax",48    "aten::sort",49    "aten::sqrt",50    "aten::sub",51    "torchvision::nms",  # TODO estimate flop for nms52}53 54 55class FlopCountAnalysis(fvcore.nn.FlopCountAnalysis):56    """57    Same as :class:`fvcore.nn.FlopCountAnalysis`, but supports detectron2 models.58    """59 60    def __init__(self, model, inputs):61        """62        Args:63            model (nn.Module):64            inputs (Any): inputs of the given model. Does not have to be tuple of tensors.65        """66        wrapper = TracingAdapter(model, inputs, allow_non_tensor=True)67        super().__init__(wrapper, wrapper.flattened_inputs)68        self.set_op_handle(**{k: None for k in _IGNORED_OPS})69 70 71def flop_count_operators(model: nn.Module, inputs: list) -> typing.DefaultDict[str, float]:72    """73    Implement operator-level flops counting using jit.74    This is a wrapper of :func:`fvcore.nn.flop_count` and adds supports for standard75    detection models in detectron2.76    Please use :class:`FlopCountAnalysis` for more advanced functionalities.77 78    Note:79        The function runs the input through the model to compute flops.80        The flops of a detection model is often input-dependent, for example,81        the flops of box & mask head depends on the number of proposals &82        the number of detected objects.83        Therefore, the flops counting using a single input may not accurately84        reflect the computation cost of a model. It's recommended to average85        across a number of inputs.86 87    Args:88        model: a detectron2 model that takes `list[dict]` as input.89        inputs (list[dict]): inputs to model, in detectron2's standard format.90            Only "image" key will be used.91        supported_ops (dict[str, Handle]): see documentation of :func:`fvcore.nn.flop_count`92 93    Returns:94        Counter: Gflop count per operator95    """96    old_train = model.training97    model.eval()98    ret = FlopCountAnalysis(model, inputs).by_operator()99    model.train(old_train)100    return {k: v / 1e9 for k, v in ret.items()}101 102 103def activation_count_operators(104    model: nn.Module, inputs: list, **kwargs105) -> typing.DefaultDict[str, float]:106    """107    Implement operator-level activations counting using jit.108    This is a wrapper of fvcore.nn.activation_count, that supports standard detection models109    in detectron2.110 111    Note:112        The function runs the input through the model to compute activations.113        The activations of a detection model is often input-dependent, for example,114        the activations of box & mask head depends on the number of proposals &115        the number of detected objects.116 117    Args:118        model: a detectron2 model that takes `list[dict]` as input.119        inputs (list[dict]): inputs to model, in detectron2's standard format.120            Only "image" key will be used.121 122    Returns:123        Counter: activation count per operator124    """125    return _wrapper_count_operators(model=model, inputs=inputs, mode=ACTIVATIONS_MODE, **kwargs)126 127 128def _wrapper_count_operators(129    model: nn.Module, inputs: list, mode: str, **kwargs130) -> typing.DefaultDict[str, float]:131    # ignore some ops132    supported_ops = {k: lambda *args, **kwargs: {} for k in _IGNORED_OPS}133    supported_ops.update(kwargs.pop("supported_ops", {}))134    kwargs["supported_ops"] = supported_ops135 136    assert len(inputs) == 1, "Please use batch size=1"137    tensor_input = inputs[0]["image"]138    inputs = [{"image": tensor_input}]  # remove other keys, in case there are any139 140    old_train = model.training141    if isinstance(model, (nn.parallel.distributed.DistributedDataParallel, nn.DataParallel)):142        model = model.module143    wrapper = TracingAdapter(model, inputs)144    wrapper.eval()145    if mode == FLOPS_MODE:146        ret = flop_count(wrapper, (tensor_input,), **kwargs)147    elif mode == ACTIVATIONS_MODE:148        ret = activation_count(wrapper, (tensor_input,), **kwargs)149    else:150        raise NotImplementedError("Count for mode {} is not supported yet.".format(mode))151    # compatible with change in fvcore152    if isinstance(ret, tuple):153        ret = ret[0]154    model.train(old_train)155    return ret156 157 158def find_unused_parameters(model: nn.Module, inputs: Any) -> List[str]:159    """160    Given a model, find parameters that do not contribute161    to the loss.162 163    Args:164        model: a model in training mode that returns losses165        inputs: argument or a tuple of arguments. Inputs of the model166 167    Returns:168        list[str]: the name of unused parameters169    """170    assert model.training171    for _, prm in model.named_parameters():172        prm.grad = None173 174    if isinstance(inputs, tuple):175        losses = model(*inputs)176    else:177        losses = model(inputs)178 179    if isinstance(losses, dict):180        losses = sum(losses.values())181    losses.backward()182 183    unused: List[str] = []184    for name, prm in model.named_parameters():185        if prm.grad is None:186            unused.append(name)187        prm.grad = None188    return unused189