Covariance Matrix for the. Dataset

I want to calculate the mean,standard deviation and Covariance matrix of MNIST Dataset for only digits 1,2,3…etc separately. (http://www-inst.eecs.berkeley.edu/~cs70/sp15/notes/n21.pdf) Page 12

I calculated the mean and SD using the below. I need covariance matrix tooo


Any help would be really great.

Screen Shot 2020-04-02 at 11.56.26 PM
Screen Shot 2020-04-02 at 11.56.26 PM
2414×1014 295 KB

def online_mean_and_sd(trainLoader):
    """Compute the mean and sd in an online fashion

        Var[x] = E[X^2] - E^2[X]
    """
    cnt = 0
    fst_moment = torch.empty(3)
    snd_moment = torch.empty(3)

    for images, _ in trainLoader:

        b, c, h, w = images.shape
        nb_pixels = b * h * w
        sum_ = torch.sum(images, dim=[0, 2, 3])
        sum_of_square = torch.sum(images ** 2, dim=[0, 2, 3])
        fst_moment = (cnt * fst_moment + sum_) / (cnt + nb_pixels)
        snd_moment = (cnt * snd_moment + sum_of_square) / (cnt + nb_pixels)

        cnt += nb_pixels

    return fst_moment, torch.sqrt(snd_moment - fst_moment ** 2)