解决ValueError与TypeError：PyTorch数据可视化报错问题

问题背景 #

运行PyTorch结合Matplotlib的绘图代码时，出现两类错误：

• ValueError: object __array__ method not producing an array
• TypeError: Invalid shape (1, 1, 6, 6) for image data

涉事代码 #

import torch
import matplotlib.pyplot as plt

data = torch.tensor([[0, 0, 1, 1, 0, 0],
                     [0, 1, 0, 0, 1, 0],
                     [0, 0, 0, 0, 1, 0],
                     [0, 0, 1, 1, 0, 0],
                     [0, 1, 0, 0, 0, 0],
                     [0, 1, 1, 1, 1, 0]]).view(-1, 1, 6, 6).float()

# Convert data to a 2D array
data_np = data[0, 0].cpu().detach().numpy()

plt.imshow(data_np, cmap='gray')
plt.axis('off')
plt.show()

环境配置 #

• torch==2.3.1+cu121
• numpy==1.26.4
• matplotlib==3.9.0
• python==3.12.4

报错原因分析 #

1. TypeError: Invalid shape (1,1,6,6)：直接将4D张量（批量+通道+高+宽）传入plt.imshow导致，imshow要求灰度图为2D数组（高×宽），RGB/A图为3D数组（高×宽×通道）。
2. ValueError: object array method not producing an array：属于Matplotlib 3.9.0与Numpy 1.26.4在Python 3.12环境下的兼容问题，张量转numpy数组后，Matplotlib无法正确识别数组格式。

解决方案 #

方案1：修复数组转换逻辑，兼容版本问题 #

显式使用numpy.asarray确保转换后的对象是标准numpy数组，同时调整张量操作顺序（先分离计算图再移到CPU，避免潜在问题）：

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

data = torch.tensor([[0, 0, 1, 1, 0, 0],
                     [0, 1, 0, 0, 1, 0],
                     [0, 0, 0, 0, 1, 0],
                     [0, 0, 1, 1, 0, 0],
                     [0, 1, 0, 0, 0, 0],
                     [0, 1, 1, 1, 1, 0]]).view(-1, 1, 6, 6).float()

# 调整转换顺序并显式转为numpy数组
data_np = np.asarray(data[0, 0].detach().cpu())

plt.imshow(data_np, cmap='gray')
plt.axis('off')
plt.show()

方案2：降级Matplotlib版本 #

如果方案1无效，直接降级Matplotlib到兼容版本（如3.8.4），执行以下命令：

pip install matplotlib==3.8.4 --force-reinstall

验证说明 #

两种方案都能解决两类报错：方案1通过强制数组格式适配当前版本，方案2通过版本回退规避兼容问题，均可正常显示6×6的灰度图像。

内容的提问来源于stack exchange，提问作者M.M