示例#1
0
    def network_initialization(self, in_channels, out_channels, D):
        self.inplanes = self.init_dim
        self.conv1 = ME.MinkowskiConvolution(
            in_channels, self.inplanes, kernel_size=5, stride=2, dimension=D)

        self.bn1 = ME.MinkowskiBatchNorm(self.inplanes)
        self.relu = ME.MinkowskiReLU(inplace=True)

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

        self.layer1 = self._make_layer(
            self.block, self.planes[0], self.layers[0], stride=2)
        self.layer2 = self._make_layer(
            self.block, self.planes[1], self.layers[1], stride=2)
        self.layer3 = self._make_layer(
            self.block, self.planes[2], self.layers[2], stride=2)
        self.layer4 = self._make_layer(
            self.block, self.planes[3], self.layers[3], stride=2)

        self.conv5 = ME.MinkowskiConvolution(
            self.inplanes, self.inplanes, kernel_size=3, stride=3, dimension=D)
        self.bn5 = ME.MinkowskiBatchNorm(self.inplanes)

        self.glob_avg = ME.MinkowskiGlobalMaxPooling()

        self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True)
示例#2
0
    def network_initialization(self, in_channels, out_channels, D):

        self.inplanes = self.INIT_DIM
        self.conv1 = nn.Sequential(
            ME.MinkowskiConvolution(in_channels,
                                    self.inplanes,
                                    kernel_size=3,
                                    stride=2,
                                    dimension=D),
            ME.MinkowskiBatchNorm(self.inplanes),
            ME.MinkowskiReLU(inplace=True),
            ME.MinkowskiMaxPooling(kernel_size=2, stride=2, dimension=D),
        )

        self.layer1 = self._make_layer(self.BLOCK,
                                       self.PLANES[0],
                                       self.LAYERS[0],
                                       stride=2)
        self.layer2 = self._make_layer(self.BLOCK,
                                       self.PLANES[1],
                                       self.LAYERS[1],
                                       stride=2)
        self.layer3 = self._make_layer(self.BLOCK,
                                       self.PLANES[2],
                                       self.LAYERS[2],
                                       stride=2)
        self.layer4 = self._make_layer(self.BLOCK,
                                       self.PLANES[3],
                                       self.LAYERS[3],
                                       stride=2)

        self.conv5 = nn.Sequential(
            ME.MinkowskiDropout(),
            ME.MinkowskiConvolution(self.inplanes,
                                    self.inplanes,
                                    kernel_size=3,
                                    stride=3,
                                    dimension=D),
            ME.MinkowskiBatchNorm(self.inplanes),
            ME.MinkowskiGELU(),
        )

        self.glob_pool = ME.MinkowskiGlobalMaxPooling()

        self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True)
示例#3
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),
        )
示例#4
0
    def __init__(self, fc_dim=64, in_channels=3, out_channels=40, D=3):
        super(SparseResNet18, self).__init__()
        original_resnet = ResNet18(in_channels=in_channels,
                                   out_channels=out_channels,
                                   D=D)
        checkpoint_path = os.path.join(
            os.path.split(os.path.join(os.getcwd(), __file__))[0],
            'resnet18_pretrained.pth')

        #Load modelnet40 pretrained weights
        checkpoint = torch.load(checkpoint_path)
        original_resnet.load_state_dict(checkpoint['state_dict'])

        self.features = nn.Sequential(*list(original_resnet.children())[:-4])

        self.fc = ME.MinkowskiConvolution(in_channels=512,
                                          out_channels=fc_dim,
                                          kernel_size=3,
                                          stride=1,
                                          dimension=D)
        self.glob_avg = ME.MinkowskiGlobalMaxPooling(dimension=D)
示例#5
0
    def __init__(self,
                 in_channel,
                 out_channel,
                 embedding_channel=1024,
                 dimension=3):
        ME.MinkowskiNetwork.__init__(self, dimension)
        self.conv1 = nn.Sequential(
            ME.MinkowskiLinear(3, 64, bias=False),
            ME.MinkowskiBatchNorm(64),
            ME.MinkowskiReLU(),
        )
        self.conv2 = nn.Sequential(
            ME.MinkowskiLinear(64, 64, bias=False),
            ME.MinkowskiBatchNorm(64),
            ME.MinkowskiReLU(),
        )
        self.conv3 = nn.Sequential(
            ME.MinkowskiLinear(64, 64, bias=False),
            ME.MinkowskiBatchNorm(64),
            ME.MinkowskiReLU(),
        )
        self.conv4 = nn.Sequential(
            ME.MinkowskiLinear(64, 128, bias=False),
            ME.MinkowskiBatchNorm(128),
            ME.MinkowskiReLU(),
        )
        self.conv5 = nn.Sequential(
            ME.MinkowskiLinear(128, embedding_channel, bias=False),
            ME.MinkowskiBatchNorm(embedding_channel),
            ME.MinkowskiReLU(),
        )
        self.max_pool = ME.MinkowskiGlobalMaxPooling()

        self.linear1 = nn.Sequential(
            ME.MinkowskiLinear(embedding_channel, 512, bias=False),
            ME.MinkowskiBatchNorm(512),
            ME.MinkowskiReLU(),
        )
        self.dp1 = ME.MinkowskiDropout()
        self.linear2 = ME.MinkowskiLinear(512, out_channel, bias=True)
示例#6
0
 def __init__(self):
     super().__init__()
     self.f = ME.MinkowskiGlobalMaxPooling()
示例#7
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.MinkowskiGlobalMaxPooling()
  def network_initialization(self, in_channels, out_channels, config, D):
    # Setup net_metadata
    dilations = self.DILATIONS
    bn_momentum = config['bn_momentum']

    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
    self.inplanes = self.INIT_DIM
    self.conv0p1s1 = conv(
        in_channels,
        self.inplanes,
        kernel_size=space_n_time_m(config['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.relu = MinkowskiReLU(inplace=True)
    # add a classification head here
    self.clf_glob_avg = ME.MinkowskiGlobalPooling(dimension=D)
    self.clf_glob_max=ME.MinkowskiGlobalMaxPooling(dimension=D)
    self.clf_conv0 = conv(
        256,
        512,
        kernel_size=3,
        stride=2,
        dilation=1,
        conv_type=self.NON_BLOCK_CONV_TYPE,
        D=D)
    self.clf_bn0 = get_norm(self.NORM_TYPE, 512, D, bn_momentum=bn_momentum)
    self.clf_conv1 = conv(
        512,
        512,
        kernel_size=3,
        stride=2,
        dilation=1,
        conv_type=self.NON_BLOCK_CONV_TYPE,
        D=D)
    self.clf_bn1 = get_norm(self.NORM_TYPE, 512, D, bn_momentum=bn_momentum)
    self.clf_conv2 = conv(
        512,
        config['clf_num_labels'],
        kernel_size=1,
        stride=1,
        dilation=1,
        conv_type=self.NON_BLOCK_CONV_TYPE,
        D=D)
示例#9
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)