Issues related to Point Cloud Training

[DoranLyong]

Thanks for sharing your nice work. To understand the application of point cloud training, I have been following the instructions in https://github.com/AILab-CVC/UniRepLKNet/tree/main/Point

However, I met a problem in the network flow, and I hope to solve it. The following figure shows the error message:

When I debugged the code, I found where it happened in p2p_adaptor_lkv.py:

class P2P(nn.Module):
    def __init__(self, cfg, is_test=False):
        super().__init__()
        self.cfg = cfg

       [...]

    # @autocast(dtype=torch.float16)
    def forward(self, pc, original_pc):

        # ######################### multi view start ##############################
        original_pc = torch.repeat_interleave(original_pc, self.num_views, dim=0)
        pc = self.point_transform(pc)
        # ########################## multi view end ###############################
        img = self.enc(original_pc, pc)  
        print(f"img: {img.shape}")
        out = self.base_model(img) 
        print(f"out: {out.shape}")

        out = self.cls_head(out)
        # ###################### cls head end #############################
        return out

It seems that the img variable is a tensor shaped in (Batch, 3, 224, 224) like input for ImageNet-1K. Then, self.base_model gets the img as input and returns the out tensor shaped in (batch, 1000).

I think this makes the error. From your code, self.cls_head instance is built by 'pooling_mlp' in models.layers.head.py

class PoolingClsHead(nn.Module):
    def __init__(self, cfg):  # , mid_channels, dropout_ratio
        super().__init__()

      [...]
        
    def forward(self, feats: torch.Tensor):
        print(feats.shape)
        B, C, H, W = feats.shape
        feats = feats.reshape(B // self.num_views, self.num_views * C, H, W)
        feats = self.adaptive_pooling(feats).squeeze()
        feats = self.norm(feats)
        feats = self.cls_head(feats)
        return feats

But, the required shape is different with the out tensor shape.
What should I change?

[DoranLyong]

I think that self.base_model(img) should return a tensor in (batch, dim, 7, 7) shape as input for self.cls_head(out). Am I right?