move spconv_utils.py to pcdet/utils/, update find_all_spconv_keys()

thri5ha · Nov 23, 2021 · cddcf9b · cddcf9b
1 parent 511244e
commit cddcf9b
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/pcdet/models/backbones_3d/spconv_backbone.py b/pcdet/models/backbones_3d/spconv_backbone.py
@@ -2,7 +2,7 @@
 
 import torch.nn as nn
 
-from ...spconv_utils import replace_feature, spconv
+from ...utils.spconv_utils import replace_feature, spconv
 
 
 def post_act_block(in_channels, out_channels, kernel_size, indice_key=None, stride=1, padding=0,

diff --git a/pcdet/models/backbones_3d/spconv_unet.py b/pcdet/models/backbones_3d/spconv_unet.py
@@ -3,7 +3,7 @@
 import torch
 import torch.nn as nn
 
-from ...spconv_utils import replace_feature, spconv
+from ...utils.spconv_utils import replace_feature, spconv
 from ...utils import common_utils
 from .spconv_backbone import post_act_block
 

diff --git a/pcdet/models/roi_heads/partA2_head.py b/pcdet/models/roi_heads/partA2_head.py
@@ -3,7 +3,7 @@
 import torch.nn as nn
 
 from ...ops.roiaware_pool3d import roiaware_pool3d_utils
-from ...spconv_utils import spconv
+from ...utils.spconv_utils import spconv
 from .roi_head_template import RoIHeadTemplate
 
 

diff --git a/pcdet/spconv_utils.py → pcdet/utils/spconv_utils.py b/pcdet/spconv_utils.py → pcdet/utils/spconv_utils.py
@@ -16,7 +16,7 @@ def find_all_spconv_keys(model: nn.Module, prefix="") -> Set[str]:
     for name, child in model.named_children():
         new_prefix = f"{prefix}.{name}" if prefix != "" else name
 
-        if isinstance(child, (spconv.SubMConv3d, spconv.SparseConv3d, spconv.SparseInverseConv3d)):
+        if isinstance(child, spconv.conv.SparseConvolution):
             new_prefix = f"{new_prefix}.weight"
             found_keys.add(new_prefix)