Change predicted_depth to fp32

Megvii-BaseDetection · Jan 17, 2023 · d78c7b5 · d78c7b5
1 parent a2d40de
commit d78c7b5
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 5 deletions.
diff --git a/bevdepth/exps/nuscenes/base_exp.py b/bevdepth/exps/nuscenes/base_exp.py
@@ -270,7 +270,7 @@ def get_depth_loss(self, depth_labels, depth_preds):
 
         with autocast(enabled=False):
             depth_loss = (F.binary_cross_entropy(
-                depth_preds[fg_mask].float(),
+                depth_preds[fg_mask],
                 depth_labels[fg_mask],
                 reduction='none',
             ).sum() / max(1.0, fg_mask.sum()))

diff --git a/bevdepth/layers/backbones/base_lss_fpn.py b/bevdepth/layers/backbones/base_lss_fpn.py
@@ -550,7 +550,10 @@ def _forward_single_sweep(self,
                     self.depth_channels + self.output_channels)].contiguous(),
                 self.voxel_num.cuda())
         if is_return_depth:
-            return feature_map.contiguous(), depth.float()
+            # final_depth has to be fp32, otherwise the depth
+            # loss will colapse during the traing process.
+            return feature_map.contiguous(
+            ), depth_feature[:, :self.depth_channels].softmax(dim=1)
         return feature_map.contiguous()
 
     def forward(self,

diff --git a/bevdepth/layers/backbones/bevstereo_lss_fpn.py b/bevdepth/layers/backbones/bevstereo_lss_fpn.py
@@ -878,7 +878,7 @@ def forward(self,
             mono_depth_all_sweeps.append(mono_depth)
             range_score_all_sweeps.append(range_score)
         depth_score_all_sweeps = list()
-
+        final_depth = None
         for ref_idx in range(num_sweeps):
             sensor2sensor_mats = list()
             for src_idx in range(num_sweeps):
@@ -955,7 +955,12 @@ def forward(self,
                     self.depth_downsample_net(stereo_depth)).softmax(
                         1, dtype=stereo_depth.dtype)
             depth_score_all_sweeps.append(depth_score)
-
+            if ref_idx == 0:
+                # final_depth has to be fp32, otherwise the
+                # depth loss will colapse during the traing process.
+                final_depth = (
+                    mono_depth_all_sweeps[ref_idx] +
+                    self.depth_downsample_net(stereo_depth)).softmax(1)
         key_frame_res = self._forward_single_sweep(
             0,
             context_all_sweeps[0].reshape(batch_size, num_cams,
@@ -985,6 +990,6 @@ def forward(self,
                 ret_feature_list.append(feature_map)
 
         if is_return_depth:
-            return torch.cat(ret_feature_list, 1), depth_score_all_sweeps[0]
+            return torch.cat(ret_feature_list, 1), final_depth
         else:
             return torch.cat(ret_feature_list, 1)