Оцените стандартное и среднее значение для большого набора данных изображения с помощью метода Велфорда с использованием pytorch

Question

Я пытаюсь оценить стандартное и среднее значение для большого набора данных изображения с помощью алгоритма Велфорда. Насколько я понимаю, этот алгоритм более надежен, чем метод, который накапливает сумму и квадрат суммы значений. Однако я столкнулся с некоторыми проблемами. Во-первых, я протестировал алгоритм на небольшом случайном наборе данных (6 изображений), и среднее и стандартное значения, оцененные методом Велфорда, не совсем совпадают со средним и стандартным значениями, рассчитанными с использованием всего набора данных одновременно. Разница очень мала, но я не уверен, что делаю что-то не так.

import torch 
torch.set_printoptions(precision=10) 
class ImgRunningStats:

        def __init__(self):
            self.n = torch.zeros(1, dtype=torch.long)
            self.old_m = torch.zeros(3, dtype=torch.float64)
            self.new_m = torch.zeros(3, dtype=torch.float64)
            self.old_s = torch.zeros(3, dtype=torch.float64)
            self.new_s = torch.zeros(3, dtype=torch.float64)

        def push(self, x):
            self.n += 1

            if self.n == 1:
                self.old_m = self.new_m = x
                self.old_s = 0
            else:
                self.new_m = self.old_m + (x - self.old_m) / self.n
                self.new_s = self.old_s + (x - self.old_m) * (x - self.new_m)

                self.old_m = self.new_m
                self.old_s = self.new_s

        def mean(self):
            return self.new_m if self.n else 0.0

        def var(self):
            return self.new_s / (self.n - 1) if self.n > 1 else 0.0

        def std(self):
            return torch.sqrt(self.var())



    rs= ImgRunningStats()

    test_img_0 = torch.rand([2,3,250,250]) test_img_1 = torch.rand([2,3,250,250]) test_img_2 = torch.rand([2,3,250,250])

    images = [test_img_0, test_img_1, test_img_2]

    test_img = torch.cat([test_img_0, test_img_1, test_img_2], dim=0)

    print('mean full: ', test_img.mean(dim=[0,2,3])) print('std full: ', test_img.std(dim=[0,2,3], unbiased=True))


    for img in images:
        pixels = img.view(-1, img.shape[1])
        for px in pixels:
            rs.push(px)


    print(f'Mean Welford: {rs.mean()}, Std Welford: {rs.std()}')

Однако основная проблема, которая у меня есть, заключается в том, что алгоритм очень медленный ... Я не знаю как я могу улучшить производительность алгоритма. Насколько я видел, невозможно запустить этот алгоритм в большом наборе данных изображений с текущей реализацией.

Можно ли распараллелить этот алгоритм?