Esempio n. 1
0
 def __init__(self, input_dim, p=3, eps=1e-6):
     super(GeM, self).__init__()
     self.input_dim = input_dim
     # Same output number of channels as input number of channels
     self.output_dim = self.input_dim
     self.p = nn.Parameter(torch.ones(1) * p)
     self.eps = eps
     self.f = ME.MinkowskiGlobalAvgPooling()
Esempio n. 2
0
    def network_initialization(
        self, in_channel, out_channel, channels, embedding_channel, kernel_size, D=3,
    ):
        self.mlp1 = self.get_mlp_block(in_channel, channels[0])
        self.conv1 = self.get_conv_block(
            channels[0], channels[1], kernel_size=kernel_size, stride=1,
        )
        self.conv2 = self.get_conv_block(
            channels[1], channels[2], kernel_size=kernel_size, stride=2,
        )

        self.conv3 = self.get_conv_block(
            channels[2], channels[3], kernel_size=kernel_size, stride=2,
        )

        self.conv4 = self.get_conv_block(
            channels[3], channels[4], kernel_size=kernel_size, stride=2,
        )
        self.conv5 = nn.Sequential(
            self.get_conv_block(
                channels[1] + channels[2] + channels[3] + channels[4],
                embedding_channel // 4,
                kernel_size=3,
                stride=2,
            ),
            self.get_conv_block(
                embedding_channel // 4, embedding_channel // 2, kernel_size=3, stride=2,
            ),
            self.get_conv_block(
                embedding_channel // 2, embedding_channel, kernel_size=3, stride=2,
            ),
        )

        self.pool = ME.MinkowskiMaxPooling(kernel_size=3, stride=2, dimension=D)

        self.global_max_pool = ME.MinkowskiGlobalMaxPooling()
        self.global_avg_pool = ME.MinkowskiGlobalAvgPooling()

        self.final = nn.Sequential(
            self.get_mlp_block(embedding_channel * 2, 512),
            ME.MinkowskiDropout(),
            self.get_mlp_block(512, 512),
            ME.MinkowskiLinear(512, out_channel, bias=True),
        )
Esempio n. 3
0
 def __init__(self, p=3, eps=1e-6):
     super(GeM, self).__init__()
     self.p = nn.Parameter(torch.ones(1) * p)
     self.eps = eps
     self.f = ME.MinkowskiGlobalAvgPooling()
Esempio n. 4
0
 def __init__(self):
     super().__init__()
     self.f = ME.MinkowskiGlobalAvgPooling()
Esempio n. 5
0
 def __init__(self, input_dim):
     super().__init__()
     self.input_dim = input_dim
     # Same output number of channels as input number of channels
     self.output_dim = self.input_dim
     self.f = ME.MinkowskiGlobalAvgPooling()
Esempio n. 6
0
    # Normalize features and create a sparse tensor
    return features, coordinates


if __name__ == '__main__':
    config = parser.parse_args()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Define a model and load the weights
    feats = [3, 8, 16, 32, 64, 128]
    features, coordinates = generate_input_sparse_tensor(
        config.file_name,
        voxel_size=config.voxel_size,
        batch_size=config.batch_size)
    pool = ME.MinkowskiGlobalAvgPooling()

    # Measure time
    print('Forward')
    for feat in feats:
        timer = Timer()
        features = torch.rand(len(coordinates), feat).to(device)

        # Feed-forward pass and get the prediction
        for i in range(20):
            sinput = ME.SparseTensor(features,
                                     coordinates=coordinates,
                                     device=device)

            timer.tic()
            soutput = pool(sinput)
