Esempio n. 1
0
    def forward(self, x, label=None, criterion=None):
        batch_size = x.size(0)
        idx, _ = knn(x, k=self.k)  # get the idx of knn in 3D space : b,n,k
        xyz = get_scorenet_input(
            x, k=self.k,
            idx=idx)  # ScoreNet input: 3D coord difference : b,6,n,k

        x = self.conv1(x)
        x = F.relu(self.bn1(x))
        ##################
        # replace the intermediate 3 MLP layers with PAConv:
        """CUDA implementation of PAConv: (presented in the supplementary material of the paper)"""
        """feature transformation:"""
        x = feat_trans_pointnet(point_input=x, kernel=self.matrice2,
                                m=self.m2)  # b,n,m1,o1
        score2 = self.scorenet2(xyz, calc_scores=self.calc_scores, bias=0)
        """assemble with scores:"""
        x = assemble_pointnet(score=score2,
                              point_input=x,
                              knn_idx=idx,
                              aggregate='sum')  # b,o1,n
        x = F.relu(self.bn2(x))

        x = feat_trans_pointnet(point_input=x, kernel=self.matrice3, m=self.m3)
        score3 = self.scorenet3(xyz, calc_scores=self.calc_scores, bias=0)
        x = assemble_pointnet(score=score3,
                              point_input=x,
                              knn_idx=idx,
                              aggregate='sum')
        x = F.relu(self.bn3(x))

        x = feat_trans_pointnet(point_input=x, kernel=self.matrice4, m=self.m4)
        score4 = self.scorenet4(xyz, calc_scores=self.calc_scores, bias=0)
        x = assemble_pointnet(score=score4,
                              point_input=x,
                              knn_idx=idx,
                              aggregate='sum')
        x = F.relu(self.bn4(x))
        ##################
        x = self.conv5(x)
        x = F.relu(self.bn5(x))

        x = F.adaptive_max_pool1d(x, 1).view(batch_size, -1)
        x = F.relu(self.bn6(self.linear1(x)))
        x = self.dp1(x)
        x = self.linear2(x)
        if criterion is not None:
            return x, criterion(x, label)
        else:
            return x
    def forward(self, x, norm_plt, cls_label, gt=None):
        B, C, N = x.size()
        idx, _ = knn(
            x, k=self.k
        )  # different with DGCNN, the knn search is only in 3D space
        xyz = get_scorenet_input(x, k=self.k, idx=idx)  # ScoreNet input
        #################
        # use MLP at the 1st layer, same with DGCNN
        x = get_graph_feature(x, k=self.k, idx=idx)
        x = x.permute(0, 3, 1, 2)  # b,2cin,n,k
        x = F.relu(self.conv1(x))
        x1 = x.max(dim=-1, keepdim=False)[0]
        #################
        # replace the last 4 DGCNN-EdgeConv with PAConv:
        """CUDA implementation of PAConv: (presented in the supplementary material of the paper)"""
        """feature transformation:"""
        x2, center2 = feat_trans_dgcnn(point_input=x1,
                                       kernel=self.matrice2,
                                       m=self.m2)
        score2 = self.scorenet2(xyz, calc_scores=self.calc_scores, bias=0)
        """assemble with scores:"""
        x = assemble_dgcnn(score=score2,
                           point_input=x2,
                           center_input=center2,
                           knn_idx=idx,
                           aggregate='sum')
        x2 = F.relu(self.bn2(x))

        x3, center3 = feat_trans_dgcnn(point_input=x2,
                                       kernel=self.matrice3,
                                       m=self.m3)
        score3 = self.scorenet3(xyz, calc_scores=self.calc_scores, bias=0)
        x = assemble_dgcnn(score=score3,
                           point_input=x3,
                           center_input=center3,
                           knn_idx=idx,
                           aggregate='sum')
        x3 = F.relu(self.bn3(x))

        x4, center4 = feat_trans_dgcnn(point_input=x3,
                                       kernel=self.matrice4,
                                       m=self.m4)
        score4 = self.scorenet4(xyz, calc_scores=self.calc_scores, bias=0)
        x = assemble_dgcnn(score=score4,
                           point_input=x4,
                           center_input=center4,
                           knn_idx=idx,
                           aggregate='sum')
        x4 = F.relu(self.bn4(x))

        x5, center5 = feat_trans_dgcnn(point_input=x4,
                                       kernel=self.matrice5,
                                       m=self.m5)
        score5 = self.scorenet5(xyz, calc_scores=self.calc_scores, bias=0)
        x = assemble_dgcnn(score=score5,
                           point_input=x5,
                           center_input=center5,
                           knn_idx=idx,
                           aggregate='sum')
        x5 = F.relu(self.bn5(x))
        ###############
        xx = torch.cat((x1, x2, x3, x4, x5), dim=1)

        xc = F.relu(self.convt(xx))
        xc = F.adaptive_max_pool1d(xc, 1).view(B, -1)

        cls_label = cls_label.view(B, 16, 1)
        cls_label = F.relu(self.convc(cls_label))
        cls = torch.cat((xc.view(B, 1024, 1), cls_label), dim=1)
        cls = cls.repeat(1, 1, N)  # B,1088,N

        x = torch.cat((xx, cls), dim=1)  # 1088+64*3
        x = F.relu(self.conv6(x))
        x = self.dp1(x)
        x = F.relu(self.conv7(x))
        x = self.dp2(x)
        x = F.relu(self.conv8(x))
        x = self.conv9(x)
        x = F.log_softmax(x, dim=1)
        x = x.permute(0, 2, 1)  # b,n,50

        if gt is not None:
            return x, F.nll_loss(x.contiguous().view(-1, self.num_part),
                                 gt.view(-1, 1)[:, 0])
        else:
            return x
