├── add_noise_to_data ├── __init__.py ├── interface.py └── random_noise.py ├── generation ├── deepul │ ├── __init__.py │ ├── utils.py │ ├── pytorch_util.py │ └── hw4_helper.py ├── test_generation.sh ├── train_latent_generator.sh ├── preprocess.sh ├── preprocess_test.json ├── preprocess_train.json ├── test_generation_config.json ├── train_latent_generator_config.json ├── preprocess.py ├── test_generation.py └── train_latent_generator.py ├── metrics_from_point_flow ├── __init__.py ├── pytorch_structural_losses │ ├── __init__.py │ ├── src │ │ ├── nndistance.cuh │ │ ├── approxmatch.cuh │ │ ├── utils.hpp │ │ ├── nndistance.cu │ │ ├── structural_loss.cpp │ │ └── approxmatch.cu │ ├── pybind │ │ ├── bind.cpp │ │ └── extern.hpp │ ├── setup.py │ ├── nn_distance.py │ ├── match_cost.py │ └── Makefile └── evaluation_metrics.py ├── criteria_comparing_sets_pcs ├── __init__.py ├── jsd_calculator.py └── all_metrics_calculator.py ├── evaluator ├── evaluator_based_on_comparing_set_pcs │ ├── __init__.py │ ├── interface.py │ ├── JSD_based_evaluator.py │ └── set_pc_comparing_based_evaluator.py ├── __init__.py └── aeevaluator.py ├── logger ├── __init__.py └── tensorboard.py ├── saver ├── __init__.py ├── interface.py └── saver.py ├── trainer ├── __init__.py ├── interface.py └── trainer.py ├── image └── teaser.png ├── models ├── __init__.py ├── utils.py ├── pointnet.py └── pcn.py ├── loss ├── __init__.py ├── emd.py ├── chamfer.py └── sw_variants.py ├── utils ├── __init__.py └── utils.py ├── registration ├── registration_config.json ├── writer.py ├── preprocess_config.json ├── preprocess.sh ├── register.sh ├── matcher.py ├── preprocessor.py ├── registration_test.py └── preprocess_data.py ├── dataset ├── __init__.py ├── download_shapenet_chair.sh ├── download_shapenet_core55_catagories.sh ├── interface.py ├── download_modelnet40_same_with_pointnet.sh ├── utils.py ├── fine3dmatch.py ├── raw3dmatch.py ├── shapenet_core55.py ├── download_3dmatch.sh └── modelnet40.py ├── classification ├── linear_svm_cls.sh ├── classify_train_test.sh ├── class_test_config.json ├── linear_svm_config.json ├── preprocess.sh ├── class_train_config.json ├── preprocess_test.json ├── preprocess_train.json ├── classifier.py ├── classification_test.py ├── linear_svm_classification.py ├── preprocess_data.py └── classification_train.py ├── reconstruction ├── test.sh ├── config.json └── reconstruction_test.py ├── train.sh ├── .gitignore ├── LICENSE ├── config.json ├── README.md └── train.py /add_noise_to_data/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /generation/deepul/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /metrics_from_point_flow/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /criteria_comparing_sets_pcs/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /evaluator/evaluator_based_on_comparing_set_pcs/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /logger/__init__.py: -------------------------------------------------------------------------------- 1 | from .tensorboard import Logger 2 | -------------------------------------------------------------------------------- /saver/__init__.py: -------------------------------------------------------------------------------- 1 | from .saver import GeneralSaver, Saver 2 | -------------------------------------------------------------------------------- /trainer/__init__.py: -------------------------------------------------------------------------------- 1 | from .trainer import AETrainer, ClassifierTrainer 2 | -------------------------------------------------------------------------------- /image/teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VinAIResearch/PointSWD/HEAD/image/teaser.png -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | from .pcn import PointCapsNet 2 | from .pointnet import PointNetAE 3 | -------------------------------------------------------------------------------- /loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .chamfer import Chamfer 2 | from .emd import EMD 3 | from .sw_variants import ASW, SWD, GenSW, MaxSW 4 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .utils import create_save_folder, evaluate_on_dataset, get_lr, initialize_main, load_model_for_evaluation 2 | -------------------------------------------------------------------------------- /generation/test_generation.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/bash 2 | python generation/test_generation.py --config='generation/test_generation_config.json' \ 3 | --logdir="logs/" -------------------------------------------------------------------------------- /registration/registration_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "seed": 1, 3 | "threshold": 0.075, 4 | "fitness_threshold": 0.3, 5 | "rmse_threshold": 0.2 6 | } -------------------------------------------------------------------------------- /generation/train_latent_generator.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/bash 2 | CUDA_VISIBLE_DEVICES=0 python generation/train_latent_generator.py --seed=1 \ 3 | --logdir="logs/" 4 | -------------------------------------------------------------------------------- /dataset/__init__.py: -------------------------------------------------------------------------------- 1 | from .fine3dmatch import Fine3dMatchDataset 2 | from .modelnet40 import ModelNet40 3 | from .shapenet_core55 import ShapeNetCore55XyzOnlyDataset 4 | -------------------------------------------------------------------------------- /classification/linear_svm_cls.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/bash 2 | python classification/linear_svm_classification.py --config='classification/linear_svm_config.json' \ 3 | --logdir="logs/" -------------------------------------------------------------------------------- /reconstruction/test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | python reconstruction/reconstruction_test.py --config="reconstruction/config.json" --logdir="logs/" --data_path="dataset/modelnet40_ply_hdf5_2048/" 3 | -------------------------------------------------------------------------------- /evaluator/__init__.py: -------------------------------------------------------------------------------- 1 | from .aeevaluator import AEEvaluator as Evaluator 2 | from .evaluator_based_on_comparing_set_pcs.set_pc_comparing_based_evaluator import SetPcsComparingBasedEvaluator 3 | -------------------------------------------------------------------------------- /train.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=0 python train.py --config="config.json" --logdir="logs/" --data_path="dataset/shapenet_core55/shapenet57448xyzonly.npz" --loss="swd" --autoencoder="pointnet" 2 | -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/__init__.py: -------------------------------------------------------------------------------- 1 | # import torch 2 | 3 | # from MakePytorchBackend import AddGPU, Foo, ApproxMatch 4 | 5 | # from Add import add_gpu, approx_match 6 | -------------------------------------------------------------------------------- /trainer/interface.py: -------------------------------------------------------------------------------- 1 | import abc 2 | 3 | 4 | class Trainer(metaclass=abc.ABCMeta): 5 | @classmethod 6 | def __subclasshook__(cls, subclass): 7 | return hasattr(subclass, "train") and callable(subclass.train) 8 | -------------------------------------------------------------------------------- /dataset/download_shapenet_chair.sh: -------------------------------------------------------------------------------- 1 | SCRIPT=`realpath $0` 2 | SCRIPTPATH=`dirname $SCRIPT` 3 | 4 | cd $SCRIPTPATH 5 | gdown https://drive.google.com/uc?id=1BIgSWmYfJov7ov0p3v7zq1QUVlZvpWr8 6 | unzip shapenet_chair.zip 7 | rm shapenet_chair.zip -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | StructuralLosses/ 3 | logs/ 4 | home1/ 5 | home2/ 6 | hotel1/ 7 | hotel2/ 8 | hotel3/ 9 | kitchen/ 10 | lab/ 11 | study/ 12 | modelnet40_ply_hdf5_2048/ 13 | shapenet_chair/ 14 | shapenet_core55/ 15 | khongcan/ 16 | -------------------------------------------------------------------------------- /add_noise_to_data/interface.py: -------------------------------------------------------------------------------- 1 | import abc 2 | 3 | 4 | class NoiseAdd_er(metaclass=abc.ABCMeta): 5 | @classmethod 6 | def __subclasshook__(cls, subclass): 7 | return hasattr(subclass, "add_noise") and callable(subclass.add_noise) 8 | -------------------------------------------------------------------------------- /dataset/download_shapenet_core55_catagories.sh: -------------------------------------------------------------------------------- 1 | SCRIPT=`realpath $0` 2 | SCRIPTPATH=`dirname $SCRIPT` 3 | 4 | cd $SCRIPTPATH 5 | gdown https://drive.google.com/uc?id=1sJd5bdCg9eOo3-FYtchUVlwDgpVdsbXB 6 | mkdir -p shapenet_core55 7 | mv shapenet57448xyzonly.npz shapenet_core55/ -------------------------------------------------------------------------------- /evaluator/evaluator_based_on_comparing_set_pcs/interface.py: -------------------------------------------------------------------------------- 1 | import abc 2 | 3 | import six 4 | 5 | 6 | @six.add_metaclass(abc.ABCMeta) 7 | class Evaluator: # (metaclass=abc.ABCMeta): 8 | @classmethod 9 | def __subclasshook__(cls, subclass): 10 | return hasattr(subclass, "evaluate") and callable(subclass.evaluate) 11 | -------------------------------------------------------------------------------- /classification/classify_train_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/bash 2 | CUDA_VISIBLE_DEVICES=0 python classification/classification_train.py --config='classification/class_train_config.json' \ 3 | --logdir="logs/" 4 | 5 | CUDA_VISIBLE_DEVICES=0 python classification/classification_test.py --config='classification/class_test_config.json' \ 6 | --logdir="logs/" 7 | -------------------------------------------------------------------------------- /generation/preprocess.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/bash 2 | CUDA_VISIBLE_DEVICES=0 python generation/preprocess.py --config='generation/preprocess_train.json' --logdir="logs/" --data_path="dataset/shapenet_chair/train.npz" 3 | 4 | CUDA_VISIBLE_DEVICES=0 python generation/preprocess.py --config='generation/preprocess_test.json' --logdir="logs/" --data_path="dataset/shapenet_chair/test.npz" 5 | -------------------------------------------------------------------------------- /saver/interface.py: -------------------------------------------------------------------------------- 1 | import abc 2 | 3 | 4 | class Saver(metaclass=abc.ABCMeta): 5 | @classmethod 6 | def __subclasshook__(cls, subclass): 7 | return ( 8 | hasattr(subclass, "save_checkpoint") 9 | and callable(subclass.save_checkpoint) 10 | and hasattr(subclass, "save_best_weights") 11 | and callable(subclass.save_best_weights) 12 | ) 13 | -------------------------------------------------------------------------------- /classification/class_test_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "save_folder": "classifier_202_148_94", 3 | "hidden_sizes": "202,148,94", 4 | "model_path":"model.pth", 5 | 6 | "seed": 3598, 7 | 8 | "device": "cuda", 9 | 10 | "batch_size": 256, 11 | "pin_memory": true, 12 | "num_workers": 0, 13 | "shuffle": false, 14 | 15 | "input_size": 256, 16 | "output_size": 40, 17 | "dropout_p": 0 18 | } -------------------------------------------------------------------------------- /classification/linear_svm_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "save_folder": "linear_svm/", 3 | "train_root": "latent_codes/model/modelnet40-train/saved_latent_vectors.npz", 4 | "test_root": "latent_codes/model/modelnet40-test/saved_latent_vectors.npz", 5 | 6 | "device": "cuda", 7 | 8 | "batch_size": 256, 9 | "pin_memory": true, 10 | "num_workers": 0, 11 | "shuffle": true, 12 | 13 | "input_size": 4096 14 | } -------------------------------------------------------------------------------- /classification/preprocess.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/bash 2 | CUDA_VISIBLE_DEVICES=0 python classification/preprocess_data.py --config='classification/preprocess_train.json' \ 3 | --logdir="logs/" \ 4 | --data_path="dataset/modelnet40_ply_hdf5_2048/train/" 5 | 6 | CUDA_VISIBLE_DEVICES=0 python classification/preprocess_data.py --config='classification/preprocess_test.json' \ 7 | --logdir="logs/" \ 8 | --data_path="dataset/modelnet40_ply_hdf5_2048/test/" -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/src/nndistance.cuh: -------------------------------------------------------------------------------- 1 | void nndistance(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream); 2 | void nndistancegrad(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream); 3 | -------------------------------------------------------------------------------- /dataset/interface.py: -------------------------------------------------------------------------------- 1 | import abc 2 | 3 | import six 4 | 5 | 6 | @six.add_metaclass(abc.ABCMeta) 7 | class Dataset: # (metaclass=abc.ABCMeta): 8 | @classmethod 9 | def __subclasshook__(cls, subclass): 10 | return ( 11 | hasattr(subclass, "__len__") 12 | and callable(subclass.__len__) 13 | and hasattr(subclass, "__getitem__") 14 | and callable(subclass.__getitem__) 15 | ) 16 | -------------------------------------------------------------------------------- /generation/preprocess_test.json: -------------------------------------------------------------------------------- 1 | { 2 | "root": "dataset/shapenet_chair/test.npz", 3 | "save_folder":"shapenet_chair/test/", 4 | "seed": 1, 5 | 6 | "model_path": "model.pth", 7 | "device": "cuda", 8 | 9 | "architecture": "pointnet", 10 | "embedding_size": 256, 11 | "point_channels": 3, 12 | "num_points": 2048, 13 | "normalize": true, 14 | 15 | "dataset": "shapenetcore55", 16 | "batch_size": 64 17 | } -------------------------------------------------------------------------------- /generation/preprocess_train.json: -------------------------------------------------------------------------------- 1 | { 2 | "root": "dataset/shapenet_chair/train.npz", 3 | "save_folder":"shapenet_chair/train/", 4 | "seed": 1, 5 | 6 | "model_path": "model.pth", 7 | "device": "cuda", 8 | 9 | "architecture": "pointnet", 10 | "embedding_size": 256, 11 | "point_channels": 3, 12 | "num_points": 2048, 13 | "normalize": true, 14 | 15 | "dataset": "shapenetcore55", 16 | "batch_size": 64 17 | } -------------------------------------------------------------------------------- /dataset/download_modelnet40_same_with_pointnet.sh: -------------------------------------------------------------------------------- 1 | SCRIPT=`realpath $0` 2 | SCRIPTPATH=`dirname $SCRIPT` 3 | 4 | cd $SCRIPTPATH/ 5 | wget https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip --no-check-certificate 6 | unzip modelnet40_ply_hdf5_2048.zip 7 | rm modelnet40_ply_hdf5_2048.zip 8 | cd $SCRIPTPATH/modelnet40_ply_hdf5_2048/ 9 | mkdir -p train 10 | rsync -avzi ply_data_train*.h5 train/ 11 | mkdir -p test 12 | rsync -avzi ply_data_test*.h5 test/ -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/pybind/bind.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | 5 | #include "pybind/extern.hpp" 6 | 7 | namespace py = pybind11; 8 | 9 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){ 10 | m.def("ApproxMatch", &ApproxMatch); 11 | m.def("MatchCost", &MatchCost); 12 | m.def("MatchCostGrad", &MatchCostGrad); 13 | m.def("NNDistance", &NNDistance); 14 | m.def("NNDistanceGrad", &NNDistanceGrad); 15 | } 16 | -------------------------------------------------------------------------------- /registration/writer.py: -------------------------------------------------------------------------------- 1 | class Writer: 2 | def __init__(self): 3 | super().__init__() 4 | 5 | @staticmethod 6 | def write_arr_to_file(txtfile, meta, arr): 7 | writer = open(txtfile, "a") 8 | writer.write(meta) 9 | n, m = arr.shape 10 | for i in range(n): 11 | for j in range(m): 12 | writer.write(str(arr[i][j])) 13 | writer.write(" ") 14 | writer.write("\n") 15 | writer.close() 16 | return 17 | -------------------------------------------------------------------------------- /dataset/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def pc_normalize(pc): 5 | centroid = np.mean(pc, axis=0) 6 | pc = pc - centroid 7 | m = np.max(np.sqrt(np.sum(pc ** 2, axis=1))) 8 | pc = pc / m 9 | return pc 10 | 11 | 12 | def sample_pc(pc, num_points): 13 | choice = np.random.choice(pc.shape[0], num_points, replace=True) 14 | return pc[choice, :] 15 | 16 | 17 | if __name__ == "__main__": 18 | opc = 15.0 * np.random.rand(2, 3) 19 | print(opc) 20 | print(pc_normalize(opc)) 21 | -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/pybind/extern.hpp: -------------------------------------------------------------------------------- 1 | std::vector ApproxMatch(at::Tensor in_a, at::Tensor in_b); 2 | at::Tensor MatchCost(at::Tensor set_d, at::Tensor set_q, at::Tensor match); 3 | std::vector MatchCostGrad(at::Tensor set_d, at::Tensor set_q, at::Tensor match); 4 | 5 | std::vector NNDistance(at::Tensor set_d, at::Tensor set_q); 6 | std::vector NNDistanceGrad(at::Tensor set_d, at::Tensor set_q, at::Tensor idx1, at::Tensor idx2, at::Tensor grad_dist1, at::Tensor grad_dist2); 7 | -------------------------------------------------------------------------------- /registration/preprocess_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "device": "cuda", 3 | 4 | "seed": 1, 5 | 6 | "model_path": "model.pth", 7 | "save_folder":"model", 8 | 9 | "autoencoder":"pointnet", 10 | "embedding_size": 256, 11 | "input_channels": 3, 12 | "output_channels": 3, 13 | "num_points": 2048, 14 | "normalize": true, 15 | 16 | "root": "dataset/", 17 | 18 | "num_workers": 0, 19 | 20 | "voxel_size": 0.1, 21 | "radius": 0.15, 22 | "batch_size": 256, 23 | 24 | "color": false 25 | } -------------------------------------------------------------------------------- /logger/tensorboard.py: -------------------------------------------------------------------------------- 1 | from torch.utils.tensorboard import SummaryWriter 2 | 3 | 4 | class Logger(object): 5 | def __init__(self, log_dir): 6 | """Create a summary writer logging to log_dir.""" 7 | self.writer = SummaryWriter(log_dir=log_dir) 8 | 9 | def scalar_summary(self, tag, value, step): 10 | """Log a scalar variable.""" 11 | self.writer.add_scalar(tag, value, step) 12 | 13 | def image_summary(self, tag, images, step): 14 | """Log a list of images.""" 15 | pass 16 | 17 | def close(self): 18 | self.writer.close() 19 | -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/src/approxmatch.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | template 3 | void AddGPUKernel(Dtype *in_a, Dtype *in_b, Dtype *out_c, int N, 4 | cudaStream_t stream); 5 | */ 6 | void approxmatch(int b,int n,int m,const float * xyz1,const float * xyz2,float * match,float * temp, cudaStream_t stream); 7 | void matchcost(int b,int n,int m,const float * xyz1,const float * xyz2,float * match, float * out, cudaStream_t stream); 8 | void matchcostgrad(int b,int n,int m,const float * xyz1,const float * xyz2,const float * match,float * grad1,float * grad2, cudaStream_t stream); 9 | -------------------------------------------------------------------------------- /evaluator/aeevaluator.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch 5 | 6 | 7 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 8 | 9 | 10 | class AEEvaluator: 11 | def __init__(self): 12 | super().__init__() 13 | 14 | @staticmethod 15 | def evaluate(autoencoder, val_data, loss_func, **kwargs): 16 | autoencoder.eval() 17 | with torch.no_grad(): 18 | latent = autoencoder.encode(val_data) 19 | reconstruction = autoencoder.decode(latent) 20 | _loss = loss_func.forward(val_data, reconstruction, latent=latent, **kwargs) 21 | return {"evaluation": _loss} 22 | -------------------------------------------------------------------------------- /classification/class_train_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "save_folder": "classifier_202_148_94", 3 | "hidden_sizes": "202,148,94", 4 | 5 | "seed": 3598, 6 | 7 | "device": "cuda", 8 | 9 | "batch_size": 256, 10 | "pin_memory": true, 11 | "num_workers": 0, 12 | "shuffle": true, 13 | 14 | "input_size": 256, 15 | "output_size": 40, 16 | "dropout_p": 0.5, 17 | 18 | "optimizer": "sgd", 19 | "learning_rate": 0.001, 20 | "momentum": 0.9, 21 | "weight_decay": 0.005, 22 | 23 | "best_loss":1e10, 24 | "best_epoch": -1, 25 | "start_epoch": 0, 26 | "num_epochs": 500, 27 | "checkpoint": "", 28 | 29 | "epoch_gap_for_save": 250 30 | } -------------------------------------------------------------------------------- /dataset/fine3dmatch.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os 3 | 4 | import numpy as np 5 | import torch.utils.data as data 6 | 7 | 8 | class Fine3dMatchDataset(data.Dataset): 9 | def __init__(self, root): 10 | super().__init__() 11 | self.root = root 12 | flistname = os.path.join(root, "*.npz") 13 | self.flistname = sorted(glob.glob(flistname), key=lambda fname: int(fname.split("/")[-1].split(".npz")[0])) 14 | self.flist = [np.load(fname, allow_pickle=True)["arr_0"].item() for fname in self.flistname] 15 | 16 | def __len__(self): 17 | return len(self.flistname) 18 | 19 | def __getitem__(self, i): 20 | return self.flist[i]["points"], self.flist[i]["features"] 21 | -------------------------------------------------------------------------------- /loss/emd.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch.nn as nn 5 | 6 | 7 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 8 | from metrics_from_point_flow.evaluation_metrics import emd_approx 9 | 10 | 11 | class EMD(nn.Module): 12 | """ 13 | Auction algorithm to estimate EMD was proposed in paper "A distributed algorithm for the assignment problem" - Dimitri P. Bertsekas, 1979. 14 | """ 15 | 16 | def __init__(self, *args, **kwargs): 17 | super(EMD, self).__init__() 18 | 19 | def forward(self, x, y, **kwargs): 20 | """ 21 | x, y: [batch size, num points in point cloud, 3] 22 | """ 23 | return {"loss": emd_approx(x, y).mean()} 24 | -------------------------------------------------------------------------------- /models/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def init_weights(m): 5 | if isinstance(m, torch.nn.Linear): 6 | torch.nn.init.xavier_uniform_(m.weight) 7 | m.bias.data.fill_(0.01) 8 | elif isinstance(m, torch.nn.Conv2d): 9 | torch.nn.init.xavier_uniform_(m.weight) 10 | if m.bias is not None: 11 | m.bias.data.fill_(0.01) 12 | 13 | 14 | def _count_parameters(model): 15 | total_params = 0 16 | for name, parameter in model.named_parameters(): 17 | if not parameter.requires_grad: 18 | continue 19 | param = parameter.numel() 20 | total_params += param 21 | print(f"Total Trainable Params: {total_params}") 22 | return total_params 23 | -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/src/utils.hpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | class Formatter { 6 | public: 7 | Formatter() {} 8 | ~Formatter() {} 9 | 10 | template Formatter &operator<<(const Type &value) { 11 | stream_ << value; 12 | return *this; 13 | } 14 | 15 | std::string str() const { return stream_.str(); } 16 | operator std::string() const { return stream_.str(); } 17 | 18 | enum ConvertToString { to_str }; 19 | 20 | std::string operator>>(ConvertToString) { return stream_.str(); } 21 | 22 | private: 23 | std::stringstream stream_; 24 | Formatter(const Formatter &); 25 | Formatter &operator=(Formatter &); 26 | }; 27 | -------------------------------------------------------------------------------- /evaluator/evaluator_based_on_comparing_set_pcs/JSD_based_evaluator.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch.nn as nn 5 | 6 | 7 | sys.path.append(osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__))))) 8 | from criteria_comparing_sets_pcs.jsd_calculator import JsdCalculator 9 | 10 | 11 | class JSDBasedEvaluator(nn.Module): 12 | def __init__(self): 13 | super().__init__() 14 | 15 | @staticmethod 16 | def evaluate(autoencoder, val_data, prior_distribution_sampler, **kwargs): 17 | num_samples = val_data.shape[0] 18 | samples = prior_distribution_sampler.sample(num_samples) 19 | synthetic_data = autoencoder.decode(samples) 20 | JSD_evaluation = JsdCalculator.forward(val_data, synthetic_data, **kwargs) 21 | return {"evaluation": JSD_evaluation} 22 | -------------------------------------------------------------------------------- /loss/chamfer.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch 5 | import torch.nn as nn 6 | 7 | 8 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 9 | from metrics_from_point_flow.evaluation_metrics import distChamferCUDA 10 | 11 | 12 | class Chamfer(nn.Module): 13 | def __init__(self, version="mean", **kwargs): 14 | super(Chamfer, self).__init__() 15 | assert version in ["mean", "max"] 16 | self.version = version 17 | 18 | def forward(self, x, y, *args, **kwargs): 19 | min_l, min_r = distChamferCUDA(x.cuda(), y.cuda()) 20 | l_dist = min_l.mean() 21 | r_dist = min_r.mean() 22 | if self.version == "mean": 23 | return {"loss": l_dist + r_dist} 24 | else: 25 | return {"loss": torch.max(l_dist, r_dist)} 26 | -------------------------------------------------------------------------------- /classification/preprocess_test.json: -------------------------------------------------------------------------------- 1 | { 2 | "autoencoder": "pointnet", 3 | "root": "dataset/modelnet40_ply_hdf5_2048/test/", 4 | "save_folder": "latent_codes/model/modelnet40-test/", 5 | 6 | "device": "cuda", 7 | "seed": 1, 8 | "num_points": 2048, 9 | "prim_caps_size": 1024, 10 | "prim_vec_size": 16, 11 | "latent_caps_size": 64, 12 | "latent_vec_size": 64, 13 | 14 | "embedding_size": 256, 15 | "input_channels": 3, 16 | "output_channels": 3, 17 | "normalize": true, 18 | 19 | "model_path":"model.pth", 20 | "dataset":"modelnet40", 21 | 22 | "batch_size": 32, 23 | "pin_memory": true, 24 | "num_workers": 0, 25 | "shuffle": false, 26 | 27 | "add_noise": false, 28 | "noise_adder":"random", 29 | "mean_noiseadder":0.0, 30 | "std_noiseadder": 0.01 31 | } -------------------------------------------------------------------------------- /classification/preprocess_train.json: -------------------------------------------------------------------------------- 1 | { 2 | "autoencoder": "pointnet", 3 | "root": "dataset/modelnet40_ply_hdf5_2048/train/", 4 | "save_folder": "latent_codes/model/modelnet40-train/", 5 | 6 | "device": "cuda", 7 | "seed": 1, 8 | "num_points": 2048, 9 | "prim_caps_size": 1024, 10 | "prim_vec_size": 16, 11 | "latent_caps_size": 64, 12 | "latent_vec_size": 64, 13 | 14 | "embedding_size": 256, 15 | "input_channels": 3, 16 | "output_channels": 3, 17 | "normalize": true, 18 | 19 | "model_path":"model.pth", 20 | "dataset":"modelnet40", 21 | 22 | "batch_size": 32, 23 | "pin_memory": true, 24 | "num_workers": 0, 25 | "shuffle": false, 26 | 27 | "add_noise": false, 28 | "noise_adder":"random", 29 | "mean_noiseadder":0.0, 30 | "std_noiseadder": 0.01 31 | } -------------------------------------------------------------------------------- /add_noise_to_data/random_noise.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch 5 | 6 | 7 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 8 | from add_noise_to_data.interface import NoiseAdd_er 9 | 10 | 11 | class RandomNoiseAdder(NoiseAdd_er): 12 | def __init__(self, mean=0.0, std=0.1, **kwargs): 13 | super().__init__() 14 | self.mean = mean 15 | self.std = std 16 | 17 | def add_noise(self, data, **kwargs): 18 | noise = torch.normal(self.mean, self.std, size=data.shape).to(data.device) 19 | perturbed_data = data + noise 20 | return perturbed_data 21 | 22 | 23 | if __name__ == "__main__": 24 | data = torch.empty((1, 2048, 3)).uniform_(0, 1).cuda() 25 | noise_adder = RandomNoiseAdder() 26 | perturbed_data = noise_adder.add_noise(data) 27 | print(perturbed_data.shape) 28 | -------------------------------------------------------------------------------- /criteria_comparing_sets_pcs/jsd_calculator.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch 5 | import torch.nn as nn 6 | 7 | 8 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 9 | from metrics_from_point_flow.evaluation_metrics import jsd_between_point_cloud_sets 10 | 11 | 12 | class JsdCalculator(nn.Module): 13 | def __init__(self): 14 | super(JsdCalculator, self).__init__() 15 | 16 | @staticmethod 17 | def forward(sample_pcs, ref_pcs, resolution=28, **kwargs): 18 | sample_pcs = sample_pcs.detach().cpu().numpy() 19 | ref_pcs = ref_pcs.detach().cpu().numpy() 20 | return jsd_between_point_cloud_sets(sample_pcs, ref_pcs, resolution) 21 | 22 | 23 | if __name__ == "__main__": 24 | sample_pcs = torch.empty(5, 2048, 3).uniform_(0, 1).numpy() 25 | ref_pcs = torch.empty(5, 2048, 3).uniform_(0, 1).numpy() 26 | print(JsdCalculator.forward(sample_pcs, ref_pcs)) 27 | -------------------------------------------------------------------------------- /generation/test_generation_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "device": "cuda", 3 | 4 | "seed": 1, 5 | 6 | "test_set_type": "shapenetcore55", 7 | "test_root": "dataset/shapenet_chair/test.npz", 8 | "batch_size": 677, 9 | "num_workers": 0, 10 | 11 | "cates": "chair", 12 | "all_points_mean": "", 13 | "all_points_std": "", 14 | 15 | "prior_distribution": "latent_codes_generator", 16 | "latent_dim": 64, 17 | "n_hidden": 3, 18 | "hidden_size": 128, 19 | "save_folder": "shapenet_chair/train/", 20 | "prior_path": "generator.pth", 21 | 22 | "architecture": "pointnet", 23 | "num_points": 2048, 24 | "point_channels": 3, 25 | "embedding_size": 256, 26 | "normalize": true, 27 | "model_path": "model.pth", 28 | 29 | "evaluator_type": "based_on_comparing_set_pcs", 30 | "eval_criteria": "all_metrics", 31 | "criteria_batch_size": 64, 32 | "use_EMD": true, 33 | "accelerated_cd": true 34 | } -------------------------------------------------------------------------------- /dataset/raw3dmatch.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os 3 | 4 | import open3d 5 | import torch.utils.data as data 6 | 7 | 8 | try: 9 | from dataset.interface import Dataset 10 | except: 11 | from interface import Dataset 12 | 13 | 14 | class ThreeDMatchRawDataset(data.Dataset): 15 | def __init__(self, root): 16 | super().__init__() 17 | self.root = root 18 | self.flistname = os.path.join(root, "*.ply") 19 | self.flistname = sorted(glob.glob(self.flistname)) 20 | self.flist = [open3d.io.read_point_cloud(fname) for fname in self.flistname] 21 | self.fnames = [fname.split("/")[-1].split(".ply")[0].split("cloud_bin_")[-1] for fname in self.flistname] 22 | 23 | def __len__(self): 24 | return len(os.listdir(self.root)) 25 | 26 | def __getitem__(self, i): 27 | return self.flist[i], self.fnames[i] 28 | 29 | 30 | Dataset.register(ThreeDMatchRawDataset) 31 | assert issubclass(ThreeDMatchRawDataset, Dataset) 32 | -------------------------------------------------------------------------------- /reconstruction/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "device": "cuda", 3 | 4 | "seed":1, 5 | 6 | "test_origin": true, 7 | 8 | "add_noise": false, 9 | "noise_adder":"random", 10 | "mean_noiseadder":0.0, 11 | "std_noiseadder": 1.0, 12 | 13 | "dataset_type": "modelnet40", 14 | "root": "dataset/modelnet40_ply_hdf5_2048/", 15 | "num_workers": 0, 16 | "shuffle": false, 17 | "num_projs": 100, 18 | "num_points": 2048, 19 | 20 | "architecture": "pointnet", 21 | "model_path": "model.pth", 22 | 23 | "embedding_size": 256, 24 | "input_channels": 3, 25 | 26 | "save_folder": "reconstruction/model/modelnet40/", 27 | "bin": 50, 28 | 29 | "output_channels": 3, 30 | "normalize": true, 31 | 32 | "prim_caps_size": 1024, 33 | "prim_vec_size": 16, 34 | "latent_caps_size": 64, 35 | "latent_vec_size": 64, 36 | 37 | "batch_size": 64, 38 | "criteria_batch_size": 64, 39 | "evaluator_type": "based_on_comparing_set_pcs" 40 | } -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | 5 | # Python interface 6 | setup( 7 | name="PyTorchStructuralLosses", 8 | version="0.1.0", 9 | install_requires=["torch"], 10 | packages=["StructuralLosses"], 11 | package_dir={"StructuralLosses": "./"}, 12 | ext_modules=[ 13 | CUDAExtension( 14 | name="StructuralLossesBackend", 15 | include_dirs=["./"], 16 | sources=[ 17 | "pybind/bind.cpp", 18 | ], 19 | libraries=["make_pytorch"], 20 | library_dirs=["objs"], 21 | # extra_compile_args=['-g'] 22 | ) 23 | ], 24 | cmdclass={"build_ext": BuildExtension}, 25 | author="Christopher B. Choy", 26 | author_email="chrischoy@ai.stanford.edu", 27 | description="Tutorial for Pytorch C++ Extension with a Makefile", 28 | keywords="Pytorch C++ Extension", 29 | url="https://github.com/chrischoy/MakePytorchPlusPlus", 30 | zip_safe=False, 31 | ) 32 | -------------------------------------------------------------------------------- /registration/preprocess.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/bash 2 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 3 | --logdir='logs/' \ 4 | --data_path='dataset/home1' 5 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 6 | --logdir='logs/' \ 7 | --data_path='dataset/home2' 8 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 9 | --logdir='logs/' \ 10 | --data_path='dataset/hotel1' 11 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 12 | --logdir='logs/' \ 13 | --data_path='dataset/hotel2' 14 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 15 | --logdir='logs/' \ 16 | --data_path='dataset/hotel3' 17 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 18 | --logdir='logs/' \ 19 | --data_path='dataset/kitchen' 20 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 21 | --logdir='logs/' \ 22 | --data_path='dataset/lab' 23 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 24 | --logdir='logs/' \ 25 | --data_path='dataset/study' 26 | echo "DONE preprocess shell" -------------------------------------------------------------------------------- /classification/classifier.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class MLPClassifier(nn.Module): 5 | def __init__(self, input_size, output_size, p, hidden_sizes): 6 | super().__init__() 7 | assert len(hidden_sizes) > 0 8 | layers = [] 9 | 10 | layer1 = nn.Linear(input_size, hidden_sizes[0], bias=True) 11 | bn1 = nn.BatchNorm1d(hidden_sizes[0]) 12 | dropout1 = nn.Dropout(p) 13 | 14 | last_layer = nn.Linear(hidden_sizes[-1], output_size, bias=True) 15 | 16 | layers.append(layer1) 17 | layers.append(bn1) 18 | layers.append(dropout1) 19 | layers.append(nn.ReLU()) 20 | 21 | if len(hidden_sizes) > 1: 22 | for i in range(len(hidden_sizes) - 1): 23 | 24 | hidden_layer = nn.Linear(hidden_sizes[i], hidden_sizes[i + 1], bias=True) 25 | bn = nn.BatchNorm1d(hidden_sizes[i + 1]) 26 | dropout = nn.Dropout(p) 27 | 28 | layers.append(hidden_layer) 29 | layers.append(bn) 30 | layers.append(dropout) 31 | layers.append(nn.ReLU()) 32 | layers.append(last_layer) 33 | self.classifier = nn.Sequential(*layers) 34 | 35 | def forward(self, x): 36 | return self.classifier(x) 37 | -------------------------------------------------------------------------------- /evaluator/evaluator_based_on_comparing_set_pcs/set_pc_comparing_based_evaluator.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch 5 | import torch.nn as nn 6 | 7 | 8 | sys.path.append(osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__))))) 9 | from evaluator.evaluator_based_on_comparing_set_pcs.interface import Evaluator 10 | 11 | 12 | class SetPcsComparingBasedEvaluator(nn.Module): 13 | def __init__(self): 14 | super(SetPcsComparingBasedEvaluator, self).__init__() 15 | 16 | @staticmethod 17 | def evaluate(autoencoder, val_data, prior_distribution_sampler, criteria_calculator, **kwargs): 18 | autoencoder.eval() 19 | with torch.no_grad(): 20 | num_samples = val_data.shape[0] 21 | if kwargs.get("model_type") in ["vae"]: 22 | _, synthetic_data = autoencoder(val_data) 23 | else: 24 | samples = prior_distribution_sampler.sample(num_samples) 25 | synthetic_data = autoencoder.decode(samples) 26 | # evaluate 27 | _evaluation = criteria_calculator.forward(val_data, synthetic_data, **kwargs) 28 | return {"evaluation": _evaluation} 29 | 30 | 31 | Evaluator.register(SetPcsComparingBasedEvaluator) 32 | assert issubclass(SetPcsComparingBasedEvaluator, Evaluator) 33 | -------------------------------------------------------------------------------- /registration/register.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | CUDA_VISIBLE_DEVICES=0 python registration/registration_test.py \ 3 | --config='registration/registration_config.json' \ 4 | --logdir='logs/model/home1/' 5 | CUDA_VISIBLE_DEVICES=0 python registration/registration_test.py \ 6 | --config='registration/registration_config.json' \ 7 | --logdir='logs/model/home2/' 8 | CUDA_VISIBLE_DEVICES=0 python registration/registration_test.py \ 9 | --config='registration/registration_config.json' \ 10 | --logdir='logs/model/hotel1/' 11 | CUDA_VISIBLE_DEVICES=0 python registration/registration_test.py \ 12 | --config='registration/registration_config.json' \ 13 | --logdir='logs/model/hotel2/' 14 | CUDA_VISIBLE_DEVICES=0 python registration/registration_test.py \ 15 | --config='registration/registration_config.json' \ 16 | --logdir='logs/model/hotel3/' 17 | CUDA_VISIBLE_DEVICES=0 python registration/registration_test.py \ 18 | --config='registration/registration_config.json' \ 19 | --logdir='logs/model/kitchen/' 20 | CUDA_VISIBLE_DEVICES=0 python registration/registration_test.py \ 21 | --config='registration/registration_config.json' \ 22 | --logdir='logs/model/lab/' 23 | CUDA_VISIBLE_DEVICES=0 python registration/registration_test.py \ 24 | --config='registration/registration_config.json' \ 25 | --logdir='logs/model/study/' 26 | echo "DONE register shell" -------------------------------------------------------------------------------- /saver/saver.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class Saver: 5 | def __init__(self): 6 | super().__init__() 7 | 8 | @staticmethod 9 | def save_checkpoint(ae, optimizer, path, **kwargs): 10 | dic = {"autoencoder": ae.state_dict(), "optimizer": optimizer.state_dict()} 11 | if "scheduler" in kwargs.keys(): 12 | dic["scheduler"] = kwargs["scheduler"].state_dict() 13 | torch.save(dic, path) 14 | return 15 | 16 | @staticmethod 17 | def save_best_weights(ae, path): 18 | torch.save(ae.state_dict(), path) 19 | return 20 | 21 | 22 | class PreprocessedDataSaver: 23 | def __init__(self): 24 | super().__init__() 25 | 26 | def save(self, path, data_dict): 27 | return 28 | 29 | 30 | class GeneralSaver: 31 | def __init__(self): 32 | super().__init__() 33 | 34 | @staticmethod 35 | def save_checkpoint(trained_model, optimizer, path, model_name, **kwargs): 36 | """ 37 | model_name: str 38 | "classifier"|"autoencoder" 39 | """ 40 | dic = {model_name: trained_model.state_dict(), "optimizer": optimizer.state_dict()} 41 | if "scheduler" in kwargs.keys(): 42 | dic["scheduler"] = kwargs["scheduler"].state_dict() 43 | torch.save(dic, path) 44 | return 45 | 46 | @staticmethod 47 | def save_best_weights(trained_model, path): 48 | torch.save(trained_model.state_dict(), path) 49 | return 50 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2021, VinAI. All rights reserved. 2 | 3 | Redistribution and use in source and binary forms, with or without 4 | modification, are permitted provided that the following conditions are met: 5 | 6 | 1. Redistributions of source code must retain the above copyright notice, 7 | this list of conditions and the following disclaimer. 8 | 9 | 2. Redistributions in binary form must reproduce the above copyright notice, 10 | this list of conditions and the following disclaimer in the documentation 11 | and/or other materials provided with the distribution. 12 | 13 | 3. Neither the name of the copyright holder nor the names of its contributors 14 | may be used to endorse or promote products derived from this software 15 | without specific prior written permission. 16 | 17 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 18 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 19 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 20 | ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 21 | LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 22 | CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 23 | SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 24 | INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 25 | CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 26 | ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 27 | POSSIBILITY OF SUCH DAMAGE. 28 | -------------------------------------------------------------------------------- /criteria_comparing_sets_pcs/all_metrics_calculator.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch 5 | import torch.nn as nn 6 | 7 | 8 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 9 | from criteria_comparing_sets_pcs.jsd_calculator import JsdCalculator 10 | from metrics_from_point_flow.evaluation_metrics import compute_all_metrics 11 | 12 | 13 | class AllMetricsCalculator(nn.Module): 14 | def __init__(self): 15 | super(AllMetricsCalculator, self).__init__() 16 | 17 | @staticmethod 18 | def forward(sample_pcs, ref_pcs, batch_size, **kwargs): 19 | results = {} 20 | results.update(compute_all_metrics(sample_pcs, ref_pcs, batch_size, **kwargs)) 21 | for key, value in results.items(): 22 | if torch.is_tensor(value): 23 | results[key] = value.item() 24 | if "save_file" in kwargs.keys(): 25 | log = "{}: {}\n" 26 | with open(kwargs["save_file"], "a") as fp: 27 | for key, value in results.items(): 28 | fp.write(log.format(key, value)) 29 | # end for 30 | # end with 31 | # end if 32 | print("\n") 33 | log = "{}: {}\n" 34 | for key, value in results.items(): 35 | print(log.format(key, value)) 36 | # end for 37 | jsd = JsdCalculator.forward(sample_pcs, ref_pcs, **kwargs) 38 | return jsd 39 | 40 | 41 | if __name__ == "__main__": 42 | sample_pcs = torch.empty(10, 2048, 3).uniform_(0, 1).cuda() 43 | ref_pcs = torch.empty(10, 2048, 3).uniform_(0, 1).cuda() 44 | batch_size = 10 45 | print(AllMetricsCalculator.forward(sample_pcs, ref_pcs, batch_size)) 46 | -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/nn_distance.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | # from extensions.StructuralLosses.StructuralLossesBackend import NNDistance, NNDistanceGrad 4 | from metrics_from_point_flow.StructuralLosses.StructuralLossesBackend import NNDistance, NNDistanceGrad 5 | from torch.autograd import Function 6 | 7 | 8 | # Inherit from Function 9 | class NNDistanceFunction(Function): 10 | # Note that both forward and backward are @staticmethods 11 | @staticmethod 12 | # bias is an optional argument 13 | def forward(ctx, seta, setb): 14 | # print("Match Cost Forward") 15 | ctx.save_for_backward(seta, setb) 16 | """ 17 | input: 18 | set1 : batch_size * #dataset_points * 3 19 | set2 : batch_size * #query_points * 3 20 | returns: 21 | dist1, idx1, dist2, idx2 22 | """ 23 | dist1, idx1, dist2, idx2 = NNDistance(seta, setb) 24 | ctx.idx1 = idx1 25 | ctx.idx2 = idx2 26 | return dist1, dist2 27 | 28 | # This function has only a single output, so it gets only one gradient 29 | @staticmethod 30 | def backward(ctx, grad_dist1, grad_dist2): 31 | # print("Match Cost Backward") 32 | # This is a pattern that is very convenient - at the top of backward 33 | # unpack saved_tensors and initialize all gradients w.r.t. inputs to 34 | # None. Thanks to the fact that additional trailing Nones are 35 | # ignored, the return statement is simple even when the function has 36 | # optional inputs. 37 | seta, setb = ctx.saved_tensors 38 | idx1 = ctx.idx1 39 | idx2 = ctx.idx2 40 | grada, gradb = NNDistanceGrad(seta, setb, idx1, idx2, grad_dist1, grad_dist2) 41 | return grada, gradb 42 | 43 | 44 | nn_distance = NNDistanceFunction.apply 45 | -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/match_cost.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from metrics_from_point_flow.StructuralLosses.StructuralLossesBackend import ApproxMatch, MatchCost, MatchCostGrad 3 | from torch.autograd import Function 4 | 5 | 6 | # Inherit from Function 7 | class MatchCostFunction(Function): 8 | # Note that both forward and backward are @staticmethods 9 | @staticmethod 10 | # bias is an optional argument 11 | def forward(ctx, seta, setb): 12 | # print("Match Cost Forward") 13 | ctx.save_for_backward(seta, setb) 14 | """ 15 | input: 16 | set1 : batch_size * #dataset_points * 3 17 | set2 : batch_size * #query_points * 3 18 | returns: 19 | match : batch_size * #query_points * #dataset_points 20 | """ 21 | match, temp = ApproxMatch(seta, setb) 22 | ctx.match = match 23 | cost = MatchCost(seta, setb, match) 24 | return cost 25 | 26 | """ 27 | grad_1,grad_2=approxmatch_module.match_cost_grad(xyz1,xyz2,match) 28 | return [grad_1*tf.expand_dims(tf.expand_dims(grad_cost,1),2),grad_2*tf.expand_dims(tf.expand_dims(grad_cost,1),2),None] 29 | """ 30 | # This function has only a single output, so it gets only one gradient 31 | @staticmethod 32 | def backward(ctx, grad_output): 33 | # print("Match Cost Backward") 34 | # This is a pattern that is very convenient - at the top of backward 35 | # unpack saved_tensors and initialize all gradients w.r.t. inputs to 36 | # None. Thanks to the fact that additional trailing Nones are 37 | # ignored, the return statement is simple even when the function has 38 | # optional inputs. 39 | seta, setb = ctx.saved_tensors 40 | # grad_input = grad_weight = grad_bias = None 41 | grada, gradb = MatchCostGrad(seta, setb, ctx.match) 42 | grad_output_expand = grad_output.unsqueeze(1).unsqueeze(2) 43 | return grada * grad_output_expand, gradb * grad_output_expand 44 | 45 | 46 | match_cost = MatchCostFunction.apply 47 | -------------------------------------------------------------------------------- /generation/train_latent_generator_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "save_folder": "latent_generator/model/chair/sinkhorn/reg_0.001/", 3 | 4 | "device": "cuda", 5 | 6 | "autoencoder": "pointnet_lcd", 7 | "embedding_size": 256, 8 | "input_channels": 3, 9 | "num_points": 2048, 10 | "normalize": true, 11 | "ae_model_path":"model.pth", 12 | 13 | "generator_type": "MLP1", 14 | "latent_dim": 64, 15 | "n_hidden": 3, 16 | "hidden_size": 128, 17 | 18 | "loss": "sinkhorn", 19 | "sink_reg": 0.001, 20 | 21 | "max_slices_origin": 100000, 22 | "init_rec_epsilon": 0.1, 23 | "next_epsilon_ratio_rec": 0.1, 24 | 25 | "max_slices_latent": 500000, 26 | "init_reg_epsilon": 0.5, 27 | "next_epsilon_ratio_reg": 0.1, 28 | 29 | "init_projs": 10, 30 | "step_projs": 100, 31 | "loop_rate_thresh": 0.05, 32 | "regularization_coef": 10.0, 33 | 34 | "latent_num_projections": 336, 35 | "origin_num_projections": 14, 36 | 37 | "train_set": "shapenetcore55", 38 | "train_root": "dataset/shapenet_chair/train.npz", 39 | "phase": "train", 40 | 41 | "cates": "chair", 42 | 43 | "have_val_set": true, 44 | "val_set": "shapenetcore55", 45 | "val_root": "dataset/shapenet_chair/val.npz", 46 | "val_batch_size": 338, 47 | 48 | "optimizer": "sgd", 49 | "learning_rate": 0.001, 50 | "momentum": 0.9, 51 | "weight_decay": 0.0005, 52 | 53 | "checkpoint": "", 54 | 55 | "batch_size": 128, 56 | "num_workers": 0, 57 | 58 | "use_scheduler": true, 59 | "scheduler": "cyclic_lr", 60 | "base_lr": 0.0001, 61 | "max_lr": 0.002, 62 | 63 | "evaluator":"based_on_comparing_set_pcs", 64 | "eval_criteria":"jsd", 65 | 66 | "best_eval_value": 1e10, 67 | "best_epoch": -1, 68 | 69 | "best_train_loss": 1e10, 70 | "best_epoch_based_on_train_loss": -1, 71 | 72 | "epoch_gap_for_evaluation": 1, 73 | 74 | "start_epoch": 0, 75 | "num_epochs": 50, 76 | 77 | "epoch_gap_for_save": 25, 78 | 79 | "empty_cache_batch": false, 80 | 81 | "empty_cache_epoch": false 82 | } -------------------------------------------------------------------------------- /config.json: -------------------------------------------------------------------------------- 1 | { 2 | "loss": "swd", 3 | "autoencoder": "pointnet", 4 | "batch_size": 256, 5 | "seed":1, 6 | 7 | "device": "cuda", 8 | "num_prev_losses": 5, 9 | "add_noise": false, 10 | "noise_adder":"random", 11 | "mean_noiseadder":0.0, 12 | "std_noiseadder": 0.03, 13 | "train_denoise": false, 14 | 15 | "embedding_size": 256, 16 | "input_channels": 3, 17 | "output_channels": 3, 18 | "normalize": true, 19 | "num_points": 2048, 20 | 21 | "prim_caps_size": 1024, 22 | "prim_vec_size": 16, 23 | "latent_caps_size": 64, 24 | "latent_vec_size": 64, 25 | 26 | "version":"mean", 27 | 28 | "max_sw_num_iters": 1, 29 | "max_sw_lr": 1e-3, 30 | 31 | "num_projs": 100, 32 | "g_type": "circular", 33 | "degree": 2, 34 | 35 | "max_slices": 500, 36 | "init_projs": 2, 37 | 38 | "fix_epsilon": true, 39 | "init_rec_epsilon": 0.5, 40 | "next_epsilon_ratio_rec": 0.01, 41 | 42 | "step_projs": 1, 43 | "loop_rate_thresh": 0.00001, 44 | 45 | "train_set": "shapenetcore55", 46 | "train_root": "dataset/shapenet_core55/shapenet57448xyzonly.npz", 47 | "phase": "train", 48 | "eval_criteria":"", 49 | "cates": "", 50 | 51 | "have_val_set": false, 52 | "val_set": "shapenetcore55", 53 | "val_root": "dataset/shapenet_core55/shapenet57448xyzonly.npz", 54 | "val_batch_size": 338, 55 | 56 | "optimizer": "sgd", 57 | "learning_rate": 0.001, 58 | "momentum": 0.9, 59 | "weight_decay": 0.0005, 60 | 61 | "checkpoint": "latest.pth", 62 | 63 | "num_workers": 0, 64 | 65 | "use_scheduler": false, 66 | "scheduler": "cyclic_lr", 67 | "base_lr": 0.0001, 68 | "max_lr": 0.01, 69 | 70 | "evaluator":"based_on_train_loss", 71 | 72 | "best_eval_value": 1e10, 73 | "best_epoch": -1, 74 | 75 | "best_train_loss": 1e10, 76 | "best_epoch_based_on_train_loss": -1, 77 | 78 | "epoch_gap_for_evaluation": 1, 79 | 80 | "start_epoch": 0, 81 | "num_epochs": 301, 82 | 83 | "epoch_gap_for_save": 50, 84 | 85 | "empty_cache_batch": false, 86 | 87 | "empty_cache_epoch": false 88 | } -------------------------------------------------------------------------------- /registration/matcher.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import open3d 3 | 4 | 5 | class Matcher: 6 | def __init__(self, threshold=0.075, fitness_thresh=0.3, rmse_thresh=0.2): 7 | super().__init__() 8 | self.threshold = threshold 9 | self.fitness_thresh = fitness_thresh 10 | self.rmse_thresh = rmse_thresh 11 | 12 | def compute_transformation(self, source_points, target_points, source_descriptors, target_descriptors): 13 | source_features = open3d.registration.Feature() 14 | source_features.data = np.transpose(source_descriptors) 15 | target_features = open3d.registration.Feature() 16 | target_features.data = np.transpose(target_descriptors) 17 | 18 | processed_source_points = open3d.geometry.PointCloud() 19 | processed_source_points.points = open3d.utility.Vector3dVector(source_points) 20 | processed_target_points = open3d.geometry.PointCloud() 21 | processed_target_points.points = open3d.utility.Vector3dVector(target_points) 22 | 23 | result = open3d.registration.registration_ransac_based_on_feature_matching( 24 | processed_source_points, 25 | processed_target_points, 26 | source_features, 27 | target_features, 28 | self.threshold, 29 | open3d.registration.TransformationEstimationPointToPoint(False), 30 | 4, 31 | [open3d.registration.CorrespondenceCheckerBasedOnDistance(self.threshold)], 32 | open3d.registration.RANSACConvergenceCriteria(4000000, 500), 33 | ) 34 | 35 | information = open3d.registration.get_information_matrix_from_point_clouds( 36 | processed_source_points, processed_target_points, self.threshold, result.transformation 37 | ) 38 | threshold_num_point = min(len(processed_source_points.points), len(processed_target_points.points)) 39 | 40 | fitness = information[5, 5] / float(threshold_num_point) 41 | 42 | if (fitness > self.fitness_thresh) and (result.inlier_rmse < self.rmse_thresh): 43 | check = True 44 | else: 45 | check = False 46 | return result.transformation, check 47 | -------------------------------------------------------------------------------- /registration/preprocessor.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import open3d 3 | import torch 4 | from open3d.open3d.geometry import voxel_down_sample 5 | from tqdm import tqdm 6 | 7 | 8 | class Preprocessor: 9 | def __init__(self, trained_ae, device="cuda"): 10 | super().__init__() 11 | self.autoencoder = trained_ae 12 | self.device = device 13 | 14 | def extract_points_and_features(self, pcd, voxel_size, radius, num_points, batch_size=256, color=False): 15 | points, patches = self._extract_uniform_patches(pcd, voxel_size, radius, num_points, color) 16 | features = self._compute_features(patches, batch_size) 17 | return points, features 18 | 19 | def _extract_uniform_patches(self, pcd, voxel_size, radius, num_points, color=False): 20 | kdtree = open3d.geometry.KDTreeFlann(pcd) 21 | downsampled_points = voxel_down_sample(pcd, voxel_size) 22 | points = np.asarray(downsampled_points.points) 23 | patches = [] 24 | for i in range(points.shape[0]): 25 | k, index, _ = kdtree.search_hybrid_vector_3d(points[i], radius, num_points) 26 | if k < num_points: 27 | index = np.random.choice(index, num_points, replace=True) 28 | xyz = np.asarray(pcd.points)[index] 29 | xyz = (xyz - points[i]) / radius # normalize to local coordinates 30 | if color: 31 | rgb = np.asarray(pcd.colors)[index] 32 | patch = np.concatenate([xyz, rgb], axis=1) 33 | else: 34 | patch = xyz 35 | patches.append(patch) 36 | patches = np.stack(patches, axis=0) 37 | return points, patches 38 | 39 | def _compute_features(self, patches, batch_size): 40 | batches = torch.tensor(patches, dtype=torch.float32) 41 | batches = torch.split(batches, batch_size) 42 | features = [] 43 | self.autoencoder.eval() 44 | with torch.no_grad(): 45 | for _, x in tqdm(enumerate(batches)): 46 | x = x.to(self.device) 47 | z = self.autoencoder.encode(x) 48 | features.append(z.cpu().numpy()) 49 | return np.concatenate(features, axis=0) 50 | -------------------------------------------------------------------------------- /dataset/shapenet_core55.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch.utils.data as data 3 | 4 | 5 | class ShapeNetCore55XyzOnlyDataset(data.Dataset): 6 | def __init__(self, root, num_points=2048, phase="train"): 7 | super(ShapeNetCore55XyzOnlyDataset, self).__init__() 8 | assert root.endswith(".npz"), "root must be .npz file" 9 | assert phase in ["train", "test"] 10 | self.data = np.load(root)["data"] 11 | self.num_points = num_points 12 | self.phase = phase 13 | 14 | def __len__(self): 15 | return self.data.shape[0] 16 | 17 | def __getitem__(self, i): 18 | pc = self._pc_normalize(self.data[i]) 19 | if self.phase == "train": 20 | choice = np.random.choice(pc.shape[0], self.num_points, replace=True) 21 | pc = pc[choice, :] 22 | return pc 23 | 24 | def _pc_normalize(self, pc): 25 | centroid = np.mean(pc, axis=0) 26 | pc = pc - centroid 27 | m = np.max(np.sqrt(np.sum(pc ** 2, axis=1))) 28 | pc = pc / m 29 | return pc 30 | 31 | 32 | class ShapeNetCore55_n_attributes(data.Dataset): 33 | def __init__(self, root, num_attr=12, num_points=2048, phase="train"): 34 | super(ShapeNetCore55_n_attributes, self).__init__() 35 | assert root.endswith(".npz"), "root must be .npz file" 36 | assert phase in ["train", "test"] 37 | self.data = np.load(root)["data"] 38 | self.num_points = num_points 39 | self.phase = phase 40 | assert num_attr % 3 == 0, "num_attr must be divisible by 3" 41 | self.num_attr = num_attr 42 | 43 | def __len__(self): 44 | return self.data.shape[0] 45 | 46 | def __getitem__(self, i): 47 | pc = self._pc_normalize(self.data[i]) 48 | if self.phase == "train": 49 | choice = np.random.choice(pc.shape[0], self.num_points, replace=True) 50 | pc = pc[choice, :] 51 | _pc = pc 52 | for i in range(2, int(self.num_attr % 3 + 1)): 53 | _pc = np.concatenate((_pc, np.power(pc, i))) 54 | return _pc 55 | 56 | def _pc_normalize(self, pc): 57 | centroid = np.mean(pc, axis=0) 58 | pc = pc - centroid 59 | m = np.max(np.sqrt(np.sum(pc ** 2, axis=1))) 60 | pc = pc / m 61 | return pc 62 | -------------------------------------------------------------------------------- /registration/registration_test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os.path as osp 4 | import random 5 | import sys 6 | import time 7 | 8 | import numpy as np 9 | import torch 10 | from tqdm import tqdm 11 | 12 | 13 | def seed_worker(worker_id): 14 | worker_seed = torch.initial_seed() % 2 ** 32 15 | np.random.seed(worker_seed) 16 | random.seed(worker_seed) 17 | 18 | 19 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 20 | from dataset import Fine3dMatchDataset 21 | from matcher import Matcher 22 | from writer import Writer 23 | 24 | 25 | def main(): 26 | parser = argparse.ArgumentParser() 27 | parser.add_argument("--config", type=str, help="registration_config.json") 28 | parser.add_argument("--logdir", type=str, help="folder containing preprocessed data") 29 | args = parser.parse_args() 30 | 31 | config = args.config 32 | logdir = args.logdir 33 | print("Check logdir in 10s: ", logdir) 34 | time.sleep(10) 35 | args = json.load(open(config)) 36 | 37 | # set seed 38 | torch.manual_seed(args["seed"]) 39 | random.seed(args["seed"]) 40 | np.random.seed(args["seed"]) 41 | 42 | fname = osp.join(logdir, "config.json") 43 | with open(fname, "w") as fp: 44 | json.dump(args, fp, indent=4) 45 | 46 | trans_file = osp.join(logdir, "transformation.log") 47 | # dataset 48 | dataset = Fine3dMatchDataset(logdir) 49 | # matcher 50 | matcher = Matcher(args["threshold"], args["fitness_threshold"], args["rmse_threshold"]) 51 | # writer 52 | writer = Writer() 53 | # process 54 | num_frag = len(dataset) 55 | for j in tqdm(range(1, num_frag)): 56 | for i in tqdm(range(0, j)): 57 | source_id = int(j) 58 | target_id = int(i) 59 | meta_data = "{} {} {}\n".format(target_id, source_id, num_frag) 60 | 61 | source_points, source_features = dataset[source_id] 62 | target_points, target_features = dataset[target_id] 63 | trans, check = matcher.compute_transformation( 64 | source_points, target_points, source_features, target_features 65 | ) 66 | if check: 67 | writer.write_arr_to_file(trans_file, meta_data, trans) 68 | 69 | 70 | if __name__ == "__main__": 71 | main() 72 | -------------------------------------------------------------------------------- /dataset/download_3dmatch.sh: -------------------------------------------------------------------------------- 1 | SCRIPT=`realpath $0` 2 | SCRIPTPATH=`dirname $SCRIPT` 3 | 4 | cd $SCRIPTPATH/ 5 | wget http://vision.princeton.edu/projects/2016/3DMatch/downloads/scene-fragments/7-scenes-redkitchen.zip --no-check-certificate 6 | unzip 7-scenes-redkitchen.zip 7 | rm 7-scenes-redkitchen.zip 8 | mv 7-scenes-redkitchen kitchen 9 | 10 | wget http://vision.princeton.edu/projects/2016/3DMatch/downloads/scene-fragments/sun3d-home_at-home_at_scan1_2013_jan_1.zip --no-check-certificate 11 | unzip sun3d-home_at-home_at_scan1_2013_jan_1.zip 12 | rm sun3d-home_at-home_at_scan1_2013_jan_1.zip 13 | mv sun3d-home_at-home_at_scan1_2013_jan_1 home1 14 | 15 | wget http://vision.princeton.edu/projects/2016/3DMatch/downloads/scene-fragments/sun3d-home_md-home_md_scan9_2012_sep_30.zip --no-check-certificate 16 | unzip sun3d-home_md-home_md_scan9_2012_sep_30.zip 17 | rm sun3d-home_md-home_md_scan9_2012_sep_30.zip 18 | mv sun3d-home_md-home_md_scan9_2012_sep_30 home2 19 | 20 | wget http://vision.princeton.edu/projects/2016/3DMatch/downloads/scene-fragments/sun3d-hotel_uc-scan3.zip --no-check-certificate 21 | unzip sun3d-hotel_uc-scan3.zip 22 | rm sun3d-hotel_uc-scan3.zip 23 | mv sun3d-hotel_uc-scan3 hotel1 24 | 25 | wget http://vision.princeton.edu/projects/2016/3DMatch/downloads/scene-fragments/sun3d-hotel_umd-maryland_hotel1.zip --no-check-certificate 26 | unzip sun3d-hotel_umd-maryland_hotel1.zip 27 | rm sun3d-hotel_umd-maryland_hotel1.zip 28 | mv sun3d-hotel_umd-maryland_hotel1 hotel2 29 | 30 | wget http://vision.princeton.edu/projects/2016/3DMatch/downloads/scene-fragments/sun3d-hotel_umd-maryland_hotel3.zip --no-check-certificate 31 | unzip sun3d-hotel_umd-maryland_hotel3.zip 32 | rm sun3d-hotel_umd-maryland_hotel3.zip 33 | mv sun3d-hotel_umd-maryland_hotel3 hotel3 34 | 35 | wget http://vision.princeton.edu/projects/2016/3DMatch/downloads/scene-fragments/sun3d-mit_76_studyroom-76-1studyroom2.zip --no-check-certificate 36 | unzip sun3d-mit_76_studyroom-76-1studyroom2.zip 37 | rm sun3d-mit_76_studyroom-76-1studyroom2.zip 38 | mv sun3d-mit_76_studyroom-76-1studyroom2 study 39 | 40 | wget http://vision.princeton.edu/projects/2016/3DMatch/downloads/scene-fragments/sun3d-mit_lab_hj-lab_hj_tea_nov_2_2012_scan1_erika.zip --no-check-certificate 41 | unzip sun3d-mit_lab_hj-lab_hj_tea_nov_2_2012_scan1_erika.zip 42 | rm sun3d-mit_lab_hj-lab_hj_tea_nov_2_2012_scan1_erika.zip 43 | mv sun3d-mit_lab_hj-lab_hj_tea_nov_2_2012_scan1_erika lab -------------------------------------------------------------------------------- /generation/preprocess.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import random 3 | import sys 4 | 5 | import numpy as np 6 | import torch 7 | from torch.utils.data import DataLoader 8 | from tqdm import tqdm 9 | 10 | 11 | def seed_worker(worker_id): 12 | worker_seed = torch.initial_seed() % 2 ** 32 13 | np.random.seed(worker_seed) 14 | random.seed(worker_seed) 15 | 16 | 17 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 18 | from dataset.shapenet_core55 import ShapeNetCore55XyzOnlyDataset 19 | from models import PointNetAE 20 | from utils.utils import create_save_folder, initialize_main, load_model_for_evaluation 21 | 22 | 23 | def main(): 24 | args, logdir, data_path = initialize_main() 25 | 26 | # set seed 27 | torch.manual_seed(args["seed"]) 28 | random.seed(args["seed"]) 29 | np.random.seed(args["seed"]) 30 | 31 | # save_folder 32 | save_folder = create_save_folder(logdir, args["save_folder"])["save_folder"] 33 | 34 | # device 35 | device = torch.device(args["device"]) 36 | 37 | # network 38 | if args["architecture"] == "pointnet": 39 | net = PointNetAE( 40 | embedding_size=args["embedding_size"], 41 | input_channels=args["point_channels"], 42 | output_channels=args["point_channels"], 43 | num_points=args["num_points"], 44 | normalize=args["normalize"], 45 | ).to(device) 46 | else: 47 | raise ValueError("Unknown architecture.") 48 | 49 | try: 50 | net = load_model_for_evaluation(net, args["model_path"]) 51 | except: 52 | net = load_model_for_evaluation(net, osp.join(logdir, args["model_path"])) 53 | 54 | # dataset 55 | if args["dataset"] == "shapenetcore55": 56 | dataset = ShapeNetCore55XyzOnlyDataset( 57 | data_path, 58 | num_points=args["num_points"], 59 | phase="test", 60 | ) 61 | else: 62 | raise ValueError("Unknown dataset type.") 63 | 64 | # dataloader 65 | dataloader = DataLoader( 66 | dataset, batch_size=args["batch_size"], shuffle=True, sampler=None, pin_memory=True, worker_init_fn=seed_worker 67 | ) 68 | 69 | # save_path 70 | save_path = osp.join(save_folder, "latent_codes.npz") 71 | 72 | # main 73 | latent_vectors_list = [] 74 | with torch.no_grad(): 75 | for _, batch in tqdm(enumerate(dataloader)): 76 | pcs = batch.to(device) 77 | latent_vectors = net.encode(pcs) 78 | latent_vectors_list.append(latent_vectors.reshape(latent_vectors.shape[0], -1).cpu().numpy()) 79 | latent_vectors_list = np.concatenate(latent_vectors_list, axis=0) 80 | np.savez(save_path, data=latent_vectors_list) 81 | 82 | 83 | if __name__ == "__main__": 84 | main() 85 | -------------------------------------------------------------------------------- /dataset/modelnet40.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os.path as osp 3 | import sys 4 | 5 | import h5py 6 | import numpy as np 7 | import torch.utils.data as data 8 | 9 | 10 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 11 | try: 12 | from utils import pc_normalize, sample_pc 13 | except: 14 | from .utils import pc_normalize, sample_pc 15 | 16 | 17 | class ModelNet40(data.Dataset): 18 | def __init__(self, root, num_points=2048): 19 | super().__init__() 20 | flistname = osp.join(root, "*.h5") 21 | flistname = sorted(glob.glob(flistname)) 22 | flist = [h5py.File(fname, "r") for fname in flistname] 23 | self.data = np.concatenate([f["data"][:] for f in flist]) 24 | self.label = np.concatenate([f["label"][:] for f in flist]) 25 | self.num_points = num_points 26 | 27 | def __len__(self): 28 | return self.label.shape[0] 29 | 30 | def __getitem__(self, i): 31 | pc = pc_normalize(self.data[i]) 32 | # pc = sample_pc(pc, self.num_points) #modelnet40 to test autoencoder, so we dont sample new pc from origin pc 33 | return pc, self.label[i] 34 | 35 | 36 | class LatentVectorsModelNet40(data.Dataset): 37 | def __init__(self, root): 38 | super().__init__() 39 | assert root.endswith(".npz"), "root must be a npz file" 40 | self.data = np.load(root)["latent_vectors"] 41 | self.labels = np.load(root)["labels"] 42 | 43 | def __len__(self): 44 | return self.labels.shape[0] 45 | 46 | def __getitem__(self, i): 47 | return self.data[i].reshape(-1), self.labels[i] 48 | 49 | 50 | class ModelNet40Train(data.Dataset): 51 | def __init__(self, root, num_points=2048): 52 | super().__init__() 53 | flistname = osp.join(root, "*.h5") 54 | flistname = sorted(glob.glob(flistname)) 55 | flist = [h5py.File(fname, "r") for fname in flistname] 56 | self.data = np.concatenate([f["data"][:] for f in flist]) 57 | self.label = np.concatenate([f["label"][:] for f in flist]) 58 | self.num_points = num_points 59 | 60 | def __len__(self): 61 | return self.label.shape[0] 62 | 63 | def __getitem__(self, i): 64 | pc = pc_normalize(self.data[i]) 65 | pc = sample_pc(pc, self.num_points) # modelnet40 to test autoencoder, so we dont sample new pc from origin pc 66 | return pc, self.label[i] 67 | 68 | 69 | class LatentCapsulesModelNet40(data.Dataset): 70 | def __init__(self, root): 71 | super().__init__() 72 | assert root.endswith(".h5"), "root must be a h5 file" 73 | fp = h5py.File(root) 74 | self.data = fp["data"] 75 | self.labels = fp["cls_label"] 76 | 77 | def __len__(self): 78 | return self.data.shape[0] 79 | 80 | def __getitem__(self, i): 81 | return self.data[i].reshape(-1), self.labels[i] 82 | -------------------------------------------------------------------------------- /registration/preprocess_data.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os 4 | import os.path as osp 5 | import random 6 | import sys 7 | 8 | import numpy as np 9 | import torch 10 | from tqdm import tqdm 11 | 12 | 13 | def seed_worker(worker_id): 14 | worker_seed = torch.initial_seed() % 2 ** 32 15 | np.random.seed(worker_seed) 16 | random.seed(worker_seed) 17 | 18 | 19 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 20 | from dataset.raw3dmatch import ThreeDMatchRawDataset 21 | from models import PointNetAE 22 | from preprocessor import Preprocessor 23 | from utils import load_model_for_evaluation 24 | 25 | 26 | def main(): 27 | parser = argparse.ArgumentParser() 28 | parser.add_argument("--config", type=str, help="config path") 29 | parser.add_argument("--logdir", type=str, help="log path") 30 | parser.add_argument("--data_path", type=str, help="home1 home2 hotel1 hotel2 hotel3 kitchen lab study") 31 | args = parser.parse_args() 32 | 33 | config = args.config 34 | logdir = args.logdir 35 | print("logdir: ", logdir) 36 | dset = args.data_path 37 | args = json.load(open(config)) 38 | 39 | # set seed 40 | torch.manual_seed(args["seed"]) 41 | random.seed(args["seed"]) 42 | np.random.seed(args["seed"]) 43 | 44 | # device 45 | device = torch.device(args["device"]) 46 | 47 | # autoencoder 48 | if args["autoencoder"] == "pointnet": 49 | autoencoder = PointNetAE( 50 | args["embedding_size"], 51 | args["input_channels"], 52 | args["output_channels"], 53 | args["num_points"], 54 | args["normalize"], 55 | ).to(device) 56 | try: 57 | autoencoder = load_model_for_evaluation(autoencoder, args["model_path"]) 58 | except: 59 | autoencoder = load_model_for_evaluation(autoencoder, osp.join(logdir, args["model_path"])) 60 | 61 | # dataloader 62 | # args["root"] = osp.join(args["root"], dset) 63 | args["root"] = dset 64 | dataset = ThreeDMatchRawDataset(args["root"]) 65 | # preprocessor 66 | preprocessor = Preprocessor(autoencoder, device) 67 | # preprocess 68 | save_folder = osp.join(logdir, args["save_folder"], args["root"].split("/")[-1]) 69 | if not osp.isdir(save_folder): 70 | os.makedirs(save_folder) 71 | for i in tqdm(range(len(dataset))): 72 | pcd, name = dataset[i] 73 | points, features = preprocessor.extract_points_and_features( 74 | pcd, args["voxel_size"], args["radius"], args["num_points"], args["batch_size"], color=args["color"] 75 | ) 76 | results = {"points": points, "features": features} 77 | fname = osp.join(save_folder, "{}.npz".format(name)) 78 | np.savez(fname, results) 79 | print(">save_folder:", save_folder) 80 | 81 | 82 | if __name__ == "__main__": 83 | main() 84 | -------------------------------------------------------------------------------- /classification/classification_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | import random 4 | import sys 5 | 6 | import numpy as np 7 | import torch 8 | import torch.nn.functional as F 9 | import torch.utils.data as data 10 | from tqdm import tqdm 11 | 12 | 13 | def seed_worker(worker_id): 14 | worker_seed = torch.initial_seed() % 2 ** 32 15 | np.random.seed(worker_seed) 16 | random.seed(worker_seed) 17 | 18 | 19 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 20 | from classifier import MLPClassifier 21 | from dataset.modelnet40 import LatentCapsulesModelNet40, LatentVectorsModelNet40 22 | from utils.utils import initialize_main, load_model_for_evaluation 23 | 24 | 25 | def main(): 26 | args, logdir = initialize_main() 27 | 28 | # set seed 29 | torch.manual_seed(args["seed"]) 30 | random.seed(args["seed"]) 31 | np.random.seed(args["seed"]) 32 | 33 | # save_results folder 34 | save_folder = osp.join(logdir, args["save_folder"]) 35 | if not osp.exists(save_folder): 36 | os.makedirs(save_folder) 37 | 38 | # device 39 | device = torch.device(args["device"]) 40 | 41 | # root 42 | args["root"] = osp.join(logdir, "latent_codes/model/modelnet40-test/saved_latent_vectors.npz") 43 | 44 | # dataloader 45 | if args["root"].endswith(".npz"): 46 | dset = LatentVectorsModelNet40(args["root"]) # root is a npz file 47 | elif args["root"].endswith(".h5"): 48 | dset = LatentCapsulesModelNet40(args["root"]) 49 | else: 50 | raise Exception("Unknown dataset.") 51 | loader = data.DataLoader( 52 | dset, 53 | batch_size=args["batch_size"], 54 | pin_memory=args["pin_memory"], 55 | num_workers=args["num_workers"], 56 | shuffle=args["shuffle"], 57 | worker_init_fn=seed_worker, 58 | ) 59 | 60 | # classifier 61 | classifier = MLPClassifier( 62 | args["input_size"], args["output_size"], args["dropout_p"], [int(i) for i in args["hidden_sizes"].split(",")] 63 | ).to(device) 64 | try: 65 | classifier = load_model_for_evaluation(classifier, args["model_path"]) 66 | except: 67 | classifier = load_model_for_evaluation(classifier, osp.join(save_folder, args["model_path"])) 68 | 69 | # test main 70 | num_true = 0 71 | with torch.no_grad(): 72 | for _, (batch, labels) in tqdm(enumerate(loader)): 73 | batch = batch.to(device) 74 | labels = labels.to(device).squeeze().type(torch.long) 75 | predicted_labels = torch.argmax(F.softmax(classifier(batch), dim=-1), dim=-1) 76 | num_true += (predicted_labels == labels).sum().item() 77 | # report 78 | print("Model: ", save_folder) 79 | log = "Accuracy: {}\n".format(num_true / len(dset)) 80 | acc = num_true * 1.0 / len(dset) 81 | print(log) 82 | with open(os.path.join(save_folder, "accuracy.txt"), "a") as fp: 83 | fp.write("{} ".format(acc)) 84 | 85 | 86 | if __name__ == "__main__": 87 | main() 88 | -------------------------------------------------------------------------------- /models/pointnet.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | 9 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 10 | 11 | 12 | class PointNetEncoder(nn.Module): 13 | def __init__(self, embedding_size, input_channels=3): 14 | super(PointNetEncoder, self).__init__() 15 | self.input_channels = input_channels 16 | self.mlp1 = nn.Sequential( 17 | nn.Conv1d(input_channels, 64, 1), 18 | nn.BatchNorm1d(64), 19 | nn.ReLU(), 20 | nn.Conv1d(64, 64, 1), 21 | nn.BatchNorm1d(64), 22 | nn.ReLU(), 23 | ) 24 | self.mlp2 = nn.Sequential( 25 | nn.Conv1d(64, 64, 1), 26 | nn.BatchNorm1d(64), 27 | nn.ReLU(), 28 | nn.Conv1d(64, 128, 1), 29 | nn.BatchNorm1d(128), 30 | nn.ReLU(), 31 | nn.Conv1d(128, 1024, 1), 32 | nn.BatchNorm1d(1024), 33 | nn.ReLU(), 34 | ) 35 | self.fc_mu = nn.Linear(1024, embedding_size) 36 | 37 | def forward(self, x): 38 | num_points = x.shape[1] 39 | x = x[:, :, : self.input_channels] 40 | x = x.transpose(2, 1) # transpose to apply 1D convolution 41 | x = self.mlp1(x) 42 | x = self.mlp2(x) 43 | x = F.max_pool1d(x, num_points).squeeze(2) # max pooling 44 | x = self.fc_mu(x) 45 | 46 | return x 47 | 48 | 49 | class PointNetDecoder(nn.Module): 50 | def __init__(self, embedding_size, output_channels=3, num_points=1024): 51 | super(PointNetDecoder, self).__init__() 52 | self.num_points = num_points 53 | self.output_channels = output_channels 54 | self.fc1 = nn.Linear(embedding_size, 1024) 55 | self.fc2 = nn.Linear(1024, 1024) 56 | self.bn1 = nn.BatchNorm1d(1024) 57 | self.bn2 = nn.BatchNorm1d(1024) 58 | self.fc3 = nn.Linear(1024, num_points * output_channels) 59 | 60 | def forward(self, x): 61 | batch_size = x.shape[0] 62 | x = F.relu(self.bn1(self.fc1(x))) 63 | x = F.relu(self.bn2(self.fc2(x))) 64 | x = torch.tanh(self.fc3(x)) 65 | x = x.view(batch_size, self.num_points, self.output_channels) 66 | x = x.contiguous() 67 | return x 68 | 69 | 70 | class PointNetAE(nn.Module): 71 | def __init__(self, embedding_size=256, input_channels=3, output_channels=3, num_points=1024, normalize=True): 72 | super(PointNetAE, self).__init__() 73 | self.normalize = normalize 74 | self.input_channels = input_channels 75 | self.output_channels = output_channels 76 | self.embedding_size = embedding_size 77 | self.encoder = PointNetEncoder(embedding_size, input_channels) 78 | self.decoder = PointNetDecoder(embedding_size, output_channels, num_points) 79 | 80 | def encode(self, x): 81 | z = self.encoder(x) 82 | if self.normalize: 83 | z = F.normalize(z) 84 | return z 85 | 86 | def decode(self, z): 87 | y = self.decoder(z) 88 | return y 89 | -------------------------------------------------------------------------------- /generation/deepul/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | from os.path import dirname, exists, join 4 | 5 | import matplotlib.pyplot as plt 6 | import numpy as np 7 | import torch 8 | from torchvision.utils import make_grid 9 | 10 | 11 | def savefig(fname, show_figure=True): 12 | if not exists(dirname(fname)): 13 | os.makedirs(dirname(fname)) 14 | plt.tight_layout() 15 | plt.savefig(fname) 16 | if show_figure: 17 | plt.show() 18 | 19 | 20 | def save_training_plot(train_losses, test_losses, title, fname): 21 | plt.figure() 22 | n_epochs = len(test_losses) - 1 23 | x_train = np.linspace(0, n_epochs, len(train_losses)) 24 | x_test = np.arange(n_epochs + 1) 25 | 26 | plt.plot(x_train, train_losses, label="train loss") 27 | plt.plot(x_test, test_losses, label="test loss") 28 | plt.legend() 29 | plt.title(title) 30 | plt.xlabel("Epoch") 31 | plt.ylabel("NLL") 32 | savefig(fname) 33 | 34 | 35 | def save_scatter_2d(data, title, fname): 36 | plt.figure() 37 | plt.title(title) 38 | plt.scatter(data[:, 0], data[:, 1]) 39 | savefig(fname) 40 | 41 | 42 | def save_distribution_1d(data, distribution, title, fname): 43 | d = len(distribution) 44 | 45 | plt.figure() 46 | plt.hist(data, bins=np.arange(d) - 0.5, label="train data", density=True) 47 | 48 | x = np.linspace(-0.5, d - 0.5, 1000) 49 | y = distribution.repeat(1000 // d) 50 | plt.plot(x, y, label="learned distribution") 51 | 52 | plt.title(title) 53 | plt.xlabel("x") 54 | plt.ylabel("Probability") 55 | plt.legend() 56 | savefig(fname) 57 | 58 | 59 | def save_distribution_2d(true_dist, learned_dist, fname): 60 | fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8)) 61 | ax1.imshow(true_dist) 62 | ax1.set_title("True Distribution") 63 | ax1.axis("off") 64 | ax2.imshow(learned_dist) 65 | ax2.set_title("Learned Distribution") 66 | ax2.axis("off") 67 | savefig(fname) 68 | 69 | 70 | def show_samples(samples, fname=None, nrow=10, title="Samples"): 71 | samples = (torch.FloatTensor(samples) / 255).permute(0, 3, 1, 2) 72 | grid_img = make_grid(samples, nrow=nrow) 73 | plt.figure() 74 | plt.title(title) 75 | plt.imshow(grid_img.permute(1, 2, 0)) 76 | plt.axis("off") 77 | 78 | if fname is not None: 79 | savefig(fname) 80 | else: 81 | plt.show() 82 | 83 | 84 | def load_pickled_data(fname, include_labels=False): 85 | with open(fname, "rb") as f: 86 | data = pickle.load(f) 87 | 88 | train_data, test_data = data["train"], data["test"] 89 | if "mnist.pkl" in fname or "shapes.pkl" in fname: 90 | # Binarize MNIST and shapes dataset 91 | train_data = (train_data > 127.5).astype("uint8") 92 | test_data = (test_data > 127.5).astype("uint8") 93 | if "celeb.pkl" in fname: 94 | train_data = train_data[:, :, :, [2, 1, 0]] 95 | test_data = test_data[:, :, :, [2, 1, 0]] 96 | if include_labels: 97 | return train_data, test_data, data["train_labels"], data["test_labels"] 98 | return train_data, test_data 99 | 100 | 101 | def get_data_dir(hw_number): 102 | return join("deepul", "homeworks", f"hw{hw_number}", "data") 103 | 104 | 105 | def quantize(images, n_bits): 106 | images = np.floor(images / 256.0 * 2 ** n_bits) 107 | return images.astype("uint8") 108 | -------------------------------------------------------------------------------- /classification/linear_svm_classification.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | import sys 4 | 5 | import numpy as np 6 | from sklearn.svm import LinearSVC 7 | from tqdm import tqdm 8 | 9 | 10 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 11 | import torch.utils.data as data 12 | from dataset.modelnet40 import LatentCapsulesModelNet40, LatentVectorsModelNet40 13 | from utils.utils import create_save_folder, initialize_main 14 | 15 | 16 | def main(): 17 | args, logdir = initialize_main() 18 | 19 | # save_results folder 20 | save_folder = create_save_folder(logdir, args["save_folder"])["save_folder"] 21 | 22 | # datasets 23 | args["train_root"] = os.path.join(logdir, args["train_root"]) 24 | args["test_root"] = os.path.join(logdir, args["test_root"]) 25 | 26 | # train loader 27 | if args["train_root"].endswith(".npz"): 28 | train_set = LatentVectorsModelNet40(args["train_root"]) # root is a npz file 29 | elif args["train_root"].endswith(".h5"): 30 | train_set = LatentCapsulesModelNet40(args["train_root"]) 31 | else: 32 | raise Exception("Unknown dataset.") 33 | train_loader = data.DataLoader( 34 | train_set, 35 | batch_size=args["batch_size"], 36 | pin_memory=args["pin_memory"], 37 | num_workers=args["num_workers"], 38 | shuffle=args["shuffle"], 39 | ) 40 | 41 | # test loader 42 | if args["test_root"].endswith(".npz"): 43 | test_set = LatentVectorsModelNet40(args["test_root"]) # root is a npz file 44 | elif args["test_root"].endswith(".h5"): 45 | test_set = LatentCapsulesModelNet40(args["test_root"]) 46 | else: 47 | raise Exception("Unknown dataset.") 48 | test_loader = data.DataLoader( 49 | test_set, 50 | batch_size=args["batch_size"], 51 | pin_memory=args["pin_memory"], 52 | num_workers=args["num_workers"], 53 | shuffle=False, 54 | ) 55 | 56 | # classifier 57 | clf = LinearSVC() 58 | 59 | # main 60 | train_feature = np.zeros((1, args["input_size"])) 61 | train_label = np.zeros((1, 1)) 62 | test_feature = np.zeros((1, args["input_size"])) 63 | test_label = np.zeros((1, 1)) 64 | for batch_id, (latents, labels) in tqdm(enumerate(train_loader)): 65 | train_label = np.concatenate((train_label, labels.numpy()), axis=None) 66 | train_label = train_label.astype(int) 67 | train_feature = np.concatenate((train_feature, latents.numpy()), axis=0) 68 | if batch_id % 10 == 0: 69 | print("add train batch: ", batch_id) 70 | 71 | for batch_id, (latents, labels) in tqdm(enumerate(test_loader)): 72 | test_label = np.concatenate((test_label, labels.numpy()), axis=None) 73 | test_label = test_label.astype(int) 74 | test_feature = np.concatenate((test_feature, latents.numpy()), axis=0) 75 | if batch_id % 10 == 0: 76 | print("add test batch: ", batch_id) 77 | 78 | train_feature = train_feature[1:, :] 79 | train_label = train_label[1:] 80 | test_feature = test_feature[1:, :] 81 | test_label = test_label[1:] 82 | 83 | print("training the linear SVM.......") 84 | clf.fit(train_feature, train_label) 85 | confidence = clf.score(test_feature, test_label) 86 | print("Accuracy: {} %".format(confidence * 100)) 87 | 88 | with open(os.path.join(save_folder, "accuracy.txt"), "a") as fp: 89 | fp.write(str(confidence) + "\n") 90 | 91 | 92 | if __name__ == "__main__": 93 | main() 94 | -------------------------------------------------------------------------------- /generation/deepul/pytorch_util.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | 5 | def soft_update_from_to(source, target, tau): 6 | for target_param, param in zip(target.parameters(), source.parameters()): 7 | target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau) 8 | 9 | 10 | def copy_model_params_from_to(source, target): 11 | for target_param, param in zip(target.parameters(), source.parameters()): 12 | target_param.data.copy_(param.data) 13 | 14 | 15 | def fanin_init(tensor): 16 | size = tensor.size() 17 | if len(size) == 2: 18 | fan_in = size[0] 19 | elif len(size) > 2: 20 | fan_in = np.prod(size[1:]) 21 | else: 22 | raise Exception("Shape must be have dimension at least 2.") 23 | bound = 1.0 / np.sqrt(fan_in) 24 | return tensor.data.uniform_(-bound, bound) 25 | 26 | 27 | def fanin_init_weights_like(tensor): 28 | size = tensor.size() 29 | if len(size) == 2: 30 | fan_in = size[0] 31 | elif len(size) > 2: 32 | fan_in = np.prod(size[1:]) 33 | else: 34 | raise Exception("Shape must be have dimension at least 2.") 35 | bound = 1.0 / np.sqrt(fan_in) 36 | new_tensor = FloatTensor(tensor.size()) 37 | new_tensor.uniform_(-bound, bound) 38 | return new_tensor 39 | 40 | 41 | """ 42 | GPU wrappers 43 | """ 44 | 45 | _use_gpu = False 46 | device = None 47 | _gpu_id = 0 48 | 49 | 50 | def set_gpu_mode(mode, gpu_id=0): 51 | global _use_gpu 52 | global device 53 | global _gpu_id 54 | _gpu_id = gpu_id 55 | _use_gpu = mode 56 | device = torch.device("cuda:" + str(gpu_id) if _use_gpu else "cpu") 57 | 58 | 59 | def gpu_enabled(): 60 | return _use_gpu 61 | 62 | 63 | def set_device(gpu_id): 64 | torch.cuda.set_device(gpu_id) 65 | 66 | 67 | # noinspection PyPep8Naming 68 | def FloatTensor(*args, torch_device=None, **kwargs): 69 | if torch_device is None: 70 | torch_device = device 71 | return torch.FloatTensor(*args, **kwargs).to(torch_device) 72 | 73 | 74 | def from_numpy(*args, **kwargs): 75 | return torch.from_numpy(*args, **kwargs).float().to(device) 76 | 77 | 78 | def get_numpy(tensor): 79 | return tensor.to("cpu").detach().numpy() 80 | 81 | 82 | def zeros(*sizes, torch_device=None, **kwargs): 83 | if torch_device is None: 84 | torch_device = device 85 | return torch.zeros(*sizes, **kwargs, device=torch_device) 86 | 87 | 88 | def ones(*sizes, torch_device=None, **kwargs): 89 | if torch_device is None: 90 | torch_device = device 91 | return torch.ones(*sizes, **kwargs, device=torch_device) 92 | 93 | 94 | def ones_like(*args, torch_device=None, **kwargs): 95 | if torch_device is None: 96 | torch_device = device 97 | return torch.ones_like(*args, **kwargs, device=torch_device) 98 | 99 | 100 | def randn(*args, torch_device=None, **kwargs): 101 | if torch_device is None: 102 | torch_device = device 103 | return torch.randn(*args, **kwargs, device=torch_device) 104 | 105 | 106 | def zeros_like(*args, torch_device=None, **kwargs): 107 | if torch_device is None: 108 | torch_device = device 109 | return torch.zeros_like(*args, **kwargs, device=torch_device) 110 | 111 | 112 | def tensor(*args, torch_device=None, **kwargs): 113 | if torch_device is None: 114 | torch_device = device 115 | return torch.tensor(*args, **kwargs, device=torch_device) 116 | 117 | 118 | def normal(*args, **kwargs): 119 | return torch.normal(*args, **kwargs).to(device) 120 | -------------------------------------------------------------------------------- /classification/preprocess_data.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import random 3 | import sys 4 | 5 | import numpy as np 6 | import torch 7 | import torch.utils.data as data 8 | from tqdm import tqdm 9 | 10 | 11 | def seed_worker(worker_id): 12 | worker_seed = torch.initial_seed() % 2 ** 32 13 | np.random.seed(worker_seed) 14 | random.seed(worker_seed) 15 | 16 | 17 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 18 | from add_noise_to_data.random_noise import RandomNoiseAdder 19 | from dataset import ModelNet40 20 | from models import PointCapsNet, PointNetAE 21 | from utils import create_save_folder, initialize_main, load_model_for_evaluation 22 | 23 | 24 | def main(): 25 | args, logdir, data_path = initialize_main() 26 | 27 | # set seed 28 | torch.manual_seed(args["seed"]) 29 | random.seed(args["seed"]) 30 | np.random.seed(args["seed"]) 31 | 32 | # save_results folder 33 | save_folder = create_save_folder(logdir, args["save_folder"])["save_folder"] 34 | 35 | # device 36 | device = torch.device(args["device"]) 37 | 38 | # autoencoder 39 | if args["autoencoder"] == "pointnet": 40 | autoencoder = PointNetAE( 41 | args["embedding_size"], 42 | args["input_channels"], 43 | args["output_channels"], 44 | args["num_points"], 45 | args["normalize"], 46 | ).to(device) 47 | 48 | elif args["autoencoder"] == "pcn": 49 | autoencoder = PointCapsNet( 50 | args["prim_caps_size"], 51 | args["prim_vec_size"], 52 | args["latent_caps_size"], 53 | args["latent_vec_size"], 54 | args["num_points"], 55 | ).to(device) 56 | 57 | else: 58 | raise Exception("Unknown autoencoder architecture.") 59 | 60 | try: 61 | autoencoder = load_model_for_evaluation(autoencoder, args["model_path"]) 62 | except: 63 | autoencoder = load_model_for_evaluation(autoencoder, osp.join(logdir, args["model_path"])) 64 | 65 | # dataset 66 | if args["dataset"] == "modelnet40": 67 | dataset = ModelNet40(data_path) # root is a folder containing h5 files 68 | else: 69 | raise ValueError("Unknown dataset type.") 70 | 71 | # dataloader 72 | loader = data.DataLoader( 73 | dataset, 74 | batch_size=args["batch_size"], 75 | pin_memory=args["pin_memory"], 76 | num_workers=args["num_workers"], 77 | shuffle=args["shuffle"], 78 | worker_init_fn=seed_worker, 79 | ) 80 | 81 | # NoiseAdder 82 | if args["add_noise"]: 83 | if args["noise_adder"] == "random": 84 | noise_adder = RandomNoiseAdder(mean=args["mean_noiseadder"], std=args["std_noiseadder"]) 85 | else: 86 | raise ValueError("Unknown noise_adder type.") 87 | 88 | # save_path 89 | save_path = osp.join(save_folder, "saved_latent_vectors.npz") 90 | 91 | # main 92 | latent_vectors_list = [] 93 | labels_list = [] 94 | with torch.no_grad(): 95 | for _, batch in tqdm(enumerate(loader)): 96 | pcs, labels = batch[0].to(device), batch[1].to(device) 97 | 98 | if args["add_noise"]: 99 | pcs = noise_adder.add_noise(pcs) 100 | 101 | try: 102 | latent_vectors = autoencoder.encode(pcs) 103 | except: 104 | latent_vectors, _ = autoencoder.forward(pcs) 105 | latent_vectors_list.append(latent_vectors.reshape(latent_vectors.shape[0], -1).cpu().numpy()) 106 | labels_list.append(labels.cpu().numpy()) 107 | latent_vectors_list = np.concatenate(latent_vectors_list, axis=0) 108 | labels_list = np.concatenate(labels_list, axis=0) 109 | results = {"latent_vectors": latent_vectors_list, "labels": labels_list} 110 | np.savez(save_path, **results) 111 | 112 | 113 | if __name__ == "__main__": 114 | main() 115 | -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/Makefile: -------------------------------------------------------------------------------- 1 | ############################################################################### 2 | # Uncomment for debugging 3 | # DEBUG := 1 4 | # Pretty build 5 | # Q ?= @ 6 | 7 | CXX := g++ 8 | PYTHON := /vinai/trungnd19/miniconda3/envs/ot-lcd/bin/python 9 | NVCC := /usr/local/cuda/bin/nvcc 10 | 11 | # PYTHON Header path 12 | PYTHON_HEADER_DIR := $(shell $(PYTHON) -c 'from distutils.sysconfig import get_python_inc; print(get_python_inc())') 13 | PYTORCH_INCLUDES := $(shell $(PYTHON) -c 'from torch.utils.cpp_extension import include_paths; [print(p) for p in include_paths()]') 14 | PYTORCH_LIBRARIES := $(shell $(PYTHON) -c 'from torch.utils.cpp_extension import library_paths; [print(p) for p in library_paths()]') 15 | 16 | # CUDA ROOT DIR that contains bin/ lib64/ and include/ 17 | # CUDA_DIR := /usr/local/cuda 18 | CUDA_DIR := $(shell $(PYTHON) -c 'from torch.utils.cpp_extension import _find_cuda_home; print(_find_cuda_home())') 19 | 20 | INCLUDE_DIRS := ./ $(CUDA_DIR)/include 21 | 22 | INCLUDE_DIRS += $(PYTHON_HEADER_DIR) 23 | INCLUDE_DIRS += $(PYTORCH_INCLUDES) 24 | 25 | # Custom (MKL/ATLAS/OpenBLAS) include and lib directories. 26 | # Leave commented to accept the defaults for your choice of BLAS 27 | # (which should work)! 28 | # BLAS_INCLUDE := /path/to/your/blas 29 | # BLAS_LIB := /path/to/your/blas 30 | 31 | ############################################################################### 32 | SRC_DIR := ./src 33 | OBJ_DIR := ./objs 34 | CPP_SRCS := $(wildcard $(SRC_DIR)/*.cpp) 35 | CU_SRCS := $(wildcard $(SRC_DIR)/*.cu) 36 | OBJS := $(patsubst $(SRC_DIR)/%.cpp,$(OBJ_DIR)/%.o,$(CPP_SRCS)) 37 | CU_OBJS := $(patsubst $(SRC_DIR)/%.cu,$(OBJ_DIR)/cuda/%.o,$(CU_SRCS)) 38 | STATIC_LIB := $(OBJ_DIR)/libmake_pytorch.a 39 | 40 | # CUDA architecture setting: going with all of them. 41 | # For CUDA < 6.0, comment the *_50 through *_61 lines for compatibility. 42 | # For CUDA < 8.0, comment the *_60 and *_61 lines for compatibility. 43 | CUDA_ARCH := -gencode arch=compute_61,code=sm_61 \ 44 | -gencode arch=compute_61,code=compute_61 \ 45 | -gencode arch=compute_52,code=sm_52 46 | 47 | # We will also explicitly add stdc++ to the link target. 48 | LIBRARIES += stdc++ cudart c10 caffe2 torch torch_python caffe2_gpu 49 | 50 | # Debugging 51 | ifeq ($(DEBUG), 1) 52 | COMMON_FLAGS += -DDEBUG -g -O0 53 | # https://gcoe-dresden.de/reaching-the-shore-with-a-fog-warning-my-eurohack-day-4-morning-session/ 54 | NVCCFLAGS += -g -G # -rdc true 55 | else 56 | COMMON_FLAGS += -DNDEBUG -O3 57 | endif 58 | 59 | WARNINGS := -Wall -Wno-sign-compare -Wcomment 60 | 61 | INCLUDE_DIRS += $(BLAS_INCLUDE) 62 | 63 | # Automatic dependency generation (nvcc is handled separately) 64 | CXXFLAGS += -MMD -MP 65 | 66 | # Complete build flags. 67 | COMMON_FLAGS += $(foreach includedir,$(INCLUDE_DIRS),-I$(includedir)) \ 68 | -DTORCH_API_INCLUDE_EXTENSION_H -D_GLIBCXX_USE_CXX11_ABI=0 69 | CXXFLAGS += -pthread -fPIC -fwrapv -std=c++11 $(COMMON_FLAGS) $(WARNINGS) 70 | NVCCFLAGS += -std=c++11 -ccbin=$(CXX) -Xcompiler -fPIC $(COMMON_FLAGS) 71 | 72 | all: $(STATIC_LIB) 73 | $(PYTHON) setup.py build 74 | @ mv build/lib.linux-x86_64-3.6/StructuralLosses .. 75 | @ mv build/lib.linux-x86_64-3.6/*.so ../StructuralLosses/ 76 | @- $(RM) -rf $(OBJ_DIR) build objs 77 | 78 | $(OBJ_DIR): 79 | @ mkdir -p $@ 80 | @ mkdir -p $@/cuda 81 | 82 | $(OBJ_DIR)/%.o: $(SRC_DIR)/%.cpp | $(OBJ_DIR) 83 | @ echo CXX $< 84 | $(Q)$(CXX) $< $(CXXFLAGS) -c -o $@ 85 | 86 | $(OBJ_DIR)/cuda/%.o: $(SRC_DIR)/%.cu | $(OBJ_DIR) 87 | @ echo NVCC $< 88 | $(Q)$(NVCC) $(NVCCFLAGS) $(CUDA_ARCH) -M $< -o ${@:.o=.d} \ 89 | -odir $(@D) 90 | $(Q)$(NVCC) $(NVCCFLAGS) $(CUDA_ARCH) -c $< -o $@ 91 | 92 | $(STATIC_LIB): $(OBJS) $(CU_OBJS) | $(OBJ_DIR) 93 | $(RM) -f $(STATIC_LIB) 94 | $(RM) -rf build dist 95 | @ echo LD -o $@ 96 | ar rc $(STATIC_LIB) $(OBJS) $(CU_OBJS) 97 | 98 | clean: 99 | @- $(RM) -rf $(OBJ_DIR) build dist ../StructuralLosses 100 | 101 | -------------------------------------------------------------------------------- /trainer/trainer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | try: 6 | from .interface import Trainer 7 | except: 8 | from interface import Trainer 9 | 10 | 11 | class AETrainer: 12 | def __init__(self): 13 | super().__init__() 14 | 15 | @staticmethod 16 | def train(ae, loss_func, optimizer, data, **kwargs): 17 | """ 18 | AUTOENCODER TRAINER 19 | ae: nn.Module 20 | the autoencoder with initial weights 21 | loss_func: nn.Module 22 | Chamfer, EMD, SWD 23 | optimizer: nn.optim 24 | optimizer, such as, SGD or Adam, with initial state 25 | dataloader: torch.utils.data.dataloader 26 | dataloader 27 | device: "cuda" or "cpu" 28 | """ 29 | ae.train() 30 | try: 31 | latent = ae.encode(data) 32 | reconstructed_data = ae.decode(latent) 33 | except: 34 | latent, reconstructed_data = ae(data) 35 | 36 | if "input" in kwargs.keys(): 37 | inp = kwargs["input"] 38 | else: 39 | inp = data 40 | 41 | result_dic = loss_func.forward(inp, reconstructed_data, latent=latent, **kwargs) 42 | 43 | optimizer.zero_grad() 44 | result_dic["loss"].backward() 45 | optimizer.step() 46 | 47 | result_dic["ae"] = ae 48 | result_dic["optimizer"] = optimizer 49 | return result_dic 50 | 51 | 52 | class LatentCodesGeneratorTrainer: 53 | def __init__(self): 54 | super().__init__() 55 | 56 | @staticmethod 57 | def train(g, ae, loss_func, optimizer, data, **kwargs): 58 | """ 59 | LATENT CODES GENERATOR TRAINER 60 | g: nn.Module 61 | the generator 62 | ae: nn.Module 63 | the pretrained autoencoder 64 | loss_func: a function that returns a dict with at least a key "loss" 65 | ChamferLoss or ASW 66 | optimizer: nn.optim 67 | optimizer for g, such as, SGD or Adam 68 | data: [batch_size, #points, #point channels] 69 | a batch of point clouds 70 | """ 71 | g.train() 72 | ae.eval() 73 | with torch.no_grad(): 74 | try: 75 | latent = ae.encode(data) 76 | except: 77 | latent, reconstructed_data = ae.forward(data) 78 | 79 | result_dic = loss_func(g, latent, **kwargs) 80 | 81 | optimizer.zero_grad() 82 | result_dic["loss"].backward() 83 | optimizer.step() 84 | 85 | result_dic["g"] = g 86 | result_dic["optimizer"] = optimizer 87 | return result_dic 88 | 89 | 90 | class ClassifierTrainer: 91 | def __init__(self): 92 | super().__init__() 93 | 94 | @staticmethod 95 | def train(classifier, loss_func, optimizer, data, gt_label): 96 | predicted_label = classifier.forward(data) 97 | loss = loss_func.forward(predicted_label, gt_label) 98 | optimizer.zero_grad() 99 | loss.backward() 100 | optimizer.step() 101 | return classifier, optimizer, loss 102 | 103 | 104 | class adasw_dynamic_eps_AETrainer: 105 | def __init__(self): 106 | super().__init__() 107 | 108 | @staticmethod 109 | def train(ae, loss_func, optimizer, data, epsilon): 110 | """ 111 | AUTOENCODER TRAINER 112 | ae: nn.Module 113 | the autoencoder with initial weights 114 | loss_func: nn.Module 115 | ChamferLoss of ASW 116 | optimizer: nn.optim 117 | optimizer, such as, SGD or Adam, with initial state 118 | dataloader: torch.utils.data.dataloader 119 | dataloader 120 | device: "cuda" or "cpu" 121 | 122 | """ 123 | ae.train() 124 | latent_vects = ae.encode(data) 125 | reconstructed_data = ae.decode(latent_vects) 126 | 127 | loss = loss_func.forward(data, reconstructed_data, epsilon) 128 | 129 | optimizer.zero_grad() 130 | loss.backward() 131 | optimizer.step() 132 | 133 | return ae, optimizer, loss 134 | 135 | 136 | Trainer.register(AETrainer) 137 | assert issubclass(AETrainer, Trainer) 138 | Trainer.register(ClassifierTrainer) 139 | assert issubclass(ClassifierTrainer, Trainer) 140 | -------------------------------------------------------------------------------- /generation/test_generation.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import random 3 | import sys 4 | 5 | import numpy as np 6 | import torch 7 | from torch.utils.data import DataLoader 8 | 9 | 10 | def seed_worker(worker_id): 11 | worker_seed = torch.initial_seed() % 2 ** 32 12 | np.random.seed(worker_seed) 13 | random.seed(worker_seed) 14 | 15 | 16 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 17 | from criteria_comparing_sets_pcs.all_metrics_calculator import AllMetricsCalculator 18 | from criteria_comparing_sets_pcs.jsd_calculator import JsdCalculator 19 | from dataset.shapenet_core55 import ShapeNetCore55XyzOnlyDataset 20 | from evaluator import SetPcsComparingBasedEvaluator 21 | from generation.train_latent_generator import MLPGenerator 22 | from models import PointNetAE 23 | from utils import create_save_folder, evaluate_on_dataset, initialize_main, load_model_for_evaluation 24 | 25 | 26 | def main(): 27 | args, logdir = initialize_main() 28 | 29 | # set seed 30 | torch.manual_seed(args["seed"]) 31 | random.seed(args["seed"]) 32 | np.random.seed(args["seed"]) 33 | 34 | # save_results folder 35 | save_folder = create_save_folder(logdir, args["save_folder"])["save_folder"] 36 | 37 | # device 38 | device = torch.device(args["device"]) 39 | 40 | # test set 41 | if args["test_set_type"] == "shapenetcore55": 42 | test_set = ShapeNetCore55XyzOnlyDataset(args["test_root"], args["num_points"], phase="test") 43 | 44 | else: 45 | raise ValueError("Unknown dataset type.") 46 | 47 | # test loader 48 | test_loader = DataLoader( 49 | test_set, 50 | batch_size=args["batch_size"], 51 | num_workers=args["num_workers"], 52 | pin_memory=True, 53 | shuffle=True, 54 | worker_init_fn=seed_worker, 55 | ) 56 | 57 | # neural net 58 | if args["architecture"] == "pointnet": 59 | model = PointNetAE( 60 | embedding_size=args["embedding_size"], 61 | input_channels=args["point_channels"], 62 | output_channels=args["point_channels"], 63 | num_points=args["num_points"], 64 | normalize=args["normalize"], 65 | ).to(device) 66 | else: 67 | raise ValueError("Unknown architecture.") 68 | 69 | try: 70 | model = load_model_for_evaluation(model, args["model_path"]) 71 | except: 72 | model = load_model_for_evaluation(model, osp.join(logdir, args["model_path"])) 73 | 74 | # evaluator 75 | save_file = osp.join(save_folder, "generation_test_results.txt") 76 | 77 | if args["evaluator_type"] == "based_on_comparing_set_pcs": 78 | evaluator = SetPcsComparingBasedEvaluator() 79 | # prior_distribution_sampler 80 | if args["prior_distribution"] == "latent_codes_generator": 81 | prior_distribution_sampler = MLPGenerator( 82 | args["latent_dim"], args["n_hidden"], args["hidden_size"], args["embedding_size"] 83 | ).to(device) 84 | 85 | try: 86 | prior_distribution_sampler = load_model_for_evaluation(prior_distribution_sampler, args["prior_path"]) 87 | except: 88 | prior_distribution_sampler = load_model_for_evaluation( 89 | prior_distribution_sampler, osp.join(save_folder, args["prior_path"]) 90 | ) 91 | 92 | else: 93 | raise ValueError("Unknown prior distribution.") 94 | 95 | if args["eval_criteria"] == "jsd": 96 | criteria_calculator = JsdCalculator() 97 | elif args["eval_criteria"] == "all_metrics": 98 | criteria_calculator = AllMetricsCalculator() 99 | args["eval_criteria"] = "jsd" 100 | else: 101 | raise ValueError("Unknown eval_criteria") 102 | 103 | eval_dic = { 104 | "prior_distribution_sampler": prior_distribution_sampler, 105 | "criteria_calculator": criteria_calculator, 106 | "batch_size": args["criteria_batch_size"], 107 | "use_EMD": args["use_EMD"], 108 | "accelerated_cd": args["accelerated_cd"], 109 | "save_file": save_file, 110 | } 111 | else: 112 | raise ValueError("Unknown evaluator type.") 113 | 114 | # main 115 | print("Evaluating...(It might take a while.)") 116 | avg_eval_value = evaluate_on_dataset(evaluator, model, test_loader, device, **eval_dic) 117 | 118 | # save results 119 | log = "{}: {}\n".format(args["eval_criteria"], avg_eval_value) 120 | with open(save_file, "a") as fp: 121 | fp.write(log) 122 | print(log) 123 | 124 | 125 | if __name__ == "__main__": 126 | main() 127 | -------------------------------------------------------------------------------- /utils/utils.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os 4 | import os.path as osp 5 | import shutil 6 | import sys 7 | import time 8 | from typing import List 9 | 10 | import torch 11 | from tqdm import tqdm 12 | 13 | 14 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 15 | 16 | 17 | def get_lr(optimizer): 18 | for param_group in optimizer.param_groups: 19 | return param_group["lr"] 20 | 21 | 22 | def initialize_main(**kwargs): 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument("--config", help="path to json config file") 25 | parser.add_argument("--logdir", help="folder to save results") 26 | parser.add_argument("--data_path", type=str, default=None, help="path to data") 27 | args = parser.parse_args() 28 | config = args.config 29 | logdir = args.logdir 30 | data_path = args.data_path 31 | print("Save results in: ", logdir) 32 | args = json.load(open(config)) 33 | 34 | if not os.path.exists(logdir): 35 | os.makedirs(logdir) 36 | print(">Logdir was created successfully at: ", logdir) 37 | else: 38 | print(">Logdir exists.") 39 | 40 | if "save_config" in kwargs.keys(): 41 | saving_config_filename = kwargs["save_config"] 42 | print(">In saving config mode.") 43 | fname = os.path.join(logdir, saving_config_filename) 44 | with open(fname, "w") as fp: 45 | json.dump(args, fp, indent=4) 46 | else: 47 | print(">Not in saving config mode.") 48 | 49 | # print hyperparameters 50 | print("You have 5s to check the hyperparameters below.") 51 | print(args) 52 | time.sleep(5) 53 | if data_path is None: 54 | return args, logdir 55 | else: 56 | return args, logdir, data_path 57 | 58 | 59 | def load_model_for_evaluation(model, model_path, **kwargs): 60 | device = next(model.parameters()).device 61 | try: 62 | model.load_state_dict(torch.load(model_path, map_location=device)) 63 | except: 64 | model.load_state_dict(torch.load(model_path, map_location=device)[kwargs["key"]]) 65 | model.eval() 66 | return model 67 | 68 | 69 | def create_save_folder(logdir, save_folder_branch): 70 | save_folder = osp.join(logdir, save_folder_branch) 71 | if not osp.exists(save_folder): 72 | os.makedirs(save_folder) 73 | print(">Save_folder was created successfully at:", save_folder) 74 | else: 75 | print(">Save_folder {} is existing.".format(save_folder)) 76 | print(">Do you want to remove it?") 77 | answer = None 78 | while answer not in ("yes", "no"): 79 | try: 80 | answer = input("Enter 'yes' or 'no': ") 81 | except: 82 | print("Please enter your answer again.") 83 | answer = raw_input("Enter 'yes' or 'no': ") 84 | if answer == "yes": 85 | shutil.rmtree(save_folder) 86 | os.makedirs(save_folder) 87 | elif answer == "no": 88 | print("SOME FILES WILL BE OVERWRITTEN OR APPENDED.") 89 | print("If you do not want this, please stop during next 5s.") 90 | time.sleep(5) 91 | else: 92 | print("Please enter yes or no.") 93 | return {"save_folder": save_folder} 94 | 95 | 96 | def evaluate_on_dataset(evaluator, model, dataloader, device="cuda", **eval_dic): 97 | eval_value_list = [] 98 | with torch.no_grad(): 99 | for _, batch in tqdm(enumerate(dataloader)): 100 | val_data = batch.to(device) 101 | result_dic = evaluator.evaluate(model, val_data, **eval_dic) 102 | if torch.is_tensor(result_dic["evaluation"]): 103 | eval_value_list.append(result_dic["evaluation"].item()) 104 | else: 105 | eval_value_list.append(result_dic["evaluation"]) 106 | # end for 107 | avg_eval_value = sum(eval_value_list) / len(eval_value_list) 108 | return avg_eval_value 109 | 110 | 111 | def load_numbers_from_file_and_compute_mean(file: str, *args, **kwargs): 112 | """ 113 | file: path to file 114 | each line in file contains only one number. 115 | """ 116 | s = 0.0 117 | count = 0 118 | with open(file, "r") as f: 119 | lines = f.readlines() 120 | for line in lines: 121 | s += float(line) 122 | count += 1 123 | print(s / count) 124 | 125 | 126 | def smooth(scalars: List[float], weight: float) -> List[float]: # Weight between 0 and 1 127 | last = scalars[0] # First value in the plot (first timestep) 128 | smoothed = list() 129 | for point in scalars: 130 | smoothed_val = last * weight + (1 - weight) * point # Calculate smoothed value 131 | smoothed.append(smoothed_val) # Save it 132 | last = smoothed_val # Anchor the last smoothed value 133 | 134 | return smoothed 135 | -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/src/nndistance.cu: -------------------------------------------------------------------------------- 1 | 2 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 3 | const int batch=512; 4 | __shared__ float buf[batch*3]; 5 | for (int i=blockIdx.x;ibest){ 117 | result[(i*n+j)]=best; 118 | result_i[(i*n+j)]=best_i; 119 | } 120 | } 121 | __syncthreads(); 122 | } 123 | } 124 | } 125 | void nndistance(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){ 126 | NmDistanceKernel<<>>(b,n,xyz,m,xyz2,result,result_i); 127 | NmDistanceKernel<<>>(b,m,xyz2,n,xyz,result2,result2_i); 128 | } 129 | __global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){ 130 | for (int i=blockIdx.x;i>>(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_xyz1,grad_xyz2); 153 | NmDistanceGradKernel<<>>(b,m,xyz2,n,xyz1,grad_dist2,idx2,grad_xyz2,grad_xyz1); 154 | } 155 | 156 | -------------------------------------------------------------------------------- /classification/classification_train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | import random 4 | import sys 5 | 6 | import numpy as np 7 | import torch 8 | import torch.utils.data as data 9 | from torch.nn import CrossEntropyLoss 10 | from torch.optim import SGD, Adam 11 | from tqdm import tqdm 12 | 13 | 14 | def seed_worker(worker_id): 15 | worker_seed = torch.initial_seed() % 2 ** 32 16 | np.random.seed(worker_seed) 17 | random.seed(worker_seed) 18 | 19 | 20 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 21 | from classifier import MLPClassifier 22 | from dataset.modelnet40 import LatentCapsulesModelNet40, LatentVectorsModelNet40 23 | from models.utils import init_weights 24 | from saver import GeneralSaver 25 | from trainer import ClassifierTrainer 26 | from utils.utils import initialize_main 27 | 28 | 29 | def main(): 30 | args, logdir = initialize_main() 31 | 32 | # set seed 33 | torch.manual_seed(args["seed"]) 34 | random.seed(args["seed"]) 35 | np.random.seed(args["seed"]) 36 | 37 | # save_results folder 38 | save_folder = osp.join(logdir, args["save_folder"]) 39 | if not osp.exists(save_folder): 40 | os.makedirs(save_folder) 41 | 42 | # device 43 | device = torch.device(args["device"]) 44 | 45 | # root 46 | args["root"] = osp.join(logdir, "latent_codes/model/modelnet40-train/saved_latent_vectors.npz") 47 | 48 | # dataloader 49 | if args["root"].endswith(".npz"): 50 | dset = LatentVectorsModelNet40(args["root"]) # root is a npz file 51 | elif args["root"].endswith(".h5"): 52 | dset = LatentCapsulesModelNet40(args["root"]) 53 | else: 54 | raise Exception("Unknown dataset.") 55 | loader = data.DataLoader( 56 | dset, 57 | batch_size=args["batch_size"], 58 | pin_memory=args["pin_memory"], 59 | num_workers=args["num_workers"], 60 | shuffle=args["shuffle"], 61 | worker_init_fn=seed_worker, 62 | ) 63 | 64 | # classifier 65 | classifier = MLPClassifier( 66 | args["input_size"], args["output_size"], args["dropout_p"], [int(i) for i in args["hidden_sizes"].split(",")] 67 | ).to(device) 68 | 69 | # optimizer 70 | if args["optimizer"] == "sgd": 71 | optimizer = SGD( 72 | classifier.parameters(), 73 | lr=args["learning_rate"], 74 | momentum=args["momentum"], 75 | weight_decay=args["weight_decay"], 76 | ) 77 | elif args["optimizer"] == "adam": 78 | optimizer = Adam( 79 | classifier.parameters(), 80 | lr=args["learning_rate"], 81 | betas=(0.9, 0.999), 82 | weight_decay=args["weight_decay"], 83 | ) 84 | else: 85 | raise Exception("Unknown optimizer.") 86 | 87 | # init weights 88 | if osp.isfile(osp.join(save_folder, args["checkpoint"])): 89 | print(">Init weights with {}".format(args["checkpoint"])) 90 | checkpoint = torch.load(osp.join(save_folder, args["checkpoint"])) 91 | classifier.load_state_dict(checkpoint["classifier"]) 92 | optimizer.load_state_dict(checkpoint["optimizer"]) 93 | else: 94 | print(">Init weights with Xavier") 95 | classifier.apply(init_weights) 96 | 97 | # loss 98 | loss_func = CrossEntropyLoss() 99 | 100 | # main 101 | best_loss = args["best_loss"] 102 | best_epoch = args["best_epoch"] 103 | start_epoch = int(args["start_epoch"]) 104 | num_epochs = int(args["num_epochs"]) 105 | model_path = osp.join(save_folder, "model.pth") 106 | train_log = osp.join(save_folder, "train.log") 107 | 108 | classifier.train() 109 | for epoch in tqdm(range(start_epoch, num_epochs)): 110 | loss_list = [] 111 | for _, (batch, labels) in tqdm(enumerate(loader)): 112 | batch = batch.to(device) 113 | labels = labels.to(device).squeeze().type(torch.long) 114 | classifier, optimizer, loss = ClassifierTrainer.train(classifier, loss_func, optimizer, batch, labels) 115 | loss_list.append(loss.item()) 116 | avg_loss = sum(loss_list) / len(loss_list) 117 | 118 | # save checkpoint 119 | checkpoint_path = osp.join(save_folder, "latest.pth") 120 | GeneralSaver.save_checkpoint(classifier, optimizer, checkpoint_path, "classifier") 121 | 122 | if epoch % args["epoch_gap_for_save"] == 0: 123 | checkpoint_path = os.path.join(save_folder, "epoch_" + str(epoch) + ".pth") 124 | GeneralSaver.save_best_weights(classifier, checkpoint_path) 125 | 126 | # save best model 127 | if avg_loss < best_loss: 128 | best_loss = avg_loss 129 | best_epoch = epoch 130 | GeneralSaver.save_best_weights(classifier, model_path) 131 | 132 | # report 133 | log = "Epoch {}| loss: {}\n".format(epoch, avg_loss) 134 | log_best = "Best epoch {}| best loss: {}".format(best_epoch, best_loss) 135 | with open(train_log, "a") as fp: 136 | fp.write(log) 137 | print(log) 138 | print(log_best) 139 | print("---------------------------------------------------------------------------------------") 140 | # end for 141 | print("save_folder: ", save_folder) 142 | 143 | 144 | if __name__ == "__main__": 145 | main() 146 | -------------------------------------------------------------------------------- /generation/deepul/hw4_helper.py: -------------------------------------------------------------------------------- 1 | # import cv2 2 | import deepul.pytorch_util as ptu 3 | import numpy as np 4 | import scipy.ndimage 5 | import torch.nn as nn 6 | import torch.utils.data 7 | import torchvision 8 | from PIL import Image as PILImage 9 | from torchvision import transforms as transforms 10 | 11 | from .hw4_utils.hw4_models import GoogLeNet 12 | from .utils import * 13 | 14 | 15 | CLASSES = ("plane", "car", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck") 16 | 17 | import math 18 | import sys 19 | 20 | import numpy as np 21 | 22 | 23 | softmax = None 24 | model = None 25 | device = torch.device("cuda:0") 26 | 27 | 28 | def plot_gan_training(losses, title, fname): 29 | plt.figure() 30 | n_itr = len(losses) 31 | xs = np.arange(n_itr) 32 | 33 | plt.plot(xs, losses, label="loss") 34 | plt.legend() 35 | plt.title(title) 36 | plt.xlabel("Training Iteration") 37 | plt.ylabel("Loss") 38 | savefig(fname) 39 | 40 | 41 | def q1_gan_plot(data, samples, xs, ys, title, fname): 42 | plt.figure() 43 | plt.hist(samples, bins=50, density=True, alpha=0.7, label="fake") 44 | plt.hist(data, bins=50, density=True, alpha=0.7, label="real") 45 | 46 | plt.plot(xs, ys, label="discrim") 47 | plt.legend() 48 | plt.title(title) 49 | savefig(fname) 50 | 51 | 52 | ###################### 53 | ##### Question 1 ##### 54 | ###################### 55 | 56 | 57 | def q1_data(n=20000): 58 | assert n % 2 == 0 59 | gaussian1 = np.random.normal(loc=-1, scale=0.25, size=(n // 2,)) 60 | gaussian2 = np.random.normal(loc=0.5, scale=0.5, size=(n // 2,)) 61 | data = (np.concatenate([gaussian1, gaussian2]) + 1).reshape([-1, 1]) 62 | scaled_data = (data - np.min(data)) / (np.max(data) - np.min(data) + 1e-8) 63 | return 2 * scaled_data - 1 64 | 65 | 66 | def visualize_q1_dataset(): 67 | data = q1_data() 68 | plt.hist(data, bins=50, alpha=0.7, label="train data") 69 | plt.legend() 70 | plt.show() 71 | 72 | 73 | def q1_save_results(part, fn): 74 | data = q1_data() 75 | losses, samples1, xs1, ys1, samples_end, xs_end, ys_end = fn(data) 76 | 77 | # loss plot 78 | plot_gan_training(losses, "Q1{} Losses".format(part), "results/q1{}_losses.png".format(part)) 79 | 80 | # samples 81 | q1_gan_plot(data, samples1, xs1, ys1, "Q1{} Epoch 1".format(part), "results/q1{}_epoch1.png".format(part)) 82 | q1_gan_plot(data, samples_end, xs_end, ys_end, "Q1{} Final".format(part), "results/q1{}_final.png".format(part)) 83 | 84 | 85 | ###################### 86 | ##### Question 2 ##### 87 | ###################### 88 | 89 | 90 | def calculate_is(samples): 91 | assert type(samples[0]) == np.ndarray 92 | assert len(samples[0].shape) == 3 93 | 94 | model = GoogLeNet().to(ptu.device) 95 | model.load_state_dict(torch.load("deepul/deepul/hw4_utils/classifier.pt")) 96 | softmax = nn.Sequential(model, nn.Softmax(dim=1)) 97 | 98 | bs = 100 99 | softmax.eval() 100 | with torch.no_grad(): 101 | preds = [] 102 | n_batches = int(math.ceil(float(len(samples)) / float(bs))) 103 | for i in range(n_batches): 104 | sys.stdout.write(".") 105 | sys.stdout.flush() 106 | inp = ptu.FloatTensor(samples[(i * bs) : min((i + 1) * bs, len(samples))]) 107 | pred = ptu.get_numpy(softmax(inp)) 108 | preds.append(pred) 109 | preds = np.concatenate(preds, 0) 110 | kl = preds * (np.log(preds) - np.log(np.expand_dims(np.mean(preds, 0), 0))) 111 | kl = np.mean(np.sum(kl, 1)) 112 | return np.exp(kl) 113 | 114 | 115 | def load_q2_data(): 116 | train_data = torchvision.datasets.CIFAR10( 117 | "./data", transform=torchvision.transforms.ToTensor(), download=True, train=True 118 | ) 119 | return train_data 120 | 121 | 122 | def visualize_q2_data(): 123 | train_data = load_q2_data() 124 | imgs = train_data.data[:100] 125 | show_samples(imgs, title=f"CIFAR-10 Samples") 126 | 127 | 128 | def q2_save_results(fn): 129 | train_data = load_q2_data() 130 | train_data = train_data.data.transpose((0, 3, 1, 2)) / 255.0 131 | train_losses, samples = fn(train_data) 132 | 133 | print("Inception score:", calculate_is(samples.transpose([0, 3, 1, 2]))) 134 | plot_gan_training(train_losses, "Q2 Losses", "results/q2_losses.png") 135 | show_samples(samples[:100] * 255.0, fname="results/q2_samples.png", title=f"CIFAR-10 generated samples") 136 | 137 | 138 | ###################### 139 | ##### Question 3 ##### 140 | ###################### 141 | 142 | 143 | def load_q3_data(): 144 | transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]) 145 | train_data = torchvision.datasets.MNIST(root="./data", train=True, download=True, transform=transform) 146 | test_data = torchvision.datasets.MNIST(root="./data", train=False, download=True, transform=transform) 147 | return train_data, test_data 148 | 149 | 150 | def visualize_q3_data(): 151 | train_data, _ = load_q3_data() 152 | imgs = train_data.data[:100] 153 | show_samples(imgs.reshape([100, 28, 28, 1]) * 255.0, title=f"MNIST samples") 154 | 155 | 156 | def plot_q3_supervised(pretrained_losses, random_losses, title, fname): 157 | plt.figure() 158 | xs = np.arange(len(pretrained_losses)) 159 | plt.plot(xs, pretrained_losses, label="bigan") 160 | xs = np.arange(len(random_losses)) 161 | plt.plot(xs, random_losses, label="random init") 162 | plt.legend() 163 | plt.title(title) 164 | savefig(fname) 165 | 166 | 167 | def q3_save_results(fn): 168 | train_data, test_data = load_q3_data() 169 | gan_losses, samples, reconstructions, pretrained_losses, random_losses = fn(train_data, test_data) 170 | 171 | plot_gan_training(gan_losses, "Q3 Losses", "results/q3_gan_losses.png") 172 | plot_q3_supervised( 173 | pretrained_losses, random_losses, "Linear classification losses", "results/q3_supervised_losses.png" 174 | ) 175 | show_samples(samples * 255.0, fname="results/q3_samples.png", title="BiGAN generated samples") 176 | show_samples( 177 | reconstructions * 255.0, nrow=20, fname="results/q3_reconstructions.png", title=f"BiGAN reconstructions" 178 | ) 179 | print("BiGAN final linear classification loss:", pretrained_losses[-1]) 180 | print("Random encoder linear classification loss:", random_losses[-1]) 181 | -------------------------------------------------------------------------------- /models/pcn.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | import os.path as osp 4 | import sys 5 | from collections import OrderedDict 6 | 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | import torch.nn.parallel 11 | import torch.utils.data 12 | from torch.autograd import Variable 13 | 14 | 15 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 16 | 17 | 18 | class ConvLayer(nn.Module): 19 | def __init__(self): 20 | super(ConvLayer, self).__init__() 21 | self.conv1 = torch.nn.Conv1d(3, 64, 1) 22 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 23 | self.bn1 = nn.BatchNorm1d(64) 24 | self.bn2 = nn.BatchNorm1d(128) 25 | 26 | def forward(self, x): 27 | x = F.relu(self.bn1(self.conv1(x))) 28 | x = F.relu(self.bn2(self.conv2(x))) 29 | return x 30 | 31 | 32 | class PrimaryPointCapsLayer(nn.Module): 33 | def __init__(self, prim_vec_size=8, num_points=2048): 34 | super(PrimaryPointCapsLayer, self).__init__() 35 | self.capsules = nn.ModuleList( 36 | [ 37 | torch.nn.Sequential( 38 | OrderedDict( 39 | [ 40 | ("conv3", torch.nn.Conv1d(128, 1024, 1)), 41 | ("bn3", nn.BatchNorm1d(1024)), 42 | ("mp1", torch.nn.MaxPool1d(num_points)), 43 | ] 44 | ) 45 | ) 46 | for _ in range(prim_vec_size) 47 | ] 48 | ) 49 | 50 | def forward(self, x): 51 | u = [capsule(x) for capsule in self.capsules] 52 | u = torch.stack(u, dim=2) 53 | return self.squash(u.squeeze()) 54 | 55 | def squash(self, input_tensor): 56 | squared_norm = (input_tensor ** 2).sum(-1, keepdim=True) 57 | output_tensor = squared_norm * input_tensor / ((1.0 + squared_norm) * torch.sqrt(squared_norm)) 58 | if output_tensor.dim() == 2: 59 | output_tensor = torch.unsqueeze(output_tensor, 0) 60 | return output_tensor 61 | 62 | 63 | class LatentCapsLayer(nn.Module): 64 | def __init__(self, latent_caps_size=16, prim_caps_size=1024, prim_vec_size=16, latent_vec_size=64, device="cuda"): 65 | super(LatentCapsLayer, self).__init__() 66 | self.prim_vec_size = prim_vec_size 67 | self.prim_caps_size = prim_caps_size 68 | self.latent_caps_size = latent_caps_size 69 | self.W = nn.Parameter(0.01 * torch.randn(latent_caps_size, prim_caps_size, latent_vec_size, prim_vec_size)) 70 | self.device = device 71 | 72 | def forward(self, x): 73 | u_hat = torch.squeeze(torch.matmul(self.W, x[:, None, :, :, None]), dim=-1) # .to(self.device) 74 | u_hat_detached = u_hat.detach() 75 | b_ij = Variable(torch.zeros(x.size(0), self.latent_caps_size, self.prim_caps_size)).to(self.device) # .cuda() 76 | num_iterations = 3 77 | for iteration in range(num_iterations): 78 | c_ij = F.softmax(b_ij, 1) 79 | if iteration == num_iterations - 1: 80 | v_j = self.squash(torch.sum(c_ij[:, :, :, None] * u_hat, dim=-2, keepdim=True)) 81 | else: 82 | v_j = self.squash(torch.sum(c_ij[:, :, :, None] * u_hat_detached, dim=-2, keepdim=True)) 83 | b_ij = b_ij + torch.sum(v_j * u_hat_detached, dim=-1) 84 | return v_j.squeeze(-2) 85 | 86 | def squash(self, input_tensor): 87 | squared_norm = (input_tensor ** 2).sum(-1, keepdim=True) 88 | output_tensor = squared_norm * input_tensor / ((1.0 + squared_norm) * torch.sqrt(squared_norm)) 89 | return output_tensor 90 | 91 | 92 | class PointGenCon(nn.Module): 93 | def __init__(self, bottleneck_size=2500): 94 | self.bottleneck_size = bottleneck_size 95 | super(PointGenCon, self).__init__() 96 | self.conv1 = torch.nn.Conv1d(self.bottleneck_size, self.bottleneck_size, 1) 97 | self.conv2 = torch.nn.Conv1d(self.bottleneck_size, int(self.bottleneck_size / 2), 1) 98 | self.conv3 = torch.nn.Conv1d(int(self.bottleneck_size / 2), int(self.bottleneck_size / 4), 1) 99 | self.conv4 = torch.nn.Conv1d(int(self.bottleneck_size / 4), 3, 1) 100 | self.th = torch.nn.Tanh() 101 | self.bn1 = torch.nn.BatchNorm1d(self.bottleneck_size) 102 | self.bn2 = torch.nn.BatchNorm1d(int(self.bottleneck_size / 2)) 103 | self.bn3 = torch.nn.BatchNorm1d(int(self.bottleneck_size / 4)) 104 | 105 | def forward(self, x): 106 | x = F.relu(self.bn1(self.conv1(x))) 107 | x = F.relu(self.bn2(self.conv2(x))) 108 | x = F.relu(self.bn3(self.conv3(x))) 109 | x = self.th(self.conv4(x)) 110 | return x 111 | 112 | 113 | class CapsDecoder(nn.Module): 114 | def __init__(self, latent_caps_size, latent_vec_size, num_points): 115 | super(CapsDecoder, self).__init__() 116 | self.latent_caps_size = latent_caps_size 117 | self.bottleneck_size = latent_vec_size 118 | self.num_points = num_points 119 | self.nb_primitives = int(num_points / latent_caps_size) 120 | self.decoder = nn.ModuleList( 121 | [PointGenCon(bottleneck_size=self.bottleneck_size + 2) for i in range(0, self.nb_primitives)] 122 | ) 123 | 124 | def forward(self, x): 125 | outs = [] 126 | for i in range(0, self.nb_primitives): 127 | rand_grid = Variable(torch.cuda.FloatTensor(x.size(0), 2, self.latent_caps_size)) 128 | rand_grid.data.uniform_(0, 1) 129 | y = torch.cat((rand_grid, x.transpose(2, 1)), 1).contiguous() 130 | outs.append(self.decoder[i](y)) 131 | return torch.cat(outs, 2).contiguous() 132 | 133 | 134 | class PointCapsNet(nn.Module): 135 | def __init__(self, prim_caps_size, prim_vec_size, latent_caps_size, latent_vec_size, num_points): 136 | super(PointCapsNet, self).__init__() 137 | self.conv_layer = ConvLayer() 138 | self.primary_point_caps_layer = PrimaryPointCapsLayer(prim_vec_size, num_points) 139 | self.latent_caps_layer = LatentCapsLayer(latent_caps_size, prim_caps_size, prim_vec_size, latent_vec_size) 140 | self.caps_decoder = CapsDecoder(latent_caps_size, latent_vec_size, num_points) 141 | 142 | def forward(self, data): # data: [batch size, num points, dim] 143 | x1 = self.conv_layer(data.transpose(2, 1)) 144 | x2 = self.primary_point_caps_layer(x1) 145 | latent_capsules = self.latent_caps_layer(x2) 146 | reconstructions = self.caps_decoder(latent_capsules) 147 | return latent_capsules, reconstructions.transpose(2, 1) 148 | -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/src/structural_loss.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "src/approxmatch.cuh" 5 | #include "src/nndistance.cuh" 6 | 7 | #include 8 | #include 9 | 10 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 11 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 12 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 13 | 14 | /* 15 | input: 16 | set1 : batch_size * #dataset_points * 3 17 | set2 : batch_size * #query_points * 3 18 | returns: 19 | match : batch_size * #query_points * #dataset_points 20 | */ 21 | // temp: TensorShape{b,(n+m)*2} 22 | std::vector ApproxMatch(at::Tensor set_d, at::Tensor set_q) { 23 | //std::cout << "[ApproxMatch] Called." << std::endl; 24 | int64_t batch_size = set_d.size(0); 25 | int64_t n_dataset_points = set_d.size(1); // n 26 | int64_t n_query_points = set_q.size(1); // m 27 | //std::cout << "[ApproxMatch] batch_size:" << batch_size << std::endl; 28 | at::Tensor match = torch::empty({batch_size, n_query_points, n_dataset_points}, torch::TensorOptions().dtype(torch::kFloat32).device(set_d.device())); 29 | at::Tensor temp = torch::empty({batch_size, (n_query_points+n_dataset_points)*2}, torch::TensorOptions().dtype(torch::kFloat32).device(set_d.device())); 30 | CHECK_INPUT(set_d); 31 | CHECK_INPUT(set_q); 32 | CHECK_INPUT(match); 33 | CHECK_INPUT(temp); 34 | 35 | approxmatch(batch_size,n_dataset_points,n_query_points,set_d.data(),set_q.data(),match.data(),temp.data(), at::cuda::getCurrentCUDAStream()); 36 | return {match, temp}; 37 | } 38 | 39 | at::Tensor MatchCost(at::Tensor set_d, at::Tensor set_q, at::Tensor match) { 40 | //std::cout << "[MatchCost] Called." << std::endl; 41 | int64_t batch_size = set_d.size(0); 42 | int64_t n_dataset_points = set_d.size(1); // n 43 | int64_t n_query_points = set_q.size(1); // m 44 | //std::cout << "[MatchCost] batch_size:" << batch_size << std::endl; 45 | at::Tensor out = torch::empty({batch_size}, torch::TensorOptions().dtype(torch::kFloat32).device(set_d.device())); 46 | CHECK_INPUT(set_d); 47 | CHECK_INPUT(set_q); 48 | CHECK_INPUT(match); 49 | CHECK_INPUT(out); 50 | matchcost(batch_size,n_dataset_points,n_query_points,set_d.data(),set_q.data(),match.data(),out.data(),at::cuda::getCurrentCUDAStream()); 51 | return out; 52 | } 53 | 54 | std::vector MatchCostGrad(at::Tensor set_d, at::Tensor set_q, at::Tensor match) { 55 | //std::cout << "[MatchCostGrad] Called." << std::endl; 56 | int64_t batch_size = set_d.size(0); 57 | int64_t n_dataset_points = set_d.size(1); // n 58 | int64_t n_query_points = set_q.size(1); // m 59 | //std::cout << "[MatchCostGrad] batch_size:" << batch_size << std::endl; 60 | at::Tensor grad1 = torch::empty({batch_size,n_dataset_points,3}, torch::TensorOptions().dtype(torch::kFloat32).device(set_d.device())); 61 | at::Tensor grad2 = torch::empty({batch_size,n_query_points,3}, torch::TensorOptions().dtype(torch::kFloat32).device(set_d.device())); 62 | CHECK_INPUT(set_d); 63 | CHECK_INPUT(set_q); 64 | CHECK_INPUT(match); 65 | CHECK_INPUT(grad1); 66 | CHECK_INPUT(grad2); 67 | matchcostgrad(batch_size,n_dataset_points,n_query_points,set_d.data(),set_q.data(),match.data(),grad1.data(),grad2.data(),at::cuda::getCurrentCUDAStream()); 68 | return {grad1, grad2}; 69 | } 70 | 71 | 72 | /* 73 | input: 74 | set_d : batch_size * #dataset_points * 3 75 | set_q : batch_size * #query_points * 3 76 | returns: 77 | dist1, idx1 : batch_size * #dataset_points 78 | dist2, idx2 : batch_size * #query_points 79 | */ 80 | std::vector NNDistance(at::Tensor set_d, at::Tensor set_q) { 81 | //std::cout << "[NNDistance] Called." << std::endl; 82 | int64_t batch_size = set_d.size(0); 83 | int64_t n_dataset_points = set_d.size(1); // n 84 | int64_t n_query_points = set_q.size(1); // m 85 | //std::cout << "[NNDistance] batch_size:" << batch_size << std::endl; 86 | at::Tensor dist1 = torch::empty({batch_size, n_dataset_points}, torch::TensorOptions().dtype(torch::kFloat32).device(set_d.device())); 87 | at::Tensor idx1 = torch::empty({batch_size, n_dataset_points}, torch::TensorOptions().dtype(torch::kInt32).device(set_d.device())); 88 | at::Tensor dist2 = torch::empty({batch_size, n_query_points}, torch::TensorOptions().dtype(torch::kFloat32).device(set_d.device())); 89 | at::Tensor idx2 = torch::empty({batch_size, n_query_points}, torch::TensorOptions().dtype(torch::kInt32).device(set_d.device())); 90 | CHECK_INPUT(set_d); 91 | CHECK_INPUT(set_q); 92 | CHECK_INPUT(dist1); 93 | CHECK_INPUT(idx1); 94 | CHECK_INPUT(dist2); 95 | CHECK_INPUT(idx2); 96 | // void nndistance(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream); 97 | nndistance(batch_size,n_dataset_points,set_d.data(),n_query_points,set_q.data(),dist1.data(),idx1.data(),dist2.data(),idx2.data(), at::cuda::getCurrentCUDAStream()); 98 | return {dist1, idx1, dist2, idx2}; 99 | } 100 | 101 | std::vector NNDistanceGrad(at::Tensor set_d, at::Tensor set_q, at::Tensor idx1, at::Tensor idx2, at::Tensor grad_dist1, at::Tensor grad_dist2) { 102 | //std::cout << "[NNDistanceGrad] Called." << std::endl; 103 | int64_t batch_size = set_d.size(0); 104 | int64_t n_dataset_points = set_d.size(1); // n 105 | int64_t n_query_points = set_q.size(1); // m 106 | //std::cout << "[NNDistanceGrad] batch_size:" << batch_size << std::endl; 107 | at::Tensor grad1 = torch::empty({batch_size,n_dataset_points,3}, torch::TensorOptions().dtype(torch::kFloat32).device(set_d.device())); 108 | at::Tensor grad2 = torch::empty({batch_size,n_query_points,3}, torch::TensorOptions().dtype(torch::kFloat32).device(set_d.device())); 109 | CHECK_INPUT(set_d); 110 | CHECK_INPUT(set_q); 111 | CHECK_INPUT(idx1); 112 | CHECK_INPUT(idx2); 113 | CHECK_INPUT(grad_dist1); 114 | CHECK_INPUT(grad_dist2); 115 | CHECK_INPUT(grad1); 116 | CHECK_INPUT(grad2); 117 | //void nndistancegrad(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream); 118 | nndistancegrad(batch_size,n_dataset_points,set_d.data(),n_query_points,set_q.data(), 119 | grad_dist1.data(),idx1.data(), 120 | grad_dist2.data(),idx2.data(), 121 | grad1.data(),grad2.data(), 122 | at::cuda::getCurrentCUDAStream()); 123 | return {grad1, grad2}; 124 | } 125 | 126 | -------------------------------------------------------------------------------- /reconstruction/reconstruction_test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os 4 | import os.path as osp 5 | import random 6 | import shutil 7 | import statistics 8 | import sys 9 | import time 10 | 11 | import matplotlib.pyplot as plt 12 | import numpy as np 13 | import torch 14 | import torch.utils.data as data 15 | from tqdm import tqdm 16 | 17 | 18 | def seed_worker(worker_id): 19 | worker_seed = torch.initial_seed() % 2 ** 32 20 | np.random.seed(worker_seed) 21 | random.seed(worker_seed) 22 | 23 | 24 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 25 | from add_noise_to_data.random_noise import RandomNoiseAdder 26 | from dataset.modelnet40 import ModelNet40 27 | from dataset.shapenet_core55 import ShapeNetCore55XyzOnlyDataset 28 | from loss import EMD, SWD, Chamfer 29 | from models import PointCapsNet, PointNetAE 30 | from utils import load_model_for_evaluation 31 | 32 | 33 | def main(): 34 | parser = argparse.ArgumentParser() 35 | parser.add_argument("--config", type=str, help="path to config file") 36 | parser.add_argument("--logdir", type=str, help="folder contains weights") 37 | parser.add_argument("--data_path", default="dataset/modelnet40_ply_hdf5_2048/", type=str, help="path to data") 38 | args = parser.parse_args() 39 | config = args.config 40 | logdir = args.logdir 41 | data_path = args.data_path 42 | args = json.load(open(config)) 43 | 44 | # set seed 45 | torch.manual_seed(args["seed"]) 46 | random.seed(args["seed"]) 47 | np.random.seed(args["seed"]) 48 | 49 | # save_results folder 50 | save_folder = osp.join(logdir, args["save_folder"]) 51 | if not osp.exists(save_folder): 52 | os.makedirs(save_folder) 53 | print(">Save_folder was created successfully at:", save_folder) 54 | else: 55 | print(">Folder {} is existing.".format(save_folder)) 56 | print(">Do you want to remove it?") 57 | answer = None 58 | while answer not in ("yes", "no"): 59 | answer = input("Enter 'yes' or 'no': ") 60 | if answer == "yes": 61 | shutil.rmtree(save_folder) 62 | os.makedirs(save_folder) 63 | elif answer == "no": 64 | print("SOME FILES WILL BE OVERWRITTEN OR APPENDED.") 65 | print("If you do not want this, please stop during next 30s.") 66 | time.sleep(30) 67 | else: 68 | print("Please enter yes or no.") 69 | fname = osp.join(save_folder, "config.json") 70 | with open(fname, "w") as fp: 71 | json.dump(args, fp, indent=4) 72 | 73 | # print hyperparameters 74 | print("You have 5s to check the hyperparameters below.") 75 | print(args) 76 | time.sleep(5) 77 | 78 | # device 79 | device = torch.device(args["device"]) 80 | 81 | # NoiseAdder 82 | if args["add_noise"]: 83 | 84 | if args["noise_adder"] == "random": 85 | noise_adder = RandomNoiseAdder(mean=args["mean_noiseadder"], std=args["std_noiseadder"]) 86 | else: 87 | raise ValueError("Unknown noise_adder type.") 88 | 89 | # dataloader 90 | if args["dataset_type"] == "shapenet55": 91 | dset = ShapeNetCore55XyzOnlyDataset(data_path, args["num_points"], "test") # root is a npz file 92 | elif args["dataset_type"] == "modelnet40": 93 | dset = ModelNet40(data_path, num_points=args["num_points"]) # root is a folder containing h5 file 94 | else: 95 | raise ValueError("Unknown dataset type.") 96 | 97 | loader = data.DataLoader( 98 | dset, 99 | batch_size=1, 100 | pin_memory=True, 101 | num_workers=args["num_workers"], 102 | shuffle=args["shuffle"], 103 | worker_init_fn=seed_worker, 104 | ) 105 | 106 | # distance 107 | chamfer = Chamfer() 108 | swd = SWD(args["num_projs"], device) 109 | emd = EMD() 110 | 111 | # architecture 112 | if args["architecture"] == "pointnet": 113 | ae = PointNetAE( 114 | args["embedding_size"], 115 | args["input_channels"], 116 | args["input_channels"], 117 | args["num_points"], 118 | args["normalize"], 119 | ).to(device) 120 | 121 | elif args["architecture"] == "pcn": 122 | ae = PointCapsNet( 123 | args["prim_caps_size"], 124 | args["prim_vec_size"], 125 | args["latent_caps_size"], 126 | args["latent_vec_size"], 127 | args["num_points"], 128 | ).to(device) 129 | 130 | else: 131 | raise ValueError("Unknown architecture.") 132 | 133 | try: 134 | ae = load_model_for_evaluation(ae, args["model_path"]) 135 | except: 136 | try: 137 | ae = load_model_for_evaluation(ae, osp.join(logdir, args["model_path"])) 138 | except: 139 | in_dic = {"key": "autoencoder"} 140 | ae = load_model_for_evaluation(ae, osp.join(logdir, args["model_path"]), **in_dic) 141 | 142 | # test 143 | chamfer_list = [] 144 | swd_list = [] 145 | emd_list = [] 146 | 147 | ae.eval() 148 | with torch.no_grad(): 149 | for _, batch in tqdm(enumerate(loader)): 150 | if args["dataset_type"] in ["modelnet40"]: 151 | batch = batch[0].to(device) 152 | else: 153 | batch = batch.to(device) 154 | 155 | inp = batch.detach().clone() 156 | 157 | if args["add_noise"]: 158 | args["test_origin"] = True 159 | batch = noise_adder.add_noise(batch) 160 | 161 | try: 162 | reconstruction = ae.decode(ae.encode(batch)) 163 | except: 164 | latent, reconstruction = ae.forward(batch) 165 | 166 | if args["test_origin"]: 167 | reconstruction, batch = reconstruction[:, :, :3], inp 168 | 169 | chamfer_list.append(chamfer.forward(reconstruction, batch)["loss"].item()) 170 | swd_list.append(np.sqrt(swd.forward(reconstruction, batch)["loss"].item())) 171 | 172 | try: 173 | emd_list.append(emd.forward(reconstruction, batch)["loss"].item()) 174 | except: 175 | # To fix bug: RuntimeError: set_d.is_contiguous() INTERNAL ASSERT FAILED at src/structural_loss.cpp:30 176 | reconstruction = reconstruction.contiguous() 177 | batch = batch.contiguous() 178 | emd_list.append(emd.forward(reconstruction, batch)["loss"].item()) 179 | # report 180 | test_log = osp.join(save_folder, "reconstruction_test.log") 181 | _log = "{}| stddev {}| mean {}\n" 182 | 183 | plt.hist(chamfer_list, density=False, bins=args["bin"]) 184 | plt.ylabel("Frequency") 185 | plt.xlabel("Chamfer value") 186 | save_path = osp.join(save_folder, "chamfer.eps") 187 | plt.savefig(save_path) 188 | plt.clf() 189 | 190 | plt.hist(swd_list, density=False, bins=args["bin"]) 191 | plt.ylabel("Frequency") 192 | plt.xlabel("SW value") 193 | save_path = osp.join(save_folder, "swd.eps") 194 | plt.savefig(save_path) 195 | plt.clf() 196 | 197 | plt.hist(emd_list, density=False, bins=args["bin"]) 198 | plt.ylabel("Frequency") 199 | plt.xlabel("EMD value") 200 | save_path = osp.join(save_folder, "emd.eps") 201 | plt.savefig(save_path) 202 | plt.clf() 203 | 204 | log = _log.format("chamfer", statistics.stdev(chamfer_list), statistics.mean(chamfer_list)) 205 | with open(test_log, "a") as fp: 206 | fp.write(log) 207 | print(log) 208 | 209 | log = _log.format("swd", statistics.stdev(swd_list), statistics.mean(swd_list)) 210 | with open(test_log, "a") as fp: 211 | fp.write(log) 212 | print(log) 213 | 214 | log = _log.format("emd", statistics.stdev(emd_list), statistics.mean(emd_list)) 215 | with open(test_log, "a") as fp: 216 | fp.write(log) 217 | print(log) 218 | 219 | 220 | if __name__ == "__main__": 221 | main() 222 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ##### Table of Content 2 | 3 | 1. [Introduction](#point-set-distances-for-learning-representations-of-3D-point-clouds) 4 | 1. [Getting Started](#getting-started) 5 | - [Datasets](#datasets) 6 | - [Installation](#installation) 7 | 1. [Experiments](#experiments) 8 | 9 | 10 | # Point-set Distances for Learning Representations of 3D Point Clouds 11 | This repository contains the implementation of our [paper](https://arxiv.org/abs/2102.04014). In particular, we release code for training a point cloud autoencoder network with different point-set distances, such as, sliced Wasserstein distance (SWD), Chamfer, ... and testing the autoencoder for classification, reconstruction, registration, and generation. 12 | 13 | | ![teaser.png](./image/teaser.png) | 14 | |:--:| 15 | | *Morph a sphere into a chair by optimizing two different loss functions: Chamfer (top, red) and SWD (bottom, blue).*| 16 | 17 | Details of the model architecture and experimental results can be found in [our following paper](https://arxiv.org/abs/2102.04014). 18 | 19 | ``` 20 | @InProceedings{Nguyen2021PointSetDistances, 21 | title={Point-set Distances for Learning Representations of 3D Point Clouds}, 22 | author={Nguyen, Trung and Pham, Quang-Hieu and Le, Tam and Pham, Tung and Ho, Nhat and Hua, Binh-Son}, 23 | booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)}, 24 | year={2021} 25 | } 26 | ``` 27 | **Please CITE** our paper whenever our model implementation is used to help produce published results or incorporated into other software. 28 | 29 | ## Getting Started 30 | 31 | ### Datasets 32 | #### ShapeNet Core with 55 categories (refered from FoldingNet.) 33 | ```bash 34 | cd dataset 35 | bash download_shapenet_core55_catagories.sh 36 | ``` 37 | #### ModelNet40 38 | ```bash 39 | cd dataset 40 | bash download_modelnet40_same_with_pointnet.sh 41 | ``` 42 | #### ShapeNet Chair 43 | ```bash 44 | cd dataset 45 | bash download_shapenet_chair.sh 46 | ``` 47 | #### 3DMatch 48 | ```bash 49 | cd dataset 50 | bash download_3dmatch.sh 51 | ``` 52 | ### Installation: 53 | The code is based on Pytorch. It has been tested with Python 3.6.9, PyTorch 1.2.0, CUDA 10.0 on Ubuntu 18.04. 54 | Other dependencies: 55 | * Tensorboard 2.3.0 56 | * Open3d 0.7.0 57 | * Tqdm 4.46.0 58 | 59 | To compile CUDA kernel for CD/EMD loss: 60 | ``` 61 | cd metrics_from_point_flow/pytorch_structural_losses/ 62 | make clean 63 | make 64 | ``` 65 | ## Experiments 66 | ### Autoencoder 67 | 68 | 69 | Available arguments for training an autoencoder: 70 | ``` 71 | train.py [-h] [--config CONFIG] [--logdir LOGDIR] 72 | [--data_path DATA_PATH] [--loss LOSS] 73 | [--autoencoder AUTOENCODER] 74 | 75 | optional arguments: 76 | -h, --help show this help message and exit 77 | --config CONFIG path to json config file 78 | --logdir LOGDIR path to the log directory 79 | --data_path DATA_PATH path to data for training 80 | --loss LOSS loss function. One of [swd, emd, chamfer, asw, msw, gsw] 81 | --autoencoder AUTOENCODER model name. One of [pointnet, pcn] 82 | ``` 83 | Example: 84 | ``` 85 | python train.py --config="config.json" \ 86 | --logdir="logs/" \ 87 | --data_path="dataset/shapenet_core55/shapenet57448xyzonly.npz" \ 88 | --loss="swd" \ 89 | --autoencoder="pointnet" 90 | 91 | # or in short, you can run 92 | bash train.sh 93 | ``` 94 | To test reconstruction: 95 | ``` 96 | python reconstruction/reconstruction_test.py --config="reconstruction/config.json" \ 97 | --logdir="logs/" \ 98 | --data_path="dataset/modelnet40_ply_hdf5_2048/" 99 | 100 | # or in short, you can run 101 | bash reconstruction/test.sh 102 | ``` 103 | ### Semi-supervised classification 104 | 111 | To generate latent codes of the train/test sets of ModelNet40 and save them into files: 112 | ``` 113 | python classification/preprocess_data.py --config='classification/preprocess_train.json' \ 114 | --logdir="logs/" \ 115 | --data_path="dataset/modelnet40_ply_hdf5_2048/train/" 116 | 117 | python classification/preprocess_data.py --config='classification/preprocess_test.json' \ 118 | --logdir="logs/" \ 119 | --data_path="dataset/modelnet40_ply_hdf5_2048/test/" 120 | 121 | # or in short, you can run 122 | bash classification/preprocess.sh 123 | ``` 124 | To get classification results: 125 | ``` 126 | python classification/classification_train.py --config='classification/class_train_config.json' \ 127 | --logdir="logs/" 128 | 129 | python classification/classification_test.py --config='classification/class_test_config.json' \ 130 | --logdir="logs/" 131 | 132 | # or in short, you can run 133 | bash classification/classify_train_test.sh 134 | ``` 135 | ### Registration 136 | To preprocess 3DMatch dataset: 137 | ``` 138 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 139 | --logdir='logs/' \ 140 | --data_path='dataset/home1' 141 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 142 | --logdir='logs/' \ 143 | --data_path='dataset/home2' 144 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 145 | --logdir='logs/' \ 146 | --data_path='dataset/hotel1' 147 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 148 | --logdir='logs/' \ 149 | --data_path='dataset/hotel2' 150 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 151 | --logdir='logs/' \ 152 | --data_path='dataset/hotel3' 153 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 154 | --logdir='logs/' \ 155 | --data_path='dataset/kitchen' 156 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 157 | --logdir='logs/' \ 158 | --data_path='dataset/lab' 159 | python registration/preprocess_data.py --config='registration/preprocess_config.json' \ 160 | --logdir='logs/' \ 161 | --data_path='dataset/study' 162 | 163 | # or in short, you can run 164 | bash registration/preprocess.sh 165 | ``` 166 | To generate transformations into log files: 167 | ``` 168 | python registration/registration_test.py --config='registration/registration_config.json' \ 169 | --logdir='logs/model/home1/' 170 | python registration/registration_test.py --config='registration/registration_config.json' \ 171 | --logdir='logs/model/home2/' 172 | python registration/registration_test.py --config='registration/registration_config.json' \ 173 | --logdir='logs/model/hotel1/' 174 | python registration/registration_test.py --config='registration/registration_config.json' \ 175 | --logdir='logs/model/hotel2/' 176 | python registration/registration_test.py --config='registration/registration_config.json' \ 177 | --logdir='logs/model/hotel3/' 178 | python registration/registration_test.py --config='registration/registration_config.json' \ 179 | --logdir='logs/model/kitchen/' 180 | python registration/registration_test.py --config='registration/registration_config.json' \ 181 | --logdir='logs/model/lab/' 182 | python registration/registration_test.py --config='registration/registration_config.json' \ 183 | --logdir='logs/model/study/' 184 | 185 | # or in short, you can run 186 | bash registration/register.sh 187 | ``` 188 | To evaluate log files, follow the instruction in the `Evaluation` section on this [page](https://3dmatch.cs.princeton.edu/#geometric-registration-benchmark). 189 | 190 | ### Generation 191 | To generate latent codes of train/test sets of ShapeNet Chair and save them into files: 192 | 193 | ``` 194 | python generation/preprocess.py --config='generation/preprocess_train.json' \ 195 | --logdir="logs/" \ 196 | --data_path="dataset/shapenet_chair/train.npz" 197 | 198 | python generation/preprocess.py --config='generation/preprocess_test.json' \ 199 | --logdir="logs/" \ 200 | --data_path="dataset/shapenet_chair/test.npz" 201 | 202 | # or in short, you can run 203 | bash generation/preprocess.sh 204 | ``` 205 | To train the generator: 206 | ``` 207 | python generation/train_latent_generator.py --seed=1 \ 208 | --logdir="logs/" 209 | 210 | # or in short, you can run 211 | bash generation/train_latent_generator.sh 212 | ``` 213 | To test the generator: 214 | ``` 215 | python generation/test_generation.py --config='generation/test_generation_config.json' \ 216 | --logdir="logs/" 217 | 218 | # or in short, you can run 219 | bash generation/test_generation.sh 220 | ``` -------------------------------------------------------------------------------- /metrics_from_point_flow/pytorch_structural_losses/src/approxmatch.cu: -------------------------------------------------------------------------------- 1 | #include "utils.hpp" 2 | 3 | __global__ void approxmatchkernel(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,float * __restrict__ match,float * temp){ 4 | float * remainL=temp+blockIdx.x*(n+m)*2, * remainR=temp+blockIdx.x*(n+m)*2+n,*ratioL=temp+blockIdx.x*(n+m)*2+n+m,*ratioR=temp+blockIdx.x*(n+m)*2+n+m+n; 5 | float multiL,multiR; 6 | if (n>=m){ 7 | multiL=1; 8 | multiR=n/m; 9 | }else{ 10 | multiL=m/n; 11 | multiR=1; 12 | } 13 | const int Block=1024; 14 | __shared__ float buf[Block*4]; 15 | for (int i=blockIdx.x;i=-2;j--){ 24 | for (int j=7;j>-2;j--){ 25 | float level=-powf(4.0f,j); 26 | if (j==-2){ 27 | level=0; 28 | } 29 | for (int k0=0;k0>>(b,n,m,xyz1,xyz2,match,out); 227 | //} 228 | 229 | __global__ void matchcostgrad2kernel(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * __restrict__ grad2){ 230 | __shared__ float sum_grad[256*3]; 231 | for (int i=blockIdx.x;i>>(b,n,m,xyz1,xyz2,match,grad2); 294 | //} 295 | 296 | /*void AddGPUKernel(Dtype *in_a, Dtype *in_b, Dtype *out_c, int N, 297 | cudaStream_t stream)*/ 298 | // temp: TensorShape{b,(n+m)*2} 299 | void approxmatch(int b,int n,int m,const float * xyz1,const float * xyz2,float * match,float * temp, cudaStream_t stream){ 300 | approxmatchkernel 301 | <<<32, 512, 0, stream>>>(b,n,m,xyz1,xyz2,match,temp); 302 | 303 | cudaError_t err = cudaGetLastError(); 304 | if (cudaSuccess != err) 305 | throw std::runtime_error(Formatter() 306 | << "CUDA kernel failed : " << std::to_string(err)); 307 | } 308 | 309 | void matchcost(int b,int n,int m,const float * xyz1,const float * xyz2,float * match, float * out, cudaStream_t stream){ 310 | matchcostkernel<<<32,512,0,stream>>>(b,n,m,xyz1,xyz2,match,out); 311 | 312 | cudaError_t err = cudaGetLastError(); 313 | if (cudaSuccess != err) 314 | throw std::runtime_error(Formatter() 315 | << "CUDA kernel failed : " << std::to_string(err)); 316 | } 317 | 318 | void matchcostgrad(int b,int n,int m,const float * xyz1,const float * xyz2,const float * match,float * grad1,float * grad2, cudaStream_t stream){ 319 | matchcostgrad1kernel<<<32,512,0,stream>>>(b,n,m,xyz1,xyz2,match,grad1); 320 | matchcostgrad2kernel<<>>(b,n,m,xyz1,xyz2,match,grad2); 321 | 322 | cudaError_t err = cudaGetLastError(); 323 | if (cudaSuccess != err) 324 | throw std::runtime_error(Formatter() 325 | << "CUDA kernel failed : " << std::to_string(err)); 326 | } 327 | -------------------------------------------------------------------------------- /metrics_from_point_flow/evaluation_metrics.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | import warnings 4 | 5 | import numpy as np 6 | import torch 7 | from numpy.linalg import norm 8 | from scipy.stats import entropy 9 | from sklearn.neighbors import NearestNeighbors 10 | 11 | 12 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 13 | 14 | # Import CUDA version of approximate EMD, from https://github.com/zekunhao1995/pcgan-pytorch/ 15 | try: 16 | from .StructuralLosses.match_cost import match_cost 17 | from .StructuralLosses.nn_distance import nn_distance 18 | except ImportError: 19 | from StructuralLosses.match_cost import match_cost 20 | from StructuralLosses.nn_distance import nn_distance 21 | 22 | # # Import CUDA version of CD, borrowed from https://github.com/ThibaultGROUEIX/AtlasNet 23 | # try: 24 | # from . chamfer_distance_ext.dist_chamfer import chamferDist 25 | # CD = chamferDist() 26 | # def distChamferCUDA(x,y): 27 | # return CD(x,y,gpu) 28 | # except: 29 | 30 | 31 | def distChamferCUDA(x, y): 32 | return nn_distance(x, y) 33 | 34 | 35 | def emd_approx(sample, ref): 36 | B, N, N_ref = sample.size(0), sample.size(1), ref.size(1) 37 | assert N == N_ref, "Not sure what would EMD do in this case" 38 | emd = match_cost(sample, ref) # (B,) 39 | emd_norm = emd / float(N) # (B,) 40 | return emd_norm 41 | 42 | 43 | # Borrow from https://github.com/ThibaultGROUEIX/AtlasNet 44 | def distChamfer(a, b): 45 | x, y = a, b 46 | bs, num_points, points_dim = x.size() 47 | xx = torch.bmm(x, x.transpose(2, 1)) 48 | yy = torch.bmm(y, y.transpose(2, 1)) 49 | zz = torch.bmm(x, y.transpose(2, 1)) 50 | diag_ind = torch.arange(0, num_points).to(a).long() 51 | rx = xx[:, diag_ind, diag_ind].unsqueeze(1).expand_as(xx) 52 | ry = yy[:, diag_ind, diag_ind].unsqueeze(1).expand_as(yy) 53 | P = rx.transpose(2, 1) + ry - 2 * zz 54 | return P.min(1)[0], P.min(2)[0] 55 | 56 | 57 | def EMD_CD(sample_pcs, ref_pcs, batch_size, accelerated_cd=False, reduced=True): 58 | N_sample = sample_pcs.shape[0] 59 | N_ref = ref_pcs.shape[0] 60 | assert N_sample == N_ref, "REF:%d SMP:%d" % (N_ref, N_sample) 61 | 62 | cd_lst = [] 63 | emd_lst = [] 64 | iterator = range(0, N_sample, batch_size) 65 | 66 | for b_start in iterator: 67 | b_end = min(N_sample, b_start + batch_size) 68 | sample_batch = sample_pcs[b_start:b_end] 69 | ref_batch = ref_pcs[b_start:b_end] 70 | 71 | if accelerated_cd: 72 | dl, dr = distChamferCUDA(sample_batch, ref_batch) 73 | else: 74 | dl, dr = distChamfer(sample_batch, ref_batch) 75 | cd_lst.append(dl.mean(dim=1) + dr.mean(dim=1)) 76 | 77 | emd_batch = emd_approx(sample_batch, ref_batch) 78 | emd_lst.append(emd_batch) 79 | 80 | if reduced: 81 | cd = torch.cat(cd_lst).mean() 82 | emd = torch.cat(emd_lst).mean() 83 | else: 84 | cd = torch.cat(cd_lst) 85 | emd = torch.cat(emd_lst) 86 | 87 | results = { 88 | "MMD-CD": cd, 89 | "MMD-EMD": emd, 90 | } 91 | return results 92 | 93 | 94 | def _pairwise_EMD_CD_(sample_pcs, ref_pcs, batch_size, accelerated_cd=True): 95 | N_sample = sample_pcs.shape[0] 96 | N_ref = ref_pcs.shape[0] 97 | all_cd = [] 98 | all_emd = [] 99 | iterator = range(N_sample) 100 | for sample_b_start in iterator: 101 | sample_batch = sample_pcs[sample_b_start] 102 | 103 | cd_lst = [] 104 | emd_lst = [] 105 | for ref_b_start in range(0, N_ref, batch_size): 106 | ref_b_end = min(N_ref, ref_b_start + batch_size) 107 | ref_batch = ref_pcs[ref_b_start:ref_b_end] 108 | 109 | batch_size_ref = ref_batch.size(0) 110 | sample_batch_exp = sample_batch.view(1, -1, 3).expand(batch_size_ref, -1, -1) 111 | sample_batch_exp = sample_batch_exp.contiguous() 112 | 113 | if accelerated_cd: 114 | dl, dr = distChamferCUDA(sample_batch_exp, ref_batch) 115 | else: 116 | dl, dr = distChamfer(sample_batch_exp, ref_batch) 117 | cd_lst.append((dl.mean(dim=1) + dr.mean(dim=1)).view(1, -1)) 118 | 119 | emd_batch = emd_approx(sample_batch_exp, ref_batch) 120 | emd_lst.append(emd_batch.view(1, -1)) 121 | 122 | cd_lst = torch.cat(cd_lst, dim=1) 123 | emd_lst = torch.cat(emd_lst, dim=1) 124 | all_cd.append(cd_lst) 125 | all_emd.append(emd_lst) 126 | 127 | all_cd = torch.cat(all_cd, dim=0) # N_sample, N_ref 128 | all_emd = torch.cat(all_emd, dim=0) # N_sample, N_ref 129 | 130 | return all_cd, all_emd 131 | 132 | 133 | # Adapted from https://github.com/xuqiantong/GAN-Metrics/blob/master/framework/metric.py 134 | def knn(Mxx, Mxy, Myy, k, sqrt=False): 135 | n0 = Mxx.size(0) 136 | n1 = Myy.size(0) 137 | label = torch.cat((torch.ones(n0), torch.zeros(n1))).to(Mxx) 138 | M = torch.cat((torch.cat((Mxx, Mxy), 1), torch.cat((Mxy.transpose(0, 1), Myy), 1)), 0) 139 | if sqrt: 140 | M = M.abs().sqrt() 141 | INFINITY = float("inf") 142 | val, idx = (M + torch.diag(INFINITY * torch.ones(n0 + n1).to(Mxx))).topk(k, 0, False) 143 | 144 | count = torch.zeros(n0 + n1).to(Mxx) 145 | for i in range(0, k): 146 | count = count + label.index_select(0, idx[i]) 147 | pred = torch.ge(count, (float(k) / 2) * torch.ones(n0 + n1).to(Mxx)).float() 148 | 149 | s = { 150 | "tp": (pred * label).sum(), 151 | "fp": (pred * (1 - label)).sum(), 152 | "fn": ((1 - pred) * label).sum(), 153 | "tn": ((1 - pred) * (1 - label)).sum(), 154 | } 155 | 156 | s.update( 157 | { 158 | "precision": s["tp"] / (s["tp"] + s["fp"] + 1e-10), 159 | "recall": s["tp"] / (s["tp"] + s["fn"] + 1e-10), 160 | "acc_t": s["tp"] / (s["tp"] + s["fn"] + 1e-10), 161 | "acc_f": s["tn"] / (s["tn"] + s["fp"] + 1e-10), 162 | "acc": torch.eq(label, pred).float().mean(), 163 | } 164 | ) 165 | return s 166 | 167 | 168 | def lgan_mmd_cov(all_dist): 169 | N_sample, N_ref = all_dist.size(0), all_dist.size(1) 170 | min_val_fromsmp, min_idx = torch.min(all_dist, dim=1) 171 | min_val, _ = torch.min(all_dist, dim=0) 172 | mmd = min_val.mean() 173 | mmd_smp = min_val_fromsmp.mean() 174 | cov = float(min_idx.unique().view(-1).size(0)) / float(N_ref) 175 | cov = torch.tensor(cov).to(all_dist) 176 | return { 177 | "lgan_mmd": mmd, 178 | "lgan_cov": cov, 179 | "lgan_mmd_smp": mmd_smp, 180 | } 181 | 182 | 183 | def compute_all_metrics(sample_pcs, ref_pcs, batch_size, accelerated_cd=False, **kwargs): 184 | results = {} 185 | 186 | M_rs_cd, M_rs_emd = _pairwise_EMD_CD_(ref_pcs, sample_pcs, batch_size, accelerated_cd=accelerated_cd) 187 | 188 | res_cd = lgan_mmd_cov(M_rs_cd.t()) 189 | results.update({"%s-CD" % k: v for k, v in res_cd.items()}) 190 | 191 | res_emd = lgan_mmd_cov(M_rs_emd.t()) 192 | results.update({"%s-EMD" % k: v for k, v in res_emd.items()}) 193 | 194 | M_rr_cd, M_rr_emd = _pairwise_EMD_CD_(ref_pcs, ref_pcs, batch_size, accelerated_cd=accelerated_cd) 195 | M_ss_cd, M_ss_emd = _pairwise_EMD_CD_(sample_pcs, sample_pcs, batch_size, accelerated_cd=accelerated_cd) 196 | 197 | # 1-NN results 198 | one_nn_cd_res = knn(M_rr_cd, M_rs_cd, M_ss_cd, 1, sqrt=False) 199 | results.update({"1-NN-CD-%s" % k: v for k, v in one_nn_cd_res.items() if "acc" in k}) 200 | one_nn_emd_res = knn(M_rr_emd, M_rs_emd, M_ss_emd, 1, sqrt=False) 201 | results.update({"1-NN-EMD-%s" % k: v for k, v in one_nn_emd_res.items() if "acc" in k}) 202 | 203 | return results 204 | 205 | 206 | ####################################################### 207 | # JSD : from https://github.com/optas/latent_3d_points 208 | ####################################################### 209 | def unit_cube_grid_point_cloud(resolution, clip_sphere=False): 210 | """Returns the center coordinates of each cell of a 3D grid with resolution^3 cells, 211 | that is placed in the unit-cube. 212 | If clip_sphere it True it drops the "corner" cells that lie outside the unit-sphere. 213 | """ 214 | grid = np.ndarray((resolution, resolution, resolution, 3), np.float32) 215 | spacing = 1.0 / float(resolution - 1) 216 | for i in range(resolution): 217 | for j in range(resolution): 218 | for k in range(resolution): 219 | grid[i, j, k, 0] = i * spacing - 0.5 220 | grid[i, j, k, 1] = j * spacing - 0.5 221 | grid[i, j, k, 2] = k * spacing - 0.5 222 | 223 | if clip_sphere: 224 | grid = grid.reshape(-1, 3) 225 | grid = grid[norm(grid, axis=1) <= 0.5] 226 | 227 | return grid, spacing 228 | 229 | 230 | def jsd_between_point_cloud_sets(sample_pcs, ref_pcs, resolution=28): 231 | """Computes the JSD between two sets of point-clouds, as introduced in the paper 232 | ```Learning Representations And Generative Models For 3D Point Clouds```. 233 | Args: 234 | sample_pcs: (np.ndarray S1xR2x3) S1 point-clouds, each of R1 points. 235 | ref_pcs: (np.ndarray S2xR2x3) S2 point-clouds, each of R2 points. 236 | resolution: (int) grid-resolution. Affects granularity of measurements. 237 | """ 238 | in_unit_sphere = True 239 | sample_grid_var = entropy_of_occupancy_grid(sample_pcs, resolution, in_unit_sphere)[1] 240 | ref_grid_var = entropy_of_occupancy_grid(ref_pcs, resolution, in_unit_sphere)[1] 241 | return jensen_shannon_divergence(sample_grid_var, ref_grid_var) 242 | 243 | 244 | def entropy_of_occupancy_grid(pclouds, grid_resolution, in_sphere=False, verbose=False): 245 | """Given a collection of point-clouds, estimate the entropy of the random variables 246 | corresponding to occupancy-grid activation patterns. 247 | Inputs: 248 | pclouds: (numpy array) #point-clouds x points per point-cloud x 3 249 | grid_resolution (int) size of occupancy grid that will be used. 250 | """ 251 | epsilon = 10e-4 252 | bound = 0.5 + epsilon 253 | if abs(np.max(pclouds)) > bound or abs(np.min(pclouds)) > bound: 254 | if verbose: 255 | warnings.warn("Point-clouds are not in unit cube.") 256 | 257 | if in_sphere and np.max(np.sqrt(np.sum(pclouds ** 2, axis=2))) > bound: 258 | if verbose: 259 | warnings.warn("Point-clouds are not in unit sphere.") 260 | 261 | grid_coordinates, _ = unit_cube_grid_point_cloud(grid_resolution, in_sphere) 262 | grid_coordinates = grid_coordinates.reshape(-1, 3) 263 | grid_counters = np.zeros(len(grid_coordinates)) 264 | grid_bernoulli_rvars = np.zeros(len(grid_coordinates)) 265 | nn = NearestNeighbors(n_neighbors=1).fit(grid_coordinates) 266 | 267 | for pc in pclouds: 268 | _, indices = nn.kneighbors(pc) 269 | indices = np.squeeze(indices) 270 | for i in indices: 271 | grid_counters[i] += 1 272 | indices = np.unique(indices) 273 | for i in indices: 274 | grid_bernoulli_rvars[i] += 1 275 | 276 | acc_entropy = 0.0 277 | n = float(len(pclouds)) 278 | for g in grid_bernoulli_rvars: 279 | if g > 0: 280 | p = float(g) / n 281 | acc_entropy += entropy([p, 1.0 - p]) 282 | 283 | return acc_entropy / len(grid_counters), grid_counters 284 | 285 | 286 | def jensen_shannon_divergence(P, Q): 287 | if np.any(P < 0) or np.any(Q < 0): 288 | raise ValueError("Negative values.") 289 | if len(P) != len(Q): 290 | raise ValueError("Non equal size.") 291 | 292 | P_ = P / np.sum(P) # Ensure probabilities. 293 | Q_ = Q / np.sum(Q) 294 | 295 | e1 = entropy(P_, base=2) 296 | e2 = entropy(Q_, base=2) 297 | e_sum = entropy((P_ + Q_) / 2.0, base=2) 298 | res = e_sum - ((e1 + e2) / 2.0) 299 | 300 | res2 = _jsdiv(P_, Q_) 301 | 302 | if not np.allclose(res, res2, atol=10e-5, rtol=0): 303 | warnings.warn("Numerical values of two JSD methods don't agree.") 304 | 305 | return res 306 | 307 | 308 | def _jsdiv(P, Q): 309 | """another way of computing JSD""" 310 | 311 | def _kldiv(A, B): 312 | a = A.copy() 313 | b = B.copy() 314 | idx = np.logical_and(a > 0, b > 0) 315 | a = a[idx] 316 | b = b[idx] 317 | return np.sum([v for v in a * np.log2(a / b)]) 318 | 319 | P_ = P / np.sum(P) 320 | Q_ = Q / np.sum(Q) 321 | 322 | M = 0.5 * (P_ + Q_) 323 | 324 | return 0.5 * (_kldiv(P_, M) + _kldiv(Q_, M)) 325 | 326 | 327 | if __name__ == "__main__": 328 | B, N = 2, 10 329 | x = torch.rand(B, N, 3) 330 | y = torch.rand(B, N, 3) 331 | 332 | min_l, min_r = distChamferCUDA(x.cuda(), y.cuda()) 333 | print(min_l.shape) 334 | print(min_r.shape) 335 | 336 | l_dist = min_l.mean().cpu().detach().item() 337 | r_dist = min_r.mean().cpu().detach().item() 338 | print(l_dist, r_dist) 339 | # print(emd_approx(x.cuda(), y.cuda())) 340 | # print(distChamferCUDA(x.cuda(), y.cuda())) 341 | -------------------------------------------------------------------------------- /generation/train_latent_generator.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | import numpy as np 4 | import torch 5 | 6 | 7 | def seed_worker(worker_id): 8 | worker_seed = torch.initial_seed() % 2 ** 32 9 | np.random.seed(worker_seed) 10 | random.seed(worker_seed) 11 | 12 | 13 | import warnings 14 | 15 | # deepul borrowed from https://github.com/rll/deepul 16 | import deepul.pytorch_util as ptu 17 | 18 | 19 | warnings.filterwarnings("ignore") 20 | 21 | ptu.set_gpu_mode(True) 22 | 23 | import os 24 | import os.path as osp 25 | import sys 26 | 27 | import matplotlib.pyplot as plt 28 | import torch.nn as nn 29 | import torch.optim as optim 30 | import torch.utils.data as data 31 | from torch.autograd import Variable 32 | from tqdm import tqdm 33 | 34 | 35 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 36 | from utils.utils import create_save_folder 37 | 38 | 39 | class L2MatrixComputer(nn.Module): 40 | def __init__(self): 41 | super(L2MatrixComputer, self).__init__() 42 | 43 | @staticmethod 44 | def compute_cost_matrix(data, generated_data, **kwargs): 45 | """ 46 | data, generated_data: shape [batch_size, dim] 47 | """ 48 | x_col = data.unsqueeze(1) 49 | y_lin = generated_data.unsqueeze(0) 50 | c = torch.sum((torch.abs(x_col - y_lin)) ** 2, 2) 51 | if ("scale" in kwargs.keys()) and (torch.max(c).item() > 100): 52 | return kwargs["scale"] * c 53 | else: 54 | return c 55 | 56 | 57 | def sink_stabilized(M, reg, numItermax=1000, tau=1e2, stopThr=1e-9, warmstart=None, print_period=20, cuda=True): 58 | 59 | if cuda: 60 | a = Variable(torch.ones((M.size()[0],)) / M.size()[0]).cuda() 61 | b = Variable(torch.ones((M.size()[1],)) / M.size()[1]).cuda() 62 | else: 63 | a = Variable(torch.ones((M.size()[0],)) / M.size()[0]) 64 | b = Variable(torch.ones((M.size()[1],)) / M.size()[1]) 65 | 66 | # init data 67 | na = len(a) 68 | nb = len(b) 69 | 70 | cpt = 0 71 | # we assume that no distances are null except those of the diagonal of 72 | # distances 73 | if warmstart is None: 74 | if cuda: 75 | alpha, beta = Variable(torch.zeros(na)).cuda(), Variable(torch.zeros(nb)).cuda() 76 | else: 77 | alpha, beta = Variable(torch.zeros(na)), Variable(torch.zeros(nb)) 78 | else: 79 | alpha, beta = warmstart 80 | 81 | if cuda: 82 | u, v = Variable(torch.ones(na) / na).cuda(), Variable(torch.ones(nb) / nb).cuda() 83 | else: 84 | u, v = Variable(torch.ones(na) / na), Variable(torch.ones(nb) / nb) 85 | 86 | def get_K(alpha, beta): 87 | return torch.exp(-(M - alpha.view((na, 1)) - beta.view((1, nb))) / reg) 88 | 89 | def get_Gamma(alpha, beta, u, v): 90 | return torch.exp( 91 | -(M - alpha.view((na, 1)) - beta.view((1, nb))) / reg 92 | + torch.log(u.view((na, 1))) 93 | + torch.log(v.view((1, nb))) 94 | ) 95 | 96 | # print(np.min(K)) 97 | 98 | K = get_K(alpha, beta) 99 | transp = K 100 | loop = 1 101 | cpt = 0 102 | err = 1 103 | while loop: 104 | 105 | uprev = u 106 | vprev = v 107 | 108 | # sinkhorn update 109 | v = torch.div(b, (K.t().matmul(u) + 1e-16)) 110 | u = torch.div(a, (K.matmul(v) + 1e-16)) 111 | 112 | # remove numerical problems and store them in K 113 | if torch.max(torch.abs(u)).item() > tau or torch.max(torch.abs(v)).item() > tau: 114 | alpha, beta = alpha + reg * torch.log(u), beta + reg * torch.log(v) 115 | 116 | if cuda: 117 | u, v = Variable(torch.ones(na) / na).cuda(), Variable(torch.ones(nb) / nb).cuda() 118 | else: 119 | u, v = Variable(torch.ones(na) / na), Variable(torch.ones(nb) / nb) 120 | 121 | K = get_K(alpha, beta) 122 | 123 | if cpt % print_period == 0: 124 | transp = get_Gamma(alpha, beta, u, v) 125 | err = (torch.sum(transp) - b).norm(1).pow(2).item() 126 | 127 | if err <= stopThr: 128 | loop = False 129 | 130 | if cpt >= numItermax: 131 | loop = False 132 | 133 | if torch.any(torch.isnan(u)) or torch.any(torch.isnan(v)): 134 | # we have reached the machine precision 135 | # come back to previous solution and quit loop 136 | print("Warning: numerical errors at iteration", cpt) 137 | u = uprev 138 | v = vprev 139 | break 140 | 141 | cpt += 1 142 | 143 | return torch.sum(get_Gamma(alpha, beta, u, v) * M) 144 | 145 | 146 | def sinkhorn(gt, pr, epsilon): 147 | cost_matrix = L2MatrixComputer.compute_cost_matrix(gt, pr) 148 | return sink_stabilized(cost_matrix, epsilon) 149 | 150 | 151 | def savefig(fname, show_figure=True): 152 | if not osp.exists(osp.dirname(fname)): 153 | os.makedirs(osp.dirname(fname)) 154 | plt.tight_layout() 155 | plt.savefig(fname) 156 | if show_figure: 157 | plt.show() 158 | 159 | 160 | def plot_gan_training(losses, title, fname): 161 | plt.figure() 162 | n_itr = len(losses) 163 | xs = np.arange(n_itr) 164 | 165 | plt.plot(xs, losses, label="loss") 166 | plt.legend() 167 | plt.title(title) 168 | plt.xlabel("Training Iteration") 169 | plt.ylabel("Loss") 170 | savefig(fname) 171 | 172 | 173 | def train( 174 | generator, 175 | critic, 176 | c_loss_fn, 177 | g_loss_fn, 178 | train_loader, 179 | g_optimizer, 180 | c_optimizer, 181 | n_critic=1, 182 | g_scheduler=None, 183 | c_scheduler=None, 184 | weight_clipping=None, 185 | ): 186 | """ 187 | generator: 188 | critic: discriminator in 1ab, general model otherwise 189 | loss_fn 190 | train_loader: instance of DataLoader class 191 | optimizer: 192 | ncritic: how many critic gradient steps to do for every generator step 193 | """ 194 | g_losses, c_losses = [], [] 195 | generator.train() 196 | critic.train() 197 | for i, x in enumerate(train_loader): 198 | x = x.to(ptu.device).float() 199 | c_loss = c_loss_fn(generator, critic, x) 200 | c_optimizer.zero_grad() 201 | c_loss.backward() 202 | c_optimizer.step() 203 | c_losses.append(c_loss.item()) 204 | if weight_clipping is not None: 205 | for param in critic.parameters(): 206 | param.data.clamp_(-weight_clipping, weight_clipping) 207 | 208 | if i % n_critic == 0: # generator step 209 | g_loss = g_loss_fn(generator, critic, x) 210 | g_optimizer.zero_grad() 211 | g_loss.backward() 212 | g_optimizer.step() 213 | g_losses.append(g_loss.item()) 214 | if g_scheduler is not None: 215 | g_scheduler.step() 216 | if c_scheduler is not None: 217 | c_scheduler.step() 218 | return dict(g_losses=g_losses, c_losses=c_losses) 219 | 220 | 221 | def train_epochs(generator, critic, g_loss_fn, c_loss_fn, train_loader, train_args): 222 | epochs, lr = train_args["epochs"], train_args["lr"] 223 | if "optim_cls" in train_args: 224 | g_optimizer = train_args["optim_cls"](generator.parameters(), lr=lr) 225 | c_optimizer = train_args["optim_cls"](critic.parameters(), lr=lr) 226 | else: 227 | g_optimizer = optim.Adam(generator.parameters(), lr=lr, betas=(0, 0.9)) 228 | c_optimizer = optim.Adam(critic.parameters(), lr=lr, betas=(0, 0.9)) 229 | 230 | if train_args.get("lr_schedule", None) is not None: 231 | g_scheduler = optim.lr_scheduler.LambdaLR(g_optimizer, train_args["lr_schedule"]) 232 | c_scheduler = optim.lr_scheduler.LambdaLR(c_optimizer, train_args["lr_schedule"]) 233 | else: 234 | g_scheduler = None 235 | c_scheduler = None 236 | 237 | train_losses = dict() 238 | for epoch in tqdm(range(epochs), desc="Epoch", leave=False): 239 | generator.train() 240 | critic.train() 241 | train_loss = train( 242 | generator, 243 | critic, 244 | c_loss_fn, 245 | g_loss_fn, 246 | train_loader, 247 | g_optimizer, 248 | c_optimizer, 249 | n_critic=train_args.get("n_critic", 0), 250 | g_scheduler=g_scheduler, 251 | c_scheduler=c_scheduler, 252 | weight_clipping=train_args.get("weight_clipping", None), 253 | ) 254 | 255 | for k in train_loss.keys(): 256 | if k not in train_losses: 257 | train_losses[k] = [] 258 | train_losses[k].extend(train_loss[k]) 259 | 260 | return {"train_losses": train_losses, "generator": generator} 261 | 262 | 263 | def get_training_snapshot(generator, critic, n_samples=5000): 264 | generator.eval() 265 | critic.eval() 266 | xs = np.linspace(-1, 1, 1000) 267 | samples = ptu.get_numpy(generator.sample(n_samples)) 268 | critic_output = ptu.get_numpy(critic(ptu.FloatTensor(xs).unsqueeze(1))) 269 | return samples, xs, critic_output 270 | 271 | 272 | class MLP(nn.Module): 273 | def __init__(self, input_size, n_hidden, hidden_size, output_size): 274 | super().__init__() 275 | layers = [] 276 | for _ in range(n_hidden): 277 | layers.append(nn.Linear(input_size, hidden_size)) 278 | layers.append(nn.LeakyReLU(0.2)) 279 | input_size = hidden_size 280 | layers.append(nn.Linear(hidden_size, output_size)) 281 | self.layers = nn.Sequential(*layers) 282 | 283 | def forward(self, x): 284 | return self.layers(x) 285 | 286 | 287 | class MLPGenerator(nn.Module): 288 | def __init__(self, latent_dim, n_hidden, hidden_size, data_dim): 289 | super().__init__() 290 | self.latent_dim = latent_dim 291 | self.mlp = MLP(latent_dim, n_hidden, hidden_size, data_dim) 292 | 293 | def forward(self, z): 294 | return torch.tanh(self.mlp(z)) 295 | 296 | def sample(self, n): 297 | # n is the number of samples to return 298 | z = ptu.normal(ptu.zeros(n, self.latent_dim), ptu.ones(n, self.latent_dim)) 299 | return self.forward(z) 300 | 301 | 302 | class MLPDiscriminator(nn.Module): 303 | def __init__(self, latent_dim, n_hidden, hidden_size, data_dim): 304 | super().__init__() 305 | self.mlp = MLP(latent_dim, n_hidden, hidden_size, data_dim) 306 | 307 | def forward(self, z): 308 | return torch.sigmoid(self.mlp(z)) 309 | 310 | 311 | def q1_b(train_data): # train OT 312 | """ 313 | train_data: An (5763, 256) numpy array: shapenet chair training data 314 | """ 315 | # create data loaders 316 | loader_args = dict(batch_size=64, shuffle=True, worker_init_fn=seed_worker) 317 | train_loader = data.DataLoader(train_data, **loader_args) 318 | 319 | # model 320 | g = MLPGenerator(64, 3, 128, 256).to(ptu.device) 321 | c = MLPDiscriminator(256, 3, 128, 1).to(ptu.device) 322 | 323 | # loss functions 324 | def g_loss_2(generator, critic, x): # have no effect on discriminator 325 | fake_data = generator.sample(x.shape[0]) 326 | loss = sinkhorn(x, fake_data, 0.001) 327 | print("loss: ", loss.item()) 328 | return loss 329 | 330 | def c_loss_2(generator, critic, x): # have no effect on generator 331 | fake_data = torch.empty(x.shape).uniform_(0, 1).cuda() 332 | return -(1 - critic(fake_data)).log().mean() - critic(x).log().mean() 333 | 334 | # train 335 | # g_loss_2, c_loss_2: sinkhorn 1e-3 336 | result_dic = train_epochs(g, c, g_loss_2, c_loss_2, train_loader, dict(epochs=25, lr=1e-4, n_critic=1, q1=True)) 337 | train_losses = result_dic["train_losses"] 338 | generator = result_dic["generator"] 339 | 340 | return {"g_losses": train_losses["g_losses"], "generator": generator} 341 | 342 | 343 | def save_results(part, fn, save_folder): 344 | npz_latent_codes_path = osp.join(save_folder, "latent_codes.npz") # path to npz file 345 | data = np.load(npz_latent_codes_path)["data"] 346 | result_dict = fn(data) 347 | losses = result_dict["g_losses"] 348 | generator = result_dict["generator"] 349 | 350 | generator_save_path = osp.join(save_folder, "generator.pth") 351 | torch.save(generator.state_dict(), generator_save_path) 352 | 353 | # loss plot 354 | plot_gan_training(losses, "Q1{} Losses".format(part), osp.join(save_folder, "q1{}_losses.png".format(part))) 355 | print(">Plot at:", osp.join(save_folder, "q1{}_losses.png".format(part))) 356 | 357 | 358 | if __name__ == "__main__": 359 | import argparse 360 | 361 | parser = argparse.ArgumentParser() 362 | parser.add_argument("--logdir") 363 | parser.add_argument("--seed", type=int, default=1) 364 | args = parser.parse_args() 365 | 366 | # set seed 367 | torch.manual_seed(args.seed) 368 | random.seed(args.seed) 369 | np.random.seed(args.seed) 370 | 371 | logdir = args.logdir 372 | 373 | save_folder = "shapenet_chair/train/" 374 | save_folder = create_save_folder(logdir, save_folder)["save_folder"] 375 | 376 | save_results("b", q1_b, save_folder) 377 | -------------------------------------------------------------------------------- /loss/sw_variants.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch 5 | import torch.nn as nn 6 | from torch.autograd import Variable 7 | 8 | 9 | sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) 10 | 11 | 12 | def minibatch_rand_projections(batchsize, dim, num_projections=1000, **kwargs): 13 | projections = torch.randn((batchsize, num_projections, dim)) 14 | projections = projections / torch.sqrt(torch.sum(projections ** 2, dim=2, keepdim=True)) 15 | return projections 16 | 17 | 18 | def proj_onto_unit_sphere(vectors): 19 | """ 20 | input: vectors: [batchsize, num_projs, dim] 21 | """ 22 | return vectors / torch.sqrt(torch.sum(vectors ** 2, dim=2, keepdim=True)) 23 | 24 | 25 | def _sample_minibatch_orthogonal_projections(batch_size, dim, num_projections): 26 | projections = torch.zeros((batch_size, num_projections, dim)) 27 | projections = torch.stack([torch.nn.init.orthogonal_(projections[i]) for i in range(projections.shape[0])], dim=0) 28 | return projections 29 | 30 | 31 | def compute_practical_moments_sw(x, y, num_projections=30, device="cuda", degree=2.0, **kwargs): 32 | """ 33 | x, y: [batch_size, num_points, dim=3] 34 | num_projections: integer number 35 | """ 36 | dim = x.size(2) 37 | batch_size = x.size(0) 38 | projections = minibatch_rand_projections(batch_size, dim, num_projections).to(device) 39 | # projs.shape: [batchsize, num_projs, dim] 40 | 41 | xproj = x.bmm(projections.transpose(1, 2)).to(device) 42 | 43 | yproj = y.bmm(projections.transpose(1, 2)).to(device) 44 | 45 | _sort = (torch.sort(xproj.transpose(1, 2))[0] - torch.sort(yproj.transpose(1, 2))[0]).to(device) 46 | 47 | _sort_pow_p_get_sum = torch.sum(torch.pow(torch.abs(_sort), degree), dim=2) 48 | 49 | first_moment = _sort_pow_p_get_sum.mean(dim=1) 50 | second_moment = _sort_pow_p_get_sum.pow(2).mean(dim=1) 51 | 52 | return first_moment, second_moment 53 | 54 | 55 | def compute_practical_moments_sw_with_predefined_projections(x, y, projections, device="cuda", degree=2.0, **kwargs): 56 | """ 57 | x, y: [batch size, num points, dim] 58 | projections: [batch size, num projs, dim] 59 | """ 60 | xproj = x.bmm(projections.transpose(1, 2)).to(device) 61 | 62 | yproj = y.bmm(projections.transpose(1, 2)).to(device) 63 | 64 | _sort = (torch.sort(xproj.transpose(1, 2))[0] - torch.sort(yproj.transpose(1, 2))[0]).to(device) 65 | 66 | _sort_pow_p_get_sum = torch.sum(torch.pow(torch.abs(_sort), degree), dim=2) 67 | 68 | first_moment = _sort_pow_p_get_sum.mean(dim=1) 69 | second_moment = _sort_pow_p_get_sum.pow(2).mean(dim=1) 70 | 71 | return first_moment, second_moment 72 | 73 | 74 | def _compute_practical_moments_sw_with_projected_data(xproj, yproj, device="cuda", degree=2.0, **kwargs): 75 | _sort = (torch.sort(xproj.transpose(1, 2))[0] - torch.sort(yproj.transpose(1, 2))[0]).to(device) 76 | 77 | _sort_pow_p_get_sum = torch.sum(torch.pow(torch.abs(_sort), degree), dim=2) 78 | 79 | first_moment = _sort_pow_p_get_sum.mean(dim=1) 80 | second_moment = _sort_pow_p_get_sum.pow(2).mean(dim=1) 81 | 82 | return first_moment, second_moment 83 | 84 | 85 | def _circular(x, theta): 86 | """The circular defining function for generalized Radon transform 87 | Inputs 88 | X: [batch size, num_points, d] - d: dim of 1 point 89 | theta: [batch size, L, d] that parameterizes for L projections 90 | """ 91 | x_s = torch.stack([x for _ in range(theta.shape[1])], dim=2) 92 | theta_s = torch.stack([theta for _ in range(x.shape[1])], dim=1) 93 | z_s = x_s - theta_s 94 | return torch.sqrt(torch.sum(z_s ** 2, dim=3)) 95 | 96 | 97 | def _linear(x, theta): 98 | """ 99 | x: [batch size, num_points, d] - d: dim of 1 point 100 | theta: [batch size, L, d] that parameterizes for L projections 101 | """ 102 | xproj = x.bmm(theta.transpose(1, 2)) 103 | return xproj 104 | 105 | 106 | class SWD(nn.Module): 107 | """ 108 | Estimate SWD with fixed number of projections 109 | """ 110 | 111 | def __init__(self, num_projs, device="cuda", **kwargs): 112 | super().__init__() 113 | self.num_projs = num_projs 114 | self.device = device 115 | 116 | def forward(self, x, y, **kwargs): 117 | """ 118 | x, y have the same shape of [batch_size, num_points_in_point_cloud, dim_of_1_point] 119 | """ 120 | squared_sw_2, _ = compute_practical_moments_sw(x, y, num_projections=self.num_projs, device=self.device) 121 | return {"loss": squared_sw_2.mean(dim=0)} 122 | 123 | 124 | class ASW(nn.Module): 125 | """ 126 | Adaptive sliced wasserstein algorithm for estimating SWD 127 | """ 128 | 129 | def __init__( 130 | self, 131 | init_projs=2, 132 | step_projs=1, 133 | k=2.0, 134 | loop_rate_thresh=0.05, 135 | projs_history="projs_history.txt", 136 | max_slices=500, 137 | **kwargs 138 | ): 139 | super().__init__() 140 | self.init_projs = init_projs 141 | self.step_projs = step_projs 142 | self.k = k 143 | self.loop_rate_thresh = loop_rate_thresh 144 | self.projs_history = projs_history 145 | self.max_slices = max_slices 146 | if "device" in kwargs.keys(): 147 | self.device = kwargs["device"] 148 | else: 149 | self.device = "cuda" 150 | 151 | def forward(self, x, y, **kwargs): 152 | """ 153 | x, y: [batch size, num points in point cloud, 3] 154 | """ 155 | # allow to adjust epsilon 156 | if "epsilon" in kwargs.keys(): 157 | epsilon = kwargs["epsilon"] 158 | else: 159 | raise ValueError("Epsilon not found.") 160 | 161 | n = self.init_projs 162 | max_slices = self.max_slices 163 | step_projs = self.step_projs 164 | 165 | first_moment_sw_p_pow_p, second_moment_sw_p_pow_p = compute_practical_moments_sw( 166 | x, y, num_projections=n, device=self.device, degree=kwargs["degree"] 167 | ) 168 | 169 | loop_conditions = (self.k ** 2 * (second_moment_sw_p_pow_p - first_moment_sw_p_pow_p ** 2)) > ( 170 | (n - 1) * epsilon ** 2 171 | ) # check ASW condition 172 | loop_rate = ( 173 | loop_conditions.sum(dim=0) * 1.0 / loop_conditions.shape[0] 174 | ) # the ratio of point clouds in the batch satifying the ASW condition. 175 | 176 | while (loop_rate > self.loop_rate_thresh) and ((n + step_projs) <= max_slices): 177 | 178 | first_moment_s_sw, second_moment_s_sw = compute_practical_moments_sw( 179 | x, y, num_projections=step_projs, device=self.device, degree=kwargs["degree"] 180 | ) # sample next s projections 181 | 182 | first_moment_sw_p_pow_p = (n * first_moment_sw_p_pow_p + step_projs * first_moment_s_sw) / ( 183 | n + step_projs 184 | ) # update first and second moments 185 | second_moment_sw_p_pow_p = (n * second_moment_sw_p_pow_p + step_projs * second_moment_s_sw) / ( 186 | n + step_projs 187 | ) 188 | n = n + step_projs 189 | loop_conditions = (self.k ** 2 * (second_moment_sw_p_pow_p - first_moment_sw_p_pow_p ** 2)) > ( 190 | (n - 1) * epsilon ** 2 191 | ) 192 | loop_rate = loop_conditions.sum(dim=0) * 1.0 / loop_conditions.shape[0] 193 | 194 | with open(self.projs_history, "a") as fp: # jot down number of sampled projections 195 | fp.write(str(n) + "\n") 196 | return {"loss": first_moment_sw_p_pow_p.mean(dim=0), "num_slices": n} 197 | 198 | 199 | class MaxSW(nn.Module): 200 | """ 201 | Max-SW distance was proposed in paper "Max-Sliced Wasserstein Distance and its use for GANs" - CVPR'19 202 | The way to estimate it was proposed in paper "Generalized Sliced Wasserstein Distance" - NeurIPS'19 203 | """ 204 | 205 | def __init__(self, device="cuda", **kwargs): 206 | super().__init__() 207 | self.device = device 208 | 209 | def forward(self, x, y, *args, **kwargs): 210 | """ 211 | x, y have the same shape of [batch_size, num_points_in_point_cloud, dim_of_1_point] 212 | """ 213 | dim = x.size(2) 214 | projections = Variable( 215 | minibatch_rand_projections(batchsize=x.size(0), dim=dim, num_projections=1).to(self.device), 216 | requires_grad=True, 217 | ) 218 | # projs.shape: [batchsize, num_projs, dim] 219 | 220 | num_iter = kwargs.get("max_sw_num_iters") if "max_sw_num_iters" in kwargs.keys() else 50 221 | lr = kwargs.get("max_sw_lr") if "max_sw_lr" in kwargs else 1e-4 222 | optimizer = torch.optim.Adam([projections], lr=lr) 223 | 224 | for _ in range(num_iter): 225 | # compute loss 226 | xproj = x.bmm(projections.transpose(1, 2)).to(self.device) 227 | 228 | yproj = y.bmm(projections.transpose(1, 2)).to(self.device) 229 | 230 | _sort = (torch.sort(xproj.transpose(1, 2))[0] - torch.sort(yproj.transpose(1, 2))[0]).to(self.device) 231 | 232 | _sort_pow_2_get_sum = torch.sum(torch.pow(_sort, 2), dim=2) 233 | 234 | negative_first_moment = -_sort_pow_2_get_sum.mean(dim=1) 235 | 236 | # perform optimization 237 | optimizer.zero_grad() 238 | negative_first_moment.mean().backward(retain_graph=True) 239 | optimizer.step() 240 | # project onto unit sphere 241 | projections = proj_onto_unit_sphere(projections) 242 | 243 | projections_no_grad = projections.detach() 244 | loss, _ = compute_practical_moments_sw_with_predefined_projections(x, y, projections_no_grad, self.device) 245 | loss = loss.mean(dim=0) 246 | 247 | return {"loss": loss} 248 | 249 | 250 | class OrtSW(nn.Module): 251 | """ 252 | Orthogonal estimation for SWD was proposed in paper "Orthogonal estimation of Wasserstein Distance - AISTATS'19" 253 | """ 254 | 255 | def __init__(self, num_projs, device="cuda", **kwargs): 256 | super().__init__() 257 | self.num_projs = num_projs 258 | self.device = device 259 | 260 | def forward(self, x, y, **kwargs): 261 | """ 262 | x, y have the same shape of [batch_size, num_points_in_point_cloud, dim_of_1_point] 263 | """ 264 | projections = torch.zeros( 265 | (x.shape[0], self.num_projs, x.shape[2]), dtype=x.dtype, layout=x.layout, device=x.device 266 | ) 267 | 268 | projections = torch.stack( 269 | [torch.nn.init.orthogonal_(projections[i]) for i in range(projections.shape[0])], dim=0 270 | ) 271 | 272 | loss, _ = compute_practical_moments_sw_with_predefined_projections(x, y, projections, device=self.device) 273 | 274 | return {"loss": loss.mean(dim=0)} 275 | 276 | 277 | class GenSW(nn.Module): 278 | """ 279 | Generalized SW distance was proposed in paper "Generalized Sliced Wasserstein Distance" - NeurIPS'19 280 | """ 281 | 282 | def __init__(self, num_projs, g_type="circular", device="cuda", **kwargs): 283 | super().__init__() 284 | self.num_projs = num_projs 285 | self.device = device 286 | self.g_type = g_type 287 | 288 | def forward(self, x, y, **kwargs): 289 | """ 290 | x, y have the same shape of [batch_size, num_points_in_point_cloud, dim_of_1_point] 291 | """ 292 | dim = x.size(2) 293 | batch_size = x.size(0) 294 | projections = minibatch_rand_projections(batch_size, dim, self.num_projs).to(self.device) 295 | 296 | if self.g_type == "circular": 297 | xproj = _circular(x, projections) 298 | yproj = _circular(y, projections) 299 | elif self.g_type == "linear": 300 | xproj = _linear(x, projections) 301 | yproj = _linear(y, projections) 302 | else: 303 | raise NotImplementedError 304 | 305 | loss, _ = _compute_practical_moments_sw_with_projected_data(xproj, yproj, self.device, kwargs["degree"]) 306 | 307 | return {"loss": loss.mean(dim=0)} 308 | 309 | 310 | class PW(nn.Module): 311 | """ 312 | Projected Wasserstein distance was proposed in paper "Orthogonal estimation of Wasserstein Distance - AISTATS'19" 313 | """ 314 | 315 | def __init__(self, num_projs, device="cuda", orthogonal=False, **kwargs): 316 | super().__init__() 317 | self.num_projs = num_projs 318 | self.device = device 319 | self.orthogonal = orthogonal 320 | 321 | def forward(self, x, y, **kwargs): 322 | """ 323 | x, y have the same shape of [batch_size, num_points_in_point_cloud, dim_of_1_point] 324 | """ 325 | 326 | dim = x.size(2) 327 | batch_size = x.size(0) 328 | if self.orthogonal: 329 | projections = _sample_minibatch_orthogonal_projections(batch_size, dim, self.num_projs).to(self.device) 330 | else: 331 | projections = minibatch_rand_projections(batch_size, dim, self.num_projs).to(self.device) 332 | # print(projections) 333 | xproj = _linear(x, projections).transpose(1, 2) # [bs, num_slices, num_points] 334 | yproj = _linear(y, projections).transpose(1, 2) # [bs, num_slices, num_points] 335 | 336 | xproj_argsort = torch.argsort(xproj, dim=2) 337 | yproj_argsort = torch.argsort(yproj, dim=2) 338 | 339 | _sorted_x = torch.stack([x[i][xproj_argsort[i]] for i in range(x.shape[0])], dim=0) 340 | _sorted_y = torch.stack([y[i][yproj_argsort[i]] for i in range(y.shape[0])], dim=0) 341 | 342 | loss = torch.mean((_sorted_x - _sorted_y) ** 2) 343 | return {"loss": loss} 344 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os 4 | import os.path as osp 5 | import random 6 | import shutil 7 | import time 8 | 9 | import numpy as np 10 | import torch 11 | from add_noise_to_data.random_noise import RandomNoiseAdder 12 | from dataset import ShapeNetCore55XyzOnlyDataset 13 | from evaluator import Evaluator 14 | from logger import Logger 15 | from loss import ASW, EMD, SWD, Chamfer, GenSW, MaxSW 16 | from models import PointCapsNet, PointNetAE 17 | from models.utils import init_weights 18 | from saver import Saver 19 | from torch.optim import SGD, Adam 20 | from torch.optim.lr_scheduler import CyclicLR 21 | from torch.utils.data import DataLoader 22 | from tqdm import tqdm 23 | from trainer import AETrainer as Trainer 24 | from utils import get_lr 25 | 26 | 27 | torch.backends.cudnn.enabled = False 28 | 29 | 30 | def seed_worker(worker_id): 31 | worker_seed = torch.initial_seed() % 2 ** 32 32 | np.random.seed(worker_seed) 33 | random.seed(worker_seed) 34 | 35 | 36 | def main(): 37 | parser = argparse.ArgumentParser() 38 | parser.add_argument("--config", help="path to json config file") 39 | parser.add_argument("--logdir", help="path to the log directory") 40 | parser.add_argument("--data_path", help="path to data") 41 | parser.add_argument("--loss", default="swd", help="[swd, emd, chamfer, asw, msw, gsw]") 42 | parser.add_argument("--autoencoder", default="pointnet", help="[pointnet, pcn]") 43 | args = parser.parse_args() 44 | config = args.config 45 | logdir = args.logdir 46 | data_path = args.data_path 47 | loss_type = args.loss 48 | ae_type = args.autoencoder 49 | print("Save checkpoints and logs in: ", logdir) 50 | args = json.load(open(config)) 51 | args["autoencoder"] = ae_type 52 | args["loss"] = loss_type 53 | 54 | # set seed 55 | torch.manual_seed(args["seed"]) 56 | random.seed(args["seed"]) 57 | np.random.seed(args["seed"]) 58 | 59 | if not os.path.exists(logdir): 60 | os.makedirs(logdir) 61 | print(">Logdir was created successfully at: ", logdir) 62 | else: 63 | print(">Folder {} is existing.".format(logdir)) 64 | print(">Do you want to remove it?") 65 | answer = None 66 | while answer not in ("yes", "no"): 67 | answer = input("Enter 'yes' or 'no': ") 68 | if answer == "yes": 69 | shutil.rmtree(logdir) 70 | os.makedirs(logdir) 71 | elif answer == "no": 72 | print("SOME FILES WILL BE OVERWRITTEN OR APPENDED.") 73 | print("If you do not want this, please stop during next 30s.") 74 | time.sleep(30) 75 | else: 76 | print("Please enter 'yes' or 'no'.") 77 | fname = os.path.join(logdir, "train_ae_config.json") 78 | with open(fname, "w") as fp: 79 | json.dump(args, fp, indent=4) 80 | 81 | # print hyperparameters 82 | print(">You have 5s to check the hyperparameters below.") 83 | print(args) 84 | time.sleep(5) 85 | 86 | # init dic of extra parameters for trainer.train 87 | dic = {} 88 | 89 | # init dic of extra parameters for evaluator.evaluate 90 | eval_dic = {} 91 | 92 | # device 93 | device = torch.device(args["device"]) 94 | 95 | # NoiseAdder 96 | if args["add_noise"]: 97 | if args["noise_adder"] == "random": 98 | noise_adder = RandomNoiseAdder(mean=args["mean_noiseadder"], std=args["std_noiseadder"]) 99 | else: 100 | raise ValueError("Unknown noise_adder type.") 101 | 102 | # autoencoder architecture 103 | if args["autoencoder"] == "pointnet": 104 | autoencoder = PointNetAE( 105 | args["embedding_size"], 106 | args["input_channels"], 107 | args["input_channels"], 108 | args["num_points"], 109 | args["normalize"], 110 | ).to(device) 111 | 112 | elif args["autoencoder"] == "pcn": 113 | autoencoder = PointCapsNet( 114 | args["prim_caps_size"], 115 | args["prim_vec_size"], 116 | args["latent_caps_size"], 117 | args["latent_vec_size"], 118 | args["num_points"], 119 | ).to(device) 120 | 121 | else: 122 | raise Exception("Unknown autoencoder.") 123 | 124 | # loss function 125 | if args["loss"] == "chamfer": 126 | loss_func = Chamfer(args["version"]) 127 | 128 | elif args["loss"] == "emd": 129 | loss_func = EMD() 130 | 131 | elif args["loss"] == "swd": 132 | loss_func = SWD(args["num_projs"], device) 133 | 134 | elif args["loss"] == "asw": 135 | sample_projs_history = os.path.join(logdir, "sample_projs_history.txt") 136 | loss_func = ASW( 137 | args["init_projs"], 138 | args["step_projs"], 139 | loop_rate_thresh=args["loop_rate_thresh"], 140 | projs_history=sample_projs_history, 141 | max_slices=args["max_slices"], 142 | ) 143 | dic = {"epsilon": args["init_epsilon"]} 144 | dic["degree"] = args["degree"] 145 | 146 | elif args["loss"] == "msw": 147 | loss_func = MaxSW(device, max_sw_num_iters=args["max_sw_num_iters"], max_sw_lr=args["max_sw_lr"]) 148 | 149 | elif args["loss"] == "gsw": 150 | loss_func = GenSW(num_projs=args["num_projs"], g_type=args["g_type"], device=device, degree=args["degree"]) 151 | 152 | else: 153 | raise Exception("Unknown loss function.") 154 | 155 | # dataset 156 | if args["train_set"] == "shapenetcore55": 157 | dataset = ShapeNetCore55XyzOnlyDataset(data_path, num_points=args["num_points"], phase="train") 158 | 159 | else: 160 | raise Exception("Unknown dataset") 161 | 162 | # optimizer 163 | if args["optimizer"] == "sgd": 164 | optimizer = SGD( 165 | autoencoder.parameters(), 166 | lr=args["learning_rate"], 167 | momentum=args["momentum"], 168 | weight_decay=args["weight_decay"], 169 | ) 170 | 171 | elif args["optimizer"] == "adam": 172 | optimizer = Adam( 173 | autoencoder.parameters(), 174 | lr=args["learning_rate"], 175 | betas=(0.5, 0.999), 176 | weight_decay=args["weight_decay"], 177 | ) 178 | 179 | else: 180 | raise Exception("Optimizer has had implementation yet.") 181 | 182 | # init weights 183 | if osp.isfile(osp.join(logdir, args["checkpoint"])): 184 | print(">Init weights with {}".format(args["checkpoint"])) 185 | checkpoint = torch.load(osp.join(logdir, args["checkpoint"])) 186 | if "autoencoder" in checkpoint.keys(): 187 | autoencoder.load_state_dict(checkpoint["autoencoder"]) 188 | else: 189 | autoencoder.load_state_dict(checkpoint) 190 | if "optimizer" in checkpoint.keys(): 191 | try: 192 | optimizer.load_state_dict(checkpoint["optimizer"]) 193 | except: 194 | print(">Found no state dict for optimizer.") 195 | 196 | elif osp.isfile(args["checkpoint"]): 197 | print(">Init weights with {}".format(args["checkpoint"])) 198 | checkpoint = torch.load(osp.join(args["checkpoint"])) 199 | if "autoencoder" in checkpoint.keys(): 200 | autoencoder.load_state_dict(checkpoint["autoencoder"]) 201 | else: 202 | autoencoder.load_state_dict(checkpoint) 203 | 204 | else: 205 | print(">Init weights with Xavier") 206 | autoencoder.apply(init_weights) 207 | 208 | # dataloader 209 | train_loader = DataLoader( 210 | dataset, 211 | batch_size=args["batch_size"], 212 | num_workers=args["num_workers"], 213 | pin_memory=True, 214 | shuffle=True, 215 | worker_init_fn=seed_worker, 216 | ) 217 | 218 | # logger 219 | tensorboard_dir = osp.join(logdir, "tensorboard") 220 | if not osp.exists(tensorboard_dir): 221 | os.makedirs(tensorboard_dir) 222 | tensorboard_logger = Logger(tensorboard_dir) 223 | 224 | # scheduler 225 | if args["use_scheduler"]: 226 | if args["scheduler"] == "cyclic_lr": 227 | scheduler = CyclicLR(optimizer, base_lr=args["base_lr"], max_lr=args["max_lr"]) 228 | else: 229 | raise Exception("Unknown learning rate scheduler.") 230 | 231 | # evaluator 232 | if args["evaluator"] == "based_on_train_loss": 233 | args["eval_criteria"] = "loss_func" 234 | args["have_val_set"] = False 235 | 236 | elif args["evaluator"] == "based_on_val_loss": 237 | args["eval_criteria"] = "loss_func" 238 | args["have_val_set"] = True 239 | 240 | else: 241 | raise ValueError("Unknown evaluator.") 242 | 243 | # val_set and val_loader 244 | if args["have_val_set"]: 245 | if args["val_set"] == "shapenetcore55": 246 | val_set = ShapeNetCore55XyzOnlyDataset(args["val_root"], num_points=args["num_points"], phase="test") 247 | 248 | else: 249 | raise Exception("Unknown dataset") 250 | 251 | val_loader = DataLoader( 252 | val_set, 253 | batch_size=args["val_batch_size"], 254 | num_workers=args["num_workers"], 255 | pin_memory=True, 256 | shuffle=False, 257 | worker_init_fn=seed_worker, 258 | ) 259 | 260 | # avg_eval_value for model selection 261 | # init avg_eval_value 262 | avg_eval_value = args["best_eval_value"] 263 | best_eval_value = float(args["best_eval_value"]) 264 | best_epoch = int(args["best_epoch"]) 265 | 266 | avg_train_loss = args["best_train_loss"] 267 | best_train_loss = float(args["best_train_loss"]) 268 | best_epoch_based_on_train_loss = int(args["best_epoch_based_on_train_loss"]) 269 | 270 | print("best eval value: ", best_eval_value) 271 | print("best epoch: ", best_epoch) 272 | 273 | # train 274 | start_epoch = args["start_epoch"] 275 | num_epochs = args["num_epochs"] 276 | 277 | model_path = os.path.join(logdir, "model.pth") 278 | best_train_loss_model_path = os.path.join(logdir, "best_train_loss_model.pth") 279 | 280 | rec_train_log_path = os.path.join(logdir, "rec_train.log") 281 | reg_train_log_path = os.path.join(logdir, "reg_train.log") 282 | 283 | train_log_path = os.path.join(logdir, "train.log") 284 | eval_log_path = os.path.join(logdir, "eval_when_train.log") 285 | 286 | best_eval_log_path = os.path.join(logdir, "best_eval_when_train.log") 287 | best_train_log_path = os.path.join(logdir, "best_train.log") 288 | 289 | start_time = time.time() 290 | 291 | dic["iter_id"] = 0 292 | prev_losses_list = [] 293 | 294 | for epoch in tqdm(range(start_epoch, num_epochs)): 295 | # Below optimizer setup as original code of 3D Point Capsule Net https://github.com/yongheng1991/3D-point-capsule-networks/blob/master/apps/AE/train_ae.py 296 | if args["autoencoder"] == "pcn": 297 | if epoch < 20: 298 | optimizer = Adam(autoencoder.parameters(), lr=0.001) 299 | elif epoch < 50: 300 | optimizer = Adam(autoencoder.parameters(), lr=0.0001) 301 | else: 302 | optimizer = Adam(autoencoder.parameters(), lr=0.00001) 303 | 304 | train_loss_list = [] 305 | rec_train_loss_list = [] 306 | reg_train_loss_list = [] 307 | 308 | for batch_id, batch in tqdm(enumerate(train_loader)): 309 | dic["iter_id"] += 1 310 | 311 | data = batch.to(device) 312 | 313 | if args["add_noise"]: 314 | if args["train_denoise"]: 315 | dic["input"] = data.detach().clone() 316 | data = noise_adder.add_noise(data) 317 | 318 | # train_on_batch 319 | result_dic = Trainer.train(autoencoder, loss_func, optimizer, data, **dic) 320 | autoencoder = result_dic["ae"] 321 | optimizer = result_dic["optimizer"] 322 | train_loss = result_dic["loss"] 323 | 324 | # 2 types of losses 325 | if "rec_loss" in result_dic.keys(): 326 | rec_train_loss_list.append(result_dic["rec_loss"].item()) 327 | if "reg_loss" in result_dic.keys(): 328 | reg_train_loss_list.append(result_dic["reg_loss"].item()) 329 | 330 | # append to loss lists 331 | train_loss_list.append(train_loss.item()) 332 | 333 | # update epsilon for adaptive sw 334 | if "epsilon" in dic.keys(): 335 | if not args["fix_epsilon"]: 336 | # updata prev_losses_list 337 | assert ("num_prev_losses" in args.keys()) and (args["num_prev_losses"] > 0) 338 | if len(prev_losses_list) == args["num_prev_losses"]: 339 | prev_losses_list.pop(0) # pop the first item 340 | prev_losses_list.append(train_loss.item()) # add item to the last 341 | dic["epsilon"] = min(prev_losses_list) * args["next_epsilon_ratio_rec"] 342 | 343 | if "rec" in dic.keys() and "epsilon" in dic["rec"].keys(): 344 | dic["rec"]["epsilon"] = result_dic["rec_loss"].item() * args["next_epsilon_ratio_rec"] 345 | if "reg" in dic.keys() and "epsilon" in dic["reg"].keys(): 346 | dic["reg"]["epsilon"] = result_dic["reg_loss"].item() * args["next_epsilon_ratio_reg"] 347 | 348 | # adjust scheduler 349 | if args["use_scheduler"]: 350 | scheduler.step() 351 | 352 | # write tensorboard log 353 | info = {"train_loss": train_loss.item(), "learning rate": get_lr(optimizer)} 354 | if "rec_loss" in result_dic.keys(): 355 | info["rec_train_loss"] = rec_train_loss_list[-1] 356 | if "reg_loss" in result_dic.keys(): 357 | info["reg_train_loss"] = reg_train_loss_list[-1] 358 | if "num_slices" in result_dic.keys(): 359 | info["num_slices"] = result_dic["num_slices"] 360 | for tag, value in info.items(): 361 | tensorboard_logger.scalar_summary(tag, value, len(train_loader) * epoch + batch_id + 1) 362 | 363 | # empty cache 364 | if ("empty_cache_batch" in args.keys()) and args["empty_cache_batch"]: 365 | torch.cuda.empty_cache() 366 | # end for 1 epoch 367 | 368 | # calculate avg_train_loss of the epoch 369 | if len(rec_train_loss_list) > 0: 370 | avg_rec_train_loss = sum(rec_train_loss_list) / len(rec_train_loss_list) 371 | if len(reg_train_loss_list) > 0: 372 | avg_reg_train_loss = sum(reg_train_loss_list) / len(reg_train_loss_list) 373 | avg_train_loss = sum(train_loss_list) / len(train_loss_list) 374 | 375 | # evaluate on validation set 376 | if args["have_val_set"] and (epoch % args["epoch_gap_for_evaluation"] == 0): 377 | eval_value_list = [] 378 | with torch.no_grad(): 379 | for batch_id, batch in tqdm(enumerate(val_loader)): 380 | val_data = batch.to(device) 381 | result_dic = Evaluator.evaluate(autoencoder, val_data, loss_func, **eval_dic) 382 | eval_value_list.append(result_dic["evaluation"].item()) 383 | # end for 384 | avg_eval_value = sum(eval_value_list) / len(eval_value_list) 385 | 386 | if not args["have_val_set"]: 387 | avg_eval_value = avg_train_loss 388 | 389 | # save checkpoint 390 | checkpoint_path = osp.join(logdir, "latest.pth") 391 | if args["use_scheduler"]: 392 | Saver.save_checkpoint(autoencoder, optimizer, checkpoint_path, scheduler=scheduler) 393 | else: 394 | Saver.save_checkpoint(autoencoder, optimizer, checkpoint_path) 395 | 396 | if epoch % args["epoch_gap_for_save"] == 0: 397 | checkpoint_path = os.path.join(logdir, "epoch_" + str(epoch) + ".pth") 398 | Saver.save_best_weights(autoencoder, checkpoint_path) 399 | 400 | # save best model based on avg_eval_value 401 | if args["eval_criteria"] in ["jsd", "loss_func", "mmd"]: 402 | better = avg_eval_value < best_eval_value 403 | elif args["eval_criteria"] in ["cov"]: 404 | better = avg_eval_value > best_eval_value 405 | else: 406 | raise Exception("Unknown eval_criteria") 407 | if better: 408 | best_eval_value = avg_eval_value 409 | best_epoch = epoch 410 | Saver.save_best_weights(autoencoder, model_path) 411 | 412 | # save best model based on avg_train_loss 413 | if avg_train_loss < best_train_loss: 414 | best_train_loss = avg_train_loss 415 | best_epoch_based_on_train_loss = epoch 416 | if args["evaluator"] != "based_on_train_loss": 417 | Saver.save_best_weights(autoencoder, best_train_loss_model_path) 418 | 419 | # report 420 | train_log = "Epoch {}| train_loss : {}\n".format(epoch, avg_train_loss) 421 | eval_log = "Epoch {}| eval_value : {}\n".format(epoch, avg_eval_value) 422 | eval_best_log = "Best epoch {}| best eval value: {}\n".format(best_epoch, best_eval_value) 423 | best_train_loss_log = "Best_train_loss epoch {}| best train loss : {}\n".format( 424 | best_epoch_based_on_train_loss, best_train_loss 425 | ) 426 | with open(train_log_path, "a") as fp: 427 | fp.write(train_log) 428 | with open(eval_log_path, "a") as fp: 429 | fp.write(eval_log) 430 | with open(best_eval_log_path, "w") as fp: 431 | fp.write(eval_best_log) 432 | with open(best_train_log_path, "w") as fp: 433 | fp.write(best_train_loss_log) 434 | print(train_log) 435 | print(eval_log) 436 | print(eval_best_log) 437 | print(best_train_loss_log) 438 | 439 | if len(rec_train_loss_list) > 0: 440 | rec_train_log = "Epoch {}| rec_train_loss : {}\n".format(epoch, avg_rec_train_loss) 441 | with open(rec_train_log_path, "a") as fp: 442 | fp.write(rec_train_log) 443 | print(rec_train_log) 444 | 445 | if len(reg_train_loss_list) > 0: 446 | reg_train_log = "Epoch {}| reg_train_loss : {}\n".format(epoch, avg_reg_train_loss) 447 | with open(reg_train_log_path, "a") as fp: 448 | fp.write(reg_train_log) 449 | print(reg_train_log) 450 | 451 | if ("empty_cache_epoch" in args.keys()) and args["empty_cache_epoch"]: 452 | torch.cuda.empty_cache() 453 | print("---------------------------------------------------------------------------------------") 454 | # end for 455 | 456 | finish_time = time.time() 457 | total_runtime = finish_time - start_time 458 | total_runtime = time.strftime("%H:%M:%S", time.gmtime(total_runtime)) 459 | runtime_log = "total runtime (hour:min:sec): {}".format(total_runtime) 460 | print("total_runtime:", total_runtime) 461 | with open(train_log_path, "a") as fp: 462 | fp.write(runtime_log) 463 | 464 | print("Saved checkpoints and logs in: ", logdir) 465 | 466 | 467 | if __name__ == "__main__": 468 | main() 469 | --------------------------------------------------------------------------------