CODEWITHSUNDEEP
pytorchconv-neural-network
Deepleeqe
Deepleeqe

Reputation: 317

Calculate variance with a kernel size in a tensor

Like what nn.Conv2d or nn.AvgPool2d do with a tensor and a kernel size, I would like to calculate the variances of a tensor with a kernel size. How can I achieve this? I guess maybe source code of pytorch should be touched?

Upvotes: 0

Views: 642

Answers (1)

Shai
Shai

Reputation: 114886

If it's only the variance you are after, you can use the fact that

var(x) = E[x^2] - E[x]^2

Using avg_pool2d you can estimate the local average of x and of x squared:

import torch.nn.functional as nnf

running_var = nnf.avg_pool2d(x**2, kernel_size=2, stride=1) - nnf.avg_pool2d(x, kernel_size=2,stride=1)**2

However, if you want a more general method of performing "sliding window" operations, you should become familiarized with unfold and fold:

u = nnf.unfold(x, kernel_size=2, stride=1)  # get all kernel_size patches as vectors
running_var2 = torch.var(u, unbiased=False, dim=1)
# reshape back to original shape ("folding")
running_var2 = running_var2.reshape(x.shape[0], 1, x.shape[2]-1, x.shape[3]-1)

Upvotes: 2

Related Questions