Векторизация scipy norm.pdf

Question

def predictDigit(img):

    prob = [0] * 10

    for digit in range(10):
        for pix in range(len(img)):
            std = pix_std[digit][pix]
            mean = pix_means[digit][pix]
            if std == 0:
                continue
            else:
                prob[digit] += np.log(norm.pdf(img[pix], mean, std))
        prob[digit] += np.log(digit_prob[digit])

    return np.argmax(prob)

Я написал эту функцию, чтобы использовать ее для реализации наивного байесовского классификатора для классификации чисел. Идея состоит в том, чтобы пройти все пиксели входного изображения, добавить np.log(norm.pdf(img[pix], mean, std)) к prob и вернуть его argmax в конце, чтобы обозначить цифру входного изображения.

Однако это занимает слишком много времени. Я успешно векторизовал получение среднего и стандартного значения, используя:

pix_means[digit] = np.mean(image_cluster[digit], axis = 0) pix_std[digit] = np.std(image_cluster[digit], axis = 0)

Но я не уверен, возможна ли векторизация с norm.pdf.

Пожалуйста, помогите.

EDIT

Цифра Проб

digit_count = {}

for digit in y_train:
    if digit not in digit_count:
        digit_count[digit] = 1
    else:
        digit_count[digit] += 1

digit_prob = {}

for digit in range(10):
    digit_prob[digit] = digit_count[digit] / len(y_train)

image_cluster

image_cluster = {}

for image, digit in zip(x_train, y_train):
    if digit not in image_cluster:
        image_cluster[digit] = [image]
    else:
        image_cluster[digit].append(image)

pix_means = {}
pix_std = {}

# get mean and sd
for digit in range(10):
    pix_means[digit] = np.mean(image_cluster[digit], axis = 0)
    pix_std[digit] = np.std(image_cluster[digit], axis = 0)