| 448 | |
| 449 | def register_hook(layer): |
| 450 | def hook(layer, input, output): |
| 451 | class_name = str(layer.__class__).split(".")[-1].split("'")[0] |
| 452 | |
| 453 | try: |
| 454 | layer_idx = int(layer._full_name.split('_')[-1]) |
| 455 | except: |
| 456 | layer_idx = len(summary) |
| 457 | |
| 458 | m_key = f"{class_name}-{layer_idx + 1}" |
| 459 | summary[m_key] = OrderedDict() |
| 460 | |
| 461 | try: |
| 462 | summary[m_key]["input_shape"] = _get_shape_from_tensor(input) |
| 463 | except: |
| 464 | warnings.warn('Get layer {} input shape failed!') |
| 465 | summary[m_key]["input_shape"] = [] |
| 466 | |
| 467 | try: |
| 468 | summary[m_key]["output_shape"] = _get_output_shape(output) |
| 469 | except: |
| 470 | warnings.warn('Get layer {} output shape failed!') |
| 471 | summary[m_key]["output_shape"] |
| 472 | |
| 473 | params = 0 |
| 474 | |
| 475 | if paddle.in_dynamic_mode(): |
| 476 | layer_state_dict = layer._parameters |
| 477 | else: |
| 478 | layer_state_dict = layer.state_dict() |
| 479 | |
| 480 | summary[m_key]["trainable_params"] = 0 |
| 481 | trainable_flag = False |
| 482 | for k, v in layer_state_dict.items(): |
| 483 | params += int(np.prod(v.shape)) |
| 484 | |
| 485 | try: |
| 486 | if (getattr(layer, k).trainable) and ( |
| 487 | not getattr(layer, k).stop_gradient |
| 488 | ): |
| 489 | summary[m_key]["trainable_params"] += int( |
| 490 | np.prod(v.shape) |
| 491 | ) |
| 492 | summary[m_key]["trainable"] = True |
| 493 | trainable_flag = True |
| 494 | elif not trainable_flag: |
| 495 | summary[m_key]["trainable"] = False |
| 496 | except: |
| 497 | summary[m_key]["trainable"] = True |
| 498 | |
| 499 | summary[m_key]["nb_params"] = params |
| 500 | |
| 501 | if ( |
| 502 | not isinstance(layer, nn.Sequential) |