image833×110 4.87 KB
As Deqin Sun mentioned, two different cost volumes exist.
I write a naive pseudo code
fmap1 = torch.ones(C, H, W) # [C, H, W]
fmap2 = torch.ones(C, H, W)


# Full cost volume (all-pairs correlation):
cost_vol = torch.ones(H, W, 1, H, W)
for i in range(H):
	for j in range(W):
		vec1 = fmap1(:, i, j)
		corr_ij = dot(vec1.T, fmap2.view(:, -1)) # [1, C] *[C, H*W] = [1, H*W]
		cost_vol(i, j, 0, corr_ij.view(1, H, W))

cost_vol = cost_vol.view(H, W, H*W)

# Partial cost volume
r = 3
dx = linespace(-r, r, 2*r+1)
dy = linespace(-r, r, 2*r+1)
grid = stack(meshgrid(dy, dx), axis=-1)

cost_vol = torch.ones(H, W, (2*r+1)**2)
for i in range(H):
	for j in range(W):
		vec1 = fmap1(:, i, j)
		fmap2_neighbor = grid_sample(fmap2, grid+[i, j]) # sample a sqare fmap2_neighbor at (i, j) in fmap2 with size(2*r+1, 2*r+1)
		corr_ij = dot(vec1.T, fmap2_neighbor.view(:, -1)) # [1, C] *[C, (2*r+1)**2] = [1, 49] 
		cost_vol(i, j, :) = corr_ij

# the first and second channel of final cost volume should be equal to the fmap1 shape. 
# We don't need to care about the order of rest of channels but the product result of the number of other channels indicates that how many feature in fmap2 we calculate the correlation with each feature in fmap1. EX: [H, W, H*W] means all-pairs correlation known as each feature in fmap1 is correlated with each feature in fmap2;   [H, W, 49] means each feature in fmap1 is correlated with a [7, 7] in fmap2.