From 1cb3d440fa6fefc444d30e8f9f0153aba175ef90 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E5=BC=BA?= Date: Sat, 6 Jun 2020 19:38:48 +0800 Subject: [PATCH 1/3] =?UTF-8?q?=E6=B7=BB=E5=8A=A0Integration=20Layer?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 36 ++++++++++++ models/pointnet2_structure_point_net.py | 76 ++++++++++++++++++------- pointnet2/utils/pointnet2_modules.py | 13 ++--- pointnet2/utils/pointnet2_utils.py | 27 +++++---- requirements.txt | 6 +- test/test_structure_points.py | 24 +++----- train/train_structure_points.py | 59 ++++++++++++------- 7 files changed, 156 insertions(+), 85 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..8215354 --- /dev/null +++ b/.gitignore @@ -0,0 +1,36 @@ +# compilation and distribution +__pycache__ +_ext +*.pyc +*.so +*.egg-info/ +build/ +dist/ +# pytorch/python/numpy formats +*.pth +*.pkl +*.npy + +# ipython/jupyter notebooks +**/.ipynb_checkpoints/ + +# Editor temporaries +*.swn +*.swo +*.swp +*~ + +# Pycharm editor settings +.idea + +# vscode editor settings +.vscode + +# MacOS +.DS_Store + +# project dirs +data +output +/.idea/ +log \ No newline at end of file diff --git a/models/pointnet2_structure_point_net.py b/models/pointnet2_structure_point_net.py index ceb5c13..ff508ba 100644 --- a/models/pointnet2_structure_point_net.py +++ b/models/pointnet2_structure_point_net.py @@ -22,7 +22,8 @@ def __init__(self): self.consistent_loss = None self.cd_loss = None - def forward(self, gt_points, structure_points, transed_gt_points=None, transed_structure_points=None, trans_func_list=None): + def forward(self, gt_points, structure_points, transed_gt_points=None, transed_structure_points=None, + trans_func_list=None): gt_points = gt_points.cuda() structure_points = structure_points.cuda() @@ -39,8 +40,9 @@ def forward(self, gt_points, structure_points, transed_gt_points=None, transed_s transed_structure_points = transed_structure_points.cuda() transed_gt_points = transed_gt_points.cuda() trans_num = transed_structure_points.shape[0] - self.cd_loss = self.cd_loss + self.cd_loss_fun(transed_structure_points.view(trans_num * batch_size, stpts_num, dim), - transed_gt_points.view(trans_num * batch_size, pts_num, dim)) + self.cd_loss = self.cd_loss + self.cd_loss_fun( + transed_structure_points.view(trans_num * batch_size, stpts_num, dim), + transed_gt_points.view(trans_num * batch_size, pts_num, dim)) self.consistent_loss = None for i in range(0, trans_num): tmp_structure_points = trans_func_list[i](structure_points) @@ -53,7 +55,6 @@ def forward(self, gt_points, structure_points, transed_gt_points=None, transed_s self.consistent_loss = self.consistent_loss + tmp_consistent_loss self.consistent_loss = self.consistent_loss / trans_num * 1000 - self.cd_loss = self.cd_loss / (trans_num + 1) self.loss = self.cd_loss @@ -69,6 +70,27 @@ def get_consistent_loss(self): return self.consistent_loss +class IntegrationLayer(nn.Module): + def __init__(self, in_channels, out_channels, out=False, kernel_size=1, dropout=0.2): + super().__init__() + self.out = out + self.dropout_conv_bn_layer = nn.Sequential( + nn.Dropout(dropout), + nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size), + nn.BatchNorm1d(num_features=out_channels), + ) + self.relu = nn.ReLU() + self.softmax = nn.Softmax(dim=2) + + def forward(self, x): + x = self.dropout_conv_bn_layer(x) + if self.out: + x = self.relu(x) + else: + x = self.softmax(x) + return x + + class Pointnet2StructurePointNet(nn.Module): def __init__(self, num_structure_points, input_channels=3, use_xyz=True): @@ -110,6 +132,8 @@ def __init__(self, num_structure_points, input_channels=3, use_xyz=True): conv1d_stpts_prob_modules = [] if num_structure_points <= 128 + 256 + 256: + # conv1d_stpts_prob_modules.append( + # IntegrationLayer(in_channels=128 + 256 + 256, out_channels=512)) conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) conv1d_stpts_prob_modules.append(nn.Conv1d(in_channels=128 + 256 + 256, out_channels=512, kernel_size=1)) conv1d_stpts_prob_modules.append(nn.BatchNorm1d(512)) @@ -117,18 +141,25 @@ def __init__(self, num_structure_points, input_channels=3, use_xyz=True): in_channels = 512 while in_channels >= self.num_structure_points * 2: out_channels = int(in_channels / 2) + # conv1d_stpts_prob_modules.append(IntegrationLayer(in_channels, out_channels)) conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) - conv1d_stpts_prob_modules.append(nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1)) + conv1d_stpts_prob_modules.append( + nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1)) conv1d_stpts_prob_modules.append(nn.BatchNorm1d(out_channels)) conv1d_stpts_prob_modules.append(nn.ReLU()) in_channels = out_channels - conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) - conv1d_stpts_prob_modules.append(nn.Conv1d(in_channels=in_channels, out_channels=self.num_structure_points, kernel_size=1)) - - conv1d_stpts_prob_modules.append(nn.BatchNorm1d(self.num_structure_points)) - conv1d_stpts_prob_modules.append(nn.Softmax(dim=2)) + conv1d_stpts_prob_modules.append(IntegrationLayer(in_channels, self.num_structure_points, out=True)) + # conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) + # conv1d_stpts_prob_modules.append( + # nn.Conv1d(in_channels=in_channels, out_channels=self.num_structure_points, kernel_size=1)) + # + # conv1d_stpts_prob_modules.append(nn.BatchNorm1d(self.num_structure_points)) + # conv1d_stpts_prob_modules.append(nn.Softmax(dim=2)) else: + + # conv1d_stpts_prob_modules.append( + # IntegrationLayer(in_channels=128 + 256 + 256, out_channels=1024, kernel_size=1)) conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) conv1d_stpts_prob_modules.append(nn.Conv1d(in_channels=128 + 256 + 256, out_channels=1024, kernel_size=1)) conv1d_stpts_prob_modules.append(nn.BatchNorm1d(1024)) @@ -137,22 +168,26 @@ def __init__(self, num_structure_points, input_channels=3, use_xyz=True): in_channels = 1024 while in_channels <= self.num_structure_points / 2: out_channels = int(in_channels * 2) + # conv1d_stpts_prob_modules.append( + # IntegrationLayer(in_channels=in_channels, out_channels=out_channels, kernel_size=1)) conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) conv1d_stpts_prob_modules.append(nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1)) conv1d_stpts_prob_modules.append(nn.BatchNorm1d(out_channels)) conv1d_stpts_prob_modules.append(nn.ReLU()) in_channels = out_channels - conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) - conv1d_stpts_prob_modules.append(nn.Conv1d(in_channels=in_channels, out_channels=self.num_structure_points, kernel_size=1)) + # conv1d_stpts_prob_modules.append(IntegrationLayer(in_channels, self.num_structure_points, True)) + conv1d_stpts_prob_modules.append(nn.dropout(0.2)) + conv1d_stpts_prob_modules.append( + nn.conv1d(in_channels=in_channels, out_channels=self.num_structure_points, kernel_size=1)) - conv1d_stpts_prob_modules.append(nn.BatchNorm1d(self.num_structure_points)) + conv1d_stpts_prob_modules.append(nn.batchnorm1d(self.num_structure_points)) conv1d_stpts_prob_modules.append(nn.Softmax(dim=2)) self.conv1d_stpts_prob = nn.Sequential(*conv1d_stpts_prob_modules) def _break_up_pc(self, pc): - xyz = pc[..., 0:3].contiguous() + xyz = pc[..., 0:3].contiguous() # 取点 features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None return xyz, features @@ -171,18 +206,17 @@ def forward(self, pointcloud, return_weighted_feature=False): self.stpts_prob_map = self.conv1d_stpts_prob(features) weighted_xyz = torch.sum(self.stpts_prob_map[:, :, :, None] * xyz[:, None, :, :], dim=2) - if return_weighted_feature: - weighted_features = torch.sum(self.stpts_prob_map[:, None, :, :] * features[:, :, None, :], dim=3) if return_weighted_feature: + weighted_features = torch.sum(self.stpts_prob_map[:, None, :, :] * features[:, :, None, :], dim=3) return weighted_xyz, weighted_features else: return weighted_xyz - - - - - +if __name__ == "__main__": + pointclouds = torch.rand(10, 1000, 3).cuda() + model = Pointnet2StructurePointNet(15).cuda() + output = model(pointclouds) + print(output.shape) diff --git a/pointnet2/utils/pointnet2_modules.py b/pointnet2/utils/pointnet2_modules.py index e863e9c..2b2e384 100644 --- a/pointnet2/utils/pointnet2_modules.py +++ b/pointnet2/utils/pointnet2_modules.py @@ -50,8 +50,8 @@ def forward(self, xyz, features=None, return_critical_index=False): pointnet2_utils.gather_operation( xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint) ) - .transpose(1, 2) - .contiguous() + .transpose(1, 2) + .contiguous() if self.npoint is not None else None ) @@ -118,8 +118,6 @@ def __init__(self, npoint, radii, nsamples, mlps, bn=True, use_xyz=True): self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn)) - - class PointnetSAModuleMSGRRI(_PointnetSAModuleBase): r"""Pointnet set abstrction layer with multiscale grouping with rigious rotation invariant @@ -154,11 +152,9 @@ def __init__(self, npoint, radii, nsamples, mlps, bn=True): ) mlp_spec = mlps[i] - self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn)) - class PointnetSAModule(PointnetSAModuleMSG): r"""Pointnet set abstrction layer @@ -177,7 +173,7 @@ class PointnetSAModule(PointnetSAModuleMSG): """ def __init__( - self, mlp, npoint=None, radius=None, nsample=None, bn=True, use_xyz=True + self, mlp, npoint=None, radius=None, nsample=None, bn=True, use_xyz=True ): # type: (PointnetSAModule, List[int], int, float, int, bool, bool) -> None super(PointnetSAModule, self).__init__( @@ -216,7 +212,7 @@ def forward(self, unknown, known, unknow_feats, known_feats): known : torch.Tensor (B, m, 3) tensor of the xyz positions of the known features unknow_feats : torch.Tensor - (B, C1, n) tensor of the features to be propigated to + (B, C1, n) tensor of the features to be propagated to known_feats : torch.Tensor (B, C2, m) tensor of features to be propigated @@ -255,7 +251,6 @@ def forward(self, unknown, known, unknow_feats, known_feats): if __name__ == "__main__": from torch.autograd import Variable - torch.manual_seed(1) torch.cuda.manual_seed_all(1) xyz = Variable(torch.randn(2, 9, 3).cuda(), requires_grad=True) diff --git a/pointnet2/utils/pointnet2_utils.py b/pointnet2/utils/pointnet2_utils.py index dd95d0d..a95c86f 100644 --- a/pointnet2/utils/pointnet2_utils.py +++ b/pointnet2/utils/pointnet2_utils.py @@ -12,10 +12,8 @@ import sys import numpy as np - import builtins - try: import pointnet2._ext as _ext except ImportError: @@ -281,7 +279,6 @@ def backward(ctx, a=None): ball_query = BallQuery.apply - class QueryAndGroupRRI(nn.Module): r""" Groups with a ball query of radius for rigious rotation invariant @@ -295,7 +292,6 @@ class QueryAndGroupRRI(nn.Module): """ def __init__(self, radius, nsample): - # type: (QueryAndGroupRRI, float, int, bool) -> None super(QueryAndGroupRRI, self).__init__() self.radius, self.nsample = radius, nsample @@ -304,27 +300,32 @@ def forward(self, xyz, new_xyz, features=None): xyz_trans = xyz.transpose(1, 2).contiguous() grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) new_xyz_reshaped = new_xyz.transpose(1, 2).unsqueeze(-1) - #grouped_xyz = grouped_xyz - new_xyz_reshaped + # grouped_xyz = grouped_xyz - new_xyz_reshaped diff_grouped2grouped = grouped_xyz[:, :, :, :, None] - grouped_xyz[:, :, :, None, :] dis_grouped2grouped = torch.sqrt(torch.sum(diff_grouped2grouped * diff_grouped2grouped, dim=1)) mean_dis_grouped2grouped = torch.mean(dis_grouped2grouped, dim=3) tip_idxs = torch.argmax(mean_dis_grouped2grouped, dim=2) - tip_idxs_viewed = tip_idxs[:, None, :].repeat(1, grouped_xyz.shape[1], 1).view(grouped_xyz.shape[0]*grouped_xyz.shape[1]*grouped_xyz.shape[2]) - tip_pts = grouped_xyz.view(grouped_xyz.shape[0]*grouped_xyz.shape[1]*grouped_xyz.shape[2], grouped_xyz.shape[3])[range(tip_idxs_viewed.shape[0]), tip_idxs_viewed] + tip_idxs_viewed = tip_idxs[:, None, :].repeat(1, grouped_xyz.shape[1], 1).view( + grouped_xyz.shape[0] * grouped_xyz.shape[1] * grouped_xyz.shape[2]) + tip_pts = \ + grouped_xyz.view(grouped_xyz.shape[0] * grouped_xyz.shape[1] * grouped_xyz.shape[2], grouped_xyz.shape[3])[ + range(tip_idxs_viewed.shape[0]), tip_idxs_viewed] tip_pts = tip_pts.view(grouped_xyz.shape[0], grouped_xyz.shape[1], grouped_xyz.shape[2]) - grouped_proj_vec = torch.cross(torch.cross(new_xyz_reshaped.repeat(1, 1, 1, grouped_xyz.shape[3]), grouped_xyz, dim=1), - new_xyz_reshaped.repeat(1, 1, 1, grouped_xyz.shape[3]), dim=1) + grouped_proj_vec = torch.cross( + torch.cross(new_xyz_reshaped.repeat(1, 1, 1, grouped_xyz.shape[3]), grouped_xyz, dim=1), + new_xyz_reshaped.repeat(1, 1, 1, grouped_xyz.shape[3]), dim=1) grouped_proj_vec = grouped_proj_vec / torch.norm(grouped_proj_vec, dim=1)[:, None, :, :] new_r = torch.sqrt(torch.sum(new_xyz_reshaped * new_xyz_reshaped, dim=1)) new_norm = new_xyz_reshaped / (new_r[:, None, :, :] + 1e-8) - tip_proj_vec = torch.cross(torch.cross(new_xyz_reshaped.squeeze(dim=3), tip_pts, dim=1), new_xyz_reshaped.squeeze(dim=3), dim=1) + tip_proj_vec = torch.cross(torch.cross(new_xyz_reshaped.squeeze(dim=3), tip_pts, dim=1), + new_xyz_reshaped.squeeze(dim=3), dim=1) tip_proj_vec = tip_proj_vec / torch.norm(tip_proj_vec, dim=1)[:, None, :] - grouped_cross_vec = torch.cross(grouped_proj_vec, tip_proj_vec[:, :, :, None].repeat(1, 1, 1, grouped_proj_vec.shape[3]), dim=1) + grouped_cross_vec = torch.cross(grouped_proj_vec, + tip_proj_vec[:, :, :, None].repeat(1, 1, 1, grouped_proj_vec.shape[3]), dim=1) grouped_proj_vec_sin = torch.sum(grouped_cross_vec * new_norm, dim=1) - dis_grouped2grouped = torch.transpose(dis_grouped2grouped, 1, 2) dis_grouped2grouped_sort, _ = torch.sort(dis_grouped2grouped, dim=1) @@ -344,7 +345,6 @@ def forward(self, xyz, new_xyz, features=None): return new_features - ''' backup def forward(self, xyz, new_xyz, features=None): @@ -401,7 +401,6 @@ def forward(self, xyz, new_xyz, features=None): ''' - class QueryAndGroup(nn.Module): r""" Groups with a ball query of radius diff --git a/requirements.txt b/requirements.txt index c0768c9..0253d60 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ git+git://github.com/erikwijmans/etw_pytorch_utils.git@v1.1.1#egg=etw_pytorch_utils numpy future -torch==1.1.0 -torchvision==0.3.0 +#torch==1.1.0 +#torchvision==0.3.0 pillow==6.0.0 -opencv-python-headless \ No newline at end of file +opencv-python-headless diff --git a/test/test_structure_points.py b/test/test_structure_points.py index 3f710ea..71ed84b 100644 --- a/test/test_structure_points.py +++ b/test/test_structure_points.py @@ -6,6 +6,7 @@ unicode_literals, ) import sys + sys.path.append("..") import torch import os @@ -17,6 +18,7 @@ import random import glob + def create_color_list(num): colors = np.ndarray(shape=(num, 3)) random.seed(30) @@ -25,11 +27,12 @@ def create_color_list(num): colors[i, 1] = random.randint(0, 255) colors[i, 2] = random.randint(0, 255) return colors + + COLOR_LIST = create_color_list(5000) def main(args): - model = Pointnet2StructurePointNet(num_structure_points=args.num_structure_points, input_channels=0, use_xyz=True) model.cuda() checkpoint_util.load_checkpoint(model_3d=model, filename=args.model_fname) @@ -40,9 +43,8 @@ def main(args): fnames = glob.glob(os.path.join(args.data_dir, '*.off')) for fname in fnames: - fname = os.path.basename(fname) - pts = point_cloud_utils.read_points_off(os.path.join(args.data_dir,fname)) + pts = point_cloud_utils.read_points_off(os.path.join(args.data_dir, fname)) batch_pts = torch.from_numpy(pts)[None, :, :].cuda() structure_points = model(batch_pts) @@ -55,7 +57,8 @@ def main(args): if __name__ == "__main__": - + print("begin") + os.environ['CUDA_VISIBLE_DEVICES'] = '1' parser = argparse.ArgumentParser( description="Arguments", formatter_class=argparse.ArgumentDefaultsHelpFormatter, @@ -74,16 +77,3 @@ def main(args): ) args = parser.parse_args() main(args) - - - - - - - - - - - - - diff --git a/train/train_structure_points.py b/train/train_structure_points.py index 55478b4..bcf89be 100644 --- a/train/train_structure_points.py +++ b/train/train_structure_points.py @@ -6,6 +6,7 @@ unicode_literals, ) import sys + sys.path.append("..") import torch import torch.optim as optim @@ -29,7 +30,8 @@ torch.backends.cudnn.benchmark = True -def train_one_epoch(model, optimizer, data_loader, lr_scheduler, bnm_scheduler, current_iter, criterions, num_of_trans, num_inpts): +def train_one_epoch(model, optimizer, data_loader, lr_scheduler, bnm_scheduler, current_iter, criterions, num_of_trans, + num_inpts): model.train() ave_loss = 0 ave_cd_loss = 0 @@ -47,8 +49,10 @@ def train_one_epoch(model, optimizer, data_loader, lr_scheduler, bnm_scheduler, batch_size = batch_points.shape[0] if num_of_trans > 0: - transed_batch_points, rot_mats, trans_func_list = data_utils.AddPCATransformsToBatchPoints(batch_points, num_of_trans=num_of_trans) - transed_batch_points = transed_batch_points.view(batch_size * num_of_trans, transed_batch_points.shape[2], transed_batch_points.shape[3]) + transed_batch_points, rot_mats, trans_func_list = data_utils.AddPCATransformsToBatchPoints(batch_points, + num_of_trans=num_of_trans) + transed_batch_points = transed_batch_points.view(batch_size * num_of_trans, transed_batch_points.shape[2], + transed_batch_points.shape[3]) batch_points_all = torch.cat((batch_points, transed_batch_points), dim=0) else: batch_points_all = batch_points @@ -58,17 +62,21 @@ def train_one_epoch(model, optimizer, data_loader, lr_scheduler, bnm_scheduler, if num_inpts > 0: batch_points_all = point_cloud_utils.farthest_pts_sampling_tensor(batch_points_all, num_inpts) + # print(batch_points_all.shape) structure_points_all = model(batch_points_all) - structure_points_all = structure_points_all.view(num_of_trans + 1, batch_size, structure_points_all.shape[1], structure_points_all.shape[2]) + structure_points_all = structure_points_all.view(num_of_trans + 1, batch_size, structure_points_all.shape[1], + structure_points_all.shape[2]) structure_points = structure_points_all[0] if num_of_trans > 0: - transed_batch_points = transed_batch_points.view(num_of_trans, batch_size, transed_batch_points.shape[1], transed_batch_points.shape[2]) + transed_batch_points = transed_batch_points.view(num_of_trans, batch_size, transed_batch_points.shape[1], + transed_batch_points.shape[2]) transed_structure_points = structure_points_all[1:(1 + num_of_trans)] if num_of_trans > 0: - loss = criterions['ComputeLoss3d'](batch_points, structure_points, transed_batch_points, transed_structure_points, trans_func_list) + loss = criterions['ComputeLoss3d'](batch_points, structure_points, transed_batch_points, + transed_structure_points, trans_func_list) else: loss = criterions['ComputeLoss3d'](batch_points, structure_points, None, None, None) @@ -80,9 +88,13 @@ def train_one_epoch(model, optimizer, data_loader, lr_scheduler, bnm_scheduler, consistent_loss = criterions['ComputeLoss3d'].get_consistent_loss() if consistent_loss is not None: - print("\rbatch {0} current_loss {1}, cd_loss {2}, consistent_loss {3}".format(count, ("%.8f" % loss.item()), ("%.8f" % cd_loss.item()), ("%.8f" % consistent_loss.item())), end=" ") + print("\rbatch {0} current_loss {1}, cd_loss {2}, consistent_loss {3}".format(count, ("%.8f" % loss.item()), + ("%.8f" % cd_loss.item()), ( + "%.8f" % consistent_loss.item())), + end=" ") else: - print("\rbatch {0} current_loss {1}, cd_loss {2}".format(count, ("%.8f" % loss.item()), ("%.8f" % cd_loss.item())), end=" ") + print("\rbatch {0} current_loss {1}, cd_loss {2}".format(count, ("%.8f" % loss.item()), + ("%.8f" % cd_loss.item())), end=" ") ave_loss += loss.item() ave_cd_loss += cd_loss.item() @@ -115,12 +127,13 @@ def train(cmd_args): trans = transforms.Compose( [ d_utils.PointcloudToTensor(), - #d_utils.PointcloudRandomPermutation(), + # d_utils.PointcloudRandomPermutation(), d_utils.PointcloudJitter() ] ) - train_set = bhcp_dataloader.bhcp_dataloader(cmd_args.data_dir, cmd_args.category, transforms=trans, is_pts_aligned=False) + train_set = bhcp_dataloader.bhcp_dataloader(cmd_args.data_dir, cmd_args.category, transforms=trans, + is_pts_aligned=False) train_loader = DataLoader( train_set, batch_size=cmd_args.batch_size, @@ -128,7 +141,8 @@ def train(cmd_args): num_workers=5, pin_memory=False, ) - model = Pointnet2StructurePointNet(num_structure_points=cmd_args.num_structure_points, input_channels=0, use_xyz=True) + model = Pointnet2StructurePointNet(num_structure_points=cmd_args.num_structure_points, input_channels=0, + use_xyz=True) model.cuda() optimizer = optim.Adam( model.parameters(), lr=cmd_args.lr, weight_decay=cmd_args.weight_decay @@ -167,18 +181,24 @@ def train(cmd_args): iters = max(iters, 0) for epoch_i in range(start_epoch, cmd_args.max_epochs): log.write('\nepoch: {0} ###########################'.format(epoch_i)) - train_loss, train_cd_loss, train_consistent_loss, iters = train_one_epoch(model, optimizer, train_loader, lr_scheduler, bnm_scheduler, iters, criterions, num_of_trans=cmd_args.num_of_transform, num_inpts=cmd_args.num_inpts) - log.write('\nave_train_loss:{0}, cd_loss:{1}, consis_loss:{2}'.format(("%.8f" % train_loss), ("%.8f" % train_cd_loss), ("%.8f" % train_consistent_loss))) + train_loss, train_cd_loss, train_consistent_loss, iters = train_one_epoch(model, optimizer, train_loader, + lr_scheduler, bnm_scheduler, iters, + criterions, + num_of_trans=cmd_args.num_of_transform, + num_inpts=cmd_args.num_inpts) + log.write( + '\nave_train_loss:{0}, cd_loss:{1}, consis_loss:{2}'.format(("%.8f" % train_loss), ("%.8f" % train_cd_loss), + ("%.8f" % train_consistent_loss))) if cmd_args.checkpoint_save_step != -1 and (epoch_i + 1) % cmd_args.checkpoint_save_step is 0: fname = os.path.join(checkpoints_dir, 'checkpoint_{}'.format(epoch_i)) - checkpoint_util.save_checkpoint(filename=fname, model_3d=model, optimizer=optimizer, iters=iters, epoch=epoch_i) + checkpoint_util.save_checkpoint(filename=fname, model_3d=model, optimizer=optimizer, iters=iters, + epoch=epoch_i) fname = os.path.join(checkpoints_dir, 'model') checkpoint_util.save_checkpoint(filename=fname, model_3d=model, optimizer=optimizer, iters=iters, epoch=epoch_i) - def parse_args(): parser = argparse.ArgumentParser( description="Arguments", @@ -210,7 +230,8 @@ def parse_args(): , help="Checkpoint to start from" ) parser.add_argument( - "-num_of_transform", type=int, default=0, help="Number of transforms for rotation data augmentation. Useful when testing on shapes without alignment" + "-num_of_transform", type=int, default=0, + help="Number of transforms for rotation data augmentation. Useful when testing on shapes without alignment" ) parser.add_argument( @@ -238,10 +259,6 @@ def parse_args(): if __name__ == "__main__": - os.environ['CUDA_VISIBLE_DEVICES'] = '0' + os.environ['CUDA_VISIBLE_DEVICES'] = '1' args = parse_args() train(cmd_args=args) - - - - From 51967e895520f24865abd5b1acf4ce8ee6248d6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E5=BC=BA?= Date: Sat, 6 Jun 2020 23:42:43 +0800 Subject: [PATCH 2/3] simplify code by using Integration Layer --- dataset/data_utils.py | 46 ++++----------- models/base.py | 24 ++++++++ models/pointnet2_structure_point_net.py | 76 ++++++------------------- train/train_structure_points.py | 23 ++++---- 4 files changed, 64 insertions(+), 105 deletions(-) create mode 100644 models/base.py diff --git a/dataset/data_utils.py b/dataset/data_utils.py index 5910154..eda5924 100644 --- a/dataset/data_utils.py +++ b/dataset/data_utils.py @@ -6,6 +6,7 @@ unicode_literals, ) import sys + sys.path.append("..") import torch import numpy as np @@ -34,14 +35,14 @@ def angle_axis_tensor(angle, axis): # yapf: disable cross_prod_mat = torch.Tensor([[0.0, -u[2], u[1]], - [u[2], 0.0, -u[0]], - [-u[1], u[0], 0.0]]).type(torch.FloatTensor) + [u[2], 0.0, -u[0]], + [-u[1], u[0], 0.0]]).type(torch.FloatTensor) R = cosval * torch.eye(3).type(torch.FloatTensor) + sinval * cross_prod_mat + (1.0 - cosval) * torch.ger(u, u) - return R + def angle_axis(angle, axis): # type: (float, np.ndarray) -> float r"""Returns a 4x4 rotation matrix that performs a rotation around axis by angle @@ -63,8 +64,8 @@ def angle_axis(angle, axis): # yapf: disable cross_prod_mat = np.array([[0.0, -u[2], u[1]], - [u[2], 0.0, -u[0]], - [-u[1], u[0], 0.0]]) + [u[2], 0.0, -u[0]], + [-u[1], u[0], 0.0]]) R = torch.from_numpy( cosval * np.eye(3) @@ -145,8 +146,8 @@ def __init__(self, std=0.01, clip=0.05): def __call__(self, points): jittered_data = ( points.new(points.size(0), 3) - .normal_(mean=0.0, std=self.std) - .clamp_(-self.clip, self.clip) + .normal_(mean=0.0, std=self.std) + .clamp_(-self.clip, self.clip) ) points[:, 0:3] += jittered_data return points @@ -154,11 +155,9 @@ def __call__(self, points): class PointcloudNormalize(object): def __init__(self, max_size=1.0): - self.max_size = max_size def __call__(self, points): - points_max, _ = torch.max(points, dim=0) points_min, _ = torch.min(points, dim=0) points_center = (points_max + points_min) / 2 @@ -167,10 +166,10 @@ def __call__(self, points): points = points / max_radius * self.max_size / 2.0 return points + class PointcloudRandomPermutation(object): def __call__(self, points): - num = points.shape[0] idxs = torch.randperm(num).type(torch.LongTensor) points = torch.index_select(points, 0, idxs).clone() @@ -208,11 +207,6 @@ def __call__(self, points): return torch.from_numpy(pc).float() - - - - - class PointcloudTranslate(object): def __init__(self, translation=np.array([0.0, 0.1, 0.0])): ''' @@ -220,7 +214,6 @@ def __init__(self, translation=np.array([0.0, 0.1, 0.0])): ''' self.translation = torch.from_numpy(translation) - def __call__(self, points): ''' @@ -242,7 +235,6 @@ def __init__(self, scaler): self.scaler = scaler def __call__(self, points): - respoints = points * self.scaler return respoints @@ -254,7 +246,6 @@ def __init__(self, angle_in_degree=np.pi, axis=np.array([0.0, 1.0, 0.0]), is_cud self.rotation_matrix_t = angle_axis(self.angle_in_degree, self.axis).t() def __call__(self, points): - ''' :param points: ... , num_of_points, 3 :return: points after rotate @@ -267,8 +258,7 @@ def __call__(self, points): return tpoints - -def GenPointcloudRandomTransformFunction(max_rot_angle=2*np.pi): +def GenPointcloudRandomTransformFunction(max_rot_angle=2 * np.pi): scale_lo = 0.8 scale_hi = 1.25 scaler = np.random.uniform(scale_lo, scale_hi) @@ -284,7 +274,7 @@ def GenPointcloudRandomTransformFunction(max_rot_angle=2*np.pi): return trans_func -def AddTransformsToBatchPoints(points, num_of_trans, max_rot_angle=2*np.pi): +def AddTransformsToBatchPoints(points, num_of_trans, max_rot_angle=2 * np.pi): ''' :param points:bn, num_of_points, 3 @@ -324,13 +314,10 @@ def __call__(self, points): else: tmp_rot = self.rot_mats - transed_poitns = torch.transpose(torch.matmul(tmp_rot, torch.transpose(points, 1, 2)), 1, 2) return transed_poitns - - def AddPCATransformsToBatchPoints(points, num_of_trans): trans_points_all = None rot_mats_all = None @@ -347,7 +334,7 @@ def AddPCATransformsToBatchPoints(points, num_of_trans): np.random.shuffle(tmp_idx) pca_axis = pca_axis_raw[tmp_idx, :] tmp_sign = np.random.randint(2, size=2) - tmp_sign[tmp_sign==0] = -1 + tmp_sign[tmp_sign == 0] = -1 pca_axis[0, :] = pca_axis[0, :] * tmp_sign[0] pca_axis[1, :] = pca_axis[1, :] * tmp_sign[1] pca_axis[2, :] = np.cross(pca_axis[0, :], pca_axis[1, :]) @@ -382,13 +369,4 @@ def AddPCATransformsToBatchPoints(points, num_of_trans): trans_func = PointcloudRotateFuns(rot_mats_all[ti, :, :, :]) transfunc_list.append(trans_func) - return trans_points_all, rot_mats_all, transfunc_list - - - - - - - - diff --git a/models/base.py b/models/base.py new file mode 100644 index 0000000..e1d7730 --- /dev/null +++ b/models/base.py @@ -0,0 +1,24 @@ +import torch.nn as nn +import numpy as np +from abc import abstractmethod + + +class BaseModel(nn.Module): + """ + Base class for all models + """ + @abstractmethod + def forward(self, *inputs): + """ + Forward pass logic + :return: Model output + """ + raise NotImplementedError + + def __str__(self): + """ + Model prints with number of trainable parameters + """ + model_parameters = filter(lambda p: p.requires_grad, self.parameters()) + params = sum([np.prod(p.size()) for p in model_parameters]) + return super().__str__() + '\nTrainable parameters: {}'.format(params) diff --git a/models/pointnet2_structure_point_net.py b/models/pointnet2_structure_point_net.py index ff508ba..9189041 100644 --- a/models/pointnet2_structure_point_net.py +++ b/models/pointnet2_structure_point_net.py @@ -7,12 +7,12 @@ ) import torch import torch.nn as nn -import etw_pytorch_utils as pt_utils from pointnet2.utils.pointnet2_modules import PointnetSAModuleMSG from models import chamfer_distance +from models.base import BaseModel -class ComputeLoss3d(nn.Module): +class ComputeLoss3d(BaseModel): def __init__(self): super(ComputeLoss3d, self).__init__() @@ -70,7 +70,7 @@ def get_consistent_loss(self): return self.consistent_loss -class IntegrationLayer(nn.Module): +class Conv1dProbLayer(BaseModel): def __init__(self, in_channels, out_channels, out=False, kernel_size=1, dropout=0.2): super().__init__() self.out = out @@ -85,14 +85,13 @@ def __init__(self, in_channels, out_channels, out=False, kernel_size=1, dropout= def forward(self, x): x = self.dropout_conv_bn_layer(x) if self.out: - x = self.relu(x) - else: x = self.softmax(x) + else: + x = self.relu(x) return x -class Pointnet2StructurePointNet(nn.Module): - +class Pointnet2StructurePointNet(BaseModel): def __init__(self, num_structure_points, input_channels=3, use_xyz=True): super(Pointnet2StructurePointNet, self).__init__() self.point_dim = 3 @@ -132,58 +131,22 @@ def __init__(self, num_structure_points, input_channels=3, use_xyz=True): conv1d_stpts_prob_modules = [] if num_structure_points <= 128 + 256 + 256: - # conv1d_stpts_prob_modules.append( - # IntegrationLayer(in_channels=128 + 256 + 256, out_channels=512)) - conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) - conv1d_stpts_prob_modules.append(nn.Conv1d(in_channels=128 + 256 + 256, out_channels=512, kernel_size=1)) - conv1d_stpts_prob_modules.append(nn.BatchNorm1d(512)) - conv1d_stpts_prob_modules.append(nn.ReLU()) + conv1d_stpts_prob_modules.append(Conv1dProbLayer(128 + 256 + 256, 512)) in_channels = 512 while in_channels >= self.num_structure_points * 2: out_channels = int(in_channels / 2) - # conv1d_stpts_prob_modules.append(IntegrationLayer(in_channels, out_channels)) - conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) - conv1d_stpts_prob_modules.append( - nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1)) - conv1d_stpts_prob_modules.append(nn.BatchNorm1d(out_channels)) - conv1d_stpts_prob_modules.append(nn.ReLU()) + conv1d_stpts_prob_modules.append(Conv1dProbLayer(in_channels, out_channels)) in_channels = out_channels + conv1d_stpts_prob_modules.append(Conv1dProbLayer(in_channels, self.num_structure_points, True)) - conv1d_stpts_prob_modules.append(IntegrationLayer(in_channels, self.num_structure_points, out=True)) - # conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) - # conv1d_stpts_prob_modules.append( - # nn.Conv1d(in_channels=in_channels, out_channels=self.num_structure_points, kernel_size=1)) - # - # conv1d_stpts_prob_modules.append(nn.BatchNorm1d(self.num_structure_points)) - # conv1d_stpts_prob_modules.append(nn.Softmax(dim=2)) else: - - # conv1d_stpts_prob_modules.append( - # IntegrationLayer(in_channels=128 + 256 + 256, out_channels=1024, kernel_size=1)) - conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) - conv1d_stpts_prob_modules.append(nn.Conv1d(in_channels=128 + 256 + 256, out_channels=1024, kernel_size=1)) - conv1d_stpts_prob_modules.append(nn.BatchNorm1d(1024)) - conv1d_stpts_prob_modules.append(nn.ReLU()) - in_channels = 1024 - while in_channels <= self.num_structure_points / 2: + conv1d_stpts_prob_modules.append(Conv1dProbLayer(128 + 256 + 256, 1024)) + while in_channels < self.num_structure_points / 2: # change <= to < out_channels = int(in_channels * 2) - # conv1d_stpts_prob_modules.append( - # IntegrationLayer(in_channels=in_channels, out_channels=out_channels, kernel_size=1)) - conv1d_stpts_prob_modules.append(nn.Dropout(0.2)) - conv1d_stpts_prob_modules.append(nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1)) - conv1d_stpts_prob_modules.append(nn.BatchNorm1d(out_channels)) - conv1d_stpts_prob_modules.append(nn.ReLU()) + conv1d_stpts_prob_modules.append(Conv1dProbLayer(in_channels, out_channels)) in_channels = out_channels - - # conv1d_stpts_prob_modules.append(IntegrationLayer(in_channels, self.num_structure_points, True)) - conv1d_stpts_prob_modules.append(nn.dropout(0.2)) - conv1d_stpts_prob_modules.append( - nn.conv1d(in_channels=in_channels, out_channels=self.num_structure_points, kernel_size=1)) - - conv1d_stpts_prob_modules.append(nn.batchnorm1d(self.num_structure_points)) - conv1d_stpts_prob_modules.append(nn.Softmax(dim=2)) - + conv1d_stpts_prob_modules.append(Conv1dProbLayer(in_channels, self.num_structure_points, True)) self.conv1d_stpts_prob = nn.Sequential(*conv1d_stpts_prob_modules) def _break_up_pc(self, pc): @@ -204,19 +167,14 @@ def forward(self, pointcloud, return_weighted_feature=False): xyz, features = module(xyz, features) self.stpts_prob_map = self.conv1d_stpts_prob(features) - weighted_xyz = torch.sum(self.stpts_prob_map[:, :, :, None] * xyz[:, None, :, :], dim=2) + # print("prob:{},xyz:{},weighted_xyz:{},features:{}".format(self.stpts_prob_map.shape, + # xyz.shape, + # weighted_xyz.shape, + # features.shape)) if return_weighted_feature: weighted_features = torch.sum(self.stpts_prob_map[:, None, :, :] * features[:, :, None, :], dim=3) return weighted_xyz, weighted_features else: - return weighted_xyz - - -if __name__ == "__main__": - pointclouds = torch.rand(10, 1000, 3).cuda() - model = Pointnet2StructurePointNet(15).cuda() - output = model(pointclouds) - print(output.shape) diff --git a/train/train_structure_points.py b/train/train_structure_points.py index bcf89be..eda289f 100644 --- a/train/train_structure_points.py +++ b/train/train_structure_points.py @@ -30,8 +30,8 @@ torch.backends.cudnn.benchmark = True -def train_one_epoch(model, optimizer, data_loader, lr_scheduler, bnm_scheduler, current_iter, criterions, num_of_trans, - num_inpts): +def train_one_epoch(model, optimizer, data_loader, lr_scheduler, bnm_scheduler, + current_iter, criterions, num_of_trans, num_inputs): model.train() ave_loss = 0 ave_cd_loss = 0 @@ -59,15 +59,12 @@ def train_one_epoch(model, optimizer, data_loader, lr_scheduler, bnm_scheduler, batch_points_all = batch_points_all.cuda() - if num_inpts > 0: - batch_points_all = point_cloud_utils.farthest_pts_sampling_tensor(batch_points_all, num_inpts) + if num_inputs > 0: + batch_points_all = point_cloud_utils.farthest_pts_sampling_tensor(batch_points_all, num_inputs) - # print(batch_points_all.shape) structure_points_all = model(batch_points_all) - structure_points_all = structure_points_all.view(num_of_trans + 1, batch_size, structure_points_all.shape[1], structure_points_all.shape[2]) - structure_points = structure_points_all[0] if num_of_trans > 0: transed_batch_points = transed_batch_points.view(num_of_trans, batch_size, transed_batch_points.shape[1], @@ -90,7 +87,7 @@ def train_one_epoch(model, optimizer, data_loader, lr_scheduler, bnm_scheduler, if consistent_loss is not None: print("\rbatch {0} current_loss {1}, cd_loss {2}, consistent_loss {3}".format(count, ("%.8f" % loss.item()), ("%.8f" % cd_loss.item()), ( - "%.8f" % consistent_loss.item())), + "%.8f" % consistent_loss.item())), end=" ") else: print("\rbatch {0} current_loss {1}, cd_loss {2}".format(count, ("%.8f" % loss.item()), @@ -143,6 +140,7 @@ def train(cmd_args): ) model = Pointnet2StructurePointNet(num_structure_points=cmd_args.num_structure_points, input_channels=0, use_xyz=True) + print(model) model.cuda() optimizer = optim.Adam( model.parameters(), lr=cmd_args.lr, weight_decay=cmd_args.weight_decay @@ -185,13 +183,13 @@ def train(cmd_args): lr_scheduler, bnm_scheduler, iters, criterions, num_of_trans=cmd_args.num_of_transform, - num_inpts=cmd_args.num_inpts) + num_inputs=cmd_args.num_inputs) log.write( '\nave_train_loss:{0}, cd_loss:{1}, consis_loss:{2}'.format(("%.8f" % train_loss), ("%.8f" % train_cd_loss), ("%.8f" % train_consistent_loss))) if cmd_args.checkpoint_save_step != -1 and (epoch_i + 1) % cmd_args.checkpoint_save_step is 0: - fname = os.path.join(checkpoints_dir, 'checkpoint_{}'.format(epoch_i)) + fname = os.path.join(checkpoints_dir, 'checkpoint_{}'.format(epoch_i+1)) checkpoint_util.save_checkpoint(filename=fname, model_3d=model, optimizer=optimizer, iters=iters, epoch=epoch_i) @@ -224,7 +222,7 @@ def parse_args(): parser.add_argument( "-checkpoint_save_step", type=int, default=10, help="Step for saving Checkpoint" ) - parser.add_argument("-batch_size", type=int, default=3, help="Batch size") + parser.add_argument("-batch_size", type=int, default=24, help="Batch size") parser.add_argument( "-checkpoint", type=str, default=None , help="Checkpoint to start from" @@ -235,7 +233,7 @@ def parse_args(): ) parser.add_argument( - "-num_inpts", type=int, default=2048, help="sample points from initial point cloud" + "-num_inputs", type=int, default=2048, help="sample points from initial point cloud" ) parser.add_argument( @@ -261,4 +259,5 @@ def parse_args(): if __name__ == "__main__": os.environ['CUDA_VISIBLE_DEVICES'] = '1' args = parse_args() + print(args) train(cmd_args=args) From f22e95b96340f8cd419c363f1521316786883718 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E5=BC=BA?= Date: Sat, 6 Jun 2020 23:55:26 +0800 Subject: [PATCH 3/3] update readme --- README.md | 1 + requirements.txt | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index fc6e5ca..390b0ce 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ This repository contains the code for our cvpr 2020 paper: [Unsupervised Learnin Current Code is tested on ubuntu16.04 with cuda9, python3.6, torch 1.1.0 and torchvision 0.3.0. We use a [pytorch version of pointnet++](https://github.com/erikwijmans/Pointnet2_PyTorch) in our pipeline. ``` +conda install pytorch==1.1.0 torchvision==0.3.0 cudatoolkit=10.0 -c pytorch pip install -r requirements.txt cd pointnet2 python setup.py build_ext --inplace diff --git a/requirements.txt b/requirements.txt index 0253d60..cb41e52 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ git+git://github.com/erikwijmans/etw_pytorch_utils.git@v1.1.1#egg=etw_pytorch_utils numpy future -#torch==1.1.0 -#torchvision==0.3.0 +# torch==1.1.0 +# torchvision==0.3.0 pillow==6.0.0 opencv-python-headless