r"""Calculate and print ``model``'s statistics by adding hook and record Module's inputs outputs size. Args: model: model that need to get stats info. inputs: user defined input data for running model and calculating stats, alternative with input_shapes. input_shapes: sh
(
model: M.Module,
inputs: Iterable[np.ndarray] = None,
input_shapes: list = None,
cal_params: bool = True,
cal_flops: bool = True,
cal_activations: bool = True,
logging_to_stdout: bool = True,
bar_length_max: int = 20,
)
| 422 | |
| 423 | |
| 424 | def module_stats( |
| 425 | model: M.Module, |
| 426 | inputs: Iterable[np.ndarray] = None, |
| 427 | input_shapes: list = None, |
| 428 | cal_params: bool = True, |
| 429 | cal_flops: bool = True, |
| 430 | cal_activations: bool = True, |
| 431 | logging_to_stdout: bool = True, |
| 432 | bar_length_max: int = 20, |
| 433 | ): |
| 434 | r"""Calculate and print ``model``'s statistics by adding hook and record Module's inputs outputs size. |
| 435 | |
| 436 | Args: |
| 437 | model: model that need to get stats info. |
| 438 | inputs: user defined input data for running model and calculating stats, alternative with input_shapes. |
| 439 | input_shapes: shapes to generate random inputs for running model and calculating stats, alternative with inputs. |
| 440 | cal_params: whether calculate and record params size. |
| 441 | cal_flops: whether calculate and record op flops. |
| 442 | cal_activations: whether calculate and record op activations. |
| 443 | logging_to_stdout: whether print all calculated statistic details. |
| 444 | bar_length_max: size of bar indicating max flops or parameter size in net stats. |
| 445 | """ |
| 446 | has_inputs = False |
| 447 | if inputs is not None: |
| 448 | has_inputs = True |
| 449 | if not isinstance(inputs, (tuple, list)): |
| 450 | inputs = [inputs] |
| 451 | |
| 452 | def load_tensor(x): |
| 453 | if isinstance(x, np.ndarray): |
| 454 | return Tensor(x) |
| 455 | elif isinstance(x, collections.abc.Mapping): |
| 456 | return {k: load_tensor(v) for k, v in x.items()} |
| 457 | elif isinstance(x, tuple) and hasattr(x, "_fields"): # nametuple |
| 458 | return type(x)(*(load_tensor(value) for value in x)) |
| 459 | elif isinstance(x, collections.abc.Sequence): |
| 460 | return [load_tensor(v) for v in x] |
| 461 | else: |
| 462 | return Tensor(x, dtype=np.float32) |
| 463 | |
| 464 | inputs = load_tensor(inputs) |
| 465 | |
| 466 | else: |
| 467 | if input_shapes: |
| 468 | if not isinstance(input_shapes[0], tuple): |
| 469 | input_shapes = [input_shapes] |
| 470 | inputs = [F.zeros(in_size, dtype=np.float32) for in_size in input_shapes] |
| 471 | else: |
| 472 | logger.error( |
| 473 | "Inputs or input_shapes is required for running model and calculating stats.", |
| 474 | exc_info=True, |
| 475 | ) |
| 476 | return |
| 477 | if not cal_activations: |
| 478 | log_activations = False |
| 479 | |
| 480 | disable_receptive_field() |
| 481 | recorded_parameters = set() |