如何通过Numpy（朴素贝叶斯）简化嵌套循环？

Question

我想在Python numpy的帮助下简化一两行内的嵌套循环。这是MNIST上朴素贝叶斯模型的最后一部分

我了解一些numpy广播函数，但这里有一个复杂的3参数循环

def Bayes_pret(test_x, test_y, prioriP, posteriorP, image_size, num_class):
    pred = np.empty(test_x.shape[0])
    for i in range(test_x.shape[0]): #10000
        prob = np.empty(num_class)
        for j in range(num_class): #10
            temp = sum([(1-test_x[i][x])*math.log(1-posteriorP[j][x]) + test_x[i][x]*math.log(posteriorP[j][x]) for x in range(image_size**2)])
            # where image_size**2 = 784  
            prob[j] = np.array(math.log(prioriP[j]) + temp)
        pred[i] = np.argmax(prob)
        # print('label=',test_y[i],'and pred=',pred[i])
    return pred, (pred == test_y).sum()/ test_y.shape[0]

这浪费了大量的时间。我试着在这里用if-else技巧简化

temp = sum([math.log(1-posteriorP[j][x]) if test_x[i][x] == 0 else math.log(posteriorP[j][x]) for x in range(image_size**2)])

但这离完美还很远。我认为不应该有循环