I would like to visualize all the attention heads together with a batch of source and target sequence (both of them are 2D arrays) in a transformer. I am following existing examples to make a nested gridspec. But existing examples only show how to put titles on the inner grids. But none of them shows how to put the title in the outer grid.
This is the closest result I can get
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.gridspec as gridspec
predictions = {}
predictions['attention'] = np.random.randn(4, 2, 5, 5)
predictions['frame'] = np.zeros((4, 2, 5, 2))
mel = np.ones((4, 2, 2, 5))
fig = plt.figure(figsize=(30, 15))
# Creating the grid for 2 attention head for the transformer
outer = gridspec.GridSpec(1, 2, wspace=0.2, hspace=0.2)
fig.suptitle("Visualizing Attention Heads", size=20)
for i in range(2):
# Creating the grid for 4 samples
inner = gridspec.GridSpecFromSubplotSpec(2, 2,
subplot_spec=outer[i], wspace=0.1, hspace=0.1)
ax = plt.Subplot(fig, outer[i])
ax.set_title(f'Head {i}', size=20) # This does not show up
for idx in range(4):
axCenter = plt.Subplot(fig, inner[idx])
fig.add_subplot(axCenter)
attentions = predictions['attention'][idx, i]
axCenter.imshow(attentions, cmap='jet', vmin=0, vmax=0.7)
attended_features = mel[idx, i]
# Create another plot on top and left of the attention map
divider = make_axes_locatable(axCenter)
axvert = divider.append_axes('left', size='30%', pad=0.5)
axhoriz = divider.append_axes('top', size='20%', pad=0.25)
axhoriz.imshow(attended_features, aspect='auto', origin='lower', cmap='jet')
axvert.imshow(predictions['frame'][idx, i], aspect='auto')
# changing axis for the center fig
axCenter.set_xticks([])
# changing axis for the output fig (left fig)
axvert.set_yticks([])
axvert.xaxis.tick_top()
axvert.set_title('Transcription')
axhoriz.set_title(f'Attended Feature (Spec)')
axhoriz.margins(x=0)
axvert.margins(y=0)
And the figure produced is like this:
As you can see, the big title for the whole plot and the titles for the inner grids are all there. But the titles for the outer grids are nowhere to be found (as indicated by the red circle). How to make the titles for outer grids appear?
