mirror of
https://github.com/wassname/Pointnet2_PyTorch.git
synced 2026-07-01 16:40:08 +08:00
Updates and some refactoring
This commit is contained in:
erikwijmans
+78
-63
@@ -43,22 +43,34 @@ class Pointnet2SSG(nn.Module):
|
||||
|
||||
self.initial_dropout = RandomDropout(0.4)
|
||||
|
||||
self.SA_module0 = PointnetSAModule(
|
||||
npoint=1024,
|
||||
radius=0.1,
|
||||
nsample=32,
|
||||
mlp=[input_channels, 32, 32, 64])
|
||||
self.SA_module1 = PointnetSAModule(
|
||||
npoint=256, radius=0.2, nsample=32, mlp=[64 + 3, 64, 64, 128])
|
||||
self.SA_module2 = PointnetSAModule(
|
||||
npoint=64, radius=0.4, nsample=32, mlp=[128 + 3, 128, 128, 256])
|
||||
self.SA_module3 = PointnetSAModule(
|
||||
npoint=16, radius=0.8, nsample=32, mlp=[256 + 3, 256, 256, 512])
|
||||
self.SA_modules = nn.ModuleList()
|
||||
self.SA_modules.append(
|
||||
PointnetSAModule(
|
||||
npoint=1024,
|
||||
radius=0.1,
|
||||
nsample=32,
|
||||
mlp=[input_channels, 32, 32, 64]))
|
||||
self.SA_modules.append(
|
||||
PointnetSAModule(
|
||||
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]))
|
||||
self.SA_modules.append(
|
||||
PointnetSAModule(
|
||||
npoint=16,
|
||||
radius=0.8,
|
||||
nsample=32,
|
||||
mlp=[256 + 3, 256, 256, 512]))
|
||||
|
||||
self.FP_module0 = PointnetFPModule(mlp=[512 + 256, 256, 256])
|
||||
self.FP_module1 = PointnetFPModule(mlp=[256 + 128, 256, 256])
|
||||
self.FP_module2 = PointnetFPModule(mlp=[256 + 64, 256, 128])
|
||||
self.FP_module3 = PointnetFPModule(mlp=[128 + 6, 128, 128, 128])
|
||||
self.FP_modules = nn.ModuleList()
|
||||
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(),
|
||||
@@ -72,18 +84,17 @@ class Pointnet2SSG(nn.Module):
|
||||
l0_xyz = self.initial_dropout(xyz)
|
||||
l0_points = None
|
||||
|
||||
l1_xyz, l1_points = self.SA_module0(l0_xyz, l0_points)
|
||||
l2_xyz, l2_points = self.SA_module1(l1_xyz, l1_points)
|
||||
l3_xyz, l3_points = self.SA_module2(l2_xyz, l2_points)
|
||||
l4_xyz, l4_points = self.SA_module3(l3_xyz, l3_points)
|
||||
l_xyz, l_points = [l0_xyz], [l0_points]
|
||||
for i in range(len(self.SA_modules)):
|
||||
li_xyz, li_points = self.SA_modules[i](l_xyz[i], l_points[i])
|
||||
l_xyz.append(li_xyz)
|
||||
l_points.append(li_points)
|
||||
|
||||
l3_points = self.FP_module0(l3_xyz, l4_xyz, l3_points, l4_points)
|
||||
l2_points = self.FP_module1(l2_xyz, l3_xyz, l2_points, l3_points)
|
||||
l1_points = self.FP_module2(l1_xyz, l2_xyz, l1_points, l2_points)
|
||||
l0_points = self.FP_module3(l0_xyz, l1_xyz, l0_points,
|
||||
l1_points).transpose(1, 2)
|
||||
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])
|
||||
|
||||
return self.FC_layer(l0_points).transpose(1, 2).contiguous()
|
||||
return self.FC_layer(l_points[0].transpose(1, 2)).transpose(1, 2).contiguous()
|
||||
|
||||
|
||||
class Pointnet2MSG(nn.Module):
|
||||
@@ -93,43 +104,50 @@ class Pointnet2MSG(nn.Module):
|
||||
self.initial_dropout = RandomDropout(0.95, inplace=True)
|
||||
self.initial_dropout = None
|
||||
|
||||
self.SA_modules = nn.ModuleList()
|
||||
c_in = input_channels
|
||||
self.SA_module0 = PointnetSAModuleMSG(
|
||||
npoint=1024,
|
||||
radii=[0.05, 0.1],
|
||||
nsamples=[16, 32],
|
||||
mlps=[[c_in, 16, 16, 32], [c_in, 32, 32, 64]])
|
||||
self.SA_modules.append(
|
||||
PointnetSAModuleMSG(
|
||||
npoint=1024,
|
||||
radii=[0.05, 0.1],
|
||||
nsamples=[16, 32],
|
||||
mlps=[[c_in, 16, 16, 32], [c_in, 32, 32, 64]]))
|
||||
c_out_0 = 32 + 64
|
||||
|
||||
c_in = c_out_0 + 3
|
||||
self.SA_module1 = PointnetSAModuleMSG(
|
||||
npoint=256,
|
||||
radii=[0.1, 0.2],
|
||||
nsamples=[16, 32],
|
||||
mlps=[[c_in, 64, 64, 128], [c_in, 64, 96, 128]])
|
||||
self.SA_modules.append(
|
||||
PointnetSAModuleMSG(
|
||||
npoint=256,
|
||||
radii=[0.1, 0.2],
|
||||
nsamples=[16, 32],
|
||||
mlps=[[c_in, 64, 64, 128], [c_in, 64, 96, 128]]))
|
||||
c_out_1 = 128 + 128
|
||||
|
||||
c_in = c_out_1 + 3
|
||||
self.SA_module2 = PointnetSAModuleMSG(
|
||||
npoint=64,
|
||||
radii=[0.2, 0.4],
|
||||
nsamples=[16, 32],
|
||||
mlps=[[c_in, 128, 196, 256], [c_in, 128, 196, 256]])
|
||||
self.SA_modules.append(
|
||||
PointnetSAModuleMSG(
|
||||
npoint=64,
|
||||
radii=[0.2, 0.4],
|
||||
nsamples=[16, 32],
|
||||
mlps=[[c_in, 128, 196, 256], [c_in, 128, 196, 256]]))
|
||||
c_out_2 = 256 + 256
|
||||
|
||||
c_in = c_out_2 + 3
|
||||
self.SA_module3 = PointnetSAModuleMSG(
|
||||
npoint=16,
|
||||
radii=[0.4, 0.8],
|
||||
nsamples=[16, 32],
|
||||
mlps=[[c_in, 256, 256, 512], [c_in, 256, 384, 512]])
|
||||
self.SA_modules.append(
|
||||
PointnetSAModuleMSG(
|
||||
npoint=16,
|
||||
radii=[0.4, 0.8],
|
||||
nsamples=[16, 32],
|
||||
mlps=[[c_in, 256, 256, 512], [c_in, 256, 384, 512]]))
|
||||
c_out_3 = 512 + 512
|
||||
|
||||
self.FP_module3 = PointnetFPModule(mlp=[c_out_3 + c_out_2, 512, 512])
|
||||
self.FP_module2 = PointnetFPModule(mlp=[512 + c_out_1, 512, 512])
|
||||
self.FP_module1 = PointnetFPModule(mlp=[512 + c_out_0, 256, 256])
|
||||
self.FP_module0 = PointnetFPModule(
|
||||
mlp=[256 + input_channels - 3, 128, 128])
|
||||
self.FP_modules = nn.ModuleList()
|
||||
self.FP_modules.append(
|
||||
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]))
|
||||
|
||||
self.FC_layer = nn.Sequential(
|
||||
pt_utils.Conv1d(128, 128, bn=True), nn.Dropout(),
|
||||
@@ -142,20 +160,17 @@ class Pointnet2MSG(nn.Module):
|
||||
elif self.initial_dropout is not None:
|
||||
xyz = self.initial_dropout(xyz)
|
||||
|
||||
l0_xyz, l0_points = xyz, points
|
||||
l_xyz, l_points = [xyz], [points]
|
||||
for i in range(len(self.SA_modules)):
|
||||
li_xyz, li_points = self.SA_modules[i](l_xyz[i], l_points[i])
|
||||
l_xyz.append(li_xyz)
|
||||
l_points.append(li_points)
|
||||
|
||||
l1_xyz, l1_points = self.SA_module0(l0_xyz, l0_points)
|
||||
l2_xyz, l2_points = self.SA_module1(l1_xyz, l1_points)
|
||||
l3_xyz, l3_points = self.SA_module2(l2_xyz, l2_points)
|
||||
l4_xyz, l4_points = self.SA_module3(l3_xyz, l3_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])
|
||||
|
||||
l3_points = self.FP_module3(l3_xyz, l4_xyz, l3_points, l4_points)
|
||||
l2_points = self.FP_module2(l2_xyz, l3_xyz, l2_points, l3_points)
|
||||
l1_points = self.FP_module1(l1_xyz, l2_xyz, l1_points, l2_points)
|
||||
l0_points = self.FP_module0(l0_xyz, l1_xyz, l0_points,
|
||||
l1_points).transpose(1, 2)
|
||||
|
||||
return self.FC_layer(l0_points).transpose(1, 2).contiguous()
|
||||
return self.FC_layer(l_points[0].transpose(1, 2)).transpose(1, 2).contiguous()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -170,7 +185,7 @@ if __name__ == "__main__":
|
||||
model = Pointnet2MSG(3)
|
||||
model.cuda()
|
||||
|
||||
optimizer = optim.Adam(model.parameters(), lr=1e-5)
|
||||
optimizer = optim.Adam(model.parameters(), lr=1e-2)
|
||||
|
||||
model_fn = model_fn_decorator(nn.CrossEntropyLoss())
|
||||
for _ in range(20):
|
||||
|
||||
Reference in New Issue
Block a user