Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
sparse_block_utils.py 4.15 KB
Newer Older
wuyuefeng's avatar
sparse_unetv2 backbone  
wuyuefeng committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143 
from torch import nn

import mmdet3d.ops.spconv as spconv
from mmdet.ops import build_norm_layer


def conv3x3(in_planes, out_planes, stride=1, indice_key=None):
    """3x3 submanifold sparse convolution with padding.

    Args:
        in_planes (int): the number of input channels
        out_planes (int): the number of output channels
        stride (int): the stride of convolution
        indice_key (str): the indice key used for sparse tensor

    Returns:
        spconv.conv.SubMConv3d: 3x3 submanifold sparse convolution ops
    """
    return spconv.SubMConv3d(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=1,
        bias=False,
        indice_key=indice_key)


def conv1x1(in_planes, out_planes, stride=1, indice_key=None):
    """1x1 submanifold sparse convolution with padding.

    Args:
        in_planes (int): the number of input channels
        out_planes (int): the number of output channels
        stride (int): the stride of convolution
        indice_key (str): the indice key used for sparse tensor

    Returns:
        spconv.conv.SubMConv3d: 1x1 submanifold sparse convolution ops
    """
    return spconv.SubMConv3d(
        in_planes,
        out_planes,
        kernel_size=1,
        stride=stride,
        padding=1,
        bias=False,
        indice_key=indice_key)


class SparseBasicBlock(spconv.SparseModule):
    expansion = 1

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 downsample=None,
                 indice_key=None,
                 norm_cfg=None):
        """Sparse basic block for PartA^2.

        Sparse basic block implemented with submanifold sparse convolution.
        """
        super(SparseBasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride, indice_key=indice_key)
        norm_name1, norm_layer1 = build_norm_layer(norm_cfg, planes)
        self.bn1 = norm_layer1
        self.relu = nn.ReLU()
        self.conv2 = conv3x3(planes, planes, indice_key=indice_key)
        norm_name2, norm_layer2 = build_norm_layer(norm_cfg, planes)
        self.bn2 = norm_layer2
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x.features

        assert x.features.dim() == 2, 'x.features.dim()=%d' % x.features.dim()

        out = self.conv1(x)
        out.features = self.bn1(out.features)
        out.features = self.relu(out.features)

        out = self.conv2(out)
        out.features = self.bn2(out.features)

        if self.downsample is not None:
            identity = self.downsample(x)

        out.features += identity
        out.features = self.relu(out.features)

        return out


class SparseBottleneck(spconv.SparseModule):
    expansion = 4

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 downsample=None,
                 indice_key=None,
                 norm_fn=None):
        """Sparse bottleneck block for PartA^2.

        Bottleneck block implemented with submanifold sparse convolution.
        """
        super(SparseBottleneck, self).__init__()
        self.conv1 = conv1x1(inplanes, planes, indice_key=indice_key)
        self.bn1 = norm_fn(planes)
        self.conv2 = conv3x3(planes, planes, stride, indice_key=indice_key)
        self.bn2 = norm_fn(planes)
        self.conv3 = conv1x1(
            planes, planes * self.expansion, indice_key=indice_key)
        self.bn3 = norm_fn(planes * self.expansion)
        self.relu = nn.ReLU()
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x.features

        out = self.conv1(x)
        out.features = self.bn1(out.features)
        out.features = self.relu(out.features)

        out = self.conv2(out)
        out.features = self.bn2(out.features)
        out.features = self.relu(out.features)

        out = self.conv3(out)
        out.features = self.bn3(out.features)

        if self.downsample is not None:
            identity = self.downsample(x)

        out.features += identity
        out.features = self.relu(out.features)

        return out