learning3d 0.0.2__tar.gz → 0.0.4__tar.gz

{learning3d-0.0.2/src/learning3d.egg-info → learning3d-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: learning3d
-Version: 0.0.2
+Version: 0.0.4
 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 → learning3d-0.0.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "learning3d"
-version = "0.0.2"
+version = "0.0.4"
 authors = [
     { name="Vinit Sarode", email="vinitsarode5@gmail.com"},
 ]

{learning3d-0.0.2 → learning3d-0.0.4}/src/learning3d/data_utils/dataloaders.py

@@ -16,9 +16,8 @@ from scipy.spatial.distance import minkowski
 from scipy.spatial import cKDTree
 from torch.utils.data import Dataset
 
-def download_modelnet40():
-	BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-	DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
+def download_modelnet40(root_dir):
+	DATA_DIR = os.path.join(root_dir, 'data')
 	if not os.path.exists(DATA_DIR):
 		os.mkdir(DATA_DIR)
 	if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')):
@@ -28,11 +27,10 @@ def download_modelnet40():
 		os.system('mv %s %s' % (zipfile[:-4], DATA_DIR))
 		os.system('rm %s' % (zipfile))
 
-def load_data(train, use_normals):
+def load_data(root_dir, train, use_normals):
 	if train: partition = 'train'
 	else: partition = 'test'
-	BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-	DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
+	DATA_DIR = os.path.join(root_dir, 'data')
 	all_data = []
 	all_label = []
 	for h5_name in glob.glob(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', 'ply_data_%s*.h5' % partition)):
@@ -186,13 +184,14 @@ class ModelNet40Data(Dataset):
 		self,
 		train=True,
 		num_points=1024,
-		download=True,
+		root_dir=None,
 		randomize_data=False,
 		use_normals=False
 	):
 		super(ModelNet40Data, self).__init__()
-		if download: download_modelnet40()
-		self.data, self.labels = load_data(train, use_normals)
+		self.root_dir = root_dir
+		if root_dir is not None: download_modelnet40(root_dir=root_dir)
+		self.data, self.labels = load_data(root_dir, train, use_normals)
 		if not train: self.shapes = self.read_classes_ModelNet40()
 		self.num_points = num_points
 		self.randomize_data = randomize_data
@@ -218,8 +217,7 @@ class ModelNet40Data(Dataset):
 		return self.shapes[label]
 
 	def read_classes_ModelNet40(self):
-		BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-		DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
+		DATA_DIR = os.path.join(self.root_dir, 'data')
 		file = open(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', 'shape_names.txt'), 'r')
 		shape_names = file.read()
 		shape_names = np.array(shape_names.split('\n')[:-1])
@@ -362,10 +360,9 @@ class FlowData(Dataset):
 
 
 class SceneflowDataset(Dataset):
-	def __init__(self, npoints=1024, root='', partition='train'):
+	def __init__(self, root_dir, npoints=1024, root='', partition='train'):
 		if root == '':
-			BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-			DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
+			DATA_DIR = os.path.join(root_dir, 'data')
 			root = os.path.join(DATA_DIR, 'data_processed_maxcut_35_20k_2k_8192')
 			if not os.path.exists(root): 
 				print("To download dataset, click here: https://drive.google.com/file/d/1CMaxdt-Tg1Wct8v8eGNwuT7qRSIyJPY-/view")

{learning3d-0.0.2 → learning3d-0.0.4}/src/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-0.0.2 → learning3d-0.0.4}/src/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 → learning3d-0.0.4/src/learning3d.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: learning3d
-Version: 0.0.2
+Version: 0.0.4
 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 → learning3d-0.0.4}/src/learning3d.egg-info/SOURCES.txt

@@ -6,7 +6,6 @@ requirements.txt
 data/modelnet40_ply_hdf5_2048/shape_names.txt
 data/modelnet40_ply_hdf5_2048/test_files.txt
 data/modelnet40_ply_hdf5_2048/train_files.txt
-pretrained/exp_prnet/args.txt
 src/learning3d/__init__.py
 src/learning3d.egg-info/PKG-INFO
 src/learning3d.egg-info/SOURCES.txt
@@ -18,7 +17,6 @@ src/learning3d/data_utils/dataloaders.py
 src/learning3d/data_utils/user_data.py
 src/learning3d/examples/test_dcp.py
 src/learning3d/examples/test_deepgmr.py
-src/learning3d/examples/test_flownet.py
 src/learning3d/examples/test_masknet.py
 src/learning3d/examples/test_masknet2.py
 src/learning3d/examples/test_pcn.py