Esempio n. 7
0
    def network_initialization(self, in_channels, out_channels, D):
        # Setup net_metadata
        dilations = self.DILATIONS
        bn_momentum = 0.02

        def space_n_time_m(n, m):
            return n if D == 3 else [n, n, n, m]

        if D == 4:
            self.OUT_PIXEL_DIST = space_n_time_m(self.OUT_PIXEL_DIST, 1)

        # Output of the first conv concated to conv6
        conv1_kernel_size = 3
        self.inplanes = self.INIT_DIM
        self.conv0p1s1 = conv(in_channels,
                              self.inplanes,
                              kernel_size=space_n_time_m(conv1_kernel_size, 1),
                              stride=1,
                              dilation=1,
                              conv_type=self.NON_BLOCK_CONV_TYPE,
                              D=D)

        self.bn0 = get_norm(self.NORM_TYPE,
                            self.inplanes,
                            D,
                            bn_momentum=bn_momentum)

        self.conv1p1s2 = conv(self.inplanes,
                              self.inplanes,
                              kernel_size=space_n_time_m(2, 1),
                              stride=space_n_time_m(2, 1),
                              dilation=1,
                              conv_type=self.NON_BLOCK_CONV_TYPE,
                              D=D)
        self.bn1 = get_norm(self.NORM_TYPE,
                            self.inplanes,
                            D,
                            bn_momentum=bn_momentum)
        self.block1 = self._make_layer(self.BLOCK,
                                       self.PLANES[0],
                                       self.LAYERS[0],
                                       dilation=dilations[0],
                                       norm_type=self.NORM_TYPE,
                                       bn_momentum=bn_momentum)

        self.conv2p2s2 = conv(self.inplanes,
                              self.inplanes,
                              kernel_size=space_n_time_m(2, 1),
                              stride=space_n_time_m(2, 1),
                              dilation=1,
                              conv_type=self.NON_BLOCK_CONV_TYPE,
                              D=D)
        self.bn2 = get_norm(self.NORM_TYPE,
                            self.inplanes,
                            D,
                            bn_momentum=bn_momentum)
        self.block2 = self._make_layer(self.BLOCK,
                                       self.PLANES[1],
                                       self.LAYERS[1],
                                       dilation=dilations[1],
                                       norm_type=self.NORM_TYPE,
                                       bn_momentum=bn_momentum)

        self.conv3p4s2 = conv(self.inplanes,
                              self.inplanes,
                              kernel_size=space_n_time_m(2, 1),
                              stride=space_n_time_m(2, 1),
                              dilation=1,
                              conv_type=self.NON_BLOCK_CONV_TYPE,
                              D=D)
        self.bn3 = get_norm(self.NORM_TYPE,
                            self.inplanes,
                            D,
                            bn_momentum=bn_momentum)
        self.block3 = self._make_layer(self.BLOCK,
                                       self.PLANES[2],
                                       self.LAYERS[2],
                                       dilation=dilations[2],
                                       norm_type=self.NORM_TYPE,
                                       bn_momentum=bn_momentum)

        self.conv4p8s2 = conv(self.inplanes,
                              self.inplanes,
                              kernel_size=space_n_time_m(2, 1),
                              stride=space_n_time_m(2, 1),
                              dilation=1,
                              conv_type=self.NON_BLOCK_CONV_TYPE,
                              D=D)
        self.bn4 = get_norm(self.NORM_TYPE,
                            self.inplanes,
                            D,
                            bn_momentum=bn_momentum)
        self.block4 = self._make_layer(self.BLOCK,
                                       self.PLANES[3],
                                       self.LAYERS[3],
                                       dilation=dilations[3],
                                       norm_type=self.NORM_TYPE,
                                       bn_momentum=bn_momentum)
        self.convtr4p16s2 = conv_tr(self.inplanes,
                                    self.PLANES[4],
                                    kernel_size=space_n_time_m(2, 1),
                                    upsample_stride=space_n_time_m(2, 1),
                                    dilation=1,
                                    bias=False,
                                    conv_type=self.NON_BLOCK_CONV_TYPE,
                                    D=D)
        self.bntr4 = get_norm(self.NORM_TYPE,
                              self.PLANES[4],
                              D,
                              bn_momentum=bn_momentum)

        self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion
        self.block5 = self._make_layer(self.BLOCK,
                                       self.PLANES[4],
                                       self.LAYERS[4],
                                       dilation=dilations[4],
                                       norm_type=self.NORM_TYPE,
                                       bn_momentum=bn_momentum)
        self.convtr5p8s2 = conv_tr(self.inplanes,
                                   self.PLANES[5],
                                   kernel_size=space_n_time_m(2, 1),
                                   upsample_stride=space_n_time_m(2, 1),
                                   dilation=1,
                                   bias=False,
                                   conv_type=self.NON_BLOCK_CONV_TYPE,
                                   D=D)
        self.bntr5 = get_norm(self.NORM_TYPE,
                              self.PLANES[5],
                              D,
                              bn_momentum=bn_momentum)

        self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion
        self.block6 = self._make_layer(self.BLOCK,
                                       self.PLANES[5],
                                       self.LAYERS[5],
                                       dilation=dilations[5],
                                       norm_type=self.NORM_TYPE,
                                       bn_momentum=bn_momentum)
        self.convtr6p4s2 = conv_tr(self.inplanes,
                                   self.PLANES[6],
                                   kernel_size=space_n_time_m(2, 1),
                                   upsample_stride=space_n_time_m(2, 1),
                                   dilation=1,
                                   bias=False,
                                   conv_type=self.NON_BLOCK_CONV_TYPE,
                                   D=D)
        self.bntr6 = get_norm(self.NORM_TYPE,
                              self.PLANES[6],
                              D,
                              bn_momentum=bn_momentum)

        self.inplanes = self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion
        self.block7 = self._make_layer(self.BLOCK,
                                       self.PLANES[6],
                                       self.LAYERS[6],
                                       dilation=dilations[6],
                                       norm_type=self.NORM_TYPE,
                                       bn_momentum=bn_momentum)
        self.convtr7p2s2 = conv_tr(self.inplanes,
                                   self.PLANES[7],
                                   kernel_size=space_n_time_m(2, 1),
                                   upsample_stride=space_n_time_m(2, 1),
                                   dilation=1,
                                   bias=False,
                                   conv_type=self.NON_BLOCK_CONV_TYPE,
                                   D=D)
        self.bntr7 = get_norm(self.NORM_TYPE,
                              self.PLANES[7],
                              D,
                              bn_momentum=bn_momentum)

        self.inplanes = self.PLANES[7] + self.INIT_DIM
        self.block8 = self._make_layer(self.BLOCK,
                                       self.PLANES[7],
                                       self.LAYERS[7],
                                       dilation=dilations[7],
                                       norm_type=self.NORM_TYPE,
                                       bn_momentum=bn_momentum)

        self.all_feat_names = [
            "en0",
            "en1",
            "en2",
            "en3",
            "en4",
            "plane4",
            "plane5",
            "plane6",
            "plane7",
        ]

        self.relu = MinkowskiReLU(inplace=True)
        self.maxpool = ME.MinkowskiGlobalMaxPooling()
        self.avgpool = ME.MinkowskiGlobalAvgPooling()
        if self.use_mlp:
            self.head = SMLP(self.mlp_dim)
    def __init__(self,
                 in_channel,
                 out_channel,
                 num_class,
                 embedding_channel=1024,
                 dimension=3):
        ME.MinkowskiNetwork.__init__(self, dimension)
        # The normal channel for Modelnet is 3, for scannet is 6, for scanobjnn is 0
        normal_channel = 3
        # in_channel = normal_channel+3 # normal ch + xyz
        self.normal_channel = normal_channel
        self.input_mlp = nn.Sequential(
            ME.MinkowskiConvolution(in_channel, 32, kernel_size=1,
                                    dimension=3), ME.MinkowskiBatchNorm(32),
            ME.MinkowskiReLU(),
            ME.MinkowskiConvolution(32, 32, kernel_size=1, dimension=3),
            ME.MinkowskiBatchNorm(32))

        self.in_dims = [32, 64, 128, 256]
        self.out_dims = [64, 128, 256, 512]
        self.neighbor_ks = [16, 32, 64, 16, 16]
        # self.neighbor_ks = [8, 8, 8, 8, 8]

        self.PTBlock0 = PTBlock(in_dim=self.in_dims[0],
                                n_sample=self.neighbor_ks[0])

        self.TDLayer1 = TDLayer(input_dim=self.in_dims[0],
                                out_dim=self.out_dims[0])
        self.PTBlock1 = PTBlock(in_dim=self.out_dims[0],
                                n_sample=self.neighbor_ks[1])

        self.TDLayer2 = TDLayer(input_dim=self.in_dims[1],
                                out_dim=self.out_dims[1])
        self.PTBlock2 = PTBlock(in_dim=self.out_dims[1],
                                n_sample=self.neighbor_ks[1])

        self.TDLayer3 = TDLayer(input_dim=self.in_dims[2],
                                out_dim=self.out_dims[2])
        self.PTBlock3 = PTBlock(in_dim=self.out_dims[2],
                                n_sample=self.neighbor_ks[2])

        self.TDLayer4 = TDLayer(input_dim=self.in_dims[3],
                                out_dim=self.out_dims[3])
        self.PTBlock4 = PTBlock(in_dim=self.out_dims[3],
                                n_sample=self.neighbor_ks[4])

        self.middle_linear = ME.MinkowskiConvolution(self.out_dims[3],
                                                     self.out_dims[3],
                                                     kernel_size=1,
                                                     dimension=3)
        self.PTBlock_middle = PTBlock(in_dim=self.out_dims[3],
                                      n_sample=self.neighbor_ks[4])

        self.TULayer5 = TULayer(input_dim=self.out_dims[3],
                                out_dim=self.in_dims[3])
        self.PTBlock5 = PTBlock(in_dim=self.in_dims[3],
                                n_sample=self.neighbor_ks[4])

        self.TULayer6 = TULayer(input_dim=self.out_dims[2],
                                out_dim=self.in_dims[2])
        self.PTBlock6 = PTBlock(in_dim=self.in_dims[2],
                                n_sample=self.neighbor_ks[3])

        self.TULayer7 = TULayer(input_dim=self.out_dims[1],
                                out_dim=self.in_dims[1])
        self.PTBlock7 = PTBlock(in_dim=self.in_dims[1],
                                n_sample=self.neighbor_ks[2])

        self.TULayer8 = TULayer(input_dim=self.out_dims[0],
                                out_dim=self.in_dims[0])
        self.PTBlock8 = PTBlock(in_dim=self.in_dims[0],
                                n_sample=self.neighbor_ks[1])

        self.fc = nn.Sequential(
            # ME.MinkowskiLinear(32, 32),
            ME.MinkowskiLinear(self.out_dims[3], 32),
            ME.MinkowskiDropout(0.4),
            ME.MinkowskiLinear(32, num_class))

        self.global_avg_pool = ME.MinkowskiGlobalAvgPooling()