added comments

This commit is contained in:
HTensor 2023-05-02 01:19:05 +08:00
parent a1cc25351d
commit e7033dabf9
2 changed files with 2 additions and 0 deletions

View File

@ -60,6 +60,7 @@ class Combined_Geo_Encoding_Volume:
@staticmethod @staticmethod
def corr(fmap1, fmap2): def corr(fmap1, fmap2):
# batch, dim, ht, wd
B, D, H, W1 = fmap1.shape B, D, H, W1 = fmap1.shape
_, _, _, W2 = fmap2.shape _, _, _, W2 = fmap2.shape
fmap1 = fmap1.view(B, D, H, W1) fmap1 = fmap1.view(B, D, H, W1)

View File

@ -167,6 +167,7 @@ class IGEVStereo(nn.Module):
match_right = self.desc(self.conv(features_right[0])) match_right = self.desc(self.conv(features_right[0]))
gwc_volume = build_gwc_volume(match_left, match_right, self.args.max_disp//4, 8) gwc_volume = build_gwc_volume(match_left, match_right, self.args.max_disp//4, 8)
gwc_volume = self.corr_stem(gwc_volume) gwc_volume = self.corr_stem(gwc_volume)
# 3d unet
gwc_volume = self.corr_feature_att(gwc_volume, features_left[0]) gwc_volume = self.corr_feature_att(gwc_volume, features_left[0])
geo_encoding_volume = self.cost_agg(gwc_volume, features_left) geo_encoding_volume = self.cost_agg(gwc_volume, features_left)