Visualize feature maps of a given model module during inference. Args: x (torch.Tensor): Features to be visualized. module_type (str): Module type. stage (int): Module stage within the model. n (int, optional): Maximum number of feature maps to plot. Default
(x, module_type, stage, n=32, save_dir=Path('runs/detect/exp'))
| 672 | |
| 673 | |
| 674 | def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detect/exp')): |
| 675 | """ |
| 676 | Visualize feature maps of a given model module during inference. |
| 677 | |
| 678 | Args: |
| 679 | x (torch.Tensor): Features to be visualized. |
| 680 | module_type (str): Module type. |
| 681 | stage (int): Module stage within the model. |
| 682 | n (int, optional): Maximum number of feature maps to plot. Defaults to 32. |
| 683 | save_dir (Path, optional): Directory to save results. Defaults to Path('runs/detect/exp'). |
| 684 | """ |
| 685 | for m in ['Detect', 'Pose', 'Segment']: |
| 686 | if m in module_type: |
| 687 | return |
| 688 | batch, channels, height, width = x.shape # batch, channels, height, width |
| 689 | if height > 1 and width > 1: |
| 690 | f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename |
| 691 | |
| 692 | blocks = torch.chunk(x[0].cpu(), channels, dim=0) # select batch index 0, block by channels |
| 693 | n = min(n, channels) # number of plots |
| 694 | fig, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True) # 8 rows x n/8 cols |
| 695 | ax = ax.ravel() |
| 696 | plt.subplots_adjust(wspace=0.05, hspace=0.05) |
| 697 | for i in range(n): |
| 698 | ax[i].imshow(blocks[i].squeeze()) # cmap='gray' |
| 699 | ax[i].axis('off') |
| 700 | |
| 701 | LOGGER.info(f'Saving {f}... ({n}/{channels})') |
| 702 | plt.savefig(f, dpi=300, bbox_inches='tight') |
| 703 | plt.close() |
| 704 | np.save(str(f.with_suffix('.npy')), x[0].cpu().numpy()) # npy save |