如何设置matplotlib add_subplot行列以实现8×2布局的16张子图？

我尝试生成16张子图，想要实现8×2的布局（每行8个，共2行），但当前代码生成的结果不符合预期。我的代码和输入数据如下：

原可视化函数代码：

def visualize_t2t(token_dict, scores):

  fig = plt.figure(figsize=(50, 50))

  for idx, scores in enumerate(scores):
      scores_np = np.array(scores)

      ax = fig.add_subplot(12, 12, idx+1)
      # append the attention weights
      im = ax.imshow(scores, cmap='viridis')

      fontdict = {'fontsize': 3}

      ax.set_xticks(range(len(all_tokens)))
      ax.set_yticks(range(len(all_tokens)))
      ax.set_xticklabels(all_tokens, fontdict=fontdict, rotation=90)
      ax.set_yticklabels(all_tokens, fontdict=fontdict)
      ax.set_xlabel('{} {}'.format('label_name', idx+1))

      fig.colorbar(im, fraction=0.046, pad=0.04)
  
  plt.tight_layout()
  name_f = str(uuid.uuid4())
  plt.savefig(f'{name_f}.pdf',
  bbox_inches='tight',
    dpi=350)

输入数据：

all_tokens = ['[CLS]',
 'what',
 'type',
 'of',
 'heart',
 'issue',
 'does',
 'the',
 'Person',
 'have',
 '[CLS]']

dummy_input = np.random.uniform(-1, 1, [16, len(all_tokens), len(all_tokens)])
visualize_t2t(all_tokens, dummy_input)

解决方法 #

核心问题是add_subplot的参数错误，当前设置的是12行12列的网格，和目标的8×2布局不匹配。需要做以下调整：

• 将add_subplot的参数改为(2, 8, idx+1)，参数顺序为行数、列数、子图的位置索引
• 调整画布尺寸figsize，8列布局不需要过大的宽度，推荐设置为(32, 8)，让子图显示更紧凑合理
• 修复循环中变量名冲突的问题：循环里的scores覆盖了函数入参的scores，建议重命名为score

修改后的完整代码：

import matplotlib.pyplot as plt
import numpy as np
import uuid

def visualize_t2t(token_dict, scores):
    # 调整画布尺寸，适配8×2布局
    fig = plt.figure(figsize=(32, 8))

    # 重命名循环变量避免覆盖入参
    for idx, score in enumerate(scores):
        score_np = np.array(score)
        # 设置子图为2行8列，当前子图的索引为idx+1
        ax = fig.add_subplot(2, 8, idx+1)
        im = ax.imshow(score_np, cmap='viridis')

        fontdict = {'fontsize': 3}
        ax.set_xticks(range(len(all_tokens)))
        ax.set_yticks(range(len(all_tokens)))
        ax.set_xticklabels(all_tokens, fontdict=fontdict, rotation=90)
        ax.set_yticklabels(all_tokens, fontdict=fontdict)
        ax.set_xlabel('{} {}'.format('label_name', idx+1))

        fig.colorbar(im, fraction=0.046, pad=0.04)
  
    plt.tight_layout()
    name_f = str(uuid.uuid4())
    plt.savefig(f'{name_f}.pdf', bbox_inches='tight', dpi=350)

# 输入数据
all_tokens = ['[CLS]',
 'what',
 'type',
 'of',
 'heart',
 'issue',
 'does',
 'the',
 'Person',
 'have',
 '[CLS]']

dummy_input = np.random.uniform(-1, 1, [16, len(all_tokens), len(all_tokens)])
visualize_t2t(all_tokens, dummy_input)

内容的提问来源于stack exchange，提问作者Aaditya Ura