如何为整个图像数据集的RGB值计算3x3协方差矩阵?

问题描述

我需要为整个图像数据集计算RGB值的协方差矩阵,然后将Cholesky分解应用于最终结果。

RGB值的协方差矩阵是3x3矩阵M,其中M_(i,i)是通道i的方差,M_(i,j)是通道i和j之间的协方差。

最终结果应该是这样的:

([[0.26,0.09,0.02],[0.27,0.00,-0.05],-0.09,0.03]])

即使Numpy具有Cov功能,我还是希望坚持使用PyTorch功能

我尝试根据其他Cov实现和克隆在PyTorch中重新创建numpy Cov函数

def pytorch_cov(tensor,tensor2=None,rowvar=True):
    if tensor2 is not None:
        tensor = torch.cat((tensor,tensor2),dim=0)
    tensor = tensor.view(1,-1) if tensor.dim() < 2 else tensor
    tensor = tensor.t() if not rowvar and tensor.size(0) != 1 else tensor
    tensor = tensor - torch.mean(tensor,dim=1,keepdim=True)
    return 1 / (tensor.size(1) - 1) * tensor.mm(tensor.t())

def cov_vec(x):
    c = x.size(0)
    m1 = x - torch.sum(x,dim=[1],keepdims=True)/ c
    out = torch.einsum('ijk,ilk->ijl',m1,m1)  / (c - 1)
    return out

数据集加载如下:

dataset = torchvision.datasets.ImageFolder(data_path)
loader = torch.utils.data.DataLoader(dataset)

for images,_ in loader:
    batch_size = images.size(0) 
    ...

目前,我仅在尝试使用torch.randn(batch_size,3,height,width)创建的图像。

编辑:

我正在尝试从Tensorflow的Lucid here复制矩阵,并在distill.pub here上作了一些解释。

第二次修改

为了使输出类似于示例一,您必须执行此操作而不是使用Cholesky:

rgb_cov_tensor = rgb_cov_tensor / len(loader.dataset)
U,S,V = torch.svd(rgb_cov_tensor)
epsilon = 1e-10
svd_sqrt = U @ torch.diag(torch.sqrt(S + epsilon))

然后,可以将所得矩阵用于执行颜色去相关,这对于可视化要素(DeepDream)很有用。我已经在我的项目here中实现了它。

解决方法

这是一个用于在名为rgb_cov的3通道图像上计算(无偏)样本协方差矩阵的函数。借助torch.cholesky,Cholesky分解非常简单:

import torch
def rgb_cov(im):
    '''
    Assuming im a torch.Tensor of shape (H,W,3):
    '''
    im_re = im.reshape(-1,3)
    im_re -= im_re.mean(0,keepdim=True)
    return 1/(im_re.shape[0]-1) * im_re.T @ im_re

#Test:
im = torch.randn(50,50,3)
cov = rgb_cov(im)
L_cholesky = torch.cholesky(cov)