learning3d 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

learning3d/__init__.py

@@ -1,2 +0,0 @@
-from .models import PointNet, create_pointconv, DGCNN, PPFNet, Pooling, Classifier, Segmentation
-from .models import DCP, PRNet, iPCRNet, PointNetLK, RPMNet, PCN, DeepGMR, MaskNet, MaskNet2

learning3d/models/__init__.py

@@ -15,9 +15,4 @@ from .rpmnet import RPMNet
 from .pcn import PCN
 from .deepgmr import DeepGMR
 from .masknet import MaskNet
-from .masknet2 import MaskNet2
-
-try:
-	from .flownet3d import FlowNet3D
-except:
-	print("Error raised in pointnet2 module for FlowNet3D Network!\nEither don't use pointnet2_utils or retry it's setup.")
+from .masknet2 import MaskNet2

learning3d/utils/__init__.py

@@ -1,9 +1,4 @@
 from .svd import SVDHead
 from .transformer import Transformer, Identity
 from .ppfnet_util import angle_difference, square_distance, index_points, farthest_point_sample, query_ball_point, sample_and_group, sample_and_group_multi
-from .pointconv_util import PointConvDensitySetAbstraction
-
-try:
-	from .lib import pointnet2_utils
-except:
-	print("Error raised in pointnet2 module in utils!\nEither don't use pointnet2_utils or retry it's setup.")
+from .pointconv_util import PointConvDensitySetAbstraction