@@ -31,7 +29,6 @@ src/learning3d/examples/test_rpmnet.py
 src/learning3d/examples/train_PointNetLK.py
 src/learning3d/examples/train_dcp.py
 src/learning3d/examples/train_deepgmr.py
-src/learning3d/examples/train_flownet.py
 src/learning3d/examples/train_masknet.py
 src/learning3d/examples/train_pcn.py
 src/learning3d/examples/train_pcrnet.py
@@ -64,7 +61,6 @@ src/learning3d/models/classifier.py
 src/learning3d/models/dcp.py
 src/learning3d/models/deepgmr.py
 src/learning3d/models/dgcnn.py
-src/learning3d/models/flownet3d.py
 src/learning3d/models/masknet.py
 src/learning3d/models/masknet2.py
 src/learning3d/models/pcn.py
@@ -89,36 +85,4 @@ src/learning3d/utils/__init__.py
 src/learning3d/utils/pointconv_util.py
 src/learning3d/utils/ppfnet_util.py
 src/learning3d/utils/svd.py
-src/learning3d/utils/transformer.py
-src/learning3d/utils/lib/pointnet2_modules.py
-src/learning3d/utils/lib/pointnet2_utils.py
-src/learning3d/utils/lib/pytorch_utils.py
-src/learning3d/utils/lib/setup.py
-src/learning3d/utils/lib/build/lib.linux-x86_64-3.5/pointnet2_cuda.cpython-35m-x86_64-linux-gnu.so
-src/learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/ball_query.o
-src/learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/ball_query_gpu.o
-src/learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/group_points.o
-src/learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/group_points_gpu.o
-src/learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/interpolate.o
-src/learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/interpolate_gpu.o
-src/learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/pointnet2_api.o
-src/learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/sampling.o
-src/learning3d/utils/lib/build/temp.linux-x86_64-3.5/src/sampling_gpu.o
-src/learning3d/utils/lib/dist/pointnet2-0.0.0-py3.5-linux-x86_64.egg
-src/learning3d/utils/lib/pointnet2.egg-info/SOURCES.txt
-src/learning3d/utils/lib/pointnet2.egg-info/dependency_links.txt
-src/learning3d/utils/lib/pointnet2.egg-info/top_level.txt
-src/learning3d/utils/lib/src/ball_query.cpp
-src/learning3d/utils/lib/src/ball_query_gpu.cu
-src/learning3d/utils/lib/src/ball_query_gpu.h
-src/learning3d/utils/lib/src/cuda_utils.h
-src/learning3d/utils/lib/src/group_points.cpp
-src/learning3d/utils/lib/src/group_points_gpu.cu
-src/learning3d/utils/lib/src/group_points_gpu.h
-src/learning3d/utils/lib/src/interpolate.cpp
-src/learning3d/utils/lib/src/interpolate_gpu.cu
-src/learning3d/utils/lib/src/interpolate_gpu.h
-src/learning3d/utils/lib/src/pointnet2_api.cpp
-src/learning3d/utils/lib/src/sampling.cpp
-src/learning3d/utils/lib/src/sampling_gpu.cu
-src/learning3d/utils/lib/src/sampling_gpu.h
+src/learning3d/utils/transformer.py

learning3d-0.0.2/pretrained/exp_prnet/args.txt

@@ -1,35 +0,0 @@
-{
-  "exp_name": "exp_multi_gpu_v1",
-  "model": "prnet",
-  "emb_nn": "dgcnn",
-  "attention": "transformer",
-  "head": "svd",
-  "n_emb_dims": 512,
-  "n_blocks": 1,
-  "n_heads": 4,
-  "n_iters": 3,
-  "discount_factor": 0.9,
-  "n_ff_dims": 1024,
-  "n_keypoints": 512,
-  "temp_factor": 100,
-  "cat_sampler": "softmax",
-  "dropout": 0.0,
-  "batch_size": 16,
-  "test_batch_size": 8,
-  "epochs": 100,
-  "use_sgd": false,
-  "lr": 0.001,
-  "momentum": 0.9,
-  "no_cuda": false,
-  "seed": 1234,
-  "eval": false,
-  "cycle_consistency_loss": 0.1,
-  "feature_alignment_loss": 0.1,
-  "gaussian_noise": false,
-  "unseen": false,
-  "n_points": 1024,
-  "n_subsampled_points": 768,
-  "dataset": "modelnet40",
-  "rot_factor": 4,
-  "model_path": ""
-}

learning3d-0.0.2/src/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-0.0.2/src/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-0.0.2/src/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()