Python实现Grad-CAM时批量循环报错：TypeError类型错误排查

问题背景 #

我正在尝试实现Grad-CAM方法，代码如下：

for idx, (data, target, meta) in enumerate(tqdm(dataloader_test)):
        print('idx', idx)
        data, target = map(lambda x: x.to(device), (data, target))
        output = model(data)
        output[:,5].sum().backward()
        grads=grad[0].cpu().data.numpy().squeeze()
        fmap=activation[0].cpu().data.numpy().squeeze()
        tmp=grads.reshape([grads.shape[0],-1])
        # Get the mean value of the gradients of every featuremap
        weights=np.mean(tmp,axis=1)
        cam = np.zeros(grads.shape[1:])
        for i,w in enumerate(weights):
            cam += w*fmap[i,:]
        #relu
        cam=(cam>0)*cam
        #cam = np.maximum(cam, 0)
        #print("cam.shape",cam.shape)
        #normalize heatmap
        cam=cam/cam.max()*224
        
        
        for data in range(idx):
            #print('shape data', data.size)
            npic = np.array(torchvision.transforms.ToPILImage()(data).convert('RGB'))
            cam = cv2.resize(cam,(npic.shape[1], npic.shape[0]))
            heatmap=cv2.applyColorMap(np.uint8(cam),cv2.COLORMAP_JET)
            cam_img=npic*0.7+heatmap*0.3
            print(cam_img.shape)
            cv2.imwrite('./visualize/map'+str(count)+'.jpg', cam_img)
            count = count + 1

运行时触发错误：

raise TypeError(f"pic should be Tensor or ndarray. Got {type(pic)}.")

TypeError: pic should be Tensor or ndarray. Got <class 'int'>.

我认为问题出在批量数据的遍历方式上，无法正确获取需要处理的样本，求解决思路。

解决思路 #

1. 核心错误定位 #

内层循环for data in range(idx)完全覆盖了外层的data批量张量，这里的data变成了从0到idx-1的整数，传给torchvision.transforms.ToPILImage()时自然不符合输入要求（需要张量或数组），导致类型错误。

2. 正确的批量样本遍历方式 #

要处理当前batch中的每个样本，有两种正确方式：

方式一：直接遍历批量张量

# 替换原内层循环
for sample in data:
    # 将单个样本张量转为PIL图像再转numpy数组
    npic = np.array(torchvision.transforms.ToPILImage()(sample).convert('RGB'))
    # 复制原始CAM避免后续修改影响其他样本
    current_cam = cv2.resize(cam, (npic.shape[1], npic.shape[0]))
    heatmap = cv2.applyColorMap(np.uint8(current_cam), cv2.COLORMAP_JET)
    cam_img = npic * 0.7 + heatmap * 0.3
    # 转成uint8格式再保存
    cv2.imwrite(f'./visualize/map{count}.jpg', np.uint8(cam_img))
    count += 1

方式二：按索引遍历批量张量

# 替换原内层循环
for i in range(data.shape[0]):
    sample = data[i]
    npic = np.array(torchvision.transforms.ToPILImage()(sample).convert('RGB'))
    current_cam = cv2.resize(cam, (npic.shape[1], npic.shape[0]))
    heatmap = cv2.applyColorMap(np.uint8(current_cam), cv2.COLORMAP_JET)
    cam_img = npic * 0.7 + heatmap * 0.3
    cv2.imwrite(f'./visualize/map{count}.jpg', np.uint8(cam_img))
    count += 1

3. 额外优化点 #

• 单样本独立CAM计算：当前代码中output[:,5].sum().backward()是对整个batch的类别5求和后反向传播，得到的是batch内所有样本的平均梯度。如果需要为每个样本生成独立的CAM，需要循环每个样本单独计算：

  for i in range(data.shape[0]):
      model.zero_grad()  # 每次反向传播前清零梯度
      output = model(data)
      output[i, 5].backward()  # 仅对第i个样本的类别5求导
      # 提取梯度、特征图、计算当前样本的CAM...
      # 后续可视化逻辑...
  

• 避免CAM变量污染：每次处理样本时复制原始CAM，不要直接修改全局的cam变量，防止后续样本使用被resize后的CAM。
• 图像保存格式修正：cam_img是浮点数数组，转成np.uint8后再保存，避免颜色失真。

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