{learning3d-0.0.2.dist-info → learning3d-0.0.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: learning3d
-Version: 0.0.2
+Version: 0.0.3
 Summary: Learning3D: A Modern Library for Deep Learning on 3D Point Clouds Data
 Author-email: Vinit Sarode <vinitsarode5@gmail.com>
 Project-URL: Homepage, https://github.com/vinits5/learning3d

{learning3d-0.0.2.dist-info → learning3d-0.0.3.dist-info}/RECORD

@@ -1,10 +1,9 @@
-learning3d/__init__.py,sha256=WsLbk69HZ7OXJwTpNFcEphgtNp4sqEVGbZ1HGeNdYZ0,189
+learning3d/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 learning3d/data_utils/__init__.py,sha256=iYAVh0FThnVlG42QIgmDYrC3NGVYuzKX8s1oRqAI1YU,261
 learning3d/data_utils/dataloaders.py,sha256=kb0wsLlMN7sB-CS_4BGSprSaZBwkUNYxS5iwUdD6JJM,14871
 learning3d/data_utils/user_data.py,sha256=ADDGeCUCr6TcXhcxvAFncIeLO71xoRHYi4H418ktvQs,4828
 learning3d/examples/test_dcp.py,sha256=o1hgl22b2xhqqYe4_k7NlNo3-rIMmSzzssCI74DCfoE,5606
 learning3d/examples/test_deepgmr.py,sha256=pb-gRE6YLLaPFmiBXQjaE-B-J314QY4RkNwLHx65bcU,5597
-learning3d/examples/test_flownet.py,sha256=52i9UOnSpCxZXrCgYJ-HNxsxRXqr_K4PydYMF90Lxv8,4603
 learning3d/examples/test_masknet.py,sha256=dkqUui9sv2SzHtvtzUzL_PxJxMcBMqCSDPAYg0BWAVU,6405
 learning3d/examples/test_masknet2.py,sha256=3_XWBQOwQjK3BCQ__BmPhCvYI_0hZMK3X4C-P2Krw6w,6859
 learning3d/examples/test_pcn.py,sha256=4eaosjJVqiFxlqnaUWu-O2Jawt4uU16UJEzitIjP2us,4342
@@ -17,7 +16,6 @@ learning3d/examples/test_rpmnet.py,sha256=oy-z7I26IQxr4TD_p0qCRnOn6H8VbQFyiWO83Z
 learning3d/examples/train_PointNetLK.py,sha256=0GgT2NYKNZl8o02rvW-nYBO_1tlfDNuakuAXtm1V16c,8773
 learning3d/examples/train_dcp.py,sha256=SQVrwnZqGmFCZv_X2tzMysBmv-HI9sllZMWw5zsW3NM,9511
 learning3d/examples/train_deepgmr.py,sha256=vxdkgfQZPtwuYryR0chegTiLuXOQag8r_ccGJ6qtw7o,9397
-learning3d/examples/train_flownet.py,sha256=V3uG7EaqsQO0BtmAFFN_aHb5bsyYoLv3JKb1_XhYKNw,10369
 learning3d/examples/train_masknet.py,sha256=XzgWsmVAm5Lk21mH9qhvNN0um4pI1fYVfsBAV4deSOM,8889
 learning3d/examples/train_pcn.py,sha256=X7MSYVXwgIMExplua1M9pG20eNhZ_0p83yTADSYrAlA,7542
 learning3d/examples/train_pcrnet.py,sha256=KQ8MiDUiR46qS9t7tc5POJ3NjMyZFBEPOVQY-7Vszpk,8198
@@ -45,12 +43,11 @@ learning3d/losses/cuda/emd_torch/pkg/layer/__init__.py,sha256=aDpfP0iZyg3Uu-ppa3
 learning3d/losses/cuda/emd_torch/pkg/layer/emd_loss_layer.py,sha256=yCEEfafLZ1Ia_BCrE7mcnDRDaaEj6je3Rj8cnQ_Zrts,1019
 learning3d/losses/cuda/emd_torch/pkg/src/emd.cpp,sha256=lQ4q2XO5K2CldYmnRJnGhKTR9hVRFTwO305bHT3Cauo,17
 learning3d/losses/cuda/emd_torch/pkg/src/cuda/emd.cu,sha256=DJXALRWyiVDzaKWsD2tQnEXrIT9GpIldkvIJ9fZMGbI,1462
-learning3d/models/__init__.py,sha256=r5upYkq5M4pu7rxuFxFI5AshVGQI-gMMsQRm_iF3TuE,642
+learning3d/models/__init__.py,sha256=WgAx7FlaUbijjl3qjQCvwyrbdhxvKHV7HbqBlq47ux8,473
 learning3d/models/classifier.py,sha256=_LUNXbLrpKNXmCkO2R1mz64dbwfrze7f_4SYT1Z6SYo,1205
 learning3d/models/dcp.py,sha256=LZFgtk9f9f9s3QvX65nFXGgC33yGIZuy4XjviwH8OGE,3377
 learning3d/models/deepgmr.py,sha256=vIxOQrZjvOCHLElJCjZ8EcZ-vm0-v71IKsPGuSF-elE,5298
 learning3d/models/dgcnn.py,sha256=Bt-dP2NwpOy4AcWrspXfVV1EKL-ZQElYUp2uaWNvE_Q,3057
-learning3d/models/flownet3d.py,sha256=2cPqzwXyw5uBNWIpHLTwRf0exSHYcW2Lyd94wOHgXZ0,17667
 learning3d/models/masknet.py,sha256=ElMF3b-JgYmgwSEf1taGQvhA7Xy7_MiHEofzc03VCd8,2705
 learning3d/models/masknet2.py,sha256=6lgukurfzUOY-6xdCpMljOYFtvADLSczAXJzRC3Jkh4,9063
 learning3d/models/pcn.py,sha256=FvpjLR6t3kFQ1I4Fhpbsaj_P8Ml6S912x36JAZ1dUKs,5346
@@ -71,45 +68,13 @@ learning3d/ops/se3.py,sha256=x6oLbQzLOXiM0xDJsVUCUE1liZ_TaJzkkHQvIyjqCqI,3957
 learning3d/ops/sinc.py,sha256=A_Ffu07RXGx3daZn4zOGdnW10_l06cmhFdAiU4NKhcw,5228
 learning3d/ops/so3.py,sha256=b0tX5nHyF2Qtp8V0ejGKaPaHJQ_G38ifQ7gSJzRU1ts,5166
 learning3d/ops/transform_functions.py,sha256=hvNjZO-uJodsGYtQwtAtDxtQ6uBpA7Lv9t-_yAg6wxo,12806
-learning3d/utils/__init__.py,sha256=iuYyToRcZ9YZMNn-ngJMjzU8p4FkaO2m8YyXSE14hN4,442
+learning3d/utils/__init__.py,sha256=lwMtW0H4LEZJuoVCi7qmfaWdOVrL2YmP_YP-ZfUoSlc,286
 learning3d/utils/pointconv_util.py,sha256=kJxGztai7X15YsGuorMOc50SPtj_k1yfkP4XCTzIWdM,14331
 learning3d/utils/ppfnet_util.py,sha256=HEoxkgUBlawKZLWspfQm3caWUyAMIrW-ECtStNYbe2Y,7989
 learning3d/utils/svd.py,sha256=yCYQt2SKqeIzCBnBEr_8xFR79m4fIoNVFnp77epn1dM,1936
 learning3d/utils/transformer.py,sha256=UDgJvnh7ekWyijaAn-a3ckeFeMxlK_chXzWlhAGDiPM,8974
-learning3d/utils/lib/pointnet2_modules.py,sha256=Tqiz32BT-fc2gHmv87xbbNUctE2mEpYpnJ3jfVmqw0w,6339
-learning3d/utils/lib/pointnet2_utils.py,sha256=TQOVhi22raBffgtKWsVf-reoEqPkN-zvv6GOh2TqHxY,10637
-learning3d/utils/lib/pytorch_utils.py,sha256=-kJfrswEu0gnnrYM1yIrj6HFW2YVyK2P9To4EyYAqNw,6173
-learning3d/utils/lib/setup.py,sha256=AF3MiatsAYYEIl26Vt5JlsmZEI19lV86k6b4fw1uxgU,679
-learning3d/utils/lib/build/lib.linux-x86_64-3.5/pointnet2_cuda.cpython-35m-x86_64-linux-gnu.so,sha256=HjZKllvXRMLn1ONxFPA-RF2athcY06clU-D91CLS4Dw,7222568
-learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/ball_query.o,sha256=_jT_UVj61Yi2eDgpn1gn595TTqbQ8hAERAazE9kcuwA,2730016
-learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/ball_query_gpu.o,sha256=GFosEyxTG6gPBUZJ7niNojLxJvt12v2cnO0lqtnO3mI,13912
-learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/group_points.o,sha256=kNMDj14OwT70OtmCf63mdtnf58tXpnKXe95DOHu63Ow,2490896
-learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/group_points_gpu.o,sha256=CCX8QD24Kp4hJqUFG34e88Co5poWc6XvkARQYpBgs88,17048
-learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/interpolate.o,sha256=xyOxGf0zrgX8k6Muvn4WgGuZWk0N1-E7r4Fnjpao-sg,2503232
-learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/interpolate_gpu.o,sha256=4kAHyU-Ho3_Sr77NWJINzVndHu8PdirIUqm0tQ7yVY4,32712
-learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/pointnet2_api.o,sha256=L8v-YcL3J53kKxFw50frNlsKl1n4HowU5krlscSZ-5E,7242784
-learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/sampling.o,sha256=lOw0IlJMXjCbTw9LONz7_TEzjWl484WFzilyzW742Io,2495440
-learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/sampling_gpu.o,sha256=ULYz2ppU_n0KU8ka0Q4FE9wdPUuZDYcmvxRWN8UgUnU,95544
-learning3d/utils/lib/dist/pointnet2-0.0.0-py3.5-linux-x86_64.egg,sha256=zwF5_pf6Id5lVow5MGkdug_upjMpbEwW7hUTIHEub0I,2241340
-learning3d/utils/lib/pointnet2.egg-info/SOURCES.txt,sha256=XtaY4SQXYOT8Hhdjw4iworVxoDjoWdmdBe5fpkwK_-c,328
-learning3d/utils/lib/pointnet2.egg-info/dependency_links.txt,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-learning3d/utils/lib/pointnet2.egg-info/top_level.txt,sha256=XVeKz4zj8Ixbw0jjiXuOqdON71dDXExnItEZVUaOdVE,15
-learning3d/utils/lib/src/ball_query.cpp,sha256=Ysx8OAn4G-dSrlPgbZS2WCDfxcV54HQTZu5FbL2QlKs,933
-learning3d/utils/lib/src/ball_query_gpu.cu,sha256=vj-Y6gUh2oVLLYTuz8JE8tcHxGrJBTC4pTRsNcWNa1o,2050
-learning3d/utils/lib/src/ball_query_gpu.h,sha256=47MXh8U46wcpvQ7wcC05QyAsE1N3GB-TExsT0hoApxE,476
-learning3d/utils/lib/src/cuda_utils.h,sha256=es4AGc6gYXileyddOxCGsmzV0Mg9ZG6LiucMUw6ee8o,353
-learning3d/utils/lib/src/group_points.cpp,sha256=koGBIklYmbwHDpCrHPURJl_f1AtsqpNB2DbgRRb4zuc,1171
-learning3d/utils/lib/src/group_points_gpu.cu,sha256=OIa_C0egafVIvJTcfu7UUI-E4ylxktcSG-fUjJQbNow,3307
-learning3d/utils/lib/src/group_points_gpu.h,sha256=lNX0VSo6xFCfWQh-ZPUIBobu2kfmFz8PF8l2zZe1NyQ,836
-learning3d/utils/lib/src/interpolate.cpp,sha256=VKMC4dveIgdfTiqpjrhe5CMqW097Hgmza5PkZVKVUfc,2521
-learning3d/utils/lib/src/interpolate_gpu.cu,sha256=yOHn2D5MFonqrfLD7K7-DFozXqSDqdetdaFHJzsPMKs,7470
-learning3d/utils/lib/src/interpolate_gpu.h,sha256=SbxS1oY6Qrp3f-r_kDtjq5_EzRpdIZpJIGVfmhhmnsI,1477
-learning3d/utils/lib/src/pointnet2_api.cpp,sha256=LRvlZoDLpaZHg2Tvb2ad5j8qVNCZ7TX9XEsNfdlSpT8,1213
-learning3d/utils/lib/src/sampling.cpp,sha256=kH75rE1jxri4v8wIksKAjaUbJlXB3FLEXYQTmb0kmII,1549
-learning3d/utils/lib/src/sampling_gpu.cu,sha256=V9dTFaYksQ-jNnoUvQgRNuFW9MmaBdZpbmXQbnvShJc,7934
-learning3d/utils/lib/src/sampling_gpu.h,sha256=STr6hTB9A8D0skTTR6TiU79j1eSc-FRqik_j0PPWDmM,1045
-learning3d-0.0.2.dist-info/LICENSE,sha256=3qY3_NeQIvalbLlsHFtOfuUKjs_U2k6u7rf6YVx6ac0,1098
-learning3d-0.0.2.dist-info/METADATA,sha256=hb8g3GL1qvPFFtblCBP_RaVPimnlrkqbv0nR4fIHLc4,15813
-learning3d-0.0.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-learning3d-0.0.2.dist-info/top_level.txt,sha256=nTmYW8NhbNV1_15DGNpl_OvvSFtQP98sy3qrrHr0eLo,11
-learning3d-0.0.2.dist-info/RECORD,,
+learning3d-0.0.3.dist-info/LICENSE,sha256=3qY3_NeQIvalbLlsHFtOfuUKjs_U2k6u7rf6YVx6ac0,1098
+learning3d-0.0.3.dist-info/METADATA,sha256=RjRxCmF_3xda4kPWDCvP2QjyPSysHZdNm3PY42-Jlhg,15813
+learning3d-0.0.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+learning3d-0.0.3.dist-info/top_level.txt,sha256=nTmYW8NhbNV1_15DGNpl_OvvSFtQP98sy3qrrHr0eLo,11
+learning3d-0.0.3.dist-info/RECORD,,

learning3d/examples/test_flownet.py

@@ -1,113 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-
-import open3d as o3d
-import os
-import gc
-import argparse
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-from torch.optim.lr_scheduler import MultiStepLR
-from learning3d.models import FlowNet3D
-from learning3d.data_utils import SceneflowDataset
-import numpy as np
-from torch.utils.data import DataLoader
-from tensorboardX import SummaryWriter
-from tqdm import tqdm
-
-def display_open3d(template, source, transformed_source):
-    template_ = o3d.geometry.PointCloud()
-    source_ = o3d.geometry.PointCloud()
-    transformed_source_ = o3d.geometry.PointCloud()
-    template_.points = o3d.utility.Vector3dVector(template)
-    source_.points = o3d.utility.Vector3dVector(source + np.array([0,0.5,0.5]))
-    transformed_source_.points = o3d.utility.Vector3dVector(transformed_source)
-    template_.paint_uniform_color([1, 0, 0])
-    source_.paint_uniform_color([0, 1, 0])
-    transformed_source_.paint_uniform_color([0, 0, 1])
-    o3d.visualization.draw_geometries([template_, source_, transformed_source_])
-
-def test_one_epoch(args, net, test_loader):
-    net.eval()
-
-    total_loss = 0
-    num_examples = 0
-    for i, data in enumerate(tqdm(test_loader)):
-        data = [d.to(args.device) for d in data]
-        pc1, pc2, color1, color2, flow, mask1 = data
-        pc1 = pc1.transpose(2,1).contiguous()
-        pc2 = pc2.transpose(2,1).contiguous()
-        color1 = color1.transpose(2,1).contiguous()
-        color2 = color2.transpose(2,1).contiguous()
-        flow = flow
-        mask1 = mask1.float()
-
-        batch_size = pc1.size(0)
-        num_examples += batch_size
-        flow_pred = net(pc1, pc2, color1, color2).permute(0,2,1)
-        loss_1 = torch.mean(mask1 * torch.sum((flow_pred - flow) * (flow_pred - flow), -1) / 2.0)
-
-        pc1, pc2 = pc1.permute(0,2,1), pc2.permute(0,2,1)
-        pc1_ = pc1 - flow_pred
-        print("Loss: ", loss_1)
-        display_open3d(pc1.detach().cpu().numpy()[0], pc2.detach().cpu().numpy()[0], pc1_.detach().cpu().numpy()[0])
-        total_loss += loss_1.item() * batch_size        
-
-    return total_loss * 1.0 / num_examples
-
-
-def test(args, net, test_loader):
-    test_loss = test_one_epoch(args, net, test_loader)
-
-def main():
-    parser = argparse.ArgumentParser(description='Point Cloud Registration')
-    parser.add_argument('--model', type=str, default='flownet', metavar='N',
-                        choices=['flownet'], help='Model to use, [flownet]')
-    parser.add_argument('--emb_dims', type=int, default=512, metavar='N',
-                        help='Dimension of embeddings')
-    parser.add_argument('--num_points', type=int, default=2048,
-                        help='Point Number [default: 2048]')
-    parser.add_argument('--test_batch_size', type=int, default=1, metavar='batch_size',
-                        help='Size of batch)')
-    
-    parser.add_argument('--gaussian_noise', type=bool, default=False, metavar='N',
-                        help='Wheter to add gaussian noise')
-    parser.add_argument('--unseen', type=bool, default=False, metavar='N',
-                        help='Whether to test on unseen category')
-    parser.add_argument('--dataset', type=str, default='SceneflowDataset',
-                        choices=['SceneflowDataset'], metavar='N',
-                        help='dataset to use')
-    parser.add_argument('--dataset_path', type=str, default='data_processed_maxcut_35_20k_2k_8192', metavar='N',
-                        help='dataset to use')
-    parser.add_argument('--pretrained', type=str, default='learning3d/pretrained/exp_flownet/models/model.best.t7', metavar='N',
-                        help='Pretrained model path')
-    parser.add_argument('--device', default='cuda:0', type=str,
-                        metavar='DEVICE', help='use CUDA if available')
-
-    args = parser.parse_args()
-    if not torch.cuda.is_available():
-        args.device = torch.device('cpu')
-    else:
-        args.device = torch.device('cuda')
-
-    if args.dataset == 'SceneflowDataset':
-        test_loader = DataLoader(
-            SceneflowDataset(npoints=args.num_points, partition='test'),
-            batch_size=args.test_batch_size, shuffle=False, drop_last=False)
-    else:
-        raise Exception("not implemented")
-
-    net = FlowNet3D()
-    assert os.path.exists(args.pretrained), "Pretrained Model Doesn't Exists!"
-    net.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
-    net = net.to(args.device)
-        
-    test(args, net, test_loader)
-    print('FINISH')
-
-
-if __name__ == '__main__':
-    main()

learning3d/examples/train_flownet.py

@@ -1,259 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-
-from __future__ import print_function
-import os
-import gc
-import argparse
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-from torch.optim.lr_scheduler import MultiStepLR
-from learning3d.models import FlowNet3D
-from learning3d.data_utils import SceneflowDataset
-import numpy as np
-from torch.utils.data import DataLoader
-from tensorboardX import SummaryWriter
-from tqdm import tqdm
-
-class IOStream:
-    def __init__(self, path):
-        self.f = open(path, 'a')
-
-    def cprint(self, text):
-        print(text)
-        self.f.write(text + '\n')
-        self.f.flush()
-
-    def close(self):
-        self.f.close()
-
-
-def _init_(args):
-    if not os.path.exists('checkpoints'):
-        os.makedirs('checkpoints')
-    if not os.path.exists('checkpoints/' + args.exp_name):
-        os.makedirs('checkpoints/' + args.exp_name)
-    if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'):
-        os.makedirs('checkpoints/' + args.exp_name + '/' + 'models')
-
-def weights_init(m):
-    classname=m.__class__.__name__
-    if classname.find('Conv2d') != -1:
-        nn.init.kaiming_normal_(m.weight.data)
-    if classname.find('Conv1d') != -1:
-        nn.init.kaiming_normal_(m.weight.data)
-
-def test_one_epoch(args, net, test_loader):
-    net.eval()
-
-    total_loss = 0
-    num_examples = 0
-    for i, data in tqdm(enumerate(test_loader), total=len(test_loader), smoothing=0.9):
-        pc1, pc2, color1, color2, flow, mask1 = data
-        pc1 = pc1.cuda().transpose(2,1).contiguous()
-        pc2 = pc2.cuda().transpose(2,1).contiguous()
-        color1 = color1.cuda().transpose(2,1).contiguous()
-        color2 = color2.cuda().transpose(2,1).contiguous()
-        flow = flow.cuda()
-        mask1 = mask1.cuda().float()
-
-        batch_size = pc1.size(0)
-        num_examples += batch_size
-        flow_pred = net(pc1, pc2, color1, color2).permute(0,2,1)
-        loss_1 = torch.mean(mask1 * torch.sum((flow_pred - flow) * (flow_pred - flow), -1) / 2.0)
-
-        pc1, pc2 = pc1.permute(0,2,1), pc2.permute(0,2,1)
-        pc1_ = pc1 + flow_pred
-
-        total_loss += loss_1.item() * batch_size
-        
-
-    return total_loss * 1.0 / num_examples
-
-
-def train_one_epoch(args, net, train_loader, opt):
-    net.train()
-    num_examples = 0
-    total_loss = 0
-    for i, data in tqdm(enumerate(train_loader), total=len(train_loader), smoothing=0.9):
-        pc1, pc2, color1, color2, flow, mask1 = data
-        pc1 = pc1.cuda().transpose(2,1).contiguous()
-        pc2 = pc2.cuda().transpose(2,1).contiguous()
-        color1 = color1.cuda().transpose(2,1).contiguous()
-        color2 = color2.cuda().transpose(2,1).contiguous()
-        flow = flow.cuda().transpose(2,1).contiguous()
-        mask1 = mask1.cuda().float()
-
-        batch_size = pc1.size(0)
-        opt.zero_grad()
-        num_examples += batch_size
-        flow_pred = net(pc1, pc2, color1, color2)
-        loss_1 = torch.mean(mask1 * torch.sum((flow_pred - flow) ** 2, 1) / 2.0)
-
-        pc1, pc2, flow_pred = pc1.permute(0,2,1), pc2.permute(0,2,1), flow_pred.permute(0,2,1)
-        pc1_ = pc1 + flow_pred
-
-        loss_1.backward()
-
-        opt.step()
-        total_loss += loss_1.item() * batch_size
-
-        # if (i+1) % 100 == 0:
-        #     print("batch: %d, mean loss: %f" % (i, total_loss / 100 / batch_size))
-        #     total_loss = 0
-    return total_loss * 1.0 / num_examples
-
-
-def test(args, net, test_loader, boardio, textio):
-
-    test_loss = test_one_epoch(args, net, test_loader)
-
-    textio.cprint('==FINAL TEST==')
-    textio.cprint('mean test loss: %f'%test_loss)
-
-
-def train(args, net, train_loader, test_loader, boardio, textio):
-    if args.use_sgd:
-        print("Use SGD")
-        opt = optim.SGD(net.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=1e-4)
-    else:
-        print("Use Adam")
-        opt = optim.Adam(net.parameters(), lr=args.lr, weight_decay=1e-4)
-    scheduler = MultiStepLR(opt, milestones=[75, 150, 200], gamma=0.1)
-
-    best_test_loss = np.inf
-    for epoch in range(args.epochs):
-        scheduler.step()
-        textio.cprint('==epoch: %d=='%epoch)
-        train_loss = train_one_epoch(args, net, train_loader, opt)
-        textio.cprint('mean train EPE loss: %f'%train_loss)
-
-        test_loss = test_one_epoch(args, net, test_loader)
-        textio.cprint('mean test EPE loss: %f'%test_loss)
-
-        if best_test_loss >= test_loss:
-            best_test_loss = test_loss
-            textio.cprint('best test loss till now: %f'%test_loss)
-            if torch.cuda.device_count() > 1:
-                torch.save(net.module.state_dict(), 'checkpoints/%s/models/model.best.t7' % args.exp_name)
-            else:
-                torch.save(net.state_dict(), 'checkpoints/%s/models/model.best.t7' % args.exp_name)
-
-        boardio.add_scalar('Train Loss', train_loss, epoch+1)
-        boardio.add_scalar('Test Loss', test_loss, epoch+1)
-        boardio.add_scalar('Best Test Loss', best_test_loss, epoch+1)
-        
-        if torch.cuda.device_count() > 1:
-            torch.save(net.module.state_dict(), 'checkpoints/%s/models/model.%d.t7' % (args.exp_name, epoch))
-        else:
-            torch.save(net.state_dict(), 'checkpoints/%s/models/model.%d.t7' % (args.exp_name, epoch))
-        gc.collect()
-
-
-def main():
-    parser = argparse.ArgumentParser(description='Point Cloud Registration')
-    parser.add_argument('--exp_name', type=str, default='exp_flownet', metavar='N',
-                        help='Name of the experiment')
-    parser.add_argument('--model', type=str, default='flownet', metavar='N',
-                        choices=['flownet'],
-                        help='Model to use, [flownet]')
-    parser.add_argument('--emb_dims', type=int, default=512, metavar='N',
-                        help='Dimension of embeddings')
-    parser.add_argument('--num_points', type=int, default=2048,
-                        help='Point Number [default: 2048]')
-    parser.add_argument('--dropout', type=float, default=0.5, metavar='N',
-                        help='Dropout ratio in transformer')
-    parser.add_argument('--batch_size', type=int, default=16, metavar='batch_size',
-                        help='Size of batch)')
-    parser.add_argument('--test_batch_size', type=int, default=10, metavar='batch_size',
-                        help='Size of batch)')
-    parser.add_argument('--epochs', type=int, default=250, metavar='N',
-                        help='number of episode to train ')
-    parser.add_argument('--use_sgd', action='store_true', default=True,
-                        help='Use SGD')
-    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
-                        help='learning rate (default: 0.001, 0.1 if using sgd)')
-    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
-                        help='SGD momentum (default: 0.9)')
-    parser.add_argument('--no_cuda', action='store_true', default=False,
-                        help='enables CUDA training')
-    parser.add_argument('--seed', type=int, default=1234, metavar='S',
-                        help='random seed (default: 1)')
-    parser.add_argument('--eval', action='store_true', default=False,
-                        help='evaluate the model')
-    parser.add_argument('--cycle', type=bool, default=False, metavar='N',
-                        help='Whether to use cycle consistency')
-    parser.add_argument('--gaussian_noise', type=bool, default=False, metavar='N',
-                        help='Wheter to add gaussian noise')
-    parser.add_argument('--unseen', type=bool, default=False, metavar='N',
-                        help='Whether to test on unseen category')
-    parser.add_argument('--dataset', type=str, default='SceneflowDataset',
-                        choices=['SceneflowDataset'], metavar='N',
-                        help='dataset to use')
-    parser.add_argument('--dataset_path', type=str, default='data_processed_maxcut_35_20k_2k_8192', metavar='N',
-                        help='dataset to use')
-    parser.add_argument('--model_path', type=str, default='', metavar='N',
-                        help='Pretrained model path')
-    parser.add_argument('--pretrained', type=str, default='', metavar='N',
-                        help='Pretrained model path')
-
-    args = parser.parse_args()
-    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
-    # CUDA settings
-    torch.backends.cudnn.deterministic = True
-    torch.manual_seed(args.seed)
-    torch.cuda.manual_seed_all(args.seed)
-    np.random.seed(args.seed)
-
-    boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name)
-    _init_(args)
-
-    textio = IOStream('checkpoints/' + args.exp_name + '/run.log')
-    textio.cprint(str(args))
-
-    if args.dataset == 'SceneflowDataset':
-        train_loader = DataLoader(
-            SceneflowDataset(npoints=args.num_points, partition='train'),
-            batch_size=args.batch_size, shuffle=True, drop_last=True)
-        test_loader = DataLoader(
-            SceneflowDataset(npoints=args.num_points, partition='test'),
-            batch_size=args.test_batch_size, shuffle=False, drop_last=False)
-    else:
-        raise Exception("not implemented")
-
-    if args.model == 'flownet':
-        net = FlowNet3D().cuda()
-        net.apply(weights_init)
-        if args.pretrained:
-            net.load_state_dict(torch.load(args.pretrained), strict=False)
-            print("Pretrained Model Loaded Successfully!")
-        if args.eval:
-            if args.model_path is '':
-                model_path = 'checkpoints' + '/' + args.exp_name + '/models/model.best.t7'
-            else:
-                model_path = args.model_path
-                print(model_path)
-            if not os.path.exists(model_path):
-                print("can't find pretrained model")
-                return
-            net.load_state_dict(torch.load(model_path), strict=False)
-        if torch.cuda.device_count() > 1:
-            net = nn.DataParallel(net)
-            print("Let's use", torch.cuda.device_count(), "GPUs!")
-    else:
-        raise Exception('Not implemented')
-    if args.eval:
-        test(args, net, test_loader, boardio, textio)
-    else:
-        train(args, net, train_loader, test_loader, boardio, textio)
-
-
-    print('FINISH')
-    # boardio.close()
-
-
-if __name__ == '__main__':
-    main()