如何优化此 NumPy 代码？

Question

以下代码是我的 Python 代码的瓶颈：

def get_payoff(self,  actual, predicted):
    if abs(actual - 1.0) < 1e-5:  # if actual == 1
        if predicted < 0.5:
            return self.fn_payoff * (0.5 - predicted)
        elif predicted > 0.5:
            return self.tp_payoff * (predicted - 0.5)
        else:
            return 0
    else:
        if predicted < 0.5:
            return self.tn_payoff * (0.5 - predicted)
        elif predicted > 0.5:
            return self.fp_payoff * (predicted - 0.5)
        else:
            return 0

def get_total_payoff(self):
    total_payoff = 0
    for target_element, prediction_element in zip(np.nditer(self.target), np.nditer(predictions)):
        total_payoff += self.get_payoff(target_element, prediction_element)

fn_payoff、tp_payoff、tn_payoff 和 fp_payoff 都是浮点数。 self.target 和 self.predictions 都是 numpy ndarrays。

我假设有某种方法可以用某种 numpy 向量化替换 get_total_payoff 中的 for 循环，但我不知道如何处理 if/then 语句以正确进行向量化。

Answer 1

根据条件使用不同表达式的矢量化函数的关键是使用np.choose。此外，在您的情况下，predict-0.5 和0.5-predict 可以替换为abs(predict-0.5)，加上对predict==0.5 的特殊处理（我猜特殊处理是为了正确处理NaN）。

import numpy as np

class A(object):
    def __init__(self):
        self.fn_payoff = 222.
        self.tn_payoff = 444.
        self.fp_payoff = 777.
        self.tp_payoff = 888.
        self.target = np.array([ 0.3, 1., 2. ])
        self.predictions = np.array([ 0.4, 0.5, 1.7 ])

    def get_payoff(self,  actual, predicted):
        if abs(actual - 1.0) < 1e-5:  # if actual == 1
            if predicted < 0.5:
                return self.fn_payoff * (0.5 - predicted)
            elif predicted > 0.5:
                return self.tp_payoff * (predicted - 0.5)
            else:
                return 0
        else:
            if predicted < 0.5:
                return self.tn_payoff * (0.5 - predicted)
            elif predicted > 0.5:
                return self.fp_payoff * (predicted - 0.5)
            else:
                return 0

    def get_total_payoff(self):
        total_payoff = 0
        for target_element, prediction_element in zip(np.nditer(self.target), np.nditer(self.predictions)):
            total_payoff += self.get_payoff(target_element, prediction_element)
        return total_payoff

    def get_total_payoff_VECTORIZED(self):
        actual_mask = np.abs(self.target - 1) < 1e-5
        predict_mask = self.predictions < 0.5
        payoff_n = np.choose(actual_mask, [ self.tn_payoff, self.fn_payoff ])
        payoff_p = np.choose(actual_mask, [ self.fp_payoff, self.tp_payoff ])
        payoff = np.choose(predict_mask, [ payoff_p, payoff_n ]) * abs(self.predictions-0.5)
        payoff[self.predictions==0.5] = 0
        return payoff.sum()

a = A()
print a.get_total_payoff()
=> 976.8
print a.get_total_payoff_VECTORIZED()
=> 976.8

Answer 2

def _get_payoff(self, actual, predicted):
    pred_factor = numpy.abs(0.5 - predicted)
    payoff_selector = 2*numpy.isclose(actual, 1) + (predicted < 0.5)
    payoff = numpy.choose(payoff_selector,
                          [
                              self.fp_payoff,
                              self.tn_payoff,
                              self.tp_payoff,
                              self.fn_payoff,
                          ])
    return numpy.sum(payoff * pred_factor)

def get_total_payoff(self):
    return self._get_payoff(self.target, predictions)

我们使用numpy.choose 生成一组收益选择，并将其与一组0.5 和预测值之间的绝对差值相乘，然后求和。 numpy.isclose 用于测试 actual 的值是否接近 1。我们可以忽略 predicted == 0.5 的情况，因为乘以 numpy.abs(0.5 - predicted) 无论如何都会得到正确的结果 0。如果保证self.target 和predictions 是一维的，那么numpy.dot 的性能可能比单独的相乘和相加要好。