如何使用Dataiku训练的Efficient B4模型Pickle文件进行图像分类预测？

问题 #

在Dataiku上训练EfficientNet B4图像分类模型并导出为Pickle文件后，使用自定义预处理得到的numpy数组调用pred.predict()时，触发以下错误：

ERROR:dataiku.doctor.deephub.deephub_torch_datasets:couldn't get item at index 0
Traceback (most recent call last):
  File "/home/utshav/dataiku-dss-11.3.2/python/dataiku/doctor/deephub/deephub_torch_datasets.py", line 35, in __getitem__
    return self._getitem(idx)
  File "/home/utshav/dataiku-dss-11.3.2/python/dataiku/doctor/deephub/deephub_torch_datasets.py", line 131, in _getitem
    row = self.get_row(idx)
  File "/home/utshav/dataiku-dss-11.3.2/python/dataiku/doctor/deephub/deephub_torch_datasets.py", line 140, in get_row
    return self.df.iloc[idx]
AttributeError: 'numpy.ndarray' object has no attribute 'iloc'

报错原因 #

Dataiku封装的predictor.predict()方法不接受直接传入numpy数组，它期望输入为Pandas DataFrame格式。内部代码会将输入视为DataFrame并调用iloc方法提取数据，而numpy数组没有该属性，因此触发错误。

解决方法 #

提供两种可行的解决思路：

方法1：构造符合要求的DataFrame输入 #

根据训练时的输入结构，构造对应列名的DataFrame传入predict()：

• 如果训练时使用图像路径列（列名例如image_path），直接传入包含路径的DataFrame：

import pandas as pd

# test_image为图像路径字符串
df = pd.DataFrame({"image_path": [test_image]})
logits = pred.predict(df)

• 如果训练时使用预处理后的像素数组列（列名例如image_pixels），将numpy数组放入DataFrame对应列：

import pandas as pd

# image1是预处理后的(1, 380, 380, 3)格式numpy数组，去掉batch维度后放入DataFrame
df = pd.DataFrame({"image_pixels": [image1[0]]})
logits = pred.predict(df)

注：列名必须与训练时的输入列名完全一致

方法2：提取底层PyTorch模型直接预测 #

绕过Dataiku的封装层，直接使用训练好的PyTorch模型进行预测，更灵活适配自定义预处理：

import torch

# 从Dataiku加载的模型中提取原始PyTorch模型
torch_model = model.get_raw_model()
torch_model.eval()  # 设置为评估模式

# 将numpy数组转换为PyTorch张量，匹配模型输入要求
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
image_tensor = torch.from_numpy(image1).permute(0, 3, 1, 2).float().to(device)
# 注：EfficientNet输入格式为(batch, channels, height, width)，需调整维度顺序

# 执行预测
with torch.no_grad():
    logits = torch_model(image_tensor)
    # 如需输出概率，添加softmax处理
    probabilities = torch.nn.functional.softmax(logits, dim=1).cpu().numpy()

完整修正示例（方法2） #

import cv2
import numpy as np
import pickle
import torch

def preprocess_image(image_path, target_size=(380, 380)):
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, target_size)
    img_array = np.array(img)
    img_array = img_array / 255.0
    img_batch = np.expand_dims(img_array, axis=0)
    return img_batch

def load_model(pickle_path):
    with open(pickle_path, 'rb') as f:
        model = pickle.load(f)
    return model

# 加载模型并提取底层PyTorch模型
model = load_model('model.pkl')
torch_model = model.get_raw_model()
torch_model.eval()

# 预处理图像
test_image = "your_test_image_path.jpg"
image1 = preprocess_image(test_image)

# 转换为模型接受的张量格式
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
image_tensor = torch.from_numpy(image1).permute(0, 3, 1, 2).float().to(device)

# 预测
with torch.no_grad():
    logits = torch_model(image_tensor)
    probabilities = torch.nn.functional.softmax(logits, dim=1).cpu().numpy()
    print("预测概率：", probabilities)

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