Esempio n. 3
0
    def forward(self, x, label=None, criterion=None):
        B, C, N = x.size()
        idx, _ = knn(
            x, k=self.k
        )  # different with DGCNN, the knn search is only in 3D space
        xyz = get_scorenet_input(
            x, idx=idx, k=self.k
        )  # ScoreNet input: 3D coord difference concat with coord: b,6,n,k

        ##################
        # replace all the DGCNN-EdgeConv with PAConv:
        """CUDA implementation of PAConv: (presented in the supplementary material of the paper)"""
        """feature transformation:"""
        point1, center1 = feat_trans_dgcnn(point_input=x,
                                           kernel=self.matrice1,
                                           m=self.m1)  # b,n,m1,o1
        score1 = self.scorenet1(xyz, calc_scores=self.calc_scores, bias=0.5)
        """assemble with scores:"""
        point1 = assemble_dgcnn(score=score1,
                                point_input=point1,
                                center_input=center1,
                                knn_idx=idx,
                                aggregate='sum')  # b,o1,n
        point1 = F.relu(self.bn1(point1))

        point2, center2 = feat_trans_dgcnn(point_input=point1,
                                           kernel=self.matrice2,
                                           m=self.m2)
        score2 = self.scorenet2(xyz, calc_scores=self.calc_scores, bias=0.5)
        point2 = assemble_dgcnn(score=score2,
                                point_input=point2,
                                center_input=center2,
                                knn_idx=idx,
                                aggregate='sum')
        point2 = F.relu(self.bn2(point2))

        point3, center3 = feat_trans_dgcnn(point_input=point2,
                                           kernel=self.matrice3,
                                           m=self.m3)
        score3 = self.scorenet3(xyz, calc_scores=self.calc_scores, bias=0.5)
        point3 = assemble_dgcnn(score=score3,
                                point_input=point3,
                                center_input=center3,
                                knn_idx=idx,
                                aggregate='sum')
        point3 = F.relu(self.bn3(point3))

        point4, center4 = feat_trans_dgcnn(point_input=point3,
                                           kernel=self.matrice4,
                                           m=self.m4)
        score4 = self.scorenet4(xyz, calc_scores=self.calc_scores, bias=0.5)
        point4 = assemble_dgcnn(score=score4,
                                point_input=point4,
                                center_input=center4,
                                knn_idx=idx,
                                aggregate='sum')
        point4 = F.relu(self.bn4(point4))
        ##################

        point = torch.cat((point1, point2, point3, point4), dim=1)
        point = F.relu(self.conv5(point))
        point11 = F.adaptive_max_pool1d(point, 1).view(B, -1)
        point22 = F.adaptive_avg_pool1d(point, 1).view(B, -1)
        point = torch.cat((point11, point22), 1)

        point = F.relu(self.bn11(self.linear1(point)))
        point = self.dp1(point)
        point = F.relu(self.bn22(self.linear2(point)))
        point = self.dp2(point)
        point = self.linear3(point)

        if criterion is not None:
            return point, criterion(point, label)  # return output and loss
        else:
            return point