mirror of
https://github.com/wassname/Pointnet2_PyTorch.git
synced 2026-06-27 16:00:07 +08:00
Updates and some refactoring
This commit is contained in:
erikwijmans
+103
-66
@@ -1,96 +1,133 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.autograd import Variable
|
||||
|
||||
import os, sys
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.append(BASE_DIR)
|
||||
sys.path.append(os.path.join(BASE_DIR, "..", "utils"))
|
||||
sys.path.append(os.path.join(BASE_DIR, "../utils"))
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.autograd import Variable
|
||||
import pytorch_utils as pt_utils
|
||||
from TransformNets import TransformNet, TranslationNet
|
||||
from pointnet2_modules import PointnetSAModule, PointnetFPModule, PointnetSAModuleMSG
|
||||
from pointnet2_utils import RandomDropout
|
||||
from collections import namedtuple
|
||||
|
||||
|
||||
def model_fn_decorator(criterion):
|
||||
transform_reg = 1e-3
|
||||
ModelReturn = namedtuple("ModelReturn", ['preds', 'loss', 'acc'])
|
||||
|
||||
def ortho_loss(matrix):
|
||||
return torch.dist(
|
||||
matrix.bmm(matrix.transpose(1, 2)),
|
||||
Variable(
|
||||
torch.eye(matrix.size(1), matrix.size(2)).type(
|
||||
torch.cuda.FloatTensor)))
|
||||
def model_fn(model, data, epoch=0, eval=False):
|
||||
inputs, labels = data
|
||||
inputs = Variable(inputs.cuda(async=True), volatile=eval)
|
||||
labels = Variable(labels.cuda(async=True), volatile=eval)
|
||||
|
||||
def wrapped(model, inputs, labels):
|
||||
labels = labels.squeeze()
|
||||
preds, end_points = model(inputs)
|
||||
xyz = inputs[..., :3]
|
||||
if inputs.size(2) > 3:
|
||||
points = inputs[..., 3:]
|
||||
else:
|
||||
points = None
|
||||
|
||||
transform_loss = 0.0
|
||||
for _, T in end_points.items():
|
||||
transform_loss += ortho_loss(T)
|
||||
preds = model(xyz, points)
|
||||
labels = labels.view(-1)
|
||||
loss = criterion(preds, labels)
|
||||
|
||||
preds_loss = criterion(preds, labels)
|
||||
loss = preds_loss + transform_reg * transform_loss
|
||||
_, classes = torch.max(preds.data, -1)
|
||||
acc = (classes == labels.data).sum() / labels.numel()
|
||||
|
||||
_, classes = torch.max(preds, 1)
|
||||
acc = (classes == labels).sum()
|
||||
return ModelReturn(preds, loss, {"acc": acc})
|
||||
|
||||
return preds, loss, acc.data[0]
|
||||
|
||||
return wrapped
|
||||
return model_fn
|
||||
|
||||
|
||||
class PointnetCls(nn.Module):
|
||||
def __init__(self):
|
||||
class Pointnet2SSG(nn.Module):
|
||||
def __init__(self, num_classes, input_channels=9):
|
||||
super().__init__()
|
||||
|
||||
self.translation_net = TranslationNet()
|
||||
self.t_net = TransformNet(1, 3, 3, scale=False)
|
||||
self.f_net = TransformNet(64, 1, 64, scale=False)
|
||||
self.SA_modules = nn.ModuleList()
|
||||
self.SA_modules.append(
|
||||
PointnetSAModule(
|
||||
npoint=512,
|
||||
radius=0.2,
|
||||
nsample=64,
|
||||
mlp=[input_channels, 64, 64, 128]))
|
||||
self.SA_modules.append(
|
||||
PointnetSAModule(
|
||||
npoint=128,
|
||||
radius=0.4,
|
||||
nsample=64,
|
||||
mlp=[128 + 3, 128, 128, 256]))
|
||||
self.SA_modules.append(PointnetSAModule(mlp=[256 + 3, 256, 512, 1024]))
|
||||
|
||||
self.input_mlp = nn.Sequential(
|
||||
pt_utils.Conv2d(1, 64, [1, 3], bn=True),
|
||||
pt_utils.Conv2d(64, 64, bn=True))
|
||||
|
||||
self.second_mlp = pt_utils.SharedMLP([64, 64, 128, 1024], bn=True)
|
||||
|
||||
self.final_mlp = nn.Sequential(
|
||||
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(0.3), pt_utils.FC(256, 40, activation=None))
|
||||
nn.Dropout(p=0.5),
|
||||
pt_utils.FC(256, num_classes, activation=None))
|
||||
|
||||
def forward(self, points: torch.Tensor):
|
||||
batch_size, n_points, _ = points.size()
|
||||
end_points = {}
|
||||
def forward(self, xyz, points=None):
|
||||
for module in self.SA_modules:
|
||||
xyz, points = module(xyz, points)
|
||||
|
||||
points = points + self.translation_net(points).unsqueeze(1)
|
||||
points, transform = self.apply_transform(
|
||||
points, *self.t_net(points.unsqueeze(1)))
|
||||
|
||||
points = self.input_mlp(points.unsqueeze(1))
|
||||
|
||||
points, transform = self.apply_transform(points.squeeze().transpose(
|
||||
1, 2), *self.f_net(points))
|
||||
end_points['trans2'] = transform
|
||||
|
||||
points = F.max_pool2d(
|
||||
self.second_mlp(points.transpose(1, 2).unsqueeze(-1)),
|
||||
kernel_size=[n_points, 1])
|
||||
return self.final_mlp(points.view(-1, 1024)), end_points
|
||||
return self.FC_layer(points.squeeze(1))
|
||||
|
||||
|
||||
def apply_transform(self, points, rotation, scale=None):
|
||||
points = points @ rotation
|
||||
if scale is not None:
|
||||
points = points * scale.contiguous().view(-1, 1, 1).repeat(
|
||||
1, points.size(1), points.size(2))
|
||||
class Pointnet2MSG(nn.Module):
|
||||
def __init__(self, num_classes, input_channels=9):
|
||||
super().__init__()
|
||||
|
||||
return points, rotation
|
||||
self.SA_modules = nn.ModuleList()
|
||||
self.SA_modules.append(
|
||||
PointnetSAModuleMSG(
|
||||
npoint=512,
|
||||
radii=[0.1, 0.2, 0.4],
|
||||
nsamples=[32, 64, 128],
|
||||
mlps=[[input_channels, 32, 32,
|
||||
64], [input_channels, 64, 64, 128],
|
||||
[input_channels, 64, 96, 128]]))
|
||||
|
||||
input_channels = 64 + 128 + 128 + 3
|
||||
self.SA_modules.append(
|
||||
PointnetSAModuleMSG(
|
||||
npoint=128,
|
||||
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]]))
|
||||
self.SA_modules.append(
|
||||
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))
|
||||
|
||||
def forward(self, xyz, points=None):
|
||||
for module in self.SA_modules:
|
||||
xyz, points = module(xyz, points)
|
||||
|
||||
return self.FC_layer(points.squeeze(1))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from torch.autograd import Variable
|
||||
model = PointnetCls()
|
||||
data = Variable(torch.randn(2, 10, 3))
|
||||
print(model(data))
|
||||
import numpy as np
|
||||
import torch.optim as optim
|
||||
B = 2
|
||||
N = 32
|
||||
inputs = torch.randn(B, N, 9).cuda()
|
||||
labels = torch.from_numpy(np.random.randint(0, 3, size=B)).cuda()
|
||||
model = Pointnet2SSG(3)
|
||||
model.cuda()
|
||||
|
||||
optimizer = optim.Adam(model.parameters(), lr=1e-2)
|
||||
|
||||
model_fn = model_fn_decorator(nn.CrossEntropyLoss())
|
||||
for _ in range(20):
|
||||
optimizer.zero_grad()
|
||||
_, loss, _ = model_fn(model, (inputs, labels))
|
||||
loss.backward()
|
||||
print(loss.data[0])
|
||||
optimizer.step()
|
||||
|
||||
+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):
|
||||
|
||||
@@ -1,75 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.autograd import Variable
|
||||
import torch.nn.functional as F
|
||||
|
||||
import os, sys
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.append(BASE_DIR)
|
||||
|
||||
import pytorch_utils as pt_utils
|
||||
|
||||
|
||||
class TransformNet(nn.Module):
|
||||
def __init__(self, in_size, channels, K, scale=False):
|
||||
super().__init__()
|
||||
self.K, self.scale = K, scale
|
||||
|
||||
self.convs = nn.Sequential()
|
||||
self.convs.add_module('conv0',
|
||||
pt_utils.Conv2d(
|
||||
in_size, 64, kernel_size=[1, channels], bn=True))
|
||||
self.convs.add_module('rest',
|
||||
pt_utils.SharedMLP([64, 128, 1024], bn=True))
|
||||
|
||||
self.fc = nn.Sequential(
|
||||
pt_utils.FC(1024, 512, bn=True), pt_utils.FC(512, 256, bn=True))
|
||||
|
||||
outsize = K * K
|
||||
if scale:
|
||||
outsize += 1
|
||||
|
||||
self.final_W = nn.Parameter(torch.FloatTensor(256, outsize))
|
||||
self.final_b = nn.Parameter(torch.FloatTensor(outsize))
|
||||
|
||||
self.init_weights()
|
||||
|
||||
def forward(self, X):
|
||||
X = self.convs(X)
|
||||
X = F.adaptive_max_pool2d(X, [1, 1])
|
||||
X = self.fc(X.view(-1, 1024))
|
||||
X = X @ self.final_W + self.final_b
|
||||
|
||||
rotation = X[:, 0:self.K * self.K].contiguous().view(
|
||||
-1, self.K, self.K)
|
||||
|
||||
if not self.scale:
|
||||
return rotation, None
|
||||
|
||||
scale = X[:, -1].contiguous()
|
||||
|
||||
return rotation, scale
|
||||
|
||||
def init_weights(self):
|
||||
torch.nn.init.constant(self.final_W, 0)
|
||||
self.final_b.data[:self.K * self.K] = (torch.eye(
|
||||
self.K, self.K) + 1e-1 * torch.randn(self.K, self.K)).view(-1)
|
||||
if self.scale:
|
||||
self.final_b.data[-1] = 1.0
|
||||
|
||||
|
||||
class TranslationNet(nn.Module):
|
||||
def forward(self, X):
|
||||
return -torch.mean(X, dim=1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from torch.autograd import Variable
|
||||
net = TransformNet(5, 1, 3, True)
|
||||
net.init_weights()
|
||||
data = Variable(torch.FloatTensor(1, 5, 10, 1))
|
||||
print(net(data))
|
||||
|
||||
net = TranslationNet(5, 1, 3)
|
||||
net.init_weights()
|
||||
print(net(data))
|
||||
+5
-1
@@ -1 +1,5 @@
|
||||
from .Pointnet2SemSeg import Pointnet2MSG, Pointnet2SSG
|
||||
from .Pointnet2SemSeg import Pointnet2MSG as Pointnet2SemMSG
|
||||
from .Pointnet2SemSeg import Pointnet2SSG as Pointnet2SemSSG
|
||||
|
||||
from .Pointnet2Cls import Pointnet2MSG as Pointnet2ClsMSG
|
||||
from .Pointnet2Cls import Pointnet2SSG as Pointnet2ClsSSG
|
||||
|
||||
Reference in New Issue
Block a user