Update `feature_visualization()` (#3807)

* Update `feature_visualization()`

Only plot for data with height, width > 1

* cleanup

* Cleanup
This commit is contained in:
Glenn Jocher 2021-06-28 13:48:14 +02:00 committed by GitHub
parent 20d45aa4f1
commit 02719dde52
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 19 additions and 17 deletions

View File

@ -448,26 +448,28 @@ def plot_results(start=0, stop=0, bucket='', id=(), labels=(), save_dir=''):
fig.savefig(Path(save_dir) / 'results.png', dpi=200) fig.savefig(Path(save_dir) / 'results.png', dpi=200)
def feature_visualization(features, module_type, module_idx, n=64): def feature_visualization(x, module_type, stage, n=64):
""" """
features: Features to be visualized x: Features to be visualized
module_type: Module type module_type: Module type
module_idx: Module layer index within model stage: Module stage within model
n: Maximum number of feature maps to plot n: Maximum number of feature maps to plot
""" """
project, name = 'runs/features', 'exp' batch, channels, height, width = x.shape # batch, channels, height, width
save_dir = increment_path(Path(project) / name) # increment run if height > 1 and width > 1:
save_dir.mkdir(parents=True, exist_ok=True) # make dir project, name = 'runs/features', 'exp'
save_dir = increment_path(Path(project) / name) # increment run
save_dir.mkdir(parents=True, exist_ok=True) # make dir
plt.figure(tight_layout=True) plt.figure(tight_layout=True)
blocks = torch.chunk(features, features.shape[1], dim=1) # block by channel dimension blocks = torch.chunk(x, channels, dim=1) # block by channel dimension
n = min(n, len(blocks)) n = min(n, len(blocks))
for i in range(n): for i in range(n):
feature = transforms.ToPILImage()(blocks[i].squeeze()) feature = transforms.ToPILImage()(blocks[i].squeeze())
ax = plt.subplot(int(math.sqrt(n)), int(math.sqrt(n)), i + 1) ax = plt.subplot(int(math.sqrt(n)), int(math.sqrt(n)), i + 1)
ax.axis('off') ax.axis('off')
plt.imshow(feature) # cmap='gray' plt.imshow(feature) # cmap='gray'
f = f"layer_{module_idx}_{module_type.split('.')[-1]}_features.png" f = f"stage_{stage}_{module_type.split('.')[-1]}_features.png"
print(f'Saving {save_dir / f}...') print(f'Saving {save_dir / f}...')
plt.savefig(save_dir / f, dpi=300) plt.savefig(save_dir / f, dpi=300)