Missing titles for the outer grids in a nested gridspec using matplotlib

Viewed 357

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:

enter image description here

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?

0 Answers
Related