chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,664 @@
|
||||
# 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 numbers
|
||||
import warnings
|
||||
from collections import OrderedDict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
import paddle
|
||||
from paddle import Tensor, nn
|
||||
from paddle.autograd import no_grad
|
||||
from paddle.static import InputSpec
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class ModelSummary(TypedDict):
|
||||
total_params: int
|
||||
trainable_params: int
|
||||
|
||||
|
||||
def summary(
|
||||
net: nn.Layer,
|
||||
input_size: (
|
||||
int
|
||||
| tuple[int, ...]
|
||||
| InputSpec
|
||||
| list[tuple[int, ...] | InputSpec]
|
||||
| None
|
||||
) = None,
|
||||
dtypes: str | Sequence[str] | None = None,
|
||||
input: Tensor | Sequence[Tensor] | dict[str, Tensor] | None = None,
|
||||
) -> ModelSummary:
|
||||
"""Prints a string summary of the network.
|
||||
|
||||
Args:
|
||||
net (Layer): The network which must be a subinstance of Layer.
|
||||
input_size (tuple|InputSpec|list[tuple|InputSpec]|None, optional): Size of input tensor. if model only
|
||||
have one input, input_size can be tuple or InputSpec. if model
|
||||
have multiple input, input_size must be a list which contain
|
||||
every input's shape. Note that input_size only dim of
|
||||
batch_size can be None or -1. Default: None. Note that
|
||||
input_size and input cannot be None at the same time.
|
||||
dtypes (str|Sequence[str]|None, optional): If dtypes is None, 'float32' will be used, Default: None.
|
||||
input (Tensor|Sequence[paddle.Tensor]|dict[str, paddle.Tensor]|None, optional): If input is given, input_size and dtype will be ignored, Default: None.
|
||||
|
||||
Returns:
|
||||
dict: A summary of the network including total params and total trainable params.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
:name: code-example-1
|
||||
|
||||
>>> # example 1: Single Input Demo
|
||||
>>> import paddle
|
||||
>>> import paddle.nn as nn
|
||||
>>> # Define Network
|
||||
>>> 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()
|
||||
>>> params_info = paddle.summary(lenet, (1, 1, 28, 28)) # doctest: +NORMALIZE_WHITESPACE
|
||||
---------------------------------------------------------------------------
|
||||
Layer (type) Input Shape Output Shape Param #
|
||||
===========================================================================
|
||||
Conv2D-1 [[1, 1, 28, 28]] [1, 6, 28, 28] 60
|
||||
ReLU-1 [[1, 6, 28, 28]] [1, 6, 28, 28] 0
|
||||
MaxPool2D-1 [[1, 6, 28, 28]] [1, 6, 14, 14] 0
|
||||
Conv2D-2 [[1, 6, 14, 14]] [1, 16, 10, 10] 2,416
|
||||
ReLU-2 [[1, 16, 10, 10]] [1, 16, 10, 10] 0
|
||||
MaxPool2D-2 [[1, 16, 10, 10]] [1, 16, 5, 5] 0
|
||||
Linear-1 [[1, 400]] [1, 120] 48,120
|
||||
Linear-2 [[1, 120]] [1, 84] 10,164
|
||||
Linear-3 [[1, 84]] [1, 10] 850
|
||||
===========================================================================
|
||||
Total params: 61,610
|
||||
Trainable params: 61,610
|
||||
Non-trainable params: 0
|
||||
---------------------------------------------------------------------------
|
||||
Input size (MB): 0.00
|
||||
Forward/backward pass size (MB): 0.11
|
||||
Params size (MB): 0.24
|
||||
Estimated Total Size (MB): 0.35
|
||||
---------------------------------------------------------------------------
|
||||
<BLANKLINE>
|
||||
>>> print(params_info)
|
||||
{'total_params': 61610, 'trainable_params': 61610}
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: code-example-2
|
||||
|
||||
>>> # example 2: multi input demo
|
||||
>>> import paddle
|
||||
>>> import paddle.nn as nn
|
||||
>>> class LeNetMultiInput(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, y):
|
||||
... x = self.features(inputs)
|
||||
...
|
||||
... if self.num_classes > 0:
|
||||
... x = paddle.flatten(x, 1)
|
||||
... x = self.fc(x + y)
|
||||
... return x
|
||||
>>> lenet_multi_input = LeNetMultiInput()
|
||||
|
||||
>>> params_info = paddle.summary(
|
||||
... lenet_multi_input, [(1, 1, 28, 28), (1, 400)], dtypes=['float32', 'float32']
|
||||
... ) # doctest: +NORMALIZE_WHITESPACE
|
||||
---------------------------------------------------------------------------
|
||||
Layer (type) Input Shape Output Shape Param #
|
||||
===========================================================================
|
||||
Conv2D-1 [[1, 1, 28, 28]] [1, 6, 28, 28] 60
|
||||
ReLU-1 [[1, 6, 28, 28]] [1, 6, 28, 28] 0
|
||||
MaxPool2D-1 [[1, 6, 28, 28]] [1, 6, 14, 14] 0
|
||||
Conv2D-2 [[1, 6, 14, 14]] [1, 16, 10, 10] 2,416
|
||||
ReLU-2 [[1, 16, 10, 10]] [1, 16, 10, 10] 0
|
||||
MaxPool2D-2 [[1, 16, 10, 10]] [1, 16, 5, 5] 0
|
||||
Linear-1 [[1, 400]] [1, 120] 48,120
|
||||
Linear-2 [[1, 120]] [1, 84] 10,164
|
||||
Linear-3 [[1, 84]] [1, 10] 850
|
||||
===========================================================================
|
||||
Total params: 61,610
|
||||
Trainable params: 61,610
|
||||
Non-trainable params: 0
|
||||
---------------------------------------------------------------------------
|
||||
Input size (MB): 0.00
|
||||
Forward/backward pass size (MB): 0.11
|
||||
Params size (MB): 0.24
|
||||
Estimated Total Size (MB): 0.35
|
||||
---------------------------------------------------------------------------
|
||||
<BLANKLINE>
|
||||
>>> print(params_info)
|
||||
{'total_params': 61610, 'trainable_params': 61610}
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: code-example-3
|
||||
|
||||
>>> # example 3: List Input Demo
|
||||
>>> import paddle
|
||||
>>> import paddle.nn as nn
|
||||
|
||||
>>> # list input demo
|
||||
>>> class LeNetListInput(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[0])
|
||||
...
|
||||
... if self.num_classes > 0:
|
||||
... x = paddle.flatten(x, 1)
|
||||
... x = self.fc(x + inputs[1])
|
||||
... return x
|
||||
>>> lenet_list_input = LeNetListInput()
|
||||
>>> input_data = [paddle.rand([1, 1, 28, 28]), paddle.rand([1, 400])]
|
||||
>>> params_info = paddle.summary(lenet_list_input, input=input_data) # doctest: +NORMALIZE_WHITESPACE
|
||||
---------------------------------------------------------------------------
|
||||
Layer (type) Input Shape Output Shape Param #
|
||||
===========================================================================
|
||||
Conv2D-1 [[1, 1, 28, 28]] [1, 6, 28, 28] 60
|
||||
ReLU-1 [[1, 6, 28, 28]] [1, 6, 28, 28] 0
|
||||
MaxPool2D-1 [[1, 6, 28, 28]] [1, 6, 14, 14] 0
|
||||
Conv2D-2 [[1, 6, 14, 14]] [1, 16, 10, 10] 2,416
|
||||
ReLU-2 [[1, 16, 10, 10]] [1, 16, 10, 10] 0
|
||||
MaxPool2D-2 [[1, 16, 10, 10]] [1, 16, 5, 5] 0
|
||||
Linear-1 [[1, 400]] [1, 120] 48,120
|
||||
Linear-2 [[1, 120]] [1, 84] 10,164
|
||||
Linear-3 [[1, 84]] [1, 10] 850
|
||||
===========================================================================
|
||||
Total params: 61,610
|
||||
Trainable params: 61,610
|
||||
Non-trainable params: 0
|
||||
---------------------------------------------------------------------------
|
||||
Input size (MB): 0.00
|
||||
Forward/backward pass size (MB): 0.11
|
||||
Params size (MB): 0.24
|
||||
Estimated Total Size (MB): 0.35
|
||||
---------------------------------------------------------------------------
|
||||
<BLANKLINE>
|
||||
>>> print(params_info)
|
||||
{'total_params': 61610, 'trainable_params': 61610}
|
||||
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: code-example-4
|
||||
|
||||
>>> # example 4: Dict Input Demo
|
||||
>>> import paddle
|
||||
>>> import paddle.nn as nn
|
||||
|
||||
>>> # Dict input demo
|
||||
>>> class LeNetDictInput(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['x1'])
|
||||
...
|
||||
... if self.num_classes > 0:
|
||||
... x = paddle.flatten(x, 1)
|
||||
... x = self.fc(x + inputs['x2'])
|
||||
... return x
|
||||
>>> lenet_dict_input = LeNetDictInput()
|
||||
>>> input_data = {'x1': paddle.rand([1, 1, 28, 28]), 'x2': paddle.rand([1, 400])}
|
||||
>>> # The module suffix number indicates its sequence in modules of the same type, used for differentiation identification
|
||||
>>> params_info = paddle.summary(lenet_dict_input, input=input_data) # doctest: +NORMALIZE_WHITESPACE
|
||||
---------------------------------------------------------------------------
|
||||
Layer (type) Input Shape Output Shape Param #
|
||||
===========================================================================
|
||||
Conv2D-1 [[1, 1, 28, 28]] [1, 6, 28, 28] 60
|
||||
ReLU-1 [[1, 6, 28, 28]] [1, 6, 28, 28] 0
|
||||
MaxPool2D-1 [[1, 6, 28, 28]] [1, 6, 14, 14] 0
|
||||
Conv2D-2 [[1, 6, 14, 14]] [1, 16, 10, 10] 2,416
|
||||
ReLU-2 [[1, 16, 10, 10]] [1, 16, 10, 10] 0
|
||||
MaxPool2D-2 [[1, 16, 10, 10]] [1, 16, 5, 5] 0
|
||||
Linear-1 [[1, 400]] [1, 120] 48,120
|
||||
Linear-2 [[1, 120]] [1, 84] 10,164
|
||||
Linear-3 [[1, 84]] [1, 10] 850
|
||||
===========================================================================
|
||||
Total params: 61,610
|
||||
Trainable params: 61,610
|
||||
Non-trainable params: 0
|
||||
---------------------------------------------------------------------------
|
||||
Input size (MB): 0.00
|
||||
Forward/backward pass size (MB): 0.11
|
||||
Params size (MB): 0.24
|
||||
Estimated Total Size (MB): 0.35
|
||||
---------------------------------------------------------------------------
|
||||
<BLANKLINE>
|
||||
>>> print(params_info)
|
||||
{'total_params': 61610, 'trainable_params': 61610}
|
||||
"""
|
||||
if input_size is None and input is None:
|
||||
raise ValueError("input_size and input cannot be None at the same time")
|
||||
|
||||
if input_size is None and input is not None:
|
||||
if paddle.is_tensor(input):
|
||||
input_size = tuple(input.shape)
|
||||
elif isinstance(input, (list, tuple)):
|
||||
input_size = []
|
||||
for x in input:
|
||||
input_size.append(tuple(x.shape))
|
||||
elif isinstance(input, dict):
|
||||
input_size = []
|
||||
for key in input.keys():
|
||||
input_size.append(tuple(input[key].shape))
|
||||
elif isinstance(input, paddle.base.framework.Variable):
|
||||
input_size = tuple(input.shape)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Input is not tensor, list, tuple and dict, unable to determine input_size, please input input_size."
|
||||
)
|
||||
|
||||
if isinstance(input_size, InputSpec):
|
||||
_input_size = tuple(input_size.shape)
|
||||
elif isinstance(input_size, list):
|
||||
_input_size = []
|
||||
for item in input_size:
|
||||
if isinstance(item, int):
|
||||
item = (item,)
|
||||
assert isinstance(item, (tuple, InputSpec)), (
|
||||
f'When input_size is list, \
|
||||
expect item in input_size is a tuple or InputSpec, but got {type(item)}'
|
||||
)
|
||||
|
||||
if isinstance(item, InputSpec):
|
||||
_input_size.append(tuple(item.shape))
|
||||
else:
|
||||
_input_size.append(item)
|
||||
elif isinstance(input_size, int):
|
||||
_input_size = (input_size,)
|
||||
else:
|
||||
_input_size = input_size
|
||||
|
||||
if not paddle.in_dynamic_mode():
|
||||
warnings.warn(
|
||||
"Your model was created in static graph mode, this may not get correct summary information!"
|
||||
)
|
||||
in_train_mode = False
|
||||
else:
|
||||
in_train_mode = net.training
|
||||
|
||||
if in_train_mode:
|
||||
net.eval()
|
||||
|
||||
def _is_shape(shape):
|
||||
for item in shape:
|
||||
if isinstance(item, (list, tuple)):
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_shape(shape):
|
||||
num_unknown = 0
|
||||
new_shape = []
|
||||
for i in range(len(shape)):
|
||||
item = shape[i]
|
||||
if item is None or item == -1:
|
||||
num_unknown += 1
|
||||
if num_unknown > 1:
|
||||
raise ValueError(
|
||||
'Option input_size only the dim of batch_size can be None or -1.'
|
||||
)
|
||||
item = 1
|
||||
elif isinstance(item, numbers.Number):
|
||||
if item <= 0:
|
||||
raise ValueError(
|
||||
f"Expected element in input size greater than zero, but got {item}"
|
||||
)
|
||||
new_shape.append(item)
|
||||
return tuple(new_shape)
|
||||
|
||||
def _check_input(input_size):
|
||||
if isinstance(input_size, (list, tuple)) and _is_shape(input_size):
|
||||
return _check_shape(input_size)
|
||||
else:
|
||||
return [_check_input(i) for i in input_size]
|
||||
|
||||
_input_size = _check_input(_input_size)
|
||||
|
||||
result, params_info = summary_string(net, _input_size, dtypes, input)
|
||||
print(result)
|
||||
|
||||
if in_train_mode:
|
||||
net.train()
|
||||
|
||||
return params_info
|
||||
|
||||
|
||||
@no_grad()
|
||||
def summary_string(model, input_size=None, dtypes=None, input=None):
|
||||
def _all_is_number(items):
|
||||
for item in items:
|
||||
if not isinstance(item, numbers.Number):
|
||||
return False
|
||||
return True
|
||||
|
||||
def _build_dtypes(input_size, dtype):
|
||||
if dtype is None:
|
||||
dtype = 'float32'
|
||||
|
||||
if isinstance(input_size, (list, tuple)) and _all_is_number(input_size):
|
||||
return [dtype]
|
||||
else:
|
||||
return [_build_dtypes(i, dtype) for i in input_size]
|
||||
|
||||
if not isinstance(dtypes, (list, tuple)):
|
||||
dtypes = _build_dtypes(input_size, dtypes)
|
||||
|
||||
summary_str = ''
|
||||
|
||||
depth = len(list(model.sublayers()))
|
||||
|
||||
def _get_shape_from_tensor(x):
|
||||
if isinstance(x, (paddle.base.Variable, paddle.base.core.eager.Tensor)):
|
||||
return list(x.shape)
|
||||
elif isinstance(x, (list, tuple)):
|
||||
return [_get_shape_from_tensor(xx) for xx in x]
|
||||
|
||||
def _get_output_shape(output):
|
||||
if isinstance(output, (list, tuple)):
|
||||
output_shape = [_get_output_shape(o) for o in output]
|
||||
elif hasattr(output, 'shape'):
|
||||
output_shape = list(output.shape)
|
||||
else:
|
||||
output_shape = []
|
||||
return output_shape
|
||||
|
||||
def register_hook(layer):
|
||||
def hook(layer, input, output):
|
||||
class_name = str(layer.__class__).split(".")[-1].split("'")[0]
|
||||
|
||||
try:
|
||||
layer_idx = int(layer._full_name.split('_')[-1])
|
||||
except:
|
||||
layer_idx = len(summary)
|
||||
|
||||
m_key = f"{class_name}-{layer_idx + 1}"
|
||||
summary[m_key] = OrderedDict()
|
||||
|
||||
try:
|
||||
summary[m_key]["input_shape"] = _get_shape_from_tensor(input)
|
||||
except:
|
||||
warnings.warn('Get layer {} input shape failed!')
|
||||
summary[m_key]["input_shape"] = []
|
||||
|
||||
try:
|
||||
summary[m_key]["output_shape"] = _get_output_shape(output)
|
||||
except:
|
||||
warnings.warn('Get layer {} output shape failed!')
|
||||
summary[m_key]["output_shape"]
|
||||
|
||||
params = 0
|
||||
|
||||
if paddle.in_dynamic_mode():
|
||||
layer_state_dict = layer._parameters
|
||||
else:
|
||||
layer_state_dict = layer.state_dict()
|
||||
|
||||
summary[m_key]["trainable_params"] = 0
|
||||
trainable_flag = False
|
||||
for k, v in layer_state_dict.items():
|
||||
params += int(np.prod(v.shape))
|
||||
|
||||
try:
|
||||
if (getattr(layer, k).trainable) and (
|
||||
not getattr(layer, k).stop_gradient
|
||||
):
|
||||
summary[m_key]["trainable_params"] += int(
|
||||
np.prod(v.shape)
|
||||
)
|
||||
summary[m_key]["trainable"] = True
|
||||
trainable_flag = True
|
||||
elif not trainable_flag:
|
||||
summary[m_key]["trainable"] = False
|
||||
except:
|
||||
summary[m_key]["trainable"] = True
|
||||
|
||||
summary[m_key]["nb_params"] = params
|
||||
|
||||
if (
|
||||
not isinstance(layer, nn.Sequential)
|
||||
and not isinstance(layer, nn.LayerList)
|
||||
and (not (layer == model) or depth < 1)
|
||||
):
|
||||
hooks.append(layer.register_forward_post_hook(hook))
|
||||
# For rnn, gru and lstm layer
|
||||
elif hasattr(layer, 'could_use_cudnn') and layer.could_use_cudnn:
|
||||
hooks.append(layer.register_forward_post_hook(hook))
|
||||
|
||||
if isinstance(input_size, tuple):
|
||||
input_size = [input_size]
|
||||
|
||||
def build_input(input_size, dtypes):
|
||||
if isinstance(input_size, (list, tuple)) and _all_is_number(input_size):
|
||||
if isinstance(dtypes, (list, tuple)):
|
||||
dtype = dtypes[0]
|
||||
else:
|
||||
dtype = dtypes
|
||||
return paddle.cast(paddle.rand(list(input_size)), dtype)
|
||||
else:
|
||||
return [
|
||||
build_input(i, dtype) for i, dtype in zip(input_size, dtypes)
|
||||
]
|
||||
|
||||
# create properties
|
||||
summary = OrderedDict()
|
||||
hooks = []
|
||||
# register hook
|
||||
model.apply(register_hook)
|
||||
if input is not None:
|
||||
x = input
|
||||
model(x)
|
||||
else:
|
||||
x = build_input(input_size, dtypes)
|
||||
# make a forward pass
|
||||
model(*x)
|
||||
|
||||
# remove these hooks
|
||||
for h in hooks:
|
||||
h.remove()
|
||||
|
||||
def _get_str_length(summary):
|
||||
head_length = {
|
||||
'layer_width': 15,
|
||||
'input_shape_width': 20,
|
||||
'output_shape_width': 20,
|
||||
'params_width': 15,
|
||||
'table_width': 75,
|
||||
}
|
||||
|
||||
for layer in summary:
|
||||
if head_length['output_shape_width'] < len(
|
||||
str(summary[layer]["output_shape"])
|
||||
):
|
||||
head_length['output_shape_width'] = len(
|
||||
str(summary[layer]["output_shape"])
|
||||
)
|
||||
if head_length['input_shape_width'] < len(
|
||||
str(summary[layer]["input_shape"])
|
||||
):
|
||||
head_length['input_shape_width'] = len(
|
||||
str(summary[layer]["input_shape"])
|
||||
)
|
||||
if head_length['layer_width'] < len(str(layer)):
|
||||
head_length['layer_width'] = len(str(layer))
|
||||
if head_length['params_width'] < len(
|
||||
str(summary[layer]["nb_params"])
|
||||
):
|
||||
head_length['params_width'] = len(
|
||||
str(summary[layer]["nb_params"])
|
||||
)
|
||||
|
||||
_temp_width = 0
|
||||
for k, v in head_length.items():
|
||||
if k != 'table_width':
|
||||
_temp_width += v
|
||||
|
||||
if head_length['table_width'] < _temp_width + 5:
|
||||
head_length['table_width'] = _temp_width + 5
|
||||
|
||||
return head_length
|
||||
|
||||
table_width = _get_str_length(summary)
|
||||
|
||||
summary_str += "-" * table_width['table_width'] + "\n"
|
||||
line_new = "{:^{}} {:^{}} {:^{}} {:^{}}".format(
|
||||
"Layer (type)",
|
||||
table_width['layer_width'],
|
||||
"Input Shape",
|
||||
table_width['input_shape_width'],
|
||||
"Output Shape",
|
||||
table_width['output_shape_width'],
|
||||
"Param #",
|
||||
table_width['params_width'],
|
||||
)
|
||||
summary_str += line_new + "\n"
|
||||
summary_str += "=" * table_width['table_width'] + "\n"
|
||||
total_params = 0
|
||||
total_output = 0
|
||||
trainable_params = 0
|
||||
max_length = 0
|
||||
for layer in summary:
|
||||
# input_shape, output_shape, trainable, nb_params
|
||||
line_new = "{:^{}} {:^{}} {:^{}} {:^{}}".format(
|
||||
layer,
|
||||
table_width['layer_width'],
|
||||
str(summary[layer]["input_shape"]),
|
||||
table_width['input_shape_width'],
|
||||
str(summary[layer]["output_shape"]),
|
||||
table_width['output_shape_width'],
|
||||
"{:,}".format(summary[layer]["nb_params"]),
|
||||
table_width['params_width'],
|
||||
)
|
||||
total_params += summary[layer]["nb_params"]
|
||||
|
||||
try:
|
||||
total_output += int(
|
||||
np.sum(np.prod(summary[layer]["output_shape"], axis=-1))
|
||||
)
|
||||
except:
|
||||
for output_shape in summary[layer]["output_shape"]:
|
||||
total_output += int(np.sum(np.prod(output_shape, axis=-1)))
|
||||
|
||||
if "trainable" in summary[layer]:
|
||||
if summary[layer]["trainable"]:
|
||||
trainable_params += summary[layer]["trainable_params"]
|
||||
summary_str += line_new + "\n"
|
||||
|
||||
def _get_input_size(input_size, size):
|
||||
if isinstance(input_size, (list, tuple)) and _all_is_number(input_size):
|
||||
size = abs(int(np.prod(input_size)) * 4.0 / (1024**2.0))
|
||||
else:
|
||||
size = sum([_get_input_size(i, size) for i in input_size])
|
||||
return size
|
||||
|
||||
total_input_size = _get_input_size(input_size, 0)
|
||||
|
||||
total_output_size = abs(
|
||||
2.0 * total_output * 4.0 / (1024**2.0)
|
||||
) # x2 for gradients
|
||||
total_params_size = abs(total_params * 4.0 / (1024**2.0))
|
||||
total_size = total_params_size + total_output_size + total_input_size
|
||||
|
||||
summary_str += "=" * table_width['table_width'] + "\n"
|
||||
summary_str += f"Total params: {total_params:,}" + "\n"
|
||||
summary_str += f"Trainable params: {trainable_params:,}" + "\n"
|
||||
summary_str += (
|
||||
f"Non-trainable params: {total_params - trainable_params:,}" + "\n"
|
||||
)
|
||||
summary_str += "-" * table_width['table_width'] + "\n"
|
||||
summary_str += f"Input size (MB): {total_input_size:0.2f}" + "\n"
|
||||
summary_str += (
|
||||
f"Forward/backward pass size (MB): {total_output_size:0.2f}" + "\n"
|
||||
)
|
||||
summary_str += f"Params size (MB): {total_params_size:0.2f}" + "\n"
|
||||
summary_str += f"Estimated Total Size (MB): {total_size:0.2f}" + "\n"
|
||||
summary_str += "-" * table_width['table_width'] + "\n"
|
||||
|
||||
# return summary
|
||||
return summary_str, {
|
||||
'total_params': total_params,
|
||||
'trainable_params': trainable_params,
|
||||
}
|
||||
Reference in New Issue
Block a user