|
|
import torch.nn as nn |
|
|
import torchsparse.nn as spnn |
|
|
from torchsparse.point_tensor import PointTensor |
|
|
|
|
|
from lib.spvcnn_utils import * |
|
|
__all__ = ['SPVCNN_CLASSIFICATION'] |
|
|
|
|
|
|
|
|
|
|
|
class BasicConvolutionBlock(nn.Module): |
|
|
def __init__(self, inc, outc, ks=3, stride=1, dilation=1): |
|
|
super().__init__() |
|
|
self.net = nn.Sequential( |
|
|
spnn.Conv3d(inc, |
|
|
outc, |
|
|
kernel_size=ks, |
|
|
dilation=dilation, |
|
|
stride=stride), |
|
|
spnn.BatchNorm(outc), |
|
|
spnn.ReLU(True)) |
|
|
|
|
|
def forward(self, x): |
|
|
out = self.net(x) |
|
|
return out |
|
|
|
|
|
|
|
|
class BasicDeconvolutionBlock(nn.Module): |
|
|
def __init__(self, inc, outc, ks=3, stride=1): |
|
|
super().__init__() |
|
|
self.net = nn.Sequential( |
|
|
spnn.Conv3d(inc, |
|
|
outc, |
|
|
kernel_size=ks, |
|
|
stride=stride, |
|
|
transpose=True), |
|
|
spnn.BatchNorm(outc), |
|
|
spnn.ReLU(True)) |
|
|
|
|
|
def forward(self, x): |
|
|
return self.net(x) |
|
|
|
|
|
|
|
|
class ResidualBlock(nn.Module): |
|
|
def __init__(self, inc, outc, ks=3, stride=1, dilation=1): |
|
|
super().__init__() |
|
|
self.net = nn.Sequential( |
|
|
spnn.Conv3d(inc, |
|
|
outc, |
|
|
kernel_size=ks, |
|
|
dilation=dilation, |
|
|
stride=stride), spnn.BatchNorm(outc), |
|
|
spnn.ReLU(True), |
|
|
spnn.Conv3d(outc, |
|
|
outc, |
|
|
kernel_size=ks, |
|
|
dilation=dilation, |
|
|
stride=1), |
|
|
spnn.BatchNorm(outc) |
|
|
) |
|
|
|
|
|
self.downsample = nn.Sequential() if (inc == outc and stride == 1) else \ |
|
|
nn.Sequential( |
|
|
spnn.Conv3d(inc, outc, kernel_size=1, dilation=1, stride=stride), |
|
|
spnn.BatchNorm(outc) |
|
|
) |
|
|
|
|
|
self.relu = spnn.ReLU(True) |
|
|
|
|
|
def forward(self, x): |
|
|
out = self.relu(self.net(x) + self.downsample(x)) |
|
|
return out |
|
|
|
|
|
|
|
|
class SPVCNN_CLASSIFICATION(nn.Module): |
|
|
def __init__(self, **kwargs): |
|
|
super().__init__() |
|
|
|
|
|
cr = kwargs.get('cr', 1.0) |
|
|
cs = [32, 32, 64, 128, 256, 256, 128, 96, 96] |
|
|
cs = [int(cr * x) for x in cs] |
|
|
|
|
|
if 'pres' in kwargs and 'vres' in kwargs: |
|
|
self.pres = kwargs['pres'] |
|
|
self.vres = kwargs['vres'] |
|
|
|
|
|
self.stem = nn.Sequential( |
|
|
spnn.Conv3d(kwargs['input_channel'], cs[0], kernel_size=3, stride=1), |
|
|
spnn.BatchNorm(cs[0]), |
|
|
spnn.ReLU(True), |
|
|
spnn.Conv3d(cs[0], cs[0], kernel_size=3, stride=1), |
|
|
spnn.BatchNorm(cs[0]), |
|
|
spnn.ReLU(True)) |
|
|
|
|
|
self.stage1 = nn.Sequential( |
|
|
BasicConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1), |
|
|
ResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1), |
|
|
ResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1), |
|
|
) |
|
|
|
|
|
self.stage2 = nn.Sequential( |
|
|
BasicConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1), |
|
|
ResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1), |
|
|
ResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1), |
|
|
) |
|
|
|
|
|
self.stage3 = nn.Sequential( |
|
|
BasicConvolutionBlock(cs[2], cs[2], ks=2, stride=2, dilation=1), |
|
|
ResidualBlock(cs[2], cs[3], ks=3, stride=1, dilation=1), |
|
|
ResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1), |
|
|
) |
|
|
|
|
|
self.stage4 = nn.Sequential( |
|
|
BasicConvolutionBlock(cs[3], cs[3], ks=2, stride=2, dilation=1), |
|
|
ResidualBlock(cs[3], cs[4], ks=3, stride=1, dilation=1), |
|
|
ResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1), |
|
|
) |
|
|
self.avg_pool = spnn.GlobalAveragePooling() |
|
|
self.classifier = nn.Sequential(nn.Linear(cs[4], kwargs['num_classes'])) |
|
|
self.point_transforms = nn.ModuleList([ |
|
|
nn.Sequential( |
|
|
nn.Linear(cs[0], cs[4]), |
|
|
nn.BatchNorm1d(cs[4]), |
|
|
nn.ReLU(True), |
|
|
), |
|
|
]) |
|
|
|
|
|
self.weight_initialization() |
|
|
self.dropout = nn.Dropout(0.3, True) |
|
|
|
|
|
def weight_initialization(self): |
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.BatchNorm1d): |
|
|
nn.init.constant_(m.weight, 1) |
|
|
nn.init.constant_(m.bias, 0) |
|
|
|
|
|
def forward(self, x): |
|
|
|
|
|
z = PointTensor(x.F, x.C.float()) |
|
|
|
|
|
x0 = initial_voxelize(z, self.pres, self.vres) |
|
|
|
|
|
x0 = self.stem(x0) |
|
|
z0 = voxel_to_point(x0, z, nearest=False) |
|
|
z0.F = z0.F |
|
|
|
|
|
x1 = point_to_voxel(x0, z0) |
|
|
x1 = self.stage1(x1) |
|
|
x2 = self.stage2(x1) |
|
|
x3 = self.stage3(x2) |
|
|
x4 = self.stage4(x3) |
|
|
z1 = voxel_to_point(x4, z0) |
|
|
z1.F = z1.F + self.point_transforms[0](z0.F) |
|
|
y1 = point_to_voxel(x4, z1) |
|
|
pool = self.avg_pool(y1) |
|
|
out = self.classifier(pool) |
|
|
|
|
|
|
|
|
return out |
|
|
|
|
|
|
|
|
|
