This commit is contained in:
erikwijmans
2018-01-06 12:13:52 -05:00
parent 7e746ba72a
commit 5a5adc2b77
20 changed files with 650 additions and 494 deletions
+21 -8
View File
@@ -39,6 +39,7 @@ def model_fn_decorator(criterion):
class Pointnet2SSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
super().__init__()
@@ -48,13 +49,17 @@ class Pointnet2SSG(nn.Module):
npoint=512,
radius=0.2,
nsample=64,
mlp=[input_channels, 64, 64, 128]))
mlp=[input_channels, 64, 64, 128]
)
)
self.SA_modules.append(
PointnetSAModule(
npoint=128,
radius=0.4,
nsample=64,
mlp=[128 + 3, 128, 128, 256]))
mlp=[128 + 3, 128, 128, 256]
)
)
self.SA_modules.append(PointnetSAModule(mlp=[256 + 3, 256, 512, 1024]))
self.FC_layer = nn.Sequential(
@@ -62,7 +67,8 @@ class Pointnet2SSG(nn.Module):
nn.Dropout(p=0.5),
pt_utils.FC(512, 256, bn=True),
nn.Dropout(p=0.5),
pt_utils.FC(256, num_classes, activation=None))
pt_utils.FC(256, num_classes, activation=None)
)
def forward(self, xyz, points=None):
for module in self.SA_modules:
@@ -72,6 +78,7 @@ class Pointnet2SSG(nn.Module):
class Pointnet2MSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
super().__init__()
@@ -83,7 +90,9 @@ class Pointnet2MSG(nn.Module):
nsamples=[32, 64, 128],
mlps=[[input_channels, 32, 32,
64], [input_channels, 64, 64, 128],
[input_channels, 64, 96, 128]]))
[input_channels, 64, 96, 128]]
)
)
input_channels = 64 + 128 + 128 + 3
self.SA_modules.append(
@@ -92,17 +101,21 @@ class Pointnet2MSG(nn.Module):
radii=[0.2, 0.4, 0.8],
nsamples=[16, 32, 64],
mlps=[[input_channels, 64, 64,
128], [input_channels, 128, 128, 256],
[input_channels, 128, 128, 256]]))
128], [input_channels, 128, 128, 256],
[input_channels, 128, 128, 256]]
)
)
self.SA_modules.append(
PointnetSAModule(mlp=[128 + 256 + 256 + 3, 256, 512, 1024]))
PointnetSAModule(mlp=[128 + 256 + 256 + 3, 256, 512, 1024])
)
self.FC_layer = nn.Sequential(
pt_utils.FC(1024, 512, bn=True),
nn.Dropout(p=0.5),
pt_utils.FC(512, 256, bn=True),
nn.Dropout(p=0.5),
pt_utils.FC(256, num_classes, activation=None))
pt_utils.FC(256, num_classes, activation=None)
)
def forward(self, xyz, points=None):
for module in self.SA_modules:
+49 -25
View File
@@ -38,6 +38,7 @@ def model_fn_decorator(criterion):
class Pointnet2SSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
super().__init__()
@@ -49,32 +50,37 @@ class Pointnet2SSG(nn.Module):
npoint=1024,
radius=0.1,
nsample=32,
mlp=[input_channels, 32, 32, 64]))
mlp=[input_channels, 32, 32, 64]
)
)
self.SA_modules.append(
PointnetSAModule(
npoint=256, radius=0.2, nsample=32, mlp=[64 + 3, 64, 64, 128]))
npoint=256, radius=0.2, nsample=32, mlp=[64 + 3, 64, 64, 128]
)
)
self.SA_modules.append(
PointnetSAModule(
npoint=64,
radius=0.4,
nsample=32,
mlp=[128 + 3, 128, 128, 256]))
npoint=64, radius=0.4, nsample=32, mlp=[128 + 3, 128, 128, 256]
)
)
self.SA_modules.append(
PointnetSAModule(
npoint=16,
radius=0.8,
nsample=32,
mlp=[256 + 3, 256, 256, 512]))
npoint=16, radius=0.8, nsample=32, mlp=[256 + 3, 256, 256, 512]
)
)
self.FP_modules = nn.ModuleList()
self.FP_modules.append(PointnetFPModule(mlp=[128 + input_channels - 3, 128, 128, 128]))
self.FP_modules.append(
PointnetFPModule(mlp=[128 + input_channels - 3, 128, 128, 128])
)
self.FP_modules.append(PointnetFPModule(mlp=[256 + 64, 256, 128]))
self.FP_modules.append(PointnetFPModule(mlp=[256 + 128, 256, 256]))
self.FP_modules.append(PointnetFPModule(mlp=[512 + 256, 256, 256]))
self.FC_layer = nn.Sequential(
pt_utils.Conv1d(128, 128, bn=True), nn.Dropout(),
pt_utils.Conv1d(128, num_classes, activation=None))
pt_utils.Conv1d(128, num_classes, activation=None)
)
def forward(self, xyz, points=None):
if points is not None:
@@ -91,13 +97,17 @@ class Pointnet2SSG(nn.Module):
l_points.append(li_points)
for i in range(-1, -(len(self.FP_modules + 1) - 1), -1):
l_points[i - 1] = self.FP_modules[i](l_xyz[i - 1], l_xyz[i],
l_points[i - 1], l_points[i])
l_points[i - 1] = self.FP_modules[i](
l_xyz[i - 1], l_xyz[i], l_points[i - 1], l_points[i]
)
return self.FC_layer(l_points[0].transpose(1, 2)).transpose(1, 2).contiguous()
return self.FC_layer(l_points[0].transpose(1,
2)).transpose(1,
2).contiguous()
class Pointnet2MSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
super().__init__()
@@ -111,7 +121,9 @@ class Pointnet2MSG(nn.Module):
npoint=1024,
radii=[0.05, 0.1],
nsamples=[16, 32],
mlps=[[c_in, 16, 16, 32], [c_in, 32, 32, 64]]))
mlps=[[c_in, 16, 16, 32], [c_in, 32, 32, 64]]
)
)
c_out_0 = 32 + 64
c_in = c_out_0 + 3
@@ -120,7 +132,9 @@ class Pointnet2MSG(nn.Module):
npoint=256,
radii=[0.1, 0.2],
nsamples=[16, 32],
mlps=[[c_in, 64, 64, 128], [c_in, 64, 96, 128]]))
mlps=[[c_in, 64, 64, 128], [c_in, 64, 96, 128]]
)
)
c_out_1 = 128 + 128
c_in = c_out_1 + 3
@@ -129,7 +143,9 @@ class Pointnet2MSG(nn.Module):
npoint=64,
radii=[0.2, 0.4],
nsamples=[16, 32],
mlps=[[c_in, 128, 196, 256], [c_in, 128, 196, 256]]))
mlps=[[c_in, 128, 196, 256], [c_in, 128, 196, 256]]
)
)
c_out_2 = 256 + 256
c_in = c_out_2 + 3
@@ -138,20 +154,25 @@ class Pointnet2MSG(nn.Module):
npoint=16,
radii=[0.4, 0.8],
nsamples=[16, 32],
mlps=[[c_in, 256, 256, 512], [c_in, 256, 384, 512]]))
mlps=[[c_in, 256, 256, 512], [c_in, 256, 384, 512]]
)
)
c_out_3 = 512 + 512
self.FP_modules = nn.ModuleList()
self.FP_modules.append(
PointnetFPModule(mlp=[256 + input_channels - 3, 128, 128]))
PointnetFPModule(mlp=[256 + input_channels - 3, 128, 128])
)
self.FP_modules.append(PointnetFPModule(mlp=[512 + c_out_0, 256, 256]))
self.FP_modules.append(PointnetFPModule(mlp=[512 + c_out_1, 512, 512]))
self.FP_modules.append(
PointnetFPModule(mlp=[c_out_3 + c_out_2, 512, 512]))
PointnetFPModule(mlp=[c_out_3 + c_out_2, 512, 512])
)
self.FC_layer = nn.Sequential(
pt_utils.Conv1d(128, 128, bn=True), nn.Dropout(),
pt_utils.Conv1d(128, num_classes, activation=None))
pt_utils.Conv1d(128, num_classes, activation=None)
)
def forward(self, xyz, points=None):
if points is not None and self.initial_dropout is not None:
@@ -167,10 +188,13 @@ class Pointnet2MSG(nn.Module):
l_points.append(li_points)
for i in range(-1, -(len(self.FP_modules) + 1), -1):
l_points[i - 1] = self.FP_modules[i](l_xyz[i - 1], l_xyz[i],
l_points[i - 1], l_points[i])
l_points[i - 1] = self.FP_modules[i](
l_xyz[i - 1], l_xyz[i], l_points[i - 1], l_points[i]
)
return self.FC_layer(l_points[0].transpose(1, 2)).transpose(1, 2).contiguous()
return self.FC_layer(l_points[0].transpose(1,
2)).transpose(1,
2).contiguous()
if __name__ == "__main__":