chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,338 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING, TypeAlias
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle import nn
|
||||
from paddle.jit.dy2static.program_translator import unwrap_decorators
|
||||
|
||||
from .static_flops import Table, static_flops
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from paddle import Tensor
|
||||
from paddle.nn import Layer
|
||||
from paddle.static import Program
|
||||
|
||||
_CustomOpsAlias: TypeAlias = dict[type[Layer], Callable[..., None]]
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
def flops(
|
||||
net: Layer | Program,
|
||||
input_size: list[int],
|
||||
custom_ops: _CustomOpsAlias | None = None,
|
||||
print_detail: bool = False,
|
||||
) -> int:
|
||||
"""Print a table about the FLOPs of network.
|
||||
|
||||
Args:
|
||||
net (paddle.nn.Layer||paddle.static.Program): The network which could be a instance of paddle.nn.Layer in
|
||||
dygraph or paddle.static.Program in static graph.
|
||||
input_size (list): size of input tensor. Note that the batch_size in argument ``input_size`` only support 1.
|
||||
custom_ops (A dict of function, optional): A dictionary which key is the class of specific operation such as
|
||||
paddle.nn.Conv2D and the value is the function used to count the FLOPs of this operation. This
|
||||
argument only work when argument ``net`` is an instance of paddle.nn.Layer. The details could be found
|
||||
in following example code. Default is None.
|
||||
print_detail (bool, optional): Whether to print the detail information, like FLOPs per layer, about the net FLOPs.
|
||||
Default is False.
|
||||
|
||||
Returns:
|
||||
Int: A number about the FLOPs of total network.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> import paddle.nn as nn
|
||||
|
||||
>>> class LeNet(nn.Layer):
|
||||
... def __init__(self, num_classes=10):
|
||||
... super().__init__()
|
||||
... self.num_classes = num_classes
|
||||
... self.features = nn.Sequential(
|
||||
... nn.Conv2D(1, 6, 3, stride=1, padding=1),
|
||||
... nn.ReLU(),
|
||||
... nn.MaxPool2D(2, 2),
|
||||
... nn.Conv2D(6, 16, 5, stride=1, padding=0),
|
||||
... nn.ReLU(),
|
||||
... nn.MaxPool2D(2, 2),
|
||||
... )
|
||||
...
|
||||
... if num_classes > 0:
|
||||
... self.fc = nn.Sequential(
|
||||
... nn.Linear(400, 120),
|
||||
... nn.Linear(120, 84),
|
||||
... nn.Linear(84, 10),
|
||||
... )
|
||||
...
|
||||
... def forward(self, inputs):
|
||||
... x = self.features(inputs)
|
||||
...
|
||||
... if self.num_classes > 0:
|
||||
... x = paddle.flatten(x, 1)
|
||||
... x = self.fc(x)
|
||||
... return x
|
||||
>>> lenet = LeNet()
|
||||
>>> # m is the instance of nn.Layer, x is the input of layer, y is the output of layer.
|
||||
>>> def count_leaky_relu(m, x, y):
|
||||
... x = x[0]
|
||||
... nelements = x.numel()
|
||||
... m.total_ops += int(nelements)
|
||||
>>> FLOPs = paddle.flops(
|
||||
... lenet,
|
||||
... [1, 1, 28, 28],
|
||||
... custom_ops={nn.LeakyReLU: count_leaky_relu},
|
||||
... print_detail=True,
|
||||
... )
|
||||
>>> # doctest: +SKIP('numpy print with different version')
|
||||
<class 'paddle.nn.layer.conv.Conv2D'>'s flops has been counted
|
||||
<class 'paddle.nn.layer.activation.ReLU'>'s flops has been counted
|
||||
Cannot find suitable count function for <class 'paddle.nn.layer.pooling.MaxPool2D'>. Treat it as zero FLOPs.
|
||||
<class 'paddle.nn.layer.common.Linear'>'s flops has been counted
|
||||
+--------------+-----------------+-----------------+--------+--------+
|
||||
| Layer Name | Input Shape | Output Shape | Params | Flops |
|
||||
+--------------+-----------------+-----------------+--------+--------+
|
||||
| conv2d_0 | [1, 1, 28, 28] | [1, 6, 28, 28] | 60 | 47040 |
|
||||
| re_lu_0 | [1, 6, 28, 28] | [1, 6, 28, 28] | 0 | 0 |
|
||||
| max_pool2d_0 | [1, 6, 28, 28] | [1, 6, 14, 14] | 0 | 0 |
|
||||
| conv2d_1 | [1, 6, 14, 14] | [1, 16, 10, 10] | 2416 | 241600 |
|
||||
| re_lu_1 | [1, 16, 10, 10] | [1, 16, 10, 10] | 0 | 0 |
|
||||
| max_pool2d_1 | [1, 16, 10, 10] | [1, 16, 5, 5] | 0 | 0 |
|
||||
| linear_0 | [1, 400] | [1, 120] | 48120 | 48000 |
|
||||
| linear_1 | [1, 120] | [1, 84] | 10164 | 10080 |
|
||||
| linear_2 | [1, 84] | [1, 10] | 850 | 840 |
|
||||
+--------------+-----------------+-----------------+--------+--------+
|
||||
Total Flops: 347560 Total Params: 61610
|
||||
>>> # doctest: -SKIP
|
||||
>>> print(FLOPs)
|
||||
347560
|
||||
"""
|
||||
if isinstance(net, nn.Layer):
|
||||
# If net is a dy2stat model, net.forward is StaticFunction instance,
|
||||
# we set net.forward to original forward function.
|
||||
_, net.forward = unwrap_decorators(net.forward)
|
||||
|
||||
inputs = paddle.randn(input_size)
|
||||
return dynamic_flops(
|
||||
net, inputs=inputs, custom_ops=custom_ops, print_detail=print_detail
|
||||
)
|
||||
elif isinstance(net, paddle.static.Program):
|
||||
return static_flops(net, print_detail=print_detail)
|
||||
else:
|
||||
warnings.warn(
|
||||
"Your model must be an instance of paddle.nn.Layer or paddle.static.Program."
|
||||
)
|
||||
return -1
|
||||
|
||||
|
||||
def count_convNd(m, x, y):
|
||||
x = x[0]
|
||||
kernel_ops = np.prod(m.weight.shape[2:])
|
||||
bias_ops = 1 if m.bias is not None else 0
|
||||
total_ops = int(y.numel()) * (
|
||||
x.shape[1] / m._groups * kernel_ops + bias_ops
|
||||
)
|
||||
m.total_ops += abs(int(total_ops))
|
||||
|
||||
|
||||
def count_leaky_relu(m, x, y):
|
||||
x = x[0]
|
||||
nelements = x.numel()
|
||||
m.total_ops += int(nelements)
|
||||
|
||||
|
||||
def count_bn(m, x, y):
|
||||
x = x[0]
|
||||
nelements = x.numel()
|
||||
if not m.training:
|
||||
total_ops = 2 * nelements
|
||||
m.total_ops += abs(int(total_ops))
|
||||
|
||||
|
||||
def count_linear(m, x, y):
|
||||
total_mul = m.weight.shape[0]
|
||||
num_elements = y.numel()
|
||||
total_ops = total_mul * num_elements
|
||||
m.total_ops += abs(int(total_ops))
|
||||
|
||||
|
||||
def count_avgpool(m, x, y):
|
||||
kernel_ops = 1
|
||||
num_elements = y.numel()
|
||||
total_ops = kernel_ops * num_elements
|
||||
|
||||
m.total_ops += int(total_ops)
|
||||
|
||||
|
||||
def count_adap_avgpool(m, x, y):
|
||||
kernel = np.array(x[0].shape[2:]) // np.array(y.shape[2:])
|
||||
total_add = np.prod(kernel)
|
||||
total_div = 1
|
||||
kernel_ops = total_add + total_div
|
||||
num_elements = y.numel()
|
||||
total_ops = kernel_ops * num_elements
|
||||
m.total_ops += abs(int(total_ops))
|
||||
|
||||
|
||||
def count_zero_ops(m, x, y):
|
||||
m.total_ops += 0
|
||||
|
||||
|
||||
def count_parameters(m, x, y):
|
||||
total_params = 0
|
||||
for p in m.parameters():
|
||||
total_params += p.numel()
|
||||
m.total_params[0] = abs(int(total_params))
|
||||
|
||||
|
||||
def count_io_info(m, x, y):
|
||||
m.register_buffer('input_shape', paddle.to_tensor(x[0].shape))
|
||||
if isinstance(y, (list, tuple)):
|
||||
m.register_buffer('output_shape', paddle.to_tensor(y[0].shape))
|
||||
else:
|
||||
m.register_buffer('output_shape', paddle.to_tensor(y.shape))
|
||||
|
||||
|
||||
register_hooks = {
|
||||
nn.Conv1D: count_convNd,
|
||||
nn.Conv2D: count_convNd,
|
||||
nn.Conv3D: count_convNd,
|
||||
nn.Conv1DTranspose: count_convNd,
|
||||
nn.Conv2DTranspose: count_convNd,
|
||||
nn.Conv3DTranspose: count_convNd,
|
||||
nn.layer.norm.BatchNorm2D: count_bn,
|
||||
nn.BatchNorm: count_bn,
|
||||
nn.ReLU: count_zero_ops,
|
||||
nn.ReLU6: count_zero_ops,
|
||||
nn.LeakyReLU: count_leaky_relu,
|
||||
nn.Linear: count_linear,
|
||||
nn.Dropout: count_zero_ops,
|
||||
nn.AvgPool1D: count_avgpool,
|
||||
nn.AvgPool2D: count_avgpool,
|
||||
nn.AvgPool3D: count_avgpool,
|
||||
nn.AdaptiveAvgPool1D: count_adap_avgpool,
|
||||
nn.AdaptiveAvgPool2D: count_adap_avgpool,
|
||||
nn.AdaptiveAvgPool3D: count_adap_avgpool,
|
||||
}
|
||||
|
||||
|
||||
def dynamic_flops(
|
||||
model: Layer,
|
||||
inputs: Tensor,
|
||||
custom_ops: _CustomOpsAlias | None = None,
|
||||
print_detail: bool = False,
|
||||
) -> int:
|
||||
handler_collection = []
|
||||
types_collection = set()
|
||||
if custom_ops is None:
|
||||
custom_ops = {}
|
||||
|
||||
def add_hooks(m):
|
||||
if len(list(m.children())) > 0:
|
||||
return
|
||||
m.register_buffer('total_ops', paddle.zeros([1], dtype='int64'))
|
||||
m.register_buffer('total_params', paddle.zeros([1], dtype='int64'))
|
||||
m_type = type(m)
|
||||
|
||||
flops_fn = None
|
||||
if m_type in custom_ops:
|
||||
flops_fn = custom_ops[m_type]
|
||||
if m_type not in types_collection:
|
||||
print(f"Customize Function has been applied to {m_type}")
|
||||
elif m_type in register_hooks:
|
||||
flops_fn = register_hooks[m_type]
|
||||
if m_type not in types_collection:
|
||||
print(f"{m_type}'s flops has been counted")
|
||||
else:
|
||||
if m_type not in types_collection:
|
||||
print(
|
||||
f"Cannot find suitable count function for {m_type}. Treat it as zero FLOPs."
|
||||
)
|
||||
|
||||
if flops_fn is not None:
|
||||
flops_handler = m.register_forward_post_hook(flops_fn)
|
||||
handler_collection.append(flops_handler)
|
||||
params_handler = m.register_forward_post_hook(count_parameters)
|
||||
io_handler = m.register_forward_post_hook(count_io_info)
|
||||
handler_collection.append(params_handler)
|
||||
handler_collection.append(io_handler)
|
||||
types_collection.add(m_type)
|
||||
|
||||
training = model.training
|
||||
|
||||
model.eval()
|
||||
model.apply(add_hooks)
|
||||
|
||||
with paddle.framework.no_grad():
|
||||
model(inputs)
|
||||
|
||||
total_ops = 0
|
||||
total_params = 0
|
||||
for m in model.sublayers():
|
||||
if len(list(m.children())) > 0:
|
||||
continue
|
||||
if {
|
||||
'total_ops',
|
||||
'total_params',
|
||||
'input_shape',
|
||||
'output_shape',
|
||||
}.issubset(set(m._buffers.keys())):
|
||||
total_ops += m.total_ops
|
||||
total_params += m.total_params
|
||||
|
||||
if training:
|
||||
model.train()
|
||||
for handler in handler_collection:
|
||||
handler.remove()
|
||||
|
||||
table = Table(
|
||||
["Layer Name", "Input Shape", "Output Shape", "Params", "Flops"]
|
||||
)
|
||||
|
||||
for n, m in model.named_sublayers():
|
||||
if len(list(m.children())) > 0:
|
||||
continue
|
||||
if {
|
||||
'total_ops',
|
||||
'total_params',
|
||||
'input_shape',
|
||||
'output_shape',
|
||||
}.issubset(set(m._buffers.keys())):
|
||||
table.add_row(
|
||||
[
|
||||
m.full_name(),
|
||||
m.input_shape.numpy().tolist(),
|
||||
m.output_shape.numpy().tolist(),
|
||||
int(m.total_params),
|
||||
int(m.total_ops),
|
||||
]
|
||||
)
|
||||
m._buffers.pop("total_ops")
|
||||
m._buffers.pop("total_params")
|
||||
m._buffers.pop('input_shape')
|
||||
m._buffers.pop('output_shape')
|
||||
if print_detail:
|
||||
table.print_table()
|
||||
print(
|
||||
f'Total Flops: {int(total_ops)} Total Params: {int(total_params)}'
|
||||
)
|
||||
return int(total_ops)
|
||||
Reference in New Issue
Block a user