PerceptualSimilarity
PerceptualSimilarity copied to clipboard
richzhang
A new `forward_distmat` method to accelerate distance matrix computing among images.
Summary
Add another forward function forward_distmat to compute the distance matrix with one or two stacks of images.
Rationale
The currentforward method of LPIPS supports computing the distance between two images (A,B), or between one image to a batch of images (As,Bs) or the distance between corresponding pairs in two batches (As,Bs).
Because of this, computing the distance matrix cannot be directly performed using the current forward function
It needs to be computed row by row, for example given a imgtsrs of shape [B,C,H,W]
Dist = LPIPS(net="squeeze", ).cuda()
Dist.requires_grad_(False)
distmat_bin = []
for i in range(imgtsrs.shape[0]):
dist_in_bin = Dist(imgtsrs.cuda(), imgtsrs[i:i + 1].cuda()).cpu().squeeze()
distmat_bin.append(dist_in_bin)
distmat_bin = torch.stack(distmat_bin, dim=0)
A simple fix to accelerate this is to move this loop into the forward function. Then it can reuse the features computed for images, so it largely accelerates the speed of computing a large distance matrix among a set of images.
distmat_bin = Dist.forward_distmat(imgtsrs_rf.cuda(), None).cpu().squeeze()
Memory usage.
We choose to compute the matrix row by row, such that the memory footprint shall be still linear in the number of images instead of quadratic.
TODO
It could be extended to compute a few rows in a batch. The batch size needs to be chosen w.r.t. memory. Currently, it's fast enough for most of my works, but there is space for further optimization. If there is interest in more efficiency I can help implement more.
