wassname/Pointnet2_PyTorch
1
0
Fork 0
mirror of https://github.com/wassname/Pointnet2_PyTorch.git synced 2026-07-24 12:50:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

Some faster/better kernels. Tensors with points are now kept in (b, c, ...) format as this is easier for pytorch

Browse Source
This commit is contained in:
erikwijmans
2018-02-10 20:32:52 -05:00
parent 8bce353da4
commit 65a127f3d2
19 changed files with 372 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/Pointnet2Cls.py
+16 -13
View File
@@ -40,7 +40,7 @@ def model_fn_decorator(criterion):
class Pointnet2SSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
def __init__(self, num_classes, input_channels=3):
super().__init__()
self.SA_modules = nn.ModuleList()
@@ -54,13 +54,10 @@ class Pointnet2SSG(nn.Module):
)
self.SA_modules.append(
PointnetSAModule(
npoint=128,
radius=0.4,
nsample=64,
mlp=[128 + 3, 128, 128, 256]
npoint=128, radius=0.4, nsample=64, mlp=[128, 128, 128, 256]
)
)
self.SA_modules.append(PointnetSAModule(mlp=[256 + 3, 256, 512, 1024]))
self.SA_modules.append(PointnetSAModule(mlp=[256, 256, 512, 1024]))
self.FC_layer = nn.Sequential(
pt_utils.FC(1024, 512, bn=True),
@@ -71,15 +68,18 @@ class Pointnet2SSG(nn.Module):
)
def forward(self, xyz, points=None):
xyz = xyz.contiguous()
points = points.transpose(1, 2
).contiguous() if points is not None else None
for module in self.SA_modules:
xyz, points = module(xyz, points)
return self.FC_layer(points.squeeze(1))
return self.FC_layer(points.squeeze(-1))
class Pointnet2MSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
def __init__(self, num_classes, input_channels=3):
super().__init__()
self.SA_modules = nn.ModuleList()
@@ -93,7 +93,7 @@ class Pointnet2MSG(nn.Module):
)
)
input_channels = 64 + 128 + 128 + 3
input_channels = 64 + 128 + 128
self.SA_modules.append(
PointnetSAModuleMSG(
npoint=128,
@@ -104,7 +104,7 @@ class Pointnet2MSG(nn.Module):
)
)
self.SA_modules.append(
PointnetSAModule(mlp=[128 + 256 + 256 + 3, 256, 512, 1024])
PointnetSAModule(mlp=[128 + 256 + 256, 256, 512, 1024])
)
self.FC_layer = nn.Sequential(
@@ -116,10 +116,13 @@ class Pointnet2MSG(nn.Module):
)
def forward(self, xyz, points=None):
xyz = xyz.contiguous()
points = points.transpose(1, 2
).contiguous() if points is not None else None
for module in self.SA_modules:
xyz, points = module(xyz, points)
return self.FC_layer(points.squeeze(1))
return self.FC_layer(points.squeeze(-1))
if __name__ == "__main__":
@@ -129,9 +132,9 @@ if __name__ == "__main__":
import torch.autograd.profiler as profiler
B = 2
N = 2048
inputs = torch.randn(B, N, 9).cuda()
inputs = torch.randn(B, N, 6).cuda()
labels = torch.from_numpy(np.random.randint(0, 3, size=B)).cuda()
model = Pointnet2MSG(3)
model = Pointnet2MSG(3, input_channels=3)
model.cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-2)