diff --git a/distance/chamfer/__pycache__/dist_chamfer.cpython-36.pyc b/distance/chamfer/__pycache__/dist_chamfer.cpython-36.pyc deleted file mode 100644 index 2577e52..0000000 Binary files a/distance/chamfer/__pycache__/dist_chamfer.cpython-36.pyc and /dev/null differ diff --git a/distance/chamfer/chamfer.cu b/distance/chamfer/chamfer.cu deleted file mode 100644 index d5b886d..0000000 --- a/distance/chamfer/chamfer.cu +++ /dev/null @@ -1,196 +0,0 @@ - -#include -#include - -#include -#include - -#include - - - -__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ - const int batch=512; - __shared__ float buf[batch*3]; - for (int i=blockIdx.x;ibest){ - result[(i*n+j)]=best; - result_i[(i*n+j)]=best_i; - } - } - __syncthreads(); - } - } -} -// int chamfer_cuda_forward(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){ -int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){ - - const auto batch_size = xyz1.size(0); - const auto n = xyz1.size(1); //num_points point cloud A - const auto m = xyz2.size(1); //num_points point cloud B - - NmDistanceKernel<<>>(batch_size, n, xyz1.data(), m, xyz2.data(), dist1.data(), idx1.data()); - NmDistanceKernel<<>>(batch_size, m, xyz2.data(), n, xyz1.data(), dist2.data(), idx2.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err)); - //THError("aborting"); - return 0; - } - return 1; - - -} -__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){ - for (int i=blockIdx.x;i>>(batch_size,n,xyz1.data(),m,xyz2.data(),graddist1.data(),idx1.data(),gradxyz1.data(),gradxyz2.data()); - NmDistanceGradKernel<<>>(batch_size,m,xyz2.data(),n,xyz1.data(),graddist2.data(),idx2.data(),gradxyz2.data(),gradxyz1.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); - //THError("aborting"); - return 0; - } - return 1; - -} - diff --git a/distance/chamfer/chamfer.egg-info/PKG-INFO b/distance/chamfer/chamfer.egg-info/PKG-INFO deleted file mode 100644 index 91901d7..0000000 --- a/distance/chamfer/chamfer.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: chamfer -Version: 0.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/distance/chamfer/chamfer.egg-info/SOURCES.txt b/distance/chamfer/chamfer.egg-info/SOURCES.txt deleted file mode 100644 index 0e582cb..0000000 --- a/distance/chamfer/chamfer.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -chamfer.cu -chamfer_cuda.cpp -setup.py -chamfer.egg-info/PKG-INFO -chamfer.egg-info/SOURCES.txt -chamfer.egg-info/dependency_links.txt -chamfer.egg-info/top_level.txt \ No newline at end of file diff --git a/distance/chamfer/chamfer.egg-info/dependency_links.txt b/distance/chamfer/chamfer.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/distance/chamfer/chamfer.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/distance/chamfer/chamfer.egg-info/top_level.txt b/distance/chamfer/chamfer.egg-info/top_level.txt deleted file mode 100644 index 27aeea6..0000000 --- a/distance/chamfer/chamfer.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -chamfer diff --git a/distance/chamfer/chamfer_cuda.cpp b/distance/chamfer/chamfer_cuda.cpp deleted file mode 100644 index 67574e2..0000000 --- a/distance/chamfer/chamfer_cuda.cpp +++ /dev/null @@ -1,33 +0,0 @@ -#include -#include - -///TMP -//#include "common.h" -/// NOT TMP - - -int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2); - - -int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2); - - - - -int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) { - return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2); -} - - -int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, - at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) { - - return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2); -} - - - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &chamfer_forward, "chamfer forward (CUDA)"); - m.def("backward", &chamfer_backward, "chamfer backward (CUDA)"); -} \ No newline at end of file diff --git a/distance/chamfer/dist/chamfer-0.0.0-py2.7-linux-x86_64.egg b/distance/chamfer/dist/chamfer-0.0.0-py2.7-linux-x86_64.egg deleted file mode 100644 index f19a5fb..0000000 Binary files a/distance/chamfer/dist/chamfer-0.0.0-py2.7-linux-x86_64.egg and /dev/null differ diff --git a/distance/chamfer/dist/chamfer-0.0.0-py3.6-linux-x86_64.egg b/distance/chamfer/dist/chamfer-0.0.0-py3.6-linux-x86_64.egg deleted file mode 100644 index 84e5d13..0000000 Binary files a/distance/chamfer/dist/chamfer-0.0.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/distance/chamfer/dist_chamfer.py b/distance/chamfer/dist_chamfer.py deleted file mode 100644 index 3ffc8fc..0000000 --- a/distance/chamfer/dist_chamfer.py +++ /dev/null @@ -1,54 +0,0 @@ -import math -from torch import nn -from torch.autograd import Function -import torch -import sys -from numbers import Number -from collections import Set, Mapping, deque -import chamfer - - -# Chamfer's distance module @thibaultgroueix -# GPU tensors only -class chamferFunction(Function): - @staticmethod - def forward(ctx, xyz1, xyz2): - batchsize, n, _ = xyz1.size() - _, m, _ = xyz2.size() - - dist1 = torch.zeros(batchsize, n) - dist2 = torch.zeros(batchsize, m) - - idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) - idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) - - dist1 = dist1.cuda() - dist2 = dist2.cuda() - idx1 = idx1.cuda() - idx2 = idx2.cuda() - - chamfer.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) - ctx.save_for_backward(xyz1, xyz2, idx1, idx2) - return dist1, dist2, idx1, idx2 - - @staticmethod - def backward(ctx, graddist1, graddist2, useless1, useless2): - xyz1, xyz2, idx1, idx2 = ctx.saved_tensors - graddist1 = graddist1.contiguous() - graddist2 = graddist2.contiguous() - - gradxyz1 = torch.zeros(xyz1.size()) - gradxyz2 = torch.zeros(xyz2.size()) - - gradxyz1 = gradxyz1.cuda() - gradxyz2 = gradxyz2.cuda() - chamfer.backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2) - return gradxyz1, gradxyz2 - -class chamferDist(nn.Module): - def __init__(self): - super(chamferDist, self).__init__() - - def forward(self, input1, input2): - return chamferFunction.apply(input1, input2) - diff --git a/distance/chamfer/setup.py b/distance/chamfer/setup.py deleted file mode 100644 index 9055958..0000000 --- a/distance/chamfer/setup.py +++ /dev/null @@ -1,14 +0,0 @@ -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup( - name='chamfer', - ext_modules=[ - CUDAExtension('chamfer', [ - 'chamfer_cuda.cpp', - 'chamfer.cu', - ]), - ], - cmdclass={ - 'build_ext': BuildExtension - }) \ No newline at end of file diff --git a/distance/chamfer_multidim/.gitignore b/distance/chamfer_multidim/.gitignore deleted file mode 100644 index abd32e8..0000000 --- a/distance/chamfer_multidim/.gitignore +++ /dev/null @@ -1 +0,0 @@ -*__pycache__* \ No newline at end of file diff --git a/distance/chamfer_multidim/LICENSE b/distance/chamfer_multidim/LICENSE deleted file mode 100644 index 794e2df..0000000 --- a/distance/chamfer_multidim/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2019 ThibaultGROUEIX - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/distance/chamfer_multidim/README.md b/distance/chamfer_multidim/README.md deleted file mode 100644 index 9b4df4f..0000000 --- a/distance/chamfer_multidim/README.md +++ /dev/null @@ -1,101 +0,0 @@ -# Pytorch Chamfer Distance. - -Include a **CUDA** version, and a **PYTHON** version with pytorch standard operations. -NB : In this depo, dist1 and dist2 are squared pointcloud euclidean distances, so you should adapt thresholds accordingly. - -- [x] F - Score - - - -### CUDA VERSION - -- [x] JIT compilation -- [x] Supports multi-gpu -- [x] 2D point clouds. -- [x] 3D point clouds. -- [x] 5D point clouds. -- [x] Contiguous() safe. - - - -### Python Version - -- [x] Supports any dimension - - - -### Usage - -```python -import torch, chamfer3D.dist_chamfer_3D, fscore -chamLoss = chamfer3D.dist_chamfer_3D.chamfer_3DDist() -points1 = torch.rand(32, 1000, 3).cuda() -points2 = torch.rand(32, 2000, 3, requires_grad=True).cuda() -dist1, dist2, idx1, idx2 = chamLoss(points1, points2) -f_score, precision, recall = fscore.fscore(dist1, dist2) -``` - - - -### Add it to your project as a submodule - -```shell -git submodule add https://github.com/ThibaultGROUEIX/ChamferDistancePytorch -``` - - - -### Benchmark: [forward + backward] pass -- [x] CUDA 10.1, NVIDIA 435, Pytorch 1.4 -- [x] p1 : 32 x 2000 x dim -- [x] p2 : 32 x 1000 x dim - -| *Timing (sec * 1000)* | 2D | 3D | 5D | -| ---------- | -------- | ------- | ------- | -| **Cuda Compiled** | **1.2** | 1.4 |1.8 | -| **Cuda JIT** | 1.3 | **1.4** |**1.5** | -| **Python** | 37 | 37 | 37 | - - -| *Memory (MB)* | 2D | 3D | 5D | -| ---------- | -------- | ------- | ------- | -| **Cuda Compiled** | 529 | 529 | 549 | -| **Cuda JIT** | **520** | **529** |**549** | -| **Python** | 2495 | 2495 | 2495 | - - - -### What is the chamfer distance ? - -[Stanford course](http://graphics.stanford.edu/courses/cs468-17-spring/LectureSlides/L14%20-%203d%20deep%20learning%20on%20point%20cloud%20representation%20(analysis).pdf) on 3D deep Learning - - - -### Aknowledgment - -Original backbone from [Fei Xia](https://github.com/fxia22/pointGAN/blob/master/nndistance/src/nnd_cuda.cu). - -JIT cool trick from [Christian Diller](https://github.com/chrdiller) - -### Troubleshoot - -- `Undefined symbol: Zxxxxxxxxxxxxxxxxx `: - ---> Fix: Make sure to `import torch` before you `import chamfer`. ---> Use pytorch.version >= 1.1.0 - -- [RuntimeError: Ninja is required to load C++ extension](https://github.com/zhanghang1989/PyTorch-Encoding/issues/167) - -```shell -wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip -sudo unzip ninja-linux.zip -d /usr/local/bin/ -sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force -``` - - - - - -#### TODO: - -* Discuss behaviour of torch.min() and tensor.min() which causes issues in some pytorch versions diff --git a/distance/chamfer_multidim/chamfer2D/chamfer2D.cu b/distance/chamfer_multidim/chamfer2D/chamfer2D.cu deleted file mode 100644 index 567dd1a..0000000 --- a/distance/chamfer_multidim/chamfer2D/chamfer2D.cu +++ /dev/null @@ -1,182 +0,0 @@ - -#include -#include - -#include -#include - -#include - - - -__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ - const int batch=512; - __shared__ float buf[batch*2]; - for (int i=blockIdx.x;ibest){ - result[(i*n+j)]=best; - result_i[(i*n+j)]=best_i; - } - } - __syncthreads(); - } - } -} -// int chamfer_cuda_forward(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){ -int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){ - - const auto batch_size = xyz1.size(0); - const auto n = xyz1.size(1); //num_points point cloud A - const auto m = xyz2.size(1); //num_points point cloud B - - NmDistanceKernel<<>>(batch_size, n, xyz1.data(), m, xyz2.data(), dist1.data(), idx1.data()); - NmDistanceKernel<<>>(batch_size, m, xyz2.data(), n, xyz1.data(), dist2.data(), idx2.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err)); - //THError("aborting"); - return 0; - } - return 1; - - -} -__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){ - for (int i=blockIdx.x;i>>(batch_size,n,xyz1.data(),m,xyz2.data(),graddist1.data(),idx1.data(),gradxyz1.data(),gradxyz2.data()); - NmDistanceGradKernel<<>>(batch_size,m,xyz2.data(),n,xyz1.data(),graddist2.data(),idx2.data(),gradxyz2.data(),gradxyz1.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); - //THError("aborting"); - return 0; - } - return 1; - -} - diff --git a/distance/chamfer_multidim/chamfer2D/chamfer_cuda.cpp b/distance/chamfer_multidim/chamfer2D/chamfer_cuda.cpp deleted file mode 100644 index 67574e2..0000000 --- a/distance/chamfer_multidim/chamfer2D/chamfer_cuda.cpp +++ /dev/null @@ -1,33 +0,0 @@ -#include -#include - -///TMP -//#include "common.h" -/// NOT TMP - - -int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2); - - -int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2); - - - - -int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) { - return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2); -} - - -int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, - at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) { - - return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2); -} - - - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &chamfer_forward, "chamfer forward (CUDA)"); - m.def("backward", &chamfer_backward, "chamfer backward (CUDA)"); -} \ No newline at end of file diff --git a/distance/chamfer_multidim/chamfer2D/dist_chamfer_2D.py b/distance/chamfer_multidim/chamfer2D/dist_chamfer_2D.py deleted file mode 100644 index f92e6f1..0000000 --- a/distance/chamfer_multidim/chamfer2D/dist_chamfer_2D.py +++ /dev/null @@ -1,73 +0,0 @@ -from torch import nn -from torch.autograd import Function -import torch -import importlib -import os -chamfer_found = importlib.find_loader("chamfer_2D") is not None -if not chamfer_found: - ## Cool trick from https://github.com/chrdiller - print("Jitting Chamfer 2D") - - from torch.utils.cpp_extension import load - chamfer_2D = load(name="chamfer_2D", - sources=[ - "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]), - "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer2D.cu"]), - ]) - print("Loaded JIT 2D CUDA chamfer distance") - -else: - import chamfer_2D - print("Loaded compiled 2D CUDA chamfer distance") - -# Chamfer's distance module @thibaultgroueix -# GPU tensors only -class chamfer_2DFunction(Function): - @staticmethod - def forward(ctx, xyz1, xyz2): - batchsize, n, _ = xyz1.size() - _, m, _ = xyz2.size() - device = xyz1.device - - dist1 = torch.zeros(batchsize, n) - dist2 = torch.zeros(batchsize, m) - - idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) - idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) - - dist1 = dist1.to(device) - dist2 = dist2.to(device) - idx1 = idx1.to(device) - idx2 = idx2.to(device) - torch.cuda.set_device(device) - - chamfer_2D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) - ctx.save_for_backward(xyz1, xyz2, idx1, idx2) - return dist1, dist2, idx1, idx2 - - @staticmethod - def backward(ctx, graddist1, graddist2, gradidx1, gradidx2): - xyz1, xyz2, idx1, idx2 = ctx.saved_tensors - graddist1 = graddist1.contiguous() - graddist2 = graddist2.contiguous() - device = graddist1.device - - gradxyz1 = torch.zeros(xyz1.size()) - gradxyz2 = torch.zeros(xyz2.size()) - - gradxyz1 = gradxyz1.to(device) - gradxyz2 = gradxyz2.to(device) - chamfer_2D.backward( - xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2 - ) - return gradxyz1, gradxyz2 - - -class chamfer_2DDist(nn.Module): - def __init__(self): - super(chamfer_2DDist, self).__init__() - - def forward(self, input1, input2): - input1 = input1.contiguous() - input2 = input2.contiguous() - return chamfer_2DFunction.apply(input1, input2) diff --git a/distance/chamfer_multidim/chamfer2D/setup.py b/distance/chamfer_multidim/chamfer2D/setup.py deleted file mode 100644 index 11d0123..0000000 --- a/distance/chamfer_multidim/chamfer2D/setup.py +++ /dev/null @@ -1,14 +0,0 @@ -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup( - name='chamfer_2D', - ext_modules=[ - CUDAExtension('chamfer_2D', [ - "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']), - "/".join(__file__.split('/')[:-1] + ['chamfer2D.cu']), - ]), - ], - cmdclass={ - 'build_ext': BuildExtension - }) \ No newline at end of file diff --git a/distance/chamfer_multidim/chamfer3D/chamfer3D.cu b/distance/chamfer_multidim/chamfer3D/chamfer3D.cu deleted file mode 100644 index d5b886d..0000000 --- a/distance/chamfer_multidim/chamfer3D/chamfer3D.cu +++ /dev/null @@ -1,196 +0,0 @@ - -#include -#include - -#include -#include - -#include - - - -__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ - const int batch=512; - __shared__ float buf[batch*3]; - for (int i=blockIdx.x;ibest){ - result[(i*n+j)]=best; - result_i[(i*n+j)]=best_i; - } - } - __syncthreads(); - } - } -} -// int chamfer_cuda_forward(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){ -int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){ - - const auto batch_size = xyz1.size(0); - const auto n = xyz1.size(1); //num_points point cloud A - const auto m = xyz2.size(1); //num_points point cloud B - - NmDistanceKernel<<>>(batch_size, n, xyz1.data(), m, xyz2.data(), dist1.data(), idx1.data()); - NmDistanceKernel<<>>(batch_size, m, xyz2.data(), n, xyz1.data(), dist2.data(), idx2.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err)); - //THError("aborting"); - return 0; - } - return 1; - - -} -__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){ - for (int i=blockIdx.x;i>>(batch_size,n,xyz1.data(),m,xyz2.data(),graddist1.data(),idx1.data(),gradxyz1.data(),gradxyz2.data()); - NmDistanceGradKernel<<>>(batch_size,m,xyz2.data(),n,xyz1.data(),graddist2.data(),idx2.data(),gradxyz2.data(),gradxyz1.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); - //THError("aborting"); - return 0; - } - return 1; - -} - diff --git a/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/PKG-INFO b/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/PKG-INFO deleted file mode 100644 index 8897803..0000000 --- a/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 1.0 -Name: chamfer-3D -Version: 0.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -Author: UNKNOWN -Author-email: UNKNOWN -License: UNKNOWN -Description: UNKNOWN -Platform: UNKNOWN diff --git a/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/SOURCES.txt b/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/SOURCES.txt deleted file mode 100644 index bd5467d..0000000 --- a/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -chamfer3D.cu -chamfer_cuda.cpp -setup.py -chamfer_3D.egg-info/PKG-INFO -chamfer_3D.egg-info/SOURCES.txt -chamfer_3D.egg-info/dependency_links.txt -chamfer_3D.egg-info/top_level.txt \ No newline at end of file diff --git a/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/dependency_links.txt b/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/top_level.txt b/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/top_level.txt deleted file mode 100644 index dd1297b..0000000 --- a/distance/chamfer_multidim/chamfer3D/chamfer_3D.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -chamfer_3D diff --git a/distance/chamfer_multidim/chamfer3D/chamfer_cuda.cpp b/distance/chamfer_multidim/chamfer3D/chamfer_cuda.cpp deleted file mode 100644 index 67574e2..0000000 --- a/distance/chamfer_multidim/chamfer3D/chamfer_cuda.cpp +++ /dev/null @@ -1,33 +0,0 @@ -#include -#include - -///TMP -//#include "common.h" -/// NOT TMP - - -int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2); - - -int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2); - - - - -int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) { - return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2); -} - - -int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, - at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) { - - return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2); -} - - - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &chamfer_forward, "chamfer forward (CUDA)"); - m.def("backward", &chamfer_backward, "chamfer backward (CUDA)"); -} \ No newline at end of file diff --git a/distance/chamfer_multidim/chamfer3D/dist/chamfer_3D-0.0.0-py3.6-linux-x86_64.egg b/distance/chamfer_multidim/chamfer3D/dist/chamfer_3D-0.0.0-py3.6-linux-x86_64.egg deleted file mode 100644 index dc3ce4d..0000000 Binary files a/distance/chamfer_multidim/chamfer3D/dist/chamfer_3D-0.0.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/distance/chamfer_multidim/chamfer3D/dist/chamfer_3D-0.0.0-py3.7-linux-x86_64.egg b/distance/chamfer_multidim/chamfer3D/dist/chamfer_3D-0.0.0-py3.7-linux-x86_64.egg deleted file mode 100644 index 8289746..0000000 Binary files a/distance/chamfer_multidim/chamfer3D/dist/chamfer_3D-0.0.0-py3.7-linux-x86_64.egg and /dev/null differ diff --git a/distance/chamfer_multidim/chamfer3D/dist_chamfer_3D.py b/distance/chamfer_multidim/chamfer3D/dist_chamfer_3D.py deleted file mode 100644 index 4471f72..0000000 --- a/distance/chamfer_multidim/chamfer3D/dist_chamfer_3D.py +++ /dev/null @@ -1,75 +0,0 @@ -from torch import nn -from torch.autograd import Function -import torch -import importlib -import os -chamfer_found = importlib.find_loader("chamfer_3D") is not None -if not chamfer_found: - ## Cool trick from https://github.com/chrdiller - print("Jitting Chamfer 3D") - - from torch.utils.cpp_extension import load - chamfer_3D = load(name="chamfer_3D", - sources=[ - "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]), - "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer3D.cu"]), - ]) - print("Loaded JIT 3D CUDA chamfer distance") - -else: - import chamfer_3D - print("Loaded compiled 3D CUDA chamfer distance") - - -# Chamfer's distance module @thibaultgroueix -# GPU tensors only -class chamfer_3DFunction(Function): - @staticmethod - def forward(ctx, xyz1, xyz2): - batchsize, n, _ = xyz1.size() - _, m, _ = xyz2.size() - device = xyz1.device - - dist1 = torch.zeros(batchsize, n) - dist2 = torch.zeros(batchsize, m) - - idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) - idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) - - dist1 = dist1.to(device) - dist2 = dist2.to(device) - idx1 = idx1.to(device) - idx2 = idx2.to(device) - torch.cuda.set_device(device) - - chamfer_3D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) - ctx.save_for_backward(xyz1, xyz2, idx1, idx2) - return dist1, dist2, idx1, idx2 - - @staticmethod - def backward(ctx, graddist1, graddist2, gradidx1, gradidx2): - xyz1, xyz2, idx1, idx2 = ctx.saved_tensors - graddist1 = graddist1.contiguous() - graddist2 = graddist2.contiguous() - device = graddist1.device - - gradxyz1 = torch.zeros(xyz1.size()) - gradxyz2 = torch.zeros(xyz2.size()) - - gradxyz1 = gradxyz1.to(device) - gradxyz2 = gradxyz2.to(device) - chamfer_3D.backward( - xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2 - ) - return gradxyz1, gradxyz2 - - -class chamfer_3DDist(nn.Module): - def __init__(self): - super(chamfer_3DDist, self).__init__() - - def forward(self, input1, input2): - input1 = input1.contiguous() - input2 = input2.contiguous() - return chamfer_3DFunction.apply(input1, input2) - diff --git a/distance/chamfer_multidim/chamfer3D/setup.py b/distance/chamfer_multidim/chamfer3D/setup.py deleted file mode 100644 index 9a23aad..0000000 --- a/distance/chamfer_multidim/chamfer3D/setup.py +++ /dev/null @@ -1,14 +0,0 @@ -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup( - name='chamfer_3D', - ext_modules=[ - CUDAExtension('chamfer_3D', [ - "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']), - "/".join(__file__.split('/')[:-1] + ['chamfer3D.cu']), - ]), - ], - cmdclass={ - 'build_ext': BuildExtension - }) \ No newline at end of file diff --git a/distance/chamfer_multidim/chamfer5D/chamfer5D.cu b/distance/chamfer_multidim/chamfer5D/chamfer5D.cu deleted file mode 100644 index 650e889..0000000 --- a/distance/chamfer_multidim/chamfer5D/chamfer5D.cu +++ /dev/null @@ -1,223 +0,0 @@ - -#include -#include - -#include -#include - -#include - - - -__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ - const int batch=2048; - __shared__ float buf[batch*5]; - for (int i=blockIdx.x;ibest){ - result[(i*n+j)]=best; - result_i[(i*n+j)]=best_i; - } - } - __syncthreads(); - } - } -} -// int chamfer_cuda_forward(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){ -int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){ - - const auto batch_size = xyz1.size(0); - const auto n = xyz1.size(1); //num_points point cloud A - const auto m = xyz2.size(1); //num_points point cloud B - - NmDistanceKernel<<>>(batch_size, n, xyz1.data(), m, xyz2.data(), dist1.data(), idx1.data()); - NmDistanceKernel<<>>(batch_size, m, xyz2.data(), n, xyz1.data(), dist2.data(), idx2.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err)); - //THError("aborting"); - return 0; - } - return 1; - - -} -__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){ - for (int i=blockIdx.x;i>>(batch_size,n,xyz1.data(),m,xyz2.data(),graddist1.data(),idx1.data(),gradxyz1.data(),gradxyz2.data()); - NmDistanceGradKernel<<>>(batch_size,m,xyz2.data(),n,xyz1.data(),graddist2.data(),idx2.data(),gradxyz2.data(),gradxyz1.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); - //THError("aborting"); - return 0; - } - return 1; - -} diff --git a/distance/chamfer_multidim/chamfer5D/chamfer_cuda.cpp b/distance/chamfer_multidim/chamfer5D/chamfer_cuda.cpp deleted file mode 100644 index 67574e2..0000000 --- a/distance/chamfer_multidim/chamfer5D/chamfer_cuda.cpp +++ /dev/null @@ -1,33 +0,0 @@ -#include -#include - -///TMP -//#include "common.h" -/// NOT TMP - - -int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2); - - -int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2); - - - - -int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) { - return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2); -} - - -int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, - at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) { - - return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2); -} - - - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &chamfer_forward, "chamfer forward (CUDA)"); - m.def("backward", &chamfer_backward, "chamfer backward (CUDA)"); -} \ No newline at end of file diff --git a/distance/chamfer_multidim/chamfer5D/dist_chamfer_5D.py b/distance/chamfer_multidim/chamfer5D/dist_chamfer_5D.py deleted file mode 100644 index 3730a1f..0000000 --- a/distance/chamfer_multidim/chamfer5D/dist_chamfer_5D.py +++ /dev/null @@ -1,75 +0,0 @@ -from torch import nn -from torch.autograd import Function -import torch -import importlib -import os - -chamfer_found = importlib.find_loader("chamfer_5D") is not None -if not chamfer_found: - ## Cool trick from https://github.com/chrdiller - print("Jitting Chamfer 5D") - - from torch.utils.cpp_extension import load - chamfer_5D = load(name="chamfer_5D", - sources=[ - "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]), - "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer5D.cu"]), - ]) - print("Loaded JIT 5D CUDA chamfer distance") - -else: - import chamfer_5D - print("Loaded compiled 5D CUDA chamfer distance") - - -# Chamfer's distance module @thibaultgroueix -# GPU tensors only -class chamfer_5DFunction(Function): - @staticmethod - def forward(ctx, xyz1, xyz2): - batchsize, n, _ = xyz1.size() - _, m, _ = xyz2.size() - device = xyz1.device - - dist1 = torch.zeros(batchsize, n) - dist2 = torch.zeros(batchsize, m) - - idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) - idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) - - dist1 = dist1.to(device) - dist2 = dist2.to(device) - idx1 = idx1.to(device) - idx2 = idx2.to(device) - torch.cuda.set_device(device) - - chamfer_5D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) - ctx.save_for_backward(xyz1, xyz2, idx1, idx2) - return dist1, dist2, idx1, idx2 - - @staticmethod - def backward(ctx, graddist1, graddist2, gradidx1, gradidx2): - xyz1, xyz2, idx1, idx2 = ctx.saved_tensors - graddist1 = graddist1.contiguous() - graddist2 = graddist2.contiguous() - device = graddist1.device - - gradxyz1 = torch.zeros(xyz1.size()) - gradxyz2 = torch.zeros(xyz2.size()) - - gradxyz1 = gradxyz1.to(device) - gradxyz2 = gradxyz2.to(device) - chamfer_5D.backward( - xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2 - ) - return gradxyz1, gradxyz2 - - -class chamfer_5DDist(nn.Module): - def __init__(self): - super(chamfer_5DDist, self).__init__() - - def forward(self, input1, input2): - input1 = input1.contiguous() - input2 = input2.contiguous() - return chamfer_5DFunction.apply(input1, input2) diff --git a/distance/chamfer_multidim/chamfer5D/setup.py b/distance/chamfer_multidim/chamfer5D/setup.py deleted file mode 100644 index 2429235..0000000 --- a/distance/chamfer_multidim/chamfer5D/setup.py +++ /dev/null @@ -1,14 +0,0 @@ -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup( - name='chamfer_5D', - ext_modules=[ - CUDAExtension('chamfer_5D', [ - "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']), - "/".join(__file__.split('/')[:-1] + ['chamfer5D.cu']), - ]), - ], - cmdclass={ - 'build_ext': BuildExtension - }) \ No newline at end of file diff --git a/distance/chamfer_multidim/chamfer_python.py b/distance/chamfer_multidim/chamfer_python.py deleted file mode 100644 index ce0aeaa..0000000 --- a/distance/chamfer_multidim/chamfer_python.py +++ /dev/null @@ -1,40 +0,0 @@ -import torch - - -def pairwise_dist(x, y): - xx, yy, zz = torch.mm(x, x.t()), torch.mm(y, y.t()), torch.mm(x, y.t()) - rx = xx.diag().unsqueeze(0).expand_as(xx) - ry = yy.diag().unsqueeze(0).expand_as(yy) - P = rx.t() + ry - 2 * zz - return P - - -def NN_loss(x, y, dim=0): - dist = pairwise_dist(x, y) - values, indices = dist.min(dim=dim) - return values.mean() - - -def distChamfer(a, b): - """ - :param a: Pointclouds Batch x nul_points x dim - :param b: Pointclouds Batch x nul_points x dim - :return: - -closest point on b of points from a - -closest point on a of points from b - -idx of closest point on b of points from a - -idx of closest point on a of points from b - Works for pointcloud of any dimension - """ - x, y = a.double(), b.double() - bs, num_points_x, points_dim = x.size() - bs, num_points_y, points_dim = y.size() - - xx = torch.pow(x, 2).sum(2) - yy = torch.pow(y, 2).sum(2) - zz = torch.bmm(x, y.transpose(2, 1)) - rx = xx.unsqueeze(1).expand(bs, num_points_y, num_points_x) # Diagonal elements xx - ry = yy.unsqueeze(1).expand(bs, num_points_x, num_points_y) # Diagonal elements yy - P = rx.transpose(2, 1) + ry - 2 * zz - return torch.min(P, 2)[0].float(), torch.min(P, 1)[0].float(), torch.min(P, 2)[1].int(), torch.min(P, 1)[1].int() - diff --git a/distance/chamfer_multidim/fscore.py b/distance/chamfer_multidim/fscore.py deleted file mode 100644 index 265378b..0000000 --- a/distance/chamfer_multidim/fscore.py +++ /dev/null @@ -1,17 +0,0 @@ -import torch - -def fscore(dist1, dist2, threshold=0.001): - """ - Calculates the F-score between two point clouds with the corresponding threshold value. - :param dist1: Batch, N-Points - :param dist2: Batch, N-Points - :param th: float - :return: fscore, precision, recall - """ - # NB : In this depo, dist1 and dist2 are squared pointcloud euclidean distances, so you should adapt the threshold accordingly. - precision_1 = torch.mean((dist1 < threshold).float(), dim=1) - precision_2 = torch.mean((dist2 < threshold).float(), dim=1) - fscore = 2 * precision_1 * precision_2 / (precision_1 + precision_2) - fscore[torch.isnan(fscore)] = 0 - return fscore, precision_1, precision_2 - diff --git a/distance/chamfer_multidim/unit_test.py b/distance/chamfer_multidim/unit_test.py deleted file mode 100644 index 13af6a3..0000000 --- a/distance/chamfer_multidim/unit_test.py +++ /dev/null @@ -1,69 +0,0 @@ -import torch, time -import chamfer2D.dist_chamfer_2D -import chamfer3D.dist_chamfer_3D -import chamfer5D.dist_chamfer_5D -import chamfer_python - -cham2D = chamfer2D.dist_chamfer_2D.chamfer_2DDist() -cham3D = chamfer3D.dist_chamfer_3D.chamfer_3DDist() -cham5D = chamfer5D.dist_chamfer_5D.chamfer_5DDist() - -from torch.autograd import Variable -from fscore import fscore - -def test_chamfer(distChamfer, dim): - points1 = torch.rand(4, 100, dim).cuda() - points2 = torch.rand(4, 200, dim, requires_grad=True).cuda() - dist1, dist2, idx1, idx2= distChamfer(points1, points2) - - loss = torch.sum(dist1) - loss.backward() - - mydist1, mydist2, myidx1, myidx2 = chamfer_python.distChamfer(points1, points2) - d1 = (dist1 - mydist1) ** 2 - d2 = (dist2 - mydist2) ** 2 - assert ( - torch.mean(d1) + torch.mean(d2) < 0.00000001 - ), "chamfer cuda and chamfer normal are not giving the same results" - - xd1 = idx1 - myidx1 - xd2 = idx2 - myidx2 - assert ( - torch.norm(xd1.float()) + torch.norm(xd2.float()) == 0 - ), "chamfer cuda and chamfer normal are not giving the same results" - print(f"fscore :", fscore(dist1, dist2)) - print("Unit test passed") - - -def timings(distChamfer, dim): - p1 = torch.rand(32, 2000, dim).cuda() - p2 = torch.rand(32, 1000, dim).cuda() - print("Timings : Start CUDA version") - start = time.time() - num_it = 100 - for i in range(num_it): - points1 = Variable(p1, requires_grad=True) - points2 = Variable(p2) - mydist1, mydist2, idx1, idx2 = distChamfer(points1, points2) - loss = torch.sum(mydist1) - loss.backward() - print(f"Ellapsed time forward backward is {(time.time() - start)/num_it} seconds.") - - - print("Timings : Start Pythonic version") - start = time.time() - for i in range(num_it): - points1 = Variable(p1, requires_grad=True) - points2 = Variable(p2) - mydist1, mydist2, idx1, idx2 = chamfer_python.distChamfer(points1, points2) - loss = torch.sum(mydist1) - loss.backward() - print(f"Ellapsed time forward backward is {(time.time() - start)/num_it} seconds.") - - - -dims = [2,3,5] -for i,cham in enumerate([cham2D, cham3D, cham5D]): - print(f"testing Chamfer {dims[i]}D") - test_chamfer(cham, dims[i]) - timings(cham, dims[i]) diff --git a/distance/emd/CDEMD.png b/distance/emd/CDEMD.png deleted file mode 100644 index 5136f07..0000000 Binary files a/distance/emd/CDEMD.png and /dev/null differ diff --git a/distance/emd/README.md b/distance/emd/README.md deleted file mode 100644 index df02cc7..0000000 --- a/distance/emd/README.md +++ /dev/null @@ -1,24 +0,0 @@ -## Earth Mover's Distance of point clouds - -![](/emd/CDEMD.png) - -Compared to the Chamfer Distance (CD), the Earth Mover's Distance (EMD) is more reliable to distinguish the visual quality of the point clouds. See our [paper](http://cseweb.ucsd.edu/~mil070/projects/AAAI2020/paper.pdf) for more details. - -We provide an EMD implementation for point cloud comparison, which only needs $O(n)$ memory and thus enables dense point clouds (with 10,000 points or over) and large batch size. It is based on an approximated algorithm (auction algorithm) and cannot guarantee a (but near) bijection assignment. It employs a parameter $\epsilon$ to balance the error rate and the speed of convergence. Smaller $\epsilon$ achieves more accurate results, but needs a longer time for convergence. The time complexity is $O(n^2k)$, where $k$ is the number of iterations. We set a $\epsilon = 0.005, k = 50$ during training and a $\epsilon = 0.002, k = 10000$ during testing. - -### Compile -Run `python3 setup.py install` to compile. - -### Example -See `emd_module.py/test_emd()` for examples. - -### Input - -- **xyz1, xyz2**: float tensors with shape `[#batch, #points, 3]`. xyz1 is the predicted point cloud and xyz2 is the ground truth point cloud. Two point clouds should have same size and be normalized to [0, 1]. The number of points should be a multiple of 1024. The batch size should be no greater than 512. Since we only calculate gradients for xyz1, please do not swap xyz1 and xyz2. -- **eps**: a float tensor, the parameter balances the error rate and the speed of convergence. -- **iters**: a int tensor, the number of iterations. - -### Output - -- **dist**: a float tensor with shape `[#batch, #points]`. sqrt(dist) are the L2 distances between the pairs of points. -- **assignment**: a int tensor with shape `[#batch, #points]`. The index of the matched point in the ground truth point cloud. diff --git a/distance/emd/__pycache__/emd_module.cpython-36.pyc b/distance/emd/__pycache__/emd_module.cpython-36.pyc deleted file mode 100644 index f0e384e..0000000 Binary files a/distance/emd/__pycache__/emd_module.cpython-36.pyc and /dev/null differ diff --git a/distance/emd/dist/emd-0.0.0-py2.7-linux-x86_64.egg b/distance/emd/dist/emd-0.0.0-py2.7-linux-x86_64.egg deleted file mode 100644 index 45e6165..0000000 Binary files a/distance/emd/dist/emd-0.0.0-py2.7-linux-x86_64.egg and /dev/null differ diff --git a/distance/emd/dist/emd-0.0.0-py3.6-linux-x86_64.egg b/distance/emd/dist/emd-0.0.0-py3.6-linux-x86_64.egg deleted file mode 100644 index 56e3cdd..0000000 Binary files a/distance/emd/dist/emd-0.0.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/distance/emd/emd.cpp b/distance/emd/emd.cpp deleted file mode 100644 index 5036016..0000000 --- a/distance/emd/emd.cpp +++ /dev/null @@ -1,31 +0,0 @@ -// EMD approximation module (based on auction algorithm) -// author: Minghua Liu -#include -#include - -int emd_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, - at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments, - at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters); - -int emd_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx); - - - -int emd_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, - at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments, - at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters) { - return emd_cuda_forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx, unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters); -} - -int emd_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx) { - - return emd_cuda_backward(xyz1, xyz2, gradxyz, graddist, idx); -} - - - - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &emd_forward, "emd forward (CUDA)"); - m.def("backward", &emd_backward, "emd backward (CUDA)"); -} \ No newline at end of file diff --git a/distance/emd/emd.egg-info/PKG-INFO b/distance/emd/emd.egg-info/PKG-INFO deleted file mode 100644 index cfdc103..0000000 --- a/distance/emd/emd.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: emd -Version: 0.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/distance/emd/emd.egg-info/SOURCES.txt b/distance/emd/emd.egg-info/SOURCES.txt deleted file mode 100644 index 3560814..0000000 --- a/distance/emd/emd.egg-info/SOURCES.txt +++ /dev/null @@ -1,8 +0,0 @@ -README.md -emd.cpp -emd_cuda.cu -setup.py -emd.egg-info/PKG-INFO -emd.egg-info/SOURCES.txt -emd.egg-info/dependency_links.txt -emd.egg-info/top_level.txt \ No newline at end of file diff --git a/distance/emd/emd.egg-info/dependency_links.txt b/distance/emd/emd.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/distance/emd/emd.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/distance/emd/emd.egg-info/top_level.txt b/distance/emd/emd.egg-info/top_level.txt deleted file mode 100644 index 9f69f84..0000000 --- a/distance/emd/emd.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -emd diff --git a/distance/emd/emd_cuda.cu b/distance/emd/emd_cuda.cu deleted file mode 100644 index 08999bb..0000000 --- a/distance/emd/emd_cuda.cu +++ /dev/null @@ -1,316 +0,0 @@ -// EMD approximation module (based on auction algorithm) -// author: Minghua Liu -#include -#include - -#include -#include -#include - -__device__ __forceinline__ float atomicMax(float *address, float val) -{ - int ret = __float_as_int(*address); - while(val > __int_as_float(ret)) - { - int old = ret; - if((ret = atomicCAS((int *)address, old, __float_as_int(val))) == old) - break; - } - return __int_as_float(ret); -} - - -__global__ void clear(int b, int * cnt_tmp, int * unass_cnt) { - for (int i = threadIdx.x; i < b; i += blockDim.x) { - cnt_tmp[i] = 0; - unass_cnt[i] = 0; - } -} - -__global__ void calc_unass_cnt(int b, int n, int * assignment, int * unass_cnt) { - // count the number of unassigned points in each batch - const int BLOCK_SIZE = 1024; - __shared__ int scan_array[BLOCK_SIZE]; - for (int i = blockIdx.x; i < b; i += gridDim.x) { - scan_array[threadIdx.x] = assignment[i * n + blockIdx.y * BLOCK_SIZE + threadIdx.x] == -1 ? 1 : 0; - __syncthreads(); - - int stride = 1; - while(stride <= BLOCK_SIZE / 2) { - int index = (threadIdx.x + 1) * stride * 2 - 1; - if(index < BLOCK_SIZE) - scan_array[index] += scan_array[index - stride]; - stride = stride * 2; - __syncthreads(); - } - __syncthreads(); - - if (threadIdx.x == BLOCK_SIZE - 1) { - atomicAdd(&unass_cnt[i], scan_array[threadIdx.x]); - } - __syncthreads(); - } -} - -__global__ void calc_unass_cnt_sum(int b, int * unass_cnt, int * unass_cnt_sum) { - // count the cumulative sum over over unass_cnt - const int BLOCK_SIZE = 512; // batch_size <= 512 - __shared__ int scan_array[BLOCK_SIZE]; - scan_array[threadIdx.x] = unass_cnt[threadIdx.x]; - __syncthreads(); - - int stride = 1; - while(stride <= BLOCK_SIZE / 2) { - int index = (threadIdx.x + 1) * stride * 2 - 1; - if(index < BLOCK_SIZE) - scan_array[index] += scan_array[index - stride]; - stride = stride * 2; - __syncthreads(); - } - __syncthreads(); - stride = BLOCK_SIZE / 4; - while(stride > 0) { - int index = (threadIdx.x + 1) * stride * 2 - 1; - if((index + stride) < BLOCK_SIZE) - scan_array[index + stride] += scan_array[index]; - stride = stride / 2; - __syncthreads(); - } - __syncthreads(); - - //printf("%d\n", unass_cnt_sum[b - 1]); - unass_cnt_sum[threadIdx.x] = scan_array[threadIdx.x]; -} - -__global__ void calc_unass_idx(int b, int n, int * assignment, int * unass_idx, int * unass_cnt, int * unass_cnt_sum, int * cnt_tmp) { - // list all the unassigned points - for (int i = blockIdx.x; i < b; i += gridDim.x) { - if (assignment[i * n + blockIdx.y * 1024 + threadIdx.x] == -1) { - int idx = atomicAdd(&cnt_tmp[i], 1); - unass_idx[unass_cnt_sum[i] - unass_cnt[i] + idx] = blockIdx.y * 1024 + threadIdx.x; - } - } -} - -__global__ void Bid(int b, int n, const float * xyz1, const float * xyz2, float eps, int * assignment, int * assignment_inv, float * price, - int * bid, float * bid_increments, float * max_increments, int * unass_cnt, int * unass_cnt_sum, int * unass_idx) { - const int batch = 2048, block_size = 1024, block_cnt = n / 1024; - __shared__ float xyz2_buf[batch * 3]; - __shared__ float price_buf[batch]; - __shared__ float best_buf[block_size]; - __shared__ float better_buf[block_size]; - __shared__ int best_i_buf[block_size]; - for (int i = blockIdx.x; i < b; i += gridDim.x) { - int _unass_cnt = unass_cnt[i]; - if (_unass_cnt == 0) - continue; - int _unass_cnt_sum = unass_cnt_sum[i]; - int unass_per_block = (_unass_cnt + block_cnt - 1) / block_cnt; - int thread_per_unass = block_size / unass_per_block; - int unass_this_block = max(min(_unass_cnt - (int) blockIdx.y * unass_per_block, unass_per_block), 0); - - float x1, y1, z1, best = -1e9, better = -1e9; - int best_i = -1, _unass_id = -1, thread_in_unass; - - if (threadIdx.x < thread_per_unass * unass_this_block) { - _unass_id = unass_per_block * blockIdx.y + threadIdx.x / thread_per_unass + _unass_cnt_sum - _unass_cnt; - _unass_id = unass_idx[_unass_id]; - thread_in_unass = threadIdx.x % thread_per_unass; - - x1 = xyz1[(i * n + _unass_id) * 3 + 0]; - y1 = xyz1[(i * n + _unass_id) * 3 + 1]; - z1 = xyz1[(i * n + _unass_id) * 3 + 2]; - } - - for (int k2 = 0; k2 < n; k2 += batch) { - int end_k = min(n, k2 + batch) - k2; - for (int j = threadIdx.x; j < end_k * 3; j += blockDim.x) { - xyz2_buf[j] = xyz2[(i * n + k2) * 3 + j]; - } - for (int j = threadIdx.x; j < end_k; j += blockDim.x) { - price_buf[j] = price[i * n + k2 + j]; - } - __syncthreads(); - - if (_unass_id != -1) { - int delta = (end_k + thread_per_unass - 1) / thread_per_unass; - int l = thread_in_unass * delta; - int r = min((thread_in_unass + 1) * delta, end_k); - for (int k = l; k < r; k++) - //if (!last || assignment_inv[i * n + k + k2] == -1) - { - float x2 = xyz2_buf[k * 3 + 0] - x1; - float y2 = xyz2_buf[k * 3 + 1] - y1; - float z2 = xyz2_buf[k * 3 + 2] - z1; - // the coordinates of points should be normalized to [0, 1] - float d = 3.0 - sqrtf(x2 * x2 + y2 * y2 + z2 * z2) - price_buf[k]; - if (d > best) { - better = best; - best = d; - best_i = k + k2; - } - else if (d > better) { - better = d; - } - } - } - __syncthreads(); - } - - best_buf[threadIdx.x] = best; - better_buf[threadIdx.x] = better; - best_i_buf[threadIdx.x] = best_i; - __syncthreads(); - - if (_unass_id != -1 && thread_in_unass == 0) { - for (int j = threadIdx.x + 1; j < threadIdx.x + thread_per_unass; j++) { - if (best_buf[j] > best) { - better = max(best, better_buf[j]); - best = best_buf[j]; - best_i = best_i_buf[j]; - } - else better = max(better, best_buf[j]); - } - bid[i * n + _unass_id] = best_i; - bid_increments[i * n + _unass_id] = best - better + eps; - atomicMax(&max_increments[i * n + best_i], best - better + eps); - } - } -} - -__global__ void GetMax(int b, int n, int * assignment, int * bid, float * bid_increments, float * max_increments, int * max_idx) { - for (int i = blockIdx.x; i < b; i += gridDim.x) { - int j = threadIdx.x + blockIdx.y * blockDim.x; - if (assignment[i * n + j] == -1) { - int bid_id = bid[i * n + j]; - float bid_inc = bid_increments[i * n + j]; - float max_inc = max_increments[i * n + bid_id]; - if (bid_inc - 1e-6 <= max_inc && max_inc <= bid_inc + 1e-6) - { - max_idx[i * n + bid_id] = j; - } - } - } -} - -__global__ void Assign(int b, int n, int * assignment, int * assignment_inv, float * price, int * bid, float * bid_increments, float * max_increments, int * max_idx, bool last) { - for (int i = blockIdx.x; i < b; i += gridDim.x) { - int j = threadIdx.x + blockIdx.y * blockDim.x; - if (assignment[i * n + j] == -1) { - int bid_id = bid[i * n + j]; - if (last || max_idx[i * n + bid_id] == j) - { - float bid_inc = bid_increments[i * n + j]; - int ass_inv = assignment_inv[i * n + bid_id]; - if (!last && ass_inv != -1) { - assignment[i * n + ass_inv] = -1; - } - assignment_inv[i * n + bid_id] = j; - assignment[i * n + j] = bid_id; - price[i * n + bid_id] += bid_inc; - max_increments[i * n + bid_id] = -1e9; - } - } - } -} - -__global__ void CalcDist(int b, int n, float * xyz1, float * xyz2, float * dist, int * assignment) { - for (int i = blockIdx.x; i < b; i += gridDim.x) { - int j = threadIdx.x + blockIdx.y * blockDim.x; - int k = assignment[i * n + j]; - float deltax = xyz1[(i * n + j) * 3 + 0] - xyz2[(i * n + k) * 3 + 0]; - float deltay = xyz1[(i * n + j) * 3 + 1] - xyz2[(i * n + k) * 3 + 1]; - float deltaz = xyz1[(i * n + j) * 3 + 2] - xyz2[(i * n + k) * 3 + 2]; - dist[i * n + j] = deltax * deltax + deltay * deltay + deltaz * deltaz; - } -} - -int emd_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, - at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments, - at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters) { - - const auto batch_size = xyz1.size(0); - const auto n = xyz1.size(1); //num_points point cloud A - const auto m = xyz2.size(1); //num_points point cloud B - - if (n != m) { - printf("Input Error! The two point clouds should have the same size.\n"); - return -1; - } - - if (batch_size > 512) { - printf("Input Error! The batch size should be less than 512.\n"); - return -1; - } - - if (n % 1024 != 0) { - printf("Input Error! The size of the point clouds should be a multiple of 1024.\n"); - return -1; - } - - //cudaEvent_t start,stop; - //cudaEventCreate(&start); - //cudaEventCreate(&stop); - //cudaEventRecord(start); - //int iters = 50; - for (int i = 0; i < iters; i++) { - clear<<<1, batch_size>>>(batch_size, cnt_tmp.data(), unass_cnt.data()); - calc_unass_cnt<<>>(batch_size, n, assignment.data(), unass_cnt.data()); - calc_unass_cnt_sum<<<1, batch_size>>>(batch_size, unass_cnt.data(), unass_cnt_sum.data()); - calc_unass_idx<<>>(batch_size, n, assignment.data(), unass_idx.data(), unass_cnt.data(), - unass_cnt_sum.data(), cnt_tmp.data()); - Bid<<>>(batch_size, n, xyz1.data(), xyz2.data(), eps, assignment.data(), assignment_inv.data(), - price.data(), bid.data(), bid_increments.data(), max_increments.data(), - unass_cnt.data(), unass_cnt_sum.data(), unass_idx.data()); - GetMax<<>>(batch_size, n, assignment.data(), bid.data(), bid_increments.data(), max_increments.data(), max_idx.data()); - Assign<<>>(batch_size, n, assignment.data(), assignment_inv.data(), price.data(), bid.data(), - bid_increments.data(), max_increments.data(), max_idx.data(), i == iters - 1); - } - CalcDist<<>>(batch_size, n, xyz1.data(), xyz2.data(), dist.data(), assignment.data()); - //cudaEventRecord(stop); - //cudaEventSynchronize(stop); - //float elapsedTime; - //cudaEventElapsedTime(&elapsedTime,start,stop); - //printf("%lf\n", elapsedTime); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd Output: %s\n", cudaGetErrorString(err)); - return 0; - } - return 1; -} - -__global__ void NmDistanceGradKernel(int b, int n, const float * xyz1, const float * xyz2, const float * grad_dist, const int * idx, float * grad_xyz){ - for (int i = blockIdx.x; i < b; i += gridDim.x) { - for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; j += blockDim.x * gridDim.y) { - float x1 = xyz1[(i * n + j) * 3 + 0]; - float y1 = xyz1[(i * n + j) * 3 + 1]; - float z1 = xyz1[(i * n + j) * 3 + 2]; - int j2 = idx[i * n + j]; - float x2 = xyz2[(i * n + j2) * 3 + 0]; - float y2 = xyz2[(i * n + j2) * 3 + 1]; - float z2 = xyz2[(i * n + j2) * 3 + 2]; - float g = grad_dist[i * n + j] * 2; - atomicAdd(&(grad_xyz[(i * n + j) * 3 + 0]), g * (x1 - x2)); - atomicAdd(&(grad_xyz[(i * n + j) * 3 + 1]), g * (y1 - y2)); - atomicAdd(&(grad_xyz[(i * n + j) * 3 + 2]), g * (z1 - z2)); - } - } -} - -int emd_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx){ - const auto batch_size = xyz1.size(0); - const auto n = xyz1.size(1); - const auto m = xyz2.size(1); - - NmDistanceGradKernel<<>>(batch_size, n, xyz1.data(), xyz2.data(), graddist.data(), idx.data(), gradxyz.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); - return 0; - } - return 1; - -} diff --git a/distance/emd/emd_module.py b/distance/emd/emd_module.py deleted file mode 100644 index a81a63e..0000000 --- a/distance/emd/emd_module.py +++ /dev/null @@ -1,98 +0,0 @@ -# EMD approximation module (based on auction algorithm) -# memory complexity: O(n) -# time complexity: O(n^2 * iter) -# author: Minghua Liu - -# Input: -# xyz1, xyz2: [#batch, #points, 3] -# where xyz1 is the predicted point cloud and xyz2 is the ground truth point cloud -# two point clouds should have same size and be normalized to [0, 1] -# #points should be a multiple of 1024 -# #batch should be no greater than 512 -# eps is a parameter which balances the error rate and the speed of convergence -# iters is the number of iteration -# we only calculate gradient for xyz1 - -# Output: -# dist: [#batch, #points], sqrt(dist) -> L2 distance -# assignment: [#batch, #points], index of the matched point in the ground truth point cloud -# the result is an approximation and the assignment is not guranteed to be a bijection - -import time -import numpy as np -import torch -from torch import nn -from torch.autograd import Function -import emd - - - - -class emdFunction(Function): - @staticmethod - def forward(ctx, xyz1, xyz2, eps, iters): - - batchsize, n, _ = xyz1.size() - _, m, _ = xyz2.size() - - assert(n == m) - assert(xyz1.size()[0] == xyz2.size()[0]) - assert(n % 1024 == 0) - assert(batchsize <= 512) - - xyz1 = xyz1.contiguous().float().cuda() - xyz2 = xyz2.contiguous().float().cuda() - dist = torch.zeros(batchsize, n, device='cuda').contiguous() - assignment = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous() - 1 - assignment_inv = torch.zeros(batchsize, m, device='cuda', dtype=torch.int32).contiguous() - 1 - price = torch.zeros(batchsize, m, device='cuda').contiguous() - bid = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous() - bid_increments = torch.zeros(batchsize, n, device='cuda').contiguous() - max_increments = torch.zeros(batchsize, m, device='cuda').contiguous() - unass_idx = torch.zeros(batchsize * n, device='cuda', dtype=torch.int32).contiguous() - max_idx = torch.zeros(batchsize * m, device='cuda', dtype=torch.int32).contiguous() - unass_cnt = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous() - unass_cnt_sum = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous() - cnt_tmp = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous() - - emd.forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx, unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters) - - ctx.save_for_backward(xyz1, xyz2, assignment) - return dist, assignment - - @staticmethod - def backward(ctx, graddist, gradidx): - xyz1, xyz2, assignment = ctx.saved_tensors - graddist = graddist.contiguous() - - gradxyz1 = torch.zeros(xyz1.size(), device='cuda').contiguous() - gradxyz2 = torch.zeros(xyz2.size(), device='cuda').contiguous() - - emd.backward(xyz1, xyz2, gradxyz1, graddist, assignment) - return gradxyz1, gradxyz2, None, None - -class emdModule(nn.Module): - def __init__(self): - super(emdModule, self).__init__() - - def forward(self, input1, input2, eps, iters): - return emdFunction.apply(input1, input2, eps, iters) - -def test_emd(): - x1 = torch.rand(20, 8192, 3).cuda() - x2 = torch.rand(20, 8192, 3).cuda() - emd = emdModule() - start_time = time.perf_counter() - dis, assigment = emd(x1, x2, 0.05, 3000) - print("Input_size: ", x1.shape) - print("Runtime: %lfs" % (time.perf_counter() - start_time)) - print("EMD: %lf" % np.sqrt(dis.cpu()).mean()) - print("|set(assignment)|: %d" % assigment.unique().numel()) - assigment = assigment.cpu().numpy() - assigment = np.expand_dims(assigment, -1) - x2 = np.take_along_axis(x2, assigment, axis = 1) - d = (x1 - x2) * (x1 - x2) - print("Verified EMD: %lf" % np.sqrt(d.cpu().sum(-1)).mean()) - -#test_emd() - diff --git a/distance/emd/setup.py b/distance/emd/setup.py deleted file mode 100644 index 8588de9..0000000 --- a/distance/emd/setup.py +++ /dev/null @@ -1,14 +0,0 @@ -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup( - name='emd', - ext_modules=[ - CUDAExtension('emd', [ - 'emd.cpp', - 'emd_cuda.cu', - ]), - ], - cmdclass={ - 'build_ext': BuildExtension - }) \ No newline at end of file diff --git a/other_models/GRNet/__pycache__/grnet.cpython-36.pyc b/other_models/GRNet/__pycache__/grnet.cpython-36.pyc deleted file mode 100644 index 392e549..0000000 Binary files a/other_models/GRNet/__pycache__/grnet.cpython-36.pyc and /dev/null differ diff --git a/other_models/GRNet/extensions/__init__.py b/other_models/GRNet/extensions/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/other_models/GRNet/extensions/__pycache__/__init__.cpython-36.pyc b/other_models/GRNet/extensions/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index dcc6f47..0000000 Binary files a/other_models/GRNet/extensions/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/GRNet/extensions/chamfer_dist/__init__.py b/other_models/GRNet/extensions/chamfer_dist/__init__.py deleted file mode 100644 index 4aaf4a3..0000000 --- a/other_models/GRNet/extensions/chamfer_dist/__init__.py +++ /dev/null @@ -1,42 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Thibault GROUEIX -# @Date: 2019-08-07 20:54:24 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-18 15:06:25 -# @Email: cshzxie@gmail.com - -import torch - -import chamfer - - -class ChamferFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, xyz1, xyz2): - dist1, dist2, idx1, idx2 = chamfer.forward(xyz1, xyz2) - ctx.save_for_backward(xyz1, xyz2, idx1, idx2) - - return dist1, dist2 - - @staticmethod - def backward(ctx, grad_dist1, grad_dist2): - xyz1, xyz2, idx1, idx2 = ctx.saved_tensors - grad_xyz1, grad_xyz2 = chamfer.backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2) - return grad_xyz1, grad_xyz2 - - -class ChamferDistance(torch.nn.Module): - def __init__(self, ignore_zeros=False): - super(ChamferDistance, self).__init__() - self.ignore_zeros = ignore_zeros - - def forward(self, xyz1, xyz2): - batch_size = xyz1.size(0) - if batch_size == 1 and self.ignore_zeros: - non_zeros1 = torch.sum(xyz1, dim=2).ne(0) - non_zeros2 = torch.sum(xyz2, dim=2).ne(0) - xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) - xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) - - dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) - return torch.mean(dist1) + torch.mean(dist2) diff --git a/other_models/GRNet/extensions/chamfer_dist/chamfer.cu b/other_models/GRNet/extensions/chamfer_dist/chamfer.cu deleted file mode 100644 index 4bde058..0000000 --- a/other_models/GRNet/extensions/chamfer_dist/chamfer.cu +++ /dev/null @@ -1,229 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-08-07 20:54:24 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:58:55 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include - -#include - -__global__ void chamfer_dist_kernel(int batch_size, - int n, - const float* xyz1, - int m, - const float* xyz2, - float* dist, - int* indexes) { - const int batch = 512; - __shared__ float buf[batch * 3]; - for (int i = blockIdx.x; i < batch_size; i += gridDim.x) { - for (int k2 = 0; k2 < m; k2 += batch) { - int end_k = min(m, k2 + batch) - k2; - for (int j = threadIdx.x; j < end_k * 3; j += blockDim.x) { - buf[j] = xyz2[(i * m + k2) * 3 + j]; - } - __syncthreads(); - for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; - j += blockDim.x * gridDim.y) { - float x1 = xyz1[(i * n + j) * 3 + 0]; - float y1 = xyz1[(i * n + j) * 3 + 1]; - float z1 = xyz1[(i * n + j) * 3 + 2]; - float best_dist = 0; - int best_dist_index = 0; - int end_ka = end_k - (end_k & 3); - if (end_ka == batch) { - for (int k = 0; k < batch; k += 4) { - { - float x2 = buf[k * 3 + 0] - x1; - float y2 = buf[k * 3 + 1] - y1; - float z2 = buf[k * 3 + 2] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - - if (k == 0 || dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2; - } - } - { - float x2 = buf[k * 3 + 3] - x1; - float y2 = buf[k * 3 + 4] - y1; - float z2 = buf[k * 3 + 5] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 1; - } - } - { - float x2 = buf[k * 3 + 6] - x1; - float y2 = buf[k * 3 + 7] - y1; - float z2 = buf[k * 3 + 8] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 2; - } - } - { - float x2 = buf[k * 3 + 9] - x1; - float y2 = buf[k * 3 + 10] - y1; - float z2 = buf[k * 3 + 11] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 3; - } - } - } - } else { - for (int k = 0; k < end_ka; k += 4) { - { - float x2 = buf[k * 3 + 0] - x1; - float y2 = buf[k * 3 + 1] - y1; - float z2 = buf[k * 3 + 2] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (k == 0 || dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2; - } - } - { - float x2 = buf[k * 3 + 3] - x1; - float y2 = buf[k * 3 + 4] - y1; - float z2 = buf[k * 3 + 5] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 1; - } - } - { - float x2 = buf[k * 3 + 6] - x1; - float y2 = buf[k * 3 + 7] - y1; - float z2 = buf[k * 3 + 8] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 2; - } - } - { - float x2 = buf[k * 3 + 9] - x1; - float y2 = buf[k * 3 + 10] - y1; - float z2 = buf[k * 3 + 11] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 3; - } - } - } - } - for (int k = end_ka; k < end_k; k++) { - float x2 = buf[k * 3 + 0] - x1; - float y2 = buf[k * 3 + 1] - y1; - float z2 = buf[k * 3 + 2] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (k == 0 || dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2; - } - } - if (k2 == 0 || dist[(i * n + j)] > best_dist) { - dist[(i * n + j)] = best_dist; - indexes[(i * n + j)] = best_dist_index; - } - } - __syncthreads(); - } - } -} - -std::vector chamfer_cuda_forward(torch::Tensor xyz1, - torch::Tensor xyz2) { - const int batch_size = xyz1.size(0); - const int n = xyz1.size(1); // num_points point cloud A - const int m = xyz2.size(1); // num_points point cloud B - torch::Tensor dist1 = - torch::zeros({batch_size, n}, torch::CUDA(torch::kFloat)); - torch::Tensor dist2 = - torch::zeros({batch_size, m}, torch::CUDA(torch::kFloat)); - torch::Tensor idx1 = torch::zeros({batch_size, n}, torch::CUDA(torch::kInt)); - torch::Tensor idx2 = torch::zeros({batch_size, m}, torch::CUDA(torch::kInt)); - - chamfer_dist_kernel<<>>( - batch_size, n, xyz1.data_ptr(), m, xyz2.data_ptr(), - dist1.data_ptr(), idx1.data_ptr()); - chamfer_dist_kernel<<>>( - batch_size, m, xyz2.data_ptr(), n, xyz1.data_ptr(), - dist2.data_ptr(), idx2.data_ptr()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in chamfer_cuda_forward: %s\n", cudaGetErrorString(err)); - } - return {dist1, dist2, idx1, idx2}; -} - -__global__ void chamfer_dist_grad_kernel(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) { - for (int i = blockIdx.x; i < b; i += gridDim.x) { - for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; - j += blockDim.x * gridDim.y) { - float x1 = xyz1[(i * n + j) * 3 + 0]; - float y1 = xyz1[(i * n + j) * 3 + 1]; - float z1 = xyz1[(i * n + j) * 3 + 2]; - int j2 = idx1[i * n + j]; - float x2 = xyz2[(i * m + j2) * 3 + 0]; - float y2 = xyz2[(i * m + j2) * 3 + 1]; - float z2 = xyz2[(i * m + j2) * 3 + 2]; - float g = grad_dist1[i * n + j] * 2; - atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 0]), g * (x1 - x2)); - atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 1]), g * (y1 - y2)); - atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 2]), g * (z1 - z2)); - atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 0]), -(g * (x1 - x2))); - atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 1]), -(g * (y1 - y2))); - atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 2]), -(g * (z1 - z2))); - } - } -} - -std::vector chamfer_cuda_backward(torch::Tensor xyz1, - torch::Tensor xyz2, - torch::Tensor idx1, - torch::Tensor idx2, - torch::Tensor grad_dist1, - torch::Tensor grad_dist2) { - const int batch_size = xyz1.size(0); - const int n = xyz1.size(1); // num_points point cloud A - const int m = xyz2.size(1); // num_points point cloud B - torch::Tensor grad_xyz1 = torch::zeros_like(xyz1, torch::CUDA(torch::kFloat)); - torch::Tensor grad_xyz2 = torch::zeros_like(xyz2, torch::CUDA(torch::kFloat)); - - chamfer_dist_grad_kernel<<>>( - batch_size, n, xyz1.data_ptr(), m, xyz2.data_ptr(), - grad_dist1.data_ptr(), idx1.data_ptr(), - grad_xyz1.data_ptr(), grad_xyz2.data_ptr()); - chamfer_dist_grad_kernel<<>>( - batch_size, m, xyz2.data_ptr(), n, xyz1.data_ptr(), - grad_dist2.data_ptr(), idx2.data_ptr(), - grad_xyz2.data_ptr(), grad_xyz1.data_ptr()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in chamfer_cuda_backward: %s\n", cudaGetErrorString(err)); - } - return {grad_xyz1, grad_xyz2}; -} diff --git a/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/PKG-INFO b/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/PKG-INFO deleted file mode 100644 index 7ba0882..0000000 --- a/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 1.0 -Name: chamfer -Version: 2.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -Author: UNKNOWN -Author-email: UNKNOWN -License: UNKNOWN -Description: UNKNOWN -Platform: UNKNOWN diff --git a/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/SOURCES.txt b/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/SOURCES.txt deleted file mode 100644 index 0e582cb..0000000 --- a/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -chamfer.cu -chamfer_cuda.cpp -setup.py -chamfer.egg-info/PKG-INFO -chamfer.egg-info/SOURCES.txt -chamfer.egg-info/dependency_links.txt -chamfer.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/dependency_links.txt b/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/top_level.txt b/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/top_level.txt deleted file mode 100644 index 27aeea6..0000000 --- a/other_models/GRNet/extensions/chamfer_dist/chamfer.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -chamfer diff --git a/other_models/GRNet/extensions/chamfer_dist/chamfer_cuda.cpp b/other_models/GRNet/extensions/chamfer_dist/chamfer_cuda.cpp deleted file mode 100644 index 9fca161..0000000 --- a/other_models/GRNet/extensions/chamfer_dist/chamfer_cuda.cpp +++ /dev/null @@ -1,39 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-08-07 20:54:24 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2019-12-10 10:33:50 - * @Email: cshzxie@gmail.com - */ - -#include -#include - -std::vector chamfer_cuda_forward(torch::Tensor xyz1, - torch::Tensor xyz2); - -std::vector chamfer_cuda_backward(torch::Tensor xyz1, - torch::Tensor xyz2, - torch::Tensor idx1, - torch::Tensor idx2, - torch::Tensor grad_dist1, - torch::Tensor grad_dist2); - -std::vector chamfer_forward(torch::Tensor xyz1, - torch::Tensor xyz2) { - return chamfer_cuda_forward(xyz1, xyz2); -} - -std::vector chamfer_backward(torch::Tensor xyz1, - torch::Tensor xyz2, - torch::Tensor idx1, - torch::Tensor idx2, - torch::Tensor grad_dist1, - torch::Tensor grad_dist2) { - return chamfer_cuda_backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &chamfer_forward, "Chamfer forward (CUDA)"); - m.def("backward", &chamfer_backward, "Chamfer backward (CUDA)"); -} diff --git a/other_models/GRNet/extensions/chamfer_dist/setup.py b/other_models/GRNet/extensions/chamfer_dist/setup.py deleted file mode 100644 index 04c6589..0000000 --- a/other_models/GRNet/extensions/chamfer_dist/setup.py +++ /dev/null @@ -1,19 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-08-07 20:54:24 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-10 10:04:25 -# @Email: cshzxie@gmail.com - -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup(name='chamfer', - version='2.0.0', - ext_modules=[ - CUDAExtension('chamfer', [ - 'chamfer_cuda.cpp', - 'chamfer.cu', - ]), - ], - cmdclass={'build_ext': BuildExtension}) diff --git a/other_models/GRNet/extensions/chamfer_dist/test.py b/other_models/GRNet/extensions/chamfer_dist/test.py deleted file mode 100644 index edf53ba..0000000 --- a/other_models/GRNet/extensions/chamfer_dist/test.py +++ /dev/null @@ -1,32 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-10 10:38:01 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 14:21:36 -# @Email: cshzxie@gmail.com -# -# Note: -# - Replace float -> double, kFloat -> kDouble in chamfer.cu - -import os -import sys -import torch -import unittest - -from torch.autograd import gradcheck - -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) -from extensions.chamfer_dist import ChamferFunction - - -class ChamferDistanceTestCase(unittest.TestCase): - def test_chamfer_dist(self): - x = torch.rand(4, 64, 3).double() - y = torch.rand(4, 128, 3).double() - x.requires_grad = True - y.requires_grad = True - print(gradcheck(ChamferFunction.apply, [x.cuda(), y.cuda()])) - - -if __name__ == '__main__': - unittest.main() diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/__init__.py b/other_models/GRNet/extensions/cubic_feature_sampling/__init__.py deleted file mode 100644 index 835582e..0000000 --- a/other_models/GRNet/extensions/cubic_feature_sampling/__init__.py +++ /dev/null @@ -1,40 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-19 16:55:15 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 13:15:14 -# @Email: cshzxie@gmail.com - -import torch - -import cubic_feature_sampling - - -class CubicFeatureSamplingFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, ptcloud, cubic_features, neighborhood_size=1): - scale = cubic_features.size(2) - point_features, grid_pt_indexes = cubic_feature_sampling.forward(scale, neighborhood_size, ptcloud, - cubic_features) - ctx.save_for_backward(torch.Tensor([scale]), torch.Tensor([neighborhood_size]), grid_pt_indexes) - return point_features - - @staticmethod - def backward(ctx, grad_point_features): - scale, neighborhood_size, grid_pt_indexes = ctx.saved_tensors - scale = int(scale.item()) - neighborhood_size = int(neighborhood_size.item()) - grad_point_features = grad_point_features.contiguous() - grad_ptcloud, grad_cubic_features = cubic_feature_sampling.backward(scale, neighborhood_size, - grad_point_features, grid_pt_indexes) - return grad_ptcloud, grad_cubic_features, None - - -class CubicFeatureSampling(torch.nn.Module): - def __init__(self): - super(CubicFeatureSampling, self).__init__() - - def forward(self, ptcloud, cubic_features, neighborhood_size=1): - h_scale = cubic_features.size(2) / 2 - ptcloud = ptcloud * h_scale + h_scale - return CubicFeatureSamplingFunction.apply(ptcloud, cubic_features, neighborhood_size) diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/__pycache__/__init__.cpython-36.pyc b/other_models/GRNet/extensions/cubic_feature_sampling/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index 9203640..0000000 Binary files a/other_models/GRNet/extensions/cubic_feature_sampling/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.cu b/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.cu deleted file mode 100644 index 423e962..0000000 --- a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.cu +++ /dev/null @@ -1,205 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-12-19 20:36:36 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:55:41 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include -#include - -#define CUDA_NUM_THREADS 512 - -// Computer the number of threads needed in GPU -inline int get_n_threads(int n) { - const int pow_2 = std::log(static_cast(n)) / std::log(2.0); - return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); -} - -__device__ int compute_index(int offset_x, - int offset_y, - int offset_z, - int scale) { - return offset_x * scale * scale + offset_y * scale + offset_z; -} - -__global__ void cubic_feature_sampling_kernel( - int scale, - int neighborhood_size, - int n_vertices, - int n_pts, - int n_cubic_channels, - const float *__restrict__ ptcloud, - const float *__restrict__ cubic_features, - float *__restrict__ point_features, - int *__restrict__ grid_pt_indexes) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - int cub_scale = scale * scale * scale; - - ptcloud += batch_index * n_pts * 3; - cubic_features += batch_index * n_cubic_channels * cub_scale; - point_features += batch_index * n_pts * n_vertices * n_cubic_channels; - grid_pt_indexes += batch_index * n_pts * n_vertices; - - for (int i = index; i < n_pts; i += stride) { - float pt_x = ptcloud[i * 3 + 0]; - float pt_y = ptcloud[i * 3 + 1]; - float pt_z = ptcloud[i * 3 + 2]; - - int lower_x = std::floor(pt_x); - int upper_x = std::ceil(pt_x); - if (lower_x == upper_x) { - upper_x += 1; - } - int lower_y = std::floor(pt_y); - int upper_y = std::ceil(pt_y); - if (lower_y == upper_y) { - upper_y += 1; - } - int lower_z = std::floor(pt_z); - int upper_z = std::ceil(pt_z); - if (lower_z == upper_z) { - upper_z += 1; - } - - int ns = neighborhood_size - 1; - int vertex_idx = 0; - for (int j = lower_x - ns; j <= upper_x + ns; ++j) { - for (int k = lower_y - ns; k <= upper_y + ns; ++k) { - for (int m = lower_z - ns; m <= upper_z + ns; ++m) { - if (j < 0 || j >= scale || k < 0 || k >= scale || m < 0 || - m >= scale) { - // Ignore points lies out of the grid - grid_pt_indexes[i * n_vertices + vertex_idx++] = -1; - } else { - // Calcuating indexes for adjacent vertices - grid_pt_indexes[i * n_vertices + vertex_idx++] = - compute_index(j, k, m, scale); - } - } - } - } - - // Gather Features - for (int j = 0; j < n_vertices; ++j) { - for (int k = 0; k < n_cubic_channels; ++k) { - int vertex_idx = grid_pt_indexes[i * n_vertices + j]; - if (vertex_idx == -1) { - continue; - } - int feature_idx = - i * n_vertices * n_cubic_channels + j * n_cubic_channels + k; - float feature_val = cubic_features[k * cub_scale + vertex_idx]; - point_features[feature_idx] = feature_val; - } - } - } -} - -std::vector cubic_feature_sampling_cuda_forward( - int scale, - int neighborhood_size, - torch::Tensor ptcloud, - torch::Tensor cubic_features, - cudaStream_t stream) { - int batch_size = ptcloud.size(0); - int n_pts = ptcloud.size(1); - int n_cubic_channels = cubic_features.size(1); - - int n_vertices = std::pow(neighborhood_size * 2, 3); - torch::Tensor point_features = - torch::zeros({batch_size, n_pts, n_vertices, n_cubic_channels}, - torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_indexes = - torch::zeros({batch_size, n_pts, n_vertices}, torch::CUDA(torch::kInt)); - - cubic_feature_sampling_kernel<<>>( - scale, neighborhood_size, n_vertices, n_pts, n_cubic_channels, - ptcloud.data(), cubic_features.data(), - point_features.data(), grid_pt_indexes.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in cubic_feature_sampling_cuda_forward: %s\n", - cudaGetErrorString(err)); - } - return {point_features, grid_pt_indexes}; -} - -__global__ void cubic_feature_sampling_grad_kernel( - int scale, - int neighborhood_size, - int n_vertices, - int n_pts, - int n_cubic_channels, - const float *__restrict__ grad_point_features, - const int *__restrict__ grid_pt_indexes, - float *__restrict__ grad_ptcloud, - float *__restrict__ grad_cubic_features) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - int cub_scale = scale * scale * scale; - - grad_point_features += batch_index * n_pts * n_vertices * n_cubic_channels; - grid_pt_indexes += batch_index * n_pts * n_vertices; - grad_ptcloud += batch_index * n_pts * 3; - grad_cubic_features += batch_index * n_cubic_channels * cub_scale; - - for (int i = index; i < n_pts; i += stride) { - for (int j = 0; j < n_vertices; ++j) { - int vertex_idx = grid_pt_indexes[i * n_vertices + j]; - if (vertex_idx == -1) { - continue; - } - for (int k = 0; k < n_cubic_channels; ++k) { - int grad_idx = - i * n_vertices * n_cubic_channels + j * n_cubic_channels + k; - float grad_val = grad_point_features[grad_idx]; - // Fix bugs: the gradients of ceil and floor functions are zeros. - // Ref: https://github.com/tensorflow/tensorflow/issues/897 - // atomicAdd(&(grad_ptcloud[i * 3 + 0]), grad_val); - // atomicAdd(&(grad_ptcloud[i * 3 + 1]), grad_val); - // atomicAdd(&(grad_ptcloud[i * 3 + 2]), grad_val); - atomicAdd(&(grad_cubic_features[k * cub_scale + vertex_idx]), grad_val); - } - } - } -} - -std::vector cubic_feature_sampling_cuda_backward( - int scale, - int neighborhood_size, - torch::Tensor grad_point_features, - torch::Tensor grid_pt_indexes, - cudaStream_t stream) { - int batch_size = grad_point_features.size(0); - int n_cubic_channels = grad_point_features.size(3); - int n_pts = grid_pt_indexes.size(1); - int n_vertices = std::pow(neighborhood_size * 2, 3); - - torch::Tensor grad_ptcloud = - torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); - torch::Tensor grad_cubic_features = - torch::zeros({batch_size, n_cubic_channels, scale, scale, scale}, - torch::CUDA(torch::kFloat)); - - cubic_feature_sampling_grad_kernel<<>>( - scale, neighborhood_size, n_vertices, n_pts, n_cubic_channels, - grad_point_features.data(), grid_pt_indexes.data(), - grad_ptcloud.data(), grad_cubic_features.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in cubic_feature_sampling_cuda_backward: %s\n", - cudaGetErrorString(err)); - } - return {grad_ptcloud, grad_cubic_features}; -} diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/PKG-INFO b/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/PKG-INFO deleted file mode 100644 index e2bce51..0000000 --- a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: cubic-feature-sampling -Version: 1.1.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/SOURCES.txt b/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/SOURCES.txt deleted file mode 100644 index a795ed8..0000000 --- a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -cubic_feature_sampling.cu -cubic_feature_sampling_cuda.cpp -setup.py -cubic_feature_sampling.egg-info/PKG-INFO -cubic_feature_sampling.egg-info/SOURCES.txt -cubic_feature_sampling.egg-info/dependency_links.txt -cubic_feature_sampling.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/dependency_links.txt b/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/top_level.txt b/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/top_level.txt deleted file mode 100644 index 5870961..0000000 --- a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -cubic_feature_sampling diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp b/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp deleted file mode 100644 index 2bcbe7e..0000000 --- a/other_models/GRNet/extensions/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp +++ /dev/null @@ -1,65 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-12-19 17:04:38 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:50:22 - * @Email: cshzxie@gmail.com - */ - -#include -#include - -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) \ - AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -std::vector cubic_feature_sampling_cuda_forward( - int scale, - int neighborhood_size, - torch::Tensor ptcloud, - torch::Tensor cubic_features, - cudaStream_t stream); - -std::vector cubic_feature_sampling_cuda_backward( - int scale, - int neighborhood_size, - torch::Tensor grad_point_features, - torch::Tensor grid_pt_indexes, - cudaStream_t stream); - -std::vector cubic_feature_sampling_forward( - int scale, - int neighborhood_size, - torch::Tensor ptcloud, - torch::Tensor cubic_features) { - CHECK_INPUT(ptcloud); - CHECK_INPUT(cubic_features); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return cubic_feature_sampling_cuda_forward(scale, neighborhood_size, ptcloud, - cubic_features, stream); -} - -std::vector cubic_feature_sampling_backward( - int scale, - int neighborhood_size, - torch::Tensor grad_point_features, - torch::Tensor grid_pt_indexes) { - CHECK_INPUT(grad_point_features); - CHECK_INPUT(grid_pt_indexes); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return cubic_feature_sampling_cuda_backward( - scale, neighborhood_size, grad_point_features, grid_pt_indexes, stream); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &cubic_feature_sampling_forward, - "Cubic Feature Sampling forward (CUDA)"); - m.def("backward", &cubic_feature_sampling_backward, - "Cubic Feature Sampling backward (CUDA)"); -} \ No newline at end of file diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.6-linux-x86_64.egg b/other_models/GRNet/extensions/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.6-linux-x86_64.egg deleted file mode 100644 index a2c5cd8..0000000 Binary files a/other_models/GRNet/extensions/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/setup.py b/other_models/GRNet/extensions/cubic_feature_sampling/setup.py deleted file mode 100644 index 9a7783b..0000000 --- a/other_models/GRNet/extensions/cubic_feature_sampling/setup.py +++ /dev/null @@ -1,16 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-19 17:03:06 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 14:02:06 -# @Email: cshzxie@gmail.com - -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup(name='cubic_feature_sampling', - version='1.1.0', - ext_modules=[ - CUDAExtension('cubic_feature_sampling', ['cubic_feature_sampling_cuda.cpp', 'cubic_feature_sampling.cu']), - ], - cmdclass={'build_ext': BuildExtension}) diff --git a/other_models/GRNet/extensions/cubic_feature_sampling/test.py b/other_models/GRNet/extensions/cubic_feature_sampling/test.py deleted file mode 100644 index 7f51017..0000000 --- a/other_models/GRNet/extensions/cubic_feature_sampling/test.py +++ /dev/null @@ -1,52 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-20 11:50:50 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 13:52:33 -# @Email: cshzxie@gmail.com -# -# Note: -# - Replace float -> double, kFloat -> kDouble in cubic_feature_sampling.cu - -import os -import sys -import torch -import unittest - -from torch.autograd import gradcheck - -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) -from extensions.cubic_feature_sampling import CubicFeatureSamplingFunction - - -class CubicFeatureSamplingTestCase(unittest.TestCase): - def test_neighborhood_size_1(self): - ptcloud = torch.rand(2, 64, 3) * 2 - 1 - cubic_features = torch.rand(2, 4, 8, 8, 8) - ptcloud.requires_grad = True - cubic_features.requires_grad = True - self.assertTrue( - gradcheck(CubicFeatureSamplingFunction.apply, - [ptcloud.double().cuda(), cubic_features.double().cuda()])) - - def test_neighborhood_size_2(self): - ptcloud = torch.rand(2, 32, 3) * 2 - 1 - cubic_features = torch.rand(2, 2, 8, 8, 8) - ptcloud.requires_grad = True - cubic_features.requires_grad = True - self.assertTrue( - gradcheck(CubicFeatureSamplingFunction.apply, - [ptcloud.double().cuda(), cubic_features.double().cuda(), 2])) - - def test_neighborhood_size_3(self): - ptcloud = torch.rand(1, 32, 3) * 2 - 1 - cubic_features = torch.rand(1, 2, 16, 16, 16) - ptcloud.requires_grad = True - cubic_features.requires_grad = True - self.assertTrue( - gradcheck(CubicFeatureSamplingFunction.apply, - [ptcloud.double().cuda(), cubic_features.double().cuda(), 3])) - - -if __name__ == '__main__': - unittest.main() diff --git a/other_models/GRNet/extensions/gridding/__init__.py b/other_models/GRNet/extensions/gridding/__init__.py deleted file mode 100644 index 4c7ff45..0000000 --- a/other_models/GRNet/extensions/gridding/__init__.py +++ /dev/null @@ -1,74 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-11-15 20:33:52 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-30 09:55:53 -# @Email: cshzxie@gmail.com - -import torch - -import gridding - - -class GriddingFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, scale, ptcloud): - grid, grid_pt_weights, grid_pt_indexes = gridding.forward(-scale, scale - 1, -scale, scale - 1, -scale, - scale - 1, ptcloud) - # print(grid.size()) # torch.Size(batch_size, n_grid_vertices) - # print(grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) - # print(grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) - ctx.save_for_backward(grid_pt_weights, grid_pt_indexes) - - return grid - - @staticmethod - def backward(ctx, grad_grid): - grid_pt_weights, grid_pt_indexes = ctx.saved_tensors - grad_ptcloud = gridding.backward(grid_pt_weights, grid_pt_indexes, grad_grid) - # print(grad_ptcloud.size()) # torch.Size(batch_size, n_pts, 3) - - return None, grad_ptcloud - - -class Gridding(torch.nn.Module): - def __init__(self, scale=1): - super(Gridding, self).__init__() - self.scale = scale // 2 - - def forward(self, ptcloud): - ptcloud = ptcloud * self.scale - _ptcloud = torch.split(ptcloud, 1, dim=0) - grids = [] - for p in _ptcloud: - non_zeros = torch.sum(p, dim=2).ne(0) - p = p[non_zeros].unsqueeze(dim=0) - grids.append(GriddingFunction.apply(self.scale, p)) - - return torch.cat(grids, dim=0).contiguous() - - -class GriddingReverseFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, scale, grid): - ptcloud = gridding.rev_forward(scale, grid) - ctx.save_for_backward(torch.Tensor([scale]), grid, ptcloud) - return ptcloud - - @staticmethod - def backward(ctx, grad_ptcloud): - scale, grid, ptcloud = ctx.saved_tensors - scale = int(scale.item()) - grad_grid = gridding.rev_backward(ptcloud, grid, grad_ptcloud) - grad_grid = grad_grid.view(-1, scale, scale, scale) - return None, grad_grid - - -class GriddingReverse(torch.nn.Module): - def __init__(self, scale=1): - super(GriddingReverse, self).__init__() - self.scale = scale - - def forward(self, grid): - ptcloud = GriddingReverseFunction.apply(self.scale, grid) - return ptcloud / self.scale * 2 diff --git a/other_models/GRNet/extensions/gridding/__pycache__/__init__.cpython-36.pyc b/other_models/GRNet/extensions/gridding/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index 70c07cd..0000000 Binary files a/other_models/GRNet/extensions/gridding/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/GRNet/extensions/gridding/dist/gridding-2.1.0-py3.6-linux-x86_64.egg b/other_models/GRNet/extensions/gridding/dist/gridding-2.1.0-py3.6-linux-x86_64.egg deleted file mode 100644 index 29834ab..0000000 Binary files a/other_models/GRNet/extensions/gridding/dist/gridding-2.1.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/other_models/GRNet/extensions/gridding/gridding.cu b/other_models/GRNet/extensions/gridding/gridding.cu deleted file mode 100644 index 2d60880..0000000 --- a/other_models/GRNet/extensions/gridding/gridding.cu +++ /dev/null @@ -1,335 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-11-13 10:53:22 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 15:00:10 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include -#include - -#define CUDA_NUM_THREADS 512 - -// Computer the number of threads needed in GPU -inline int get_n_threads(int n) { - const int pow_2 = std::log(static_cast(n)) / std::log(2.0); - return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); -} - -__device__ int compute_index( - int offset_x, int offset_y, int offset_z, int len_y, int len_z) { - return offset_x * len_y * len_z + offset_y * len_z + offset_z; -} - -__device__ float compute_weight(float x, float x0) { return 1 - abs(x - x0); } - -__global__ void gridding_kernel(int n_grid_vertices, - int n_pts, - float min_x, - float min_y, - float min_z, - int len_y, - int len_z, - const float *__restrict__ ptcloud, - float *__restrict__ grid_weights, - float *__restrict__ grid_pt_weights, - int *__restrict__ grid_pt_indexes) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - ptcloud += batch_index * n_pts * 3; - grid_weights += batch_index * n_grid_vertices; - grid_pt_weights += batch_index * n_pts * 24; - grid_pt_indexes += batch_index * n_pts * 8; - - for (int j = index; j < n_pts; j += stride) { - float pt_x = ptcloud[j * 3 + 0]; - float pt_y = ptcloud[j * 3 + 1]; - float pt_z = ptcloud[j * 3 + 2]; - - int lower_x = std::floor(pt_x); - int upper_x = std::ceil(pt_x); - if (lower_x == upper_x) { - upper_x += 1; - } - int lower_y = std::floor(pt_y); - int upper_y = std::ceil(pt_y); - if (lower_y == upper_y) { - upper_y += 1; - } - int lower_z = std::floor(pt_z); - int upper_z = std::ceil(pt_z); - if (lower_z == upper_z) { - upper_z += 1; - } - - int lx_offset = lower_x - min_x, ux_offset = upper_x - min_x; - int ly_offset = lower_y - min_y, uy_offset = upper_y - min_y; - int lz_offset = lower_z - min_z, uz_offset = upper_z - min_z; - - // Compute weights and corresponding positions, a loop for 8 points - // LLL -> Lower X, Lower Y, Lower Z - grid_pt_indexes[j * 8 + 0] = - compute_index(lx_offset, ly_offset, lz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 0] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 1] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 2] = compute_weight(pt_z, lower_z); - - // LLU -> Lower X, Lower Y, Upper Z - grid_pt_indexes[j * 8 + 1] = - compute_index(lx_offset, ly_offset, uz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 3] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 4] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 5] = compute_weight(pt_z, upper_z); - - // LUL -> Lower X, Upper Y, Lower Z - grid_pt_indexes[j * 8 + 2] = - compute_index(lx_offset, uy_offset, lz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 6] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 7] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 8] = compute_weight(pt_z, lower_z); - - // LUU -> Lower X, Upper Y, Upper Z - grid_pt_indexes[j * 8 + 3] = - compute_index(lx_offset, uy_offset, uz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 9] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 10] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 11] = compute_weight(pt_z, upper_z); - - // ULL -> Upper X, Lower Y, Lower Z - grid_pt_indexes[j * 8 + 4] = - compute_index(ux_offset, ly_offset, lz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 12] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 13] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 14] = compute_weight(pt_z, lower_z); - - // ULU -> Upper X, Lower Y, Upper Z - grid_pt_indexes[j * 8 + 5] = - compute_index(ux_offset, ly_offset, uz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 15] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 16] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 17] = compute_weight(pt_z, upper_z); - - // UUL -> Upper X, Upper Y, Lower Z - grid_pt_indexes[j * 8 + 6] = - compute_index(ux_offset, uy_offset, lz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 18] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 19] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 20] = compute_weight(pt_z, lower_z); - - // UUU -> Upper X, Upper Y, Upper Z - grid_pt_indexes[j * 8 + 7] = - compute_index(ux_offset, uy_offset, uz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 21] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 22] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 23] = compute_weight(pt_z, upper_z); - } - - __syncthreads(); - - int gvtx_idx = 0; - for (int j = index; j < n_pts; j += stride) { - // LLL -> Lower X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 0]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 0] * - grid_pt_weights[j * 24 + 1] * - grid_pt_weights[j * 24 + 2]); - // LLU -> Lower X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 1]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 3] * - grid_pt_weights[j * 24 + 4] * - grid_pt_weights[j * 24 + 5]); - // LUL -> Lower X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 2]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 6] * - grid_pt_weights[j * 24 + 7] * - grid_pt_weights[j * 24 + 8]); - // LUU -> Lower X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 3]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 9] * - grid_pt_weights[j * 24 + 10] * - grid_pt_weights[j * 24 + 11]); - // ULL -> Upper X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 4]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 12] * - grid_pt_weights[j * 24 + 13] * - grid_pt_weights[j * 24 + 14]); - // ULU -> Upper X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 5]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 15] * - grid_pt_weights[j * 24 + 16] * - grid_pt_weights[j * 24 + 17]); - // UUL -> Upper X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 6]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 18] * - grid_pt_weights[j * 24 + 19] * - grid_pt_weights[j * 24 + 20]); - // UUU -> Upper X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 7]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 21] * - grid_pt_weights[j * 24 + 22] * - grid_pt_weights[j * 24 + 23]); - } -} - -std::vector gridding_cuda_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud, - cudaStream_t stream) { - int batch_size = ptcloud.size(0); - int n_pts = ptcloud.size(1); - int len_x = max_x - min_x + 1; - int len_y = max_y - min_y + 1; - int len_z = max_z - min_z + 1; - int n_grid_vertices = len_x * len_y * len_z; - - torch::Tensor grid_weights = - torch::zeros({batch_size, n_grid_vertices}, torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_weights = - torch::zeros({batch_size, n_pts, 8, 3}, torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_indexes = - torch::zeros({batch_size, n_pts, 8}, torch::CUDA(torch::kInt)); - - gridding_kernel<<>>( - n_grid_vertices, n_pts, min_x, min_y, min_z, len_y, len_z, - ptcloud.data(), grid_weights.data(), - grid_pt_weights.data(), grid_pt_indexes.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_cuda_forward: %s\n", cudaGetErrorString(err)); - } - return {grid_weights, grid_pt_weights, grid_pt_indexes}; -} - -__global__ void gridding_grad_kernel(int n_grid_vertices, - int n_pts, - const float *__restrict__ grid_pt_weights, - const int *__restrict__ grid_pt_indexes, - const float *__restrict__ grad_grid, - float *__restrict__ grad_ptcloud) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - grid_pt_weights += batch_index * n_pts * 24; - grid_pt_indexes += batch_index * n_pts * 8; - grad_grid += batch_index * n_grid_vertices; - grad_ptcloud += batch_index * n_pts * 3; - - int gvtx_idx = 0; - float grad_vtx = 0, x_weights = 0, y_weights = 0, z_weights = 0; - for (int j = index; j < n_pts; j += stride) { - // Compute gradient for the corresponding positions, a loop for 8 points - // LLL -> Lower X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 0]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 0]; - y_weights = grid_pt_weights[j * 24 + 1]; - z_weights = grid_pt_weights[j * 24 + 2]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // LLU -> Lower X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 1]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 3]; - y_weights = grid_pt_weights[j * 24 + 4]; - z_weights = grid_pt_weights[j * 24 + 5]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // LUL -> Lower X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 2]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 6]; - y_weights = grid_pt_weights[j * 24 + 7]; - z_weights = grid_pt_weights[j * 24 + 8]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // LUU -> Lower X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 3]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 9]; - y_weights = grid_pt_weights[j * 24 + 10]; - z_weights = grid_pt_weights[j * 24 + 11]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // ULL -> Upper X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 4]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 12]; - y_weights = grid_pt_weights[j * 24 + 13]; - z_weights = grid_pt_weights[j * 24 + 14]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // ULU -> Upper X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 5]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 15]; - y_weights = grid_pt_weights[j * 24 + 16]; - z_weights = grid_pt_weights[j * 24 + 17]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // UUL -> Upper X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 6]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 18]; - y_weights = grid_pt_weights[j * 24 + 19]; - z_weights = grid_pt_weights[j * 24 + 20]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // UUU -> Upper X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 7]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 21]; - y_weights = grid_pt_weights[j * 24 + 22]; - z_weights = grid_pt_weights[j * 24 + 23]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - } -} - -torch::Tensor gridding_cuda_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid, - cudaStream_t stream) { - int batch_size = grad_grid.size(0); - int n_grid_vertices = grad_grid.size(1); - int n_pts = grid_pt_indexes.size(1); - - torch::Tensor grad_ptcloud = - torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); - - gridding_grad_kernel<<>>( - n_grid_vertices, n_pts, grid_pt_weights.data(), - grid_pt_indexes.data(), grad_grid.data(), - grad_ptcloud.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_cuda_backward: %s\n", cudaGetErrorString(err)); - } - return grad_ptcloud; -} diff --git a/other_models/GRNet/extensions/gridding/gridding.egg-info/PKG-INFO b/other_models/GRNet/extensions/gridding/gridding.egg-info/PKG-INFO deleted file mode 100644 index 72f20bd..0000000 --- a/other_models/GRNet/extensions/gridding/gridding.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: gridding -Version: 2.1.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/GRNet/extensions/gridding/gridding.egg-info/SOURCES.txt b/other_models/GRNet/extensions/gridding/gridding.egg-info/SOURCES.txt deleted file mode 100644 index f241b60..0000000 --- a/other_models/GRNet/extensions/gridding/gridding.egg-info/SOURCES.txt +++ /dev/null @@ -1,8 +0,0 @@ -gridding.cu -gridding_cuda.cpp -gridding_reverse.cu -setup.py -gridding.egg-info/PKG-INFO -gridding.egg-info/SOURCES.txt -gridding.egg-info/dependency_links.txt -gridding.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/GRNet/extensions/gridding/gridding.egg-info/dependency_links.txt b/other_models/GRNet/extensions/gridding/gridding.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/GRNet/extensions/gridding/gridding.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/GRNet/extensions/gridding/gridding.egg-info/top_level.txt b/other_models/GRNet/extensions/gridding/gridding.egg-info/top_level.txt deleted file mode 100644 index 7ee56cc..0000000 --- a/other_models/GRNet/extensions/gridding/gridding.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -gridding diff --git a/other_models/GRNet/extensions/gridding/gridding_cuda.cpp b/other_models/GRNet/extensions/gridding/gridding_cuda.cpp deleted file mode 100644 index e4201c0..0000000 --- a/other_models/GRNet/extensions/gridding/gridding_cuda.cpp +++ /dev/null @@ -1,94 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-11-13 10:52:53 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:52:32 - * @Email: cshzxie@gmail.com - */ - -#include -#include - -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) \ - AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -std::vector gridding_cuda_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud, - cudaStream_t stream); - -torch::Tensor gridding_cuda_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid, - cudaStream_t stream); - -torch::Tensor gridding_reverse_cuda_forward(int scale, - torch::Tensor grid, - cudaStream_t stream); - -torch::Tensor gridding_reverse_cuda_backward(torch::Tensor ptcloud, - torch::Tensor grid, - torch::Tensor grad_ptcloud, - cudaStream_t stream); - -std::vector gridding_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud) { - CHECK_INPUT(ptcloud); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_cuda_forward(min_x, max_x, min_y, max_y, min_z, max_z, - ptcloud, stream); -} - -torch::Tensor gridding_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid) { - CHECK_INPUT(grid_pt_weights); - CHECK_INPUT(grid_pt_indexes); - CHECK_INPUT(grad_grid); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_cuda_backward(grid_pt_weights, grid_pt_indexes, grad_grid, - stream); -} - -torch::Tensor gridding_reverse_forward(int scale, torch::Tensor grid) { - CHECK_INPUT(grid); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_reverse_cuda_forward(scale, grid, stream); -} - -torch::Tensor gridding_reverse_backward(torch::Tensor ptcloud, - torch::Tensor grid, - torch::Tensor grad_ptcloud) { - CHECK_INPUT(ptcloud); - CHECK_INPUT(grid); - CHECK_INPUT(grad_ptcloud); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_reverse_cuda_backward(ptcloud, grid, grad_ptcloud, stream); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &gridding_forward, "Gridding forward (CUDA)"); - m.def("backward", &gridding_backward, "Gridding backward (CUDA)"); - m.def("rev_forward", &gridding_reverse_forward, - "Gridding Reverse forward (CUDA)"); - m.def("rev_backward", &gridding_reverse_backward, - "Gridding Reverse backward (CUDA)"); -} diff --git a/other_models/GRNet/extensions/gridding/gridding_reverse.cu b/other_models/GRNet/extensions/gridding/gridding_reverse.cu deleted file mode 100644 index 63b594c..0000000 --- a/other_models/GRNet/extensions/gridding/gridding_reverse.cu +++ /dev/null @@ -1,236 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-11-21 16:42:18 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 15:00:21 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include -#include - -#define CUDA_NUM_THREADS 512 -#define EPS 1e-6 - -// Computer the number of threads needed in GPU -inline int get_n_threads(int n) { - const int pow_2 = std::log(static_cast(n)) / std::log(2.0); - return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); -} - -__device__ int compute_index(int offset_x, - int offset_y, - int offset_z, - int scale) { - return offset_x * scale * scale + offset_y * scale + offset_z; -} - -__global__ void gridding_reverse_kernel(int scale, - int n_pts, - const float *__restrict__ grid, - float *__restrict__ ptcloud) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - ptcloud += batch_index * n_pts * 3; - grid += batch_index * n_pts; - - for (int j = index; j < n_pts; j += stride) { - int sqr_scale = scale * scale; - int x_offset = j / sqr_scale; - int y_offset = j % sqr_scale / scale; - int z_offset = j % sqr_scale % scale; - if (x_offset == 0 || y_offset == 0 || z_offset == 0) { - continue; - } - - // assert j == compute_index(x_offset, y_offset, z_offset, scale) - float weights[8] = { - grid[compute_index(x_offset - 1, y_offset - 1, z_offset - 1, scale)], - grid[compute_index(x_offset - 1, y_offset - 1, z_offset, scale)], - grid[compute_index(x_offset - 1, y_offset, z_offset - 1, scale)], - grid[compute_index(x_offset - 1, y_offset, z_offset, scale)], - grid[compute_index(x_offset, y_offset - 1, z_offset - 1, scale)], - grid[compute_index(x_offset, y_offset - 1, z_offset, scale)], - grid[compute_index(x_offset, y_offset, z_offset - 1, scale)], - grid[j]}; - - float weights_sum = 0; - for (size_t i = 0; i < 8; ++i) { - weights_sum += weights[i]; - } - if (weights_sum < EPS) { - continue; - } - for (size_t i = 0; i < 8; ++i) { - weights[i] /= weights_sum; - } - - x_offset -= scale / 2; - y_offset -= scale / 2; - z_offset -= scale / 2; - - // clang-format off - ptcloud[j * 3 + 0] = weights[0] * (x_offset - 1) + - weights[1] * (x_offset - 1) + - weights[2] * (x_offset - 1) + - weights[3] * (x_offset - 1) + - weights[4] * x_offset + - weights[5] * x_offset + - weights[6] * x_offset + - weights[7] * x_offset; - ptcloud[j * 3 + 1] = weights[0] * (y_offset - 1) + - weights[1] * (y_offset - 1) + - weights[2] * y_offset + - weights[3] * y_offset + - weights[4] * (y_offset - 1) + - weights[5] * (y_offset - 1) + - weights[6] * y_offset + - weights[7] * y_offset; - ptcloud[j * 3 + 2] = weights[0] * (z_offset - 1) + - weights[1] * z_offset + - weights[2] * (z_offset - 1) + - weights[3] * z_offset + - weights[4] * (z_offset - 1) + - weights[5] * z_offset + - weights[6] * (z_offset - 1) + - weights[7] * z_offset; - // clang-format on - } -} - -torch::Tensor gridding_reverse_cuda_forward(int scale, - torch::Tensor grid, - cudaStream_t stream) { - int batch_size = grid.size(0); - int n_pts = scale * scale * scale; - - torch::Tensor ptcloud = - torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); - - gridding_reverse_kernel<<>>( - scale, n_pts, grid.data(), ptcloud.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_cuda_forward: %s\n", cudaGetErrorString(err)); - } - return ptcloud; -} - -__global__ void gridding_reverse_grad_kernel( - int scale, - int n_pts, - const float *__restrict__ ptcloud, - const float *__restrict__ grid, - const float *__restrict__ grad_ptcloud, - float *__restrict__ grad_grid) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - ptcloud += batch_index * n_pts * 3; - grid += batch_index * n_pts; - grad_ptcloud += batch_index * n_pts * 3; - grad_grid += batch_index * n_pts; - - for (int j = index; j < n_pts; j += stride) { - int sqr_scale = scale * scale; - int x_offset = j / sqr_scale; - int y_offset = j % sqr_scale / scale; - int z_offset = j % sqr_scale % scale; - if (x_offset == 0 || y_offset == 0 || z_offset == 0) { - continue; - } - - int gvtx_indexes[8] = { - compute_index(x_offset - 1, y_offset - 1, z_offset - 1, scale), - compute_index(x_offset - 1, y_offset - 1, z_offset, scale), - compute_index(x_offset - 1, y_offset, z_offset - 1, scale), - compute_index(x_offset - 1, y_offset, z_offset, scale), - compute_index(x_offset, y_offset - 1, z_offset - 1, scale), - compute_index(x_offset, y_offset - 1, z_offset, scale), - compute_index(x_offset, y_offset, z_offset - 1, scale), - j}; - float weights[8] = {grid[gvtx_indexes[0]], grid[gvtx_indexes[1]], - grid[gvtx_indexes[2]], grid[gvtx_indexes[3]], - grid[gvtx_indexes[4]], grid[gvtx_indexes[5]], - grid[gvtx_indexes[6]], grid[gvtx_indexes[7]]}; - - float weights_sum = 0; - for (size_t i = 0; i < 8; ++i) { - weights_sum += weights[i]; - } - - if (weights_sum < EPS) { - continue; - } - for (size_t i = 0; i < 8; ++i) { - weights[i] /= weights_sum; - } - - x_offset -= scale / 2; - y_offset -= scale / 2; - z_offset -= scale / 2; - - // clang-format off - atomicAdd(&(grad_grid[gvtx_indexes[0]]), - grad_ptcloud[j * 3 + 0] * ((x_offset - 1) - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * ((y_offset - 1) - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * ((z_offset - 1) - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[1]]), - grad_ptcloud[j * 3 + 0] * ((x_offset - 1) - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * ((y_offset - 1) - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * (z_offset - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[2]]), - grad_ptcloud[j * 3 + 0] * ((x_offset - 1) - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * (y_offset - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * ((z_offset - 1) - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[3]]), - grad_ptcloud[j * 3 + 0] * ((x_offset - 1) - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * (y_offset - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * (z_offset - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[4]]), - grad_ptcloud[j * 3 + 0] * (x_offset - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * ((y_offset - 1) - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * ((z_offset - 1) - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[5]]), - grad_ptcloud[j * 3 + 0] * (x_offset - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * ((y_offset - 1) - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * (z_offset - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[6]]), - grad_ptcloud[j * 3 + 0] * (x_offset - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * (y_offset - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * ((z_offset - 1) - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[7]]), - grad_ptcloud[j * 3 + 0] * (x_offset - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * (y_offset - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * (z_offset - ptcloud[j * 3 + 2]) / weights_sum); - // clang-format on - } -} - -torch::Tensor gridding_reverse_cuda_backward(torch::Tensor ptcloud, - torch::Tensor grid, - torch::Tensor grad_ptcloud, - cudaStream_t stream) { - int batch_size = ptcloud.size(0); - int n_pts = ptcloud.size(1); - int scale = static_cast(std::cbrt(n_pts)); - - torch::Tensor grad_grid = - torch::zeros({batch_size, n_pts}, torch::CUDA(torch::kFloat)); - - gridding_reverse_grad_kernel<<>>( - scale, n_pts, ptcloud.data(), grid.data(), - grad_ptcloud.data(), grad_grid.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_cuda_forward: %s\n", cudaGetErrorString(err)); - } - return grad_grid; -} diff --git a/other_models/GRNet/extensions/gridding/setup.py b/other_models/GRNet/extensions/gridding/setup.py deleted file mode 100644 index eaaa6ef..0000000 --- a/other_models/GRNet/extensions/gridding/setup.py +++ /dev/null @@ -1,16 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-11-13 10:51:33 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-02 17:02:16 -# @Email: cshzxie@gmail.com - -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup(name='gridding', - version='2.1.0', - ext_modules=[ - CUDAExtension('gridding', ['gridding_cuda.cpp', 'gridding.cu', 'gridding_reverse.cu']), - ], - cmdclass={'build_ext': BuildExtension}) diff --git a/other_models/GRNet/extensions/gridding/test.py b/other_models/GRNet/extensions/gridding/test.py deleted file mode 100644 index f47cf6b..0000000 --- a/other_models/GRNet/extensions/gridding/test.py +++ /dev/null @@ -1,45 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-10 10:48:55 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 14:20:42 -# @Email: cshzxie@gmail.com -# -# Note: -# - Replace float -> double, kFloat -> kDouble in gridding.cu and gridding_reverse.cu - -import os -import sys -import torch -import unittest - -from torch.autograd import gradcheck - -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) -from extensions.gridding import GriddingFunction, GriddingReverseFunction - - -class GriddingTestCase(unittest.TestCase): - def test_gridding_reverse_function_4(self): - x = torch.rand(2, 4, 4, 4) - x.requires_grad = True - self.assertTrue(gradcheck(GriddingReverseFunction.apply, [4, x.double().cuda()])) - - def test_gridding_reverse_function_8(self): - x = torch.rand(4, 8, 8, 8) - x.requires_grad = True - self.assertTrue(gradcheck(GriddingReverseFunction.apply, [8, x.double().cuda()])) - - def test_gridding_reverse_function_16(self): - x = torch.rand(1, 16, 16, 16) - x.requires_grad = True - self.assertTrue(gradcheck(GriddingReverseFunction.apply, [16, x.double().cuda()])) - - def test_gridding_function_32pts(self): - x = torch.rand(1, 32, 3) - x.requires_grad = True - self.assertTrue(gradcheck(GriddingFunction.apply, [x.double().cuda()])) - - -if __name__ == '__main__': - unittest.main() diff --git a/other_models/GRNet/extensions/gridding_loss/__init__.py b/other_models/GRNet/extensions/gridding_loss/__init__.py deleted file mode 100644 index c7d3ad7..0000000 --- a/other_models/GRNet/extensions/gridding_loss/__init__.py +++ /dev/null @@ -1,114 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-30 09:56:06 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2020-02-22 19:19:43 -# @Email: cshzxie@gmail.com - -import torch - -import gridding_distance - - -class GriddingDistanceFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud, gt_cloud): - pred_grid, pred_grid_pt_weights, pred_grid_pt_indexes = gridding_distance.forward( - min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud) - # print(pred_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) - # print(pred_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) - # print(pred_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) - gt_grid, gt_grid_pt_weights, gt_grid_pt_indexes = gridding_distance.forward( - min_x, max_x, min_y, max_y, min_z, max_z, gt_cloud) - # print(gt_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) - # print(gt_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) - # print(gt_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) - - ctx.save_for_backward(pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes) - return pred_grid, gt_grid - - @staticmethod - def backward(ctx, grad_pred_grid, grad_gt_grid): - pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes = ctx.saved_tensors - - grad_pred_cloud = gridding_distance.backward(pred_grid_pt_weights, pred_grid_pt_indexes, grad_pred_grid) - # print(grad_pred_cloud.size()) # torch.Size(batch_size, n_pts, 3) - grad_gt_cloud = gridding_distance.backward(gt_grid_pt_weights, gt_grid_pt_indexes, grad_gt_grid) - # print(grad_gt_cloud.size()) # torch.Size(batch_size, n_pts, 3) - - return None, None, None, None, None, None, grad_pred_cloud, grad_gt_cloud - - -class GriddingDistance(torch.nn.Module): - def __init__(self, scale=1): - super(GriddingDistance, self).__init__() - self.scale = scale - - def forward(self, pred_cloud, gt_cloud): - ''' - pred_cloud(b, n_pts1, 3) - gt_cloud(b, n_pts2, 3) - ''' - pred_cloud = pred_cloud * self.scale / 2 - gt_cloud = gt_cloud * self.scale / 2 - - min_pred_x = torch.min(pred_cloud[:, :, 0]) - max_pred_x = torch.max(pred_cloud[:, :, 0]) - min_pred_y = torch.min(pred_cloud[:, :, 1]) - max_pred_y = torch.max(pred_cloud[:, :, 1]) - min_pred_z = torch.min(pred_cloud[:, :, 2]) - max_pred_z = torch.max(pred_cloud[:, :, 2]) - - min_gt_x = torch.min(gt_cloud[:, :, 0]) - max_gt_x = torch.max(gt_cloud[:, :, 0]) - min_gt_y = torch.min(gt_cloud[:, :, 1]) - max_gt_y = torch.max(gt_cloud[:, :, 1]) - min_gt_z = torch.min(gt_cloud[:, :, 2]) - max_gt_z = torch.max(gt_cloud[:, :, 2]) - - min_x = torch.floor(torch.min(min_pred_x, min_gt_x)) - 1 - max_x = torch.ceil(torch.max(max_pred_x, max_gt_x)) + 1 - min_y = torch.floor(torch.min(min_pred_y, min_gt_y)) - 1 - max_y = torch.ceil(torch.max(max_pred_y, max_gt_y)) + 1 - min_z = torch.floor(torch.min(min_pred_z, min_gt_z)) - 1 - max_z = torch.ceil(torch.max(max_pred_z, max_gt_z)) + 1 - - _pred_clouds = torch.split(pred_cloud, 1, dim=0) - _gt_clouds = torch.split(gt_cloud, 1, dim=0) - pred_grids = [] - gt_grids = [] - for pc, gc in zip(_pred_clouds, _gt_clouds): - non_zeros = torch.sum(pc, dim=2).ne(0) - pc = pc[non_zeros].unsqueeze(dim=0) - non_zeros = torch.sum(gc, dim=2).ne(0) - gc = gc[non_zeros].unsqueeze(dim=0) - pred_grid, gt_grid = GriddingDistanceFunction.apply(min_x, max_x, min_y, max_y, min_z, max_z, pc, gc) - pred_grids.append(pred_grid) - gt_grids.append(gt_grid) - - return torch.cat(pred_grids, dim=0).contiguous(), torch.cat(gt_grids, dim=0).contiguous() - - -class GriddingLoss(torch.nn.Module): - def __init__(self, scales=[], alphas=[]): - super(GriddingLoss, self).__init__() - self.scales = scales - self.alphas = alphas - self.gridding_dists = [GriddingDistance(scale=s) for s in scales] - self.l1_loss = torch.nn.L1Loss() - - def forward(self, pred_cloud, gt_cloud): - gridding_loss = None - n_dists = len(self.scales) - - for i in range(n_dists): - alpha = self.alphas[i] - gdist = self.gridding_dists[i] - pred_grid, gt_grid = gdist(pred_cloud, gt_cloud) - - if gridding_loss is None: - gridding_loss = alpha * self.l1_loss(pred_grid, gt_grid) - else: - gridding_loss += alpha * self.l1_loss(pred_grid, gt_grid) - - return gridding_loss diff --git a/other_models/GRNet/extensions/gridding_loss/__pycache__/__init__.cpython-36.pyc b/other_models/GRNet/extensions/gridding_loss/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index e932c3e..0000000 Binary files a/other_models/GRNet/extensions/gridding_loss/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/GRNet/extensions/gridding_loss/dist/gridding_distance-1.0.0-py3.6-linux-x86_64.egg b/other_models/GRNet/extensions/gridding_loss/dist/gridding_distance-1.0.0-py3.6-linux-x86_64.egg deleted file mode 100644 index 0187bfa..0000000 Binary files a/other_models/GRNet/extensions/gridding_loss/dist/gridding_distance-1.0.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/other_models/GRNet/extensions/gridding_loss/gridding_distance.cu b/other_models/GRNet/extensions/gridding_loss/gridding_distance.cu deleted file mode 100644 index f0a9b0f..0000000 --- a/other_models/GRNet/extensions/gridding_loss/gridding_distance.cu +++ /dev/null @@ -1,338 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-12-30 11:35:30 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 15:00:45 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include -#include - -#define CUDA_NUM_THREADS 512 - -// Computer the number of threads needed in GPU -inline int get_n_threads(int n) { - const int pow_2 = std::log(static_cast(n)) / std::log(2.0); - return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); -} - -__device__ int compute_index( - int offset_x, int offset_y, int offset_z, int len_y, int len_z) { - return offset_x * len_y * len_z + offset_y * len_z + offset_z; -} - -__device__ float compute_weight(float x, float x0) { return 1 - abs(x - x0); } - -__global__ void gridding_dist_kernel(int n_grid_vertices, - int n_pts, - float min_x, - float min_y, - float min_z, - int len_y, - int len_z, - const float *__restrict__ ptcloud, - float *__restrict__ grid_weights, - float *__restrict__ grid_pt_weights, - int *__restrict__ grid_pt_indexes) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - ptcloud += batch_index * n_pts * 3; - grid_weights += batch_index * n_grid_vertices * 8; - grid_pt_weights += batch_index * n_pts * 24; - grid_pt_indexes += batch_index * n_pts * 8; - - for (int j = index; j < n_pts; j += stride) { - float pt_x = ptcloud[j * 3 + 0]; - float pt_y = ptcloud[j * 3 + 1]; - float pt_z = ptcloud[j * 3 + 2]; - - int lower_x = std::floor(pt_x); - int upper_x = std::ceil(pt_x); - if (lower_x == upper_x) { - upper_x += 1; - } - int lower_y = std::floor(pt_y); - int upper_y = std::ceil(pt_y); - if (lower_y == upper_y) { - upper_y += 1; - } - int lower_z = std::floor(pt_z); - int upper_z = std::ceil(pt_z); - if (lower_z == upper_z) { - upper_z += 1; - } - - int lx_offset = lower_x - min_x, ux_offset = upper_x - min_x; - int ly_offset = lower_y - min_y, uy_offset = upper_y - min_y; - int lz_offset = lower_z - min_z, uz_offset = upper_z - min_z; - - // Compute weights and corresponding positions, a loop for 8 points - // LLL -> Lower X, Lower Y, Lower Z - grid_pt_indexes[j * 8 + 0] = - compute_index(lx_offset, ly_offset, lz_offset, len_y, len_z) * 8 + 0; - grid_pt_weights[j * 24 + 0] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 1] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 2] = compute_weight(pt_z, lower_z); - - // LLU -> Lower X, Lower Y, Upper Z - grid_pt_indexes[j * 8 + 1] = - compute_index(lx_offset, ly_offset, uz_offset, len_y, len_z) * 8 + 1; - grid_pt_weights[j * 24 + 3] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 4] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 5] = compute_weight(pt_z, upper_z); - - // LUL -> Lower X, Upper Y, Lower Z - grid_pt_indexes[j * 8 + 2] = - compute_index(lx_offset, uy_offset, lz_offset, len_y, len_z) * 8 + 2; - grid_pt_weights[j * 24 + 6] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 7] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 8] = compute_weight(pt_z, lower_z); - - // LUU -> Lower X, Upper Y, Upper Z - grid_pt_indexes[j * 8 + 3] = - compute_index(lx_offset, uy_offset, uz_offset, len_y, len_z) * 8 + 3; - grid_pt_weights[j * 24 + 9] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 10] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 11] = compute_weight(pt_z, upper_z); - - // ULL -> Upper X, Lower Y, Lower Z - grid_pt_indexes[j * 8 + 4] = - compute_index(ux_offset, ly_offset, lz_offset, len_y, len_z) * 8 + 4; - grid_pt_weights[j * 24 + 12] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 13] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 14] = compute_weight(pt_z, lower_z); - - // ULU -> Upper X, Lower Y, Upper Z - grid_pt_indexes[j * 8 + 5] = - compute_index(ux_offset, ly_offset, uz_offset, len_y, len_z) * 8 + 5; - grid_pt_weights[j * 24 + 15] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 16] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 17] = compute_weight(pt_z, upper_z); - - // UUL -> Upper X, Upper Y, Lower Z - grid_pt_indexes[j * 8 + 6] = - compute_index(ux_offset, uy_offset, lz_offset, len_y, len_z) * 8 + 6; - grid_pt_weights[j * 24 + 18] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 19] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 20] = compute_weight(pt_z, lower_z); - - // UUU -> Upper X, Upper Y, Upper Z - grid_pt_indexes[j * 8 + 7] = - compute_index(ux_offset, uy_offset, uz_offset, len_y, len_z) * 8 + 7; - grid_pt_weights[j * 24 + 21] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 22] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 23] = compute_weight(pt_z, upper_z); - } - - __syncthreads(); - - int gvtx_idx = 0; - for (int j = index; j < n_pts; j += stride) { - // LLL -> Lower X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 0]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 0] * - grid_pt_weights[j * 24 + 1] * - grid_pt_weights[j * 24 + 2]); - // LLU -> Lower X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 1]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 3] * - grid_pt_weights[j * 24 + 4] * - grid_pt_weights[j * 24 + 5]); - // LUL -> Lower X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 2]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 6] * - grid_pt_weights[j * 24 + 7] * - grid_pt_weights[j * 24 + 8]); - // LUU -> Lower X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 3]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 9] * - grid_pt_weights[j * 24 + 10] * - grid_pt_weights[j * 24 + 11]); - // ULL -> Upper X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 4]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 12] * - grid_pt_weights[j * 24 + 13] * - grid_pt_weights[j * 24 + 14]); - // ULU -> Upper X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 5]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 15] * - grid_pt_weights[j * 24 + 16] * - grid_pt_weights[j * 24 + 17]); - // UUL -> Upper X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 6]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 18] * - grid_pt_weights[j * 24 + 19] * - grid_pt_weights[j * 24 + 20]); - // UUU -> Upper X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 7]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 21] * - grid_pt_weights[j * 24 + 22] * - grid_pt_weights[j * 24 + 23]); - } -} - -std::vector gridding_distance_cuda_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud, - cudaStream_t stream) { - int batch_size = ptcloud.size(0); - int n_pts = ptcloud.size(1); - int len_x = max_x - min_x + 1; - int len_y = max_y - min_y + 1; - int len_z = max_z - min_z + 1; - int n_grid_vertices = len_x * len_y * len_z; - - torch::Tensor grid_weights = - torch::zeros({batch_size, n_grid_vertices, 8}, torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_weights = - torch::zeros({batch_size, n_pts, 8, 3}, torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_indexes = - torch::zeros({batch_size, n_pts, 8}, torch::CUDA(torch::kInt)); - - gridding_dist_kernel<<>>( - n_grid_vertices, n_pts, min_x, min_y, min_z, len_y, len_z, - ptcloud.data(), grid_weights.data(), - grid_pt_weights.data(), grid_pt_indexes.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_distance_cuda_forward: %s\n", - cudaGetErrorString(err)); - } - return {grid_weights, grid_pt_weights, grid_pt_indexes}; -} - -__global__ void gridding_dist_grad_kernel( - int n_grid_vertices, - int n_pts, - const float *__restrict__ grid_pt_weights, - const int *__restrict__ grid_pt_indexes, - const float *__restrict__ grad_grid, - float *__restrict__ grad_ptcloud) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - grid_pt_weights += batch_index * n_pts * 24; - grid_pt_indexes += batch_index * n_pts * 8; - grad_grid += batch_index * n_grid_vertices * 8; - grad_ptcloud += batch_index * n_pts * 3; - - int gvtx_idx = 0; - float grad_vtx = 0, x_weights = 0, y_weights = 0, z_weights = 0; - for (int j = index; j < n_pts; j += stride) { - // Compute gradient for the corresponding positions, a loop for 8 points - // LLL -> Lower X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 0]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 0]; - y_weights = grid_pt_weights[j * 24 + 1]; - z_weights = grid_pt_weights[j * 24 + 2]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // LLU -> Lower X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 1]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 3]; - y_weights = grid_pt_weights[j * 24 + 4]; - z_weights = grid_pt_weights[j * 24 + 5]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // LUL -> Lower X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 2]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 6]; - y_weights = grid_pt_weights[j * 24 + 7]; - z_weights = grid_pt_weights[j * 24 + 8]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // LUU -> Lower X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 3]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 9]; - y_weights = grid_pt_weights[j * 24 + 10]; - z_weights = grid_pt_weights[j * 24 + 11]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // ULL -> Upper X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 4]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 12]; - y_weights = grid_pt_weights[j * 24 + 13]; - z_weights = grid_pt_weights[j * 24 + 14]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // ULU -> Upper X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 5]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 15]; - y_weights = grid_pt_weights[j * 24 + 16]; - z_weights = grid_pt_weights[j * 24 + 17]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // UUL -> Upper X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 6]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 18]; - y_weights = grid_pt_weights[j * 24 + 19]; - z_weights = grid_pt_weights[j * 24 + 20]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // UUU -> Upper X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 7]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 21]; - y_weights = grid_pt_weights[j * 24 + 22]; - z_weights = grid_pt_weights[j * 24 + 23]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - } -} - -torch::Tensor gridding_distance_cuda_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid, - cudaStream_t stream) { - int batch_size = grad_grid.size(0); - int n_grid_vertices = grad_grid.size(1); - int n_pts = grid_pt_indexes.size(1); - - torch::Tensor grad_ptcloud = - torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); - - gridding_dist_grad_kernel<<>>( - n_grid_vertices, n_pts, grid_pt_weights.data(), - grid_pt_indexes.data(), grad_grid.data(), - grad_ptcloud.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_distance_cuda_backward: %s\n", - cudaGetErrorString(err)); - } - return grad_ptcloud; -} diff --git a/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/PKG-INFO b/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/PKG-INFO deleted file mode 100644 index 5453aab..0000000 --- a/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: gridding-distance -Version: 1.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/SOURCES.txt b/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/SOURCES.txt deleted file mode 100644 index ad3fcfb..0000000 --- a/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -gridding_distance.cu -gridding_distance_cuda.cpp -setup.py -gridding_distance.egg-info/PKG-INFO -gridding_distance.egg-info/SOURCES.txt -gridding_distance.egg-info/dependency_links.txt -gridding_distance.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/dependency_links.txt b/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/top_level.txt b/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/top_level.txt deleted file mode 100644 index cdcb49c..0000000 --- a/other_models/GRNet/extensions/gridding_loss/gridding_distance.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -gridding_distance diff --git a/other_models/GRNet/extensions/gridding_loss/gridding_distance_cuda.cpp b/other_models/GRNet/extensions/gridding_loss/gridding_distance_cuda.cpp deleted file mode 100644 index 889b5a9..0000000 --- a/other_models/GRNet/extensions/gridding_loss/gridding_distance_cuda.cpp +++ /dev/null @@ -1,65 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-12-30 10:59:31 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:52:52 - * @Email: cshzxie@gmail.com - */ - -#include -#include - -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) \ - AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -std::vector gridding_distance_cuda_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud, - cudaStream_t stream); - -torch::Tensor gridding_distance_cuda_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid, - cudaStream_t stream); - -std::vector gridding_distance_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud) { - CHECK_INPUT(ptcloud); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_distance_cuda_forward(min_x, max_x, min_y, max_y, min_z, - max_z, ptcloud, stream); -} - -torch::Tensor gridding_distance_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid) { - CHECK_INPUT(grid_pt_weights); - CHECK_INPUT(grid_pt_indexes); - CHECK_INPUT(grad_grid); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_distance_cuda_backward(grid_pt_weights, grid_pt_indexes, - grad_grid, stream); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &gridding_distance_forward, - "Gridding Distance Forward (CUDA)"); - m.def("backward", &gridding_distance_backward, - "Gridding Distance Backward (CUDA)"); -} diff --git a/other_models/GRNet/extensions/gridding_loss/setup.py b/other_models/GRNet/extensions/gridding_loss/setup.py deleted file mode 100644 index f237fd7..0000000 --- a/other_models/GRNet/extensions/gridding_loss/setup.py +++ /dev/null @@ -1,16 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-30 11:03:55 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-30 11:13:39 -# @Email: cshzxie@gmail.com - -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup(name='gridding_distance', - version='1.0.0', - ext_modules=[ - CUDAExtension('gridding_distance', ['gridding_distance_cuda.cpp', 'gridding_distance.cu']), - ], - cmdclass={'build_ext': BuildExtension}) diff --git a/other_models/GRNet/extensions_cloudconv/__init__.py b/other_models/GRNet/extensions_cloudconv/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/other_models/GRNet/extensions_cloudconv/__pycache__/__init__.cpython-36.pyc b/other_models/GRNet/extensions_cloudconv/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index 2eb9654..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/chamfer_dist/__init__.py b/other_models/GRNet/extensions_cloudconv/chamfer_dist/__init__.py deleted file mode 100644 index 4aaf4a3..0000000 --- a/other_models/GRNet/extensions_cloudconv/chamfer_dist/__init__.py +++ /dev/null @@ -1,42 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Thibault GROUEIX -# @Date: 2019-08-07 20:54:24 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-18 15:06:25 -# @Email: cshzxie@gmail.com - -import torch - -import chamfer - - -class ChamferFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, xyz1, xyz2): - dist1, dist2, idx1, idx2 = chamfer.forward(xyz1, xyz2) - ctx.save_for_backward(xyz1, xyz2, idx1, idx2) - - return dist1, dist2 - - @staticmethod - def backward(ctx, grad_dist1, grad_dist2): - xyz1, xyz2, idx1, idx2 = ctx.saved_tensors - grad_xyz1, grad_xyz2 = chamfer.backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2) - return grad_xyz1, grad_xyz2 - - -class ChamferDistance(torch.nn.Module): - def __init__(self, ignore_zeros=False): - super(ChamferDistance, self).__init__() - self.ignore_zeros = ignore_zeros - - def forward(self, xyz1, xyz2): - batch_size = xyz1.size(0) - if batch_size == 1 and self.ignore_zeros: - non_zeros1 = torch.sum(xyz1, dim=2).ne(0) - non_zeros2 = torch.sum(xyz2, dim=2).ne(0) - xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) - xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) - - dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) - return torch.mean(dist1) + torch.mean(dist2) diff --git a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.cu b/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.cu deleted file mode 100644 index 4bde058..0000000 --- a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.cu +++ /dev/null @@ -1,229 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-08-07 20:54:24 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:58:55 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include - -#include - -__global__ void chamfer_dist_kernel(int batch_size, - int n, - const float* xyz1, - int m, - const float* xyz2, - float* dist, - int* indexes) { - const int batch = 512; - __shared__ float buf[batch * 3]; - for (int i = blockIdx.x; i < batch_size; i += gridDim.x) { - for (int k2 = 0; k2 < m; k2 += batch) { - int end_k = min(m, k2 + batch) - k2; - for (int j = threadIdx.x; j < end_k * 3; j += blockDim.x) { - buf[j] = xyz2[(i * m + k2) * 3 + j]; - } - __syncthreads(); - for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; - j += blockDim.x * gridDim.y) { - float x1 = xyz1[(i * n + j) * 3 + 0]; - float y1 = xyz1[(i * n + j) * 3 + 1]; - float z1 = xyz1[(i * n + j) * 3 + 2]; - float best_dist = 0; - int best_dist_index = 0; - int end_ka = end_k - (end_k & 3); - if (end_ka == batch) { - for (int k = 0; k < batch; k += 4) { - { - float x2 = buf[k * 3 + 0] - x1; - float y2 = buf[k * 3 + 1] - y1; - float z2 = buf[k * 3 + 2] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - - if (k == 0 || dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2; - } - } - { - float x2 = buf[k * 3 + 3] - x1; - float y2 = buf[k * 3 + 4] - y1; - float z2 = buf[k * 3 + 5] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 1; - } - } - { - float x2 = buf[k * 3 + 6] - x1; - float y2 = buf[k * 3 + 7] - y1; - float z2 = buf[k * 3 + 8] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 2; - } - } - { - float x2 = buf[k * 3 + 9] - x1; - float y2 = buf[k * 3 + 10] - y1; - float z2 = buf[k * 3 + 11] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 3; - } - } - } - } else { - for (int k = 0; k < end_ka; k += 4) { - { - float x2 = buf[k * 3 + 0] - x1; - float y2 = buf[k * 3 + 1] - y1; - float z2 = buf[k * 3 + 2] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (k == 0 || dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2; - } - } - { - float x2 = buf[k * 3 + 3] - x1; - float y2 = buf[k * 3 + 4] - y1; - float z2 = buf[k * 3 + 5] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 1; - } - } - { - float x2 = buf[k * 3 + 6] - x1; - float y2 = buf[k * 3 + 7] - y1; - float z2 = buf[k * 3 + 8] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 2; - } - } - { - float x2 = buf[k * 3 + 9] - x1; - float y2 = buf[k * 3 + 10] - y1; - float z2 = buf[k * 3 + 11] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2 + 3; - } - } - } - } - for (int k = end_ka; k < end_k; k++) { - float x2 = buf[k * 3 + 0] - x1; - float y2 = buf[k * 3 + 1] - y1; - float z2 = buf[k * 3 + 2] - z1; - float dist = x2 * x2 + y2 * y2 + z2 * z2; - if (k == 0 || dist < best_dist) { - best_dist = dist; - best_dist_index = k + k2; - } - } - if (k2 == 0 || dist[(i * n + j)] > best_dist) { - dist[(i * n + j)] = best_dist; - indexes[(i * n + j)] = best_dist_index; - } - } - __syncthreads(); - } - } -} - -std::vector chamfer_cuda_forward(torch::Tensor xyz1, - torch::Tensor xyz2) { - const int batch_size = xyz1.size(0); - const int n = xyz1.size(1); // num_points point cloud A - const int m = xyz2.size(1); // num_points point cloud B - torch::Tensor dist1 = - torch::zeros({batch_size, n}, torch::CUDA(torch::kFloat)); - torch::Tensor dist2 = - torch::zeros({batch_size, m}, torch::CUDA(torch::kFloat)); - torch::Tensor idx1 = torch::zeros({batch_size, n}, torch::CUDA(torch::kInt)); - torch::Tensor idx2 = torch::zeros({batch_size, m}, torch::CUDA(torch::kInt)); - - chamfer_dist_kernel<<>>( - batch_size, n, xyz1.data_ptr(), m, xyz2.data_ptr(), - dist1.data_ptr(), idx1.data_ptr()); - chamfer_dist_kernel<<>>( - batch_size, m, xyz2.data_ptr(), n, xyz1.data_ptr(), - dist2.data_ptr(), idx2.data_ptr()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in chamfer_cuda_forward: %s\n", cudaGetErrorString(err)); - } - return {dist1, dist2, idx1, idx2}; -} - -__global__ void chamfer_dist_grad_kernel(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) { - for (int i = blockIdx.x; i < b; i += gridDim.x) { - for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; - j += blockDim.x * gridDim.y) { - float x1 = xyz1[(i * n + j) * 3 + 0]; - float y1 = xyz1[(i * n + j) * 3 + 1]; - float z1 = xyz1[(i * n + j) * 3 + 2]; - int j2 = idx1[i * n + j]; - float x2 = xyz2[(i * m + j2) * 3 + 0]; - float y2 = xyz2[(i * m + j2) * 3 + 1]; - float z2 = xyz2[(i * m + j2) * 3 + 2]; - float g = grad_dist1[i * n + j] * 2; - atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 0]), g * (x1 - x2)); - atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 1]), g * (y1 - y2)); - atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 2]), g * (z1 - z2)); - atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 0]), -(g * (x1 - x2))); - atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 1]), -(g * (y1 - y2))); - atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 2]), -(g * (z1 - z2))); - } - } -} - -std::vector chamfer_cuda_backward(torch::Tensor xyz1, - torch::Tensor xyz2, - torch::Tensor idx1, - torch::Tensor idx2, - torch::Tensor grad_dist1, - torch::Tensor grad_dist2) { - const int batch_size = xyz1.size(0); - const int n = xyz1.size(1); // num_points point cloud A - const int m = xyz2.size(1); // num_points point cloud B - torch::Tensor grad_xyz1 = torch::zeros_like(xyz1, torch::CUDA(torch::kFloat)); - torch::Tensor grad_xyz2 = torch::zeros_like(xyz2, torch::CUDA(torch::kFloat)); - - chamfer_dist_grad_kernel<<>>( - batch_size, n, xyz1.data_ptr(), m, xyz2.data_ptr(), - grad_dist1.data_ptr(), idx1.data_ptr(), - grad_xyz1.data_ptr(), grad_xyz2.data_ptr()); - chamfer_dist_grad_kernel<<>>( - batch_size, m, xyz2.data_ptr(), n, xyz1.data_ptr(), - grad_dist2.data_ptr(), idx2.data_ptr(), - grad_xyz2.data_ptr(), grad_xyz1.data_ptr()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in chamfer_cuda_backward: %s\n", cudaGetErrorString(err)); - } - return {grad_xyz1, grad_xyz2}; -} diff --git a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/PKG-INFO b/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/PKG-INFO deleted file mode 100644 index 9b7a6c5..0000000 --- a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: chamfer -Version: 2.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/SOURCES.txt b/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/SOURCES.txt deleted file mode 100644 index 0e582cb..0000000 --- a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -chamfer.cu -chamfer_cuda.cpp -setup.py -chamfer.egg-info/PKG-INFO -chamfer.egg-info/SOURCES.txt -chamfer.egg-info/dependency_links.txt -chamfer.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/dependency_links.txt b/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/top_level.txt b/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/top_level.txt deleted file mode 100644 index 27aeea6..0000000 --- a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -chamfer diff --git a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer_cuda.cpp b/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer_cuda.cpp deleted file mode 100644 index 9fca161..0000000 --- a/other_models/GRNet/extensions_cloudconv/chamfer_dist/chamfer_cuda.cpp +++ /dev/null @@ -1,39 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-08-07 20:54:24 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2019-12-10 10:33:50 - * @Email: cshzxie@gmail.com - */ - -#include -#include - -std::vector chamfer_cuda_forward(torch::Tensor xyz1, - torch::Tensor xyz2); - -std::vector chamfer_cuda_backward(torch::Tensor xyz1, - torch::Tensor xyz2, - torch::Tensor idx1, - torch::Tensor idx2, - torch::Tensor grad_dist1, - torch::Tensor grad_dist2); - -std::vector chamfer_forward(torch::Tensor xyz1, - torch::Tensor xyz2) { - return chamfer_cuda_forward(xyz1, xyz2); -} - -std::vector chamfer_backward(torch::Tensor xyz1, - torch::Tensor xyz2, - torch::Tensor idx1, - torch::Tensor idx2, - torch::Tensor grad_dist1, - torch::Tensor grad_dist2) { - return chamfer_cuda_backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &chamfer_forward, "Chamfer forward (CUDA)"); - m.def("backward", &chamfer_backward, "Chamfer backward (CUDA)"); -} diff --git a/other_models/GRNet/extensions_cloudconv/chamfer_dist/setup.py b/other_models/GRNet/extensions_cloudconv/chamfer_dist/setup.py deleted file mode 100644 index 04c6589..0000000 --- a/other_models/GRNet/extensions_cloudconv/chamfer_dist/setup.py +++ /dev/null @@ -1,19 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-08-07 20:54:24 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-10 10:04:25 -# @Email: cshzxie@gmail.com - -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup(name='chamfer', - version='2.0.0', - ext_modules=[ - CUDAExtension('chamfer', [ - 'chamfer_cuda.cpp', - 'chamfer.cu', - ]), - ], - cmdclass={'build_ext': BuildExtension}) diff --git a/other_models/GRNet/extensions_cloudconv/chamfer_dist/test.py b/other_models/GRNet/extensions_cloudconv/chamfer_dist/test.py deleted file mode 100644 index edf53ba..0000000 --- a/other_models/GRNet/extensions_cloudconv/chamfer_dist/test.py +++ /dev/null @@ -1,32 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-10 10:38:01 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 14:21:36 -# @Email: cshzxie@gmail.com -# -# Note: -# - Replace float -> double, kFloat -> kDouble in chamfer.cu - -import os -import sys -import torch -import unittest - -from torch.autograd import gradcheck - -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) -from extensions.chamfer_dist import ChamferFunction - - -class ChamferDistanceTestCase(unittest.TestCase): - def test_chamfer_dist(self): - x = torch.rand(4, 64, 3).double() - y = torch.rand(4, 128, 3).double() - x.requires_grad = True - y.requires_grad = True - print(gradcheck(ChamferFunction.apply, [x.cuda(), y.cuda()])) - - -if __name__ == '__main__': - unittest.main() diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/__init__.py b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/__init__.py deleted file mode 100644 index 835582e..0000000 --- a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/__init__.py +++ /dev/null @@ -1,40 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-19 16:55:15 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 13:15:14 -# @Email: cshzxie@gmail.com - -import torch - -import cubic_feature_sampling - - -class CubicFeatureSamplingFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, ptcloud, cubic_features, neighborhood_size=1): - scale = cubic_features.size(2) - point_features, grid_pt_indexes = cubic_feature_sampling.forward(scale, neighborhood_size, ptcloud, - cubic_features) - ctx.save_for_backward(torch.Tensor([scale]), torch.Tensor([neighborhood_size]), grid_pt_indexes) - return point_features - - @staticmethod - def backward(ctx, grad_point_features): - scale, neighborhood_size, grid_pt_indexes = ctx.saved_tensors - scale = int(scale.item()) - neighborhood_size = int(neighborhood_size.item()) - grad_point_features = grad_point_features.contiguous() - grad_ptcloud, grad_cubic_features = cubic_feature_sampling.backward(scale, neighborhood_size, - grad_point_features, grid_pt_indexes) - return grad_ptcloud, grad_cubic_features, None - - -class CubicFeatureSampling(torch.nn.Module): - def __init__(self): - super(CubicFeatureSampling, self).__init__() - - def forward(self, ptcloud, cubic_features, neighborhood_size=1): - h_scale = cubic_features.size(2) / 2 - ptcloud = ptcloud * h_scale + h_scale - return CubicFeatureSamplingFunction.apply(ptcloud, cubic_features, neighborhood_size) diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/__pycache__/__init__.cpython-36.pyc b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index f19a394..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.cu b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.cu deleted file mode 100644 index 423e962..0000000 --- a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.cu +++ /dev/null @@ -1,205 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-12-19 20:36:36 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:55:41 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include -#include - -#define CUDA_NUM_THREADS 512 - -// Computer the number of threads needed in GPU -inline int get_n_threads(int n) { - const int pow_2 = std::log(static_cast(n)) / std::log(2.0); - return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); -} - -__device__ int compute_index(int offset_x, - int offset_y, - int offset_z, - int scale) { - return offset_x * scale * scale + offset_y * scale + offset_z; -} - -__global__ void cubic_feature_sampling_kernel( - int scale, - int neighborhood_size, - int n_vertices, - int n_pts, - int n_cubic_channels, - const float *__restrict__ ptcloud, - const float *__restrict__ cubic_features, - float *__restrict__ point_features, - int *__restrict__ grid_pt_indexes) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - int cub_scale = scale * scale * scale; - - ptcloud += batch_index * n_pts * 3; - cubic_features += batch_index * n_cubic_channels * cub_scale; - point_features += batch_index * n_pts * n_vertices * n_cubic_channels; - grid_pt_indexes += batch_index * n_pts * n_vertices; - - for (int i = index; i < n_pts; i += stride) { - float pt_x = ptcloud[i * 3 + 0]; - float pt_y = ptcloud[i * 3 + 1]; - float pt_z = ptcloud[i * 3 + 2]; - - int lower_x = std::floor(pt_x); - int upper_x = std::ceil(pt_x); - if (lower_x == upper_x) { - upper_x += 1; - } - int lower_y = std::floor(pt_y); - int upper_y = std::ceil(pt_y); - if (lower_y == upper_y) { - upper_y += 1; - } - int lower_z = std::floor(pt_z); - int upper_z = std::ceil(pt_z); - if (lower_z == upper_z) { - upper_z += 1; - } - - int ns = neighborhood_size - 1; - int vertex_idx = 0; - for (int j = lower_x - ns; j <= upper_x + ns; ++j) { - for (int k = lower_y - ns; k <= upper_y + ns; ++k) { - for (int m = lower_z - ns; m <= upper_z + ns; ++m) { - if (j < 0 || j >= scale || k < 0 || k >= scale || m < 0 || - m >= scale) { - // Ignore points lies out of the grid - grid_pt_indexes[i * n_vertices + vertex_idx++] = -1; - } else { - // Calcuating indexes for adjacent vertices - grid_pt_indexes[i * n_vertices + vertex_idx++] = - compute_index(j, k, m, scale); - } - } - } - } - - // Gather Features - for (int j = 0; j < n_vertices; ++j) { - for (int k = 0; k < n_cubic_channels; ++k) { - int vertex_idx = grid_pt_indexes[i * n_vertices + j]; - if (vertex_idx == -1) { - continue; - } - int feature_idx = - i * n_vertices * n_cubic_channels + j * n_cubic_channels + k; - float feature_val = cubic_features[k * cub_scale + vertex_idx]; - point_features[feature_idx] = feature_val; - } - } - } -} - -std::vector cubic_feature_sampling_cuda_forward( - int scale, - int neighborhood_size, - torch::Tensor ptcloud, - torch::Tensor cubic_features, - cudaStream_t stream) { - int batch_size = ptcloud.size(0); - int n_pts = ptcloud.size(1); - int n_cubic_channels = cubic_features.size(1); - - int n_vertices = std::pow(neighborhood_size * 2, 3); - torch::Tensor point_features = - torch::zeros({batch_size, n_pts, n_vertices, n_cubic_channels}, - torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_indexes = - torch::zeros({batch_size, n_pts, n_vertices}, torch::CUDA(torch::kInt)); - - cubic_feature_sampling_kernel<<>>( - scale, neighborhood_size, n_vertices, n_pts, n_cubic_channels, - ptcloud.data(), cubic_features.data(), - point_features.data(), grid_pt_indexes.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in cubic_feature_sampling_cuda_forward: %s\n", - cudaGetErrorString(err)); - } - return {point_features, grid_pt_indexes}; -} - -__global__ void cubic_feature_sampling_grad_kernel( - int scale, - int neighborhood_size, - int n_vertices, - int n_pts, - int n_cubic_channels, - const float *__restrict__ grad_point_features, - const int *__restrict__ grid_pt_indexes, - float *__restrict__ grad_ptcloud, - float *__restrict__ grad_cubic_features) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - int cub_scale = scale * scale * scale; - - grad_point_features += batch_index * n_pts * n_vertices * n_cubic_channels; - grid_pt_indexes += batch_index * n_pts * n_vertices; - grad_ptcloud += batch_index * n_pts * 3; - grad_cubic_features += batch_index * n_cubic_channels * cub_scale; - - for (int i = index; i < n_pts; i += stride) { - for (int j = 0; j < n_vertices; ++j) { - int vertex_idx = grid_pt_indexes[i * n_vertices + j]; - if (vertex_idx == -1) { - continue; - } - for (int k = 0; k < n_cubic_channels; ++k) { - int grad_idx = - i * n_vertices * n_cubic_channels + j * n_cubic_channels + k; - float grad_val = grad_point_features[grad_idx]; - // Fix bugs: the gradients of ceil and floor functions are zeros. - // Ref: https://github.com/tensorflow/tensorflow/issues/897 - // atomicAdd(&(grad_ptcloud[i * 3 + 0]), grad_val); - // atomicAdd(&(grad_ptcloud[i * 3 + 1]), grad_val); - // atomicAdd(&(grad_ptcloud[i * 3 + 2]), grad_val); - atomicAdd(&(grad_cubic_features[k * cub_scale + vertex_idx]), grad_val); - } - } - } -} - -std::vector cubic_feature_sampling_cuda_backward( - int scale, - int neighborhood_size, - torch::Tensor grad_point_features, - torch::Tensor grid_pt_indexes, - cudaStream_t stream) { - int batch_size = grad_point_features.size(0); - int n_cubic_channels = grad_point_features.size(3); - int n_pts = grid_pt_indexes.size(1); - int n_vertices = std::pow(neighborhood_size * 2, 3); - - torch::Tensor grad_ptcloud = - torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); - torch::Tensor grad_cubic_features = - torch::zeros({batch_size, n_cubic_channels, scale, scale, scale}, - torch::CUDA(torch::kFloat)); - - cubic_feature_sampling_grad_kernel<<>>( - scale, neighborhood_size, n_vertices, n_pts, n_cubic_channels, - grad_point_features.data(), grid_pt_indexes.data(), - grad_ptcloud.data(), grad_cubic_features.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in cubic_feature_sampling_cuda_backward: %s\n", - cudaGetErrorString(err)); - } - return {grad_ptcloud, grad_cubic_features}; -} diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/PKG-INFO b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/PKG-INFO deleted file mode 100644 index e2bce51..0000000 --- a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: cubic-feature-sampling -Version: 1.1.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/SOURCES.txt b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/SOURCES.txt deleted file mode 100644 index a795ed8..0000000 --- a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -cubic_feature_sampling.cu -cubic_feature_sampling_cuda.cpp -setup.py -cubic_feature_sampling.egg-info/PKG-INFO -cubic_feature_sampling.egg-info/SOURCES.txt -cubic_feature_sampling.egg-info/dependency_links.txt -cubic_feature_sampling.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/dependency_links.txt b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/top_level.txt b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/top_level.txt deleted file mode 100644 index 5870961..0000000 --- a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -cubic_feature_sampling diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp deleted file mode 100644 index 2bcbe7e..0000000 --- a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp +++ /dev/null @@ -1,65 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-12-19 17:04:38 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:50:22 - * @Email: cshzxie@gmail.com - */ - -#include -#include - -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) \ - AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -std::vector cubic_feature_sampling_cuda_forward( - int scale, - int neighborhood_size, - torch::Tensor ptcloud, - torch::Tensor cubic_features, - cudaStream_t stream); - -std::vector cubic_feature_sampling_cuda_backward( - int scale, - int neighborhood_size, - torch::Tensor grad_point_features, - torch::Tensor grid_pt_indexes, - cudaStream_t stream); - -std::vector cubic_feature_sampling_forward( - int scale, - int neighborhood_size, - torch::Tensor ptcloud, - torch::Tensor cubic_features) { - CHECK_INPUT(ptcloud); - CHECK_INPUT(cubic_features); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return cubic_feature_sampling_cuda_forward(scale, neighborhood_size, ptcloud, - cubic_features, stream); -} - -std::vector cubic_feature_sampling_backward( - int scale, - int neighborhood_size, - torch::Tensor grad_point_features, - torch::Tensor grid_pt_indexes) { - CHECK_INPUT(grad_point_features); - CHECK_INPUT(grid_pt_indexes); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return cubic_feature_sampling_cuda_backward( - scale, neighborhood_size, grad_point_features, grid_pt_indexes, stream); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &cubic_feature_sampling_forward, - "Cubic Feature Sampling forward (CUDA)"); - m.def("backward", &cubic_feature_sampling_backward, - "Cubic Feature Sampling backward (CUDA)"); -} \ No newline at end of file diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py2.7-linux-x86_64.egg b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py2.7-linux-x86_64.egg deleted file mode 100644 index aa6a419..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py2.7-linux-x86_64.egg and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.6-linux-x86_64.egg b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.6-linux-x86_64.egg deleted file mode 100644 index cb94e76..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/setup.py b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/setup.py deleted file mode 100644 index 9a7783b..0000000 --- a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/setup.py +++ /dev/null @@ -1,16 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-19 17:03:06 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 14:02:06 -# @Email: cshzxie@gmail.com - -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup(name='cubic_feature_sampling', - version='1.1.0', - ext_modules=[ - CUDAExtension('cubic_feature_sampling', ['cubic_feature_sampling_cuda.cpp', 'cubic_feature_sampling.cu']), - ], - cmdclass={'build_ext': BuildExtension}) diff --git a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/test.py b/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/test.py deleted file mode 100644 index 7f51017..0000000 --- a/other_models/GRNet/extensions_cloudconv/cubic_feature_sampling/test.py +++ /dev/null @@ -1,52 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-20 11:50:50 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 13:52:33 -# @Email: cshzxie@gmail.com -# -# Note: -# - Replace float -> double, kFloat -> kDouble in cubic_feature_sampling.cu - -import os -import sys -import torch -import unittest - -from torch.autograd import gradcheck - -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) -from extensions.cubic_feature_sampling import CubicFeatureSamplingFunction - - -class CubicFeatureSamplingTestCase(unittest.TestCase): - def test_neighborhood_size_1(self): - ptcloud = torch.rand(2, 64, 3) * 2 - 1 - cubic_features = torch.rand(2, 4, 8, 8, 8) - ptcloud.requires_grad = True - cubic_features.requires_grad = True - self.assertTrue( - gradcheck(CubicFeatureSamplingFunction.apply, - [ptcloud.double().cuda(), cubic_features.double().cuda()])) - - def test_neighborhood_size_2(self): - ptcloud = torch.rand(2, 32, 3) * 2 - 1 - cubic_features = torch.rand(2, 2, 8, 8, 8) - ptcloud.requires_grad = True - cubic_features.requires_grad = True - self.assertTrue( - gradcheck(CubicFeatureSamplingFunction.apply, - [ptcloud.double().cuda(), cubic_features.double().cuda(), 2])) - - def test_neighborhood_size_3(self): - ptcloud = torch.rand(1, 32, 3) * 2 - 1 - cubic_features = torch.rand(1, 2, 16, 16, 16) - ptcloud.requires_grad = True - cubic_features.requires_grad = True - self.assertTrue( - gradcheck(CubicFeatureSamplingFunction.apply, - [ptcloud.double().cuda(), cubic_features.double().cuda(), 3])) - - -if __name__ == '__main__': - unittest.main() diff --git a/other_models/GRNet/extensions_cloudconv/gridding/__init__.py b/other_models/GRNet/extensions_cloudconv/gridding/__init__.py deleted file mode 100644 index 4c7ff45..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/__init__.py +++ /dev/null @@ -1,74 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-11-15 20:33:52 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-30 09:55:53 -# @Email: cshzxie@gmail.com - -import torch - -import gridding - - -class GriddingFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, scale, ptcloud): - grid, grid_pt_weights, grid_pt_indexes = gridding.forward(-scale, scale - 1, -scale, scale - 1, -scale, - scale - 1, ptcloud) - # print(grid.size()) # torch.Size(batch_size, n_grid_vertices) - # print(grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) - # print(grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) - ctx.save_for_backward(grid_pt_weights, grid_pt_indexes) - - return grid - - @staticmethod - def backward(ctx, grad_grid): - grid_pt_weights, grid_pt_indexes = ctx.saved_tensors - grad_ptcloud = gridding.backward(grid_pt_weights, grid_pt_indexes, grad_grid) - # print(grad_ptcloud.size()) # torch.Size(batch_size, n_pts, 3) - - return None, grad_ptcloud - - -class Gridding(torch.nn.Module): - def __init__(self, scale=1): - super(Gridding, self).__init__() - self.scale = scale // 2 - - def forward(self, ptcloud): - ptcloud = ptcloud * self.scale - _ptcloud = torch.split(ptcloud, 1, dim=0) - grids = [] - for p in _ptcloud: - non_zeros = torch.sum(p, dim=2).ne(0) - p = p[non_zeros].unsqueeze(dim=0) - grids.append(GriddingFunction.apply(self.scale, p)) - - return torch.cat(grids, dim=0).contiguous() - - -class GriddingReverseFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, scale, grid): - ptcloud = gridding.rev_forward(scale, grid) - ctx.save_for_backward(torch.Tensor([scale]), grid, ptcloud) - return ptcloud - - @staticmethod - def backward(ctx, grad_ptcloud): - scale, grid, ptcloud = ctx.saved_tensors - scale = int(scale.item()) - grad_grid = gridding.rev_backward(ptcloud, grid, grad_ptcloud) - grad_grid = grad_grid.view(-1, scale, scale, scale) - return None, grad_grid - - -class GriddingReverse(torch.nn.Module): - def __init__(self, scale=1): - super(GriddingReverse, self).__init__() - self.scale = scale - - def forward(self, grid): - ptcloud = GriddingReverseFunction.apply(self.scale, grid) - return ptcloud / self.scale * 2 diff --git a/other_models/GRNet/extensions_cloudconv/gridding/__pycache__/__init__.cpython-36.pyc b/other_models/GRNet/extensions_cloudconv/gridding/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index 6cf9831..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/gridding/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/gridding/dist/gridding-2.1.0-py2.7-linux-x86_64.egg b/other_models/GRNet/extensions_cloudconv/gridding/dist/gridding-2.1.0-py2.7-linux-x86_64.egg deleted file mode 100644 index 596a8d3..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/gridding/dist/gridding-2.1.0-py2.7-linux-x86_64.egg and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/gridding/dist/gridding-2.1.0-py3.6-linux-x86_64.egg b/other_models/GRNet/extensions_cloudconv/gridding/dist/gridding-2.1.0-py3.6-linux-x86_64.egg deleted file mode 100644 index 32594c9..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/gridding/dist/gridding-2.1.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/gridding/gridding.cu b/other_models/GRNet/extensions_cloudconv/gridding/gridding.cu deleted file mode 100644 index 2d60880..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/gridding.cu +++ /dev/null @@ -1,335 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-11-13 10:53:22 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 15:00:10 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include -#include - -#define CUDA_NUM_THREADS 512 - -// Computer the number of threads needed in GPU -inline int get_n_threads(int n) { - const int pow_2 = std::log(static_cast(n)) / std::log(2.0); - return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); -} - -__device__ int compute_index( - int offset_x, int offset_y, int offset_z, int len_y, int len_z) { - return offset_x * len_y * len_z + offset_y * len_z + offset_z; -} - -__device__ float compute_weight(float x, float x0) { return 1 - abs(x - x0); } - -__global__ void gridding_kernel(int n_grid_vertices, - int n_pts, - float min_x, - float min_y, - float min_z, - int len_y, - int len_z, - const float *__restrict__ ptcloud, - float *__restrict__ grid_weights, - float *__restrict__ grid_pt_weights, - int *__restrict__ grid_pt_indexes) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - ptcloud += batch_index * n_pts * 3; - grid_weights += batch_index * n_grid_vertices; - grid_pt_weights += batch_index * n_pts * 24; - grid_pt_indexes += batch_index * n_pts * 8; - - for (int j = index; j < n_pts; j += stride) { - float pt_x = ptcloud[j * 3 + 0]; - float pt_y = ptcloud[j * 3 + 1]; - float pt_z = ptcloud[j * 3 + 2]; - - int lower_x = std::floor(pt_x); - int upper_x = std::ceil(pt_x); - if (lower_x == upper_x) { - upper_x += 1; - } - int lower_y = std::floor(pt_y); - int upper_y = std::ceil(pt_y); - if (lower_y == upper_y) { - upper_y += 1; - } - int lower_z = std::floor(pt_z); - int upper_z = std::ceil(pt_z); - if (lower_z == upper_z) { - upper_z += 1; - } - - int lx_offset = lower_x - min_x, ux_offset = upper_x - min_x; - int ly_offset = lower_y - min_y, uy_offset = upper_y - min_y; - int lz_offset = lower_z - min_z, uz_offset = upper_z - min_z; - - // Compute weights and corresponding positions, a loop for 8 points - // LLL -> Lower X, Lower Y, Lower Z - grid_pt_indexes[j * 8 + 0] = - compute_index(lx_offset, ly_offset, lz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 0] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 1] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 2] = compute_weight(pt_z, lower_z); - - // LLU -> Lower X, Lower Y, Upper Z - grid_pt_indexes[j * 8 + 1] = - compute_index(lx_offset, ly_offset, uz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 3] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 4] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 5] = compute_weight(pt_z, upper_z); - - // LUL -> Lower X, Upper Y, Lower Z - grid_pt_indexes[j * 8 + 2] = - compute_index(lx_offset, uy_offset, lz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 6] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 7] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 8] = compute_weight(pt_z, lower_z); - - // LUU -> Lower X, Upper Y, Upper Z - grid_pt_indexes[j * 8 + 3] = - compute_index(lx_offset, uy_offset, uz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 9] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 10] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 11] = compute_weight(pt_z, upper_z); - - // ULL -> Upper X, Lower Y, Lower Z - grid_pt_indexes[j * 8 + 4] = - compute_index(ux_offset, ly_offset, lz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 12] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 13] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 14] = compute_weight(pt_z, lower_z); - - // ULU -> Upper X, Lower Y, Upper Z - grid_pt_indexes[j * 8 + 5] = - compute_index(ux_offset, ly_offset, uz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 15] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 16] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 17] = compute_weight(pt_z, upper_z); - - // UUL -> Upper X, Upper Y, Lower Z - grid_pt_indexes[j * 8 + 6] = - compute_index(ux_offset, uy_offset, lz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 18] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 19] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 20] = compute_weight(pt_z, lower_z); - - // UUU -> Upper X, Upper Y, Upper Z - grid_pt_indexes[j * 8 + 7] = - compute_index(ux_offset, uy_offset, uz_offset, len_y, len_z); - grid_pt_weights[j * 24 + 21] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 22] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 23] = compute_weight(pt_z, upper_z); - } - - __syncthreads(); - - int gvtx_idx = 0; - for (int j = index; j < n_pts; j += stride) { - // LLL -> Lower X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 0]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 0] * - grid_pt_weights[j * 24 + 1] * - grid_pt_weights[j * 24 + 2]); - // LLU -> Lower X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 1]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 3] * - grid_pt_weights[j * 24 + 4] * - grid_pt_weights[j * 24 + 5]); - // LUL -> Lower X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 2]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 6] * - grid_pt_weights[j * 24 + 7] * - grid_pt_weights[j * 24 + 8]); - // LUU -> Lower X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 3]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 9] * - grid_pt_weights[j * 24 + 10] * - grid_pt_weights[j * 24 + 11]); - // ULL -> Upper X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 4]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 12] * - grid_pt_weights[j * 24 + 13] * - grid_pt_weights[j * 24 + 14]); - // ULU -> Upper X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 5]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 15] * - grid_pt_weights[j * 24 + 16] * - grid_pt_weights[j * 24 + 17]); - // UUL -> Upper X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 6]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 18] * - grid_pt_weights[j * 24 + 19] * - grid_pt_weights[j * 24 + 20]); - // UUU -> Upper X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 7]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 21] * - grid_pt_weights[j * 24 + 22] * - grid_pt_weights[j * 24 + 23]); - } -} - -std::vector gridding_cuda_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud, - cudaStream_t stream) { - int batch_size = ptcloud.size(0); - int n_pts = ptcloud.size(1); - int len_x = max_x - min_x + 1; - int len_y = max_y - min_y + 1; - int len_z = max_z - min_z + 1; - int n_grid_vertices = len_x * len_y * len_z; - - torch::Tensor grid_weights = - torch::zeros({batch_size, n_grid_vertices}, torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_weights = - torch::zeros({batch_size, n_pts, 8, 3}, torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_indexes = - torch::zeros({batch_size, n_pts, 8}, torch::CUDA(torch::kInt)); - - gridding_kernel<<>>( - n_grid_vertices, n_pts, min_x, min_y, min_z, len_y, len_z, - ptcloud.data(), grid_weights.data(), - grid_pt_weights.data(), grid_pt_indexes.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_cuda_forward: %s\n", cudaGetErrorString(err)); - } - return {grid_weights, grid_pt_weights, grid_pt_indexes}; -} - -__global__ void gridding_grad_kernel(int n_grid_vertices, - int n_pts, - const float *__restrict__ grid_pt_weights, - const int *__restrict__ grid_pt_indexes, - const float *__restrict__ grad_grid, - float *__restrict__ grad_ptcloud) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - grid_pt_weights += batch_index * n_pts * 24; - grid_pt_indexes += batch_index * n_pts * 8; - grad_grid += batch_index * n_grid_vertices; - grad_ptcloud += batch_index * n_pts * 3; - - int gvtx_idx = 0; - float grad_vtx = 0, x_weights = 0, y_weights = 0, z_weights = 0; - for (int j = index; j < n_pts; j += stride) { - // Compute gradient for the corresponding positions, a loop for 8 points - // LLL -> Lower X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 0]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 0]; - y_weights = grid_pt_weights[j * 24 + 1]; - z_weights = grid_pt_weights[j * 24 + 2]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // LLU -> Lower X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 1]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 3]; - y_weights = grid_pt_weights[j * 24 + 4]; - z_weights = grid_pt_weights[j * 24 + 5]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // LUL -> Lower X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 2]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 6]; - y_weights = grid_pt_weights[j * 24 + 7]; - z_weights = grid_pt_weights[j * 24 + 8]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // LUU -> Lower X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 3]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 9]; - y_weights = grid_pt_weights[j * 24 + 10]; - z_weights = grid_pt_weights[j * 24 + 11]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // ULL -> Upper X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 4]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 12]; - y_weights = grid_pt_weights[j * 24 + 13]; - z_weights = grid_pt_weights[j * 24 + 14]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // ULU -> Upper X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 5]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 15]; - y_weights = grid_pt_weights[j * 24 + 16]; - z_weights = grid_pt_weights[j * 24 + 17]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // UUL -> Upper X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 6]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 18]; - y_weights = grid_pt_weights[j * 24 + 19]; - z_weights = grid_pt_weights[j * 24 + 20]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // UUU -> Upper X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 7]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 21]; - y_weights = grid_pt_weights[j * 24 + 22]; - z_weights = grid_pt_weights[j * 24 + 23]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - } -} - -torch::Tensor gridding_cuda_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid, - cudaStream_t stream) { - int batch_size = grad_grid.size(0); - int n_grid_vertices = grad_grid.size(1); - int n_pts = grid_pt_indexes.size(1); - - torch::Tensor grad_ptcloud = - torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); - - gridding_grad_kernel<<>>( - n_grid_vertices, n_pts, grid_pt_weights.data(), - grid_pt_indexes.data(), grad_grid.data(), - grad_ptcloud.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_cuda_backward: %s\n", cudaGetErrorString(err)); - } - return grad_ptcloud; -} diff --git a/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/PKG-INFO b/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/PKG-INFO deleted file mode 100644 index 72f20bd..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: gridding -Version: 2.1.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/SOURCES.txt b/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/SOURCES.txt deleted file mode 100644 index f241b60..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/SOURCES.txt +++ /dev/null @@ -1,8 +0,0 @@ -gridding.cu -gridding_cuda.cpp -gridding_reverse.cu -setup.py -gridding.egg-info/PKG-INFO -gridding.egg-info/SOURCES.txt -gridding.egg-info/dependency_links.txt -gridding.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/dependency_links.txt b/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/top_level.txt b/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/top_level.txt deleted file mode 100644 index 7ee56cc..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/gridding.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -gridding diff --git a/other_models/GRNet/extensions_cloudconv/gridding/gridding_cuda.cpp b/other_models/GRNet/extensions_cloudconv/gridding/gridding_cuda.cpp deleted file mode 100644 index e4201c0..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/gridding_cuda.cpp +++ /dev/null @@ -1,94 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-11-13 10:52:53 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:52:32 - * @Email: cshzxie@gmail.com - */ - -#include -#include - -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) \ - AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -std::vector gridding_cuda_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud, - cudaStream_t stream); - -torch::Tensor gridding_cuda_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid, - cudaStream_t stream); - -torch::Tensor gridding_reverse_cuda_forward(int scale, - torch::Tensor grid, - cudaStream_t stream); - -torch::Tensor gridding_reverse_cuda_backward(torch::Tensor ptcloud, - torch::Tensor grid, - torch::Tensor grad_ptcloud, - cudaStream_t stream); - -std::vector gridding_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud) { - CHECK_INPUT(ptcloud); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_cuda_forward(min_x, max_x, min_y, max_y, min_z, max_z, - ptcloud, stream); -} - -torch::Tensor gridding_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid) { - CHECK_INPUT(grid_pt_weights); - CHECK_INPUT(grid_pt_indexes); - CHECK_INPUT(grad_grid); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_cuda_backward(grid_pt_weights, grid_pt_indexes, grad_grid, - stream); -} - -torch::Tensor gridding_reverse_forward(int scale, torch::Tensor grid) { - CHECK_INPUT(grid); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_reverse_cuda_forward(scale, grid, stream); -} - -torch::Tensor gridding_reverse_backward(torch::Tensor ptcloud, - torch::Tensor grid, - torch::Tensor grad_ptcloud) { - CHECK_INPUT(ptcloud); - CHECK_INPUT(grid); - CHECK_INPUT(grad_ptcloud); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_reverse_cuda_backward(ptcloud, grid, grad_ptcloud, stream); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &gridding_forward, "Gridding forward (CUDA)"); - m.def("backward", &gridding_backward, "Gridding backward (CUDA)"); - m.def("rev_forward", &gridding_reverse_forward, - "Gridding Reverse forward (CUDA)"); - m.def("rev_backward", &gridding_reverse_backward, - "Gridding Reverse backward (CUDA)"); -} diff --git a/other_models/GRNet/extensions_cloudconv/gridding/gridding_reverse.cu b/other_models/GRNet/extensions_cloudconv/gridding/gridding_reverse.cu deleted file mode 100644 index 63b594c..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/gridding_reverse.cu +++ /dev/null @@ -1,236 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-11-21 16:42:18 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 15:00:21 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include -#include - -#define CUDA_NUM_THREADS 512 -#define EPS 1e-6 - -// Computer the number of threads needed in GPU -inline int get_n_threads(int n) { - const int pow_2 = std::log(static_cast(n)) / std::log(2.0); - return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); -} - -__device__ int compute_index(int offset_x, - int offset_y, - int offset_z, - int scale) { - return offset_x * scale * scale + offset_y * scale + offset_z; -} - -__global__ void gridding_reverse_kernel(int scale, - int n_pts, - const float *__restrict__ grid, - float *__restrict__ ptcloud) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - ptcloud += batch_index * n_pts * 3; - grid += batch_index * n_pts; - - for (int j = index; j < n_pts; j += stride) { - int sqr_scale = scale * scale; - int x_offset = j / sqr_scale; - int y_offset = j % sqr_scale / scale; - int z_offset = j % sqr_scale % scale; - if (x_offset == 0 || y_offset == 0 || z_offset == 0) { - continue; - } - - // assert j == compute_index(x_offset, y_offset, z_offset, scale) - float weights[8] = { - grid[compute_index(x_offset - 1, y_offset - 1, z_offset - 1, scale)], - grid[compute_index(x_offset - 1, y_offset - 1, z_offset, scale)], - grid[compute_index(x_offset - 1, y_offset, z_offset - 1, scale)], - grid[compute_index(x_offset - 1, y_offset, z_offset, scale)], - grid[compute_index(x_offset, y_offset - 1, z_offset - 1, scale)], - grid[compute_index(x_offset, y_offset - 1, z_offset, scale)], - grid[compute_index(x_offset, y_offset, z_offset - 1, scale)], - grid[j]}; - - float weights_sum = 0; - for (size_t i = 0; i < 8; ++i) { - weights_sum += weights[i]; - } - if (weights_sum < EPS) { - continue; - } - for (size_t i = 0; i < 8; ++i) { - weights[i] /= weights_sum; - } - - x_offset -= scale / 2; - y_offset -= scale / 2; - z_offset -= scale / 2; - - // clang-format off - ptcloud[j * 3 + 0] = weights[0] * (x_offset - 1) + - weights[1] * (x_offset - 1) + - weights[2] * (x_offset - 1) + - weights[3] * (x_offset - 1) + - weights[4] * x_offset + - weights[5] * x_offset + - weights[6] * x_offset + - weights[7] * x_offset; - ptcloud[j * 3 + 1] = weights[0] * (y_offset - 1) + - weights[1] * (y_offset - 1) + - weights[2] * y_offset + - weights[3] * y_offset + - weights[4] * (y_offset - 1) + - weights[5] * (y_offset - 1) + - weights[6] * y_offset + - weights[7] * y_offset; - ptcloud[j * 3 + 2] = weights[0] * (z_offset - 1) + - weights[1] * z_offset + - weights[2] * (z_offset - 1) + - weights[3] * z_offset + - weights[4] * (z_offset - 1) + - weights[5] * z_offset + - weights[6] * (z_offset - 1) + - weights[7] * z_offset; - // clang-format on - } -} - -torch::Tensor gridding_reverse_cuda_forward(int scale, - torch::Tensor grid, - cudaStream_t stream) { - int batch_size = grid.size(0); - int n_pts = scale * scale * scale; - - torch::Tensor ptcloud = - torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); - - gridding_reverse_kernel<<>>( - scale, n_pts, grid.data(), ptcloud.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_cuda_forward: %s\n", cudaGetErrorString(err)); - } - return ptcloud; -} - -__global__ void gridding_reverse_grad_kernel( - int scale, - int n_pts, - const float *__restrict__ ptcloud, - const float *__restrict__ grid, - const float *__restrict__ grad_ptcloud, - float *__restrict__ grad_grid) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - ptcloud += batch_index * n_pts * 3; - grid += batch_index * n_pts; - grad_ptcloud += batch_index * n_pts * 3; - grad_grid += batch_index * n_pts; - - for (int j = index; j < n_pts; j += stride) { - int sqr_scale = scale * scale; - int x_offset = j / sqr_scale; - int y_offset = j % sqr_scale / scale; - int z_offset = j % sqr_scale % scale; - if (x_offset == 0 || y_offset == 0 || z_offset == 0) { - continue; - } - - int gvtx_indexes[8] = { - compute_index(x_offset - 1, y_offset - 1, z_offset - 1, scale), - compute_index(x_offset - 1, y_offset - 1, z_offset, scale), - compute_index(x_offset - 1, y_offset, z_offset - 1, scale), - compute_index(x_offset - 1, y_offset, z_offset, scale), - compute_index(x_offset, y_offset - 1, z_offset - 1, scale), - compute_index(x_offset, y_offset - 1, z_offset, scale), - compute_index(x_offset, y_offset, z_offset - 1, scale), - j}; - float weights[8] = {grid[gvtx_indexes[0]], grid[gvtx_indexes[1]], - grid[gvtx_indexes[2]], grid[gvtx_indexes[3]], - grid[gvtx_indexes[4]], grid[gvtx_indexes[5]], - grid[gvtx_indexes[6]], grid[gvtx_indexes[7]]}; - - float weights_sum = 0; - for (size_t i = 0; i < 8; ++i) { - weights_sum += weights[i]; - } - - if (weights_sum < EPS) { - continue; - } - for (size_t i = 0; i < 8; ++i) { - weights[i] /= weights_sum; - } - - x_offset -= scale / 2; - y_offset -= scale / 2; - z_offset -= scale / 2; - - // clang-format off - atomicAdd(&(grad_grid[gvtx_indexes[0]]), - grad_ptcloud[j * 3 + 0] * ((x_offset - 1) - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * ((y_offset - 1) - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * ((z_offset - 1) - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[1]]), - grad_ptcloud[j * 3 + 0] * ((x_offset - 1) - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * ((y_offset - 1) - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * (z_offset - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[2]]), - grad_ptcloud[j * 3 + 0] * ((x_offset - 1) - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * (y_offset - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * ((z_offset - 1) - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[3]]), - grad_ptcloud[j * 3 + 0] * ((x_offset - 1) - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * (y_offset - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * (z_offset - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[4]]), - grad_ptcloud[j * 3 + 0] * (x_offset - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * ((y_offset - 1) - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * ((z_offset - 1) - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[5]]), - grad_ptcloud[j * 3 + 0] * (x_offset - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * ((y_offset - 1) - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * (z_offset - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[6]]), - grad_ptcloud[j * 3 + 0] * (x_offset - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * (y_offset - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * ((z_offset - 1) - ptcloud[j * 3 + 2]) / weights_sum); - atomicAdd(&(grad_grid[gvtx_indexes[7]]), - grad_ptcloud[j * 3 + 0] * (x_offset - ptcloud[j * 3 + 0]) / weights_sum + - grad_ptcloud[j * 3 + 1] * (y_offset - ptcloud[j * 3 + 1]) / weights_sum + - grad_ptcloud[j * 3 + 2] * (z_offset - ptcloud[j * 3 + 2]) / weights_sum); - // clang-format on - } -} - -torch::Tensor gridding_reverse_cuda_backward(torch::Tensor ptcloud, - torch::Tensor grid, - torch::Tensor grad_ptcloud, - cudaStream_t stream) { - int batch_size = ptcloud.size(0); - int n_pts = ptcloud.size(1); - int scale = static_cast(std::cbrt(n_pts)); - - torch::Tensor grad_grid = - torch::zeros({batch_size, n_pts}, torch::CUDA(torch::kFloat)); - - gridding_reverse_grad_kernel<<>>( - scale, n_pts, ptcloud.data(), grid.data(), - grad_ptcloud.data(), grad_grid.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_cuda_forward: %s\n", cudaGetErrorString(err)); - } - return grad_grid; -} diff --git a/other_models/GRNet/extensions_cloudconv/gridding/log.txt b/other_models/GRNet/extensions_cloudconv/gridding/log.txt deleted file mode 100644 index 267a362..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/log.txt +++ /dev/null @@ -1,15 +0,0 @@ -running install -running bdist_egg -running egg_info -writing gridding.egg-info/PKG-INFO -writing dependency_links to gridding.egg-info/dependency_links.txt -writing top-level names to gridding.egg-info/top_level.txt -reading manifest file 'gridding.egg-info/SOURCES.txt' -writing manifest file 'gridding.egg-info/SOURCES.txt' -installing library code to build/bdist.linux-x86_64/egg -running install_lib -running build_ext -building 'gridding' extension -Emitting ninja build file /home/wangyida/Documents/gitfarm/cloudconv/other_models/GRNet/extensions/gridding/build/temp.linux-x86_64-3.6/build.ninja... -Compiling objects... -Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N) diff --git a/other_models/GRNet/extensions_cloudconv/gridding/setup.py b/other_models/GRNet/extensions_cloudconv/gridding/setup.py deleted file mode 100644 index eaaa6ef..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/setup.py +++ /dev/null @@ -1,16 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-11-13 10:51:33 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-02 17:02:16 -# @Email: cshzxie@gmail.com - -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup(name='gridding', - version='2.1.0', - ext_modules=[ - CUDAExtension('gridding', ['gridding_cuda.cpp', 'gridding.cu', 'gridding_reverse.cu']), - ], - cmdclass={'build_ext': BuildExtension}) diff --git a/other_models/GRNet/extensions_cloudconv/gridding/test.py b/other_models/GRNet/extensions_cloudconv/gridding/test.py deleted file mode 100644 index f47cf6b..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding/test.py +++ /dev/null @@ -1,45 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-10 10:48:55 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-26 14:20:42 -# @Email: cshzxie@gmail.com -# -# Note: -# - Replace float -> double, kFloat -> kDouble in gridding.cu and gridding_reverse.cu - -import os -import sys -import torch -import unittest - -from torch.autograd import gradcheck - -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) -from extensions.gridding import GriddingFunction, GriddingReverseFunction - - -class GriddingTestCase(unittest.TestCase): - def test_gridding_reverse_function_4(self): - x = torch.rand(2, 4, 4, 4) - x.requires_grad = True - self.assertTrue(gradcheck(GriddingReverseFunction.apply, [4, x.double().cuda()])) - - def test_gridding_reverse_function_8(self): - x = torch.rand(4, 8, 8, 8) - x.requires_grad = True - self.assertTrue(gradcheck(GriddingReverseFunction.apply, [8, x.double().cuda()])) - - def test_gridding_reverse_function_16(self): - x = torch.rand(1, 16, 16, 16) - x.requires_grad = True - self.assertTrue(gradcheck(GriddingReverseFunction.apply, [16, x.double().cuda()])) - - def test_gridding_function_32pts(self): - x = torch.rand(1, 32, 3) - x.requires_grad = True - self.assertTrue(gradcheck(GriddingFunction.apply, [x.double().cuda()])) - - -if __name__ == '__main__': - unittest.main() diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/__init__.py b/other_models/GRNet/extensions_cloudconv/gridding_loss/__init__.py deleted file mode 100644 index c7d3ad7..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding_loss/__init__.py +++ /dev/null @@ -1,114 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-30 09:56:06 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2020-02-22 19:19:43 -# @Email: cshzxie@gmail.com - -import torch - -import gridding_distance - - -class GriddingDistanceFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud, gt_cloud): - pred_grid, pred_grid_pt_weights, pred_grid_pt_indexes = gridding_distance.forward( - min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud) - # print(pred_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) - # print(pred_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) - # print(pred_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) - gt_grid, gt_grid_pt_weights, gt_grid_pt_indexes = gridding_distance.forward( - min_x, max_x, min_y, max_y, min_z, max_z, gt_cloud) - # print(gt_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) - # print(gt_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) - # print(gt_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) - - ctx.save_for_backward(pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes) - return pred_grid, gt_grid - - @staticmethod - def backward(ctx, grad_pred_grid, grad_gt_grid): - pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes = ctx.saved_tensors - - grad_pred_cloud = gridding_distance.backward(pred_grid_pt_weights, pred_grid_pt_indexes, grad_pred_grid) - # print(grad_pred_cloud.size()) # torch.Size(batch_size, n_pts, 3) - grad_gt_cloud = gridding_distance.backward(gt_grid_pt_weights, gt_grid_pt_indexes, grad_gt_grid) - # print(grad_gt_cloud.size()) # torch.Size(batch_size, n_pts, 3) - - return None, None, None, None, None, None, grad_pred_cloud, grad_gt_cloud - - -class GriddingDistance(torch.nn.Module): - def __init__(self, scale=1): - super(GriddingDistance, self).__init__() - self.scale = scale - - def forward(self, pred_cloud, gt_cloud): - ''' - pred_cloud(b, n_pts1, 3) - gt_cloud(b, n_pts2, 3) - ''' - pred_cloud = pred_cloud * self.scale / 2 - gt_cloud = gt_cloud * self.scale / 2 - - min_pred_x = torch.min(pred_cloud[:, :, 0]) - max_pred_x = torch.max(pred_cloud[:, :, 0]) - min_pred_y = torch.min(pred_cloud[:, :, 1]) - max_pred_y = torch.max(pred_cloud[:, :, 1]) - min_pred_z = torch.min(pred_cloud[:, :, 2]) - max_pred_z = torch.max(pred_cloud[:, :, 2]) - - min_gt_x = torch.min(gt_cloud[:, :, 0]) - max_gt_x = torch.max(gt_cloud[:, :, 0]) - min_gt_y = torch.min(gt_cloud[:, :, 1]) - max_gt_y = torch.max(gt_cloud[:, :, 1]) - min_gt_z = torch.min(gt_cloud[:, :, 2]) - max_gt_z = torch.max(gt_cloud[:, :, 2]) - - min_x = torch.floor(torch.min(min_pred_x, min_gt_x)) - 1 - max_x = torch.ceil(torch.max(max_pred_x, max_gt_x)) + 1 - min_y = torch.floor(torch.min(min_pred_y, min_gt_y)) - 1 - max_y = torch.ceil(torch.max(max_pred_y, max_gt_y)) + 1 - min_z = torch.floor(torch.min(min_pred_z, min_gt_z)) - 1 - max_z = torch.ceil(torch.max(max_pred_z, max_gt_z)) + 1 - - _pred_clouds = torch.split(pred_cloud, 1, dim=0) - _gt_clouds = torch.split(gt_cloud, 1, dim=0) - pred_grids = [] - gt_grids = [] - for pc, gc in zip(_pred_clouds, _gt_clouds): - non_zeros = torch.sum(pc, dim=2).ne(0) - pc = pc[non_zeros].unsqueeze(dim=0) - non_zeros = torch.sum(gc, dim=2).ne(0) - gc = gc[non_zeros].unsqueeze(dim=0) - pred_grid, gt_grid = GriddingDistanceFunction.apply(min_x, max_x, min_y, max_y, min_z, max_z, pc, gc) - pred_grids.append(pred_grid) - gt_grids.append(gt_grid) - - return torch.cat(pred_grids, dim=0).contiguous(), torch.cat(gt_grids, dim=0).contiguous() - - -class GriddingLoss(torch.nn.Module): - def __init__(self, scales=[], alphas=[]): - super(GriddingLoss, self).__init__() - self.scales = scales - self.alphas = alphas - self.gridding_dists = [GriddingDistance(scale=s) for s in scales] - self.l1_loss = torch.nn.L1Loss() - - def forward(self, pred_cloud, gt_cloud): - gridding_loss = None - n_dists = len(self.scales) - - for i in range(n_dists): - alpha = self.alphas[i] - gdist = self.gridding_dists[i] - pred_grid, gt_grid = gdist(pred_cloud, gt_cloud) - - if gridding_loss is None: - gridding_loss = alpha * self.l1_loss(pred_grid, gt_grid) - else: - gridding_loss += alpha * self.l1_loss(pred_grid, gt_grid) - - return gridding_loss diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/__pycache__/__init__.cpython-36.pyc b/other_models/GRNet/extensions_cloudconv/gridding_loss/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index 648536f..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/gridding_loss/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/dist/gridding_distance-1.0.0-py2.7-linux-x86_64.egg b/other_models/GRNet/extensions_cloudconv/gridding_loss/dist/gridding_distance-1.0.0-py2.7-linux-x86_64.egg deleted file mode 100644 index ba5e845..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/gridding_loss/dist/gridding_distance-1.0.0-py2.7-linux-x86_64.egg and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/dist/gridding_distance-1.0.0-py3.6-linux-x86_64.egg b/other_models/GRNet/extensions_cloudconv/gridding_loss/dist/gridding_distance-1.0.0-py3.6-linux-x86_64.egg deleted file mode 100644 index 27266ed..0000000 Binary files a/other_models/GRNet/extensions_cloudconv/gridding_loss/dist/gridding_distance-1.0.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.cu b/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.cu deleted file mode 100644 index f0a9b0f..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.cu +++ /dev/null @@ -1,338 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-12-30 11:35:30 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 15:00:45 - * @Email: cshzxie@gmail.com - */ - -#include -#include -#include -#include - -#define CUDA_NUM_THREADS 512 - -// Computer the number of threads needed in GPU -inline int get_n_threads(int n) { - const int pow_2 = std::log(static_cast(n)) / std::log(2.0); - return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); -} - -__device__ int compute_index( - int offset_x, int offset_y, int offset_z, int len_y, int len_z) { - return offset_x * len_y * len_z + offset_y * len_z + offset_z; -} - -__device__ float compute_weight(float x, float x0) { return 1 - abs(x - x0); } - -__global__ void gridding_dist_kernel(int n_grid_vertices, - int n_pts, - float min_x, - float min_y, - float min_z, - int len_y, - int len_z, - const float *__restrict__ ptcloud, - float *__restrict__ grid_weights, - float *__restrict__ grid_pt_weights, - int *__restrict__ grid_pt_indexes) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - ptcloud += batch_index * n_pts * 3; - grid_weights += batch_index * n_grid_vertices * 8; - grid_pt_weights += batch_index * n_pts * 24; - grid_pt_indexes += batch_index * n_pts * 8; - - for (int j = index; j < n_pts; j += stride) { - float pt_x = ptcloud[j * 3 + 0]; - float pt_y = ptcloud[j * 3 + 1]; - float pt_z = ptcloud[j * 3 + 2]; - - int lower_x = std::floor(pt_x); - int upper_x = std::ceil(pt_x); - if (lower_x == upper_x) { - upper_x += 1; - } - int lower_y = std::floor(pt_y); - int upper_y = std::ceil(pt_y); - if (lower_y == upper_y) { - upper_y += 1; - } - int lower_z = std::floor(pt_z); - int upper_z = std::ceil(pt_z); - if (lower_z == upper_z) { - upper_z += 1; - } - - int lx_offset = lower_x - min_x, ux_offset = upper_x - min_x; - int ly_offset = lower_y - min_y, uy_offset = upper_y - min_y; - int lz_offset = lower_z - min_z, uz_offset = upper_z - min_z; - - // Compute weights and corresponding positions, a loop for 8 points - // LLL -> Lower X, Lower Y, Lower Z - grid_pt_indexes[j * 8 + 0] = - compute_index(lx_offset, ly_offset, lz_offset, len_y, len_z) * 8 + 0; - grid_pt_weights[j * 24 + 0] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 1] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 2] = compute_weight(pt_z, lower_z); - - // LLU -> Lower X, Lower Y, Upper Z - grid_pt_indexes[j * 8 + 1] = - compute_index(lx_offset, ly_offset, uz_offset, len_y, len_z) * 8 + 1; - grid_pt_weights[j * 24 + 3] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 4] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 5] = compute_weight(pt_z, upper_z); - - // LUL -> Lower X, Upper Y, Lower Z - grid_pt_indexes[j * 8 + 2] = - compute_index(lx_offset, uy_offset, lz_offset, len_y, len_z) * 8 + 2; - grid_pt_weights[j * 24 + 6] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 7] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 8] = compute_weight(pt_z, lower_z); - - // LUU -> Lower X, Upper Y, Upper Z - grid_pt_indexes[j * 8 + 3] = - compute_index(lx_offset, uy_offset, uz_offset, len_y, len_z) * 8 + 3; - grid_pt_weights[j * 24 + 9] = compute_weight(pt_x, lower_x); - grid_pt_weights[j * 24 + 10] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 11] = compute_weight(pt_z, upper_z); - - // ULL -> Upper X, Lower Y, Lower Z - grid_pt_indexes[j * 8 + 4] = - compute_index(ux_offset, ly_offset, lz_offset, len_y, len_z) * 8 + 4; - grid_pt_weights[j * 24 + 12] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 13] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 14] = compute_weight(pt_z, lower_z); - - // ULU -> Upper X, Lower Y, Upper Z - grid_pt_indexes[j * 8 + 5] = - compute_index(ux_offset, ly_offset, uz_offset, len_y, len_z) * 8 + 5; - grid_pt_weights[j * 24 + 15] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 16] = compute_weight(pt_y, lower_y); - grid_pt_weights[j * 24 + 17] = compute_weight(pt_z, upper_z); - - // UUL -> Upper X, Upper Y, Lower Z - grid_pt_indexes[j * 8 + 6] = - compute_index(ux_offset, uy_offset, lz_offset, len_y, len_z) * 8 + 6; - grid_pt_weights[j * 24 + 18] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 19] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 20] = compute_weight(pt_z, lower_z); - - // UUU -> Upper X, Upper Y, Upper Z - grid_pt_indexes[j * 8 + 7] = - compute_index(ux_offset, uy_offset, uz_offset, len_y, len_z) * 8 + 7; - grid_pt_weights[j * 24 + 21] = compute_weight(pt_x, upper_x); - grid_pt_weights[j * 24 + 22] = compute_weight(pt_y, upper_y); - grid_pt_weights[j * 24 + 23] = compute_weight(pt_z, upper_z); - } - - __syncthreads(); - - int gvtx_idx = 0; - for (int j = index; j < n_pts; j += stride) { - // LLL -> Lower X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 0]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 0] * - grid_pt_weights[j * 24 + 1] * - grid_pt_weights[j * 24 + 2]); - // LLU -> Lower X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 1]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 3] * - grid_pt_weights[j * 24 + 4] * - grid_pt_weights[j * 24 + 5]); - // LUL -> Lower X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 2]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 6] * - grid_pt_weights[j * 24 + 7] * - grid_pt_weights[j * 24 + 8]); - // LUU -> Lower X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 3]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 9] * - grid_pt_weights[j * 24 + 10] * - grid_pt_weights[j * 24 + 11]); - // ULL -> Upper X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 4]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 12] * - grid_pt_weights[j * 24 + 13] * - grid_pt_weights[j * 24 + 14]); - // ULU -> Upper X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 5]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 15] * - grid_pt_weights[j * 24 + 16] * - grid_pt_weights[j * 24 + 17]); - // UUL -> Upper X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 6]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 18] * - grid_pt_weights[j * 24 + 19] * - grid_pt_weights[j * 24 + 20]); - // UUU -> Upper X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 7]; - atomicAdd(&(grid_weights[gvtx_idx]), grid_pt_weights[j * 24 + 21] * - grid_pt_weights[j * 24 + 22] * - grid_pt_weights[j * 24 + 23]); - } -} - -std::vector gridding_distance_cuda_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud, - cudaStream_t stream) { - int batch_size = ptcloud.size(0); - int n_pts = ptcloud.size(1); - int len_x = max_x - min_x + 1; - int len_y = max_y - min_y + 1; - int len_z = max_z - min_z + 1; - int n_grid_vertices = len_x * len_y * len_z; - - torch::Tensor grid_weights = - torch::zeros({batch_size, n_grid_vertices, 8}, torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_weights = - torch::zeros({batch_size, n_pts, 8, 3}, torch::CUDA(torch::kFloat)); - torch::Tensor grid_pt_indexes = - torch::zeros({batch_size, n_pts, 8}, torch::CUDA(torch::kInt)); - - gridding_dist_kernel<<>>( - n_grid_vertices, n_pts, min_x, min_y, min_z, len_y, len_z, - ptcloud.data(), grid_weights.data(), - grid_pt_weights.data(), grid_pt_indexes.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_distance_cuda_forward: %s\n", - cudaGetErrorString(err)); - } - return {grid_weights, grid_pt_weights, grid_pt_indexes}; -} - -__global__ void gridding_dist_grad_kernel( - int n_grid_vertices, - int n_pts, - const float *__restrict__ grid_pt_weights, - const int *__restrict__ grid_pt_indexes, - const float *__restrict__ grad_grid, - float *__restrict__ grad_ptcloud) { - int batch_index = blockIdx.x; - int index = threadIdx.x; - int stride = blockDim.x; - - grid_pt_weights += batch_index * n_pts * 24; - grid_pt_indexes += batch_index * n_pts * 8; - grad_grid += batch_index * n_grid_vertices * 8; - grad_ptcloud += batch_index * n_pts * 3; - - int gvtx_idx = 0; - float grad_vtx = 0, x_weights = 0, y_weights = 0, z_weights = 0; - for (int j = index; j < n_pts; j += stride) { - // Compute gradient for the corresponding positions, a loop for 8 points - // LLL -> Lower X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 0]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 0]; - y_weights = grid_pt_weights[j * 24 + 1]; - z_weights = grid_pt_weights[j * 24 + 2]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // LLU -> Lower X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 1]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 3]; - y_weights = grid_pt_weights[j * 24 + 4]; - z_weights = grid_pt_weights[j * 24 + 5]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // LUL -> Lower X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 2]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 6]; - y_weights = grid_pt_weights[j * 24 + 7]; - z_weights = grid_pt_weights[j * 24 + 8]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // LUU -> Lower X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 3]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 9]; - y_weights = grid_pt_weights[j * 24 + 10]; - z_weights = grid_pt_weights[j * 24 + 11]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), -grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // ULL -> Upper X, Lower Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 4]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 12]; - y_weights = grid_pt_weights[j * 24 + 13]; - z_weights = grid_pt_weights[j * 24 + 14]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // ULU -> Upper X, Lower Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 5]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 15]; - y_weights = grid_pt_weights[j * 24 + 16]; - z_weights = grid_pt_weights[j * 24 + 17]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), -grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - - // UUL -> Upper X, Upper Y, Lower Z - gvtx_idx = grid_pt_indexes[j * 8 + 6]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 18]; - y_weights = grid_pt_weights[j * 24 + 19]; - z_weights = grid_pt_weights[j * 24 + 20]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), -grad_vtx * x_weights * y_weights); - - // UUU -> Upper X, Upper Y, Upper Z - gvtx_idx = grid_pt_indexes[j * 8 + 7]; - grad_vtx = grad_grid[gvtx_idx]; - x_weights = grid_pt_weights[j * 24 + 21]; - y_weights = grid_pt_weights[j * 24 + 22]; - z_weights = grid_pt_weights[j * 24 + 23]; - atomicAdd(&(grad_ptcloud[j * 3 + 0]), grad_vtx * y_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 1]), grad_vtx * x_weights * z_weights); - atomicAdd(&(grad_ptcloud[j * 3 + 2]), grad_vtx * x_weights * y_weights); - } -} - -torch::Tensor gridding_distance_cuda_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid, - cudaStream_t stream) { - int batch_size = grad_grid.size(0); - int n_grid_vertices = grad_grid.size(1); - int n_pts = grid_pt_indexes.size(1); - - torch::Tensor grad_ptcloud = - torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); - - gridding_dist_grad_kernel<<>>( - n_grid_vertices, n_pts, grid_pt_weights.data(), - grid_pt_indexes.data(), grad_grid.data(), - grad_ptcloud.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Error in gridding_distance_cuda_backward: %s\n", - cudaGetErrorString(err)); - } - return grad_ptcloud; -} diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/PKG-INFO b/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/PKG-INFO deleted file mode 100644 index 5453aab..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: gridding-distance -Version: 1.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/SOURCES.txt b/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/SOURCES.txt deleted file mode 100644 index ad3fcfb..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -gridding_distance.cu -gridding_distance_cuda.cpp -setup.py -gridding_distance.egg-info/PKG-INFO -gridding_distance.egg-info/SOURCES.txt -gridding_distance.egg-info/dependency_links.txt -gridding_distance.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/dependency_links.txt b/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/top_level.txt b/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/top_level.txt deleted file mode 100644 index cdcb49c..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -gridding_distance diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance_cuda.cpp b/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance_cuda.cpp deleted file mode 100644 index 9b8829f..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding_loss/gridding_distance_cuda.cpp +++ /dev/null @@ -1,63 +0,0 @@ -/* - * @Author: Haozhe Xie - * @Date: 2019-12-30 10:59:31 - * @Last Modified by: Haozhe Xie - * @Last Modified time: 2020-06-17 14:52:52 - * @Email: cshzxie@gmail.com - */ - -#include -// #include -#include - -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) - -std::vector gridding_distance_cuda_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud, - cudaStream_t stream); - -torch::Tensor gridding_distance_cuda_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid, - cudaStream_t stream); - -std::vector gridding_distance_forward(float min_x, - float max_x, - float min_y, - float max_y, - float min_z, - float max_z, - torch::Tensor ptcloud) { - CHECK_INPUT(ptcloud); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_distance_cuda_forward(min_x, max_x, min_y, max_y, min_z, - max_z, ptcloud, stream); -} - -torch::Tensor gridding_distance_backward(torch::Tensor grid_pt_weights, - torch::Tensor grid_pt_indexes, - torch::Tensor grad_grid) { - CHECK_INPUT(grid_pt_weights); - CHECK_INPUT(grid_pt_indexes); - CHECK_INPUT(grad_grid); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return gridding_distance_cuda_backward(grid_pt_weights, grid_pt_indexes, - grad_grid, stream); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &gridding_distance_forward, - "Gridding Distance Forward (CUDA)"); - m.def("backward", &gridding_distance_backward, - "Gridding Distance Backward (CUDA)"); -} diff --git a/other_models/GRNet/extensions_cloudconv/gridding_loss/setup.py b/other_models/GRNet/extensions_cloudconv/gridding_loss/setup.py deleted file mode 100644 index f237fd7..0000000 --- a/other_models/GRNet/extensions_cloudconv/gridding_loss/setup.py +++ /dev/null @@ -1,16 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-12-30 11:03:55 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2019-12-30 11:13:39 -# @Email: cshzxie@gmail.com - -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup(name='gridding_distance', - version='1.0.0', - ext_modules=[ - CUDAExtension('gridding_distance', ['gridding_distance_cuda.cpp', 'gridding_distance.cu']), - ], - cmdclass={'build_ext': BuildExtension}) diff --git a/other_models/GRNet/grnet.py b/other_models/GRNet/grnet.py deleted file mode 100644 index 12b8c86..0000000 --- a/other_models/GRNet/grnet.py +++ /dev/null @@ -1,205 +0,0 @@ -# -*- coding: utf-8 -*- -# @Author: Haozhe Xie -# @Date: 2019-09-06 11:35:30 -# @Last Modified by: Haozhe Xie -# @Last Modified time: 2020-02-22 19:20:36 -# @Email: cshzxie@gmail.com - -import torch -import mcubes -import numpy as np - -from GRNet.extensions.gridding import Gridding, GriddingReverse -from GRNet.extensions.cubic_feature_sampling import CubicFeatureSampling - - -class RandomPointSampling(torch.nn.Module): - def __init__(self, n_points): - super(RandomPointSampling, self).__init__() - self.n_points = n_points - - def forward(self, pred_cloud, partial_cloud=None): - if partial_cloud is not None: - pred_cloud = torch.cat([partial_cloud, pred_cloud], dim=1) - - _ptcloud = torch.split(pred_cloud, 1, dim=0) - ptclouds = [] - for p in _ptcloud: - non_zeros = torch.sum(p, dim=2).ne(0) - p = p[non_zeros].unsqueeze(dim=0) - n_pts = p.size(1) - if n_pts < self.n_points: - rnd_idx = torch.cat( - [torch.randint(0, n_pts, (self.n_points, ))]) - else: - rnd_idx = torch.randperm(p.size(1))[:self.n_points] - ptclouds.append(p[:, rnd_idx, :]) - - return torch.cat(ptclouds, dim=0).contiguous() - - -class GRNet(torch.nn.Module): - def __init__(self): - super(GRNet, self).__init__() - self.scale = 64 - self.gridding = Gridding(scale=self.scale) - self.conv1 = torch.nn.Sequential( - torch.nn.Conv3d(1, 32, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(32), torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d( - kernel_size=2)) # Take care of the gridding scale!!! - self.conv2 = torch.nn.Sequential( - torch.nn.Conv3d(32, 64, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(64), torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d(kernel_size=2)) - self.conv3 = torch.nn.Sequential( - torch.nn.Conv3d(64, 128, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(128), torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d(kernel_size=2)) - self.conv4 = torch.nn.Sequential( - torch.nn.Conv3d(128, 256, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(256), torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d(kernel_size=2)) - if self.scale >= 128: - self.conv5 = torch.nn.Sequential( - torch.nn.Conv3d(256, 256, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(256), torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d(kernel_size=2)) - if self.scale >= 256: - self.conv6 = torch.nn.Sequential( - torch.nn.Conv3d(256, 256, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(256), torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d(kernel_size=2)) - """ - self.fc5 = torch.nn.Sequential( - torch.nn.Linear(1024, 2048), torch.nn.ReLU()) - """ - self.fc5 = torch.nn.Sequential( - torch.nn.Linear(16384, 2048), torch.nn.ReLU()) - self.fc6 = torch.nn.Sequential( - torch.nn.Linear(2048, 16384), torch.nn.ReLU()) - if self.scale >= 128: - self.dconv5 = torch.nn.Sequential( - torch.nn.ConvTranspose3d( - 256, 256, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(256), torch.nn.ReLU()) - if self.scale >= 256: - self.dconv6 = torch.nn.Sequential( - torch.nn.ConvTranspose3d( - 256, 256, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(256), torch.nn.ReLU()) - self.dconv7 = torch.nn.Sequential( - torch.nn.ConvTranspose3d( - 256, 128, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(128), torch.nn.ReLU()) - self.dconv8 = torch.nn.Sequential( - torch.nn.ConvTranspose3d( - 128, 64, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(64), torch.nn.ReLU()) - self.dconv9 = torch.nn.Sequential( - torch.nn.ConvTranspose3d( - 64, 32, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(32), torch.nn.ReLU()) - self.dconv10 = torch.nn.Sequential( - torch.nn.ConvTranspose3d( - 32, 1, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(1), torch.nn.ReLU()) - self.gridding_rev = GriddingReverse(scale=self.scale) - self.point_sampling = RandomPointSampling(n_points=2048) - self.feature_sampling = CubicFeatureSampling() - from activations.siren import SirenNet - """ - self.sirennet = SirenNet(dim_in=1792, dim_hidden=256, dim_out=112, num_layers=5) - """ - self.fc11 = torch.nn.Sequential( - torch.nn.Linear(1792, 1792), torch.nn.ReLU()) - self.fc12 = torch.nn.Sequential( - torch.nn.Linear(1792, 448), torch.nn.ReLU()) - self.fc13 = torch.nn.Sequential( - torch.nn.Linear(448, 112), torch.nn.ReLU()) - self.fc14 = torch.nn.Linear(112, 3 * 8) - # self.fc14_seg = torch.nn.Linear(112, 12 * 8) - self.unet = 1.0 - - def forward(self, partial_cloud, global_feature=[]): - # print(partial_cloud.size()) # torch.Size([batch_size, 2048, 3]) - pt_features_64_l = self.gridding(partial_cloud).view( - -1, 1, self.scale, self.scale, self.scale) - # print(pt_features_64_l.size()) # torch.Size([batch_size, 1, self.scale, self.scale, self.scale]) - pt_features_32_l = self.conv1(pt_features_64_l) - # print(pt_features_32_l.size()) # torch.Size([batch_size, 32, 32, 32, 32]) - pt_features_16_l = self.conv2(pt_features_32_l) - # print(pt_features_16_l.size()) # torch.Size([batch_size, 64, 16, 16, 16]) - pt_features_8_l = self.conv3(pt_features_16_l) - # print(pt_features_8_l.size()) # torch.Size([batch_size, 124, 4, 4, 8]) - pt_features_4_l = self.conv4(pt_features_8_l) - # print(pt_features_4_l.size()) # torch.Size([batch_size, 256, 4, 4, 4]) - if self.scale >= 128: - pt_features_4_l = self.conv5(pt_features_4_l) - if self.scale >= 256: - pt_features_4_l = self.conv6(pt_features_4_l) - if global_feature != []: - features = self.fc5(global_feature) - else: - features = self.fc5(pt_features_4_l.view(-1, 16384)) - # print(features.size()) # torch.Size([batch_size, 2048]) - pt_features_4_r = self.fc6(features).view( - -1, 256, 4, 4, 4) + pt_features_4_l * self.unet - # print(pt_features_4_r.size()) # torch.Size([batch_size, 256, 4, 4, 4]) - if self.scale >= 128: - pt_features_4_r = self.dconv5(pt_features_4_r) - if self.scale >= 256: - pt_features_4_r = self.dconv5(pt_features_4_r) - pt_features_8_r = self.dconv7( - pt_features_4_r) + pt_features_8_l * self.unet - # print(pt_features_8_r.size()) # torch.Size([batch_size, 124, 4, 4, 8]) - pt_features_16_r = self.dconv8( - pt_features_8_r) + pt_features_16_l * self.unet - # print(pt_features_16_r.size()) # torch.Size([batch_size, 64, 16, 16, 16]) - pt_features_32_r = self.dconv9( - pt_features_16_r) + pt_features_32_l * self.unet - # print(pt_features_32_r.size()) # torch.Size([batch_size, 32, 32, 32, 32]) - pt_features_64_r = self.dconv10( - pt_features_32_r) + pt_features_64_l * self.unet - # print(pt_features_64_r.size()) # torch.Size([batch_size, 1, 64, 64, 64]) - sparse_cloud = self.gridding_rev(pt_features_64_r.squeeze(dim=1)) - # print(sparse_cloud.size()) # torch.Size([batch_size, 262144, 3]) - sparse_cloud_samp = self.point_sampling(sparse_cloud) - # sparse_cloud_samp = self.point_sampling(sparse_cloud, partial_cloud) - # print(sparse_cloud_samp.size()) # torch.Size([batch_size, 2048, 3]) - point_features_32 = self.feature_sampling(sparse_cloud_samp, - pt_features_32_r).view( - -1, 2048, 256) - # print(point_features_32.size()) # torch.Size([batch_size, 2048, 256]) - point_features_16 = self.feature_sampling(sparse_cloud_samp, - pt_features_16_r).view( - -1, 2048, 512) - # print(point_features_16.size()) # torch.Size([batch_size, 2048, 512]) - point_features_8 = self.feature_sampling(sparse_cloud_samp, - pt_features_8_r).view( - -1, 2048, 1024) - # print(point_features_8.size()) # torch.Size([batch_size, 2048, 1024]) - # freq = self.fourier_map(sparse_cloud_samp.transpose(1, 2)).transpose(1, 2) - point_features = torch.cat( - [point_features_32, point_features_16, point_features_8], dim=2) - # print(point_features.size()) # torch.Size([batch_size, 2048, 1792]) - point_features = self.fc11(point_features) - # print(point_features.size()) # torch.Size([batch_size, 2048, 1792]) - point_features = self.fc12(point_features) - # print(point_features.size()) # torch.Size([batch_size, 2048, 448]) - point_features = self.fc13(point_features) - # print(point_features.size()) # torch.Size([batch_size, 2048, 112]) - """ - point_features = self.sirennet(point_features) - """ - point_f_offset = self.fc14(point_features).view(-1, 16384, 3) - # point_f_segs = self.fc14_seg(point_features).view(-1, 16384, 12) - # print(point_features.size()) # torch.Size([batch_size, 16384, 3]) - dense_cloud = sparse_cloud_samp.unsqueeze(dim=2).repeat(1, 1, 8, 1).view( - -1, 16384, 3) + point_f_offset - """ - dense_cloud = point_f_offset - """ - # print(dense_cloud.size()) # torch.Size([batch_size, 16384, 3]) - - return sparse_cloud, dense_cloud, pt_features_64_r diff --git a/other_models/MSN/.DS_Store b/other_models/MSN/.DS_Store deleted file mode 100644 index 3a6893b..0000000 Binary files a/other_models/MSN/.DS_Store and /dev/null differ diff --git a/other_models/MSN/MDS/MDS.cpp b/other_models/MSN/MDS/MDS.cpp deleted file mode 100644 index 902d0fb..0000000 --- a/other_models/MSN/MDS/MDS.cpp +++ /dev/null @@ -1,127 +0,0 @@ -#include -#include - -#ifndef AT_CHECK -#define AT_CHECK TORCH_CHECK -#endif - -void gather_points_kernel_wrapper(int b, int c, int n, int npoints, - const float *points, const int *idx, - float *out); -void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, - const float *grad_out, const int *idx, - float *grad_points); - -void minimum_density_sampling_kernel_wrapper(int b, int n, int m, - const float *dataset, float *temp, - int *idxs, float *mean_mst_length); - - -#define CUDA_CHECK_ERRORS() \ - do { \ - cudaError_t err = cudaGetLastError(); \ - if (cudaSuccess != err) { \ - fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ - cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ - __FILE__); \ - exit(-1); \ - } \ - } while (0) - -#define CHECK_CUDA(x) \ - do { \ - AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ - } while (0) - -#define CHECK_CONTIGUOUS(x) \ - do { \ - AT_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ - } while (0) - -#define CHECK_IS_INT(x) \ - do { \ - AT_CHECK(x.scalar_type() == at::ScalarType::Int, \ - #x " must be an int tensor"); \ - } while (0) - -#define CHECK_IS_FLOAT(x) \ - do { \ - AT_CHECK(x.scalar_type() == at::ScalarType::Float, \ - #x " must be a float tensor"); \ - } while (0) - - -at::Tensor gather_points(at::Tensor points, at::Tensor idx) { - CHECK_CONTIGUOUS(points); - CHECK_CONTIGUOUS(idx); - CHECK_IS_FLOAT(points); - CHECK_IS_INT(idx); - - if (points.type().is_cuda()) { - CHECK_CUDA(idx); - } - - at::Tensor output = - torch::zeros({points.size(0), points.size(1), idx.size(1)}, - at::device(points.device()).dtype(at::ScalarType::Float)); - - if (points.type().is_cuda()) { - gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), - idx.size(1), points.data(), - idx.data(), output.data()); - } else { - AT_CHECK(false, "CPU not supported"); - } - - return output; -} - -at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, - const int n) { - CHECK_CONTIGUOUS(grad_out); - CHECK_CONTIGUOUS(idx); - CHECK_IS_FLOAT(grad_out); - CHECK_IS_INT(idx); - - if (grad_out.type().is_cuda()) { - CHECK_CUDA(idx); - } - - at::Tensor output = - torch::zeros({grad_out.size(0), grad_out.size(1), n}, - at::device(grad_out.device()).dtype(at::ScalarType::Float)); - - if (grad_out.type().is_cuda()) { - gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, - idx.size(1), grad_out.data(), - idx.data(), output.data()); - } else { - AT_CHECK(false, "CPU not supported"); - } - - return output; -} -at::Tensor minimum_density_sampling(at::Tensor points, const int nsamples, at::Tensor mean_mst_length, at::Tensor output) { - CHECK_CONTIGUOUS(points); - CHECK_IS_FLOAT(points); - - at::Tensor tmp = - torch::zeros({points.size(0), points.size(1)}, - at::device(points.device()).dtype(at::ScalarType::Float)); - - if (points.type().is_cuda()) { - minimum_density_sampling_kernel_wrapper( - points.size(0), points.size(1), nsamples, points.data(), - tmp.data(), output.data(), mean_mst_length.data()); - } else { - AT_CHECK(false, "CPU not supported"); - } - - return output; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("minimum_density_sampling", &minimum_density_sampling, "minimum_density_sampling (CUDA)"); - m.def("gather_points_grad", &gather_points_grad, "gather_points_grad (CUDA)"); - m.def("gather_points", &gather_points, "gather_points (CUDA)"); -} diff --git a/other_models/MSN/MDS/MDS.egg-info/PKG-INFO b/other_models/MSN/MDS/MDS.egg-info/PKG-INFO deleted file mode 100644 index 4ec3a97..0000000 --- a/other_models/MSN/MDS/MDS.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: MDS -Version: 0.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/MSN/MDS/MDS.egg-info/SOURCES.txt b/other_models/MSN/MDS/MDS.egg-info/SOURCES.txt deleted file mode 100644 index 4815378..0000000 --- a/other_models/MSN/MDS/MDS.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -MDS.cpp -MDS_cuda.cu -setup.py -MDS.egg-info/PKG-INFO -MDS.egg-info/SOURCES.txt -MDS.egg-info/dependency_links.txt -MDS.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/MSN/MDS/MDS.egg-info/dependency_links.txt b/other_models/MSN/MDS/MDS.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/MSN/MDS/MDS.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/MSN/MDS/MDS.egg-info/top_level.txt b/other_models/MSN/MDS/MDS.egg-info/top_level.txt deleted file mode 100644 index 0ed512a..0000000 --- a/other_models/MSN/MDS/MDS.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -MDS diff --git a/other_models/MSN/MDS/MDS_cuda.cu b/other_models/MSN/MDS/MDS_cuda.cu deleted file mode 100644 index 62c168d..0000000 --- a/other_models/MSN/MDS/MDS_cuda.cu +++ /dev/null @@ -1,271 +0,0 @@ -#include -#include -#include -#include - -#define TOTAL_THREADS 1024 - -inline int opt_n_threads(int work_size) { - const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); - - return max(min(1 << pow_2, TOTAL_THREADS), 1); -} - - -#define CUDA_CHECK_ERRORS() \ - do { \ - cudaError_t err = cudaGetLastError(); \ - if (cudaSuccess != err) { \ - fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ - cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ - __FILE__); \ - exit(-1); \ - } \ - } while (0) - - -// input: points(b, c, n) idx(b, m) -// output: out(b, c, m) -__global__ void gather_points_kernel(int b, int c, int n, int m, - const float *__restrict__ points, - const int *__restrict__ idx, - float *__restrict__ out) { - for (int i = blockIdx.x; i < b; i += gridDim.x) { - for (int l = blockIdx.y; l < c; l += gridDim.y) { - for (int j = threadIdx.x; j < m; j += blockDim.x) { - int a = idx[i * m + j]; - out[(i * c + l) * m + j] = points[(i * c + l) * n + a]; - } - } - } -} - -void gather_points_kernel_wrapper(int b, int c, int n, int npoints, - const float *points, const int *idx, - float *out) { - gather_points_kernel<<>>(b, c, n, npoints, - points, idx, out); - - CUDA_CHECK_ERRORS(); -} - -// input: grad_out(b, c, m) idx(b, m) -// output: grad_points(b, c, n) -__global__ void gather_points_grad_kernel(int b, int c, int n, int m, - const float *__restrict__ grad_out, - const int *__restrict__ idx, - float *__restrict__ grad_points) { - for (int i = blockIdx.x; i < b; i += gridDim.x) { - for (int l = blockIdx.y; l < c; l += gridDim.y) { - for (int j = threadIdx.x; j < m; j += blockDim.x) { - int a = idx[i * m + j]; - //atomicAdd(grad_points + (i * c + l) * n + a, - // grad_out[(i * c + l) * m + j]); - grad_points[(i * c + l) * n + a] += grad_out[(i * c + l) * m + j]; - } - } - } -} - -void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, - const float *grad_out, const int *idx, - float *grad_points) { - gather_points_grad_kernel<<>>( - b, c, n, npoints, grad_out, idx, grad_points); - - CUDA_CHECK_ERRORS(); -} - -__device__ void __update(float *__restrict__ cnt, int *__restrict__ cnt_i, - int idx1, int idx2) { - const float v1 = cnt[idx1], v2 = cnt[idx2]; - const int i1 = cnt_i[idx1], i2 = cnt_i[idx2]; - cnt[idx1] = min(v1, v2); - cnt_i[idx1] = v2 < v1 ? i2 : i1; -} - -// Input dataset: (b, n, 3), tmp: (b, n) -// Ouput idxs (b, m) -template -__global__ void minimum_density_sampling_kernel ( - int b, int n, int m, const float *__restrict__ dataset, - float *__restrict__ temp, int *__restrict__ idxs, float *__restrict__ mean_mst_length) { - if (m <= 0) return; - __shared__ float cnt[block_size]; - __shared__ int cnt_i[block_size]; - - int batch_index = blockIdx.x; - dataset += batch_index * n * 3; - temp += batch_index * n; - idxs += batch_index * m; - - int tid = threadIdx.x; - const int stride = block_size; - - int old = 0; - if (threadIdx.x == 0) { - idxs[0] = old; - temp[old] = 1e9; - } - - __syncthreads(); - float t = 5.0 * mean_mst_length[batch_index] * mean_mst_length[batch_index]; - - float x1 = dataset[old * 3 + 0]; - float y1 = dataset[old * 3 + 1]; - float z1 = dataset[old * 3 + 2]; - for (int j = 1; j < m; j++) { - int besti = 0; - float best = 1e9; - for (int k = tid; k < n; k += stride) { - float x2, y2, z2; - x2 = dataset[k * 3 + 0]; - y2 = dataset[k * 3 + 1]; - z2 = dataset[k * 3 + 2]; - - float d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); - - temp[k] += k < 8192 ? exp(-d / t) : exp(-d / t) * 2.0; - besti = temp[k] < best ? k : besti; - best = temp[k] < best ? temp[k] : best; - } - cnt[tid] = best; - cnt_i[tid] = besti; - __syncthreads(); - - - if (block_size >= 1024) { - if (tid < 512) { - __update(cnt, cnt_i, tid, tid + 512); - } - __syncthreads(); - } - if (block_size >= 512) { - if (tid < 256) { - __update(cnt, cnt_i, tid, tid + 256); - } - __syncthreads(); - } - if (block_size >= 256) { - if (tid < 128) { - __update(cnt, cnt_i, tid, tid + 128); - } - __syncthreads(); - } - if (block_size >= 128) { - if (tid < 64) { - __update(cnt, cnt_i, tid, tid + 64); - } - __syncthreads(); - } - if (block_size >= 64) { - if (tid < 32) { - __update(cnt, cnt_i, tid, tid + 32); - } - __syncthreads(); - } - if (block_size >= 32) { - if (tid < 16) { - __update(cnt, cnt_i, tid, tid + 16); - } - __syncthreads(); - } - if (block_size >= 16) { - if (tid < 8) { - __update(cnt, cnt_i, tid, tid + 8); - } - __syncthreads(); - } - if (block_size >= 8) { - if (tid < 4) { - __update(cnt, cnt_i, tid, tid + 4); - } - __syncthreads(); - } - if (block_size >= 4) { - if (tid < 2) { - __update(cnt, cnt_i, tid, tid + 2); - } - __syncthreads(); - } - if (block_size >= 2) { - if (tid < 1) { - __update(cnt, cnt_i, tid, tid + 1); - } - __syncthreads(); - } - - old = cnt_i[0]; - if (tid == 0) { - idxs[j] = old; - temp[old] = 1e9; - } - - x1 = dataset[old * 3 + 0]; - y1 = dataset[old * 3 + 1]; - z1 = dataset[old * 3 + 2]; - - } -} - -void minimum_density_sampling_kernel_wrapper(int b, int n, int m, - const float *dataset, float *temp, - int *idxs, float * mean_mst_length) { - unsigned int n_threads = opt_n_threads(n); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - switch (n_threads) { - case 1024: - minimum_density_sampling_kernel<1024> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 512: - minimum_density_sampling_kernel<512> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 256: - minimum_density_sampling_kernel<256> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 128: - minimum_density_sampling_kernel<128> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 64: - minimum_density_sampling_kernel<64> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 32: - minimum_density_sampling_kernel<32> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 16: - minimum_density_sampling_kernel<16> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 8: - minimum_density_sampling_kernel<8> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 4: - minimum_density_sampling_kernel<4> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 2: - minimum_density_sampling_kernel<2> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - case 1: - minimum_density_sampling_kernel<1> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - break; - default: - minimum_density_sampling_kernel<1024> - <<>>(b, n, m, dataset, temp, idxs, mean_mst_length); - } - - CUDA_CHECK_ERRORS(); -} diff --git a/other_models/MSN/MDS/MDS_module.py b/other_models/MSN/MDS/MDS_module.py deleted file mode 100644 index 824fb85..0000000 --- a/other_models/MSN/MDS/MDS_module.py +++ /dev/null @@ -1,75 +0,0 @@ -import torch -from torch import nn -from torch.autograd import Function -import MDS - -class MinimumDensitySampling(Function): - @staticmethod - def forward(ctx, xyz, npoint, mean_mst_length): - # type: (Any, torch.Tensor, int) -> torch.Tensor - r""" - Uses iterative radius point sampling to select a set of npoint features that have the largest - minimum distance - - Parameters - ---------- - xyz : torch.Tensor - (B, N, 3) tensor where N > npoint - npoint : int32 - number of features in the sampled set - mean_mst_length : torch.Tensor - (B) the average edge length from expansion penalty module - - Returns - ------- - torch.Tensor - (B, npoint) tensor containing the set - """ - idx = torch.zeros(xyz.shape[0], npoint, requires_grad= False, device='cuda', dtype=torch.int32).contiguous() - MDS.minimum_density_sampling(xyz, npoint, mean_mst_length, idx) - return idx - - @staticmethod - def backward(grad_idx, a=None): - return None, None, None - - -minimum_density_sample = MinimumDensitySampling.apply - - -class GatherOperation(Function): - @staticmethod - def forward(ctx, features, idx): - # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor - r""" - - Parameters - ---------- - features : torch.Tensor - (B, C, N) tensor - - idx : torch.Tensor - (B, npoint) tensor of the features to gather - - Returns - ------- - torch.Tensor - (B, C, npoint) tensor - """ - - _, C, N = features.size() - - ctx.for_backwards = (idx, C, N) - - return MDS.gather_points(features, idx) - - @staticmethod - def backward(ctx, grad_out): - idx, C, N = ctx.for_backwards - - grad_features = MDS.gather_points_grad(grad_out.contiguous(), idx, N) - return grad_features, None - - -gather_operation = GatherOperation.apply - diff --git a/other_models/MSN/MDS/__pycache__/MDS_module.cpython-36.pyc b/other_models/MSN/MDS/__pycache__/MDS_module.cpython-36.pyc deleted file mode 100644 index ec551f0..0000000 Binary files a/other_models/MSN/MDS/__pycache__/MDS_module.cpython-36.pyc and /dev/null differ diff --git a/other_models/MSN/MDS/dist/MDS-0.0.0-py3.6-linux-x86_64.egg b/other_models/MSN/MDS/dist/MDS-0.0.0-py3.6-linux-x86_64.egg deleted file mode 100644 index fe6d4de..0000000 Binary files a/other_models/MSN/MDS/dist/MDS-0.0.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/other_models/MSN/MDS/setup.py b/other_models/MSN/MDS/setup.py deleted file mode 100644 index 31de885..0000000 --- a/other_models/MSN/MDS/setup.py +++ /dev/null @@ -1,14 +0,0 @@ -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup( - name='MDS', - ext_modules=[ - CUDAExtension('MDS', [ - 'MDS_cuda.cu', - 'MDS.cpp', - ]), - ], - cmdclass={ - 'build_ext': BuildExtension - }) \ No newline at end of file diff --git a/other_models/MSN/__pycache__/msn.cpython-36.pyc b/other_models/MSN/__pycache__/msn.cpython-36.pyc deleted file mode 100644 index 5d40e64..0000000 Binary files a/other_models/MSN/__pycache__/msn.cpython-36.pyc and /dev/null differ diff --git a/other_models/MSN/expansion_penalty/__pycache__/expansion_penalty_module.cpython-36.pyc b/other_models/MSN/expansion_penalty/__pycache__/expansion_penalty_module.cpython-36.pyc deleted file mode 100644 index da305eb..0000000 Binary files a/other_models/MSN/expansion_penalty/__pycache__/expansion_penalty_module.cpython-36.pyc and /dev/null differ diff --git a/other_models/MSN/expansion_penalty/dist/expansion_penalty-0.0.0-py3.6-linux-x86_64.egg b/other_models/MSN/expansion_penalty/dist/expansion_penalty-0.0.0-py3.6-linux-x86_64.egg deleted file mode 100644 index 8f62aa0..0000000 Binary files a/other_models/MSN/expansion_penalty/dist/expansion_penalty-0.0.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/other_models/MSN/expansion_penalty/expansion_penalty.cpp b/other_models/MSN/expansion_penalty/expansion_penalty.cpp deleted file mode 100644 index bb7adc7..0000000 --- a/other_models/MSN/expansion_penalty/expansion_penalty.cpp +++ /dev/null @@ -1,20 +0,0 @@ -#include -#include - -int expansion_penalty_cuda_forward(at::Tensor xyz, int primitive_size, at::Tensor father, at::Tensor dist, double alpha, at::Tensor neighbor, at::Tensor cost, at::Tensor mean_mst_length); - -int expansion_penalty_cuda_backward(at::Tensor xyz, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx); - -int expansion_penalty_forward(at::Tensor xyz, int primitive_size, at::Tensor father, at::Tensor dist, double alpha, at::Tensor neighbor, at::Tensor cost, at::Tensor mean_mst_length) { - return expansion_penalty_cuda_forward(xyz, primitive_size, father, dist, alpha, neighbor, cost, mean_mst_length); -} - -int expansion_penalty_backward(at::Tensor xyz, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx) { - - return expansion_penalty_cuda_backward(xyz, gradxyz, graddist, idx); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &expansion_penalty_forward, "expansion_penalty forward (CUDA)"); - m.def("backward", &expansion_penalty_backward, "expansion_penalty backward (CUDA)"); -} \ No newline at end of file diff --git a/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/PKG-INFO b/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/PKG-INFO deleted file mode 100644 index 267ea3a..0000000 --- a/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: expansion-penalty -Version: 0.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/SOURCES.txt b/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/SOURCES.txt deleted file mode 100644 index 0c15831..0000000 --- a/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/SOURCES.txt +++ /dev/null @@ -1,7 +0,0 @@ -expansion_penalty.cpp -expansion_penalty_cuda.cu -setup.py -expansion_penalty.egg-info/PKG-INFO -expansion_penalty.egg-info/SOURCES.txt -expansion_penalty.egg-info/dependency_links.txt -expansion_penalty.egg-info/top_level.txt \ No newline at end of file diff --git a/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/dependency_links.txt b/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/top_level.txt b/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/top_level.txt deleted file mode 100644 index 5055922..0000000 --- a/other_models/MSN/expansion_penalty/expansion_penalty.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -expansion_penalty diff --git a/other_models/MSN/expansion_penalty/expansion_penalty_cuda.cu b/other_models/MSN/expansion_penalty/expansion_penalty_cuda.cu deleted file mode 100644 index e163b20..0000000 --- a/other_models/MSN/expansion_penalty/expansion_penalty_cuda.cu +++ /dev/null @@ -1,198 +0,0 @@ -#include -#include - -#include -#include - -__global__ void calc_penalty(int b, int n, int primitive_size, const float * xyz, int * idx, float * dist, float alpha, int * neighbor, float * cost, float * mean_mst_length) { - const int batch = 512; // primitive_size should be less than 512 - __shared__ float xyz_buf[batch * 3]; - __shared__ bool vis[batch]; - __shared__ float cur_dis[batch]; - __shared__ int cur_idx[batch]; - __shared__ float min_dis[batch]; - __shared__ int min_idx[batch]; - __shared__ float sum_dis[batch]; - __shared__ int cnt[batch]; - __shared__ int degree[batch]; - - for (int i = blockIdx.x; i < b; i += gridDim.x) { - vis[threadIdx.x] = false; - cur_dis[threadIdx.x] = 1e9; - cnt[threadIdx.x] = 0; - degree[threadIdx.x] = 0; - - for (int j = threadIdx.x; j < primitive_size * 3; j += blockDim.x) { - xyz_buf[j] = xyz[(i * n + blockIdx.y * primitive_size) * 3 + j]; - } - __syncthreads(); - - __shared__ int last; - __shared__ float x_last; - __shared__ float y_last; - __shared__ float z_last; - - if (threadIdx.x == 0) { - vis[0] = true; - sum_dis[0] = 0; - last = 0; - x_last = xyz_buf[0]; - y_last = xyz_buf[1]; - z_last = xyz_buf[2]; - } - __syncthreads(); - - for (int j = 0; j < primitive_size - 1; j++) { - if (vis[threadIdx.x] == false) { - float delta_x = xyz_buf[threadIdx.x * 3 + 0] - x_last; - float delta_y = xyz_buf[threadIdx.x * 3 + 1] - y_last; - float delta_z = xyz_buf[threadIdx.x * 3 + 2] - z_last; - float d = sqrtf(delta_x * delta_x + delta_y * delta_y + delta_z * delta_z); - - if (d < cur_dis[threadIdx.x]) { - cur_dis[threadIdx.x] = d; - cur_idx[threadIdx.x] = last; - } - min_dis[threadIdx.x] = cur_dis[threadIdx.x]; - } - else { - min_dis[threadIdx.x] = 1e9; - } - min_idx[threadIdx.x] = threadIdx.x; - __syncthreads(); - - int stride = 1; - while(stride <= primitive_size / 2) { - int index = (threadIdx.x + 1) * stride * 2 - 1; - if(index < primitive_size && min_dis[index - stride] < min_dis[index]) { - min_dis[index] = min_dis[index - stride]; - min_idx[index] = min_idx[index - stride]; - } - stride = stride * 2; - __syncthreads(); - } - __syncthreads(); - - if (threadIdx.x == primitive_size - 1) { - last = min_idx[threadIdx.x]; - int u = cur_idx[last]; - vis[last] = true; - x_last = xyz_buf[last * 3 + 0]; - y_last = xyz_buf[last * 3 + 1]; - z_last = xyz_buf[last * 3 + 2]; - - cnt[last] += 1; - degree[last] += 1; - neighbor[(i * n + blockIdx.y * primitive_size + last) * 512 + cnt[last]] = u; - cost[(i * n + blockIdx.y * primitive_size + last) * 512 + cnt[last]] = cur_dis[last]; - cnt[u] += 1; - degree[u] += 1; - neighbor[(i * n + blockIdx.y * primitive_size + u) * 512 + cnt[u]] = last; - cost[(i * n + blockIdx.y * primitive_size + u) * 512 + cnt[u]] = cur_dis[last]; - - if (cnt[last] >= 512 || cnt[u] >= 512) { - printf("MST Error: Too many neighbors! %d %d %d %d\n", cnt[last], cnt[u], last, u); - } - - sum_dis[last] = cur_dis[last]; - } - __syncthreads(); - } - - __syncthreads(); - int stride = 1; - while(stride <= primitive_size / 2) { - int index = (threadIdx.x + 1) * stride * 2 - 1; - if (index < primitive_size) - sum_dis[index] += sum_dis[index - stride]; - stride = stride * 2; - __syncthreads(); - } - __syncthreads(); - - __shared__ float mean_dis; - if (threadIdx.x == 0) { - mean_dis = sum_dis[primitive_size - 1] / (primitive_size - 1); - atomicAdd(&mean_mst_length[i], mean_dis); - } - - dist[i * n + blockIdx.y * primitive_size + threadIdx.x] = 0; - idx[i * n + blockIdx.y * primitive_size + threadIdx.x] = -1; - __syncthreads(); - - while (true) { - __shared__ int flag; - flag = 0; - int tmp = cnt[threadIdx.x]; - __syncthreads(); - if (tmp == 1) { - atomicAdd(&flag, 1); - for (int j = 1; j <= degree[threadIdx.x]; j++) { - int u = neighbor[(i * n + blockIdx.y * primitive_size + threadIdx.x) * 512 + j]; - if (cnt[u] > 1 || (cnt[u] == 1 && threadIdx.x > u)) { - float c = cost[(i * n + blockIdx.y * primitive_size + threadIdx.x) * 512 + j]; - atomicAdd(&cnt[threadIdx.x], -1); - atomicAdd(&cnt[u], -1); - if (c > mean_dis * alpha) { - dist[i * n + blockIdx.y * primitive_size + threadIdx.x] = c; - idx[i * n + blockIdx.y * primitive_size + threadIdx.x] = blockIdx.y * primitive_size + u; - } - } - } - } - __syncthreads(); - if (flag == 0) break; - __syncthreads(); - } - __syncthreads(); - } -} - -int expansion_penalty_cuda_forward(at::Tensor xyz, int primitive_size, at::Tensor idx, at::Tensor dist, double alpha, at::Tensor neighbor, at::Tensor cost, at::Tensor mean_mst_length) { - - const auto batch_size = xyz.size(0); - const auto n = xyz.size(1); - - calc_penalty<<>>(batch_size, n, primitive_size, xyz.data(), idx.data(), dist.data(), - alpha, neighbor.data(), cost.data(), mean_mst_length.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd Output: %s\n", cudaGetErrorString(err)); - return 0; - } - return 1; -} - -__global__ void calc_grad(int b, int n, const float * xyz, const float * grad_dist, const int * idx, float * grad_xyz) { - for (int i = blockIdx.x; i < b; i += gridDim.x) { - for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; j += blockDim.x * gridDim.y) - if (idx[i * n + j] != -1) { - float x1 = xyz[(i * n + j) * 3 + 0]; - float y1 = xyz[(i * n + j) * 3 + 1]; - float z1 = xyz[(i * n + j) * 3 + 2]; - int j2 = idx[i * n + j]; - float x2 = xyz[(i * n + j2) * 3 + 0]; - float y2 = xyz[(i * n + j2) * 3 + 1]; - float z2 = xyz[(i * n + j2) * 3 + 2]; - float g = grad_dist[i * n + j] * 2; - atomicAdd(&(grad_xyz[(i * n + j) * 3 + 0]), g * (x1 - x2)); - atomicAdd(&(grad_xyz[(i * n + j) * 3 + 1]), g * (y1 - y2)); - atomicAdd(&(grad_xyz[(i * n + j) * 3 + 2]), g * (z1 - z2)); - } - } -} - -int expansion_penalty_cuda_backward(at::Tensor xyz, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx) { - const auto batch_size = xyz.size(0); - const auto n = xyz.size(1); - - calc_grad<<>>(batch_size, n, xyz.data(), graddist.data(), idx.data(), gradxyz.data()); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); - return 0; - } - return 1; -} diff --git a/other_models/MSN/expansion_penalty/expansion_penalty_module.py b/other_models/MSN/expansion_penalty/expansion_penalty_module.py deleted file mode 100644 index 4213077..0000000 --- a/other_models/MSN/expansion_penalty/expansion_penalty_module.py +++ /dev/null @@ -1,63 +0,0 @@ -# Expansion penalty module (based on minimum spanning tree) -# author: Minghua Liu - -# Input: -# xyz: [#batch, #point] -# primitive_size: int, the number of points of sampled points for each surface elements, which should be no greater than 512 -# in each point cloud, the points from the same surface element should be successive -# alpha: float, > 1, only penalize those edges whose length are greater than (alpha * mean_length) - -#Output: -# dist: [#batch, #point], if the point u is penalized then dist[u] is the distance between u and its neighbor in the MST, otherwise dist[u] is 0 -# assignment: [#batch, #point], if the point u is penalized then assignment[u] is its neighbor in the MST, otherwise assignment[u] is -1 -# mean_mst_length: [#batch], the average length of the edeges in each point clouds - - -import time -import numpy as np -import torch -from torch import nn -from torch.autograd import Function -import expansion_penalty - -# GPU tensors only -class expansionPenaltyFunction(Function): - @staticmethod - def forward(ctx, xyz, primitive_size, alpha): - assert(primitive_size <= 512) - batchsize, n, _ = xyz.size() - assert(n % primitive_size == 0) - xyz = xyz.contiguous().float().cuda() - dist = torch.zeros(batchsize, n, device='cuda').contiguous() - assignment = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous() - 1 - neighbor = torch.zeros(batchsize, n * 512, device='cuda', dtype=torch.int32).contiguous() - cost = torch.zeros(batchsize, n * 512, device='cuda').contiguous() - mean_mst_length = torch.zeros(batchsize, device='cuda').contiguous() - expansion_penalty.forward(xyz, primitive_size, assignment, dist, alpha, neighbor, cost, mean_mst_length) - ctx.save_for_backward(xyz, assignment) - return dist, assignment, mean_mst_length / (n / primitive_size) - - @staticmethod - def backward(ctx, grad_dist, grad_idx, grad_mml): - xyz, assignment = ctx.saved_tensors - grad_dist = grad_dist.contiguous() - grad_xyz = torch.zeros(xyz.size(), device='cuda').contiguous() - expansion_penalty.backward(xyz, grad_xyz, grad_dist, assignment) - return grad_xyz, None, None - -class expansionPenaltyModule(nn.Module): - def __init__(self): - super(expansionPenaltyModule, self).__init__() - - def forward(self, input, primitive_size, alpha): - return expansionPenaltyFunction.apply(input, primitive_size, alpha) - -def test_expansion_penalty(): - x = torch.rand(20, 8192, 3).cuda() - print("Input_size: ", x.shape) - expansion = expansionPenaltyModule() - start_time = time.perf_counter() - dis, ass, mean_length = expansion(x, 512, 1.5) - print("Runtime: %lfs" % (time.perf_counter() - start_time)) - -#test_expansion_penalty() \ No newline at end of file diff --git a/other_models/MSN/expansion_penalty/setup.py b/other_models/MSN/expansion_penalty/setup.py deleted file mode 100644 index 1e1992f..0000000 --- a/other_models/MSN/expansion_penalty/setup.py +++ /dev/null @@ -1,14 +0,0 @@ -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup( - name='expansion_penalty', - ext_modules=[ - CUDAExtension('expansion_penalty', [ - 'expansion_penalty.cpp', - 'expansion_penalty_cuda.cu', - ]), - ], - cmdclass={ - 'build_ext': BuildExtension - }) \ No newline at end of file diff --git a/other_models/MSN/msn.py b/other_models/MSN/msn.py deleted file mode 100644 index 98a2950..0000000 --- a/other_models/MSN/msn.py +++ /dev/null @@ -1,133 +0,0 @@ -from __future__ import print_function -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -from torch.autograd import Variable -import numpy as np -import torch.nn.functional as F -import sys -sys.path.append("./other_models/") -import MSN.expansion_penalty.expansion_penalty_module as expansion -import MSN.MDS.MDS_module as MDS_module - - -class PointGenCon(nn.Module): - def __init__(self, bottleneck_size=8192): - self.bottleneck_size = bottleneck_size - super(PointGenCon, self).__init__() - self.conv1 = torch.nn.Conv1d(self.bottleneck_size, - self.bottleneck_size, 1) - self.conv2 = torch.nn.Conv1d(self.bottleneck_size, - self.bottleneck_size // 2, 1) - self.conv3 = torch.nn.Conv1d(self.bottleneck_size // 2, - self.bottleneck_size // 4, 1) - self.conv4 = torch.nn.Conv1d(self.bottleneck_size // 4, 3, 1) - - self.th = nn.Tanh() - self.bn1 = torch.nn.BatchNorm1d(self.bottleneck_size) - self.bn2 = torch.nn.BatchNorm1d(self.bottleneck_size // 2) - self.bn3 = torch.nn.BatchNorm1d(self.bottleneck_size // 4) - - def forward(self, x): - x = F.relu(self.bn1(self.conv1(x))) - x = F.relu(self.bn2(self.conv2(x))) - x = F.relu(self.bn3(self.conv3(x))) - # x = self.th(self.conv4(x)) - x = self.conv4(x) - return x - - -class PointNetRes(nn.Module): - def __init__(self): - super(PointNetRes, self).__init__() - self.conv1 = torch.nn.Conv1d( - 4, 64, kernel_size=1, padding=0, padding_mode='replicate') - self.conv2 = torch.nn.Conv1d( - 64, 128, kernel_size=1, padding=0, padding_mode='replicate') - self.conv3 = torch.nn.Conv1d( - 128, 1024, kernel_size=1, padding=0, padding_mode='replicate') - self.conv4 = torch.nn.Conv1d( - 1088, 512, kernel_size=1, padding=0, padding_mode='replicate') - self.conv5 = torch.nn.Conv1d( - 512, 256, kernel_size=1, padding=0, padding_mode='replicate') - self.conv6 = torch.nn.Conv1d( - 256, 128, kernel_size=1, padding=0, padding_mode='replicate') - self.conv7 = torch.nn.Conv1d( - 128, 3, kernel_size=1, padding=0, padding_mode='replicate') - - self.bn1 = torch.nn.BatchNorm1d(64) - self.bn2 = torch.nn.BatchNorm1d(128) - self.bn3 = torch.nn.BatchNorm1d(1024) - self.bn4 = torch.nn.BatchNorm1d(512) - self.bn5 = torch.nn.BatchNorm1d(256) - self.bn6 = torch.nn.BatchNorm1d(128) - self.bn7 = torch.nn.BatchNorm1d(3) - self.th = nn.Tanh() - - def forward(self, x): - npoints = x.size()[2] - x = F.relu(self.bn1(self.conv1(x))) - pointfeat = x - x = F.relu(self.bn2(self.conv2(x))) - x = self.bn3(self.conv3(x)) - x, _ = torch.max(x, 2) - x = x.view(-1, 1024) - x = x.view(-1, 1024, 1).repeat(1, 1, npoints) - x = torch.cat([x, pointfeat], 1) - x = F.relu(self.bn4(self.conv4(x))) - x = F.relu(self.bn5(self.conv5(x))) - x = F.relu(self.bn6(self.conv6(x))) - x = self.th(self.conv7(x)) - return x - - -class MSN(nn.Module): - def __init__(self, num_points=8192, bottleneck_size=256, n_primitives=16): - super(MSN, self).__init__() - self.num_points = num_points - self.bottleneck_size = bottleneck_size - self.n_primitives = n_primitives - self.decoder = nn.ModuleList([ - PointGenCon(bottleneck_size=2 + self.bottleneck_size) - for i in range(0, self.n_primitives) - ]) - self.res = PointNetRes() - self.expansion = expansion.expansionPenaltyModule() - - def forward(self, partial, feat): - outs = [] - x = feat - for i in range(0, self.n_primitives): - # x = feat[:, i, :] - rand_grid = Variable( - torch.cuda.FloatTensor( - x.size(0), 2, self.num_points // self.n_primitives)) - rand_grid.data.uniform_(0, 1) - y = x.unsqueeze(2).expand(x.size(0), x.size(1), - rand_grid.size(2)).contiguous() - y = torch.cat((rand_grid, y), 1).contiguous() - outs.append(self.decoder[i](y)) - - outs = torch.cat(outs, 2).contiguous() - out1 = outs.transpose(1, 2).contiguous() - - dist, _, mean_mst_dis = self.expansion( - out1, self.num_points // self.n_primitives, 1.5) - loss_mst = torch.mean(dist) - - id0 = torch.zeros(outs.shape[0], 1, outs.shape[2]).cuda().contiguous() - outs = torch.cat((outs, id0), 1) - id1 = torch.ones(partial.shape[0], 1, - partial.shape[2]).cuda().contiguous() - partial = torch.cat((partial, id1), 1) - xx = torch.cat((outs, partial), 2) - - resampled_idx = MDS_module.minimum_density_sample( - xx[:, 0:3, :].transpose(1, 2).contiguous(), out1.shape[1], - mean_mst_dis) - xx = MDS_module.gather_operation(xx, resampled_idx) - delta = self.res(xx) - xx = xx[:, 0:3, :] - out2 = (xx + delta).transpose(2, 1).contiguous() - return [out1, out2, loss_mst, mean_mst_dis] diff --git a/other_models/VRCNet/__pycache__/vrcnet.cpython-36.pyc b/other_models/VRCNet/__pycache__/vrcnet.cpython-36.pyc deleted file mode 100644 index 7b49af2..0000000 Binary files a/other_models/VRCNet/__pycache__/vrcnet.cpython-36.pyc and /dev/null differ diff --git a/other_models/VRCNet/models/__init__.py b/other_models/VRCNet/models/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/other_models/VRCNet/models/__pycache__/__init__.cpython-36.pyc b/other_models/VRCNet/models/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index 11908a1..0000000 Binary files a/other_models/VRCNet/models/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/VRCNet/models/__pycache__/pcn.cpython-36.pyc b/other_models/VRCNet/models/__pycache__/pcn.cpython-36.pyc deleted file mode 100644 index e2531f8..0000000 Binary files a/other_models/VRCNet/models/__pycache__/pcn.cpython-36.pyc and /dev/null differ diff --git a/other_models/VRCNet/models/cascade.py b/other_models/VRCNet/models/cascade.py deleted file mode 100644 index 051381f..0000000 --- a/other_models/VRCNet/models/cascade.py +++ /dev/null @@ -1,203 +0,0 @@ -from __future__ import print_function -import math -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -import torch.nn.functional as F -import sys -import os -from utils.model_utils import calc_emd, calc_cd, gen_grid_up - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "utils/Pointnet2.PyTorch/pointnet2")) -import pointnet2_utils as pn2 -from pointnet2_modules import PointnetSAModuleMSG - - -def symmetric_sample(points, num): - p1_idx = pn2.furthest_point_sample(points, num) - input_fps = pn2.gather_operation(points.transpose(1, 2).contiguous(), p1_idx).transpose(1, 2).contiguous() - x = torch.unsqueeze(input_fps[:, :, 0], dim=2) - y = torch.unsqueeze(input_fps[:, :, 1], dim=2) - z = torch.unsqueeze(-input_fps[:, :, 2], dim=2) - input_fps_flip = torch.cat([x, y, z], dim=2) - input_fps = torch.cat([input_fps, input_fps_flip], dim=1) - return input_fps - - -class MLP(nn.Module): - def __init__(self, dims): - super().__init__() - self.model = nn.Sequential() - for i, num_channels in enumerate(dims[:-1]): - self.model.add_module('fc_%d' % (i+1), nn.Linear(num_channels, dims[i+1])) - if i != len(dims) - 2: - self.model.add_module('relu_%d' % (i+1), nn.ReLU()) - - def forward(self, features): - return self.model(features) - - -class MLPConv(nn.Module): - def __init__(self, dims): - super().__init__() - self.model = nn.Sequential() - for i, num_channels in enumerate(dims[:-1]): - self.model.add_module('conv1d_%d' % (i+1), nn.Conv1d(num_channels, dims[i+1], kernel_size=1)) - if i != len(dims) - 2: - self.model.add_module('relu_%d' % (i+1), nn.ReLU()) - - def forward(self, inputs): - return self.model(inputs) - - -class ContractExpandOperation(nn.Module): - def __init__(self, num_input_channels, up_ratio): - super().__init__() - self.up_ratio = up_ratio - # PyTorch default padding is 'VALID' - # !!! rmb to add in L2 loss for conv2d weights - self.conv2d_1 = nn.Conv2d(num_input_channels, 64, kernel_size=(1, self.up_ratio), stride=(1, 1)) - self.conv2d_2 = nn.Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1)) - self.conv2d_3 = nn.Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1)) - - def forward(self, inputs): # (32, 64, 2048) - net = inputs.view(inputs.shape[0], inputs.shape[1], self.up_ratio, -1) # (32, 64, 2, 1024) - net = net.permute(0, 1, 3, 2).contiguous() # (32, 64, 1024, 2) - net = F.relu(self.conv2d_1(net)) # (32, 64, 1024, 1) - net = F.relu(self.conv2d_2(net)) # (32, 128, 1024, 1) - net = net.permute(0, 2, 3, 1).contiguous() # (32, 1024, 1, 128) - net = net.view(net.shape[0], -1, self.up_ratio, 64) # (32, 1024, 2, 64) - net = net.permute(0, 3, 1, 2).contiguous() # (32, 64, 1024, 2) - net = F.relu(self.conv2d_3(net)) # (32, 64, 1024, 2) - net = net.view(net.shape[0], 64, -1) # (32, 64, 2048) - return net - - -class Encoder(nn.Module): - def __init__(self, embed_size=1024): - super().__init__() - self.conv1 = MLPConv([3, 128, 256]) - self.conv2 = MLPConv([512, 512, embed_size]) - - def forward(self, inputs): - ''' - :param inputs: B * C * N - :return: B * C - ''' - features = self.conv1(inputs) # [32, 256, 2048] - features_global, _ = torch.max(features, 2, keepdim=True) # [32, 256, 1] - features_global_tiled = features_global.repeat(1, 1, inputs.shape[2]) # [32, 256, 2048] - features = torch.cat([features, features_global_tiled], dim=1) # [32, 512, 2048] - features = self.conv2(features) # [32, 1024, 2048] - features, _ = torch.max(features, 2) # [32, 1024] - return features - - -class Decoder(nn.Module): - def __init__(self): - super().__init__() - self.coarse_mlp = MLP([1024, 1024, 1024, 512 * 3]) - self.mean_fc = nn.Linear(1024, 128) - self.up_branch_mlp_conv_mf = MLPConv([1157, 128, 64]) - self.up_branch_mlp_conv_nomf = MLPConv([1029, 128, 64]) - self.contract_expand = ContractExpandOperation(64, 2) - self.fine_mlp_conv = MLPConv([64, 512, 512, 3]) - - def forward(self, code, inputs, step_ratio, num_extract=512, mean_feature=None): - ''' - :param code: B * C - :param inputs: B * C * N - :param step_ratio: int - :param num_extract: int - :param mean_feature: B * C - :return: coarse(B * N * C), fine(B, N, C) - ''' - coarse = torch.tanh(self.coarse_mlp(code)) # (32, 1536) - coarse = coarse.view(-1, 512, 3) # (32, 512, 3) - coarse = coarse.transpose(2, 1).contiguous() # (32, 3, 512) - - inputs_new = inputs.transpose(2, 1).contiguous() # (32, 2048, 3) - input_fps = symmetric_sample(inputs_new, int(num_extract/2)) # [32, 512, 3] - input_fps = input_fps.transpose(2, 1).contiguous() # [32, 3, 512] - level0 = torch.cat([input_fps, coarse], 2) # (32, 3, 1024) - if num_extract > 512: - level0_flipped = level0.transpose(2, 1).contiguous() - level0 = pn2.gather_operation(level0, pn2.furthest_point_sample(level0_flipped, 1024)) - - for i in range(int(math.log2(step_ratio))): - num_fine = 2 ** (i + 1) * 1024 - grid = gen_grid_up(2 ** (i + 1)).cuda().contiguous() - grid = torch.unsqueeze(grid, 0) # (1, 2, 2) - grid_feat = grid.repeat(level0.shape[0], 1, 1024) # (32, 2, 2048) - point_feat = torch.unsqueeze(level0, 3).repeat(1, 1, 1, 2) # (32, 3, 1024, 2) - point_feat = point_feat.view(-1, 3, num_fine) # (32, 3, 2048) - global_feat = torch.unsqueeze(code, 2).repeat(1, 1, num_fine) # (32, 1024, 2048) - - if mean_feature is not None: - mean_feature_use = F.relu(self.mean_fc(mean_feature)) #(32, 128) - mean_feature_use = torch.unsqueeze(mean_feature_use, 2).repeat(1, 1, num_fine) #(32, 128, 2048) - feat = torch.cat([grid_feat, point_feat, global_feat, mean_feature_use], dim=1) # (32, 1157, 2048) - feat1 = F.relu(self.up_branch_mlp_conv_mf(feat)) # (32, 64, 2048) - else: - feat = torch.cat([grid_feat, point_feat, global_feat], dim=1) - feat1 = F.relu(self.up_branch_mlp_conv_nomf(feat)) # (32, 64, 2048) - - feat2 = self.contract_expand(feat1) # (32, 64, 2048) - feat = feat1 + feat2 # (32, 64, 2048) - - fine = self.fine_mlp_conv(feat) + point_feat # (32, 3, 2048) - level0 = fine - - return coarse.transpose(1, 2).contiguous(), fine.transpose(1, 2).contiguous() - - -class Discriminator(nn.Module): - def __init__(self, args, divide_ratio=2): - super(Discriminator, self).__init__() - self.num_points = args.num_points - self.pointnet_sa_module = PointnetSAModuleMSG(npoint=int(self.num_points/8), radii=[0.1, 0.2, 0.4], nsamples=[16, 32, 128], - mlps=[[3, 32 // divide_ratio, 32 // divide_ratio, 64 // divide_ratio], - [3, 64 // divide_ratio, 64 // divide_ratio, 128 // divide_ratio], - [3, 64 // divide_ratio, 96 // divide_ratio, 128 // divide_ratio]],) - self.patch_mlp_conv = MLPConv([(64//divide_ratio + 128 // divide_ratio + 128 // divide_ratio), 1]) - - def forward(self, xyz): - _, l1_points = self.pointnet_sa_module(xyz, features=None) - patch_values = self.patch_mlp_conv(l1_points) - return patch_values - - -class Model(nn.Module): - def __init__(self, args): - super(Model, self).__init__() - self.step_ratio = 2 * (args.num_points // 2048) - self.train_loss = args.loss - self.encoder = Encoder() - self.decoder = Decoder() - - def forward(self, x, gt, is_training=True, mean_feature=None, alpha=None): - features_partial_pt = self.encoder(x) - out1, out2 = self.decoder(features_partial_pt, x, self.step_ratio, mean_feature=mean_feature) - - if is_training: - if self.train_loss == 'emd': - loss1 = calc_emd(out1, gt) - loss2 = calc_emd(out2, gt) - elif self.train_loss == 'cd': - loss1, _ = calc_cd(out1, gt) - loss2, _ = calc_cd(out2, gt) - total_train_loss = loss1.mean() + loss2.mean() * alpha - return out2, loss2, total_train_loss - else: - emd = calc_emd(out2, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(out2, gt, calc_f1=True) - return {'out1': out1, 'out2': out2, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} - - -def weights_init(m): - if isinstance(m, nn.Conv2d): - torch.nn.init.xavier_uniform_(m.weight) - if m.bias is not None: - torch.nn.init.zeros_(m.bias) diff --git a/other_models/VRCNet/models/ecg.py b/other_models/VRCNet/models/ecg.py deleted file mode 100644 index 1a299a9..0000000 --- a/other_models/VRCNet/models/ecg.py +++ /dev/null @@ -1,244 +0,0 @@ -from __future__ import print_function -import numpy as np -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -import torch.nn.functional as F -from utils.model_utils import * -from models.pcn import PCN_encoder - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "utils/Pointnet2.PyTorch/pointnet2")) -import pointnet2_utils as pn2 - - -class Stack_conv(nn.Module): - def __init__(self, input_size, output_size, act=None): - super(Stack_conv, self).__init__() - self.model = nn.Sequential() - self.model.add_module('conv', nn.Conv2d(input_size, output_size, 1)) - - if act is not None: - self.model.add_module('act', act) - - def forward(self, x): - y = self.model(x) - y = torch.cat((x, y), 1) - return y - - -class Dense_conv(nn.Module): - def __init__(self, input_size, growth_rate=64, dense_n=3, k=16): - super(Dense_conv, self).__init__() - self.growth_rate = growth_rate - self.dense_n = dense_n - self.k = k - self.comp = growth_rate * 2 - self.input_size = input_size - - self.first_conv = nn.Conv2d(self.input_size * 2, growth_rate, 1) - - self.input_size += self.growth_rate - - self.model = nn.Sequential() - for i in range(dense_n - 1): - if i == dense_n - 2: - self.model.add_module('stack_conv_%d' % (i + 1), Stack_conv(self.input_size, self.growth_rate, None)) - else: - self.model.add_module('stack_conv_%d' % (i + 1), - Stack_conv(self.input_size, self.growth_rate, nn.ReLU())) - self.input_size += growth_rate - - def forward(self, x): - y = get_graph_feature(x, k=self.k) - y = F.relu(self.first_conv(y)) - y = torch.cat((y, x.unsqueeze(3).repeat(1, 1, 1, self.k)), 1) - - y = self.model(y) - y, _ = torch.max(y, 3) - return y - - -class EF_encoder(nn.Module): - def __init__(self, growth_rate=24, dense_n=3, k=16, hierarchy=[1024, 256, 64], input_size=3, output_size=256): - super(EF_encoder, self).__init__() - self.growth_rate = growth_rate - self.comp = growth_rate * 2 - self.dense_n = dense_n - self.k = k - self.hierarchy = hierarchy - - self.init_channel = 24 - - self.conv1 = nn.Conv1d(input_size, self.init_channel, 1) - self.dense_conv1 = Dense_conv(self.init_channel, self.growth_rate, self.dense_n, self.k) - - out_channel_size_1 = (self.init_channel * 2 + self.growth_rate * self.dense_n) # 24*2 + 24*3 = 120 - self.conv2 = nn.Conv1d(out_channel_size_1 * 2, self.comp, 1) - self.dense_conv2 = Dense_conv(self.comp, self.growth_rate, self.dense_n, self.k) - - out_channel_size_2 = ( - out_channel_size_1 * 2 + self.comp + self.growth_rate * self.dense_n) # 120*2 + 48 + 24*3 = 360 - self.conv3 = nn.Conv1d(out_channel_size_2 * 2, self.comp, 1) - self.dense_conv3 = Dense_conv(self.comp, self.growth_rate, self.dense_n, self.k) - - out_channel_size_3 = ( - out_channel_size_2 * 2 + self.comp + self.growth_rate * self.dense_n) # 360*2 + 48 + 24*3 = 840 - self.conv4 = nn.Conv1d(out_channel_size_3 * 2, self.comp, 1) - self.dense_conv4 = Dense_conv(self.comp, self.growth_rate, self.dense_n, self.k) - - out_channel_size_4 = out_channel_size_3 * 2 + self.comp + self.growth_rate * self.dense_n # 840*2 + 48 + 24*3 = 1800 - self.gf_conv = nn.Conv1d(out_channel_size_4, 1024, 1) - - self.fc1 = nn.Linear(1024, 512) - self.fc2 = nn.Linear(512, 1024) - - out_channel_size = out_channel_size_4 + 1024 - self.conv5 = nn.Conv1d(out_channel_size, 1024, 1) - - out_channel_size = out_channel_size_3 + 1024 - self.conv6 = nn.Conv1d(out_channel_size, 768, 1) - - out_channel_size = out_channel_size_2 + 768 - self.conv7 = nn.Conv1d(out_channel_size, 512, 1) - - out_channel_size = out_channel_size_1 + 512 - self.conv8 = nn.Conv1d(out_channel_size, output_size, 1) - - def forward(self, x): - point_cloud1 = x[:, 0:3, :] - point_cloud1 = point_cloud1.transpose(1, 2).contiguous() - - x0 = F.relu(self.conv1(x)) # 24 - x1 = F.relu(self.dense_conv1(x0)) # 24 + 24 * 3 = 96 - x1 = torch.cat((x1, x0), 1) # 120 - x1d, _, _, point_cloud2 = edge_preserve_sampling(x1, point_cloud1, self.hierarchy[0], self.k) # 240 - - x2 = F.relu(self.conv2(x1d)) # 48 - x2 = F.relu(self.dense_conv2(x2)) # 48 + 24 * 3 = 120 - x2 = torch.cat((x2, x1d), 1) # 120 + 240 = 360 - x2d, _, _, point_cloud3 = edge_preserve_sampling(x2, point_cloud2, self.hierarchy[1], self.k) # 720 - - x3 = F.relu(self.conv3(x2d)) - x3 = F.relu(self.dense_conv3(x3)) - x3 = torch.cat((x3, x2d), 1) - x3d, _, _, point_cloud4 = edge_preserve_sampling(x3, point_cloud3, self.hierarchy[2], self.k) - - x4 = F.relu(self.conv4(x3d)) - x4 = F.relu(self.dense_conv4(x4)) - x4 = torch.cat((x4, x3d), 1) - - global_feat = self.gf_conv(x4) - global_feat, _ = torch.max(global_feat, -1) - global_feat = F.relu(self.fc1(global_feat)) - global_feat = F.relu(self.fc2(global_feat)).unsqueeze(2).repeat(1, 1, self.hierarchy[2]) - - x4 = torch.cat((global_feat, x4), 1) - x4 = F.relu(self.conv5(x4)) - idx, weight = three_nn_upsampling(point_cloud3, point_cloud4) - x4 = pn2.three_interpolate(x4, idx, weight) - - x3 = torch.cat((x3, x4), 1) - x3 = F.relu(self.conv6(x3)) - idx, weight = three_nn_upsampling(point_cloud2, point_cloud3) - x3 = pn2.three_interpolate(x3, idx, weight) - - x2 = torch.cat((x2, x3), 1) - x2 = F.relu(self.conv7(x2)) - idx, weight = three_nn_upsampling(point_cloud1, point_cloud2) - x2 = pn2.three_interpolate(x2, idx, weight) - - x1 = torch.cat((x1, x2), 1) - x1 = self.conv8(x1) - return x1 - - -class ECG_decoder(nn.Module): - def __init__(self, num_coarse, num_fine, num_input): - super(ECG_decoder, self).__init__() - self.num_coarse = num_coarse - self.num_fine = num_fine - - self.scale = int(np.ceil(num_fine / (num_coarse + num_input))) - - self.fc1 = nn.Linear(1024, 1024) - self.fc2 = nn.Linear(1024, 1024) - self.fc3 = nn.Linear(1024, num_coarse * 3) - - self.dense_feature_size = 256 - self.expand_feature_size = 64 - self.input_size = 3 - - self.encoder = EF_encoder(growth_rate=24, dense_n=3, k=16, hierarchy=[1024, 256, 64], - input_size=self.input_size, output_size=self.dense_feature_size) - - if self.scale >= 2: - self.expansion = EF_expansion(input_size=self.dense_feature_size, output_size=self.expand_feature_size, - step_ratio=self.scale, k=4) - self.conv1 = nn.Conv1d(self.expand_feature_size, self.expand_feature_size, 1) - else: - self.expansion = None - self.conv1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - self.conv2 = nn.Conv1d(self.expand_feature_size, 3, 1) - - def forward(self, global_feat, point_input): - batch_size = global_feat.size()[0] - coarse = F.relu(self.fc1(global_feat)) - coarse = F.relu(self.fc2(coarse)) - coarse = self.fc3(coarse).view(batch_size, 3, self.num_coarse) - org_points_input = point_input - points = torch.cat((coarse, org_points_input), 2) - - dense_feat = self.encoder(points) - - if self.scale >= 2: - dense_feat = self.expansion(dense_feat) - - point_feat = F.relu(self.conv1(dense_feat)) - fine = self.conv2(point_feat) - - num_out = fine.size()[2] - if num_out > self.num_fine: - fine = pn2.gather_operation(fine, - pn2.furthest_point_sample(fine.transpose(1, 2).contiguous(), self.num_fine)) - - return coarse, fine - - -class Model(nn.Module): - def __init__(self, args, num_coarse=1024, num_input=2048): - super(Model, self).__init__() - self.num_coarse = num_coarse - self.num_points = args.num_points - self.train_loss = args.loss - self.encoder = PCN_encoder() - self.decoder = ECG_decoder(num_coarse, self.num_points, num_input) - - def forward(self, x, gt, is_training=True, mean_feature=None, alpha=None): - if mean_feature: - raise NotImplementedError - feat = self.encoder(x) - out1, out2 = self.decoder(feat, x) - out1 = out1.transpose(1, 2).contiguous() - out2 = out2.transpose(1, 2).contiguous() - uniform_loss1 = get_uniform_loss(out1) - uniform_loss2 = get_uniform_loss(out2) - - if is_training: - if self.train_loss == 'emd': - loss1 = calc_emd(out1, gt) - loss2 = calc_emd(out2, gt) - elif self.train_loss == 'cd': - loss1, _ = calc_cd(out1, gt) - loss2, _ = calc_cd(out2, gt) - else: - raise NotImplementedError('Train loss is either CD or EMD!') - - total_train_loss = loss1.mean() + uniform_loss1.mean() * 0.1 + \ - (loss2.mean() + uniform_loss2.mean() * 0.1) * alpha - return out2, loss2, total_train_loss - else: - emd = calc_emd(out2, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(out2, gt, calc_f1=True) - return {'out1': out1, 'out2': out2, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} diff --git a/other_models/VRCNet/models/grnet.py b/other_models/VRCNet/models/grnet.py deleted file mode 100644 index 75eee66..0000000 --- a/other_models/VRCNet/models/grnet.py +++ /dev/null @@ -1,153 +0,0 @@ -import torch -import os - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "utils/Pointnet2.PyTorch/pointnet2")) -import pointnet2_utils as pn2 - -from utils.gridding import Gridding, GriddingReverse -from utils.cubic_feature_sampling import CubicFeatureSampling -from utils.model_utils import calc_cd, calc_emd - - -class RandomPointSampling(torch.nn.Module): - def __init__(self, n_points): - super(RandomPointSampling, self).__init__() - self.n_points = n_points - - def forward(self, pred_cloud, partial_cloud=None): - if partial_cloud is not None: - pred_cloud = torch.cat([partial_cloud, pred_cloud], dim=1) - - _ptcloud = torch.split(pred_cloud, 1, dim=0) - ptclouds = [] - for p in _ptcloud: - non_zeros = torch.sum(p, dim=2).ne(0) - p = p[non_zeros].unsqueeze(dim=0) - n_pts = p.size(1) - if n_pts < self.n_points: - rnd_idx = torch.cat([torch.randint(0, n_pts, (self.n_points, ))]) - else: - rnd_idx = torch.randperm(p.size(1))[:self.n_points] - ptclouds.append(p[:, rnd_idx, :]) - - return torch.cat(ptclouds, dim=0).contiguous() - - -class Model(torch.nn.Module): - def __init__(self, args): - super(Model, self).__init__() - self.num_points = args.num_points - self.gridding = Gridding(scale=64) - self.conv1 = torch.nn.Sequential( - torch.nn.Conv3d(1, 32, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(32), - torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d(kernel_size=2) - ) - self.conv2 = torch.nn.Sequential( - torch.nn.Conv3d(32, 64, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(64), - torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d(kernel_size=2) - ) - self.conv3 = torch.nn.Sequential( - torch.nn.Conv3d(64, 128, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(128), - torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d(kernel_size=2) - ) - self.conv4 = torch.nn.Sequential( - torch.nn.Conv3d(128, 256, kernel_size=4, padding=2), - torch.nn.BatchNorm3d(256), - torch.nn.LeakyReLU(0.2), - torch.nn.MaxPool3d(kernel_size=2) - ) - self.fc5 = torch.nn.Sequential( - torch.nn.Linear(16384, 2048), - torch.nn.ReLU() - ) - self.fc6 = torch.nn.Sequential( - torch.nn.Linear(2048, 16384), - torch.nn.ReLU() - ) - self.dconv7 = torch.nn.Sequential( - torch.nn.ConvTranspose3d(256, 128, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(128), - torch.nn.ReLU() - ) - self.dconv8 = torch.nn.Sequential( - torch.nn.ConvTranspose3d(128, 64, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(64), - torch.nn.ReLU() - ) - self.dconv9 = torch.nn.Sequential( - torch.nn.ConvTranspose3d(64, 32, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(32), - torch.nn.ReLU() - ) - self.dconv10 = torch.nn.Sequential( - torch.nn.ConvTranspose3d(32, 1, kernel_size=4, stride=2, bias=False, padding=1), - torch.nn.BatchNorm3d(1), - torch.nn.ReLU() - ) - self.gridding_rev = GriddingReverse(scale=64) - self.point_sampling = RandomPointSampling(n_points=2048) - self.feature_sampling = CubicFeatureSampling() - self.fc11 = torch.nn.Sequential( - torch.nn.Linear(1792, 1792), - torch.nn.ReLU() - ) - self.fc12 = torch.nn.Sequential( - torch.nn.Linear(1792, 448), - torch.nn.ReLU() - ) - self.fc13 = torch.nn.Sequential( - torch.nn.Linear(448, 112), - torch.nn.ReLU() - ) - self.fc14 = torch.nn.Linear(112, 24) - - def forward(self, partial_cloud, gt, is_training=True, mean_feature=None, alpha=None): - pt_features_64_l = self.gridding(partial_cloud).view(-1, 1, 64, 64, 64) - pt_features_32_l = self.conv1(pt_features_64_l) - pt_features_16_l = self.conv2(pt_features_32_l) - pt_features_8_l = self.conv3(pt_features_16_l) - pt_features_4_l = self.conv4(pt_features_8_l) - features = self.fc5(pt_features_4_l.view(-1, 16384)) - pt_features_4_r = self.fc6(features).view(-1, 256, 4, 4, 4) + pt_features_4_l - pt_features_8_r = self.dconv7(pt_features_4_r) + pt_features_8_l - pt_features_16_r = self.dconv8(pt_features_8_r) + pt_features_16_l - pt_features_32_r = self.dconv9(pt_features_16_r) + pt_features_32_l - pt_features_64_r = self.dconv10(pt_features_32_r) + pt_features_64_l - sparse_cloud = self.gridding_rev(pt_features_64_r.squeeze(dim=1)) - sparse_cloud = self.point_sampling(sparse_cloud, partial_cloud) - point_features_32 = self.feature_sampling(sparse_cloud, pt_features_32_r).view(-1, 2048, 256) - point_features_16 = self.feature_sampling(sparse_cloud, pt_features_16_r).view(-1, 2048, 512) - point_features_8 = self.feature_sampling(sparse_cloud, pt_features_8_r).view(-1, 2048, 1024) - point_features = torch.cat([point_features_32, point_features_16, point_features_8], dim=2) - point_features = self.fc11(point_features) - point_features = self.fc12(point_features) - point_features = self.fc13(point_features) - point_offset = self.fc14(point_features).view(-1, 16384, 3) - dense_cloud = sparse_cloud.unsqueeze(dim=2).repeat(1, 1, 8, 1).view(-1, 16384, 3) + point_offset - if self.num_points < 16384: - idx_fps = pn2.furthest_point_sample(dense_cloud, self.num_points) - dense_cloud = pn2.gather_operation(dense_cloud, idx_fps) - - if is_training: - if self.train_loss == 'emd': - loss1 = calc_emd(sparse_cloud, gt) - loss2 = calc_emd(dense_cloud, gt) - elif self.train_loss == 'cd': - _, loss1 = calc_cd(sparse_cloud, gt) # cd_t - _, loss2 = calc_cd(dense_cloud, gt) # cd_t - else: - raise NotImplementedError('Train loss is either CD or EMD!') - - total_train_loss = loss1.mean() + loss2.mean() - return dense_cloud, loss2, total_train_loss - else: - emd = calc_emd(dense_cloud, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(dense_cloud, gt, calc_f1=True) - return {'out1': dense_cloud, 'out2': dense_cloud, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} \ No newline at end of file diff --git a/other_models/VRCNet/models/msn.py b/other_models/VRCNet/models/msn.py deleted file mode 100644 index 6fedc91..0000000 --- a/other_models/VRCNet/models/msn.py +++ /dev/null @@ -1,215 +0,0 @@ -from __future__ import print_function -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -from torch.autograd import Variable -import numpy as np -import torch.nn.functional as F -import sys -import os -from utils.model_utils import calc_emd, calc_cd - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "utils/expansion_penalty")) -import expansion_penalty_module as expansion -sys.path.append(os.path.join(proj_dir, "utils/MDS")) -import MDS_module - - -class STN3d(nn.Module): - def __init__(self, num_points = 2500): - super(STN3d, self).__init__() - self.num_points = num_points - self.conv1 = torch.nn.Conv1d(3, 64, 1) - self.conv2 = torch.nn.Conv1d(64, 128, 1) - self.conv3 = torch.nn.Conv1d(128, 1024, 1) - self.fc1 = nn.Linear(1024, 512) - self.fc2 = nn.Linear(512, 256) - self.fc3 = nn.Linear(256, 9) - self.relu = nn.ReLU() - - def forward(self, x): - batchsize = x.size()[0] - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - x = F.relu(self.conv3(x)) - x,_ = torch.max(x, 2) - x = x.view(-1, 1024) - - x = F.relu(self.fc1(x)) - x = F.relu(self.fc2(x)) - x = self.fc3(x) - - iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1) - if x.is_cuda: - iden = iden.cuda() - x = x + iden - x = x.view(-1, 3, 3) - return x - -class PointNetfeat(nn.Module): - def __init__(self, num_points = 8192, global_feat = True): - super(PointNetfeat, self).__init__() - self.stn = STN3d(num_points = num_points) - self.conv1 = torch.nn.Conv1d(3, 64, 1) - self.conv2 = torch.nn.Conv1d(64, 128, 1) - self.conv3 = torch.nn.Conv1d(128, 1024, 1) - - self.bn1 = torch.nn.BatchNorm1d(64) - self.bn2 = torch.nn.BatchNorm1d(128) - self.bn3 = torch.nn.BatchNorm1d(1024) - - self.num_points = num_points - self.global_feat = global_feat - def forward(self, x): - batchsize = x.size()[0] - x = F.relu(self.bn1(self.conv1(x))) - x = F.relu(self.bn2(self.conv2(x))) - x = self.bn3(self.conv3(x)) - x,_ = torch.max(x, 2) - x = x.view(-1, 1024) - return x - -class PointGenCon(nn.Module): - def __init__(self, bottleneck_size = 8192): - self.bottleneck_size = bottleneck_size - super(PointGenCon, self).__init__() - self.conv1 = torch.nn.Conv1d(self.bottleneck_size, self.bottleneck_size, 1) - self.conv2 = torch.nn.Conv1d(self.bottleneck_size, self.bottleneck_size//2, 1) - self.conv3 = torch.nn.Conv1d(self.bottleneck_size//2, self.bottleneck_size//4, 1) - self.conv4 = torch.nn.Conv1d(self.bottleneck_size//4, 3, 1) - - self.th = nn.Tanh() - self.bn1 = torch.nn.BatchNorm1d(self.bottleneck_size) - self.bn2 = torch.nn.BatchNorm1d(self.bottleneck_size//2) - self.bn3 = torch.nn.BatchNorm1d(self.bottleneck_size//4) - - def forward(self, x): - batchsize = x.size()[0] - x = F.relu(self.bn1(self.conv1(x))) - x = F.relu(self.bn2(self.conv2(x))) - x = F.relu(self.bn3(self.conv3(x))) - x = self.th(self.conv4(x)) - return x - -class PointNetRes(nn.Module): - def __init__(self): - super(PointNetRes, self).__init__() - self.conv1 = torch.nn.Conv1d(4, 64, 1) - self.conv2 = torch.nn.Conv1d(64, 128, 1) - self.conv3 = torch.nn.Conv1d(128, 1024, 1) - self.conv4 = torch.nn.Conv1d(1088, 512, 1) - self.conv5 = torch.nn.Conv1d(512, 256, 1) - self.conv6 = torch.nn.Conv1d(256, 128, 1) - self.conv7 = torch.nn.Conv1d(128, 3, 1) - - - self.bn1 = torch.nn.BatchNorm1d(64) - self.bn2 = torch.nn.BatchNorm1d(128) - self.bn3 = torch.nn.BatchNorm1d(1024) - self.bn4 = torch.nn.BatchNorm1d(512) - self.bn5 = torch.nn.BatchNorm1d(256) - self.bn6 = torch.nn.BatchNorm1d(128) - self.bn7 = torch.nn.BatchNorm1d(3) - self.th = nn.Tanh() - - def forward(self, x): - batchsize = x.size()[0] - npoints = x.size()[2] - x = F.relu(self.bn1(self.conv1(x))) - pointfeat = x - x = F.relu(self.bn2(self.conv2(x))) - x = self.bn3(self.conv3(x)) - x,_ = torch.max(x, 2) - x = x.view(-1, 1024) - x = x.view(-1, 1024, 1).repeat(1, 1, npoints) - x = torch.cat([x, pointfeat], 1) - x = F.relu(self.bn4(self.conv4(x))) - x = F.relu(self.bn5(self.conv5(x))) - x = F.relu(self.bn6(self.conv6(x))) - x = self.th(self.conv7(x)) - return x - - -class Model(nn.Module): - def __init__(self, args, bottleneck_size=1024): - super(Model, self).__init__() - self.train_loss = args.loss - self.num_points = args.num_points - self.bottleneck_size = bottleneck_size - self.n_primitives = args.n_primitives - self.encoder = nn.Sequential( - PointNetfeat(self.num_points, global_feat=True), - nn.Linear(1024, self.bottleneck_size), - nn.BatchNorm1d(self.bottleneck_size), - nn.ReLU() - ) - self.decoder = nn.ModuleList([PointGenCon(bottleneck_size = 2 +self.bottleneck_size) for i in range(0,self.n_primitives)]) - self.res = PointNetRes() - self.expansion = expansion.expansionPenaltyModule() - - def forward(self, x, gt, is_training=True, mean_feature=None, alpha=None): - if mean_feature: - raise NotImplementedError - partial = x - x = self.encoder(x) - outs = [] - for i in range(0, self.n_primitives): - rand_grid = Variable(torch.cuda.FloatTensor(x.size(0), 2, self.num_points // self.n_primitives)) - rand_grid.data.uniform_(0, 1) - y = x.unsqueeze(2).expand(x.size(0), x.size(1), rand_grid.size(2)).contiguous() - y = torch.cat((rand_grid, y), 1).contiguous() - outs.append(self.decoder[i](y)) - - outs = torch.cat(outs, 2).contiguous() - out1 = outs.transpose(1, 2).contiguous() - - dist, _, mean_mst_dis = self.expansion(out1, self.num_points // self.n_primitives, 1.5) - expansion_penalty = torch.mean(dist) - - id0 = torch.zeros(outs.shape[0], 1, outs.shape[2]).cuda().contiguous() - outs = torch.cat((outs, id0), 1) - id1 = torch.ones(partial.shape[0], 1, partial.shape[2]).cuda().contiguous() - partial = torch.cat((partial, id1), 1) - xx = torch.cat((outs, partial), 2) - - resampled_idx = MDS_module.minimum_density_sample(xx[:, 0:3, :].transpose(1, 2).contiguous(), out1.shape[1], - mean_mst_dis) - xx = MDS_module.gather_operation(xx, resampled_idx) - delta = self.res(xx) - xx = xx[:, 0:3, :] - out2 = (xx + delta).transpose(2, 1).contiguous() - - if is_training: - if self.train_loss == 'emd': - loss1 = calc_emd(out1, gt) - loss2 = calc_emd(out2, gt) - elif self.train_loss == 'cd': - loss1, _ = calc_cd(out1, gt) - loss2, _ = calc_cd(out2, gt) - else: - raise NotImplementedError('Train loss is either CD or EMD!') - - total_train_loss = loss1.mean() + loss2.mean() + expansion_penalty*0.1 - return out2, loss2, total_train_loss - else: - emd = calc_emd(out2, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(out2, gt, calc_f1=True) - return {'out1': out1, 'out2': out2, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} - - -def weights_init(m): - classname = m.__class__.__name__ - if classname.find('Conv') != -1 and classname.find('Conv2d') == -1: - m.weight.data.normal_(0.0, 0.02) - elif classname.find('BatchNorm') != -1 and classname.find('BatchNorm2d') == -1: - m.weight.data.normal_(1.0, 0.02) - m.bias.data.fill_(0) - - - - - - - diff --git a/other_models/VRCNet/models/pcn.py b/other_models/VRCNet/models/pcn.py deleted file mode 100644 index be5d8f9..0000000 --- a/other_models/VRCNet/models/pcn.py +++ /dev/null @@ -1,106 +0,0 @@ -from __future__ import print_function -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -import torch.nn.functional as F -import math -from VRCNet.utils.model_utils import gen_grid_up, calc_emd, calc_cd - - -class PCN_encoder(nn.Module): - def __init__(self, output_size=1024): - super(PCN_encoder, self).__init__() - self.conv1 = nn.Conv1d(3, 128, 1) - self.conv2 = nn.Conv1d(128, 256, 1) - self.conv3 = nn.Conv1d(512, 512, 1) - self.conv4 = nn.Conv1d(512, output_size, 1) - - def forward(self, x): - batch_size, _, num_points = x.size() - x = F.relu(self.conv1(x)) - x = self.conv2(x) - global_feature, _ = torch.max(x, 2) - x = torch.cat((x, global_feature.view(batch_size, -1, 1).repeat(1, 1, num_points).contiguous()), 1) - x = F.relu(self.conv3(x)) - x = self.conv4(x) - global_feature, _ = torch.max(x, 2) - return global_feature.view(batch_size, -1) - - -class PCN_decoder(nn.Module): - def __init__(self, num_coarse, num_fine, scale, cat_feature_num): - super(PCN_decoder, self).__init__() - self.num_coarse = num_coarse - self.num_fine = num_fine - self.fc1 = nn.Linear(1024, 1024) - self.fc2 = nn.Linear(1024, 1024) - self.fc3 = nn.Linear(1024, num_coarse * 3) - - self.scale = scale - self.grid = gen_grid_up(2 ** (int(math.log2(scale))), 0.05).cuda().contiguous() - self.conv1 = nn.Conv1d(cat_feature_num, 512, 1) - self.conv2 = nn.Conv1d(512, 512, 1) - self.conv3 = nn.Conv1d(512, 3, 1) - - def forward(self, x): - batch_size = x.size()[0] - coarse = F.relu(self.fc1(x)) - coarse = F.relu(self.fc2(coarse)) - coarse = self.fc3(coarse).view(-1, 3, self.num_coarse) - - grid = self.grid.clone().detach() - grid_feat = grid.unsqueeze(0).repeat(batch_size, 1, self.num_coarse).contiguous().cuda() - - point_feat = ( - (coarse.transpose(1, 2).contiguous()).unsqueeze(2).repeat(1, 1, self.scale, 1).view(-1, self.num_fine, - 3)).transpose(1, - 2).contiguous() - - global_feat = x.unsqueeze(2).repeat(1, 1, self.num_fine) - - feat = torch.cat((grid_feat, point_feat, global_feat), 1) - - center = ((coarse.transpose(1, 2).contiguous()).unsqueeze(2).repeat(1, 1, self.scale, 1).view(-1, self.num_fine, - 3)).transpose(1, - 2).contiguous() - - fine = self.conv3(F.relu(self.conv2(F.relu(self.conv1(feat))))) + center - return coarse, fine - - -class Model(nn.Module): - def __init__(self, args, num_coarse=1024): - super(Model, self).__init__() - - self.num_coarse = num_coarse - self.num_points = args.num_points - self.train_loss = args.loss - self.scale = self.num_points // num_coarse - self.cat_feature_num = 2 + 3 + 1024 - - self.encoder = PCN_encoder() - self.decoder = PCN_decoder(num_coarse, self.num_points, self.scale, self.cat_feature_num) - - def forward(self, x, gt, is_training=True, mean_feature=None, alpha=None): - feat = self.encoder(x) - out1, out2 = self.decoder(feat) - out1 = out1.transpose(1, 2).contiguous() - out2 = out2.transpose(1, 2).contiguous() - - if is_training: - if self.train_loss == 'emd': - loss1 = calc_emd(out1, gt) - loss2 = calc_emd(out2, gt) - elif self.train_loss == 'cd': - loss1, _ = calc_cd(out1, gt) - loss2, _ = calc_cd(out2, gt) - else: - raise NotImplementedError('Train loss is either CD or EMD!') - - total_train_loss = loss1.mean() + loss2.mean() * alpha - return out2, loss2, total_train_loss - else: - emd = calc_emd(out2, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(out2, gt, calc_f1=True) - return {'out1': out1, 'out2': out2, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} diff --git a/other_models/VRCNet/models/topnet.py b/other_models/VRCNet/models/topnet.py deleted file mode 100644 index f442b2c..0000000 --- a/other_models/VRCNet/models/topnet.py +++ /dev/null @@ -1,168 +0,0 @@ -from __future__ import print_function -import math -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -import torch.nn.functional as F -import numpy as np -from copy import deepcopy -from utils.model_utils import calc_emd, calc_cd - - -# Number of children per tree levels for 2048 output points -tree_arch = {} -tree_arch[2] = [32, 64] -tree_arch[4] = [4, 8, 8, 8] -tree_arch[6] = [2, 4, 4, 4, 4, 4] -tree_arch[8] = [2, 2, 2, 2, 2, 4, 4, 4] - - -def get_arch(nlevels, npts): - logmult = int(math.log2(npts/2048)) - assert 2048*(2**(logmult)) == npts, "Number of points is %d, expected 2048x(2^n)" % (npts) - arch = deepcopy(tree_arch[nlevels]) - while logmult > 0: - last_min_pos = np.where(arch==np.min(arch))[0][-1] - arch[last_min_pos]*=2 - logmult -= 1 - return arch - - -class MLP(nn.Module): - def __init__(self, dims, bn=None): - super().__init__() - self.model = nn.Sequential() - for i, num_channels in enumerate(dims[:-2]): - self.model.add_module('fc_%d' % (i+1), nn.Linear(num_channels, dims[i+1])) - self.bn = bn - if self.bn: - self.batch_norm = nn.BatchNorm1d(dims[-2]) - - self.output_layer = nn.Linear(dims[-2], dims[-1]) - - def forward(self, features): - features = self.model(features) - if self.bn: - features = self.batch_norm(features) - features = F.relu(features) - outputs = self.output_layer(features) - return outputs - - -class MLPConv(nn.Module): - def __init__(self, dims, bn=None): - super().__init__() - self.model = nn.Sequential() - for i, num_channels in enumerate(dims[:-2]): - self.model.add_module('conv1d_%d' % (i+1), nn.Conv1d(num_channels, dims[i+1], kernel_size=1)) - self.bn = bn - if self.bn: - self.batch_norm = nn.BatchNorm1d(dims[-2]) - - self.output_layer = nn.Conv1d(dims[-2], dims[-1], kernel_size=1) - - def forward(self, inputs): - inputs = self.model(inputs) - if self.bn: - self.batch_norm.cuda() - inputs = self.batch_norm(inputs) - inputs = F.relu(inputs) - outputs = self.output_layer(inputs) - return outputs - - -class CreateLevel(nn.Module): - def __init__(self, level, input_channels, output_channels, bn, tarch): - super().__init__() - self.output_channels = output_channels - self.mlp_conv = MLPConv([input_channels, input_channels, int(input_channels / 2), int(input_channels / 4), - int(input_channels / 8), output_channels * int(tarch[level])], bn=bn) - - def forward(self, inputs): - features = self.mlp_conv(inputs) - features = features.view(features.shape[0], self.output_channels, -1) - return features - - -class PCNEncoder(nn.Module): - def __init__(self, embed_size=1024): - super().__init__() - self.conv1 = MLPConv([3, 128, 256]) # no bn - self.conv2 = MLPConv([512, 512, embed_size]) # no bn - - def forward(self, inputs): - ''' - :param inputs: B * C * N - :return: B * C - ''' - features = self.conv1(inputs) # [32, 256, 2048] - features_global, _ = torch.max(features, 2, keepdim=True) # [32, 256, 1] - features_global_tiled = features_global.repeat(1, 1, inputs.shape[2]) # [32, 256, 2048] - features = torch.cat([features, features_global_tiled], dim=1) # [32, 512, 2048] - features = self.conv2(features) # [32, 1024, 2048] - features, _ = torch.max(features, 2) # [32, 1024] - return features - - -class TopnetDecoder(nn.Module): - def __init__(self, npts): - super().__init__() - self.tarch = get_arch(6, npts) - self.N = int(np.prod([int(k) for k in self.tarch])) - assert self.N == npts, "Number of tree outputs is %d, expected %d" % (self.N, npts) - self.NFEAT = 8 - self.CODE_NFTS = 1024 - self.Nin = self.NFEAT + self.CODE_NFTS - self.Nout = self.NFEAT - self.N0 = int(self.tarch[0]) - self.nlevels = len(self.tarch) - self.mlp = MLP([1024, 256, 64, self.NFEAT * self.N0], bn=True) - self.mlp_conv_list = nn.ModuleList() - bn = True - for i in range(1, self.nlevels): - if i == self.nlevels - 1: - self.Nout = 3 - bn = False - self.mlp_conv_list.append(CreateLevel(i, self.Nin, self.Nout, bn, self.tarch)) - - def forward(self, code): - level0 = self.mlp(code) # - level0 = torch.tanh(level0) - level0 = level0.view(-1, self.NFEAT, self.N0) # (32, 8, 2) - outs = [level0, ] - for i in range(self.nlevels-1): - inp = outs[-1] - y = torch.unsqueeze(code, dim=2) # (32, 1024, 1) - y = y.repeat(1, 1, inp.shape[2]) # (32, 1024, 2) - y = torch.cat([inp, y], dim=1) # (32, 1032, 2) - conv_outs = self.mlp_conv_list[i](y) - outs.append(torch.tanh(conv_outs)) - return outs[-1] - - -class Model(nn.Module): - def __init__(self, args): - super(Model, self).__init__() - self.encoder = PCNEncoder() - self.decoder = TopnetDecoder(args.num_points) - self.train_loss = args.loss - - def forward(self, x, gt, is_training=True, mean_feature=None, alpha=None): - features = self.encoder(x) - out = self.decoder(features) - out = out.transpose(1, 2).contiguous() - - if is_training: - if self.train_loss == 'emd': - loss = calc_emd(out, gt) - elif self.train_loss == 'cd': - loss, _ = calc_cd(out, gt) - else: - raise NotImplementedError('Train loss is either CD or EMD!') - - return out, loss, loss.mean() - else: - emd = calc_emd(out, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(out, gt, calc_f1=True) - return {'out1': None, 'out2': out, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} diff --git a/other_models/VRCNet/models/vrc.py b/other_models/VRCNet/models/vrc.py deleted file mode 100644 index 309dc55..0000000 --- a/other_models/VRCNet/models/vrc.py +++ /dev/null @@ -1,515 +0,0 @@ -import numpy as np -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -import torch.nn.functional as F -from utils.model_utils import * -from models.pcn import PCN_encoder - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "utils/Pointnet2.PyTorch/pointnet2")) -import pointnet2_utils as pn2 - - -class SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=16): - super(SA_module, self).__init__() - self.share_planes = share_planes - self.k = k - self.conv1 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv2 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv3 = nn.Conv2d(in_planes, mid_planes, kernel_size=1) - - self.conv_w = nn.Sequential(nn.ReLU(inplace=False), - nn.Conv2d(rel_planes * (k + 1), mid_planes // share_planes, kernel_size=1, - bias=False), - nn.ReLU(inplace=False), - nn.Conv2d(mid_planes // share_planes, k * mid_planes // share_planes, - kernel_size=1)) - self.activation_fn = nn.ReLU(inplace=False) - - self.conv_out = nn.Conv2d(mid_planes, out_planes, kernel_size=1) - - def forward(self, input): - x, idx = input - batch_size, _, _, num_points = x.size() - identity = x # B C 1 N - x = self.activation_fn(x) - xn = get_edge_features(x, idx) # B C K N - x1, x2, x3 = self.conv1(x), self.conv2(xn), self.conv3(xn) - - x2 = x2.view(batch_size, -1, 1, num_points).contiguous() # B kC 1 N - w = self.conv_w(torch.cat([x1, x2], 1)).view(batch_size, -1, self.k, num_points) - w = w.repeat(1, self.share_planes, 1, 1) - out = w * x3 - out = torch.sum(out, dim=2, keepdim=True) - - out = self.activation_fn(out) - out = self.conv_out(out) # B C 1 N - out += identity - return [out, idx] - - -class Folding(nn.Module): - def __init__(self, input_size, output_size, step_ratio, global_feature_size=1024, num_models=1): - super(Folding, self).__init__() - self.input_size = input_size - self.output_size = output_size - self.step_ratio = step_ratio - self.num_models = num_models - - self.conv = nn.Conv1d(input_size + global_feature_size + 2, output_size, 1, bias=True) - - sqrted = int(math.sqrt(step_ratio)) + 1 - for i in range(1, sqrted + 1).__reversed__(): - if (step_ratio % i) == 0: - num_x = i - num_y = step_ratio // i - break - - grid_x = torch.linspace(-0.2, 0.2, steps=num_x) - grid_y = torch.linspace(-0.2, 0.2, steps=num_y) - - x, y = torch.meshgrid(grid_x, grid_y) # x, y shape: (2, 1) - self.grid = torch.stack([x, y], dim=-1).view(-1, 2) # (2, 2) - - def forward(self, point_feat, global_feat): - batch_size, num_features, num_points = point_feat.size() - point_feat = point_feat.transpose(1, 2).contiguous().unsqueeze(2).repeat(1, 1, self.step_ratio, 1).view( - batch_size, - -1, num_features).transpose(1, 2).contiguous() - global_feat = global_feat.unsqueeze(2).repeat(1, 1, num_points * self.step_ratio).repeat(self.num_models, 1, 1) - grid_feat = self.grid.unsqueeze(0).repeat(batch_size, num_points, 1).transpose(1, 2).contiguous().cuda() - features = torch.cat([global_feat, point_feat, grid_feat], axis=1) - features = F.relu(self.conv(features)) - return features - - -class Linear_ResBlock(nn.Module): - def __init__(self, input_size=1024, output_size=256): - super(Linear_ResBlock, self).__init__() - self.conv1 = nn.Linear(input_size, input_size) - self.conv2 = nn.Linear(input_size, output_size) - self.conv_res = nn.Linear(input_size, output_size) - - self.af = nn.ReLU(inplace=True) - - def forward(self, feature): - return self.conv2(self.af(self.conv1(self.af(feature)))) + self.conv_res(feature) - - -class SK_SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=[10, 20], r=2, L=32): - super(SK_SA_module, self).__init__() - - self.num_kernels = len(k) - d = max(int(out_planes / r), L) - - self.sams = nn.ModuleList([]) - - for i in range(len(k)): - self.sams.append(SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k[i])) - - self.fc = nn.Linear(out_planes, d) - self.fcs = nn.ModuleList([]) - - for i in range(len(k)): - self.fcs.append(nn.Linear(d, out_planes)) - - self.softmax = nn.Softmax(dim=1) - self.af = nn.ReLU(inplace=False) - - def forward(self, input): - x, idxs = input - assert (self.num_kernels == len(idxs)) - for i, sam in enumerate(self.sams): - fea, _ = sam([x, idxs[i]]) - fea = self.af(fea) - fea = fea.unsqueeze(dim=1) - if i == 0: - feas = fea - else: - feas = torch.cat([feas, fea], dim=1) - - fea_U = torch.sum(feas, dim=1) - fea_s = fea_U.mean(-1).mean(-1) - fea_z = self.fc(fea_s) - - for i, fc in enumerate(self.fcs): - vector = fc(fea_z).unsqueeze_(dim=1) - if i == 0: - attention_vectors = vector - else: - attention_vectors = torch.cat([attention_vectors, vector], dim=1) - - attention_vectors = self.softmax(attention_vectors) - attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1) - fea_v = (feas * attention_vectors).sum(dim=1) - return [fea_v, idxs] - - -class SKN_Res_unit(nn.Module): - def __init__(self, input_size, output_size, k=[10, 20], layers=1): - super(SKN_Res_unit, self).__init__() - self.conv1 = nn.Conv2d(input_size, output_size, 1, bias=False) - self.sam = self._make_layer(output_size, output_size // 16, output_size // 4, output_size, int(layers), 8, k=k) - self.conv2 = nn.Conv2d(output_size, output_size, 1, bias=False) - self.conv_res = nn.Conv2d(input_size, output_size, 1, bias=False) - self.af = nn.ReLU(inplace=False) - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def forward(self, feat, idx): - x, _ = self.sam([self.conv1(feat), idx]) - x = self.conv2(self.af(x)) - return x + self.conv_res(feat) - - -class SA_SKN_Res_encoder(nn.Module): - def __init__(self, input_size=3, k=[10, 20], pk=16, output_size=64, layers=[2, 2, 2, 2], - pts_num=[3072, 1536, 768, 384]): - super(SA_SKN_Res_encoder, self).__init__() - self.init_channel = 64 - - c1 = self.init_channel - self.sam_res1 = SKN_Res_unit(input_size, c1, k, int(layers[0])) - - c2 = c1 * 2 - self.sam_res2 = SKN_Res_unit(c2, c2, k, int(layers[1])) - - c3 = c2 * 2 - self.sam_res3 = SKN_Res_unit(c3, c3, k, int(layers[2])) - - c4 = c3 * 2 - self.sam_res4 = SKN_Res_unit(c4, c4, k, int(layers[3])) - - self.conv5 = nn.Conv2d(c4, 1024, 1) - - self.fc1 = nn.Linear(1024, 512) - self.fc2 = nn.Linear(512, 1024) - - self.conv6 = nn.Conv2d(c4 + 1024, c4, 1) - self.conv7 = nn.Conv2d(c3 + c4, c3, 1) - self.conv8 = nn.Conv2d(c2 + c3, c2, 1) - self.conv9 = nn.Conv2d(c1 + c2, c1, 1) - - self.conv_out = nn.Conv2d(c1, output_size, 1) - self.dropout = nn.Dropout() - self.af = nn.ReLU(inplace=False) - self.k = k - self.pk = pk - self.rate = 2 - - self.pts_num = pts_num - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def _edge_pooling(self, features, points, rate=2, k=16, sample_num=None): - features = features.squeeze(2) - - if sample_num is None: - input_points_num = int(features.size()[2]) - sample_num = input_points_num // rate - - ds_features, p_idx, pn_idx, ds_points = edge_preserve_sampling(features, points, sample_num, k) - ds_features = ds_features.unsqueeze(2) - return ds_features, p_idx, pn_idx, ds_points - - def _edge_unpooling(self, features, src_pts, tgt_pts): - features = features.squeeze(2) - idx, weight = three_nn_upsampling(tgt_pts, src_pts) - features = pn2.three_interpolate(features, idx, weight) - features = features.unsqueeze(2) - return features - - def forward(self, features): - batch_size, _, num_points = features.size() - pt1 = features[:, 0:3, :] - - idx1 = [] - for i in range(len(self.k)): - idx = knn(pt1, self.k[i]) - idx1.append(idx) - - pt1 = pt1.transpose(1, 2).contiguous() - - x = features.unsqueeze(2) - x = self.sam_res1(x, idx1) - x1 = self.af(x) - - x, _, _, pt2 = self._edge_pooling(x1, pt1, self.rate, self.pk, self.pts_num[1]) - idx2 = [] - for i in range(len(self.k)): - idx = knn(pt2.transpose(1, 2).contiguous(), self.k[i]) - idx2.append(idx) - - x = self.sam_res2(x, idx2) - x2 = self.af(x) - - x, _, _, pt3 = self._edge_pooling(x2, pt2, self.rate, self.pk, self.pts_num[2]) - idx3 = [] - for i in range(len(self.k)): - idx = knn(pt3.transpose(1, 2).contiguous(), self.k[i]) - idx3.append(idx) - - x = self.sam_res3(x, idx3) - x3 = self.af(x) - - x, _, _, pt4 = self._edge_pooling(x3, pt3, self.rate, self.pk, self.pts_num[3]) - idx4 = [] - for i in range(len(self.k)): - idx = knn(pt4.transpose(1, 2).contiguous(), self.k[i]) - idx4.append(idx) - - x = self.sam_res4(x, idx4) - x4 = self.af(x) - x = self.conv5(x4) - x, _ = torch.max(x, -1) - x = x.view(batch_size, -1) - x = self.dropout(self.af(self.fc2(self.dropout(self.af(self.fc1(x)))))) - - x = x.unsqueeze(2).repeat(1, 1, self.pts_num[3]).unsqueeze(2) - x = self.af(self.conv6(torch.cat([x, x4], 1))) - x = self._edge_unpooling(x, pt4, pt3) - x = self.af(self.conv7(torch.cat([x, x3], 1))) - x = self._edge_unpooling(x, pt3, pt2) - x = self.af(self.conv8(torch.cat([x, x2], 1))) - x = self._edge_unpooling(x, pt2, pt1) - x = self.af(self.conv9(torch.cat([x, x1], 1))) - x = self.conv_out(x) - x = x.squeeze(2) - return x - - -class MSAP_SKN_decoder(nn.Module): - def __init__(self, num_coarse_raw, num_fps, num_coarse, num_fine, layers=[2, 2, 2, 2], knn_list=[10, 20], pk=10, - points_label=False, local_folding=False): - super(MSAP_SKN_decoder, self).__init__() - self.num_coarse_raw = num_coarse_raw - self.num_fps = num_fps - self.num_coarse = num_coarse - self.num_fine = num_fine - self.points_label = points_label - self.local_folding = local_folding - - self.fc1 = nn.Linear(1024, 1024) - self.fc2 = nn.Linear(1024, 1024) - self.fc3 = nn.Linear(1024, num_coarse_raw * 3) - - self.dense_feature_size = 256 - self.expand_feature_size = 64 - - if points_label: - self.input_size = 4 - else: - self.input_size = 3 - - - self.encoder = SA_SKN_Res_encoder(input_size=self.input_size, k=knn_list, pk=pk, - output_size=self.dense_feature_size, layers=layers) - - self.up_scale = int(np.ceil(num_fine / (num_coarse_raw + 2048))) - - if self.up_scale >= 2: - self.expansion1 = EF_expansion(input_size=self.dense_feature_size, output_size=self.expand_feature_size, - step_ratio=self.up_scale, k=4) - self.conv_cup1 = nn.Conv1d(self.expand_feature_size, self.expand_feature_size, 1) - else: - self.expansion1 = None - self.conv_cup1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - - self.conv_cup2 = nn.Conv1d(self.expand_feature_size, 3, 1, bias=True) - - self.conv_s1 = nn.Conv1d(self.expand_feature_size, 16, 1, bias=True) - self.conv_s2 = nn.Conv1d(16, 8, 1, bias=True) - self.conv_s3 = nn.Conv1d(8, 1, 1, bias=True) - - if self.local_folding: - self.expansion2 = Folding(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(num_fine // num_coarse)) - else: - self.expansion2 = EF_expansion(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(num_fine // num_coarse), k=4) - - self.conv_f1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - self.conv_f2 = nn.Conv1d(self.expand_feature_size, 3, 1) - - self.af = nn.ReLU(inplace=False) - - def forward(self, global_feat, point_input): - batch_size = global_feat.size()[0] - - coarse_raw = self.fc3(self.af(self.fc2(self.af(self.fc1(global_feat))))).view(batch_size, 3, - self.num_coarse_raw) - - input_points_num = point_input.size()[2] - org_points_input = point_input - - if self.points_label: - id0 = torch.zeros(coarse_raw.shape[0], 1, coarse_raw.shape[2]).cuda().contiguous() - coarse_input = torch.cat( (coarse_raw, id0), 1) - id1 = torch.ones(org_points_input.shape[0], 1, org_points_input.shape[2]).cuda().contiguous() - org_points_input = torch.cat( (org_points_input, id1), 1) - else: - coarse_input = coarse_raw - - points = torch.cat((coarse_input, org_points_input), 2) - dense_feat = self.encoder(points) - - if self.up_scale >= 2: - dense_feat = self.expansion1(dense_feat) - - coarse_features = self.af(self.conv_cup1(dense_feat)) - coarse_high = self.conv_cup2(coarse_features) - - if coarse_high.size()[2] > self.num_fps: - idx_fps = pn2.furthest_point_sample(coarse_high.transpose(1, 2).contiguous(), self.num_fps) - coarse_fps = pn2.gather_operation(coarse_high, idx_fps) - coarse_features = pn2.gather_operation(coarse_features, idx_fps) - else: - coarse_fps = coarse_high - - if coarse_fps.size()[2] > self.num_coarse: - scores = F.softplus(self.conv_s3(self.af(self.conv_s2(self.af(self.conv_s1(coarse_features)))))) - idx_scores = scores.topk(k=self.num_coarse, dim=2)[1].view(batch_size, -1).int() - coarse = pn2.gather_operation(coarse_fps, idx_scores) - coarse_features = pn2.gather_operation(coarse_features, idx_scores) - else: - coarse = coarse_fps - - if coarse.size()[2] < self.num_fine: - if self.local_folding: - up_features = self.expansion2(coarse_features, global_feat) - center = coarse.transpose(2, 1).contiguous().unsqueeze(2).repeat(1, 1, self.num_fine // self.num_coarse, - 1).view(batch_size, self.num_fine, - 3).transpose(2, 1).contiguous() - fine = self.conv_f2(self.af(self.conv_f1(up_features))) + center - else: - up_features = self.expansion2(coarse_features) - fine = self.conv_f2(self.af(self.conv_f1(up_features))) - else: - assert (coarse.size()[2] == self.num_fine) - fine = coarse - - return coarse_raw, coarse_high, coarse, fine - - -class Model(nn.Module): - def __init__(self, args, size_z=128, global_feature_size=1024): - super(Model, self).__init__() - - layers = [int(i) for i in args.layers.split(',')] - knn_list = [int(i) for i in args.knn_list.split(',')] - - self.size_z = size_z - self.distribution_loss = args.distribution_loss - self.train_loss = args.loss - self.encoder = PCN_encoder(output_size=global_feature_size) - self.posterior_infer1 = Linear_ResBlock(input_size=global_feature_size, output_size=global_feature_size) - self.posterior_infer2 = Linear_ResBlock(input_size=global_feature_size, output_size=size_z * 2) - self.prior_infer = Linear_ResBlock(input_size=global_feature_size, output_size=size_z * 2) - self.generator = Linear_ResBlock(input_size=size_z, output_size=global_feature_size) - self.decoder = MSAP_SKN_decoder(num_fps=args.num_fps, num_fine=args.num_points, num_coarse=args.num_coarse, - num_coarse_raw=args.num_coarse_raw, layers=layers, knn_list=knn_list, - pk=args.pk, local_folding=args.local_folding, points_label=args.points_label) - - def compute_kernel(self, x, y): - x_size = x.size()[0] - y_size = y.size()[0] - dim = x.size()[1] - - tiled_x = x.unsqueeze(1).repeat(1, y_size, 1) - tiled_y = y.unsqueeze(0).repeat(x_size, 1, 1) - return torch.exp(-torch.mean((tiled_x - tiled_y)**2, dim=2) / float(dim)) - - def mmd_loss(self, x, y): - x_kernel = self.compute_kernel(x, x) - y_kernel = self.compute_kernel(y, y) - xy_kernel = self.compute_kernel(x, y) - return torch.mean(x_kernel) + torch.mean(y_kernel) - 2 * torch.mean(xy_kernel) - - def forward(self, x, gt, is_training=True, mean_feature=None, alpha=None): - num_input = x.size()[2] - - if is_training: - y = pn2.gather_operation(gt.transpose(1, 2).contiguous(), pn2.furthest_point_sample(gt, num_input)) - gt = torch.cat([gt, gt], dim=0) - points = torch.cat([x, y], dim=0) - x = torch.cat([x, x], dim=0) - else: - points = x - feat = self.encoder(points) - - if is_training: - feat_x, feat_y = feat.chunk(2) - o_x = self.posterior_infer2(self.posterior_infer1(feat_x)) - q_mu, q_std = torch.split(o_x, self.size_z, dim=1) - o_y = self.prior_infer(feat_y) - p_mu, p_std = torch.split(o_y, self.size_z, dim=1) - q_std = F.softplus(q_std) - p_std = F.softplus(p_std) - q_distribution = torch.distributions.Normal(q_mu, q_std) - p_distribution = torch.distributions.Normal(p_mu, p_std) - p_distribution_fix = torch.distributions.Normal(p_mu.detach(), p_std.detach()) - m_distribution = torch.distributions.Normal(torch.zeros_like(p_mu), torch.ones_like(p_std)) - z_q = q_distribution.rsample() - z_p = p_distribution.rsample() - z = torch.cat([z_q, z_p], dim=0) - feat = torch.cat([feat_x, feat_x], dim=0) - - else: - o_x = self.posterior_infer2(self.posterior_infer1(feat)) - q_mu, q_std = torch.split(o_x, self.size_z, dim=1) - q_std = F.softplus(q_std) - q_distribution = torch.distributions.Normal(q_mu, q_std) - p_distribution = q_distribution - p_distribution_fix = p_distribution - m_distribution = p_distribution - z = q_distribution.rsample() - - feat += self.generator(z) - - coarse_raw, coarse_high, coarse, fine = self.decoder(feat, x) - coarse_raw = coarse_raw.transpose(1, 2).contiguous() - coarse_high = coarse_high.transpose(1, 2).contiguous() - coarse = coarse.transpose(1, 2).contiguous() - fine = fine.transpose(1, 2).contiguous() - - if is_training: - if self.distribution_loss == 'MMD': - z_m = m_distribution.rsample() - z_q = q_distribution.rsample() - z_p = p_distribution.rsample() - z_p_fix = p_distribution_fix.rsample() - dl_rec = self.mmd_loss(z_m, z_p) - dl_g = self.mmd_loss2(z_q, z_p_fix) - elif self.distribution_loss == 'KLD': - dl_rec = torch.distributions.kl_divergence(m_distribution, p_distribution) - dl_g = torch.distributions.kl_divergence(p_distribution_fix, q_distribution) - else: - raise NotImplementedError('Distribution loss is either MMD or KLD') - - if self.train_loss == 'cd': - loss1, _ = calc_cd(coarse_raw, gt) - loss2, _ = calc_cd(coarse_high, gt) - loss3, _ = calc_cd(coarse, gt) - loss4, _ = calc_cd(fine, gt) - else: - raise NotImplementedError('Only CD is supported') - - total_train_loss = loss1.mean() * 10 + loss2.mean() * 0.5 + loss3.mean() + loss4.mean() * alpha - total_train_loss += (dl_rec.mean() + dl_g.mean()) * 20 - return fine, loss4, total_train_loss - else: - emd = calc_emd(fine, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(fine, gt, calc_f1=True) - return {'out1': coarse_raw, 'out2': fine, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} diff --git a/other_models/VRCNet/models/vrcnet.py b/other_models/VRCNet/models/vrcnet.py deleted file mode 100644 index 309dc55..0000000 --- a/other_models/VRCNet/models/vrcnet.py +++ /dev/null @@ -1,515 +0,0 @@ -import numpy as np -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -import torch.nn.functional as F -from utils.model_utils import * -from models.pcn import PCN_encoder - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "utils/Pointnet2.PyTorch/pointnet2")) -import pointnet2_utils as pn2 - - -class SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=16): - super(SA_module, self).__init__() - self.share_planes = share_planes - self.k = k - self.conv1 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv2 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv3 = nn.Conv2d(in_planes, mid_planes, kernel_size=1) - - self.conv_w = nn.Sequential(nn.ReLU(inplace=False), - nn.Conv2d(rel_planes * (k + 1), mid_planes // share_planes, kernel_size=1, - bias=False), - nn.ReLU(inplace=False), - nn.Conv2d(mid_planes // share_planes, k * mid_planes // share_planes, - kernel_size=1)) - self.activation_fn = nn.ReLU(inplace=False) - - self.conv_out = nn.Conv2d(mid_planes, out_planes, kernel_size=1) - - def forward(self, input): - x, idx = input - batch_size, _, _, num_points = x.size() - identity = x # B C 1 N - x = self.activation_fn(x) - xn = get_edge_features(x, idx) # B C K N - x1, x2, x3 = self.conv1(x), self.conv2(xn), self.conv3(xn) - - x2 = x2.view(batch_size, -1, 1, num_points).contiguous() # B kC 1 N - w = self.conv_w(torch.cat([x1, x2], 1)).view(batch_size, -1, self.k, num_points) - w = w.repeat(1, self.share_planes, 1, 1) - out = w * x3 - out = torch.sum(out, dim=2, keepdim=True) - - out = self.activation_fn(out) - out = self.conv_out(out) # B C 1 N - out += identity - return [out, idx] - - -class Folding(nn.Module): - def __init__(self, input_size, output_size, step_ratio, global_feature_size=1024, num_models=1): - super(Folding, self).__init__() - self.input_size = input_size - self.output_size = output_size - self.step_ratio = step_ratio - self.num_models = num_models - - self.conv = nn.Conv1d(input_size + global_feature_size + 2, output_size, 1, bias=True) - - sqrted = int(math.sqrt(step_ratio)) + 1 - for i in range(1, sqrted + 1).__reversed__(): - if (step_ratio % i) == 0: - num_x = i - num_y = step_ratio // i - break - - grid_x = torch.linspace(-0.2, 0.2, steps=num_x) - grid_y = torch.linspace(-0.2, 0.2, steps=num_y) - - x, y = torch.meshgrid(grid_x, grid_y) # x, y shape: (2, 1) - self.grid = torch.stack([x, y], dim=-1).view(-1, 2) # (2, 2) - - def forward(self, point_feat, global_feat): - batch_size, num_features, num_points = point_feat.size() - point_feat = point_feat.transpose(1, 2).contiguous().unsqueeze(2).repeat(1, 1, self.step_ratio, 1).view( - batch_size, - -1, num_features).transpose(1, 2).contiguous() - global_feat = global_feat.unsqueeze(2).repeat(1, 1, num_points * self.step_ratio).repeat(self.num_models, 1, 1) - grid_feat = self.grid.unsqueeze(0).repeat(batch_size, num_points, 1).transpose(1, 2).contiguous().cuda() - features = torch.cat([global_feat, point_feat, grid_feat], axis=1) - features = F.relu(self.conv(features)) - return features - - -class Linear_ResBlock(nn.Module): - def __init__(self, input_size=1024, output_size=256): - super(Linear_ResBlock, self).__init__() - self.conv1 = nn.Linear(input_size, input_size) - self.conv2 = nn.Linear(input_size, output_size) - self.conv_res = nn.Linear(input_size, output_size) - - self.af = nn.ReLU(inplace=True) - - def forward(self, feature): - return self.conv2(self.af(self.conv1(self.af(feature)))) + self.conv_res(feature) - - -class SK_SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=[10, 20], r=2, L=32): - super(SK_SA_module, self).__init__() - - self.num_kernels = len(k) - d = max(int(out_planes / r), L) - - self.sams = nn.ModuleList([]) - - for i in range(len(k)): - self.sams.append(SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k[i])) - - self.fc = nn.Linear(out_planes, d) - self.fcs = nn.ModuleList([]) - - for i in range(len(k)): - self.fcs.append(nn.Linear(d, out_planes)) - - self.softmax = nn.Softmax(dim=1) - self.af = nn.ReLU(inplace=False) - - def forward(self, input): - x, idxs = input - assert (self.num_kernels == len(idxs)) - for i, sam in enumerate(self.sams): - fea, _ = sam([x, idxs[i]]) - fea = self.af(fea) - fea = fea.unsqueeze(dim=1) - if i == 0: - feas = fea - else: - feas = torch.cat([feas, fea], dim=1) - - fea_U = torch.sum(feas, dim=1) - fea_s = fea_U.mean(-1).mean(-1) - fea_z = self.fc(fea_s) - - for i, fc in enumerate(self.fcs): - vector = fc(fea_z).unsqueeze_(dim=1) - if i == 0: - attention_vectors = vector - else: - attention_vectors = torch.cat([attention_vectors, vector], dim=1) - - attention_vectors = self.softmax(attention_vectors) - attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1) - fea_v = (feas * attention_vectors).sum(dim=1) - return [fea_v, idxs] - - -class SKN_Res_unit(nn.Module): - def __init__(self, input_size, output_size, k=[10, 20], layers=1): - super(SKN_Res_unit, self).__init__() - self.conv1 = nn.Conv2d(input_size, output_size, 1, bias=False) - self.sam = self._make_layer(output_size, output_size // 16, output_size // 4, output_size, int(layers), 8, k=k) - self.conv2 = nn.Conv2d(output_size, output_size, 1, bias=False) - self.conv_res = nn.Conv2d(input_size, output_size, 1, bias=False) - self.af = nn.ReLU(inplace=False) - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def forward(self, feat, idx): - x, _ = self.sam([self.conv1(feat), idx]) - x = self.conv2(self.af(x)) - return x + self.conv_res(feat) - - -class SA_SKN_Res_encoder(nn.Module): - def __init__(self, input_size=3, k=[10, 20], pk=16, output_size=64, layers=[2, 2, 2, 2], - pts_num=[3072, 1536, 768, 384]): - super(SA_SKN_Res_encoder, self).__init__() - self.init_channel = 64 - - c1 = self.init_channel - self.sam_res1 = SKN_Res_unit(input_size, c1, k, int(layers[0])) - - c2 = c1 * 2 - self.sam_res2 = SKN_Res_unit(c2, c2, k, int(layers[1])) - - c3 = c2 * 2 - self.sam_res3 = SKN_Res_unit(c3, c3, k, int(layers[2])) - - c4 = c3 * 2 - self.sam_res4 = SKN_Res_unit(c4, c4, k, int(layers[3])) - - self.conv5 = nn.Conv2d(c4, 1024, 1) - - self.fc1 = nn.Linear(1024, 512) - self.fc2 = nn.Linear(512, 1024) - - self.conv6 = nn.Conv2d(c4 + 1024, c4, 1) - self.conv7 = nn.Conv2d(c3 + c4, c3, 1) - self.conv8 = nn.Conv2d(c2 + c3, c2, 1) - self.conv9 = nn.Conv2d(c1 + c2, c1, 1) - - self.conv_out = nn.Conv2d(c1, output_size, 1) - self.dropout = nn.Dropout() - self.af = nn.ReLU(inplace=False) - self.k = k - self.pk = pk - self.rate = 2 - - self.pts_num = pts_num - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def _edge_pooling(self, features, points, rate=2, k=16, sample_num=None): - features = features.squeeze(2) - - if sample_num is None: - input_points_num = int(features.size()[2]) - sample_num = input_points_num // rate - - ds_features, p_idx, pn_idx, ds_points = edge_preserve_sampling(features, points, sample_num, k) - ds_features = ds_features.unsqueeze(2) - return ds_features, p_idx, pn_idx, ds_points - - def _edge_unpooling(self, features, src_pts, tgt_pts): - features = features.squeeze(2) - idx, weight = three_nn_upsampling(tgt_pts, src_pts) - features = pn2.three_interpolate(features, idx, weight) - features = features.unsqueeze(2) - return features - - def forward(self, features): - batch_size, _, num_points = features.size() - pt1 = features[:, 0:3, :] - - idx1 = [] - for i in range(len(self.k)): - idx = knn(pt1, self.k[i]) - idx1.append(idx) - - pt1 = pt1.transpose(1, 2).contiguous() - - x = features.unsqueeze(2) - x = self.sam_res1(x, idx1) - x1 = self.af(x) - - x, _, _, pt2 = self._edge_pooling(x1, pt1, self.rate, self.pk, self.pts_num[1]) - idx2 = [] - for i in range(len(self.k)): - idx = knn(pt2.transpose(1, 2).contiguous(), self.k[i]) - idx2.append(idx) - - x = self.sam_res2(x, idx2) - x2 = self.af(x) - - x, _, _, pt3 = self._edge_pooling(x2, pt2, self.rate, self.pk, self.pts_num[2]) - idx3 = [] - for i in range(len(self.k)): - idx = knn(pt3.transpose(1, 2).contiguous(), self.k[i]) - idx3.append(idx) - - x = self.sam_res3(x, idx3) - x3 = self.af(x) - - x, _, _, pt4 = self._edge_pooling(x3, pt3, self.rate, self.pk, self.pts_num[3]) - idx4 = [] - for i in range(len(self.k)): - idx = knn(pt4.transpose(1, 2).contiguous(), self.k[i]) - idx4.append(idx) - - x = self.sam_res4(x, idx4) - x4 = self.af(x) - x = self.conv5(x4) - x, _ = torch.max(x, -1) - x = x.view(batch_size, -1) - x = self.dropout(self.af(self.fc2(self.dropout(self.af(self.fc1(x)))))) - - x = x.unsqueeze(2).repeat(1, 1, self.pts_num[3]).unsqueeze(2) - x = self.af(self.conv6(torch.cat([x, x4], 1))) - x = self._edge_unpooling(x, pt4, pt3) - x = self.af(self.conv7(torch.cat([x, x3], 1))) - x = self._edge_unpooling(x, pt3, pt2) - x = self.af(self.conv8(torch.cat([x, x2], 1))) - x = self._edge_unpooling(x, pt2, pt1) - x = self.af(self.conv9(torch.cat([x, x1], 1))) - x = self.conv_out(x) - x = x.squeeze(2) - return x - - -class MSAP_SKN_decoder(nn.Module): - def __init__(self, num_coarse_raw, num_fps, num_coarse, num_fine, layers=[2, 2, 2, 2], knn_list=[10, 20], pk=10, - points_label=False, local_folding=False): - super(MSAP_SKN_decoder, self).__init__() - self.num_coarse_raw = num_coarse_raw - self.num_fps = num_fps - self.num_coarse = num_coarse - self.num_fine = num_fine - self.points_label = points_label - self.local_folding = local_folding - - self.fc1 = nn.Linear(1024, 1024) - self.fc2 = nn.Linear(1024, 1024) - self.fc3 = nn.Linear(1024, num_coarse_raw * 3) - - self.dense_feature_size = 256 - self.expand_feature_size = 64 - - if points_label: - self.input_size = 4 - else: - self.input_size = 3 - - - self.encoder = SA_SKN_Res_encoder(input_size=self.input_size, k=knn_list, pk=pk, - output_size=self.dense_feature_size, layers=layers) - - self.up_scale = int(np.ceil(num_fine / (num_coarse_raw + 2048))) - - if self.up_scale >= 2: - self.expansion1 = EF_expansion(input_size=self.dense_feature_size, output_size=self.expand_feature_size, - step_ratio=self.up_scale, k=4) - self.conv_cup1 = nn.Conv1d(self.expand_feature_size, self.expand_feature_size, 1) - else: - self.expansion1 = None - self.conv_cup1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - - self.conv_cup2 = nn.Conv1d(self.expand_feature_size, 3, 1, bias=True) - - self.conv_s1 = nn.Conv1d(self.expand_feature_size, 16, 1, bias=True) - self.conv_s2 = nn.Conv1d(16, 8, 1, bias=True) - self.conv_s3 = nn.Conv1d(8, 1, 1, bias=True) - - if self.local_folding: - self.expansion2 = Folding(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(num_fine // num_coarse)) - else: - self.expansion2 = EF_expansion(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(num_fine // num_coarse), k=4) - - self.conv_f1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - self.conv_f2 = nn.Conv1d(self.expand_feature_size, 3, 1) - - self.af = nn.ReLU(inplace=False) - - def forward(self, global_feat, point_input): - batch_size = global_feat.size()[0] - - coarse_raw = self.fc3(self.af(self.fc2(self.af(self.fc1(global_feat))))).view(batch_size, 3, - self.num_coarse_raw) - - input_points_num = point_input.size()[2] - org_points_input = point_input - - if self.points_label: - id0 = torch.zeros(coarse_raw.shape[0], 1, coarse_raw.shape[2]).cuda().contiguous() - coarse_input = torch.cat( (coarse_raw, id0), 1) - id1 = torch.ones(org_points_input.shape[0], 1, org_points_input.shape[2]).cuda().contiguous() - org_points_input = torch.cat( (org_points_input, id1), 1) - else: - coarse_input = coarse_raw - - points = torch.cat((coarse_input, org_points_input), 2) - dense_feat = self.encoder(points) - - if self.up_scale >= 2: - dense_feat = self.expansion1(dense_feat) - - coarse_features = self.af(self.conv_cup1(dense_feat)) - coarse_high = self.conv_cup2(coarse_features) - - if coarse_high.size()[2] > self.num_fps: - idx_fps = pn2.furthest_point_sample(coarse_high.transpose(1, 2).contiguous(), self.num_fps) - coarse_fps = pn2.gather_operation(coarse_high, idx_fps) - coarse_features = pn2.gather_operation(coarse_features, idx_fps) - else: - coarse_fps = coarse_high - - if coarse_fps.size()[2] > self.num_coarse: - scores = F.softplus(self.conv_s3(self.af(self.conv_s2(self.af(self.conv_s1(coarse_features)))))) - idx_scores = scores.topk(k=self.num_coarse, dim=2)[1].view(batch_size, -1).int() - coarse = pn2.gather_operation(coarse_fps, idx_scores) - coarse_features = pn2.gather_operation(coarse_features, idx_scores) - else: - coarse = coarse_fps - - if coarse.size()[2] < self.num_fine: - if self.local_folding: - up_features = self.expansion2(coarse_features, global_feat) - center = coarse.transpose(2, 1).contiguous().unsqueeze(2).repeat(1, 1, self.num_fine // self.num_coarse, - 1).view(batch_size, self.num_fine, - 3).transpose(2, 1).contiguous() - fine = self.conv_f2(self.af(self.conv_f1(up_features))) + center - else: - up_features = self.expansion2(coarse_features) - fine = self.conv_f2(self.af(self.conv_f1(up_features))) - else: - assert (coarse.size()[2] == self.num_fine) - fine = coarse - - return coarse_raw, coarse_high, coarse, fine - - -class Model(nn.Module): - def __init__(self, args, size_z=128, global_feature_size=1024): - super(Model, self).__init__() - - layers = [int(i) for i in args.layers.split(',')] - knn_list = [int(i) for i in args.knn_list.split(',')] - - self.size_z = size_z - self.distribution_loss = args.distribution_loss - self.train_loss = args.loss - self.encoder = PCN_encoder(output_size=global_feature_size) - self.posterior_infer1 = Linear_ResBlock(input_size=global_feature_size, output_size=global_feature_size) - self.posterior_infer2 = Linear_ResBlock(input_size=global_feature_size, output_size=size_z * 2) - self.prior_infer = Linear_ResBlock(input_size=global_feature_size, output_size=size_z * 2) - self.generator = Linear_ResBlock(input_size=size_z, output_size=global_feature_size) - self.decoder = MSAP_SKN_decoder(num_fps=args.num_fps, num_fine=args.num_points, num_coarse=args.num_coarse, - num_coarse_raw=args.num_coarse_raw, layers=layers, knn_list=knn_list, - pk=args.pk, local_folding=args.local_folding, points_label=args.points_label) - - def compute_kernel(self, x, y): - x_size = x.size()[0] - y_size = y.size()[0] - dim = x.size()[1] - - tiled_x = x.unsqueeze(1).repeat(1, y_size, 1) - tiled_y = y.unsqueeze(0).repeat(x_size, 1, 1) - return torch.exp(-torch.mean((tiled_x - tiled_y)**2, dim=2) / float(dim)) - - def mmd_loss(self, x, y): - x_kernel = self.compute_kernel(x, x) - y_kernel = self.compute_kernel(y, y) - xy_kernel = self.compute_kernel(x, y) - return torch.mean(x_kernel) + torch.mean(y_kernel) - 2 * torch.mean(xy_kernel) - - def forward(self, x, gt, is_training=True, mean_feature=None, alpha=None): - num_input = x.size()[2] - - if is_training: - y = pn2.gather_operation(gt.transpose(1, 2).contiguous(), pn2.furthest_point_sample(gt, num_input)) - gt = torch.cat([gt, gt], dim=0) - points = torch.cat([x, y], dim=0) - x = torch.cat([x, x], dim=0) - else: - points = x - feat = self.encoder(points) - - if is_training: - feat_x, feat_y = feat.chunk(2) - o_x = self.posterior_infer2(self.posterior_infer1(feat_x)) - q_mu, q_std = torch.split(o_x, self.size_z, dim=1) - o_y = self.prior_infer(feat_y) - p_mu, p_std = torch.split(o_y, self.size_z, dim=1) - q_std = F.softplus(q_std) - p_std = F.softplus(p_std) - q_distribution = torch.distributions.Normal(q_mu, q_std) - p_distribution = torch.distributions.Normal(p_mu, p_std) - p_distribution_fix = torch.distributions.Normal(p_mu.detach(), p_std.detach()) - m_distribution = torch.distributions.Normal(torch.zeros_like(p_mu), torch.ones_like(p_std)) - z_q = q_distribution.rsample() - z_p = p_distribution.rsample() - z = torch.cat([z_q, z_p], dim=0) - feat = torch.cat([feat_x, feat_x], dim=0) - - else: - o_x = self.posterior_infer2(self.posterior_infer1(feat)) - q_mu, q_std = torch.split(o_x, self.size_z, dim=1) - q_std = F.softplus(q_std) - q_distribution = torch.distributions.Normal(q_mu, q_std) - p_distribution = q_distribution - p_distribution_fix = p_distribution - m_distribution = p_distribution - z = q_distribution.rsample() - - feat += self.generator(z) - - coarse_raw, coarse_high, coarse, fine = self.decoder(feat, x) - coarse_raw = coarse_raw.transpose(1, 2).contiguous() - coarse_high = coarse_high.transpose(1, 2).contiguous() - coarse = coarse.transpose(1, 2).contiguous() - fine = fine.transpose(1, 2).contiguous() - - if is_training: - if self.distribution_loss == 'MMD': - z_m = m_distribution.rsample() - z_q = q_distribution.rsample() - z_p = p_distribution.rsample() - z_p_fix = p_distribution_fix.rsample() - dl_rec = self.mmd_loss(z_m, z_p) - dl_g = self.mmd_loss2(z_q, z_p_fix) - elif self.distribution_loss == 'KLD': - dl_rec = torch.distributions.kl_divergence(m_distribution, p_distribution) - dl_g = torch.distributions.kl_divergence(p_distribution_fix, q_distribution) - else: - raise NotImplementedError('Distribution loss is either MMD or KLD') - - if self.train_loss == 'cd': - loss1, _ = calc_cd(coarse_raw, gt) - loss2, _ = calc_cd(coarse_high, gt) - loss3, _ = calc_cd(coarse, gt) - loss4, _ = calc_cd(fine, gt) - else: - raise NotImplementedError('Only CD is supported') - - total_train_loss = loss1.mean() * 10 + loss2.mean() * 0.5 + loss3.mean() + loss4.mean() * alpha - total_train_loss += (dl_rec.mean() + dl_g.mean()) * 20 - return fine, loss4, total_train_loss - else: - emd = calc_emd(fine, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(fine, gt, calc_f1=True) - return {'out1': coarse_raw, 'out2': fine, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} diff --git a/other_models/VRCNet/utils/__pycache__/model_utils.cpython-36.pyc b/other_models/VRCNet/utils/__pycache__/model_utils.cpython-36.pyc deleted file mode 100644 index a943edc..0000000 Binary files a/other_models/VRCNet/utils/__pycache__/model_utils.cpython-36.pyc and /dev/null differ diff --git a/other_models/VRCNet/utils/model_utils.py b/other_models/VRCNet/utils/model_utils.py deleted file mode 100644 index 5ee50e3..0000000 --- a/other_models/VRCNet/utils/model_utils.py +++ /dev/null @@ -1,285 +0,0 @@ -import torch -import math -import os -import sys -import torch.nn as nn -import torch.nn.functional as F - - - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "utils/pointnet2")) -import pointnet2_utils as pn2 -sys.path.append(os.path.join(proj_dir, "../../distance/emd")) -import emd_module as emd -sys.path.append(os.path.join(proj_dir, "../../distance/chamfer_multidim")) -from chamfer3D import dist_chamfer_3D -from fscore import fscore - - -class EF_expansion(nn.Module): - def __init__(self, input_size, output_size=64, step_ratio=2, k=4): - super(EF_expansion, self).__init__() - self.step_ratio = step_ratio - self.k = k - self.input_size = input_size - self.output_size = output_size - - self.conv1 = nn.Conv2d(input_size * 2, output_size, 1) - self.conv2 = nn.Conv2d(input_size * 2 + output_size, output_size * step_ratio, 1) - self.conv3 = nn.Conv2d(output_size, output_size, 1) - - def forward(self, x): - batch_size, _, num_points = x.size() - - input_edge_feature = get_graph_feature(x, self.k, minus_center=False).permute(0, 1, 3, - 2).contiguous() # B C K N - edge_feature = self.conv1(input_edge_feature) - edge_feature = F.relu(torch.cat((edge_feature, input_edge_feature), 1)) - - edge_feature = F.relu(self.conv2(edge_feature)) # B C K N - edge_feature = edge_feature.permute(0, 2, 3, 1).contiguous().view(batch_size, self.k, - num_points * self.step_ratio, - self.output_size).permute(0, 3, 1, - 2) # B C K N - - edge_feature = self.conv3(edge_feature) - edge_feature, _ = torch.max(edge_feature, 2) - - return edge_feature - - -def attention(query, key, value, mask=None): - d_k = query.size(-1) - scores = torch.matmul(query, key.transpose(-2, -1).contiguous()) / math.sqrt(d_k) # B x 4 x points x points - if mask is not None: - scores = scores.masked_fill(mask == 0, -1e9) - p_attn = F.softmax(scores, dim=-1) - return torch.matmul(p_attn, value), p_attn - - -def calc_cd(output, gt, calc_f1=False): - cham_loss = dist_chamfer_3D.chamfer_3DDist() - dist1, dist2, _, _ = cham_loss(gt, output) - cd_p = (torch.sqrt(dist1).mean(1) + torch.sqrt(dist2).mean(1)) / 2 - cd_t = (dist1.mean(1) + dist2.mean(1)) - if calc_f1: - f1, _, _ = fscore(dist1, dist2) - return cd_p, cd_t, f1 - else: - return cd_p, cd_t - - -def calc_emd(output, gt, eps=0.005, iterations=50): - emd_loss = emd.emdModule() - dist, _ = emd_loss(output, gt, eps, iterations) - emd_out = torch.sqrt(dist).mean(1) - return emd_out - - -def edge_preserve_sampling(feature_input, point_input, num_samples, k=10): - batch_size = feature_input.size()[0] - feature_size = feature_input.size()[1] - num_points = feature_input.size()[2] - - p_idx = pn2.furthest_point_sample(point_input, num_samples) - point_output = pn2.gather_operation(point_input.transpose(1, 2).contiguous(), p_idx).transpose(1, - 2).contiguous() - - pk = int(min(k, num_points)) - _, pn_idx = knn_point(pk, point_input, point_output) - pn_idx = pn_idx.detach().int() - neighbor_feature = pn2.gather_operation(feature_input, pn_idx.view(batch_size, num_samples * pk)).view(batch_size, - feature_size, - num_samples, - pk) - neighbor_feature, _ = torch.max(neighbor_feature, 3) - - center_feature = pn2.grouping_operation(feature_input, p_idx.unsqueeze(2)).view(batch_size, -1, num_samples) - - net = torch.cat((center_feature, neighbor_feature), 1) - - return net, p_idx, pn_idx, point_output - - -def get_edge_features(x, idx): - batch_size, num_points, k = idx.size() - device = torch.device('cuda') - idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points - idx = idx + idx_base - idx = idx.view(-1) - x = x.squeeze(2) - _, num_dims, _ = x.size() - x = x.transpose(2, 1).contiguous() - feature = x.view(batch_size * num_points, -1)[idx, :] - feature = feature.view(batch_size, num_points, k, num_dims).permute(0, 3, 2, 1) # B, C, K, N - return feature - - -def gen_grid(num_grid_point): - x = torch.linspace(-0.05, 0.05, steps=num_grid_point) - x, y = torch.meshgrid(x, x) - grid = torch.stack([x, y], axis=-1).view(2, num_grid_point ** 2) - return grid - - -def gen_1d_grid(num_grid_point): - x = torch.linspace(-0.05, 0.05, num_grid_point) - grid = x.view(1, num_grid_point) - return grid - - -def gen_grid_up(up_ratio, grid_size=0.2): - sqrted = int(math.sqrt(up_ratio)) + 1 - for i in range(1, sqrted + 1).__reversed__(): - if (up_ratio % i) == 0: - num_x = i - num_y = up_ratio // i - break - - grid_x = torch.linspace(-grid_size, grid_size, steps=num_x) - grid_y = torch.linspace(-grid_size, grid_size, steps=num_y) - - x, y = torch.meshgrid(grid_x, grid_y) # x, y shape: (2, 1) - grid = torch.stack([x, y], dim=-1).view(-1, 2).transpose(0, 1).contiguous() - return grid - - -def get_graph_feature(x, k=20, minus_center=True): - idx = knn(x, k=k) - batch_size, num_points, _ = idx.size() - device = torch.device('cuda') - - idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points - - idx = idx + idx_base - - idx = idx.view(-1) - - _, num_dims, _ = x.size() - - x = x.transpose(2, 1).contiguous() - feature = x.view(batch_size * num_points, -1)[idx, :] - feature = feature.view(batch_size, num_points, k, num_dims) - x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1) - - if minus_center: - feature = torch.cat((x, feature - x), dim=3).permute(0, 3, 1, 2) - else: - feature = torch.cat((x, feature), dim=3).permute(0, 3, 1, 2) - return feature - - -def get_repulsion_loss(pred, nsample=20, radius=0.07): - # pred: (batch_size, npoint,3) - # idx = pn2.ball_query(radius, nsample, pred, pred) - idx = knn(pred.transpose(1, 2).contiguous(), nsample).int() - pred_flipped = pred.transpose(1, 2).contiguous() - grouped_pred = pn2.grouping_operation(pred_flipped, idx) # (B, C, npoint, nsample) - grouped_pred -= pred_flipped.unsqueeze(-1) - - # get the uniform loss - h = 0.03 - dist_square = torch.sum(grouped_pred ** 2, dim=1) - dist_square, idx = torch.topk(-dist_square, 5) - dist_square = -dist_square[:, :, 1:] # remove the first one - dist_square = torch.max(torch.FloatTensor([1e-12]).expand_as(dist_square).cuda(), dist_square) - dist = torch.sqrt(dist_square) - weight = torch.exp(-dist_square / h ** 2) - uniform_loss = torch.mean(radius - dist * weight) - return uniform_loss - - -def get_uniform_loss(pcd, percentages=[0.004, 0.006, 0.008, 0.010, 0.012], radius=1.0): - B, N, C = pcd.size() - npoint = int(N * 0.05) - loss = 0 - for p in percentages: - nsample = int(N*p) - r = math.sqrt(p*radius) - disk_area = math.pi * (radius ** 2) * p/nsample - new_xyz = pn2.gather_operation(pcd.transpose(1, 2).contiguous(), - pn2.furthest_point_sample(pcd, npoint)).transpose(1, 2).contiguous() - idx = pn2.ball_query(r, nsample, pcd, new_xyz) - expect_len = math.sqrt(disk_area) - - grouped_pcd = pn2.grouping_operation(pcd.transpose(1,2).contiguous(), idx) - grouped_pcd = grouped_pcd.permute(0, 2, 3, 1).contiguous().view(-1, nsample, 3) - - var, _ = knn_point(2, grouped_pcd, grouped_pcd) - uniform_dis = -var[:, :, 1:] - - uniform_dis = torch.sqrt(torch.abs(uniform_dis+1e-8)) - uniform_dis = torch.mean(uniform_dis, dim=-1) - uniform_dis = ((uniform_dis - expect_len)**2 / (expect_len + 1e-8)) - - mean = torch.mean(uniform_dis) - mean = mean*math.pow(p*100,2) - loss += mean - return loss/len(percentages) - - -def index_points(points, idx): - device = points.device - B = points.shape[0] - view_shape = list(idx.shape) - view_shape[1:] = [1] * (len(view_shape) - 1) - repeat_shape = list(idx.shape) - repeat_shape[0] = 1 - batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape) - new_points = points[batch_indices, idx, :] - return new_points - - -def knn(x, k): - inner = -2 * torch.matmul(x.transpose(2, 1).contiguous(), x) - xx = torch.sum(x ** 2, dim=1, keepdim=True) - pairwise_distance = -xx - inner - xx.transpose(2, 1).contiguous() - idx = pairwise_distance.topk(k=k, dim=-1)[1] - return idx - - -def knn_point(pk, point_input, point_output): - m = point_output.size()[1] - n = point_input.size()[1] - - inner = -2 * torch.matmul(point_output, point_input.transpose(2, 1).contiguous()) - xx = torch.sum(point_output ** 2, dim=2, keepdim=True).repeat(1, 1, n) - yy = torch.sum(point_input ** 2, dim=2, keepdim=False).unsqueeze(1).repeat(1, m, 1) - pairwise_distance = -xx - inner - yy - dist, idx = pairwise_distance.topk(k=pk, dim=-1) - return dist, idx - - -def knn_point_all(pk, point_input, point_output): - m = point_output.size()[1] - n = point_input.size()[1] - - inner = -2 * torch.matmul(point_output, point_input.transpose(2, 1).contiguous()) - xx = torch.sum(point_output ** 2, dim=2, keepdim=True).repeat(1, 1, n) - yy = torch.sum(point_input ** 2, dim=2, keepdim=False).unsqueeze(1).repeat(1, m, 1) - pairwise_distance = -xx - inner - yy - dist, idx = pairwise_distance.topk(k=pk, dim=-1) - - return dist, idx - - -def symmetric_sample(points, num=512): - p1_idx = pn2.furthest_point_sample(points, num) - input_fps = pn2.gather_operation(points.transpose(1, 2).contiguous(), p1_idx).transpose(1, 2).contiguous() - x = torch.unsqueeze(input_fps[:, :, 0], dim=2) - y = torch.unsqueeze(input_fps[:, :, 1], dim=2) - z = torch.unsqueeze(-input_fps[:, :, 2], dim=2) - input_fps_flip = torch.cat([x, y, z], dim=2) - input_fps = torch.cat([input_fps, input_fps_flip], dim=1) - return input_fps - - -def three_nn_upsampling(target_points, source_points): - dist, idx = pn2.three_nn(target_points, source_points) - dist = torch.max(dist, torch.ones(1).cuda() * 1e-10) - norm = torch.sum((1.0 / dist), 2, keepdim=True) - norm = norm.repeat(1, 1, 3) - weight = (1.0 / dist) / norm - - return idx, weight diff --git a/other_models/VRCNet/utils/pointnet2/__pycache__/pointnet2_utils.cpython-36.pyc b/other_models/VRCNet/utils/pointnet2/__pycache__/pointnet2_utils.cpython-36.pyc deleted file mode 100644 index 2ef6891..0000000 Binary files a/other_models/VRCNet/utils/pointnet2/__pycache__/pointnet2_utils.cpython-36.pyc and /dev/null differ diff --git a/other_models/VRCNet/utils/pointnet2/dist/pointnet2-0.0.0-py2.7-linux-x86_64.egg b/other_models/VRCNet/utils/pointnet2/dist/pointnet2-0.0.0-py2.7-linux-x86_64.egg deleted file mode 100644 index d3b83d3..0000000 Binary files a/other_models/VRCNet/utils/pointnet2/dist/pointnet2-0.0.0-py2.7-linux-x86_64.egg and /dev/null differ diff --git a/other_models/VRCNet/utils/pointnet2/dist/pointnet2-0.0.0-py3.6-linux-x86_64.egg b/other_models/VRCNet/utils/pointnet2/dist/pointnet2-0.0.0-py3.6-linux-x86_64.egg deleted file mode 100644 index 57ee19f..0000000 Binary files a/other_models/VRCNet/utils/pointnet2/dist/pointnet2-0.0.0-py3.6-linux-x86_64.egg and /dev/null differ diff --git a/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/PKG-INFO b/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/PKG-INFO deleted file mode 100644 index a33dd2c..0000000 --- a/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 2.1 -Name: pointnet2 -Version: 0.0.0 -Summary: UNKNOWN -Home-page: UNKNOWN -License: UNKNOWN -Platform: UNKNOWN - -UNKNOWN - diff --git a/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/SOURCES.txt b/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/SOURCES.txt deleted file mode 100644 index 21253f6..0000000 --- a/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/SOURCES.txt +++ /dev/null @@ -1,14 +0,0 @@ -setup.py -pointnet2.egg-info/PKG-INFO -pointnet2.egg-info/SOURCES.txt -pointnet2.egg-info/dependency_links.txt -pointnet2.egg-info/top_level.txt -src/ball_query.cpp -src/ball_query_gpu.cu -src/group_points.cpp -src/group_points_gpu.cu -src/interpolate.cpp -src/interpolate_gpu.cu -src/pointnet2_api.cpp -src/sampling.cpp -src/sampling_gpu.cu \ No newline at end of file diff --git a/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/dependency_links.txt b/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/top_level.txt b/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/top_level.txt deleted file mode 100644 index 2d59a59..0000000 --- a/other_models/VRCNet/utils/pointnet2/pointnet2.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -pointnet2_cuda diff --git a/other_models/VRCNet/utils/pointnet2/pointnet2_modules.py b/other_models/VRCNet/utils/pointnet2/pointnet2_modules.py deleted file mode 100644 index 674e6f9..0000000 --- a/other_models/VRCNet/utils/pointnet2/pointnet2_modules.py +++ /dev/null @@ -1,156 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -import pointnet2_utils -import pytorch_utils as pt_utils -from typing import List - - -class _PointnetSAModuleBase(nn.Module): - - def __init__(self): - super().__init__() - self.npoint = None - self.groupers = None - self.mlps = None - self.pool_method = 'max_pool' - - def forward(self, xyz: torch.Tensor, features: torch.Tensor = None, new_xyz=None) -> (torch.Tensor, torch.Tensor): - """ - :param xyz: (B, N, 3) tensor of the xyz coordinates of the features - :param features: (B, N, C) tensor of the descriptors of the the features - :param new_xyz: - :return: - new_xyz: (B, npoint, 3) tensor of the new features' xyz - new_features: (B, npoint, \sum_k(mlps[k][-1])) tensor of the new_features descriptors - """ - new_features_list = [] - xyz_flipped = xyz.transpose(1, 2).contiguous() - if new_xyz is None: - new_xyz = pointnet2_utils.gather_operation( - xyz_flipped, - pointnet2_utils.furthest_point_sample(xyz, self.npoint) - ).transpose(1, 2).contiguous() if self.npoint is not None else None - for i in range(len(self.groupers)): - new_features = self.groupers[i](xyz, new_xyz, features) # (B, C, npoint, nsample) - new_features = self.mlps[i](new_features) # (B, mlp[-1], npoint, nsample) - if self.pool_method == 'max_pool': - new_features = F.max_pool2d( - new_features, kernel_size=[1, new_features.size(3)] - ) # (B, mlp[-1], npoint, 1) - elif self.pool_method == 'avg_pool': - new_features = F.avg_pool2d( - new_features, kernel_size=[1, new_features.size(3)] - ) # (B, mlp[-1], npoint, 1) - else: - raise NotImplementedError - - new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) - new_features_list.append(new_features) - return new_xyz, torch.cat(new_features_list, dim=1) - - -class PointnetSAModuleMSG(_PointnetSAModuleBase): - """Pointnet set abstraction layer with multiscale grouping""" - - def __init__(self, *, npoint: int, radii: List[float], nsamples: List[int], mlps: List[List[int]], bn: bool = True, - use_xyz: bool = True, pool_method='max_pool', instance_norm=False, features: torch.Tensor = None): - """ - :param npoint: int - :param radii: list of float, list of radii to group with - :param nsamples: list of int, number of samples in each ball query - :param mlps: list of list of int, spec of the pointnet before the global pooling for each scale - :param bn: whether to use batchnorm - :param use_xyz: - :param pool_method: max_pool / avg_pool - :param instance_norm: whether to use instance_norm - """ - super().__init__() - - assert len(radii) == len(nsamples) == len(mlps) - - self.npoint = npoint - self.groupers = nn.ModuleList() - self.mlps = nn.ModuleList() - for i in range(len(radii)): - radius = radii[i] - nsample = nsamples[i] - self.groupers.append( - pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz) - if npoint is not None else pointnet2_utils.GroupAll(use_xyz) - ) - mlp_spec = mlps[i] - - if features is not None: - if use_xyz: - mlp_spec[0]+=3 - self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn, instance_norm=instance_norm)) - self.pool_method = pool_method - - -class PointnetSAModule(PointnetSAModuleMSG): - """Pointnet set abstraction layer""" - - def __init__(self, *, mlp: List[int], npoint: int = None, radius: float = None, nsample: int = None, - bn: bool = True, use_xyz: bool = True, pool_method='max_pool', instance_norm=False): - """ - :param mlp: list of int, spec of the pointnet before the global max_pool - :param npoint: int, number of features - :param radius: float, radius of ball - :param nsample: int, number of samples in the ball query - :param bn: whether to use batchnorm - :param use_xyz: - :param pool_method: max_pool / avg_pool - :param instance_norm: whether to use instance_norm - """ - super().__init__( - mlps=[mlp], npoint=npoint, radii=[radius], nsamples=[nsample], bn=bn, use_xyz=use_xyz, - pool_method=pool_method, instance_norm=instance_norm - ) - - -class PointnetFPModule(nn.Module): - r"""Propigates the features of one set to another""" - - def __init__(self, *, mlp: List[int], bn: bool = True): - """ - :param mlp: list of int - :param bn: whether to use batchnorm - """ - super().__init__() - self.mlp = pt_utils.SharedMLP(mlp, bn=bn) - - def forward( - self, unknown: torch.Tensor, known: torch.Tensor, unknow_feats: torch.Tensor, known_feats: torch.Tensor - ) -> torch.Tensor: - """ - :param unknown: (B, n, 3) tensor of the xyz positions of the unknown features - :param known: (B, m, 3) tensor of the xyz positions of the known features - :param unknow_feats: (B, C1, n) tensor of the features to be propigated to - :param known_feats: (B, C2, m) tensor of features to be propigated - :return: - new_features: (B, mlp[-1], n) tensor of the features of the unknown features - """ - if known is not None: - dist, idx = pointnet2_utils.three_nn(unknown, known) - dist_recip = 1.0 / (dist + 1e-8) - norm = torch.sum(dist_recip, dim=2, keepdim=True) - weight = dist_recip / norm - - interpolated_feats = pointnet2_utils.three_interpolate(known_feats, idx, weight) - else: - interpolated_feats = known_feats.expand(*known_feats.size()[0:2], unknown.size(1)) - - if unknow_feats is not None: - new_features = torch.cat([interpolated_feats, unknow_feats], dim=1) # (B, C2 + C1, n) - else: - new_features = interpolated_feats - - new_features = new_features.unsqueeze(-1) - new_features = self.mlp(new_features) - - return new_features.squeeze(-1) - - -if __name__ == "__main__": - pass diff --git a/other_models/VRCNet/utils/pointnet2/pointnet2_utils.py b/other_models/VRCNet/utils/pointnet2/pointnet2_utils.py deleted file mode 100644 index e814102..0000000 --- a/other_models/VRCNet/utils/pointnet2/pointnet2_utils.py +++ /dev/null @@ -1,290 +0,0 @@ -import torch -from torch.autograd import Variable -from torch.autograd import Function -import torch.nn as nn -from typing import Tuple - -import pointnet2_cuda as pointnet2 - - -class FurthestPointSampling(Function): - @staticmethod - def forward(ctx, xyz: torch.Tensor, npoint: int) -> torch.Tensor: - """ - Uses iterative furthest point sampling to select a set of npoint features that have the largest - minimum distance - :param ctx: - :param xyz: (B, N, 3) where N > npoint - :param npoint: int, number of features in the sampled set - :return: - output: (B, npoint) tensor containing the set - """ - assert xyz.is_contiguous() - - B, N, _ = xyz.size() - output = torch.cuda.IntTensor(B, npoint) - temp = torch.cuda.FloatTensor(B, N).fill_(1e10) - - pointnet2.furthest_point_sampling_wrapper(B, N, npoint, xyz, temp, output) - return output - - @staticmethod - def backward(xyz, a=None): - return None, None - - -furthest_point_sample = FurthestPointSampling.apply - - -class GatherOperation(Function): - - @staticmethod - def forward(ctx, features: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: - """ - :param ctx: - :param features: (B, C, N) - :param idx: (B, npoint) index tensor of the features to gather - :return: - output: (B, C, npoint) - """ - assert features.is_contiguous() - assert idx.is_contiguous() - - B, npoint = idx.size() - _, C, N = features.size() - output = torch.cuda.FloatTensor(B, C, npoint) - - pointnet2.gather_points_wrapper(B, C, N, npoint, features, idx, output) - - ctx.for_backwards = (idx, C, N) - return output - - @staticmethod - def backward(ctx, grad_out): - idx, C, N = ctx.for_backwards - B, npoint = idx.size() - - grad_features = Variable(torch.cuda.FloatTensor(B, C, N).zero_()) - grad_out_data = grad_out.data.contiguous() - pointnet2.gather_points_grad_wrapper(B, C, N, npoint, grad_out_data, idx, grad_features.data) - return grad_features, None - - -gather_operation = GatherOperation.apply - - -class ThreeNN(Function): - - @staticmethod - def forward(ctx, unknown: torch.Tensor, known: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Find the three nearest neighbors of unknown in known - :param ctx: - :param unknown: (B, N, 3) - :param known: (B, M, 3) - :return: - dist: (B, N, 3) l2 distance to the three nearest neighbors - idx: (B, N, 3) index of 3 nearest neighbors - """ - assert unknown.is_contiguous() - assert known.is_contiguous() - - B, N, _ = unknown.size() - m = known.size(1) - dist2 = torch.cuda.FloatTensor(B, N, 3) - idx = torch.cuda.IntTensor(B, N, 3) - - pointnet2.three_nn_wrapper(B, N, m, unknown, known, dist2, idx) - return torch.sqrt(dist2), idx - - @staticmethod - def backward(ctx, a=None, b=None): - return None, None - - -three_nn = ThreeNN.apply - - -class ThreeInterpolate(Function): - - @staticmethod - def forward(ctx, features: torch.Tensor, idx: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - """ - Performs weight linear interpolation on 3 features - :param ctx: - :param features: (B, C, M) Features descriptors to be interpolated from - :param idx: (B, n, 3) three nearest neighbors of the target features in features - :param weight: (B, n, 3) weights - :return: - output: (B, C, N) tensor of the interpolated features - """ - assert features.is_contiguous() - assert idx.is_contiguous() - assert weight.is_contiguous() - - B, c, m = features.size() - n = idx.size(1) - ctx.three_interpolate_for_backward = (idx, weight, m) - output = torch.cuda.FloatTensor(B, c, n) - - pointnet2.three_interpolate_wrapper(B, c, m, n, features, idx, weight, output) - return output - - @staticmethod - def backward(ctx, grad_out: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - :param ctx: - :param grad_out: (B, C, N) tensor with gradients of outputs - :return: - grad_features: (B, C, M) tensor with gradients of features - None: - None: - """ - idx, weight, m = ctx.three_interpolate_for_backward - B, c, n = grad_out.size() - - grad_features = Variable(torch.cuda.FloatTensor(B, c, m).zero_()) - grad_out_data = grad_out.data.contiguous() - - pointnet2.three_interpolate_grad_wrapper(B, c, n, m, grad_out_data, idx, weight, grad_features.data) - return grad_features, None, None - - -three_interpolate = ThreeInterpolate.apply - - -class GroupingOperation(Function): - - @staticmethod - def forward(ctx, features: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: - """ - :param ctx: - :param features: (B, C, N) tensor of features to group - :param idx: (B, npoint, nsample) tensor containing the indicies of features to group with - :return: - output: (B, C, npoint, nsample) tensor - """ - assert features.is_contiguous() - assert idx.is_contiguous() - - B, nfeatures, nsample = idx.size() - _, C, N = features.size() - output = torch.cuda.FloatTensor(B, C, nfeatures, nsample) - - pointnet2.group_points_wrapper(B, C, N, nfeatures, nsample, features, idx, output) - - ctx.for_backwards = (idx, N) - return output - - @staticmethod - def backward(ctx, grad_out: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """ - :param ctx: - :param grad_out: (B, C, npoint, nsample) tensor of the gradients of the output from forward - :return: - grad_features: (B, C, N) gradient of the features - """ - idx, N = ctx.for_backwards - - B, C, npoint, nsample = grad_out.size() - grad_features = Variable(torch.cuda.FloatTensor(B, C, N).zero_()) - - grad_out_data = grad_out.data.contiguous() - pointnet2.group_points_grad_wrapper(B, C, N, npoint, nsample, grad_out_data, idx, grad_features.data) - return grad_features, None - - -grouping_operation = GroupingOperation.apply - - -class BallQuery(Function): - - @staticmethod - def forward(ctx, radius: float, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor) -> torch.Tensor: - """ - :param ctx: - :param radius: float, radius of the balls - :param nsample: int, maximum number of features in the balls - :param xyz: (B, N, 3) xyz coordinates of the features - :param new_xyz: (B, npoint, 3) centers of the ball query - :return: - idx: (B, npoint, nsample) tensor with the indicies of the features that form the query balls - """ - assert new_xyz.is_contiguous() - assert xyz.is_contiguous() - - B, N, _ = xyz.size() - npoint = new_xyz.size(1) - idx = torch.cuda.IntTensor(B, npoint, nsample).zero_() - - pointnet2.ball_query_wrapper(B, N, npoint, radius, nsample, new_xyz, xyz, idx) - return idx - - @staticmethod - def backward(ctx, a=None): - return None, None, None, None - - -ball_query = BallQuery.apply - - -class QueryAndGroup(nn.Module): - def __init__(self, radius: float, nsample: int, use_xyz: bool = True): - """ - :param radius: float, radius of ball - :param nsample: int, maximum number of features to gather in the ball - :param use_xyz: - """ - super().__init__() - self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz - - def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor, features: torch.Tensor = None) -> Tuple[torch.Tensor]: - """ - :param xyz: (B, N, 3) xyz coordinates of the features - :param new_xyz: (B, npoint, 3) centroids - :param features: (B, C, N) descriptors of the features - :return: - new_features: (B, 3 + C, npoint, nsample) - """ - idx = ball_query(self.radius, self.nsample, xyz, new_xyz) - xyz_trans = xyz.transpose(1, 2).contiguous() - grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) - grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1) - - if features is not None: - grouped_features = grouping_operation(features, idx) - if self.use_xyz: - new_features = torch.cat([grouped_xyz, grouped_features], dim=1) # (B, C + 3, npoint, nsample) - else: - new_features = grouped_features - else: - assert self.use_xyz, "Cannot have not features and not use xyz as a feature!" - new_features = grouped_xyz - - return new_features - - -class GroupAll(nn.Module): - def __init__(self, use_xyz: bool = True): - super().__init__() - self.use_xyz = use_xyz - - def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor, features: torch.Tensor = None): - """ - :param xyz: (B, N, 3) xyz coordinates of the features - :param new_xyz: ignored - :param features: (B, C, N) descriptors of the features - :return: - new_features: (B, C + 3, 1, N) - """ - grouped_xyz = xyz.transpose(1, 2).unsqueeze(2) - if features is not None: - grouped_features = features.unsqueeze(2) - if self.use_xyz: - new_features = torch.cat([grouped_xyz, grouped_features], dim=1) # (B, 3 + C, 1, N) - else: - new_features = grouped_features - else: - new_features = grouped_xyz - - return new_features diff --git a/other_models/VRCNet/utils/pointnet2/pytorch_utils.py b/other_models/VRCNet/utils/pointnet2/pytorch_utils.py deleted file mode 100644 index 09cb7bc..0000000 --- a/other_models/VRCNet/utils/pointnet2/pytorch_utils.py +++ /dev/null @@ -1,236 +0,0 @@ -import torch.nn as nn -from typing import List, Tuple - - -class SharedMLP(nn.Sequential): - - def __init__( - self, - args: List[int], - *, - bn: bool = False, - activation=nn.ReLU(inplace=True), - preact: bool = False, - first: bool = False, - name: str = "", - instance_norm: bool = False, - ): - super().__init__() - - for i in range(len(args) - 1): - self.add_module( - name + 'layer{}'.format(i), - Conv2d( - args[i], - args[i + 1], - bn=(not first or not preact or (i != 0)) and bn, - activation=activation - if (not first or not preact or (i != 0)) else None, - preact=preact, - instance_norm=instance_norm - ) - ) - - -class _ConvBase(nn.Sequential): - - def __init__( - self, - in_size, - out_size, - kernel_size, - stride, - padding, - activation, - bn, - init, - conv=None, - batch_norm=None, - bias=True, - preact=False, - name="", - instance_norm=False, - instance_norm_func=None - ): - super().__init__() - - bias = bias and (not bn) - conv_unit = conv( - in_size, - out_size, - kernel_size=kernel_size, - stride=stride, - padding=padding, - bias=bias - ) - init(conv_unit.weight) - if bias: - nn.init.constant_(conv_unit.bias, 0) - - if bn: - if not preact: - bn_unit = batch_norm(out_size) - else: - bn_unit = batch_norm(in_size) - if instance_norm: - if not preact: - in_unit = instance_norm_func(out_size, affine=False, track_running_stats=False) - else: - in_unit = instance_norm_func(in_size, affine=False, track_running_stats=False) - - if preact: - if bn: - self.add_module(name + 'bn', bn_unit) - - if activation is not None: - self.add_module(name + 'activation', activation) - - if not bn and instance_norm: - self.add_module(name + 'in', in_unit) - - self.add_module(name + 'conv', conv_unit) - - if not preact: - if bn: - self.add_module(name + 'bn', bn_unit) - - if activation is not None: - self.add_module(name + 'activation', activation) - - if not bn and instance_norm: - self.add_module(name + 'in', in_unit) - - -class _BNBase(nn.Sequential): - - def __init__(self, in_size, batch_norm=None, name=""): - super().__init__() - self.add_module(name + "bn", batch_norm(in_size)) - - nn.init.constant_(self[0].weight, 1.0) - nn.init.constant_(self[0].bias, 0) - - -class BatchNorm1d(_BNBase): - - def __init__(self, in_size: int, *, name: str = ""): - super().__init__(in_size, batch_norm=nn.BatchNorm1d, name=name) - - -class BatchNorm2d(_BNBase): - - def __init__(self, in_size: int, name: str = ""): - super().__init__(in_size, batch_norm=nn.BatchNorm2d, name=name) - - -class Conv1d(_ConvBase): - - def __init__( - self, - in_size: int, - out_size: int, - *, - kernel_size: int = 1, - stride: int = 1, - padding: int = 0, - activation=nn.ReLU(inplace=True), - bn: bool = False, - init=nn.init.kaiming_normal_, - bias: bool = True, - preact: bool = False, - name: str = "", - instance_norm=False - ): - super().__init__( - in_size, - out_size, - kernel_size, - stride, - padding, - activation, - bn, - init, - conv=nn.Conv1d, - batch_norm=BatchNorm1d, - bias=bias, - preact=preact, - name=name, - instance_norm=instance_norm, - instance_norm_func=nn.InstanceNorm1d - ) - - -class Conv2d(_ConvBase): - - def __init__( - self, - in_size: int, - out_size: int, - *, - kernel_size: Tuple[int, int] = (1, 1), - stride: Tuple[int, int] = (1, 1), - padding: Tuple[int, int] = (0, 0), - activation=nn.ReLU(inplace=True), - bn: bool = False, - init=nn.init.kaiming_normal_, - bias: bool = True, - preact: bool = False, - name: str = "", - instance_norm=False - ): - super().__init__( - in_size, - out_size, - kernel_size, - stride, - padding, - activation, - bn, - init, - conv=nn.Conv2d, - batch_norm=BatchNorm2d, - bias=bias, - preact=preact, - name=name, - instance_norm=instance_norm, - instance_norm_func=nn.InstanceNorm2d - ) - - -class FC(nn.Sequential): - - def __init__( - self, - in_size: int, - out_size: int, - *, - activation=nn.ReLU(inplace=True), - bn: bool = False, - init=None, - preact: bool = False, - name: str = "" - ): - super().__init__() - - fc = nn.Linear(in_size, out_size, bias=not bn) - if init is not None: - init(fc.weight) - if not bn: - nn.init.constant(fc.bias, 0) - - if preact: - if bn: - self.add_module(name + 'bn', BatchNorm1d(in_size)) - - if activation is not None: - self.add_module(name + 'activation', activation) - - self.add_module(name + 'fc', fc) - - if not preact: - if bn: - self.add_module(name + 'bn', BatchNorm1d(out_size)) - - if activation is not None: - self.add_module(name + 'activation', activation) - diff --git a/other_models/VRCNet/utils/pointnet2/setup.py b/other_models/VRCNet/utils/pointnet2/setup.py deleted file mode 100644 index 99e59e3..0000000 --- a/other_models/VRCNet/utils/pointnet2/setup.py +++ /dev/null @@ -1,23 +0,0 @@ -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup( - name='pointnet2', - ext_modules=[ - CUDAExtension('pointnet2_cuda', [ - 'src/pointnet2_api.cpp', - - 'src/ball_query.cpp', - 'src/ball_query_gpu.cu', - 'src/group_points.cpp', - 'src/group_points_gpu.cu', - 'src/interpolate.cpp', - 'src/interpolate_gpu.cu', - 'src/sampling.cpp', - 'src/sampling_gpu.cu', - ], - extra_compile_args={'cxx': ['-g'], - 'nvcc': ['-O2']}) - ], - cmdclass={'build_ext': BuildExtension} -) diff --git a/other_models/VRCNet/utils/pointnet2/src/ball_query.cpp b/other_models/VRCNet/utils/pointnet2/src/ball_query.cpp deleted file mode 100644 index 953a5d1..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/ball_query.cpp +++ /dev/null @@ -1,36 +0,0 @@ -#include -#include -#include -#include -#include -#include "ball_query_gpu.h" - -extern THCState *state; - -#ifndef AT_CHECK -#define AT_CHECK TORCH_CHECK -#endif - -/* // pytorch 1.7 -#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ") -#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x, " must be contiguous ") -#define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) -*/ - -// pytorch 1.9 -#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) - -int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, - at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) { - CHECK_INPUT(new_xyz_tensor); - CHECK_INPUT(xyz_tensor); - const float *new_xyz = new_xyz_tensor.data(); - const float *xyz = xyz_tensor.data(); - int *idx = idx_tensor.data(); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - ball_query_kernel_launcher_fast(b, n, m, radius, nsample, new_xyz, xyz, idx, stream); - return 1; -} diff --git a/other_models/VRCNet/utils/pointnet2/src/ball_query_gpu.cu b/other_models/VRCNet/utils/pointnet2/src/ball_query_gpu.cu deleted file mode 100644 index f8840aa..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/ball_query_gpu.cu +++ /dev/null @@ -1,67 +0,0 @@ -#include -#include -#include - -#include "ball_query_gpu.h" -#include "cuda_utils.h" - - -__global__ void ball_query_kernel_fast(int b, int n, int m, float radius, int nsample, - const float *__restrict__ new_xyz, const float *__restrict__ xyz, int *__restrict__ idx) { - // new_xyz: (B, M, 3) - // xyz: (B, N, 3) - // output: - // idx: (B, M, nsample) - int bs_idx = blockIdx.y; - int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (bs_idx >= b || pt_idx >= m) return; - - new_xyz += bs_idx * m * 3 + pt_idx * 3; - xyz += bs_idx * n * 3; - idx += bs_idx * m * nsample + pt_idx * nsample; - - float radius2 = radius * radius; - float new_x = new_xyz[0]; - float new_y = new_xyz[1]; - float new_z = new_xyz[2]; - - int cnt = 0; - for (int k = 0; k < n; ++k) { - float x = xyz[k * 3 + 0]; - float y = xyz[k * 3 + 1]; - float z = xyz[k * 3 + 2]; - float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); - if (d2 < radius2){ - if (cnt == 0){ - for (int l = 0; l < nsample; ++l) { - idx[l] = k; - } - } - idx[cnt] = k; - ++cnt; - if (cnt >= nsample) break; - } - } -} - - -void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample, \ - const float *new_xyz, const float *xyz, int *idx, cudaStream_t stream) { - // new_xyz: (B, M, 3) - // xyz: (B, N, 3) - // output: - // idx: (B, M, nsample) - - cudaError_t err; - - dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) - dim3 threads(THREADS_PER_BLOCK); - - ball_query_kernel_fast<<>>(b, n, m, radius, nsample, new_xyz, xyz, idx); - // cudaDeviceSynchronize(); // for using printf in kernel function - err = cudaGetLastError(); - if (cudaSuccess != err) { - fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } -} \ No newline at end of file diff --git a/other_models/VRCNet/utils/pointnet2/src/ball_query_gpu.h b/other_models/VRCNet/utils/pointnet2/src/ball_query_gpu.h deleted file mode 100644 index ffc831a..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/ball_query_gpu.h +++ /dev/null @@ -1,15 +0,0 @@ -#ifndef _BALL_QUERY_GPU_H -#define _BALL_QUERY_GPU_H - -#include -#include -#include -#include - -int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, - at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor); - -void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample, - const float *xyz, const float *new_xyz, int *idx, cudaStream_t stream); - -#endif diff --git a/other_models/VRCNet/utils/pointnet2/src/cuda_utils.h b/other_models/VRCNet/utils/pointnet2/src/cuda_utils.h deleted file mode 100644 index 7fe2796..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/cuda_utils.h +++ /dev/null @@ -1,15 +0,0 @@ -#ifndef _CUDA_UTILS_H -#define _CUDA_UTILS_H - -#include - -#define TOTAL_THREADS 1024 -#define THREADS_PER_BLOCK 256 -#define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) - -inline int opt_n_threads(int work_size) { - const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); - - return max(min(1 << pow_2, TOTAL_THREADS), 1); -} -#endif diff --git a/other_models/VRCNet/utils/pointnet2/src/group_points.cpp b/other_models/VRCNet/utils/pointnet2/src/group_points.cpp deleted file mode 100644 index b28830f..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/group_points.cpp +++ /dev/null @@ -1,36 +0,0 @@ -#include -#include -#include -#include -#include -#include "group_points_gpu.h" - -extern THCState *state; - - -int group_points_grad_wrapper_fast(int b, int c, int n, int npoints, int nsample, - at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) { - - float *grad_points = grad_points_tensor.data(); - const int *idx = idx_tensor.data(); - const float *grad_out = grad_out_tensor.data(); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - group_points_grad_kernel_launcher_fast(b, c, n, npoints, nsample, grad_out, idx, grad_points, stream); - return 1; -} - - -int group_points_wrapper_fast(int b, int c, int n, int npoints, int nsample, - at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor) { - - const float *points = points_tensor.data(); - const int *idx = idx_tensor.data(); - float *out = out_tensor.data(); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - group_points_kernel_launcher_fast(b, c, n, npoints, nsample, points, idx, out, stream); - return 1; -} diff --git a/other_models/VRCNet/utils/pointnet2/src/group_points_gpu.cu b/other_models/VRCNet/utils/pointnet2/src/group_points_gpu.cu deleted file mode 100644 index c015a81..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/group_points_gpu.cu +++ /dev/null @@ -1,86 +0,0 @@ -#include -#include - -#include "cuda_utils.h" -#include "group_points_gpu.h" - - -__global__ void group_points_grad_kernel_fast(int b, int c, int n, int npoints, int nsample, - const float *__restrict__ grad_out, const int *__restrict__ idx, float *__restrict__ grad_points) { - // grad_out: (B, C, npoints, nsample) - // idx: (B, npoints, nsample) - // output: - // grad_points: (B, C, N) - int bs_idx = blockIdx.z; - int c_idx = blockIdx.y; - int index = blockIdx.x * blockDim.x + threadIdx.x; - int pt_idx = index / nsample; - if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return; - - int sample_idx = index % nsample; - grad_out += bs_idx * c * npoints * nsample + c_idx * npoints * nsample + pt_idx * nsample + sample_idx; - idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx; - - atomicAdd(grad_points + bs_idx * c * n + c_idx * n + idx[0] , grad_out[0]); -} - -void group_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample, - const float *grad_out, const int *idx, float *grad_points, cudaStream_t stream) { - // grad_out: (B, C, npoints, nsample) - // idx: (B, npoints, nsample) - // output: - // grad_points: (B, C, N) - cudaError_t err; - dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) - dim3 threads(THREADS_PER_BLOCK); - - group_points_grad_kernel_fast<<>>(b, c, n, npoints, nsample, grad_out, idx, grad_points); - - err = cudaGetLastError(); - if (cudaSuccess != err) { - fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } -} - - -__global__ void group_points_kernel_fast(int b, int c, int n, int npoints, int nsample, - const float *__restrict__ points, const int *__restrict__ idx, float *__restrict__ out) { - // points: (B, C, N) - // idx: (B, npoints, nsample) - // output: - // out: (B, C, npoints, nsample) - int bs_idx = blockIdx.z; - int c_idx = blockIdx.y; - int index = blockIdx.x * blockDim.x + threadIdx.x; - int pt_idx = index / nsample; - if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return; - - int sample_idx = index % nsample; - - idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx; - int in_idx = bs_idx * c * n + c_idx * n + idx[0]; - int out_idx = bs_idx * c * npoints * nsample + c_idx * npoints * nsample + pt_idx * nsample + sample_idx; - - out[out_idx] = points[in_idx]; -} - - -void group_points_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample, - const float *points, const int *idx, float *out, cudaStream_t stream) { - // points: (B, C, N) - // idx: (B, npoints, nsample) - // output: - // out: (B, C, npoints, nsample) - cudaError_t err; - dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) - dim3 threads(THREADS_PER_BLOCK); - - group_points_kernel_fast<<>>(b, c, n, npoints, nsample, points, idx, out); - // cudaDeviceSynchronize(); // for using printf in kernel function - err = cudaGetLastError(); - if (cudaSuccess != err) { - fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } -} diff --git a/other_models/VRCNet/utils/pointnet2/src/group_points_gpu.h b/other_models/VRCNet/utils/pointnet2/src/group_points_gpu.h deleted file mode 100644 index 76c73ca..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/group_points_gpu.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef _GROUP_POINTS_GPU_H -#define _GROUP_POINTS_GPU_H - -#include -#include -#include -#include - - -int group_points_wrapper_fast(int b, int c, int n, int npoints, int nsample, - at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor); - -void group_points_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample, - const float *points, const int *idx, float *out, cudaStream_t stream); - -int group_points_grad_wrapper_fast(int b, int c, int n, int npoints, int nsample, - at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor); - -void group_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints, int nsample, - const float *grad_out, const int *idx, float *grad_points, cudaStream_t stream); - -#endif diff --git a/other_models/VRCNet/utils/pointnet2/src/interpolate.cpp b/other_models/VRCNet/utils/pointnet2/src/interpolate.cpp deleted file mode 100644 index 8c1b8df..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/interpolate.cpp +++ /dev/null @@ -1,54 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include -#include "interpolate_gpu.h" - -extern THCState *state; - - -void three_nn_wrapper_fast(int b, int n, int m, at::Tensor unknown_tensor, - at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor) { - const float *unknown = unknown_tensor.data(); - const float *known = known_tensor.data(); - float *dist2 = dist2_tensor.data(); - int *idx = idx_tensor.data(); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - three_nn_kernel_launcher_fast(b, n, m, unknown, known, dist2, idx, stream); -} - - -void three_interpolate_wrapper_fast(int b, int c, int m, int n, - at::Tensor points_tensor, - at::Tensor idx_tensor, - at::Tensor weight_tensor, - at::Tensor out_tensor) { - - const float *points = points_tensor.data(); - const float *weight = weight_tensor.data(); - float *out = out_tensor.data(); - const int *idx = idx_tensor.data(); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - three_interpolate_kernel_launcher_fast(b, c, m, n, points, idx, weight, out, stream); -} - -void three_interpolate_grad_wrapper_fast(int b, int c, int n, int m, - at::Tensor grad_out_tensor, - at::Tensor idx_tensor, - at::Tensor weight_tensor, - at::Tensor grad_points_tensor) { - - const float *grad_out = grad_out_tensor.data(); - const float *weight = weight_tensor.data(); - float *grad_points = grad_points_tensor.data(); - const int *idx = idx_tensor.data(); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - three_interpolate_grad_kernel_launcher_fast(b, c, n, m, grad_out, idx, weight, grad_points, stream); -} diff --git a/other_models/VRCNet/utils/pointnet2/src/interpolate_gpu.cu b/other_models/VRCNet/utils/pointnet2/src/interpolate_gpu.cu deleted file mode 100644 index a123dd8..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/interpolate_gpu.cu +++ /dev/null @@ -1,161 +0,0 @@ -#include -#include -#include - -#include "cuda_utils.h" -#include "interpolate_gpu.h" - - -__global__ void three_nn_kernel_fast(int b, int n, int m, const float *__restrict__ unknown, - const float *__restrict__ known, float *__restrict__ dist2, int *__restrict__ idx) { - // unknown: (B, N, 3) - // known: (B, M, 3) - // output: - // dist2: (B, N, 3) - // idx: (B, N, 3) - - int bs_idx = blockIdx.y; - int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (bs_idx >= b || pt_idx >= n) return; - - unknown += bs_idx * n * 3 + pt_idx * 3; - known += bs_idx * m * 3; - dist2 += bs_idx * n * 3 + pt_idx * 3; - idx += bs_idx * n * 3 + pt_idx * 3; - - float ux = unknown[0]; - float uy = unknown[1]; - float uz = unknown[2]; - - double best1 = 1e40, best2 = 1e40, best3 = 1e40; - int besti1 = 0, besti2 = 0, besti3 = 0; - for (int k = 0; k < m; ++k) { - float x = known[k * 3 + 0]; - float y = known[k * 3 + 1]; - float z = known[k * 3 + 2]; - float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); - if (d < best1) { - best3 = best2; besti3 = besti2; - best2 = best1; besti2 = besti1; - best1 = d; besti1 = k; - } - else if (d < best2) { - best3 = best2; besti3 = besti2; - best2 = d; besti2 = k; - } - else if (d < best3) { - best3 = d; besti3 = k; - } - } - dist2[0] = best1; dist2[1] = best2; dist2[2] = best3; - idx[0] = besti1; idx[1] = besti2; idx[2] = besti3; -} - - -void three_nn_kernel_launcher_fast(int b, int n, int m, const float *unknown, - const float *known, float *dist2, int *idx, cudaStream_t stream) { - // unknown: (B, N, 3) - // known: (B, M, 3) - // output: - // dist2: (B, N, 3) - // idx: (B, N, 3) - - cudaError_t err; - dim3 blocks(DIVUP(n, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) - dim3 threads(THREADS_PER_BLOCK); - - three_nn_kernel_fast<<>>(b, n, m, unknown, known, dist2, idx); - - err = cudaGetLastError(); - if (cudaSuccess != err) { - fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } -} - - -__global__ void three_interpolate_kernel_fast(int b, int c, int m, int n, const float *__restrict__ points, - const int *__restrict__ idx, const float *__restrict__ weight, float *__restrict__ out) { - // points: (B, C, M) - // idx: (B, N, 3) - // weight: (B, N, 3) - // output: - // out: (B, C, N) - - int bs_idx = blockIdx.z; - int c_idx = blockIdx.y; - int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (bs_idx >= b || c_idx >= c || pt_idx >= n) return; - - weight += bs_idx * n * 3 + pt_idx * 3; - points += bs_idx * c * m + c_idx * m; - idx += bs_idx * n * 3 + pt_idx * 3; - out += bs_idx * c * n + c_idx * n; - - out[pt_idx] = weight[0] * points[idx[0]] + weight[1] * points[idx[1]] + weight[2] * points[idx[2]]; -} - -void three_interpolate_kernel_launcher_fast(int b, int c, int m, int n, - const float *points, const int *idx, const float *weight, float *out, cudaStream_t stream) { - // points: (B, C, M) - // idx: (B, N, 3) - // weight: (B, N, 3) - // output: - // out: (B, C, N) - - cudaError_t err; - dim3 blocks(DIVUP(n, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) - dim3 threads(THREADS_PER_BLOCK); - three_interpolate_kernel_fast<<>>(b, c, m, n, points, idx, weight, out); - - err = cudaGetLastError(); - if (cudaSuccess != err) { - fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } -} - - -__global__ void three_interpolate_grad_kernel_fast(int b, int c, int n, int m, const float *__restrict__ grad_out, - const int *__restrict__ idx, const float *__restrict__ weight, float *__restrict__ grad_points) { - // grad_out: (B, C, N) - // weight: (B, N, 3) - // output: - // grad_points: (B, C, M) - - int bs_idx = blockIdx.z; - int c_idx = blockIdx.y; - int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (bs_idx >= b || c_idx >= c || pt_idx >= n) return; - - grad_out += bs_idx * c * n + c_idx * n + pt_idx; - weight += bs_idx * n * 3 + pt_idx * 3; - grad_points += bs_idx * c * m + c_idx * m; - idx += bs_idx * n * 3 + pt_idx * 3; - - - atomicAdd(grad_points + idx[0], grad_out[0] * weight[0]); - atomicAdd(grad_points + idx[1], grad_out[0] * weight[1]); - atomicAdd(grad_points + idx[2], grad_out[0] * weight[2]); -} - -void three_interpolate_grad_kernel_launcher_fast(int b, int c, int n, int m, const float *grad_out, - const int *idx, const float *weight, float *grad_points, cudaStream_t stream) { - // grad_out: (B, C, N) - // weight: (B, N, 3) - // output: - // grad_points: (B, C, M) - - cudaError_t err; - dim3 blocks(DIVUP(n, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) - dim3 threads(THREADS_PER_BLOCK); - three_interpolate_grad_kernel_fast<<>>(b, c, n, m, grad_out, idx, weight, grad_points); - - err = cudaGetLastError(); - if (cudaSuccess != err) { - fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } -} \ No newline at end of file diff --git a/other_models/VRCNet/utils/pointnet2/src/interpolate_gpu.h b/other_models/VRCNet/utils/pointnet2/src/interpolate_gpu.h deleted file mode 100644 index f177108..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/interpolate_gpu.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef _INTERPOLATE_GPU_H -#define _INTERPOLATE_GPU_H - -#include -#include -#include -#include - - -void three_nn_wrapper_fast(int b, int n, int m, at::Tensor unknown_tensor, - at::Tensor known_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor); - -void three_nn_kernel_launcher_fast(int b, int n, int m, const float *unknown, - const float *known, float *dist2, int *idx, cudaStream_t stream); - - -void three_interpolate_wrapper_fast(int b, int c, int m, int n, at::Tensor points_tensor, - at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor out_tensor); - -void three_interpolate_kernel_launcher_fast(int b, int c, int m, int n, - const float *points, const int *idx, const float *weight, float *out, cudaStream_t stream); - - -void three_interpolate_grad_wrapper_fast(int b, int c, int n, int m, at::Tensor grad_out_tensor, - at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_points_tensor); - -void three_interpolate_grad_kernel_launcher_fast(int b, int c, int n, int m, const float *grad_out, - const int *idx, const float *weight, float *grad_points, cudaStream_t stream); - -#endif diff --git a/other_models/VRCNet/utils/pointnet2/src/pointnet2_api.cpp b/other_models/VRCNet/utils/pointnet2/src/pointnet2_api.cpp deleted file mode 100644 index d91f0f2..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/pointnet2_api.cpp +++ /dev/null @@ -1,24 +0,0 @@ -#include -#include - -#include "ball_query_gpu.h" -#include "group_points_gpu.h" -#include "sampling_gpu.h" -#include "interpolate_gpu.h" - - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("ball_query_wrapper", &ball_query_wrapper_fast, "ball_query_wrapper_fast"); - - m.def("group_points_wrapper", &group_points_wrapper_fast, "group_points_wrapper_fast"); - m.def("group_points_grad_wrapper", &group_points_grad_wrapper_fast, "group_points_grad_wrapper_fast"); - - m.def("gather_points_wrapper", &gather_points_wrapper_fast, "gather_points_wrapper_fast"); - m.def("gather_points_grad_wrapper", &gather_points_grad_wrapper_fast, "gather_points_grad_wrapper_fast"); - - m.def("furthest_point_sampling_wrapper", &furthest_point_sampling_wrapper, "furthest_point_sampling_wrapper"); - - m.def("three_nn_wrapper", &three_nn_wrapper_fast, "three_nn_wrapper_fast"); - m.def("three_interpolate_wrapper", &three_interpolate_wrapper_fast, "three_interpolate_wrapper_fast"); - m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper_fast, "three_interpolate_grad_wrapper_fast"); -} diff --git a/other_models/VRCNet/utils/pointnet2/src/sampling.cpp b/other_models/VRCNet/utils/pointnet2/src/sampling.cpp deleted file mode 100644 index 4f0d62c..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/sampling.cpp +++ /dev/null @@ -1,46 +0,0 @@ -#include -#include -#include -#include - -#include "sampling_gpu.h" - -extern THCState *state; - - -int gather_points_wrapper_fast(int b, int c, int n, int npoints, - at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor){ - const float *points = points_tensor.data(); - const int *idx = idx_tensor.data(); - float *out = out_tensor.data(); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - gather_points_kernel_launcher_fast(b, c, n, npoints, points, idx, out, stream); - return 1; -} - - -int gather_points_grad_wrapper_fast(int b, int c, int n, int npoints, - at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor) { - - const float *grad_out = grad_out_tensor.data(); - const int *idx = idx_tensor.data(); - float *grad_points = grad_points_tensor.data(); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - gather_points_grad_kernel_launcher_fast(b, c, n, npoints, grad_out, idx, grad_points, stream); - return 1; -} - - -int furthest_point_sampling_wrapper(int b, int n, int m, - at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) { - - const float *points = points_tensor.data(); - float *temp = temp_tensor.data(); - int *idx = idx_tensor.data(); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - furthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx, stream); - return 1; -} diff --git a/other_models/VRCNet/utils/pointnet2/src/sampling_gpu.cu b/other_models/VRCNet/utils/pointnet2/src/sampling_gpu.cu deleted file mode 100644 index 9e49a60..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/sampling_gpu.cu +++ /dev/null @@ -1,253 +0,0 @@ -#include -#include - -#include "cuda_utils.h" -#include "sampling_gpu.h" - - -__global__ void gather_points_kernel_fast(int b, int c, int n, int m, - const float *__restrict__ points, const int *__restrict__ idx, float *__restrict__ out) { - // points: (B, C, N) - // idx: (B, M) - // output: - // out: (B, C, M) - - int bs_idx = blockIdx.z; - int c_idx = blockIdx.y; - int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (bs_idx >= b || c_idx >= c || pt_idx >= m) return; - - out += bs_idx * c * m + c_idx * m + pt_idx; - idx += bs_idx * m + pt_idx; - points += bs_idx * c * n + c_idx * n; - out[0] = points[idx[0]]; -} - -void gather_points_kernel_launcher_fast(int b, int c, int n, int npoints, - const float *points, const int *idx, float *out, cudaStream_t stream) { - // points: (B, C, N) - // idx: (B, npoints) - // output: - // out: (B, C, npoints) - - cudaError_t err; - dim3 blocks(DIVUP(npoints, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) - dim3 threads(THREADS_PER_BLOCK); - - gather_points_kernel_fast<<>>(b, c, n, npoints, points, idx, out); - - err = cudaGetLastError(); - if (cudaSuccess != err) { - fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } -} - -__global__ void gather_points_grad_kernel_fast(int b, int c, int n, int m, const float *__restrict__ grad_out, - const int *__restrict__ idx, float *__restrict__ grad_points) { - // grad_out: (B, C, M) - // idx: (B, M) - // output: - // grad_points: (B, C, N) - - int bs_idx = blockIdx.z; - int c_idx = blockIdx.y; - int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (bs_idx >= b || c_idx >= c || pt_idx >= m) return; - - grad_out += bs_idx * c * m + c_idx * m + pt_idx; - idx += bs_idx * m + pt_idx; - grad_points += bs_idx * c * n + c_idx * n; - - atomicAdd(grad_points + idx[0], grad_out[0]); -} - -void gather_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints, - const float *grad_out, const int *idx, float *grad_points, cudaStream_t stream) { - // grad_out: (B, C, npoints) - // idx: (B, npoints) - // output: - // grad_points: (B, C, N) - - cudaError_t err; - dim3 blocks(DIVUP(npoints, THREADS_PER_BLOCK), c, b); // blockIdx.x(col), blockIdx.y(row) - dim3 threads(THREADS_PER_BLOCK); - - gather_points_grad_kernel_fast<<>>(b, c, n, npoints, grad_out, idx, grad_points); - - err = cudaGetLastError(); - if (cudaSuccess != err) { - fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } -} - - -__device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i, int idx1, int idx2){ - const float v1 = dists[idx1], v2 = dists[idx2]; - const int i1 = dists_i[idx1], i2 = dists_i[idx2]; - dists[idx1] = max(v1, v2); - dists_i[idx1] = v2 > v1 ? i2 : i1; -} - -template -__global__ void furthest_point_sampling_kernel(int b, int n, int m, - const float *__restrict__ dataset, float *__restrict__ temp, int *__restrict__ idxs) { - // dataset: (B, N, 3) - // tmp: (B, N) - // output: - // idx: (B, M) - - if (m <= 0) return; - __shared__ float dists[block_size]; - __shared__ int dists_i[block_size]; - - int batch_index = blockIdx.x; - dataset += batch_index * n * 3; - temp += batch_index * n; - idxs += batch_index * m; - - int tid = threadIdx.x; - const int stride = block_size; - - int old = 0; - if (threadIdx.x == 0) - idxs[0] = old; - - __syncthreads(); - for (int j = 1; j < m; j++) { - int besti = 0; - float best = -1; - float x1 = dataset[old * 3 + 0]; - float y1 = dataset[old * 3 + 1]; - float z1 = dataset[old * 3 + 2]; - for (int k = tid; k < n; k += stride) { - float x2, y2, z2; - x2 = dataset[k * 3 + 0]; - y2 = dataset[k * 3 + 1]; - z2 = dataset[k * 3 + 2]; - // float mag = (x2 * x2) + (y2 * y2) + (z2 * z2); - // if (mag <= 1e-3) - // continue; - - float d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); - float d2 = min(d, temp[k]); - temp[k] = d2; - besti = d2 > best ? k : besti; - best = d2 > best ? d2 : best; - } - dists[tid] = best; - dists_i[tid] = besti; - __syncthreads(); - - if (block_size >= 1024) { - if (tid < 512) { - __update(dists, dists_i, tid, tid + 512); - } - __syncthreads(); - } - - if (block_size >= 512) { - if (tid < 256) { - __update(dists, dists_i, tid, tid + 256); - } - __syncthreads(); - } - if (block_size >= 256) { - if (tid < 128) { - __update(dists, dists_i, tid, tid + 128); - } - __syncthreads(); - } - if (block_size >= 128) { - if (tid < 64) { - __update(dists, dists_i, tid, tid + 64); - } - __syncthreads(); - } - if (block_size >= 64) { - if (tid < 32) { - __update(dists, dists_i, tid, tid + 32); - } - __syncthreads(); - } - if (block_size >= 32) { - if (tid < 16) { - __update(dists, dists_i, tid, tid + 16); - } - __syncthreads(); - } - if (block_size >= 16) { - if (tid < 8) { - __update(dists, dists_i, tid, tid + 8); - } - __syncthreads(); - } - if (block_size >= 8) { - if (tid < 4) { - __update(dists, dists_i, tid, tid + 4); - } - __syncthreads(); - } - if (block_size >= 4) { - if (tid < 2) { - __update(dists, dists_i, tid, tid + 2); - } - __syncthreads(); - } - if (block_size >= 2) { - if (tid < 1) { - __update(dists, dists_i, tid, tid + 1); - } - __syncthreads(); - } - - old = dists_i[0]; - if (tid == 0) - idxs[j] = old; - } -} - -void furthest_point_sampling_kernel_launcher(int b, int n, int m, - const float *dataset, float *temp, int *idxs, cudaStream_t stream) { - // dataset: (B, N, 3) - // tmp: (B, N) - // output: - // idx: (B, M) - - cudaError_t err; - unsigned int n_threads = opt_n_threads(n); - - switch (n_threads) { - case 1024: - furthest_point_sampling_kernel<1024><<>>(b, n, m, dataset, temp, idxs); break; - case 512: - furthest_point_sampling_kernel<512><<>>(b, n, m, dataset, temp, idxs); break; - case 256: - furthest_point_sampling_kernel<256><<>>(b, n, m, dataset, temp, idxs); break; - case 128: - furthest_point_sampling_kernel<128><<>>(b, n, m, dataset, temp, idxs); break; - case 64: - furthest_point_sampling_kernel<64><<>>(b, n, m, dataset, temp, idxs); break; - case 32: - furthest_point_sampling_kernel<32><<>>(b, n, m, dataset, temp, idxs); break; - case 16: - furthest_point_sampling_kernel<16><<>>(b, n, m, dataset, temp, idxs); break; - case 8: - furthest_point_sampling_kernel<8><<>>(b, n, m, dataset, temp, idxs); break; - case 4: - furthest_point_sampling_kernel<4><<>>(b, n, m, dataset, temp, idxs); break; - case 2: - furthest_point_sampling_kernel<2><<>>(b, n, m, dataset, temp, idxs); break; - case 1: - furthest_point_sampling_kernel<1><<>>(b, n, m, dataset, temp, idxs); break; - default: - furthest_point_sampling_kernel<512><<>>(b, n, m, dataset, temp, idxs); - } - - err = cudaGetLastError(); - if (cudaSuccess != err) { - fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } -} diff --git a/other_models/VRCNet/utils/pointnet2/src/sampling_gpu.h b/other_models/VRCNet/utils/pointnet2/src/sampling_gpu.h deleted file mode 100644 index 6200c59..0000000 --- a/other_models/VRCNet/utils/pointnet2/src/sampling_gpu.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef _SAMPLING_GPU_H -#define _SAMPLING_GPU_H - -#include -#include -#include - - -int gather_points_wrapper_fast(int b, int c, int n, int npoints, - at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor); - -void gather_points_kernel_launcher_fast(int b, int c, int n, int npoints, - const float *points, const int *idx, float *out, cudaStream_t stream); - - -int gather_points_grad_wrapper_fast(int b, int c, int n, int npoints, - at::Tensor grad_out_tensor, at::Tensor idx_tensor, at::Tensor grad_points_tensor); - -void gather_points_grad_kernel_launcher_fast(int b, int c, int n, int npoints, - const float *grad_out, const int *idx, float *grad_points, cudaStream_t stream); - - -int furthest_point_sampling_wrapper(int b, int n, int m, - at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor); - -void furthest_point_sampling_kernel_launcher(int b, int n, int m, - const float *dataset, float *temp, int *idxs, cudaStream_t stream); - -#endif diff --git a/other_models/VRCNet/utils/train_utils.py b/other_models/VRCNet/utils/train_utils.py deleted file mode 100644 index f66530d..0000000 --- a/other_models/VRCNet/utils/train_utils.py +++ /dev/null @@ -1,60 +0,0 @@ -import torch - -class AverageValueMeter(object): - def __init__(self): - self.reset() - - def reset(self): - self.val = 0 - self.avg = 0 - self.sum = 0 - self.count = 0.0 - - def update(self, val, n=1): - self.val = val - self.sum += val * n - self.count += n - self.avg = self.sum / self.count - - -def set_requires_grad(nets, requires_grad=False): - if not isinstance(nets, list): - nets = [nets] - for net in nets: - if net is not None: - for param in net.parameters(): - param.requires_grad = requires_grad - - -def save_model(path, net, net_d=None): - if net_d is not None: - torch.save({'net_state_dict': net.module.state_dict(), - 'D_state_dict': net_d.module.state_dict()}, path) - else: - torch.save({'net_state_dict': net.module.state_dict()}, path) - - -def generator_step(net_d, out2, net_loss, optimizer): - set_requires_grad(net_d, False) - d_fake = net_d(out2[:, 0:2048, :]) - errG_loss_batch = torch.mean((d_fake - 1) ** 2) - total_gen_loss_batch = errG_loss_batch + net_loss * 200 - total_gen_loss_batch.backward(torch.ones(torch.cuda.device_count()).cuda(), retain_graph=True, ) - optimizer.step() - return d_fake - - -def discriminator_step(net_d, gt, d_fake, optimizer_d): - set_requires_grad(net_d, True) - d_real = net_d(gt[:, 0:2048, :]) - d_loss_fake = torch.mean(d_fake ** 2) - d_loss_real = torch.mean((d_real - 1) ** 2) - errD_loss_batch = 0.5 * (d_loss_real + d_loss_fake) - total_dis_loss_batch = errD_loss_batch - total_dis_loss_batch.backward(torch.ones(torch.cuda.device_count()).cuda()) - optimizer_d.step() - - - - - diff --git a/other_models/VRCNet/utils/vis_utils.py b/other_models/VRCNet/utils/vis_utils.py deleted file mode 100644 index bcd9084..0000000 --- a/other_models/VRCNet/utils/vis_utils.py +++ /dev/null @@ -1,60 +0,0 @@ -import open3d as o3d -import numpy as np -from mpl_toolkits.mplot3d import Axes3D -import matplotlib.pyplot as plt - - -def get_pts(pcd): - points = np.asarray(pcd.points) - X = [] - Y = [] - Z = [] - for pt in range(points.shape[0]): - X.append(points[pt][0]) - Y.append(points[pt][1]) - Z.append(points[pt][2]) - return np.asarray(X), np.asarray(Y), np.asarray(Z) - - -def set_axes_equal(ax): - x_limits = ax.get_xlim3d() - y_limits = ax.get_ylim3d() - z_limits = ax.get_zlim3d() - x_range = abs(x_limits[1] - x_limits[0]) - x_middle = np.mean(x_limits) - y_range = abs(y_limits[1] - y_limits[0]) - y_middle = np.mean(y_limits) - z_range = abs(z_limits[1] - z_limits[0]) - z_middle = np.mean(z_limits) - plot_radius = 0.5*max([x_range, y_range, z_range]) - ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius]) - ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius]) - ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius]) - - -def plot_single_pcd(points, save_path): - fig = plt.figure() - ax = fig.add_subplot(111, projection='3d') - ax.set_aspect('equal') - pcd = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(points)) - rotation_matrix = np.asarray([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]) - pcd = pcd.transform(rotation_matrix) - X, Y, Z = get_pts(pcd) - t = Z - ax.scatter(X, Y, Z, c=t, cmap='jet', marker='o', s=0.5, linewidths=0) - ax.grid(False) - ax.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0)) - ax.w_yaxis.set_pane_color((1.0, 1.0, 1.0, 1.0)) - ax.w_zaxis.set_pane_color((1.0, 1.0, 1.0, 1.0)) - set_axes_equal(ax) - plt.axis('off') - plt.savefig(save_path, format='png', dpi=600) - plt.close() - - - - - - - - diff --git a/other_models/VRCNet/vrc.py b/other_models/VRCNet/vrc.py deleted file mode 100644 index 309dc55..0000000 --- a/other_models/VRCNet/vrc.py +++ /dev/null @@ -1,515 +0,0 @@ -import numpy as np -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -import torch.nn.functional as F -from utils.model_utils import * -from models.pcn import PCN_encoder - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "utils/Pointnet2.PyTorch/pointnet2")) -import pointnet2_utils as pn2 - - -class SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=16): - super(SA_module, self).__init__() - self.share_planes = share_planes - self.k = k - self.conv1 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv2 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv3 = nn.Conv2d(in_planes, mid_planes, kernel_size=1) - - self.conv_w = nn.Sequential(nn.ReLU(inplace=False), - nn.Conv2d(rel_planes * (k + 1), mid_planes // share_planes, kernel_size=1, - bias=False), - nn.ReLU(inplace=False), - nn.Conv2d(mid_planes // share_planes, k * mid_planes // share_planes, - kernel_size=1)) - self.activation_fn = nn.ReLU(inplace=False) - - self.conv_out = nn.Conv2d(mid_planes, out_planes, kernel_size=1) - - def forward(self, input): - x, idx = input - batch_size, _, _, num_points = x.size() - identity = x # B C 1 N - x = self.activation_fn(x) - xn = get_edge_features(x, idx) # B C K N - x1, x2, x3 = self.conv1(x), self.conv2(xn), self.conv3(xn) - - x2 = x2.view(batch_size, -1, 1, num_points).contiguous() # B kC 1 N - w = self.conv_w(torch.cat([x1, x2], 1)).view(batch_size, -1, self.k, num_points) - w = w.repeat(1, self.share_planes, 1, 1) - out = w * x3 - out = torch.sum(out, dim=2, keepdim=True) - - out = self.activation_fn(out) - out = self.conv_out(out) # B C 1 N - out += identity - return [out, idx] - - -class Folding(nn.Module): - def __init__(self, input_size, output_size, step_ratio, global_feature_size=1024, num_models=1): - super(Folding, self).__init__() - self.input_size = input_size - self.output_size = output_size - self.step_ratio = step_ratio - self.num_models = num_models - - self.conv = nn.Conv1d(input_size + global_feature_size + 2, output_size, 1, bias=True) - - sqrted = int(math.sqrt(step_ratio)) + 1 - for i in range(1, sqrted + 1).__reversed__(): - if (step_ratio % i) == 0: - num_x = i - num_y = step_ratio // i - break - - grid_x = torch.linspace(-0.2, 0.2, steps=num_x) - grid_y = torch.linspace(-0.2, 0.2, steps=num_y) - - x, y = torch.meshgrid(grid_x, grid_y) # x, y shape: (2, 1) - self.grid = torch.stack([x, y], dim=-1).view(-1, 2) # (2, 2) - - def forward(self, point_feat, global_feat): - batch_size, num_features, num_points = point_feat.size() - point_feat = point_feat.transpose(1, 2).contiguous().unsqueeze(2).repeat(1, 1, self.step_ratio, 1).view( - batch_size, - -1, num_features).transpose(1, 2).contiguous() - global_feat = global_feat.unsqueeze(2).repeat(1, 1, num_points * self.step_ratio).repeat(self.num_models, 1, 1) - grid_feat = self.grid.unsqueeze(0).repeat(batch_size, num_points, 1).transpose(1, 2).contiguous().cuda() - features = torch.cat([global_feat, point_feat, grid_feat], axis=1) - features = F.relu(self.conv(features)) - return features - - -class Linear_ResBlock(nn.Module): - def __init__(self, input_size=1024, output_size=256): - super(Linear_ResBlock, self).__init__() - self.conv1 = nn.Linear(input_size, input_size) - self.conv2 = nn.Linear(input_size, output_size) - self.conv_res = nn.Linear(input_size, output_size) - - self.af = nn.ReLU(inplace=True) - - def forward(self, feature): - return self.conv2(self.af(self.conv1(self.af(feature)))) + self.conv_res(feature) - - -class SK_SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=[10, 20], r=2, L=32): - super(SK_SA_module, self).__init__() - - self.num_kernels = len(k) - d = max(int(out_planes / r), L) - - self.sams = nn.ModuleList([]) - - for i in range(len(k)): - self.sams.append(SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k[i])) - - self.fc = nn.Linear(out_planes, d) - self.fcs = nn.ModuleList([]) - - for i in range(len(k)): - self.fcs.append(nn.Linear(d, out_planes)) - - self.softmax = nn.Softmax(dim=1) - self.af = nn.ReLU(inplace=False) - - def forward(self, input): - x, idxs = input - assert (self.num_kernels == len(idxs)) - for i, sam in enumerate(self.sams): - fea, _ = sam([x, idxs[i]]) - fea = self.af(fea) - fea = fea.unsqueeze(dim=1) - if i == 0: - feas = fea - else: - feas = torch.cat([feas, fea], dim=1) - - fea_U = torch.sum(feas, dim=1) - fea_s = fea_U.mean(-1).mean(-1) - fea_z = self.fc(fea_s) - - for i, fc in enumerate(self.fcs): - vector = fc(fea_z).unsqueeze_(dim=1) - if i == 0: - attention_vectors = vector - else: - attention_vectors = torch.cat([attention_vectors, vector], dim=1) - - attention_vectors = self.softmax(attention_vectors) - attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1) - fea_v = (feas * attention_vectors).sum(dim=1) - return [fea_v, idxs] - - -class SKN_Res_unit(nn.Module): - def __init__(self, input_size, output_size, k=[10, 20], layers=1): - super(SKN_Res_unit, self).__init__() - self.conv1 = nn.Conv2d(input_size, output_size, 1, bias=False) - self.sam = self._make_layer(output_size, output_size // 16, output_size // 4, output_size, int(layers), 8, k=k) - self.conv2 = nn.Conv2d(output_size, output_size, 1, bias=False) - self.conv_res = nn.Conv2d(input_size, output_size, 1, bias=False) - self.af = nn.ReLU(inplace=False) - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def forward(self, feat, idx): - x, _ = self.sam([self.conv1(feat), idx]) - x = self.conv2(self.af(x)) - return x + self.conv_res(feat) - - -class SA_SKN_Res_encoder(nn.Module): - def __init__(self, input_size=3, k=[10, 20], pk=16, output_size=64, layers=[2, 2, 2, 2], - pts_num=[3072, 1536, 768, 384]): - super(SA_SKN_Res_encoder, self).__init__() - self.init_channel = 64 - - c1 = self.init_channel - self.sam_res1 = SKN_Res_unit(input_size, c1, k, int(layers[0])) - - c2 = c1 * 2 - self.sam_res2 = SKN_Res_unit(c2, c2, k, int(layers[1])) - - c3 = c2 * 2 - self.sam_res3 = SKN_Res_unit(c3, c3, k, int(layers[2])) - - c4 = c3 * 2 - self.sam_res4 = SKN_Res_unit(c4, c4, k, int(layers[3])) - - self.conv5 = nn.Conv2d(c4, 1024, 1) - - self.fc1 = nn.Linear(1024, 512) - self.fc2 = nn.Linear(512, 1024) - - self.conv6 = nn.Conv2d(c4 + 1024, c4, 1) - self.conv7 = nn.Conv2d(c3 + c4, c3, 1) - self.conv8 = nn.Conv2d(c2 + c3, c2, 1) - self.conv9 = nn.Conv2d(c1 + c2, c1, 1) - - self.conv_out = nn.Conv2d(c1, output_size, 1) - self.dropout = nn.Dropout() - self.af = nn.ReLU(inplace=False) - self.k = k - self.pk = pk - self.rate = 2 - - self.pts_num = pts_num - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def _edge_pooling(self, features, points, rate=2, k=16, sample_num=None): - features = features.squeeze(2) - - if sample_num is None: - input_points_num = int(features.size()[2]) - sample_num = input_points_num // rate - - ds_features, p_idx, pn_idx, ds_points = edge_preserve_sampling(features, points, sample_num, k) - ds_features = ds_features.unsqueeze(2) - return ds_features, p_idx, pn_idx, ds_points - - def _edge_unpooling(self, features, src_pts, tgt_pts): - features = features.squeeze(2) - idx, weight = three_nn_upsampling(tgt_pts, src_pts) - features = pn2.three_interpolate(features, idx, weight) - features = features.unsqueeze(2) - return features - - def forward(self, features): - batch_size, _, num_points = features.size() - pt1 = features[:, 0:3, :] - - idx1 = [] - for i in range(len(self.k)): - idx = knn(pt1, self.k[i]) - idx1.append(idx) - - pt1 = pt1.transpose(1, 2).contiguous() - - x = features.unsqueeze(2) - x = self.sam_res1(x, idx1) - x1 = self.af(x) - - x, _, _, pt2 = self._edge_pooling(x1, pt1, self.rate, self.pk, self.pts_num[1]) - idx2 = [] - for i in range(len(self.k)): - idx = knn(pt2.transpose(1, 2).contiguous(), self.k[i]) - idx2.append(idx) - - x = self.sam_res2(x, idx2) - x2 = self.af(x) - - x, _, _, pt3 = self._edge_pooling(x2, pt2, self.rate, self.pk, self.pts_num[2]) - idx3 = [] - for i in range(len(self.k)): - idx = knn(pt3.transpose(1, 2).contiguous(), self.k[i]) - idx3.append(idx) - - x = self.sam_res3(x, idx3) - x3 = self.af(x) - - x, _, _, pt4 = self._edge_pooling(x3, pt3, self.rate, self.pk, self.pts_num[3]) - idx4 = [] - for i in range(len(self.k)): - idx = knn(pt4.transpose(1, 2).contiguous(), self.k[i]) - idx4.append(idx) - - x = self.sam_res4(x, idx4) - x4 = self.af(x) - x = self.conv5(x4) - x, _ = torch.max(x, -1) - x = x.view(batch_size, -1) - x = self.dropout(self.af(self.fc2(self.dropout(self.af(self.fc1(x)))))) - - x = x.unsqueeze(2).repeat(1, 1, self.pts_num[3]).unsqueeze(2) - x = self.af(self.conv6(torch.cat([x, x4], 1))) - x = self._edge_unpooling(x, pt4, pt3) - x = self.af(self.conv7(torch.cat([x, x3], 1))) - x = self._edge_unpooling(x, pt3, pt2) - x = self.af(self.conv8(torch.cat([x, x2], 1))) - x = self._edge_unpooling(x, pt2, pt1) - x = self.af(self.conv9(torch.cat([x, x1], 1))) - x = self.conv_out(x) - x = x.squeeze(2) - return x - - -class MSAP_SKN_decoder(nn.Module): - def __init__(self, num_coarse_raw, num_fps, num_coarse, num_fine, layers=[2, 2, 2, 2], knn_list=[10, 20], pk=10, - points_label=False, local_folding=False): - super(MSAP_SKN_decoder, self).__init__() - self.num_coarse_raw = num_coarse_raw - self.num_fps = num_fps - self.num_coarse = num_coarse - self.num_fine = num_fine - self.points_label = points_label - self.local_folding = local_folding - - self.fc1 = nn.Linear(1024, 1024) - self.fc2 = nn.Linear(1024, 1024) - self.fc3 = nn.Linear(1024, num_coarse_raw * 3) - - self.dense_feature_size = 256 - self.expand_feature_size = 64 - - if points_label: - self.input_size = 4 - else: - self.input_size = 3 - - - self.encoder = SA_SKN_Res_encoder(input_size=self.input_size, k=knn_list, pk=pk, - output_size=self.dense_feature_size, layers=layers) - - self.up_scale = int(np.ceil(num_fine / (num_coarse_raw + 2048))) - - if self.up_scale >= 2: - self.expansion1 = EF_expansion(input_size=self.dense_feature_size, output_size=self.expand_feature_size, - step_ratio=self.up_scale, k=4) - self.conv_cup1 = nn.Conv1d(self.expand_feature_size, self.expand_feature_size, 1) - else: - self.expansion1 = None - self.conv_cup1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - - self.conv_cup2 = nn.Conv1d(self.expand_feature_size, 3, 1, bias=True) - - self.conv_s1 = nn.Conv1d(self.expand_feature_size, 16, 1, bias=True) - self.conv_s2 = nn.Conv1d(16, 8, 1, bias=True) - self.conv_s3 = nn.Conv1d(8, 1, 1, bias=True) - - if self.local_folding: - self.expansion2 = Folding(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(num_fine // num_coarse)) - else: - self.expansion2 = EF_expansion(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(num_fine // num_coarse), k=4) - - self.conv_f1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - self.conv_f2 = nn.Conv1d(self.expand_feature_size, 3, 1) - - self.af = nn.ReLU(inplace=False) - - def forward(self, global_feat, point_input): - batch_size = global_feat.size()[0] - - coarse_raw = self.fc3(self.af(self.fc2(self.af(self.fc1(global_feat))))).view(batch_size, 3, - self.num_coarse_raw) - - input_points_num = point_input.size()[2] - org_points_input = point_input - - if self.points_label: - id0 = torch.zeros(coarse_raw.shape[0], 1, coarse_raw.shape[2]).cuda().contiguous() - coarse_input = torch.cat( (coarse_raw, id0), 1) - id1 = torch.ones(org_points_input.shape[0], 1, org_points_input.shape[2]).cuda().contiguous() - org_points_input = torch.cat( (org_points_input, id1), 1) - else: - coarse_input = coarse_raw - - points = torch.cat((coarse_input, org_points_input), 2) - dense_feat = self.encoder(points) - - if self.up_scale >= 2: - dense_feat = self.expansion1(dense_feat) - - coarse_features = self.af(self.conv_cup1(dense_feat)) - coarse_high = self.conv_cup2(coarse_features) - - if coarse_high.size()[2] > self.num_fps: - idx_fps = pn2.furthest_point_sample(coarse_high.transpose(1, 2).contiguous(), self.num_fps) - coarse_fps = pn2.gather_operation(coarse_high, idx_fps) - coarse_features = pn2.gather_operation(coarse_features, idx_fps) - else: - coarse_fps = coarse_high - - if coarse_fps.size()[2] > self.num_coarse: - scores = F.softplus(self.conv_s3(self.af(self.conv_s2(self.af(self.conv_s1(coarse_features)))))) - idx_scores = scores.topk(k=self.num_coarse, dim=2)[1].view(batch_size, -1).int() - coarse = pn2.gather_operation(coarse_fps, idx_scores) - coarse_features = pn2.gather_operation(coarse_features, idx_scores) - else: - coarse = coarse_fps - - if coarse.size()[2] < self.num_fine: - if self.local_folding: - up_features = self.expansion2(coarse_features, global_feat) - center = coarse.transpose(2, 1).contiguous().unsqueeze(2).repeat(1, 1, self.num_fine // self.num_coarse, - 1).view(batch_size, self.num_fine, - 3).transpose(2, 1).contiguous() - fine = self.conv_f2(self.af(self.conv_f1(up_features))) + center - else: - up_features = self.expansion2(coarse_features) - fine = self.conv_f2(self.af(self.conv_f1(up_features))) - else: - assert (coarse.size()[2] == self.num_fine) - fine = coarse - - return coarse_raw, coarse_high, coarse, fine - - -class Model(nn.Module): - def __init__(self, args, size_z=128, global_feature_size=1024): - super(Model, self).__init__() - - layers = [int(i) for i in args.layers.split(',')] - knn_list = [int(i) for i in args.knn_list.split(',')] - - self.size_z = size_z - self.distribution_loss = args.distribution_loss - self.train_loss = args.loss - self.encoder = PCN_encoder(output_size=global_feature_size) - self.posterior_infer1 = Linear_ResBlock(input_size=global_feature_size, output_size=global_feature_size) - self.posterior_infer2 = Linear_ResBlock(input_size=global_feature_size, output_size=size_z * 2) - self.prior_infer = Linear_ResBlock(input_size=global_feature_size, output_size=size_z * 2) - self.generator = Linear_ResBlock(input_size=size_z, output_size=global_feature_size) - self.decoder = MSAP_SKN_decoder(num_fps=args.num_fps, num_fine=args.num_points, num_coarse=args.num_coarse, - num_coarse_raw=args.num_coarse_raw, layers=layers, knn_list=knn_list, - pk=args.pk, local_folding=args.local_folding, points_label=args.points_label) - - def compute_kernel(self, x, y): - x_size = x.size()[0] - y_size = y.size()[0] - dim = x.size()[1] - - tiled_x = x.unsqueeze(1).repeat(1, y_size, 1) - tiled_y = y.unsqueeze(0).repeat(x_size, 1, 1) - return torch.exp(-torch.mean((tiled_x - tiled_y)**2, dim=2) / float(dim)) - - def mmd_loss(self, x, y): - x_kernel = self.compute_kernel(x, x) - y_kernel = self.compute_kernel(y, y) - xy_kernel = self.compute_kernel(x, y) - return torch.mean(x_kernel) + torch.mean(y_kernel) - 2 * torch.mean(xy_kernel) - - def forward(self, x, gt, is_training=True, mean_feature=None, alpha=None): - num_input = x.size()[2] - - if is_training: - y = pn2.gather_operation(gt.transpose(1, 2).contiguous(), pn2.furthest_point_sample(gt, num_input)) - gt = torch.cat([gt, gt], dim=0) - points = torch.cat([x, y], dim=0) - x = torch.cat([x, x], dim=0) - else: - points = x - feat = self.encoder(points) - - if is_training: - feat_x, feat_y = feat.chunk(2) - o_x = self.posterior_infer2(self.posterior_infer1(feat_x)) - q_mu, q_std = torch.split(o_x, self.size_z, dim=1) - o_y = self.prior_infer(feat_y) - p_mu, p_std = torch.split(o_y, self.size_z, dim=1) - q_std = F.softplus(q_std) - p_std = F.softplus(p_std) - q_distribution = torch.distributions.Normal(q_mu, q_std) - p_distribution = torch.distributions.Normal(p_mu, p_std) - p_distribution_fix = torch.distributions.Normal(p_mu.detach(), p_std.detach()) - m_distribution = torch.distributions.Normal(torch.zeros_like(p_mu), torch.ones_like(p_std)) - z_q = q_distribution.rsample() - z_p = p_distribution.rsample() - z = torch.cat([z_q, z_p], dim=0) - feat = torch.cat([feat_x, feat_x], dim=0) - - else: - o_x = self.posterior_infer2(self.posterior_infer1(feat)) - q_mu, q_std = torch.split(o_x, self.size_z, dim=1) - q_std = F.softplus(q_std) - q_distribution = torch.distributions.Normal(q_mu, q_std) - p_distribution = q_distribution - p_distribution_fix = p_distribution - m_distribution = p_distribution - z = q_distribution.rsample() - - feat += self.generator(z) - - coarse_raw, coarse_high, coarse, fine = self.decoder(feat, x) - coarse_raw = coarse_raw.transpose(1, 2).contiguous() - coarse_high = coarse_high.transpose(1, 2).contiguous() - coarse = coarse.transpose(1, 2).contiguous() - fine = fine.transpose(1, 2).contiguous() - - if is_training: - if self.distribution_loss == 'MMD': - z_m = m_distribution.rsample() - z_q = q_distribution.rsample() - z_p = p_distribution.rsample() - z_p_fix = p_distribution_fix.rsample() - dl_rec = self.mmd_loss(z_m, z_p) - dl_g = self.mmd_loss2(z_q, z_p_fix) - elif self.distribution_loss == 'KLD': - dl_rec = torch.distributions.kl_divergence(m_distribution, p_distribution) - dl_g = torch.distributions.kl_divergence(p_distribution_fix, q_distribution) - else: - raise NotImplementedError('Distribution loss is either MMD or KLD') - - if self.train_loss == 'cd': - loss1, _ = calc_cd(coarse_raw, gt) - loss2, _ = calc_cd(coarse_high, gt) - loss3, _ = calc_cd(coarse, gt) - loss4, _ = calc_cd(fine, gt) - else: - raise NotImplementedError('Only CD is supported') - - total_train_loss = loss1.mean() * 10 + loss2.mean() * 0.5 + loss3.mean() + loss4.mean() * alpha - total_train_loss += (dl_rec.mean() + dl_g.mean()) * 20 - return fine, loss4, total_train_loss - else: - emd = calc_emd(fine, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(fine, gt, calc_f1=True) - return {'out1': coarse_raw, 'out2': fine, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} diff --git a/other_models/VRCNet/vrcnet.py b/other_models/VRCNet/vrcnet.py deleted file mode 100644 index 4e55257..0000000 --- a/other_models/VRCNet/vrcnet.py +++ /dev/null @@ -1,571 +0,0 @@ -import numpy as np -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -import torch.nn.functional as F -from other_models.VRCNet.utils.model_utils import * -from other_models.VRCNet.models.pcn import PCN_encoder - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "VRCNet/utils/pointnet2")) -import pointnet2_utils as pn2 - - -class SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=16): - super(SA_module, self).__init__() - self.share_planes = share_planes - self.k = k - self.conv1 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv2 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv3 = nn.Conv2d(in_planes, mid_planes, kernel_size=1) - - self.conv_w = nn.Sequential(nn.ReLU(inplace=False), - nn.Conv2d(rel_planes * (k + 1), mid_planes // share_planes, kernel_size=1, - bias=False), - nn.ReLU(inplace=False), - nn.Conv2d(mid_planes // share_planes, k * mid_planes // share_planes, - kernel_size=1)) - self.activation_fn = nn.ReLU(inplace=False) - - self.conv_out = nn.Conv2d(mid_planes, out_planes, kernel_size=1) - - def forward(self, input): - x, idx = input - batch_size, _, _, num_points = x.size() - identity = x # B C 1 N - x = self.activation_fn(x) - xn = get_edge_features(x, idx) # B C K N - x1, x2, x3 = self.conv1(x), self.conv2(xn), self.conv3(xn) - - x2 = x2.view(batch_size, -1, 1, num_points).contiguous() # B kC 1 N - w = self.conv_w(torch.cat([x1, x2], 1)).view(batch_size, -1, self.k, num_points) - w = w.repeat(1, self.share_planes, 1, 1) - out = w * x3 - out = torch.sum(out, dim=2, keepdim=True) - - out = self.activation_fn(out) - out = self.conv_out(out) # B C 1 N - out += identity - return [out, idx] - - -class Folding(nn.Module): - def __init__(self, input_size, output_size, step_ratio, global_feature_size=1024, num_models=1): - super(Folding, self).__init__() - self.input_size = input_size - self.output_size = output_size - self.step_ratio = step_ratio - self.num_models = num_models - - self.conv = nn.Conv1d(input_size + global_feature_size + 2, output_size, 1, bias=True) - - sqrted = int(math.sqrt(step_ratio)) + 1 - for i in range(1, sqrted + 1).__reversed__(): - if (step_ratio % i) == 0: - num_x = i - num_y = step_ratio // i - break - - grid_x = torch.linspace(-0.2, 0.2, steps=num_x) - grid_y = torch.linspace(-0.2, 0.2, steps=num_y) - - x, y = torch.meshgrid(grid_x, grid_y) # x, y shape: (2, 1) - self.grid = torch.stack([x, y], dim=-1).view(-1, 2) # (2, 2) - - def forward(self, point_feat, global_feat): - batch_size, num_features, num_points = point_feat.size() - point_feat = point_feat.transpose(1, 2).contiguous().unsqueeze(2).repeat(1, 1, self.step_ratio, 1).view( - batch_size, - -1, num_features).transpose(1, 2).contiguous() - global_feat = global_feat.unsqueeze(2).repeat(1, 1, num_points * self.step_ratio).repeat(self.num_models, 1, 1) - grid_feat = self.grid.unsqueeze(0).repeat(batch_size, num_points, 1).transpose(1, 2).contiguous().cuda() - features = torch.cat([global_feat, point_feat, grid_feat], axis=1) - features = F.relu(self.conv(features)) - return features - - -class Linear_ResBlock(nn.Module): - def __init__(self, input_size=1024, output_size=256): - super(Linear_ResBlock, self).__init__() - self.conv1 = nn.Linear(input_size, input_size) - self.conv2 = nn.Linear(input_size, output_size) - self.conv_res = nn.Linear(input_size, output_size) - - # self.af = nn.ReLU(inplace=True) - self.af = nn.ReLU(inplace=False) - - def forward(self, feature): - return self.conv2(self.af(self.conv1(self.af(feature)))) + self.conv_res(feature) - - -class SK_SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=[10, 20], r=2, L=32): - super(SK_SA_module, self).__init__() - - self.num_kernels = len(k) - d = max(int(out_planes / r), L) - - self.sams = nn.ModuleList([]) - - for i in range(len(k)): - self.sams.append(SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k[i])) - - self.fc = nn.Linear(out_planes, d) - self.fcs = nn.ModuleList([]) - - for i in range(len(k)): - self.fcs.append(nn.Linear(d, out_planes)) - - self.softmax = nn.Softmax(dim=1) - self.af = nn.ReLU(inplace=False) - - def forward(self, input): - x, idxs = input - assert (self.num_kernels == len(idxs)) - for i, sam in enumerate(self.sams): - fea, _ = sam([x, idxs[i]]) - fea = self.af(fea) - fea = fea.unsqueeze(dim=1) - if i == 0: - feas = fea - else: - feas = torch.cat([feas, fea], dim=1) - - fea_U = torch.sum(feas, dim=1) - fea_s = fea_U.mean(-1).mean(-1) - fea_z = self.fc(fea_s) - - for i, fc in enumerate(self.fcs): - vector = fc(fea_z).unsqueeze_(dim=1) - if i == 0: - attention_vectors = vector - else: - attention_vectors = torch.cat([attention_vectors, vector], dim=1) - - attention_vectors = self.softmax(attention_vectors) - attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1) - fea_v = (feas * attention_vectors).sum(dim=1) - return [fea_v, idxs] - - -class SKN_Res_unit(nn.Module): - def __init__(self, input_size, output_size, k=[10, 20], layers=1): - super(SKN_Res_unit, self).__init__() - self.conv1 = nn.Conv2d(input_size, output_size, 1, bias=False) - self.sam = self._make_layer(output_size, output_size // 16, output_size // 4, output_size, int(layers), 8, k=k) - self.conv2 = nn.Conv2d(output_size, output_size, 1, bias=False) - self.conv_res = nn.Conv2d(input_size, output_size, 1, bias=False) - self.af = nn.ReLU(inplace=False) - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def forward(self, feat, idx): - x, _ = self.sam([self.conv1(feat), idx]) - x = self.conv2(self.af(x)) - return x + self.conv_res(feat) - - -class SA_SKN_Res_encoder(nn.Module): - def __init__(self, input_size=3, k=[10, 20], pk=16, output_size=64, layers=[2, 2, 2, 2], - pts_num=[3072, 1536, 768, 384]): - super(SA_SKN_Res_encoder, self).__init__() - self.init_channel = 64 - - c1 = self.init_channel - self.sam_res1 = SKN_Res_unit(input_size, c1, k, int(layers[0])) - - c2 = c1 * 2 - self.sam_res2 = SKN_Res_unit(c2, c2, k, int(layers[1])) - - c3 = c2 * 2 - self.sam_res3 = SKN_Res_unit(c3, c3, k, int(layers[2])) - - c4 = c3 * 2 - self.sam_res4 = SKN_Res_unit(c4, c4, k, int(layers[3])) - - self.conv5 = nn.Conv2d(c4, 1024, 1) - - self.fc1 = nn.Linear(1024, 512) - self.fc2 = nn.Linear(512, 1024) - - self.conv6 = nn.Conv2d(c4 + 1024, c4, 1) - self.conv7 = nn.Conv2d(c3 + c4, c3, 1) - self.conv8 = nn.Conv2d(c2 + c3, c2, 1) - self.conv9 = nn.Conv2d(c1 + c2, c1, 1) - - self.conv_out = nn.Conv2d(c1, output_size, 1) - self.dropout = nn.Dropout() - self.af = nn.ReLU(inplace=False) - self.k = k - self.pk = pk - self.rate = 2 - - self.pts_num = pts_num - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def _edge_pooling(self, features, points, rate=2, k=16, sample_num=None): - features = features.squeeze(2) - - if sample_num is None: - input_points_num = int(features.size()[2]) - sample_num = input_points_num // rate - - ds_features, p_idx, pn_idx, ds_points = edge_preserve_sampling(features, points, sample_num, k) - ds_features = ds_features.unsqueeze(2) - return ds_features, p_idx, pn_idx, ds_points - - def _edge_unpooling(self, features, src_pts, tgt_pts): - features = features.squeeze(2) - idx, weight = three_nn_upsampling(tgt_pts, src_pts) - features = pn2.three_interpolate(features, idx, weight) - features = features.unsqueeze(2) - return features - - def forward(self, features): - batch_size, _, num_points = features.size() - pt1 = features[:, 0:3, :] - - idx1 = [] - for i in range(len(self.k)): - idx = knn(pt1, self.k[i]) - idx1.append(idx) - - pt1 = pt1.transpose(1, 2).contiguous() - - x = features.unsqueeze(2) - x = self.sam_res1(x, idx1) - x1 = self.af(x) - - x, _, _, pt2 = self._edge_pooling(x1, pt1, self.rate, self.pk, self.pts_num[1]) - idx2 = [] - for i in range(len(self.k)): - idx = knn(pt2.transpose(1, 2).contiguous(), self.k[i]) - idx2.append(idx) - - x = self.sam_res2(x, idx2) - x2 = self.af(x) - - x, _, _, pt3 = self._edge_pooling(x2, pt2, self.rate, self.pk, self.pts_num[2]) - idx3 = [] - for i in range(len(self.k)): - idx = knn(pt3.transpose(1, 2).contiguous(), self.k[i]) - idx3.append(idx) - - x = self.sam_res3(x, idx3) - x3 = self.af(x) - - x, _, _, pt4 = self._edge_pooling(x3, pt3, self.rate, self.pk, self.pts_num[3]) - idx4 = [] - for i in range(len(self.k)): - idx = knn(pt4.transpose(1, 2).contiguous(), self.k[i]) - idx4.append(idx) - - x = self.sam_res4(x, idx4) - x4 = self.af(x) - x = self.conv5(x4) - x, _ = torch.max(x, -1) - x = x.view(batch_size, -1) - x = self.dropout(self.af(self.fc2(self.dropout(self.af(self.fc1(x)))))) - - x = x.unsqueeze(2).repeat(1, 1, self.pts_num[3]).unsqueeze(2) - x = self.af(self.conv6(torch.cat([x, x4], 1))) - x = self._edge_unpooling(x, pt4, pt3) - x = self.af(self.conv7(torch.cat([x, x3], 1))) - x = self._edge_unpooling(x, pt3, pt2) - x = self.af(self.conv8(torch.cat([x, x2], 1))) - x = self._edge_unpooling(x, pt2, pt1) - x = self.af(self.conv9(torch.cat([x, x1], 1))) - x = self.conv_out(x) - x = x.squeeze(2) - return x - - -class MSAP_SKN_decoder(nn.Module): - def __init__(self, num_coarse_raw, num_fps, num_coarse, num_fine, layers=[2, 2, 2, 2], knn_list=[10, 20], pk=10, - points_label=False, local_folding=False): - super(MSAP_SKN_decoder, self).__init__() - self.num_coarse_raw = num_coarse_raw - self.num_fps = num_fps - self.num_coarse = num_coarse - self.num_fine = num_fine - self.points_label = points_label - self.local_folding = local_folding - - self.fc1 = nn.Linear(1024, 1024) - self.fc2 = nn.Linear(1024, 1024) - self.fc3 = nn.Linear(1024, num_coarse_raw * 3) - - self.dense_feature_size = 256 - self.expand_feature_size = 64 - - if points_label: - self.input_size = 4 - else: - self.input_size = 3 - - - self.encoder = SA_SKN_Res_encoder(input_size=self.input_size, k=knn_list, pk=pk, - output_size=self.dense_feature_size, layers=layers) - - self.up_scale = int(np.ceil(num_fine / (num_coarse_raw + 2048))) - - if self.up_scale >= 2: - self.expansion1 = EF_expansion(input_size=self.dense_feature_size, output_size=self.expand_feature_size, - step_ratio=self.up_scale, k=4) - self.conv_cup1 = nn.Conv1d(self.expand_feature_size, self.expand_feature_size, 1) - else: - self.expansion1 = None - self.conv_cup1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - - self.conv_cup2 = nn.Conv1d(self.expand_feature_size, 3, 1, bias=True) - - self.conv_s1 = nn.Conv1d(self.expand_feature_size, 16, 1, bias=True) - self.conv_s2 = nn.Conv1d(16, 8, 1, bias=True) - self.conv_s3 = nn.Conv1d(8, 1, 1, bias=True) - - if self.local_folding: - self.expansion2 = Folding(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(num_fine // num_coarse)) - else: - self.expansion2 = EF_expansion(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(num_fine // num_coarse), k=4) - - self.conv_f1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - self.conv_f2 = nn.Conv1d(self.expand_feature_size, 3, 1) - - self.af = nn.ReLU(inplace=False) - """ - # Here is changing the coarse output - - from GCN.gcn_decode import Generator - self.gcn_dec = Generator( - features=[1024, 256, 256, 256, 128, 128, 128, 3], - degrees=[1, 2, 2, 2, 2, 4, 32], - support=10, - root_num=1) - """ - - def forward(self, global_feat, point_input): - batch_size = global_feat.size()[0] - - coarse_raw = self.fc3(self.af(self.fc2(self.af(self.fc1(global_feat))))).view(batch_size, 3, - self.num_coarse_raw) - # Here is changing the coarse output - """ - coarse_raw = self.gcn_dec([global_feat.unsqueeze(1)]).transpose(1, 2) - """ - - input_points_num = point_input.size()[2] - org_points_input = point_input - - if self.points_label: - id0 = torch.zeros(coarse_raw.shape[0], 1, coarse_raw.shape[2]).cuda().contiguous() - coarse_input = torch.cat( (coarse_raw, id0), 1) - id1 = torch.ones(org_points_input.shape[0], 1, org_points_input.shape[2]).cuda().contiguous() - org_points_input = torch.cat( (org_points_input, id1), 1) - else: - coarse_input = coarse_raw - - points = torch.cat((coarse_input, org_points_input), 2) - dense_feat = self.encoder(points) - - if self.up_scale >= 2: - dense_feat = self.expansion1(dense_feat) - - coarse_features = self.af(self.conv_cup1(dense_feat)) - coarse_high = self.conv_cup2(coarse_features) - - if coarse_high.size()[2] > self.num_fps: - idx_fps = pn2.furthest_point_sample(coarse_high.transpose(1, 2).contiguous(), self.num_fps) - coarse_fps = pn2.gather_operation(coarse_high, idx_fps) - coarse_features = pn2.gather_operation(coarse_features, idx_fps) - else: - coarse_fps = coarse_high - - if coarse_fps.size()[2] > self.num_coarse: - scores = F.softplus(self.conv_s3(self.af(self.conv_s2(self.af(self.conv_s1(coarse_features)))))) - idx_scores = scores.topk(k=self.num_coarse, dim=2)[1].view(batch_size, -1).int() - coarse = pn2.gather_operation(coarse_fps, idx_scores) - coarse_features = pn2.gather_operation(coarse_features, idx_scores) - else: - coarse = coarse_fps - - if coarse.size()[2] < self.num_fine: - if self.local_folding: - up_features = self.expansion2(coarse_features, global_feat) - center = coarse.transpose(2, 1).contiguous().unsqueeze(2).repeat(1, 1, self.num_fine // self.num_coarse, - 1).view(batch_size, self.num_fine, - 3).transpose(2, 1).contiguous() - fine = self.conv_f2(self.af(self.conv_f1(up_features))) + center - else: - up_features = self.expansion2(coarse_features) - fine = self.conv_f2(self.af(self.conv_f1(up_features))) - else: - assert (coarse.size()[2] == self.num_fine) - fine = coarse - - return coarse_raw, coarse_high, coarse, fine - - -class Model(nn.Module): - def __init__(self, args, size_z=128, global_feature_size=1024): - super(Model, self).__init__() - - layers = [int(i) for i in args.layers.split(',')] - knn_list = [int(i) for i in args.knn_list.split(',')] - - self.size_z = size_z - self.distribution_loss = args.distribution_loss - self.train_loss = args.loss - self.encoder = PCN_encoder(output_size=global_feature_size) - """ - import GCN.gcn_encode as gcn3d - self.encoder = gcn3d.GCN3DFeatDeep(support_num=10, neighbor_num=10) - """ - self.posterior_infer1 = Linear_ResBlock(input_size=global_feature_size, output_size=global_feature_size) - self.posterior_infer2 = Linear_ResBlock(input_size=global_feature_size, output_size=size_z * 2) - self.prior_infer = Linear_ResBlock(input_size=global_feature_size, output_size=size_z * 2) - self.generator = Linear_ResBlock(input_size=size_z, output_size=global_feature_size) - self.decoder = MSAP_SKN_decoder(num_fps=args.num_fps, num_fine=args.num_points, num_coarse=args.num_coarse, - num_coarse_raw=args.num_coarse_raw, layers=layers, knn_list=knn_list, - pk=args.pk, local_folding=args.local_folding, points_label=args.points_label) - - def compute_kernel(self, x, y): - x_size = x.size()[0] - y_size = y.size()[0] - dim = x.size()[1] - - tiled_x = x.unsqueeze(1).repeat(1, y_size, 1) - tiled_y = y.unsqueeze(0).repeat(x_size, 1, 1) - return torch.exp(-torch.mean((tiled_x - tiled_y)**2, dim=2) / float(dim)) - - def mmd_loss(self, x, y): - x_kernel = self.compute_kernel(x, x) - y_kernel = self.compute_kernel(y, y) - xy_kernel = self.compute_kernel(x, y) - return torch.mean(x_kernel) + torch.mean(y_kernel) - 2 * torch.mean(xy_kernel) - - def forward(self, x, mean_feature=None, alpha=None): - num_input = x.size()[2] - - points = x - feat = self.encoder(points) - # feat = self.encoder(points).view(x.shape[0], -1) - - o_x = self.posterior_infer2(self.posterior_infer1(feat)) - q_mu, q_std = torch.split(o_x, self.size_z, dim=1) - q_std = F.softplus(q_std) - q_distribution = torch.distributions.Normal(q_mu, q_std) - p_distribution = q_distribution - p_distribution_fix = p_distribution - m_distribution = p_distribution - z = q_distribution.rsample() - - feat = feat + self.generator(z) - - coarse_raw, coarse_high, coarse, fine = self.decoder(feat, x) - coarse_raw = coarse_raw.transpose(1, 2).contiguous() - coarse_high = coarse_high.transpose(1, 2).contiguous() - coarse = coarse.transpose(1, 2).contiguous() - fine = fine.transpose(1, 2).contiguous() - - # return {'out1': coarse_raw, 'out2': fine, 'out3': coarse_high, 'out4': coarse} - return coarse_raw, fine, coarse_high, coarse - - def trainer(self, x, gt, mean_feature=None, alpha=None): - num_input = x.size()[2] - - y = pn2.gather_operation(gt.transpose(1, 2).contiguous(), pn2.furthest_point_sample(gt, num_input)) - gt = torch.cat([gt, gt], dim=0) - points = torch.cat([x, y], dim=0) - x = torch.cat([x, x], dim=0) - feat = self.encoder(points) - # feat = self.encoder(points.transpose(1, 2)).view(x.shape[0], -1) - - feat_x, feat_y = feat.chunk(2) - o_x = self.posterior_infer2(self.posterior_infer1(feat_x)) - q_mu, q_std = torch.split(o_x, self.size_z, dim=1) - o_y = self.prior_infer(feat_y) - p_mu, p_std = torch.split(o_y, self.size_z, dim=1) - q_std = F.softplus(q_std) - p_std = F.softplus(p_std) - q_distribution = torch.distributions.Normal(q_mu, q_std) - p_distribution = torch.distributions.Normal(p_mu, p_std) - p_distribution_fix = torch.distributions.Normal(p_mu.detach(), p_std.detach()) - m_distribution = torch.distributions.Normal(torch.zeros_like(p_mu), torch.ones_like(p_std)) - z_q = q_distribution.rsample() - z_p = p_distribution.rsample() - z = torch.cat([z_q, z_p], dim=0) - feat = torch.cat([feat_x, feat_x], dim=0) - - - feat += self.generator(z) - - coarse_raw, coarse_high, coarse, fine = self.decoder(feat, x) - coarse_raw = coarse_raw.transpose(1, 2).contiguous() - coarse_high = coarse_high.transpose(1, 2).contiguous() - coarse = coarse.transpose(1, 2).contiguous() - fine = fine.transpose(1, 2).contiguous() - - if self.distribution_loss == 'MMD': - z_m = m_distribution.rsample() - z_q = q_distribution.rsample() - z_p = p_distribution.rsample() - z_p_fix = p_distribution_fix.rsample() - dl_rec = self.mmd_loss(z_m, z_p) - dl_g = self.mmd_loss2(z_q, z_p_fix) - elif self.distribution_loss == 'KLD': - dl_rec = torch.distributions.kl_divergence(m_distribution, p_distribution) - dl_g = torch.distributions.kl_divergence(p_distribution_fix, q_distribution) - else: - raise NotImplementedError('Distribution loss is either MMD or KLD') - - if self.train_loss == 'cd': - loss1, _ = calc_cd(coarse_raw, gt) - loss2, _ = calc_cd(coarse_high, gt) - loss3, _ = calc_cd(coarse, gt) - loss4, _ = calc_cd(fine, gt) - else: - raise NotImplementedError('Only CD is supported') - - total_train_loss = loss1.mean() * 10 + loss2.mean() * 0.5 + loss3.mean() + loss4.mean() * alpha - total_train_loss += (dl_rec.mean() + dl_g.mean()) * 20 - return fine, loss4, total_train_loss - - def evaluate(self, x, gt, mean_feature=None, alpha=None): - num_input = x.size()[2] - - points = x - feat = self.encoder(points) - - o_x = self.posterior_infer2(self.posterior_infer1(feat)) - q_mu, q_std = torch.split(o_x, self.size_z, dim=1) - q_std = F.softplus(q_std) - q_distribution = torch.distributions.Normal(q_mu, q_std) - p_distribution = q_distribution - p_distribution_fix = p_distribution - m_distribution = p_distribution - z = q_distribution.rsample() - - feat += self.generator(z) - - coarse_raw, coarse_high, coarse, fine = self.decoder(feat, x) - coarse_raw = coarse_raw.transpose(1, 2).contiguous() - coarse_high = coarse_high.transpose(1, 2).contiguous() - coarse = coarse.transpose(1, 2).contiguous() - fine = fine.transpose(1, 2).contiguous() - - emd = calc_emd(fine, gt, eps=0.004, iterations=3000) - cd_p, cd_t, f1 = calc_cd(fine, gt, calc_f1=True) - return {'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} diff --git a/other_models/VRCNet/vrcnet.yaml b/other_models/VRCNet/vrcnet.yaml deleted file mode 100644 index dae9b6e..0000000 --- a/other_models/VRCNet/vrcnet.yaml +++ /dev/null @@ -1,41 +0,0 @@ -batch_size: 32 -workers: 0 -nepoch: 100 -model_name: vrcnet -load_model: null -start_epoch: 0 -num_points: 4096 -work_dir: log/ -flag: debug -loss: cd -manual_seed: null -use_mean_feature: False -step_interval_to_print: 500 -epoch_interval_to_save: 1 -epoch_interval_to_val: 1 -varying_constant: 0.01, 0.1, 0.5, 1 -varying_constant_epochs: 5, 15, 30 - -lr: 0.0001 -lr_decay: True -lr_decay_interval: 40 -lr_decay_rate: 0.7 -lr_step_decay_epochs: null -lr_step_decay_rates: null -lr_clip: 1.e-6 -optimizer: Adam -weight_decay: 0 -betas: 0.9, 0.999 - -layers: 1, 1, 1, 1 -distribution_loss: KLD -knn_list: "16" -pk: 10 -local_folding: True -points_label: True -num_coarse_raw: 2048 -num_fps: 2048 -num_coarse: 2048 - -# test -save_vis: True diff --git a/other_models/displace/__pycache__/decode.cpython-36.pyc b/other_models/displace/__pycache__/decode.cpython-36.pyc deleted file mode 100644 index 9436cdd..0000000 Binary files a/other_models/displace/__pycache__/decode.cpython-36.pyc and /dev/null differ diff --git a/other_models/displace/__pycache__/encode.cpython-36.pyc b/other_models/displace/__pycache__/encode.cpython-36.pyc deleted file mode 100644 index 808bba3..0000000 Binary files a/other_models/displace/__pycache__/encode.cpython-36.pyc and /dev/null differ diff --git a/other_models/displace/__pycache__/gcn_decode.cpython-36.pyc b/other_models/displace/__pycache__/gcn_decode.cpython-36.pyc deleted file mode 100644 index 8be8366..0000000 Binary files a/other_models/displace/__pycache__/gcn_decode.cpython-36.pyc and /dev/null differ diff --git a/other_models/displace/__pycache__/gcn_encode.cpython-36.pyc b/other_models/displace/__pycache__/gcn_encode.cpython-36.pyc deleted file mode 100644 index 71fbb84..0000000 Binary files a/other_models/displace/__pycache__/gcn_encode.cpython-36.pyc and /dev/null differ diff --git a/other_models/displace/__pycache__/gcn_encode_siren.cpython-36.pyc b/other_models/displace/__pycache__/gcn_encode_siren.cpython-36.pyc deleted file mode 100644 index 248be2b..0000000 Binary files a/other_models/displace/__pycache__/gcn_encode_siren.cpython-36.pyc and /dev/null differ diff --git a/other_models/displace/__pycache__/tree.cpython-36.pyc b/other_models/displace/__pycache__/tree.cpython-36.pyc deleted file mode 100644 index 9c74f2f..0000000 Binary files a/other_models/displace/__pycache__/tree.cpython-36.pyc and /dev/null differ diff --git a/other_models/displace/decode.py b/other_models/displace/decode.py deleted file mode 100644 index b05605d..0000000 --- a/other_models/displace/decode.py +++ /dev/null @@ -1,86 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -from displace.tree import TreeGCN - -from math import ceil - - -class Discriminator(nn.Module): - def __init__(self, batch_size, features): - self.batch_size = batch_size - self.layer_num = len(features) - 1 - super(Discriminator, self).__init__() - - self.fc_layers = nn.ModuleList([]) - for inx in range(self.layer_num): - self.fc_layers.append( - nn.Conv1d( - features[inx], features[inx + 1], kernel_size=1, stride=1)) - - self.leaky_relu = nn.LeakyReLU(negative_slope=0.2) - self.final_layer = nn.Sequential( - nn.Linear(features[-1], features[-1]), - nn.Linear(features[-1], features[-2]), - nn.Linear(features[-2], features[-2]), nn.Linear(features[-2], 1)) - - def forward(self, f): - feat = f.transpose(1, 2) - vertex_num = feat.size(2) - - for inx in range(self.layer_num): - feat = self.fc_layers[inx](feat) - feat = self.leaky_relu(feat) - - out = F.max_pool1d(input=feat, kernel_size=vertex_num).squeeze(-1) - out = self.final_layer(out) # (B, 1) - - return out - - -class Decoder(nn.Module): - def __init__(self, features, degrees, support, root_num): - # self.batch_size = batch_size - self.layer_num = len(features) - 1 - assert self.layer_num == len( - degrees - ), "Number of features should be one more than number of degrees." - self.pointcloud = None - super(Decoder, self).__init__() - - vertex_num = root_num - self.gcn = nn.Sequential() - for inx in range(self.layer_num): - if inx == self.layer_num - 1: - self.gcn.add_module( - 'TreeGCN_' + str(inx), - TreeGCN( - inx, - features, - degrees, - support=support, - node=vertex_num, - upsample=True, - activation=False)) - else: - self.gcn.add_module( - 'TreeGCN_' + str(inx), - TreeGCN( - inx, - features, - degrees, - support=support, - node=vertex_num, - upsample=True, - activation=True)) - vertex_num = int(vertex_num * degrees[inx]) - - def forward(self, tree): - feat = self.gcn(tree) - - self.pointcloud = feat[-1] - - return self.pointcloud - - def getPointcloud(self): - return self.pointcloud[-1] diff --git a/other_models/displace/encode.py b/other_models/displace/encode.py deleted file mode 100644 index 1f6b31c..0000000 --- a/other_models/displace/encode.py +++ /dev/null @@ -1,337 +0,0 @@ -import math -import torch -import torch.nn as nn -import torch.nn.functional as F -from FNet import FNet - - -def get_neighbor_index(vertices: "(bs, vertice_num, 3)", neighbor_num: int): - """ - Return: (bs, vertice_num, neighbor_num) - """ - bs, v, _ = vertices.size() - device = vertices.device - inner = torch.bmm(vertices, vertices.transpose(1, 2)) # (bs, v, v) - quadratic = torch.sum(vertices**2, dim=2) # (bs, v) - distance = inner * (-2) + quadratic.unsqueeze(1) + quadratic.unsqueeze(2) - neighbor_index = torch.topk( - distance, k=neighbor_num + 1, dim=-1, largest=False)[1] - neighbor_index = neighbor_index[:, :, 1:] - return neighbor_index - - -def get_nearest_index(target: "(bs, v1, 3)", source: "(bs, v2, 3)"): - """ - Return: (bs, v1, 1) - """ - inner = torch.bmm(target, source.transpose(1, 2)) # (bs, v1, v2) - s_norm_2 = torch.sum(source**2, dim=2) # (bs, v2) - t_norm_2 = torch.sum(target**2, dim=2) # (bs, v1) - d_norm_2 = s_norm_2.unsqueeze(1) + t_norm_2.unsqueeze(2) - 2 * inner - nearest_index = torch.topk(d_norm_2, k=1, dim=-1, largest=False)[1] - return nearest_index - - -def indexing_neighbor(tensor: "(bs, vertice_num, dim)", - index: "(bs, vertice_num, neighbor_num)"): - """ - Return: (bs, vertice_num, neighbor_num, dim) - """ - bs, v, n = index.size() - id_0 = torch.arange(bs).view(-1, 1, 1) - tensor_indexed = tensor[id_0, index] - return tensor_indexed - - -def get_neighbor_displacement( - vertices: "(bs, vertice_num, 3)", - neighbor_index: "(bs, vertice_num, neighbor_num)"): - """ - Return: (bs, vertice_num, neighobr_num, 3) - """ - neighbors = indexing_neighbor(vertices, neighbor_index) # (bs, v, n, 3) - neighbor_anchored = neighbors - vertices.unsqueeze(2) - return neighbor_anchored - - -class Operator3D(nn.Module): - """ - Extract structure feafure from surface, independent from vertice coordinates - """ - def __init__(self, kernel_num, support_num): - super().__init__() - self.kernel_num = kernel_num - self.support_num = support_num - - self.relu = nn.ReLU(inplace=True) - self.weights = nn.Parameter( - torch.FloatTensor(1, 1, support_num, kernel_num)) - self.displacement = nn.Parameter( - torch.FloatTensor(3, support_num * kernel_num)) - # self.fourier_map = Periodics(dim_input=3, dim_output=32) - self.initialize() - - def initialize(self): - stdv = 1. / math.sqrt(self.support_num * self.kernel_num) - self.weights.data.uniform_(-stdv, stdv) - self.displacement.data.uniform_(-stdv, stdv) - - def forward(self, neighbor_index: "(bs, vertice_num, neighbor_num)", - vertices: "(bs, vertice_num, 3)"): - """ - Return vertices with local feature: (bs, vertice_num, kernel_num) - """ - bs, vertice_num, neighbor_num = neighbor_index.size() - neighbor_displacement = get_neighbor_displacement( - vertices, neighbor_index) - - # NOTE displacements are not normalized - # with shape of (bs, vertice_num, neighbor_num, s*k) - theta = neighbor_displacement @ self.displacement - - theta = self.relu(theta) - theta = theta.contiguous().view(bs, vertice_num, neighbor_num, - self.support_num, self.kernel_num) - theta = torch.max( - theta, dim=2)[0] * self.weights - # (bs, vertice_num, support_num, kernel_num) - feature = torch.sum(theta, dim=2) - # (bs, vertice_num, kernel_num) - # freq = self.fnet_map(vertices) - # return torch.cat((feature, freq), 2) - return feature - -class OperatorND(nn.Module): - def __init__(self, in_channel, out_channel, support_num): - super().__init__() - # arguments: - self.in_channel = in_channel - self.out_channel = out_channel - self.support_num = support_num - - # parameters: - self.relu = nn.ReLU(inplace=True) - self.weights = nn.Parameter( - torch.FloatTensor(in_channel, (support_num + 1) * out_channel)) - self.bias = nn.Parameter( - torch.FloatTensor((support_num + 1) * out_channel)) - self.displacement = nn.Parameter( - torch.FloatTensor(3, support_num * out_channel)) - self.initialize() - - def initialize(self): - stdv = 1. / math.sqrt(self.out_channel * (self.support_num + 1)) - self.weights.data.uniform_(-stdv, stdv) - self.bias.data.uniform_(-stdv, stdv) - self.displacement.data.uniform_(-stdv, stdv) - - def forward(self, neighbor_index: "(bs, vertice_num, neighbor_index)", - vertices: "(bs, vertice_num, 3)", - feature_map: "(bs, vertice_num, in_channel)"): - """ - Return: output feature map: (bs, vertice_num, out_channel) - """ - bs, vertice_num, neighbor_num = neighbor_index.size() - neighbor_displacement = get_neighbor_displacement( - vertices, neighbor_index) - theta = neighbor_displacement @ self.displacement - # (bs, vertice_num, neighbor_num, support_num * out_channel) - theta = self.relu(theta) - theta = theta.contiguous().view(bs, vertice_num, neighbor_num, -1) - # (bs, vertice_num, neighbor_num, support_num * out_channel) - - feature_out = feature_map @ self.weights + self.bias - # (bs, vertice_num, (support_num + 1) * out_channel) - feature_center = feature_out[:, :, :self.out_channel] - # (bs, vertice_num, out_channel) - feature_support = feature_out[:, :, self.out_channel:] - # (bs, vertice_num, support_num * out_channel) - - # Fuse together - max among product - feature_support = indexing_neighbor( - feature_support, neighbor_index - ) - # (bs, vertice_num, neighbor_num, support_num * out_channel) - activation_support = theta * feature_support - # (bs, vertice_num, neighbor_num, support_num * out_channel) - activation_support = activation_support.view( - bs, vertice_num, neighbor_num, self.support_num, self.out_channel) - activation_support = torch.max( - activation_support, - dim=2)[0] - # (bs, vertice_num, support_num, out_channel) - activation_support = torch.sum( - activation_support, dim=2) # (bs, vertice_num, out_channel) - feature_fuse = feature_center + activation_support - # (bs, vertice_num, out_channel) - return feature_fuse - -class Pooling(nn.Module): - def __init__(self, pooling_rate: int = 4, neighbor_num: int = 4): - super().__init__() - self.pooling_rate = pooling_rate - self.neighbor_num = neighbor_num - - def forward(self, vertices: "(bs, vertice_num, 3)", - feature_map: "(bs, vertice_num, channel_num)"): - """ - Return: - vertices_pool: (bs, pool_vertice_num, 3), - feature_map_pool: (bs, pool_vertice_num, channel_num) - """ - bs, vertice_num, _ = vertices.size() - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - neighbor_feature = indexing_neighbor( - feature_map, - neighbor_index) - # (bs, vertice_num, neighbor_num, channel_num) - pooled_feature = torch.max( - neighbor_feature, dim=2)[0] - # (bs, vertice_num, channel_num) - - pool_num = int(vertice_num / self.pooling_rate) - sample_idx = torch.randperm(vertice_num)[:pool_num] - vertices_pool = vertices[:, sample_idx, :] - # (bs, pool_num, 3) - feature_map_pool = pooled_feature[:, sample_idx, :] - # (bs, pool_num, channel_num) - return vertices_pool, feature_map_pool - - -class Encoder(nn.Module): - def __init__(self, support_num: int, neighbor_num: int): - super().__init__() - self.neighbor_num = neighbor_num - - self.conv_0 = Operator3D(kernel_num=32, support_num=support_num) - self.dropping = torch.nn.Dropout(p=0.1, inplace=True) - self.conv_1 = OperatorND(32, 64, support_num=support_num) - self.pool_1 = Pooling(pooling_rate=8, neighbor_num=8) - self.conv_2 = OperatorND(64, 128, support_num=support_num) - self.conv_3 = OperatorND(128, 256, support_num=support_num) - self.pool_2 = Pooling(pooling_rate=8, neighbor_num=8) - self.conv_4 = OperatorND(256, 512, support_num=support_num) - self.conv_5 = OperatorND(512, 512, support_num=support_num) - self.pool_3 = Pooling(pooling_rate=8, neighbor_num=8) - self.conv_6 = OperatorND(512, 1024, support_num=support_num) - - def forward(self, vertices: "(bs, vertice_num, 3)"): - bs, vertice_num, _ = vertices.size() - - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - fm_0 = self.conv_0(neighbor_index, vertices) - fm_0 = self.dropping(fm_0) - fm_0 = F.relu(fm_0, inplace=True) - fm_1 = self.conv_1(neighbor_index, vertices, fm_0) - fm_1 = F.relu(fm_1, inplace=True) - vertices, fm_1 = self.pool_1(vertices, fm_1) - vertices_anchor = vertices - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - - fm_2 = self.conv_2(neighbor_index, vertices, fm_1) - fm_2 = F.relu(fm_2, inplace=True) - fm_3 = self.conv_3(neighbor_index, vertices, fm_2) - fm_3 = F.relu(fm_3, inplace=True) - vertices, fm_3 = self.pool_2(vertices, fm_3) - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - - fm_4 = self.conv_4(neighbor_index, vertices, fm_3) - fm_4 = F.relu(fm_4, inplace=True) - fm_5 = self.conv_5(neighbor_index, vertices, fm_4) - fm_5 = F.relu(fm_5, inplace=True) - vertices, fm_5 = self.pool_3(vertices, fm_5) - neighbor_index = get_neighbor_index(vertices, 3) - - fm_6 = self.conv_6(neighbor_index, vertices, fm_5) - feature_global = fm_6.max(1, keepdim=True)[0] - return feature_global, vertices_anchor - -class Disp3D(nn.Module): - def __init__(self, class_num, support_num, neighbor_num): - super().__init__() - self.neighbor_num = neighbor_num - - self.conv_0 = Operator3D(kernel_num=128, support_num=support_num) - self.conv_1 = OperatorND(128, 128, support_num=support_num) - self.pool_1 = Pooling(pooling_rate=4, neighbor_num=4) - self.conv_2 = OperatorND(128, 256, support_num=support_num) - self.conv_3 = OperatorND(256, 256, support_num=support_num) - self.pool_2 = Pooling(pooling_rate=4, neighbor_num=4) - self.conv_4 = OperatorND(256, 512, support_num=support_num) - - dim_fuse = sum([128, 128, 256, 256, 512, 512]) - self.conv1d_block = nn.Sequential( - nn.Conv1d(dim_fuse, 512, 1), - nn.ReLU(inplace=True), - nn.Conv1d(512, 512, 1), - nn.ReLU(inplace=True), - nn.Conv1d(512, class_num, 1), - ) - - def forward(self, - vertices: "tensor (bs, vetice_num, 3)"): - """ - Return: (bs, vertice_num, class_num) - """ - - bs, vertice_num, _ = vertices.size() - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - - fm_0 = F.relu(self.conv_0(neighbor_index, vertices), inplace=True) - fm_1 = F.relu(self.conv_1(neighbor_index, vertices, fm_0), inplace=True) - v_pool_1, fm_pool_1 = self.pool_1(vertices, fm_1) - neighbor_index = get_neighbor_index(v_pool_1, self.neighbor_num) - - fm_2 = F.relu(self.conv_2(neighbor_index, v_pool_1, fm_pool_1), inplace=True) - fm_3 = F.relu(self.conv_3(neighbor_index, v_pool_1, fm_2), inplace=True) - v_pool_2, fm_pool_2 = self.pool_2(v_pool_1, fm_3) - neighbor_index = get_neighbor_index(v_pool_2, self.neighbor_num) - - fm_4 = self.conv_4(neighbor_index, v_pool_2, fm_pool_2) - f_global = fm_4.max(1)[0] #(bs, f) - - nearest_pool_1 = get_nearest_index(vertices, v_pool_1) - nearest_pool_2 = get_nearest_index(vertices, v_pool_2) - fm_2 = indexing_neighbor(fm_2, nearest_pool_1).squeeze(2) - fm_3 = indexing_neighbor(fm_3, nearest_pool_1).squeeze(2) - fm_4 = indexing_neighbor(fm_4, nearest_pool_2).squeeze(2) - f_global = f_global.unsqueeze(1).repeat(1, vertice_num, 1) - # (bs, vertice_num, cat_one_hot) - fm_fuse = torch.cat([fm_0, fm_1, fm_2, fm_3, fm_4, f_global], dim=2) - - conv1d_input = fm_fuse.permute(0, 2, 1) # (bs, fuse_ch, vertice_num) - conv1d_out = self.conv1d_block(conv1d_input) - pred = conv1d_out.permute(0, 2, 1) # (bs, vertice_num, ch) - return pred - -def test(): - import time - bs = 8 - v = 1024 - dim = 3 - n = 20 - vertices = torch.randn(bs, v, dim) - neighbor_index = get_neighbor_index(vertices, n) - - s = 3 - conv_1 = Operator3D(kernel_num=32, support_num=s) - conv_2 = OperatorND(in_channel=32, out_channel=64, support_num=s) - pool = Pooling(pooling_rate=4, neighbor_num=4) - - print("Input size: {}".format(vertices.size())) - start = time.time() - f1 = conv_1(neighbor_index, vertices) - print("\n[1] Time: {}".format(time.time() - start)) - print("[1] Out shape: {}".format(f1.size())) - start = time.time() - f2 = conv_2(neighbor_index, vertices, f1) - print("\n[2] Time: {}".format(time.time() - start)) - print("[2] Out shape: {}".format(f2.size())) - start = time.time() - v_pool, f_pool = pool(vertices, f2) - print("\n[3] Time: {}".format(time.time() - start)) - print("[3] v shape: {}, f shape: {}".format(v_pool.size(), f_pool.size())) - - -if __name__ == "__main__": - test() diff --git a/other_models/displace/gcn_encode_siren.py b/other_models/displace/gcn_encode_siren.py deleted file mode 100644 index 4199c48..0000000 --- a/other_models/displace/gcn_encode_siren.py +++ /dev/null @@ -1,329 +0,0 @@ -""" -@Author: Zhi-Hao Lin -@Contact: r08942062@ntu.edu.tw -@Time: 2020/03/06 -@Document: Basic operation/blocks of 3D-GCN -""" - -import math -import torch -import torch.nn as nn -import torch.nn.functional as F -from softpool import Periodics -from FNet import FNet -from activations.siren import Siren - - -def get_neighbor_index(vertices: "(bs, vertice_num, 3)", neighbor_num: int): - """ - Return: (bs, vertice_num, neighbor_num) - """ - bs, v, _ = vertices.size() - device = vertices.device - inner = torch.bmm(vertices, vertices.transpose(1, 2)) #(bs, v, v) - quadratic = torch.sum(vertices**2, dim=2) #(bs, v) - distance = inner * (-2) + quadratic.unsqueeze(1) + quadratic.unsqueeze(2) - neighbor_index = torch.topk( - distance, k=neighbor_num + 1, dim=-1, largest=False)[1] - neighbor_index = neighbor_index[:, :, 1:] - return neighbor_index - - -def get_nearest_index(target: "(bs, v1, 3)", source: "(bs, v2, 3)"): - """ - Return: (bs, v1, 1) - """ - inner = torch.bmm(target, source.transpose(1, 2)) #(bs, v1, v2) - s_norm_2 = torch.sum(source**2, dim=2) #(bs, v2) - t_norm_2 = torch.sum(target**2, dim=2) #(bs, v1) - d_norm_2 = s_norm_2.unsqueeze(1) + t_norm_2.unsqueeze(2) - 2 * inner - nearest_index = torch.topk(d_norm_2, k=1, dim=-1, largest=False)[1] - return nearest_index - - -def indexing_neighbor(tensor: "(bs, vertice_num, dim)", - index: "(bs, vertice_num, neighbor_num)"): - """ - Return: (bs, vertice_num, neighbor_num, dim) - """ - bs, v, n = index.size() - id_0 = torch.arange(bs).view(-1, 1, 1) - tensor_indexed = tensor[id_0, index] - return tensor_indexed - - -def get_neighbor_direction_norm( - vertices: "(bs, vertice_num, 3)", - neighbor_index: "(bs, vertice_num, neighbor_num)"): - """ - Return: (bs, vertice_num, neighobr_num, 3) - """ - neighbors = indexing_neighbor(vertices, neighbor_index) # (bs, v, n, 3) - neighbor_direction = neighbors - vertices.unsqueeze(2) - neighbor_direction_norm = F.normalize(neighbor_direction, dim=-1) - return neighbor_direction_norm - - -class Conv_surface(nn.Module): - """Extract structure feafure from surface, independent from vertice coordinates""" - - def __init__(self, kernel_num, support_num): - super().__init__() - self.kernel_num = kernel_num - self.support_num = support_num - - self.relu = nn.ReLU(inplace=True) - self.weights = nn.Parameter( - torch.FloatTensor(1, 1, support_num, kernel_num)) - self.siren = Siren( - dim_in = 3, - dim_out = 128, - w0 = 30, - use_bias = True, - is_first = True - ) - self.directions = nn.Parameter( - torch.FloatTensor(128, support_num * kernel_num)) - self.initialize() - - def initialize(self): - stdv = 1. / math.sqrt(self.support_num * self.kernel_num) - self.weights.data.uniform_(-stdv, stdv) - self.directions.data.uniform_(-stdv, stdv) - - def forward(self, neighbor_index: "(bs, vertice_num, neighbor_num)", - vertices: "(bs, vertice_num, 3)"): - """ - Return vertices with local feature: (bs, vertice_num, kernel_num) - """ - vertices = self.siren(vertices) - bs, vertice_num, neighbor_num = neighbor_index.size() - neighbor_direction_norm = get_neighbor_direction_norm( - vertices, neighbor_index) - support_direction_norm = F.normalize(self.directions, dim=0) #(3, s * k) - theta = neighbor_direction_norm @ support_direction_norm # (bs, vertice_num, neighbor_num, s*k) - - theta = self.relu(theta) - theta = theta.contiguous().view(bs, vertice_num, neighbor_num, - self.support_num, self.kernel_num) - theta = torch.max( - theta, dim=2)[0] * self.weights # (bs, vertice_num, support_num, kernel_num) - feature = torch.sum(theta, dim=2) # (bs, vertice_num, kernel_num) - return feature - - -class Conv_layer(nn.Module): - def __init__(self, in_channel, out_channel, support_num): - super().__init__() - # arguments: - self.in_channel = in_channel - self.out_channel = out_channel - self.support_num = support_num - - # parameters: - self.relu = nn.ReLU(inplace=True) - self.siren = Siren( - dim_in = 3, - dim_out = 32, - w0 = 1, - use_bias = True, - is_first = False - ) - self.weights = nn.Parameter( - torch.FloatTensor(in_channel, (support_num + 1) * out_channel)) - self.bias = nn.Parameter( - torch.FloatTensor((support_num + 1) * out_channel)) - self.directions = nn.Parameter( - torch.FloatTensor(32, support_num * out_channel)) - self.initialize() - - def initialize(self): - stdv = 1. / math.sqrt(self.out_channel * (self.support_num + 1)) - self.weights.data.uniform_(-stdv, stdv) - self.bias.data.uniform_(-stdv, stdv) - self.directions.data.uniform_(-stdv, stdv) - - def forward(self, neighbor_index: "(bs, vertice_num, neighbor_index)", - vertices: "(bs, vertice_num, 3)", - feature_map: "(bs, vertice_num, in_channel)"): - """ - Return: output feature map: (bs, vertice_num, out_channel) - """ - bs, vertice_num, neighbor_num = neighbor_index.size() - vertices = self.siren(vertices) - neighbor_direction_norm = get_neighbor_direction_norm( - vertices, neighbor_index) - support_direction_norm = F.normalize(self.directions, dim=0) - theta = neighbor_direction_norm @ support_direction_norm # (bs, vertice_num, neighbor_num, support_num * out_channel) - theta = self.relu(theta) - theta = theta.contiguous().view(bs, vertice_num, neighbor_num, -1) - # (bs, vertice_num, neighbor_num, support_num * out_channel) - - feature_out = feature_map @ self.weights + self.bias # (bs, vertice_num, (support_num + 1) * out_channel) - feature_center = feature_out[:, :, :self. - out_channel] # (bs, vertice_num, out_channel) - feature_support = feature_out[:, :, self. - out_channel:] #(bs, vertice_num, support_num * out_channel) - - # Fuse together - max among product - feature_support = indexing_neighbor( - feature_support, neighbor_index - ) # (bs, vertice_num, neighbor_num, support_num * out_channel) - activation_support = theta * feature_support # (bs, vertice_num, neighbor_num, support_num * out_channel) - activation_support = activation_support.view( - bs, vertice_num, neighbor_num, self.support_num, self.out_channel) - activation_support = torch.max( - activation_support, - dim=2)[0] # (bs, vertice_num, support_num, out_channel) - activation_support = torch.sum( - activation_support, dim=2) # (bs, vertice_num, out_channel) - feature_fuse = feature_center + activation_support # (bs, vertice_num, out_channel) - return feature_fuse - - -class Pool_layer(nn.Module): - def __init__(self, pooling_rate: int = 4, neighbor_num: int = 4): - super().__init__() - self.pooling_rate = pooling_rate - self.neighbor_num = neighbor_num - - def forward(self, vertices: "(bs, vertice_num, 3)", - feature_map: "(bs, vertice_num, channel_num)"): - """ - Return: - vertices_pool: (bs, pool_vertice_num, 3), - feature_map_pool: (bs, pool_vertice_num, channel_num) - """ - bs, vertice_num, _ = vertices.size() - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - neighbor_feature = indexing_neighbor( - feature_map, - neighbor_index) #(bs, vertice_num, neighbor_num, channel_num) - pooled_feature = torch.max( - neighbor_feature, dim=2)[0] #(bs, vertice_num, channel_num) - - pool_num = int(vertice_num / self.pooling_rate) - sample_idx = torch.randperm(vertice_num)[:pool_num] - vertices_pool = vertices[:, sample_idx, :] # (bs, pool_num, 3) - feature_map_pool = pooled_feature[:, - sample_idx, :] #(bs, pool_num, channel_num) - return vertices_pool, feature_map_pool - - -class GCN3DFeat(nn.Module): - def __init__(self, support_num: int, neighbor_num: int): - super().__init__() - self.neighbor_num = neighbor_num - - self.conv_0 = Conv_surface(kernel_num=32, support_num=support_num) - self.conv_1 = Conv_layer(32, 64, support_num=support_num) - self.pool_1 = Pool_layer(pooling_rate=8, neighbor_num=8) - self.conv_2 = Conv_layer(64, 128, support_num=support_num) - self.conv_3 = Conv_layer(128, 256, support_num=support_num) - self.pool_2 = Pool_layer(pooling_rate=8, neighbor_num=8) - self.conv_4 = Conv_layer(256, 1024, support_num=support_num) - - def forward(self, vertices: "(bs, vertice_num, 3)"): - bs, vertice_num, _ = vertices.size() - - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - fm_0 = self.conv_0(neighbor_index, vertices) - fm_0 = F.relu(fm_0, inplace=True) - fm_1 = self.conv_1(neighbor_index, vertices, fm_0) - fm_1 = F.relu(fm_1, inplace=True) - vertices, fm_1 = self.pool_1(vertices, fm_1) - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - - fm_2 = self.conv_2(neighbor_index, vertices, fm_1) - fm_2 = F.relu(fm_2, inplace=True) - fm_3 = self.conv_3(neighbor_index, vertices, fm_2) - fm_3 = F.relu(fm_3, inplace=True) - vertices, fm_3 = self.pool_2(vertices, fm_3) - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - - fm_4 = self.conv_4(neighbor_index, vertices, fm_3) - feature_global = fm_4.max(1, keepdim=True)[0] - # feature_global = fm_4#.max(1)[0] - return feature_global - - -class GCN3DFeatDeep(nn.Module): - def __init__(self, support_num: int, neighbor_num: int): - super().__init__() - self.neighbor_num = neighbor_num - - self.conv_0 = Conv_surface(kernel_num=32, support_num=support_num) - self.dropping = torch.nn.Dropout(p=0.1, inplace=True) - self.conv_1 = Conv_layer(32, 64, support_num=support_num) - self.pool_1 = Pool_layer(pooling_rate=8, neighbor_num=8) - self.conv_2 = Conv_layer(64, 128, support_num=support_num) - self.conv_3 = Conv_layer(128, 256, support_num=support_num) - self.pool_2 = Pool_layer(pooling_rate=8, neighbor_num=8) - self.conv_4 = Conv_layer(256, 512, support_num=support_num) - self.conv_5 = Conv_layer(512, 512, support_num=support_num) - self.pool_3 = Pool_layer(pooling_rate=8, neighbor_num=8) - self.conv_6 = Conv_layer(512, 1024, support_num=support_num) - - def forward(self, vertices: "(bs, vertice_num, 3)"): - bs, vertice_num, _ = vertices.size() - - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - fm_0 = self.conv_0(neighbor_index, vertices) - fm_0 = self.dropping(fm_0) - fm_0 = F.relu(fm_0, inplace=True) - fm_1 = self.conv_1(neighbor_index, vertices, fm_0) - fm_1 = F.relu(fm_1, inplace=True) - vertices, fm_1 = self.pool_1(vertices, fm_1) - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - - fm_2 = self.conv_2(neighbor_index, vertices, fm_1) - fm_2 = F.relu(fm_2, inplace=True) - fm_3 = self.conv_3(neighbor_index, vertices, fm_2) - fm_3 = F.relu(fm_3, inplace=True) - vertices, fm_3 = self.pool_2(vertices, fm_3) - neighbor_index = get_neighbor_index(vertices, self.neighbor_num) - - fm_4 = self.conv_4(neighbor_index, vertices, fm_3) - fm_4 = F.relu(fm_4, inplace=True) - fm_5 = self.conv_5(neighbor_index, vertices, fm_4) - fm_5 = F.relu(fm_5, inplace=True) - vertices, fm_5 = self.pool_3(vertices, fm_5) - neighbor_index = get_neighbor_index(vertices, 3) - - fm_6 = self.conv_6(neighbor_index, vertices, fm_5) - # feature_global = fm_6[:,:2,:].max(1, keepdim=True)[0] - feature_global = fm_6.max(1, keepdim=True)[0] - return feature_global - - -def test(): - import time - bs = 8 - v = 1024 - dim = 3 - n = 20 - vertices = torch.randn(bs, v, dim) - neighbor_index = get_neighbor_index(vertices, n) - - s = 3 - conv_1 = Conv_surface(kernel_num=32, support_num=s) - conv_2 = Conv_layer(in_channel=32, out_channel=64, support_num=s) - pool = Pool_layer(pooling_rate=4, neighbor_num=4) - - print("Input size: {}".format(vertices.size())) - start = time.time() - f1 = conv_1(neighbor_index, vertices) - print("\n[1] Time: {}".format(time.time() - start)) - print("[1] Out shape: {}".format(f1.size())) - start = time.time() - f2 = conv_2(neighbor_index, vertices, f1) - print("\n[2] Time: {}".format(time.time() - start)) - print("[2] Out shape: {}".format(f2.size())) - start = time.time() - v_pool, f_pool = pool(vertices, f2) - print("\n[3] Time: {}".format(time.time() - start)) - print("[3] v shape: {}, f shape: {}".format(v_pool.size(), f_pool.size())) - - -if __name__ == "__main__": - test() diff --git a/other_models/displace/tree.py b/other_models/displace/tree.py deleted file mode 100644 index eacd39c..0000000 --- a/other_models/displace/tree.py +++ /dev/null @@ -1,68 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.init as init -import math - -class TreeGCN(nn.Module): - def __init__(self, depth, features, degrees, support=10, node=1, upsample=False, activation=True): - self.depth = depth - self.in_feature = features[depth] - self.out_feature = features[depth+1] - self.node = node - self.degree = degrees[depth] - self.upsample = upsample - self.activation = activation - super(TreeGCN, self).__init__() - - self.W_root = nn.ModuleList([nn.Linear(features[inx], self.out_feature, bias=False) for inx in range(self.depth+1)]) - - if self.upsample: - self.W_branch = nn.Parameter(torch.FloatTensor(self.node, self.in_feature, self.degree*self.in_feature)) - - self.W_loop = nn.Sequential(nn.Linear(self.in_feature, self.in_feature*support, bias=False), - nn.Linear(self.in_feature*support, self.out_feature, bias=False)) - - self.bias = nn.Parameter(torch.FloatTensor(1, self.degree, self.out_feature)) - - self.leaky_relu = nn.LeakyReLU(negative_slope=0.2) - # self.fourier_feat = Periodics(dim_input=self.out_feature, dim_output=self.out_feature, transpose=True, is_first=False) - - self.init_param() - - def init_param(self): - if self.upsample: - init.xavier_uniform_(self.W_branch.data, gain=init.calculate_gain('relu')) - - stdv = 1. / math.sqrt(self.out_feature) - self.bias.data.uniform_(-stdv, stdv) - - def forward(self, tree): - root = 0 - for inx in range(self.depth+1): - root_num = tree[inx].size(1) - repeat_num = int(self.node / root_num) - root_node = self.W_root[inx](tree[inx]) - self.batch = root_node.shape[0] - root = root + root_node.repeat(1,1,repeat_num).view(self.batch,-1,self.out_feature) - - branch = 0 - if self.upsample: - branch = tree[-1].unsqueeze(2) @ self.W_branch - branch = self.leaky_relu(branch) - branch = branch.view(self.batch,self.node*self.degree,self.in_feature) - - branch = self.W_loop(branch) - - branch = root.repeat(1,1,self.degree).view(self.batch,-1,self.out_feature) + branch - else: - branch = self.W_loop(tree[-1]) - - branch = root + branch - - if self.activation: - branch = self.leaky_relu(branch + self.bias.repeat(1,self.node,1)) - # branch = self.fourier_feat(branch) - - tree.append(branch) - - return tree diff --git a/other_models/imgpointr/PoinTr.py b/other_models/imgpointr/PoinTr.py deleted file mode 100644 index bb20796..0000000 --- a/other_models/imgpointr/PoinTr.py +++ /dev/null @@ -1,397 +0,0 @@ -import torch -from torch import nn -import os, sys -import numpy as np -import math - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "VRCNet/utils/pointnet2")) -# from pointnet2_ops import pointnet2_utils -import pointnet2_utils -# from extensions.chamfer_dist import ChamferDistanceL1 -sys.path.append("./distance/chamfer/") -import dist_chamfer as cd -CD = cd.chamferDist() -from .Transformer import PCTransformer -from .build import MODELS -from other_models.VRCNet.utils.model_utils import * - - -def fps(pc, num): - fps_idx = pointnet2_utils.furthest_point_sample(pc, num) - sub_pc = pointnet2_utils.gather_operation(pc.transpose(1, 2).contiguous(), fps_idx).transpose(1,2).contiguous() - return sub_pc - - -class Fold(nn.Module): - def __init__(self, in_channel , step , hidden_dim = 512): - super().__init__() - - self.in_channel = in_channel - self.step = step - - a = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(1, step).expand(step, step).reshape(1, -1) - b = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(step, 1).expand(step, step).reshape(1, -1) - self.folding_seed = torch.cat([a, b], dim=0).cuda() - - self.folding1 = nn.Sequential( - nn.Conv1d(in_channel + 2, hidden_dim, 1), - nn.BatchNorm1d(hidden_dim), - nn.ReLU(inplace=True), - nn.Conv1d(hidden_dim, hidden_dim//2, 1), - nn.BatchNorm1d(hidden_dim//2), - nn.ReLU(inplace=True), - nn.Conv1d(hidden_dim//2, 3, 1), - ) - - self.folding2 = nn.Sequential( - nn.Conv1d(in_channel + 3, hidden_dim, 1), - nn.BatchNorm1d(hidden_dim), - nn.ReLU(inplace=True), - nn.Conv1d(hidden_dim, hidden_dim//2, 1), - nn.BatchNorm1d(hidden_dim//2), - nn.ReLU(inplace=True), - nn.Conv1d(hidden_dim//2, 3, 1), - ) - - def forward(self, x): - num_sample = self.step * self.step - bs = x.size(0) - features = x.view(bs, self.in_channel, 1).expand(bs, self.in_channel, num_sample) - seed = self.folding_seed.view(1, 2, num_sample).expand(bs, 2, num_sample).to(x.device) - - x = torch.cat([seed, features], dim=1) - fd1 = self.folding1(x) - x = torch.cat([fd1, features], dim=1) - fd2 = self.folding2(x) - - return fd2 - -@MODELS.register_module() -class PoinTr(nn.Module): - def __init__(self, config, **kwargs): - super().__init__() - self.refiner = config.refiner - self.trans_dim = config.trans_dim - self.knn_layer = config.knn_layer - self.num_pred = config.num_pred - self.num_query = config.num_query - - self.fold_step = int(pow(self.num_pred//self.num_query, 0.5) + 0.5) - self.base_model = PCTransformer(in_chans = 3, embed_dim = self.trans_dim, depth = [6, 8], drop_rate = 0., num_query = self.num_query, knn_layer = self.knn_layer) - - self.increase_dim = nn.Sequential( - nn.Conv1d(self.trans_dim, 1024, 1), - nn.BatchNorm1d(1024), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(1024, 1024, 1) - ) - self.reduce_map = nn.Linear(self.trans_dim + 1027, self.trans_dim) - self.build_loss_func() - - self.foldingnet = Fold(self.trans_dim, step = self.fold_step, hidden_dim = 256) # rebuild a cluster point - - def build_loss_func(self): - # self.loss_func = ChamferDistanceL1() - self.loss_func = cd.chamferDist() - - def get_loss(self, ret, gt): - loss_coarse = self.loss_func(ret[0], gt) - loss_fine = self.loss_func(ret[1], gt) - return loss_coarse, loss_fine - - def forward(self, imgs_feat): - q, coarse_point_cloud = self.base_model(imgs_feat) # B M C and B M 3 - - # cat the input - """ - inp_sparse = fps(xyz, self.num_query) - querries_all = torch.cat([coarse_point_cloud, inp_sparse], dim=1).contiguous() - """ - querries_all = coarse_point_cloud - - B, M ,C = q.shape - - global_feature = self.increase_dim(q.transpose(1,2)).transpose(1,2) # B M 1024 - global_feature = torch.max(global_feature, dim=1)[0] # B 1024 - - rebuild_feature = torch.cat([ - global_feature.unsqueeze(-2).expand(-1, M, -1), - q, - coarse_point_cloud], dim=-1) # B M 1027 + C - - rebuild_feature = self.reduce_map(rebuild_feature.reshape(B*M, -1)) # BM C - - # NOTE: foldingNet - relative_xyz = self.foldingnet(rebuild_feature).reshape(B, M, 3, -1) # B M 3 S - - rebuild_points = (relative_xyz + coarse_point_cloud.unsqueeze(-1)).transpose(2,3).reshape(B, -1, 3) # B N 3 - # rebuild_points = torch.cat([rebuild_points, fps(xyz, self.num_pred - M * self.fold_step**2)],dim=1).contiguous() - ret = (querries_all, rebuild_points) - - return ret - -class Folding(nn.Module): - def __init__(self, input_size, output_size, step_ratio, global_feature_size=1024, num_models=1): - super(Folding, self).__init__() - self.input_size = input_size - self.output_size = output_size - self.step_ratio = step_ratio - self.num_models = num_models - - self.conv = nn.Conv1d(input_size + global_feature_size + 2, output_size, 1, bias=True) - - sqrted = int(math.sqrt(step_ratio)) + 1 - for i in range(1, sqrted + 1).__reversed__(): - if (step_ratio % i) == 0: - num_x = i - num_y = step_ratio // i - break - - grid_x = torch.linspace(-0.2, 0.2, steps=num_x) - grid_y = torch.linspace(-0.2, 0.2, steps=num_y) - - x, y = torch.meshgrid(grid_x, grid_y) # x, y shape: (2, 1) - self.grid = torch.stack([x, y], dim=-1).view(-1, 2) # (2, 2) - - def forward(self, point_feat, global_feat): - batch_size, num_features, num_points = point_feat.size() - point_feat = point_feat.transpose(1, 2).contiguous().unsqueeze(2).repeat(1, 1, self.step_ratio, 1).view( - batch_size, - -1, num_features).transpose(1, 2).contiguous() - global_feat = global_feat.unsqueeze(2).repeat(1, 1, num_points * self.step_ratio).repeat(self.num_models, 1, 1) - grid_feat = self.grid.unsqueeze(0).repeat(batch_size, num_points, 1).transpose(1, 2).contiguous().cuda() - features = torch.cat([global_feat, point_feat, grid_feat], axis=1) - features = F.relu(self.conv(features)) - return features - - -class SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=16): - super(SA_module, self).__init__() - self.share_planes = share_planes - self.k = k - self.conv1 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv2 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv3 = nn.Conv2d(in_planes, mid_planes, kernel_size=1) - - self.conv_w = nn.Sequential(nn.ReLU(inplace=False), - nn.Conv2d(rel_planes * (k + 1), mid_planes // share_planes, kernel_size=1, - bias=False), - nn.ReLU(inplace=False), - nn.Conv2d(mid_planes // share_planes, k * mid_planes // share_planes, - kernel_size=1)) - self.activation_fn = nn.ReLU(inplace=False) - - self.conv_out = nn.Conv2d(mid_planes, out_planes, kernel_size=1) - - def forward(self, input): - x, idx = input - batch_size, _, _, num_points = x.size() - identity = x # B C 1 N - x = self.activation_fn(x) - xn = get_edge_features(x, idx) # B C K N - x1, x2, x3 = self.conv1(x), self.conv2(xn), self.conv3(xn) - - x2 = x2.view(batch_size, -1, 1, num_points).contiguous() # B kC 1 N - w = self.conv_w(torch.cat([x1, x2], 1)).view(batch_size, -1, self.k, num_points) - w = w.repeat(1, self.share_planes, 1, 1) - out = w * x3 - out = torch.sum(out, dim=2, keepdim=True) - - out = self.activation_fn(out) - out = self.conv_out(out) # B C 1 N - out += identity - return [out, idx] - - -class SK_SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=[10, 20], r=2, L=32): - super(SK_SA_module, self).__init__() - - self.num_kernels = len(k) - d = max(int(out_planes / r), L) - - self.sams = nn.ModuleList([]) - - for i in range(len(k)): - self.sams.append(SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k[i])) - - self.fc = nn.Linear(out_planes, d) - self.fcs = nn.ModuleList([]) - - for i in range(len(k)): - self.fcs.append(nn.Linear(d, out_planes)) - - self.softmax = nn.Softmax(dim=1) - self.af = nn.ReLU(inplace=False) - - def forward(self, input): - x, idxs = input - assert (self.num_kernels == len(idxs)) - for i, sam in enumerate(self.sams): - fea, _ = sam([x, idxs[i]]) - fea = self.af(fea) - fea = fea.unsqueeze(dim=1) - if i == 0: - feas = fea - else: - feas = torch.cat([feas, fea], dim=1) - - fea_U = torch.sum(feas, dim=1) - fea_s = fea_U.mean(-1).mean(-1) - fea_z = self.fc(fea_s) - - for i, fc in enumerate(self.fcs): - vector = fc(fea_z).unsqueeze_(dim=1) - if i == 0: - attention_vectors = vector - else: - attention_vectors = torch.cat([attention_vectors, vector], dim=1) - - attention_vectors = self.softmax(attention_vectors) - attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1) - fea_v = (feas * attention_vectors).sum(dim=1) - return [fea_v, idxs] - - -class SKN_Res_unit(nn.Module): - def __init__(self, input_size, output_size, k=[10, 20], layers=1): - super(SKN_Res_unit, self).__init__() - self.conv1 = nn.Conv2d(input_size, output_size, 1, bias=False) - self.sam = self._make_layer(output_size, output_size // 16, output_size // 4, output_size, int(layers), 8, k=k) - self.conv2 = nn.Conv2d(output_size, output_size, 1, bias=False) - self.conv_res = nn.Conv2d(input_size, output_size, 1, bias=False) - self.af = nn.ReLU(inplace=False) - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def forward(self, feat, idx): - x, _ = self.sam([self.conv1(feat), idx]) - x = self.conv2(self.af(x)) - return x + self.conv_res(feat) - - -class SA_SKN_Res_encoder(nn.Module): - def __init__(self, input_size=3, k=[10, 20], pk=16, output_size=64, layers=[2, 2, 2, 2], - pts_num=[3072, 1536, 768, 384]): - super(SA_SKN_Res_encoder, self).__init__() - self.init_channel = 64 - - c1 = self.init_channel - self.sam_res1 = SKN_Res_unit(input_size, c1, k, int(layers[0])) - - c2 = c1 * 2 - self.sam_res2 = SKN_Res_unit(c2, c2, k, int(layers[1])) - - c3 = c2 * 2 - self.sam_res3 = SKN_Res_unit(c3, c3, k, int(layers[2])) - - c4 = c3 * 2 - self.sam_res4 = SKN_Res_unit(c4, c4, k, int(layers[3])) - - self.conv5 = nn.Conv2d(c4, 1024, 1) - - self.fc1 = nn.Linear(1024, 512) - self.fc2 = nn.Linear(512, 1024) - - self.conv6 = nn.Conv2d(c4 + 1024, c4, 1) - self.conv7 = nn.Conv2d(c3 + c4, c3, 1) - self.conv8 = nn.Conv2d(c2 + c3, c2, 1) - self.conv9 = nn.Conv2d(c1 + c2, c1, 1) - - self.conv_out = nn.Conv2d(c1, output_size, 1) - self.dropout = nn.Dropout() - self.af = nn.ReLU(inplace=False) - self.k = k - self.pk = pk - self.rate = 2 - - self.pts_num = pts_num - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=[16]): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def _edge_pooling(self, features, points, rate=2, k=16, sample_num=None): - features = features.squeeze(2) - - if sample_num is None: - input_points_num = int(features.size()[2]) - sample_num = input_points_num // rate - - ds_features, p_idx, pn_idx, ds_points = edge_preserve_sampling(features, points, sample_num, k) - ds_features = ds_features.unsqueeze(2) - return ds_features, p_idx, pn_idx, ds_points - - def _edge_unpooling(self, features, src_pts, tgt_pts): - features = features.squeeze(2) - idx, weight = three_nn_upsampling(tgt_pts, src_pts) - features = pn2.three_interpolate(features, idx, weight) - features = features.unsqueeze(2) - return features - - def forward(self, features): - batch_size, _, num_points = features.size() - pt1 = features[:, 0:3, :] - - idx1 = [] - for i in range(len(self.k)): - idx = knn(pt1, self.k[i]) - idx1.append(idx) - - pt1 = pt1.transpose(1, 2).contiguous() - - x = features.unsqueeze(2) - x = self.sam_res1(x, idx1) - x1 = self.af(x) - - x, _, _, pt2 = self._edge_pooling(x1, pt1, self.rate, self.pk, self.pts_num[1]) - idx2 = [] - for i in range(len(self.k)): - idx = knn(pt2.transpose(1, 2).contiguous(), self.k[i]) - idx2.append(idx) - - x = self.sam_res2(x, idx2) - x2 = self.af(x) - - x, _, _, pt3 = self._edge_pooling(x2, pt2, self.rate, self.pk, self.pts_num[2]) - idx3 = [] - for i in range(len(self.k)): - idx = knn(pt3.transpose(1, 2).contiguous(), self.k[i]) - idx3.append(idx) - - x = self.sam_res3(x, idx3) - x3 = self.af(x) - - x, _, _, pt4 = self._edge_pooling(x3, pt3, self.rate, self.pk, self.pts_num[3]) - idx4 = [] - for i in range(len(self.k)): - idx = knn(pt4.transpose(1, 2).contiguous(), self.k[i]) - idx4.append(idx) - - x = self.sam_res4(x, idx4) - x4 = self.af(x) - x = self.conv5(x4) - x, _ = torch.max(x, -1) - x = x.view(batch_size, -1) - x = self.dropout(self.af(self.fc2(self.dropout(self.af(self.fc1(x)))))) - - x = x.unsqueeze(2).repeat(1, 1, self.pts_num[3]).unsqueeze(2) - x = self.af(self.conv6(torch.cat([x, x4], 1))) - x = self._edge_unpooling(x, pt4, pt3) - x = self.af(self.conv7(torch.cat([x, x3], 1))) - x = self._edge_unpooling(x, pt3, pt2) - x = self.af(self.conv8(torch.cat([x, x2], 1))) - x = self._edge_unpooling(x, pt2, pt1) - x = self.af(self.conv9(torch.cat([x, x1], 1))) - x = self.conv_out(x) - x = x.squeeze(2) - return x - diff --git a/other_models/imgpointr/PoinTr.yaml b/other_models/imgpointr/PoinTr.yaml deleted file mode 100644 index 8b714aa..0000000 --- a/other_models/imgpointr/PoinTr.yaml +++ /dev/null @@ -1,41 +0,0 @@ -optimizer : { - type: AdamW, - kwargs: { - lr : 0.0005, - weight_decay : 0.0005 -}} - -scheduler: { - type: LambdaLR, - kwargs: { - decay_step: 21, - lr_decay: 0.9, - lowest_decay: 0.02 # min lr = lowest_decay * lr -}} - -bnmscheduler: { - type: Lambda, - kwargs: { - decay_step: 21, - bn_decay: 0.5, - bn_momentum: 0.9, - lowest_decay: 0.01 -}} - -dataset : { - train : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'train'}}, - val : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'test'}}, - test : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'test'}}} -# model : { -# NAME: PoinTr, num_pred: 8192, num_query: 224, knn_layer: 1, trans_dim: 256, refiner: 'foldingnet'} -model : { - NAME: PoinTr, num_pred: 8192, num_query: 1024, knn_layer: 1, trans_dim: 256, refiner: 'foldingnet'} - -total_bs : 48 -step_per_update : 1 -max_epoch : 300 - -consider_metric: CDL1 diff --git a/other_models/imgpointr/Transformer.py b/other_models/imgpointr/Transformer.py deleted file mode 100644 index ea916cb..0000000 --- a/other_models/imgpointr/Transformer.py +++ /dev/null @@ -1,399 +0,0 @@ -import torch -import torch.nn as nn - -from timm.models.layers import DropPath,trunc_normal_ - -from .dgcnn_group import DGCNN_Grouper -from .logger import * -import numpy as np -from knn_cuda import KNN -knn = KNN(k=8, transpose_mode=False) - -def get_knn_index(coor_q, coor_k=None): - coor_k = coor_k if coor_k is not None else coor_q - # coor: bs, 3, np - batch_size, _, num_points = coor_q.size() - num_points_k = coor_k.size(2) - - with torch.no_grad(): - _, idx = knn(coor_k, coor_q) # bs k np - idx_base = torch.arange(0, batch_size, device=coor_q.device).view(-1, 1, 1) * num_points_k - idx = idx + idx_base - idx = idx.view(-1) - - return idx # bs*k*np - -def get_graph_feature(x, knn_index, x_q=None): - - #x: bs, np, c, knn_index: bs*k*np - k = 8 - batch_size, num_points, num_dims = x.size() - num_query = x_q.size(1) if x_q is not None else num_points - feature = x.view(batch_size * num_points, num_dims)[knn_index, :] - feature = feature.view(batch_size, k, num_query, num_dims) - x = x_q if x_q is not None else x - x = x.view(batch_size, 1, num_query, num_dims).expand(-1, k, -1, -1) - feature = torch.cat((feature - x, x), dim=-1) - return feature # b k np c - -class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = act_layer() - self.fc2 = nn.Linear(hidden_features, out_features) - self.drop = nn.Dropout(drop) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - x = self.fc2(x) - x = self.drop(x) - return x - - -class Attention(nn.Module): - def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): - super().__init__() - self.num_heads = num_heads - head_dim = dim // num_heads - # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights - self.scale = qk_scale or head_dim ** -0.5 - - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.attn_drop = nn.Dropout(attn_drop) - self.proj = nn.Linear(dim, dim) - self.proj_drop = nn.Dropout(proj_drop) - - def forward(self, x): - B, N, C = x.shape - qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) - q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) - - attn = (q @ k.transpose(-2, -1)) * self.scale - attn = attn.softmax(dim=-1) - attn = self.attn_drop(attn) - - x = (attn @ v).transpose(1, 2).reshape(B, N, C) - x = self.proj(x) - x = self.proj_drop(x) - # with a shape of 64 x 256 - return x - - - -class CrossAttention(nn.Module): - def __init__(self, dim, out_dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): - super().__init__() - self.num_heads = num_heads - self.dim = dim - self.out_dim = out_dim - head_dim = out_dim // num_heads - # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights - self.scale = qk_scale or head_dim ** -0.5 - - self.q_map = nn.Linear(dim, out_dim, bias=qkv_bias) - self.k_map = nn.Linear(dim, out_dim, bias=qkv_bias) - self.v_map = nn.Linear(dim, out_dim, bias=qkv_bias) - self.attn_drop = nn.Dropout(attn_drop) - - self.proj = nn.Linear(out_dim, out_dim) - self.proj_drop = nn.Dropout(proj_drop) - - def forward(self, q, v): - B, N, _ = q.shape - C = self.out_dim - k = v - NK = k.size(1) - - q = self.q_map(q).view(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) - k = self.k_map(k).view(B, NK, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) - v = self.v_map(v).view(B, NK, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) - - attn = (q @ k.transpose(-2, -1)) * self.scale - attn = attn.softmax(dim=-1) - attn = self.attn_drop(attn) - - x = (attn @ v).transpose(1, 2).reshape(B, N, C) - x = self.proj(x) - x = self.proj_drop(x) - return x - - -class DecoderBlock(nn.Module): - def __init__(self, dim, num_heads, dim_q = None, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): - super().__init__() - self.norm1 = norm_layer(dim) - self.self_attn = Attention( - dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - dim_q = dim_q or dim - self.norm_q = norm_layer(dim_q) - self.norm_v = norm_layer(dim) - self.attn = CrossAttention( - dim, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) - - self.knn_map = nn.Sequential( - nn.Linear(dim * 2, dim), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.merge_map = nn.Linear(dim*2, dim) - - self.knn_map_cross = nn.Sequential( - nn.Linear(dim * 2, dim), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.merge_map_cross = nn.Linear(dim*2, dim) - - def forward(self, q, v, self_knn_index=None, cross_knn_index=None): - # q = q + self.drop_path(self.self_attn(self.norm1(q))) - norm_q = self.norm1(q) - q_1 = self.self_attn(norm_q) - - if self_knn_index is not None: - knn_f = get_graph_feature(norm_q, self_knn_index) - knn_f = self.knn_map(knn_f) - knn_f = knn_f.max(dim=1, keepdim=False)[0] - q_1 = torch.cat([q_1, knn_f], dim=-1) - q_1 = self.merge_map(q_1) - - q = q + self.drop_path(q_1) - - norm_q = self.norm_q(q) - norm_v = self.norm_v(v) - q_2 = self.attn(norm_q, norm_v) - - if cross_knn_index is not None: - knn_f = get_graph_feature(norm_v, cross_knn_index, norm_q) - knn_f = self.knn_map_cross(knn_f) - knn_f = knn_f.max(dim=1, keepdim=False)[0] - q_2 = torch.cat([q_2, knn_f], dim=-1) - q_2 = self.merge_map_cross(q_2) - - q = q + self.drop_path(q_2) - - # q = q + self.drop_path(self.attn(self.norm_q(q), self.norm_v(v))) - q = q + self.drop_path(self.mlp(self.norm2(q))) - return q - - -class Block(nn.Module): - - def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): - super().__init__() - self.norm1 = norm_layer(dim) - self.attn = Attention( - dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - - self.knn_map = nn.Sequential( - nn.Linear(dim * 2, dim), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.merge_map = nn.Linear(dim*2, dim) - - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) - - def forward(self, x, knn_index = None): - # x = x + self.drop_path(self.attn(self.norm1(x))) - norm_x = self.norm1(x) - x_1 = self.attn(norm_x) - - if knn_index is not None: - knn_f = get_graph_feature(norm_x, knn_index) - knn_f = self.knn_map(knn_f) - knn_f = knn_f.max(dim=1, keepdim=False)[0] - x_1 = torch.cat([x_1, knn_f], dim=-1) - x_1 = self.merge_map(x_1) - - x = x + self.drop_path(x_1) - x = x + self.drop_path(self.mlp(self.norm2(x))) - return x - - - -class PCTransformer(nn.Module): - """ Vision Transformer with support for point cloud completion - """ - def __init__(self, in_chans=3, embed_dim=768, depth=[6, 6], num_heads=8, mlp_ratio=2., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., - num_query = 224, knn_layer = -1): - super().__init__() - - self.num_features = self.embed_dim = embed_dim - - self.knn_layer = knn_layer - - print_log(' Transformer with knn_layer %d' % self.knn_layer, logger='MODEL') - - self.grouper = DGCNN_Grouper() # B 3 N to B C(3) N(128) and B C(128) N(128) - - self.pos_embed = nn.Sequential( - nn.Conv1d(in_chans, 128, 1), - nn.BatchNorm1d(128), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(128, embed_dim, 1) - ) - # self.pos_embed_wave = nn.Sequential( - # nn.Conv1d(60, 128, 1), - # nn.BatchNorm1d(128), - # nn.LeakyReLU(negative_slope=0.2), - # nn.Conv1d(128, embed_dim, 1) - # ) - - # self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) - # self.cls_pos = nn.Parameter(torch.zeros(1, 1, embed_dim)) - self.input_proj = nn.Sequential( - nn.Conv1d(128, embed_dim, 1), - nn.BatchNorm1d(embed_dim), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(embed_dim, embed_dim, 1) - ) - - self.encoder = nn.ModuleList([ - Block( - dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate) - for i in range(depth[0])]) - - # self.increase_dim = nn.Sequential( - # nn.Linear(embed_dim,1024), - # nn.ReLU(inplace=True), - # nn.Linear(1024, 1024) - # ) - - self.increase_dim = nn.Sequential( - nn.Conv1d(embed_dim, 1024, 1), - nn.BatchNorm1d(1024), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(1024, 1024, 1) - ) - - self.num_query = num_query - self.coarse_pred = nn.Sequential( - nn.Linear(1024, 1024), - nn.ReLU(inplace=True), - nn.Linear(1024, 3 * num_query) - ) - self.mlp_query = nn.Sequential( - nn.Conv1d(1024 + 3, 1024, 1), - # nn.BatchNorm1d(1024), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(1024, 1024, 1), - # nn.BatchNorm1d(1024), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(1024, embed_dim, 1) - ) - - self.decoder = nn.ModuleList([ - DecoderBlock( - dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate) - for i in range(depth[1])]) - - # trunc_normal_(self.cls_token, std=.02) - # trunc_normal_(self.cls_pos, std=.02) - self.apply(self._init_weights) - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) - if isinstance(m, nn.Linear) and m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.LayerNorm): - nn.init.constant_(m.bias, 0) - nn.init.constant_(m.weight, 1.0) - elif isinstance(m, nn.Conv1d): - nn.init.xavier_normal_(m.weight.data, gain=1) - elif isinstance(m, nn.BatchNorm1d): - nn.init.constant_(m.weight.data, 1) - nn.init.constant_(m.bias.data, 0) - - def pos_encoding_sin_wave(self, coor): - # ref to https://arxiv.org/pdf/2003.08934v2.pdf - D = 64 # - # normal the coor into [-1, 1], batch wise - normal_coor = 2 * ((coor - coor.min()) / (coor.max() - coor.min())) - 1 - - # define sin wave freq - freqs = torch.arange(D, dtype=torch.float).cuda() - freqs = np.pi * (2**freqs) - - freqs = freqs.view(*[1]*len(normal_coor.shape), -1) # 1 x 1 x 1 x D - normal_coor = normal_coor.unsqueeze(-1) # B x 3 x N x 1 - k = normal_coor * freqs # B x 3 x N x D - s = torch.sin(k) # B x 3 x N x D - c = torch.cos(k) # B x 3 x N x D - x = torch.cat([s,c], -1) # B x 3 x N x 2D - pos = x.transpose(-1,-2).reshape(coor.shape[0], -1, coor.shape[-1]) # B 6D N - # zero_pad = torch.zeros(x.size(0), 2, x.size(-1)).cuda() - # pos = torch.cat([x, zero_pad], dim = 1) - # pos = self.pos_embed_wave(x) - return pos - - def forward(self, imgs_feat): - # build point proxy - bs = imgs_feat.size(0) - # NOTE: try to use a sin wave coor B 3 N, change the pos_embed input dim - """ - pos = self.pos_encoding_sin_wave(coor).transpose(1,2) - # pos = self.pos_embed(coor).transpose(1,2) - x = self.input_proj(f).transpose(1,2) - """ - # NOTE: YIDA modfied following parts using image features - x = imgs_feat - # cls_pos = self.cls_pos.expand(bs, -1, -1) - # cls_token = self.cls_pos.expand(bs, -1, -1) - # x = torch.cat([cls_token, x], dim=1) - # pos = torch.cat([cls_pos, pos], dim=1) - # encoder - """ - for i, blk in enumerate(self.encoder): - if i < self.knn_layer: - x = blk(x + pos, knn_index) # B N C - else: - x = blk(x + pos) - """ - # NOTE: YIDA modfied following parts using image features - for i, blk in enumerate(self.encoder): - x = blk(x) - # build the query feature for decoder - # global_feature = x[:, 0] # B C - - global_feature = self.increase_dim(x.transpose(1,2)) # B 1024 N - global_feature = torch.max(global_feature, dim=-1)[0] # B 1024 - - coarse_point_cloud = self.coarse_pred(global_feature).reshape(bs, -1, 3) # B M C(3) - - new_knn_index = get_knn_index(coarse_point_cloud.transpose(1, 2).contiguous()) - """ - cross_knn_index = get_knn_index(coor_k=coor, coor_q=coarse_point_cloud.transpose(1, 2).contiguous()) - """ - - query_feature = torch.cat([ - global_feature.unsqueeze(1).expand(-1, self.num_query, -1), - coarse_point_cloud], dim=-1) # B M C+3 - q = self.mlp_query(query_feature.transpose(1,2)).transpose(1,2) # B M C - # decoder - for i, blk in enumerate(self.decoder): - if i < self.knn_layer: - q = blk(q, x, new_knn_index) # B M C - else: - q = blk(q, x) - - return q, coarse_point_cloud - diff --git a/other_models/imgpointr/__pycache__/PoinTr.cpython-36.pyc b/other_models/imgpointr/__pycache__/PoinTr.cpython-36.pyc deleted file mode 100644 index 9ddc167..0000000 Binary files a/other_models/imgpointr/__pycache__/PoinTr.cpython-36.pyc and /dev/null differ diff --git a/other_models/imgpointr/__pycache__/Transformer.cpython-36.pyc b/other_models/imgpointr/__pycache__/Transformer.cpython-36.pyc deleted file mode 100644 index 0aaf423..0000000 Binary files a/other_models/imgpointr/__pycache__/Transformer.cpython-36.pyc and /dev/null differ diff --git a/other_models/imgpointr/__pycache__/build.cpython-36.pyc b/other_models/imgpointr/__pycache__/build.cpython-36.pyc deleted file mode 100644 index a8b68ff..0000000 Binary files a/other_models/imgpointr/__pycache__/build.cpython-36.pyc and /dev/null differ diff --git a/other_models/imgpointr/__pycache__/config.cpython-36.pyc b/other_models/imgpointr/__pycache__/config.cpython-36.pyc deleted file mode 100644 index 07a770e..0000000 Binary files a/other_models/imgpointr/__pycache__/config.cpython-36.pyc and /dev/null differ diff --git a/other_models/imgpointr/__pycache__/dgcnn_group.cpython-36.pyc b/other_models/imgpointr/__pycache__/dgcnn_group.cpython-36.pyc deleted file mode 100644 index 6d00b39..0000000 Binary files a/other_models/imgpointr/__pycache__/dgcnn_group.cpython-36.pyc and /dev/null differ diff --git a/other_models/imgpointr/__pycache__/logger.cpython-36.pyc b/other_models/imgpointr/__pycache__/logger.cpython-36.pyc deleted file mode 100644 index 77beef0..0000000 Binary files a/other_models/imgpointr/__pycache__/logger.cpython-36.pyc and /dev/null differ diff --git a/other_models/imgpointr/__pycache__/registry.cpython-36.pyc b/other_models/imgpointr/__pycache__/registry.cpython-36.pyc deleted file mode 100644 index 30f4c99..0000000 Binary files a/other_models/imgpointr/__pycache__/registry.cpython-36.pyc and /dev/null differ diff --git a/other_models/imgpointr/build.py b/other_models/imgpointr/build.py deleted file mode 100644 index 12216e2..0000000 --- a/other_models/imgpointr/build.py +++ /dev/null @@ -1,17 +0,0 @@ -from . import registry - - -MODELS = registry.Registry('models') - - -def build_model_from_cfg(cfg, **kwargs): - """ - Build a dataset, defined by `dataset_name`. - Args: - cfg (eDICT): - Returns: - Dataset: a constructed dataset specified by dataset_name. - """ - return MODELS.build(cfg, **kwargs) - - diff --git a/other_models/imgpointr/config.py b/other_models/imgpointr/config.py deleted file mode 100644 index 8c86d76..0000000 --- a/other_models/imgpointr/config.py +++ /dev/null @@ -1,63 +0,0 @@ -import yaml -from easydict import EasyDict -import os -from .logger import print_log - -def log_args_to_file(args, pre='args', logger=None): - for key, val in args.__dict__.items(): - print_log(f'{pre}.{key} : {val}', logger = logger) - -def log_config_to_file(cfg, pre='cfg', logger=None): - for key, val in cfg.items(): - if isinstance(cfg[key], EasyDict): - print_log(f'{pre}.{key} = edict()', logger = logger) - log_config_to_file(cfg[key], pre=pre + '.' + key, logger=logger) - continue - print_log(f'{pre}.{key} : {val}', logger = logger) - -def merge_new_config(config, new_config): - for key, val in new_config.items(): - if not isinstance(val, dict): - if key == '_base_': - with open(new_config['_base_'], 'r') as f: - try: - val = yaml.load(f, Loader=yaml.FullLoader) - except: - val = yaml.load(f) - config[key] = EasyDict() - merge_new_config(config[key], val) - else: - config[key] = val - continue - if key not in config: - config[key] = EasyDict() - merge_new_config(config[key], val) - return config - -def cfg_from_yaml_file(cfg_file): - config = EasyDict() - with open(cfg_file, 'r') as f: - try: - new_config = yaml.load(f, Loader=yaml.FullLoader) - except: - new_config = yaml.load(f) - merge_new_config(config=config, new_config=new_config) - return config - -def get_config(args, logger=None): - if args.resume: - cfg_path = os.path.join(args.experiment_path, 'config.yaml') - if not os.path.exists(cfg_path): - print_log("Failed to resume", logger = logger) - raise FileNotFoundError() - print_log(f'Resume yaml from {cfg_path}', logger = logger) - args.config = cfg_path - config = cfg_from_yaml_file(args.config) - if not args.resume and args.local_rank == 0: - save_experiment_config(args, config, logger) - return config - -def save_experiment_config(args, config, logger = None): - config_path = os.path.join(args.experiment_path, 'config.yaml') - os.system('cp %s %s' % (args.config, config_path)) - print_log(f'Copy the Config file from {args.config} to {config_path}',logger = logger ) \ No newline at end of file diff --git a/other_models/imgpointr/dgcnn_group.py b/other_models/imgpointr/dgcnn_group.py deleted file mode 100644 index ebf8e90..0000000 --- a/other_models/imgpointr/dgcnn_group.py +++ /dev/null @@ -1,111 +0,0 @@ -import torch -from torch import nn -import sys, os -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "VRCNet/utils/pointnet2")) -# from pointnet2_ops import pointnet2_utils -import pointnet2_utils -from knn_cuda import KNN -knn = KNN(k=16, transpose_mode=False) - - -class DGCNN_Grouper(nn.Module): - def __init__(self): - super().__init__() - ''' - K has to be 16 - ''' - self.input_trans = nn.Conv1d(3, 8, 1) - - self.layer1 = nn.Sequential(nn.Conv2d(16, 32, kernel_size=1, bias=False), - nn.GroupNorm(4, 32), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.layer2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=1, bias=False), - nn.GroupNorm(4, 64), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.layer3 = nn.Sequential(nn.Conv2d(128, 64, kernel_size=1, bias=False), - nn.GroupNorm(4, 64), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.layer4 = nn.Sequential(nn.Conv2d(128, 128, kernel_size=1, bias=False), - nn.GroupNorm(4, 128), - nn.LeakyReLU(negative_slope=0.2) - ) - - - @staticmethod - def fps_downsample(coor, x, num_group): - xyz = coor.transpose(1, 2).contiguous() # b, n, 3 - fps_idx = pointnet2_utils.furthest_point_sample(xyz, num_group) - - combined_x = torch.cat([coor, x], dim=1) - - new_combined_x = ( - pointnet2_utils.gather_operation( - combined_x, fps_idx - ) - ) - - new_coor = new_combined_x[:, :3] - new_x = new_combined_x[:, 3:] - - return new_coor, new_x - - @staticmethod - def get_graph_feature(coor_q, x_q, coor_k, x_k): - - # coor: bs, 3, np, x: bs, c, np - - k = 16 - batch_size = x_k.size(0) - num_points_k = x_k.size(2) - num_points_q = x_q.size(2) - - with torch.no_grad(): - _, idx = knn(coor_k, coor_q) # bs k np - assert idx.shape[1] == k - idx_base = torch.arange(0, batch_size, device=x_q.device).view(-1, 1, 1) * num_points_k - idx = idx + idx_base - idx = idx.view(-1) - num_dims = x_k.size(1) - x_k = x_k.transpose(2, 1).contiguous() - feature = x_k.view(batch_size * num_points_k, -1)[idx, :] - feature = feature.view(batch_size, k, num_points_q, num_dims).permute(0, 3, 2, 1).contiguous() - x_q = x_q.view(batch_size, num_dims, num_points_q, 1).expand(-1, -1, -1, k) - feature = torch.cat((feature - x_q, x_q), dim=1) - return feature - - def forward(self, x): - - # x: bs, 3, np - - # bs 3 N(128) bs C(224)128 N(128) - coor = x - f = self.input_trans(x) - - f = self.get_graph_feature(coor, f, coor, f) - f = self.layer1(f) - f = f.max(dim=-1, keepdim=False)[0] - - coor_q, f_q = self.fps_downsample(coor, f, 512) - f = self.get_graph_feature(coor_q, f_q, coor, f) - f = self.layer2(f) - f = f.max(dim=-1, keepdim=False)[0] - coor = coor_q - - f = self.get_graph_feature(coor, f, coor, f) - f = self.layer3(f) - f = f.max(dim=-1, keepdim=False)[0] - - coor_q, f_q = self.fps_downsample(coor, f, 128) - f = self.get_graph_feature(coor_q, f_q, coor, f) - f = self.layer4(f) - f = f.max(dim=-1, keepdim=False)[0] - coor = coor_q - - return coor, f diff --git a/other_models/imgpointr/ftrans.yaml b/other_models/imgpointr/ftrans.yaml deleted file mode 100644 index 69565fc..0000000 --- a/other_models/imgpointr/ftrans.yaml +++ /dev/null @@ -1,39 +0,0 @@ -optimizer : { - type: AdamW, - kwargs: { - lr : 0.0005, - weight_decay : 0.0005 -}} - -scheduler: { - type: LambdaLR, - kwargs: { - decay_step: 21, - lr_decay: 0.9, - lowest_decay: 0.02 # min lr = lowest_decay * lr -}} - -bnmscheduler: { - type: Lambda, - kwargs: { - decay_step: 21, - bn_decay: 0.5, - bn_momentum: 0.9, - lowest_decay: 0.01 -}} - -dataset : { - train : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'train'}}, - val : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'test'}}, - test : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'test'}}} -model : { - NAME: PoinTr, num_pred: 2048, num_query: 224, knn_layer: 1, trans_dim: 384, refiner: 'renet'} - -total_bs : 48 -step_per_update : 1 -max_epoch : 300 - -consider_metric: CDL1 diff --git a/other_models/imgpointr/logger.py b/other_models/imgpointr/logger.py deleted file mode 100644 index 847c1c7..0000000 --- a/other_models/imgpointr/logger.py +++ /dev/null @@ -1,127 +0,0 @@ -import logging -import torch.distributed as dist - -logger_initialized = {} - -def get_root_logger(log_file=None, log_level=logging.INFO, name='main'): - """Get root logger and add a keyword filter to it. - The logger will be initialized if it has not been initialized. By default a - StreamHandler will be added. If `log_file` is specified, a FileHandler will - also be added. The name of the root logger is the top-level package name, - e.g., "mmdet3d". - Args: - log_file (str, optional): File path of log. Defaults to None. - log_level (int, optional): The level of logger. - Defaults to logging.INFO. - name (str, optional): The name of the root logger, also used as a - filter keyword. Defaults to 'mmdet3d'. - Returns: - :obj:`logging.Logger`: The obtained logger - """ - logger = get_logger(name=name, log_file=log_file, log_level=log_level) - # add a logging filter - logging_filter = logging.Filter(name) - logging_filter.filter = lambda record: record.find(name) != -1 - - return logger - - -def get_logger(name, log_file=None, log_level=logging.INFO, file_mode='w'): - """Initialize and get a logger by name. - If the logger has not been initialized, this method will initialize the - logger by adding one or two handlers, otherwise the initialized logger will - be directly returned. During initialization, a StreamHandler will always be - added. If `log_file` is specified and the process rank is 0, a FileHandler - will also be added. - Args: - name (str): Logger name. - log_file (str | None): The log filename. If specified, a FileHandler - will be added to the logger. - log_level (int): The logger level. Note that only the process of - rank 0 is affected, and other processes will set the level to - "Error" thus be silent most of the time. - file_mode (str): The file mode used in opening log file. - Defaults to 'w'. - Returns: - logging.Logger: The expected logger. - """ - logger = logging.getLogger(name) - if name in logger_initialized: - return logger - # handle hierarchical names - # e.g., logger "a" is initialized, then logger "a.b" will skip the - # initialization since it is a child of "a". - for logger_name in logger_initialized: - if name.startswith(logger_name): - return logger - - # handle duplicate logs to the console - # Starting in 1.8.0, PyTorch DDP attaches a StreamHandler (NOTSET) - # to the root logger. As logger.propagate is True by default, this root - # level handler causes logging messages from rank>0 processes to - # unexpectedly show up on the console, creating much unwanted clutter. - # To fix this issue, we set the root logger's StreamHandler, if any, to log - # at the ERROR level. - for handler in logger.root.handlers: - if type(handler) is logging.StreamHandler: - handler.setLevel(logging.ERROR) - - stream_handler = logging.StreamHandler() - handlers = [stream_handler] - - if dist.is_available() and dist.is_initialized(): - rank = dist.get_rank() - else: - rank = 0 - - # only rank 0 will add a FileHandler - if rank == 0 and log_file is not None: - # Here, the default behaviour of the official logger is 'a'. Thus, we - # provide an interface to change the file mode to the default - # behaviour. - file_handler = logging.FileHandler(log_file, file_mode) - handlers.append(file_handler) - - formatter = logging.Formatter( - '%(asctime)s - %(name)s - %(levelname)s - %(message)s') - for handler in handlers: - handler.setFormatter(formatter) - handler.setLevel(log_level) - logger.addHandler(handler) - - if rank == 0: - logger.setLevel(log_level) - else: - logger.setLevel(logging.ERROR) - - logger_initialized[name] = True - - - return logger - - -def print_log(msg, logger=None, level=logging.INFO): - """Print a log message. - Args: - msg (str): The message to be logged. - logger (logging.Logger | str | None): The logger to be used. - Some special loggers are: - - "silent": no message will be printed. - - other str: the logger obtained with `get_root_logger(logger)`. - - None: The `print()` method will be used to print log messages. - level (int): Logging level. Only available when `logger` is a Logger - object or "root". - """ - if logger is None: - print(msg) - elif isinstance(logger, logging.Logger): - logger.log(level, msg) - elif logger == 'silent': - pass - elif isinstance(logger, str): - _logger = get_logger(logger) - _logger.log(level, msg) - else: - raise TypeError( - 'logger should be either a logging.Logger object, str, ' - f'"silent" or None, but got {type(logger)}') \ No newline at end of file diff --git a/other_models/imgpointr/registry.py b/other_models/imgpointr/registry.py deleted file mode 100644 index f103ec6..0000000 --- a/other_models/imgpointr/registry.py +++ /dev/null @@ -1,288 +0,0 @@ -import inspect -import warnings -from functools import partial -from . import config - -class Registry: - """A registry to map strings to classes. - Registered object could be built from registry. - Example: - >>> MODELS = Registry('models') - >>> @MODELS.register_module() - >>> class ResNet: - >>> pass - >>> resnet = MODELS.build(dict(NAME='ResNet')) - Please refer to https://mmcv.readthedocs.io/en/latest/registry.html for - advanced useage. - Args: - name (str): Registry name. - build_func(func, optional): Build function to construct instance from - Registry, func:`build_from_cfg` is used if neither ``parent`` or - ``build_func`` is specified. If ``parent`` is specified and - ``build_func`` is not given, ``build_func`` will be inherited - from ``parent``. Default: None. - parent (Registry, optional): Parent registry. The class registered in - children registry could be built from parent. Default: None. - scope (str, optional): The scope of registry. It is the key to search - for children registry. If not specified, scope will be the name of - the package where class is defined, e.g. mmdet, mmcls, mmseg. - Default: None. - """ - - def __init__(self, name, build_func=None, parent=None, scope=None): - self._name = name - self._module_dict = dict() - self._children = dict() - self._scope = self.infer_scope() if scope is None else scope - - # self.build_func will be set with the following priority: - # 1. build_func - # 2. parent.build_func - # 3. build_from_cfg - if build_func is None: - if parent is not None: - self.build_func = parent.build_func - else: - self.build_func = build_from_cfg - else: - self.build_func = build_func - if parent is not None: - assert isinstance(parent, Registry) - parent._add_children(self) - self.parent = parent - else: - self.parent = None - - def __len__(self): - return len(self._module_dict) - - def __contains__(self, key): - return self.get(key) is not None - - def __repr__(self): - format_str = self.__class__.__name__ + \ - f'(name={self._name}, ' \ - f'items={self._module_dict})' - return format_str - - @staticmethod - def infer_scope(): - """Infer the scope of registry. - The name of the package where registry is defined will be returned. - Example: - # in mmdet/models/backbone/resnet.py - >>> MODELS = Registry('models') - >>> @MODELS.register_module() - >>> class ResNet: - >>> pass - The scope of ``ResNet`` will be ``mmdet``. - Returns: - scope (str): The inferred scope name. - """ - # inspect.stack() trace where this function is called, the index-2 - # indicates the frame where `infer_scope()` is called - filename = inspect.getmodule(inspect.stack()[2][0]).__name__ - split_filename = filename.split('.') - return split_filename[0] - - @staticmethod - def split_scope_key(key): - """Split scope and key. - The first scope will be split from key. - Examples: - >>> Registry.split_scope_key('mmdet.ResNet') - 'mmdet', 'ResNet' - >>> Registry.split_scope_key('ResNet') - None, 'ResNet' - Return: - scope (str, None): The first scope. - key (str): The remaining key. - """ - split_index = key.find('.') - if split_index != -1: - return key[:split_index], key[split_index + 1:] - else: - return None, key - - @property - def name(self): - return self._name - - @property - def scope(self): - return self._scope - - @property - def module_dict(self): - return self._module_dict - - @property - def children(self): - return self._children - - def get(self, key): - """Get the registry record. - Args: - key (str): The class name in string format. - Returns: - class: The corresponding class. - """ - scope, real_key = self.split_scope_key(key) - if scope is None or scope == self._scope: - # get from self - if real_key in self._module_dict: - return self._module_dict[real_key] - else: - # get from self._children - if scope in self._children: - return self._children[scope].get(real_key) - else: - # goto root - parent = self.parent - while parent.parent is not None: - parent = parent.parent - return parent.get(key) - - def build(self, *args, **kwargs): - return self.build_func(*args, **kwargs, registry=self) - - def _add_children(self, registry): - """Add children for a registry. - The ``registry`` will be added as children based on its scope. - The parent registry could build objects from children registry. - Example: - >>> models = Registry('models') - >>> mmdet_models = Registry('models', parent=models) - >>> @mmdet_models.register_module() - >>> class ResNet: - >>> pass - >>> resnet = models.build(dict(NAME='mmdet.ResNet')) - """ - - assert isinstance(registry, Registry) - assert registry.scope is not None - assert registry.scope not in self.children, \ - f'scope {registry.scope} exists in {self.name} registry' - self.children[registry.scope] = registry - - def _register_module(self, module_class, module_name=None, force=False): - if not inspect.isclass(module_class): - raise TypeError('module must be a class, ' - f'but got {type(module_class)}') - - if module_name is None: - module_name = module_class.__name__ - if isinstance(module_name, str): - module_name = [module_name] - for name in module_name: - if not force and name in self._module_dict: - raise KeyError(f'{name} is already registered ' - f'in {self.name}') - self._module_dict[name] = module_class - - def deprecated_register_module(self, cls=None, force=False): - warnings.warn( - 'The old API of register_module(module, force=False) ' - 'is deprecated and will be removed, please use the new API ' - 'register_module(name=None, force=False, module=None) instead.') - if cls is None: - return partial(self.deprecated_register_module, force=force) - self._register_module(cls, force=force) - return cls - - def register_module(self, name=None, force=False, module=None): - """Register a module. - A record will be added to `self._module_dict`, whose key is the class - name or the specified name, and value is the class itself. - It can be used as a decorator or a normal function. - Example: - >>> backbones = Registry('backbone') - >>> @backbones.register_module() - >>> class ResNet: - >>> pass - >>> backbones = Registry('backbone') - >>> @backbones.register_module(name='mnet') - >>> class MobileNet: - >>> pass - >>> backbones = Registry('backbone') - >>> class ResNet: - >>> pass - >>> backbones.register_module(ResNet) - Args: - name (str | None): The module name to be registered. If not - specified, the class name will be used. - force (bool, optional): Whether to override an existing class with - the same name. Default: False. - module (type): Module class to be registered. - """ - if not isinstance(force, bool): - raise TypeError(f'force must be a boolean, but got {type(force)}') - # NOTE: This is a walkaround to be compatible with the old api, - # while it may introduce unexpected bugs. - if isinstance(name, type): - return self.deprecated_register_module(name, force=force) - - # raise the error ahead of time - if not (name is None or isinstance(name, str) or misc.is_seq_of(name, str)): - raise TypeError( - 'name must be either of None, an instance of str or a sequence' - f' of str, but got {type(name)}') - - # use it as a normal method: x.register_module(module=SomeClass) - if module is not None: - self._register_module( - module_class=module, module_name=name, force=force) - return module - - # use it as a decorator: @x.register_module() - def _register(cls): - self._register_module( - module_class=cls, module_name=name, force=force) - return cls - - return _register - - -def build_from_cfg(cfg, registry, default_args=None): - """Build a module from config dict. - Args: - cfg (edict): Config dict. It should at least contain the key "NAME". - registry (:obj:`Registry`): The registry to search the type from. - Returns: - object: The constructed object. - """ - if not isinstance(cfg, dict): - raise TypeError(f'cfg must be a dict, but got {type(cfg)}') - if 'NAME' not in cfg: - if default_args is None or 'NAME' not in default_args: - raise KeyError( - '`cfg` or `default_args` must contain the key "NAME", ' - f'but got {cfg}\n{default_args}') - if not isinstance(registry, Registry): - raise TypeError('registry must be an mmcv.Registry object, ' - f'but got {type(registry)}') - - if not (isinstance(default_args, dict) or default_args is None): - raise TypeError('default_args must be a dict or None, ' - f'but got {type(default_args)}') - - if default_args is not None: - cfg = config.merge_new_config(cfg, default_args) - - obj_type = cfg.get('NAME') - - if isinstance(obj_type, str): - obj_cls = registry.get(obj_type) - if obj_cls is None: - raise KeyError( - f'{obj_type} is not in the {registry.name} registry') - elif inspect.isclass(obj_type): - obj_cls = obj_type - else: - raise TypeError( - f'type must be a str or valid type, but got {type(obj_type)}') - try: - return obj_cls(cfg) - except Exception as e: - # Normal TypeError does not print class name. - raise type(e)(f'{obj_cls.__name__}: {e}') diff --git a/other_models/pix2vox/__pycache__/attention.cpython-36.pyc b/other_models/pix2vox/__pycache__/attention.cpython-36.pyc deleted file mode 100644 index 9cf2541..0000000 Binary files a/other_models/pix2vox/__pycache__/attention.cpython-36.pyc and /dev/null differ diff --git a/other_models/pix2vox/__pycache__/encoder.cpython-36.pyc b/other_models/pix2vox/__pycache__/encoder.cpython-36.pyc deleted file mode 100644 index 4c3e3e1..0000000 Binary files a/other_models/pix2vox/__pycache__/encoder.cpython-36.pyc and /dev/null differ diff --git a/other_models/pix2vox/__pycache__/model.cpython-36.pyc b/other_models/pix2vox/__pycache__/model.cpython-36.pyc deleted file mode 100644 index 061a4a4..0000000 Binary files a/other_models/pix2vox/__pycache__/model.cpython-36.pyc and /dev/null differ diff --git a/other_models/pix2vox/attention.py b/other_models/pix2vox/attention.py deleted file mode 100644 index 905fd7e..0000000 --- a/other_models/pix2vox/attention.py +++ /dev/null @@ -1,132 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.nn.init as init - -import math - - -class AttentionConv(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, groups=1, bias=False): - super(AttentionConv, self).__init__() - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.groups = groups - - assert self.out_channels % self.groups == 0, "out_channels should be divided by groups. (example: out_channels: 40, groups: 4)" - - self.rel_h = nn.Parameter(torch.randn(out_channels // 2, 1, 1, kernel_size, 1), requires_grad=True) - self.rel_w = nn.Parameter(torch.randn(out_channels // 2, 1, 1, 1, kernel_size), requires_grad=True) - - self.key_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias) - self.query_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias) - self.value_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias) - - self.reset_parameters() - - def forward(self, x): - batch, channels, height, width = x.size() - - padded_x = F.pad(x, [self.padding, self.padding, self.padding, self.padding]) - q_out = self.query_conv(x) - k_out = self.key_conv(padded_x) - v_out = self.value_conv(padded_x) - - k_out = k_out.unfold(2, self.kernel_size, self.stride).unfold(3, self.kernel_size, self.stride) - v_out = v_out.unfold(2, self.kernel_size, self.stride).unfold(3, self.kernel_size, self.stride) - - k_out_h, k_out_w = k_out.split(self.out_channels // 2, dim=1) - k_out = torch.cat((k_out_h + self.rel_h, k_out_w + self.rel_w), dim=1) - - k_out = k_out.contiguous().view(batch, self.groups, self.out_channels // self.groups, height, width, -1) - v_out = v_out.contiguous().view(batch, self.groups, self.out_channels // self.groups, height, width, -1) - - q_out = q_out.view(batch, self.groups, self.out_channels // self.groups, height, width, 1) - - out = q_out * k_out - out = F.softmax(out, dim=-1) - out = torch.einsum('bnchwk,bnchwk -> bnchw', out, v_out).view(batch, -1, height, width) - - return out - - def reset_parameters(self): - init.kaiming_normal_(self.key_conv.weight, mode='fan_out', nonlinearity='relu') - init.kaiming_normal_(self.value_conv.weight, mode='fan_out', nonlinearity='relu') - init.kaiming_normal_(self.query_conv.weight, mode='fan_out', nonlinearity='relu') - - init.normal_(self.rel_h, 0, 1) - init.normal_(self.rel_w, 0, 1) - - -class AttentionStem(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, groups=1, m=4, bias=False): - super(AttentionStem, self).__init__() - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.groups = groups - self.m = m - - assert self.out_channels % self.groups == 0, "out_channels should be divided by groups. (example: out_channels: 40, groups: 4)" - - self.emb_a = nn.Parameter(torch.randn(out_channels // groups, kernel_size), requires_grad=True) - self.emb_b = nn.Parameter(torch.randn(out_channels // groups, kernel_size), requires_grad=True) - self.emb_mix = nn.Parameter(torch.randn(m, out_channels // groups), requires_grad=True) - - self.key_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias) - self.query_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias) - self.value_conv = nn.ModuleList([nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias) for _ in range(m)]) - - self.reset_parameters() - - def forward(self, x): - batch, channels, height, width = x.size() - - padded_x = F.pad(x, [self.padding, self.padding, self.padding, self.padding]) - - q_out = self.query_conv(x) - k_out = self.key_conv(padded_x) - v_out = torch.stack([self.value_conv[_](padded_x) for _ in range(self.m)], dim=0) - - k_out = k_out.unfold(2, self.kernel_size, self.stride).unfold(3, self.kernel_size, self.stride) - v_out = v_out.unfold(3, self.kernel_size, self.stride).unfold(4, self.kernel_size, self.stride) - - k_out = k_out[:, :, :height, :width, :, :] - v_out = v_out[:, :, :, :height, :width, :, :] - - emb_logit_a = torch.einsum('mc,ca->ma', self.emb_mix, self.emb_a) - emb_logit_b = torch.einsum('mc,cb->mb', self.emb_mix, self.emb_b) - emb = emb_logit_a.unsqueeze(2) + emb_logit_b.unsqueeze(1) - emb = F.softmax(emb.view(self.m, -1), dim=0).view(self.m, 1, 1, 1, 1, self.kernel_size, self.kernel_size) - - v_out = emb * v_out - - k_out = k_out.contiguous().view(batch, self.groups, self.out_channels // self.groups, height, width, -1) - v_out = v_out.contiguous().view(self.m, batch, self.groups, self.out_channels // self.groups, height, width, -1) - v_out = torch.sum(v_out, dim=0).view(batch, self.groups, self.out_channels // self.groups, height, width, -1) - - q_out = q_out.view(batch, self.groups, self.out_channels // self.groups, height, width, 1) - - out = q_out * k_out - out = F.softmax(out, dim=-1) - out = torch.einsum('bnchwk,bnchwk->bnchw', out, v_out).view(batch, -1, height, width) - - return out - - def reset_parameters(self): - init.kaiming_normal_(self.key_conv.weight, mode='fan_out', nonlinearity='relu') - init.kaiming_normal_(self.query_conv.weight, mode='fan_out', nonlinearity='relu') - for _ in self.value_conv: - init.kaiming_normal_(_.weight, mode='fan_out', nonlinearity='relu') - - init.normal_(self.emb_a, 0, 1) - init.normal_(self.emb_b, 0, 1) - init.normal_(self.emb_mix, 0, 1) - - -# temp = torch.randn((2, 3, 32, 32)) -# conv = AttentionConv(3, 16, kernel_size=3, padding=1) -# print(conv(temp).size()) diff --git a/other_models/pix2vox/encoder.py b/other_models/pix2vox/encoder.py deleted file mode 100644 index 2bcf23e..0000000 --- a/other_models/pix2vox/encoder.py +++ /dev/null @@ -1,79 +0,0 @@ -# -*- coding: utf-8 -*- -# -# Developed by Haozhe Xie -# -# References: -# - https://github.com/shawnxu1318/MVCNN-Multi-View-Convolutional-Neural-Networks/blob/master/mvcnn.py - -import torch -import torchvision.models - -from pix2vox.attention import AttentionConv, AttentionStem - -class Encoder(torch.nn.Module): - def __init__(self): - super(Encoder, self).__init__() - # self.cfg = cfg - - # Layer Definition - resnet = torchvision.models.resnet50(pretrained=True) - self.resnet = torch.nn.Sequential(*[ - resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1, resnet.layer2, resnet.layer3, - resnet.layer4 - ])[:6] - self.layer1 = torch.nn.Sequential( - torch.nn.Conv2d(512, 512, kernel_size=3, padding=1), - AttentionConv(512, 512, kernel_size=7, padding=3, groups=8), - torch.nn.BatchNorm2d(512), - torch.nn.ReLU() - ) - self.layer2 = torch.nn.Sequential( - torch.nn.Conv2d(512, 256, kernel_size=3, padding=1), - torch.nn.BatchNorm2d(256), - torch.nn.ReLU(), - torch.nn.MaxPool2d(kernel_size=2) - ) - self.layer3 = torch.nn.Sequential( - torch.nn.Conv2d(256, 256, kernel_size=3, padding=1), - torch.nn.BatchNorm2d(256), - torch.nn.ReLU(), - torch.nn.MaxPool2d(kernel_size=2) - ) - self.layer4 = torch.nn.Sequential( - torch.nn.Conv2d(256, 256, kernel_size=3, padding=1), - torch.nn.BatchNorm2d(256), - torch.nn.ReLU(), - torch.nn.MaxPool2d(kernel_size=2) - ) - """ - self.layer4 = torch.nn.Sequential( - torch.nn.Conv2d(256, 384, kernel_size=3, padding=1), - torch.nn.BatchNorm2d(384), - torch.nn.ReLU(), - torch.nn.MaxPool2d(kernel_size=2) - ) - """ - - def forward(self, rendering_images): - # print(rendering_images.size()) # torch.Size([batch_size, n_views, img_c, img_h, img_w]) - # rendering_images = rendering_images.permute(1, 0, 2, 3, 4).contiguous() - rendering_images = rendering_images.permute(0, 1, 4, 2, 3).contiguous() - rendering_images = torch.split(rendering_images, 1, dim=0) - image_features = [] - - for img in rendering_images: - features = self.resnet(img.squeeze(dim=0)) - # print(features.size()) # torch.Size([batch_size, 512, 28, 28]) - features = self.layer1(features) - # print(features.size()) # torch.Size([batch_size, 512, 28, 28]) - features = self.layer2(features) - # print(features.size()) # torch.Size([batch_size, 256, 14, 14]) - features = self.layer3(features) - # print(features.size()) # torch.Size([batch_size, 256, 7, 7]) - features = self.layer4(features) - image_features.append(features) - - image_features = torch.stack(image_features).contiguous() - # image_features = torch.stack(image_features).permute(1, 0, 2, 3, 4).contiguous() - # print(image_features.size()) # torch.Size([batch_size, n_views, 256, 7, 7]) - return image_features diff --git a/other_models/pix2vox/model.py b/other_models/pix2vox/model.py deleted file mode 100644 index cfb2769..0000000 --- a/other_models/pix2vox/model.py +++ /dev/null @@ -1,195 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - -from pix2vox.attention import AttentionConv, AttentionStem - - -class Bottleneck(nn.Module): - expansion = 4 - - def __init__(self, in_channels, out_channels, stride=1, groups=1, base_width=64): - super(Bottleneck, self).__init__() - self.stride = stride - width = int(out_channels * (base_width / 64.)) * groups - - self.conv1 = nn.Sequential( - nn.Conv2d(in_channels, width, kernel_size=1, bias=False), - nn.BatchNorm2d(width), - nn.ReLU(), - ) - self.conv2 = nn.Sequential( - AttentionConv(width, width, kernel_size=7, padding=3, groups=8), - nn.BatchNorm2d(width), - nn.ReLU(), - ) - self.conv3 = nn.Sequential( - nn.Conv2d(width, self.expansion * out_channels, kernel_size=1, bias=False), - nn.BatchNorm2d(self.expansion * out_channels), - ) - - self.shortcut = nn.Sequential() - if stride != 1 or in_channels != self.expansion * out_channels: - self.shortcut = nn.Sequential( - nn.Conv2d(in_channels, self.expansion * out_channels, kernel_size=1, stride=stride, bias=False), - nn.BatchNorm2d(self.expansion * out_channels) - ) - - def forward(self, x): - out = self.conv1(x) - out = self.conv2(out) - out = self.conv3(out) - if self.stride >= 2: - out = F.avg_pool2d(out, (self.stride, self.stride)) - - out += self.shortcut(x) - out = F.relu(out) - - return out - - -class Model(nn.Module): - def __init__(self, block, num_blocks, num_classes=1000, stem=False): - super(Model, self).__init__() - self.in_places = 64 - - if stem: - self.init = nn.Sequential( - # CIFAR10 - AttentionStem(in_channels=3, out_channels=64, kernel_size=4, stride=1, padding=2, groups=1), - nn.BatchNorm2d(64), - nn.ReLU(), - - # For ImageNet - # AttentionStem(in_channels=3, out_channels=64, kernel_size=4, stride=1, padding=2, groups=1), - # nn.BatchNorm2d(64), - # nn.ReLU(), - # nn.MaxPool2d(4, 4) - ) - else: - self.init = nn.Sequential( - # CIFAR10 - nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False), - nn.BatchNorm2d(64), - nn.ReLU(), - - # For ImageNet - # nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False), - # nn.BatchNorm2d(64), - # nn.ReLU(), - # nn.MaxPool2d(kernel_size=3, stride=2, padding=1), - ) - - self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1) - self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2) - self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2) - self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2) - self.dense = nn.Linear(512 * block.expansion, num_classes) - - def _make_layer(self, block, planes, num_blocks, stride): - strides = [stride] + [1] * (num_blocks - 1) - layers = [] - for stride in strides: - layers.append(block(self.in_places, planes, stride)) - self.in_places = planes * block.expansion - return nn.Sequential(*layers) - - def forward(self, x): - out = self.init(x) - out = self.layer1(out) - out = self.layer2(out) - out = self.layer3(out) - out = self.layer4(out) - out = F.avg_pool2d(out, 4) - out = out.view(out.size(0), -1) - out = self.dense(out) - - return out - -class Model_encoder(nn.Module): - def __init__(self, block, num_blocks, stem=False): - super(Model_encoder, self).__init__() - self.in_places = 64 - - if stem: - self.init = nn.Sequential( - # CIFAR10 - AttentionStem(in_channels=3, out_channels=64, kernel_size=4, stride=1, padding=2, groups=1), - nn.BatchNorm2d(64), - nn.ReLU(), - - # For ImageNet - # AttentionStem(in_channels=3, out_channels=64, kernel_size=4, stride=1, padding=2, groups=1), - # nn.BatchNorm2d(64), - # nn.ReLU(), - # nn.MaxPool2d(4, 4) - ) - else: - self.init = nn.Sequential( - # CIFAR10 - nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False), - nn.BatchNorm2d(64), - nn.ReLU(), - - # For ImageNet - # nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False), - # nn.BatchNorm2d(64), - # nn.ReLU(), - # nn.MaxPool2d(kernel_size=3, stride=2, padding=1), - ) - - self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1) - self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2) - self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2) - self.layer4 = self._make_layer(block, 64, num_blocks[3], stride=2) - # self.dense = nn.Linear(512 * block.expansion, num_classes) - - def _make_layer(self, block, planes, num_blocks, stride): - strides = [stride] + [1] * (num_blocks - 1) - layers = [] - for stride in strides: - layers.append(block(self.in_places, planes, stride)) - self.in_places = planes * block.expansion - return nn.Sequential(*layers) - - def forward(self, x): - # x = F.avg_pool2d(x, 4) - out = self.init(x) - out = self.layer1(out) - out = self.layer2(out) - out = self.layer3(out) - out = self.layer4(out) - # out = self.layer5(out) - # out = self.layer6(out) - # out = F.avg_pool2d(out, 4) - # out = out.view(out.size(0), -1) - # out = self.dense(out) - - return out - - -def ResNet26(num_classes=1000, stem=False): - return Model(Bottleneck, [1, 2, 4, 1], num_classes=num_classes, stem=stem) - - -def ResNet38(num_classes=1000, stem=False): - return Model(Bottleneck, [2, 3, 5, 2], num_classes=num_classes, stem=stem) - - -def ResNet50(num_classes=1000, stem=False): - return Model(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, stem=stem) - - -def get_model_parameters(model): - total_parameters = 0 - for layer in list(model.parameters()): - layer_parameter = 1 - for l in list(layer.size()): - layer_parameter *= l - total_parameters += layer_parameter - return total_parameters - - -# temp = torch.randn((2, 3, 224, 224)) -# model = ResNet38(num_classes=1000, stem=True) -# print(get_model_parameters(model)) diff --git a/other_models/pointcnn/PointCNN.py b/other_models/pointcnn/PointCNN.py deleted file mode 100644 index 6be6425..0000000 --- a/other_models/pointcnn/PointCNN.py +++ /dev/null @@ -1,303 +0,0 @@ -from pointcnn.Utils import knn_indices_func_cpu -import torch -from torch.autograd import Variable -import numpy as np -import torch.nn as nn -from torch import cuda, FloatTensor, LongTensor -from pointcnn.xConv import XConv, Dense -from collections import OrderedDict - - -class PointCnnLayer(nn.Module): - def __init__(self, features, settings): - """ - :param points: Points cloud - :param features: features as input, can be None if nothing known - :param settings: Settings of the network, inside there are: - :setting xcon parameters : parameters for the xconvolutions : xconv_param_name = ('K', 'D', 'P', 'C') "C = C_out" 8, 1, -1, 32 * x - : setting fc parameters : parameters for the fully convolutional part of the network : fc_param_name = ('C', 'dropout_rate') - """ - super(PointCnnLayer, self).__init__() - # N = pc - #print("There are " + str(N) + " points in the begining") - with_X_transformation = True ## investigating that - sorting_method = None ## sorting or not points along a dimension - sampling = 'fps' ## investigating that - self.nb_xconv = len(settings[0]) - self.settings = settings - - #C_mid = C_out // 2 if C_in == 0 else C_out // 4 - #depth_multiplier = min(int(np.ceil(C_out / C_in)), 4)]) - self.xvonc1 = XConv( - C_in=0, - C_out=settings[0][0].get('C'), - dims=3, - K=settings[0][0].get('K'), - P=self.settings[0][0].get('P'), - C_mid=settings[0][0].get('C') // 2, - depth_multiplier=4) ## First XConvolution - - self.dense1 = Dense( - settings[0][0].get('C'), settings[0][1].get('C') // 2, drop=0) - self.xvonc2 = XConv( - C_in=settings[0][1].get('C') // 2, - C_out=settings[0][1].get('C'), - dims=3, - K=settings[0][1].get('K'), - P=self.settings[0][1].get('P'), - C_mid=settings[0][1].get('C') // 4, - depth_multiplier=settings[0][0].get('C') // - 4) ## Second XConvolution - - self.dense2 = Dense( - settings[0][1].get('C'), settings[0][2].get('C') // 2, drop=0) - self.xvonc3 = XConv( - C_in=settings[0][2].get('C') // 2, - C_out=settings[0][2].get('C'), - dims=3, - K=settings[0][2].get('K'), - P=self.settings[0][2].get('P'), - C_mid=settings[0][2].get('C') // 4, - depth_multiplier=settings[0][1].get('C') // - 4) ## Third XConvolution - - self.dense3 = Dense( - settings[0][2].get('C'), settings[0][3].get('C') // 2, drop=0) - self.xvonc4 = XConv( - C_in=settings[0][3].get('C') // 2, - C_out=settings[0][3].get('C'), - dims=3, - K=settings[0][3].get('K'), - P=self.settings[0][3].get('P'), - C_mid=settings[0][3].get('C') // 4, - depth_multiplier=settings[0][2].get('C') // - 4) ## Third XConvolution - self.layers_conv = [self.xvonc1, self.xvonc2, self.xvonc3, self.xvonc4] - ## deconvolution inputs : pts (output of previous conv/deconv/), fts (output of previous conv/deconv/), N , K, D, P (of concatenated layer output), - #C (of concatenated layer output),C_prev (C of previous layer ) // 4 , depth_multiplier = 1 - # xdconv_param_name = ('K', 'D', 'pts_layer_idx', 'qrs_layer_idx') - #print(self.layers_conv) - deconvolutions = OrderedDict() - dense_deconv = OrderedDict() - fc = OrderedDict() - for i in range(0, len(settings[1])): - deconvolutions["deconv" + str(i)] = XConv( - C_in=self.layers_conv[settings[1][i].get('pts_layer_idx')].C, - C_out=self.layers_conv[settings[1][i].get('qrs_layer_idx')].C, - dims=3, - K=settings[1][i].get('K'), - P=self.layers_conv[settings[1][i].get('qrs_layer_idx')].P, - C_mid=self.layers_conv[settings[1][i].get('pts_layer_idx')].C - // 4, - depth_multiplier=1) ## First dXConvolution - - dense_deconv["dense" + str(i)] = Dense( - self.layers_conv[settings[1][i].get('qrs_layer_idx')].C * 2, - self.layers_conv[settings[1][i].get('qrs_layer_idx')].C) - self.deconvolutions = nn.Sequential(deconvolutions) - self.dense_deconv = nn.Sequential(dense_deconv) - - for i in range(0, len(settings[2])): - fc["fc" + str(i)] = Dense( - self.dense_deconv[-1].out - if i == 0 else fc["fc" + str(i - 1)].out, - settings[2][i].get('C'), - drop=settings[2][i].get('dropout_rate'), - acti=True) - - self.fc = nn.Sequential(fc) - - self.sigmoid = nn.Sigmoid() - - def forward(self, x): - layer_pts = [x] - outs = [None] - pts_regionals = [] - fts_regionals = [] - #print("First CONVOLUTION") - if (self.settings[0][0].get('P') != -1): - idxx = np.random.choice( - x.size()[1], self.settings[0][0].get('P'), - replace=False).tolist() ## select representative points - rep_pts = x[:, idxx, :] - else: - rep_pts = x - pts_idx = knn_indices_func_cpu(rep_pts, x, - self.settings[0][0].get('K'), - self.settings[0][0].get('D')) - #pts_idx = pts_idx[:,::self.settings[0][0].get('D'),:] - pts_regional = torch.stack([ - x[n][idx, :] for n, idx in enumerate(torch.unbind(pts_idx, dim=0)) - ], - dim=0) - out = self.xvonc1(rep_pts, pts_regional, None) ## FTS - layer_pts.append(rep_pts) - outs.append(out) - pts_regionals.append(pts_regional) - fts_regionals.append(None) - - #print("SECOND CONVOLUTION") - if (not (self.settings[0][1].get('P') == -1)): - # print("been there" + str(self.settings[0][1].get('P'))) - idxx = np.random.choice( - rep_pts.size()[1], self.settings[0][1].get('P'), - replace=False).tolist() ## select representative points - rep_pts2 = rep_pts[:, idxx, :] - else: - rep_pts2 = rep_pts - fts = self.dense1(out) - pts_idx = knn_indices_func_cpu(rep_pts2, rep_pts, - self.settings[0][1].get('K'), - self.settings[0][1].get('D')) - #pts_idx = pts_idx[:,:,::self.settings[0][1].get('D')] - pts_regional = torch.stack([ - rep_pts[n][idx, :] - for n, idx in enumerate(torch.unbind(pts_idx, dim=0)) - ], - dim=0) - fts_regional = torch.stack([ - fts[n][idx, :] - for n, idx in enumerate(torch.unbind(pts_idx, dim=0)) - ], - dim=0) - out2 = self.xvonc2(rep_pts2, pts_regional, fts_regional) - layer_pts.append(rep_pts2) - outs.append(out2) - pts_regionals.append(pts_regional) - fts_regionals.append(fts_regional) - - #print("THIRD CONVOLUTION") - if (not (self.settings[0][2].get('P') == -1)): - #print("been there" + str(self.settings[0][1].get('P'))) - idxx = np.random.choice( - rep_pts2.size()[1], - self.settings[0][2].get('P'), - replace=False).tolist() ## select representative points - rep_pts3 = rep_pts2[:, idxx, :] - else: - rep_pts3 = rep_pts2 - fts = self.dense2(out2) - pts_idx = knn_indices_func_cpu(rep_pts3, rep_pts2, - self.settings[0][2].get('K'), - self.settings[0][2].get('D')) - #pts_idx = pts_idx[:,:,::self.settings[0][2].get('D')] - pts_regional = torch.stack([ - rep_pts2[n][idx, :] - for n, idx in enumerate(torch.unbind(pts_idx, dim=0)) - ], - dim=0) - fts_regional = torch.stack([ - fts[n][idx, :] - for n, idx in enumerate(torch.unbind(pts_idx, dim=0)) - ], - dim=0) - out3 = self.xvonc3(rep_pts3, pts_regional, fts_regional) - layer_pts.append(rep_pts3) - outs.append(out3) - pts_regionals.append(pts_regional) - fts_regionals.append(fts_regional) - - #print("FOURTH CONVOLUTION") - if (not (self.settings[0][3].get('P') == -1)): - # print("been there" + str(self.settings[0][1].get('P'))) - idxx = np.random.choice( - rep_pts3.size()[1], - self.settings[0][3].get('P'), - replace=False).tolist() ## select representative points - rep_pts4 = rep_pts3[:, idxx, :] - else: - rep_pts4 = rep_pts3 - fts = self.dense3(out3) - #print("dimensions rep pts : " +str(rep_pts4.shape) +" : " + str(rep_pts3.shape)) - #print("inputs " + str(rep_pts4.shape) + " : " + str(rep_pts3.shape)) - pts_idx = knn_indices_func_cpu(rep_pts4, rep_pts3, - self.settings[0][3].get('K'), - self.settings[0][3].get('D')) - #pts_idx = pts_idx[:,:,::self.settings[0][3].get('D')] - pts_regional = torch.stack([ - rep_pts3[n][idx, :] - for n, idx in enumerate(torch.unbind(pts_idx, dim=0)) - ], - dim=0) - fts_regional = torch.stack([ - fts[n][idx, :] - for n, idx in enumerate(torch.unbind(pts_idx, dim=0)) - ], - dim=0) - out4 = self.xvonc4(rep_pts4, pts_regional, fts_regional) - layer_pts.append(rep_pts4) - outs.append(out4) - pts_regionals.append(pts_regional) - fts_regionals.append(fts_regional) - - ############################END CONVOLUTION, START DECONVOLUTIONS #################### - - for i in range(0, len(self.deconvolutions)): - - #print("DECONVOLUTION " + str(i)) - this_out = outs[self.settings[1][i].get('pts_layer_idx') + - 1] if i == 0 else outs[-1] - rep = layer_pts[self.settings[1][i].get('qrs_layer_idx') + 1] - rep2 = layer_pts[self.settings[1][i].get('pts_layer_idx') + 1] - #print("dimensions rep pts : " +str(rep.shape) + " : "+ str(rep2.shape) ) - pts_idx = knn_indices_func_cpu(rep, rep2, - self.settings[1][i].get('K'), - self.settings[1][i].get('D')) - #pts_idx = pts_idx[:,:,::self.settings[0][3].get('D')] - pts_regional = torch.stack([ - rep2[n][idx, :] - for n, idx in enumerate(torch.unbind(pts_idx, dim=0)) - ], - dim=0) - this_out = torch.stack([ - this_out[n][idx, :] - for n, idx in enumerate(torch.unbind(pts_idx, dim=0)) - ], - dim=0) - #print("features in : " + str( this_out.shape)) - out = self.deconvolutions[i](rep, pts_regional, this_out) - - out = torch.cat( - (out, outs[self.settings[1][i].get('qrs_layer_idx') + 1]), -1) - #print("OUT CONTATENATED : " + str(out.shape)) - densed = self.dense_deconv[i](out) - #print("DENSED : "+ str(densed.shape)) - outs.append(densed) - - ############################END DECONVOLUTIONS, START Fully connected #################### - #print("FULLY_CONNECTED") - output = outs[-1] - for i in range(0, len(self.fc)): - output = self.fc[i](output) - - return output # self.sigmoid(output) - - -if __name__ == "__main__": - x = 8 - xyz = torch.rand(4, 2048, 3).cuda() - xconv_param_name = ('K', 'D', 'P', 'C') - xconv_params = [ - dict(zip(xconv_param_name, xconv_param)) - for xconv_param in [(8, 1, -1, 256), (12, 2, 768, - 256), (16, 2, 384, - 512), (16, 4, 128, 1024)] - ] - - xdconv_param_name = ('K', 'D', 'pts_layer_idx', 'qrs_layer_idx') - xdconv_params = [ - dict(zip(xdconv_param_name, xdconv_param)) - for xdconv_param in [(16, 4, 3, 3), (16, 2, 3, 2), (12, 2, 2, - 1), (8, 2, 1, 0)] - ] - - fc_param_name = ('C', 'dropout_rate') - fc_params = [ - dict(zip(fc_param_name, fc_param)) - for fc_param in [(32 * x, 0.0), (32 * x, 0.5), (3, 0.5)] - ] - - model = PointCnnLayer(xyz, ["features"], - [xconv_params, xdconv_params, fc_params]).cuda() - out = model(xyz) - print(out.shape) diff --git a/other_models/pointcnn/Utils.py b/other_models/pointcnn/Utils.py deleted file mode 100644 index 3e4ae83..0000000 --- a/other_models/pointcnn/Utils.py +++ /dev/null @@ -1,48 +0,0 @@ -import open3d -import torch -from torch.autograd import Variable -import numpy as np -import torch.nn as nn -import torch.optim as optim -from torch.utils import data -from torch import cuda, FloatTensor, LongTensor -from pointcnn.xConv import XConv, Dense -import sys - -from sklearn.neighbors import NearestNeighbors -#import matplotlib.pyplot as plt -params = {'batch_size': 5, 'shuffle': False, 'num_workers': 4} - - -def knn_indices_func_cpu( - rep_pts: FloatTensor, # (N, pts, dim) - pts: FloatTensor, # (N, x, dim) - K: int, - D: int) -> LongTensor: # (N, pts, K) - """ - CPU-based Indexing function based on K-Nearest Neighbors search. - :param rep_pts: Representative points. - :param pts: Point cloud to get indices from. - :param K: Number of nearest neighbors to collect. - :param D: dilatation factor - :return: Array of indices, P_idx, into pts such that pts[n][P_idx[n],:] - is the set k-nearest neighbors for the representative points in pts[n]. - """ - if rep_pts.is_cuda: - rep_pts = rep_pts.cpu() - if pts.is_cuda: - pts = pts.cpu() - rep_pts = rep_pts.data.numpy() - pts = pts.data.numpy() - - region_idx = [] - - for n, p in enumerate(rep_pts): - P_particular = pts[n] - nbrs = NearestNeighbors(D * K + 1, algorithm="auto").fit(P_particular) - indices = nbrs.kneighbors(p)[1] - region_idx.append(indices[:, 1::D]) - - region_idx = torch.from_numpy(np.stack(region_idx, axis=0)) - - return region_idx diff --git a/other_models/pointcnn/xConv.py b/other_models/pointcnn/xConv.py deleted file mode 100644 index a354b24..0000000 --- a/other_models/pointcnn/xConv.py +++ /dev/null @@ -1,141 +0,0 @@ -import torch -from torch.autograd import Variable -import numpy as np -import torch.nn as nn -from torch import cuda, FloatTensor, LongTensor -from typing import Tuple, Callable, Optional -from typing import Union - - -class Dense(nn.Module): - def __init__(self, inn, out, drop=0, acti=True): - super(Dense, self).__init__() - self.inn = inn - self.out = out - self.acti = acti - self.drop = drop - self.linear = nn.Linear(inn, out) - self.elu = nn.ELU() - if (self.drop > 0): - self.dropout = nn.Dropout(drop) - - def forward(self, x): - out = self.linear(x.float()) - if (self.acti): - out = self.elu(out) - if (self.drop > 0): - out = self.dropout(out) - return out - return out - - -class XConv(nn.Module): - def __init__(self, C_in: int, C_out: int, dims: int, K: int, P: int, - C_mid: int, depth_multiplier: int): - """ - :param C_in: Input dimension of the points' features. - :param C_out: Output dimension of the representative point features. - :param dims: Spatial dimensionality of points. - :param K: Number of neighbors to convolve over. - :param P: Number of representative points. - :param C_mid: Dimensionality of lifted point features. - :param depth_multiplier: Depth multiplier for internal depthwise separable convolution. - """ - super(XConv, self).__init__() - self.dense1 = Dense(dims, C_mid) - self.dense2 = Dense(C_mid, C_mid) - self.C_in = C_in - self.C = C_out - self.P = P - self.K = K - ###get x ### - #x = x.permute(0,3,1,2)# - self.conv1 = nn.Sequential( - nn.Conv2d(dims, K * K, (1, K), bias=True), nn.ELU()) - #x = x.permute(0,2,3,1)# - self.x_dense1 = Dense(K * K, K * K) - self.x_dense2 = Dense(K * K, K * K, acti=False) - - ### end Conv ### - #x = x.permute(0,3,1,2)# - self.conv2 = nn.Sequential( - nn.Conv2d( - C_mid + C_in, (C_mid + C_in) * depth_multiplier, (1, K), - groups=C_mid + C_in), - nn.Conv2d((C_mid + C_in) * depth_multiplier, C_out, 1, bias=True), - nn.ELU(), nn.BatchNorm2d(C_out, momentum=0.9)) - #x = x.permute(0,2,3,1)# - def forward(self, rep_pt, pts, fts): - """ - Applies XConv to the input data. - :param x: (rep_pt, pts, fts) where - - rep_pt: Representative point. - - pts: Regional point cloud such that fts[:,p_idx,:] is the feature - associated with pts[:,p_idx,:]. - - fts: Regional features such that pts[:,p_idx,:] is the feature - associated with fts[:,p_idx,:]. - :return: Features aggregated into point rep_pt. - _, indices_dilated = pf.knn_indices_general(qrs, pts, K * D, True) - indices = indices_dilated[:, :, ::D, :] - """ - - N = len(pts) - P = rep_pt.shape[1] # (N, P, K, dims) - p_center = torch.unsqueeze(rep_pt, dim=2) # (N, P, 1, dims) - ##FIRST STEP : Move pts to local coordinates of the reference point ## - #print("COUCOU " + str(pts.shape) + " : " + str(p_center.shape)) - pts_local = pts - p_center # (N, P, K, dims) - #print("HELLO " + str(pts_local.shape)) - - ##SECOND STEP : We lift every point individually to C_mid space - fts_lifted0 = self.dense1(pts_local) - fts_lifted = self.dense2(fts_lifted0) # (N, P, K, C_mid) - ## THIRD STEP : We check if there are already features as input (first layer or not) and cocnatenate Fsigma and previous F - if fts is None: - fts_cat = fts_lifted - else: - #print("concatenation : " + str(fts_lifted.shape) + " : " + str(fts.shape) ) - fts_cat = torch.cat((fts_lifted, fts), - -1) # (N, P, K, C_mid + C_in) - - ##FOURTH STEP : We need to learn the transformation matrix - X_shape = (N, P, self.K, self.K) - X = pts_local.permute(0, 3, 1, 2) - X = self.conv1(X) - X = X.permute(0, 2, 3, 1) - X = self.x_dense1(X) - X = self.x_dense2(X) - #print("X SHAPE "+ str(X.shape)) - X = X.view(*X_shape) - #print("X SHAPE "+ str(X.shape)) - - ## FIFTH STEP : we weight and permute F* with X - fts_X = torch.matmul(X, fts_cat) - #print("FTS_X SHAPE " + str(fts_X.shape)) - #SIXTH STEP : Last convolution giving us the output of the X convolution - X2 = fts_X.permute(0, 3, 1, 2) - X2 = self.conv2(X2) - x2 = X2.permute(0, 2, 3, 1) - fts_p = X2.squeeze(dim=2) - #print(fts_p.shape) - - # tranform to (N,P,K,C/K) - fts_p = fts_p.permute(0, 2, 1, 3) - fts_shape = (N, len(fts_p[0]), 8, int(self.C / self.K)) - fts_p = fts_p.squeeze(dim=3) - #fts_p = fts_p.view(fts_shape) - # print("################# END CONV FEATURES###############") - # print(fts_p.shape) - # print("##################################################") - return fts_p - - -### TEST THE X CONVOLUTION ### -if __name__ == "__main__": - N = 4 - D = 3 - C_in = 8 - C_out = 32 - N_neighbors = 100 - convo = XConv(4, 8, 2, 10, 1, 1000, 10).cuda() - print(convo) diff --git a/other_models/pointr/PoinTr.py b/other_models/pointr/PoinTr.py deleted file mode 100644 index 4083368..0000000 --- a/other_models/pointr/PoinTr.py +++ /dev/null @@ -1,503 +0,0 @@ -import torch -from torch import nn -import os, sys -import numpy as np -import math - -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "VRCNet/utils/pointnet2")) -# from pointnet2_ops import pointnet2_utils -import pointnet2_utils -# from extensions.chamfer_dist import ChamferDistanceL1 -sys.path.append("./distance/chamfer/") -import dist_chamfer as cd -CD = cd.chamferDist() -from .Transformer import PCTransformer -from .build import MODELS -from other_models.VRCNet.utils.model_utils import * - - -def fps(pc, num): - fps_idx = pointnet2_utils.furthest_point_sample(pc, num) - sub_pc = pointnet2_utils.gather_operation(pc.transpose(1, 2).contiguous(), fps_idx).transpose(1,2).contiguous() - return sub_pc - - -class Fold(nn.Module): - def __init__(self, in_channel , step , hidden_dim = 512): - super().__init__() - - self.in_channel = in_channel - self.step = step - - a = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(1, step).expand(step, step).reshape(1, -1) - b = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(step, 1).expand(step, step).reshape(1, -1) - self.folding_seed = torch.cat([a, b], dim=0).cuda() - - self.folding1 = nn.Sequential( - nn.Conv1d(in_channel + 2, hidden_dim, 1), - nn.BatchNorm1d(hidden_dim), - nn.ReLU(inplace=True), - nn.Conv1d(hidden_dim, hidden_dim//2, 1), - nn.BatchNorm1d(hidden_dim//2), - nn.ReLU(inplace=True), - nn.Conv1d(hidden_dim//2, 3, 1), - ) - - self.folding2 = nn.Sequential( - nn.Conv1d(in_channel + 3, hidden_dim, 1), - nn.BatchNorm1d(hidden_dim), - nn.ReLU(inplace=True), - nn.Conv1d(hidden_dim, hidden_dim//2, 1), - nn.BatchNorm1d(hidden_dim//2), - nn.ReLU(inplace=True), - nn.Conv1d(hidden_dim//2, 3, 1), - ) - - def forward(self, x): - num_sample = self.step * self.step - bs = x.size(0) - features = x.view(bs, self.in_channel, 1).expand(bs, self.in_channel, num_sample) - seed = self.folding_seed.view(1, 2, num_sample).expand(bs, 2, num_sample).to(x.device) - - x = torch.cat([seed, features], dim=1) - fd1 = self.folding1(x) - x = torch.cat([fd1, features], dim=1) - fd2 = self.folding2(x) - - return fd2 - -@MODELS.register_module() -class PoinTr(nn.Module): - def __init__(self, config, **kwargs): - super().__init__() - self.refiner = config.refiner - self.trans_dim = config.trans_dim - self.knn_layer = config.knn_layer - self.num_pred = config.num_pred - self.num_query = config.num_query - - self.fold_step = int(pow(self.num_pred//self.num_query, 0.5) + 0.5) - self.base_model = PCTransformer(in_chans = 3, embed_dim = self.trans_dim, depth = [6, 8], drop_rate = 0., num_query = self.num_query, knn_layer = self.knn_layer) - - self.increase_dim = nn.Sequential( - nn.Conv1d(self.trans_dim, 1024, 1), - nn.BatchNorm1d(1024), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(1024, 1024, 1) - ) - self.reduce_map = nn.Linear(self.trans_dim + 1027, self.trans_dim) - self.build_loss_func() - - if self.refiner == 'foldingnet': - self.foldingnet = Fold(self.trans_dim, step = self.fold_step, hidden_dim = 256) # rebuild a cluster point - elif self.refiner == 'pointgcn': - from GCN.gcn_decode import Generator - self.gcn_dec = Generator( - features=[384, 256, 64, 3], - degrees=[1, 2, 2], - support=10, - root_num=1) - elif self.refiner == 'renet': - ### NOTE: Yida - self.num_coarse_raw = self.num_query*2 - self.num_fps = 2048 - self.num_coarse = 2048 - self.num_fine = self.num_pred - self.points_label = True - self.local_folding = True - - self.dense_feature_size = 256 - self.expand_feature_size = 64 - - if self.points_label: - self.input_size = 4 - else: - self.input_size = 3 - - - self.encoder = SA_SKN_Res_encoder(input_size=self.input_size, k=[16], pk=10, - output_size=self.dense_feature_size, layers=[1,1,1,1]) - - self.up_scale = int(np.ceil(self.num_fps / (self.num_coarse_raw))) - - if self.up_scale >= 2: - self.expansion1 = EF_expansion(input_size=self.dense_feature_size, output_size=self.expand_feature_size, - step_ratio=self.up_scale, k=4) - self.conv_cup1 = nn.Conv1d(self.expand_feature_size, self.expand_feature_size, 1) - else: - self.expansion1 = None - self.conv_cup1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - - self.conv_cup2 = nn.Conv1d(self.expand_feature_size, 3, 1, bias=True) - - self.conv_s1 = nn.Conv1d(self.expand_feature_size, 16, 1, bias=True) - self.conv_s2 = nn.Conv1d(16, 8, 1, bias=True) - self.conv_s3 = nn.Conv1d(8, 1, 1, bias=True) - - if self.local_folding: - self.expansion2 = Folding(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(self.num_fine // self.num_coarse)) - else: - self.expansion2 = EF_expansion(input_size=self.expand_feature_size, output_size=self.dense_feature_size, - step_ratio=(self.num_fine // self.num_coarse), k=4) - - self.conv_f1 = nn.Conv1d(self.dense_feature_size, self.expand_feature_size, 1) - self.conv_f2 = nn.Conv1d(self.expand_feature_size, 3, 1) - - self.af = nn.ReLU(inplace=False) - - def build_loss_func(self): - # self.loss_func = ChamferDistanceL1() - self.loss_func = cd.chamferDist() - - def get_loss(self, ret, gt): - loss_coarse = self.loss_func(ret[0], gt) - loss_fine = self.loss_func(ret[1], gt) - return loss_coarse, loss_fine - - def forward(self, xyz): - q, coarse_point_cloud = self.base_model(xyz) # B M C and B M 3 - - # cat the input - inp_sparse = fps(xyz, self.num_query) - querries_all = torch.cat([coarse_point_cloud, inp_sparse], dim=1).contiguous() - - B, M ,C = q.shape - - global_feature = self.increase_dim(q.transpose(1,2)).transpose(1,2) # B M 1024 - global_feature = torch.max(global_feature, dim=1)[0] # B 1024 - - rebuild_feature = torch.cat([ - global_feature.unsqueeze(-2).expand(-1, M, -1), - q, - coarse_point_cloud], dim=-1) # B M 1027 + C - - rebuild_feature = self.reduce_map(rebuild_feature.reshape(B*M, -1)) # BM C - - # NOTE: foldingNet - if self.refiner == 'foldingnet': - relative_xyz = self.foldingnet(rebuild_feature).reshape(B, M, 3, -1) # B M 3 S - - rebuild_points = (relative_xyz + coarse_point_cloud.unsqueeze(-1)).transpose(2,3).reshape(B, -1, 3) # B N 3 - rebuild_points = torch.cat([rebuild_points, fps(xyz, self.num_pred - M * self.fold_step**2)],dim=1).contiguous() - ret = (querries_all, rebuild_points, rebuild_points) - elif self.refiner == 'renet': - # NOTE: Yida - # NOTE: following codes are VRCNet's sampling strategy - querries_all = querries_all.transpose(1, 2) - coarse_point_cloud = coarse_point_cloud.transpose(1, 2) - inp_sparse = inp_sparse.transpose(1, 2) - if self.points_label: - id0 = torch.zeros(coarse_point_cloud.shape[0], 1, coarse_point_cloud.shape[2]).cuda().contiguous() - coarse_input = torch.cat( (coarse_point_cloud, id0), 1) - id1 = torch.ones(inp_sparse.shape[0], 1, inp_sparse.shape[2]).cuda().contiguous() - org_points_input = torch.cat( (inp_sparse, id1), 1) - else: - coarse_input = coarse_point_cloud - points = torch.cat((coarse_input, org_points_input), 2) - dense_feat = self.encoder(points) - - if self.up_scale >= 2: - dense_feat = self.expansion1(dense_feat) - - coarse_features = self.af(self.conv_cup1(dense_feat)) - coarse_high = self.conv_cup2(coarse_features) - - if coarse_high.size()[2] > self.num_fps: - idx_fps = pn2.furthest_point_sample(coarse_high.transpose(1, 2).contiguous(), self.num_fps) - coarse_fps = pn2.gather_operation(coarse_high, idx_fps) - coarse_features = pn2.gather_operation(coarse_features, idx_fps) - else: - coarse_fps = coarse_high - - if coarse_fps.size()[2] > self.num_coarse: - scores = F.softplus(self.conv_s3(self.af(self.conv_s2(self.af(self.conv_s1(coarse_features)))))) - idx_scores = scores.topk(k=self.num_coarse, dim=2)[1].view(B, -1).int() - coarse = pn2.gather_operation(coarse_fps, idx_scores) - coarse_features = pn2.gather_operation(coarse_features, idx_scores) - else: - coarse = coarse_fps - - if coarse.size()[2] < self.num_fine: - if self.local_folding: - # up_features = self.expansion2(coarse_features, global_feat) - up_features = self.expansion2(coarse_features, global_feature) - center = coarse.transpose(2, 1).contiguous().unsqueeze(2).repeat(1, 1, self.num_fine // self.num_coarse, - 1).view(B, self.num_fine, - 3).transpose(2, 1).contiguous() - fine = self.conv_f2(self.af(self.conv_f1(up_features))) + center - else: - up_features = self.expansion2(coarse_features) - fine = self.conv_f2(self.af(self.conv_f1(up_features))) - else: - assert (coarse.size()[2] == self.num_fine) - fine = coarse - - ret = (querries_all.transpose(1, 2), fine.transpose(1, 2), coarse_fps.transpose(1, 2)) - return ret - -class Folding(nn.Module): - def __init__(self, input_size, output_size, step_ratio, global_feature_size=1024, num_models=1): - super(Folding, self).__init__() - self.input_size = input_size - self.output_size = output_size - self.step_ratio = step_ratio - self.num_models = num_models - - self.conv = nn.Conv1d(input_size + global_feature_size + 2, output_size, 1, bias=True) - - sqrted = int(math.sqrt(step_ratio)) + 1 - for i in range(1, sqrted + 1).__reversed__(): - if (step_ratio % i) == 0: - num_x = i - num_y = step_ratio // i - break - - grid_x = torch.linspace(-0.2, 0.2, steps=num_x) - grid_y = torch.linspace(-0.2, 0.2, steps=num_y) - - x, y = torch.meshgrid(grid_x, grid_y) # x, y shape: (2, 1) - self.grid = torch.stack([x, y], dim=-1).view(-1, 2) # (2, 2) - - def forward(self, point_feat, global_feat): - batch_size, num_features, num_points = point_feat.size() - point_feat = point_feat.transpose(1, 2).contiguous().unsqueeze(2).repeat(1, 1, self.step_ratio, 1).view( - batch_size, - -1, num_features).transpose(1, 2).contiguous() - global_feat = global_feat.unsqueeze(2).repeat(1, 1, num_points * self.step_ratio).repeat(self.num_models, 1, 1) - grid_feat = self.grid.unsqueeze(0).repeat(batch_size, num_points, 1).transpose(1, 2).contiguous().cuda() - features = torch.cat([global_feat, point_feat, grid_feat], axis=1) - features = F.relu(self.conv(features)) - return features - - -class SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=16): - super(SA_module, self).__init__() - self.share_planes = share_planes - self.k = k - self.conv1 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv2 = nn.Conv2d(in_planes, rel_planes, kernel_size=1) - self.conv3 = nn.Conv2d(in_planes, mid_planes, kernel_size=1) - - self.conv_w = nn.Sequential(nn.ReLU(inplace=False), - nn.Conv2d(rel_planes * (k + 1), mid_planes // share_planes, kernel_size=1, - bias=False), - nn.ReLU(inplace=False), - nn.Conv2d(mid_planes // share_planes, k * mid_planes // share_planes, - kernel_size=1)) - self.activation_fn = nn.ReLU(inplace=False) - - self.conv_out = nn.Conv2d(mid_planes, out_planes, kernel_size=1) - - def forward(self, input): - x, idx = input - batch_size, _, _, num_points = x.size() - identity = x # B C 1 N - x = self.activation_fn(x) - xn = get_edge_features(x, idx) # B C K N - x1, x2, x3 = self.conv1(x), self.conv2(xn), self.conv3(xn) - - x2 = x2.view(batch_size, -1, 1, num_points).contiguous() # B kC 1 N - w = self.conv_w(torch.cat([x1, x2], 1)).view(batch_size, -1, self.k, num_points) - w = w.repeat(1, self.share_planes, 1, 1) - out = w * x3 - out = torch.sum(out, dim=2, keepdim=True) - - out = self.activation_fn(out) - out = self.conv_out(out) # B C 1 N - out += identity - return [out, idx] - - -class SK_SA_module(nn.Module): - def __init__(self, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, k=[10, 20], r=2, L=32): - super(SK_SA_module, self).__init__() - - self.num_kernels = len(k) - d = max(int(out_planes / r), L) - - self.sams = nn.ModuleList([]) - - for i in range(len(k)): - self.sams.append(SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k[i])) - - self.fc = nn.Linear(out_planes, d) - self.fcs = nn.ModuleList([]) - - for i in range(len(k)): - self.fcs.append(nn.Linear(d, out_planes)) - - self.softmax = nn.Softmax(dim=1) - self.af = nn.ReLU(inplace=False) - - def forward(self, input): - x, idxs = input - assert (self.num_kernels == len(idxs)) - for i, sam in enumerate(self.sams): - fea, _ = sam([x, idxs[i]]) - fea = self.af(fea) - fea = fea.unsqueeze(dim=1) - if i == 0: - feas = fea - else: - feas = torch.cat([feas, fea], dim=1) - - fea_U = torch.sum(feas, dim=1) - fea_s = fea_U.mean(-1).mean(-1) - fea_z = self.fc(fea_s) - - for i, fc in enumerate(self.fcs): - vector = fc(fea_z).unsqueeze_(dim=1) - if i == 0: - attention_vectors = vector - else: - attention_vectors = torch.cat([attention_vectors, vector], dim=1) - - attention_vectors = self.softmax(attention_vectors) - attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1) - fea_v = (feas * attention_vectors).sum(dim=1) - return [fea_v, idxs] - - -class SKN_Res_unit(nn.Module): - def __init__(self, input_size, output_size, k=[10, 20], layers=1): - super(SKN_Res_unit, self).__init__() - self.conv1 = nn.Conv2d(input_size, output_size, 1, bias=False) - self.sam = self._make_layer(output_size, output_size // 16, output_size // 4, output_size, int(layers), 8, k=k) - self.conv2 = nn.Conv2d(output_size, output_size, 1, bias=False) - self.conv_res = nn.Conv2d(input_size, output_size, 1, bias=False) - self.af = nn.ReLU(inplace=False) - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=16): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def forward(self, feat, idx): - x, _ = self.sam([self.conv1(feat), idx]) - x = self.conv2(self.af(x)) - return x + self.conv_res(feat) - - -class SA_SKN_Res_encoder(nn.Module): - def __init__(self, input_size=3, k=[10, 20], pk=16, output_size=64, layers=[2, 2, 2, 2], - pts_num=[3072, 1536, 768, 384]): - super(SA_SKN_Res_encoder, self).__init__() - self.init_channel = 64 - - c1 = self.init_channel - self.sam_res1 = SKN_Res_unit(input_size, c1, k, int(layers[0])) - - c2 = c1 * 2 - self.sam_res2 = SKN_Res_unit(c2, c2, k, int(layers[1])) - - c3 = c2 * 2 - self.sam_res3 = SKN_Res_unit(c3, c3, k, int(layers[2])) - - c4 = c3 * 2 - self.sam_res4 = SKN_Res_unit(c4, c4, k, int(layers[3])) - - self.conv5 = nn.Conv2d(c4, 1024, 1) - - self.fc1 = nn.Linear(1024, 512) - self.fc2 = nn.Linear(512, 1024) - - self.conv6 = nn.Conv2d(c4 + 1024, c4, 1) - self.conv7 = nn.Conv2d(c3 + c4, c3, 1) - self.conv8 = nn.Conv2d(c2 + c3, c2, 1) - self.conv9 = nn.Conv2d(c1 + c2, c1, 1) - - self.conv_out = nn.Conv2d(c1, output_size, 1) - self.dropout = nn.Dropout() - self.af = nn.ReLU(inplace=False) - self.k = k - self.pk = pk - self.rate = 2 - - self.pts_num = pts_num - - def _make_layer(self, in_planes, rel_planes, mid_planes, out_planes, blocks, share_planes=8, k=[16]): - layers = [] - for _ in range(0, blocks): - layers.append(SK_SA_module(in_planes, rel_planes, mid_planes, out_planes, share_planes, k)) - return nn.Sequential(*layers) - - def _edge_pooling(self, features, points, rate=2, k=16, sample_num=None): - features = features.squeeze(2) - - if sample_num is None: - input_points_num = int(features.size()[2]) - sample_num = input_points_num // rate - - ds_features, p_idx, pn_idx, ds_points = edge_preserve_sampling(features, points, sample_num, k) - ds_features = ds_features.unsqueeze(2) - return ds_features, p_idx, pn_idx, ds_points - - def _edge_unpooling(self, features, src_pts, tgt_pts): - features = features.squeeze(2) - idx, weight = three_nn_upsampling(tgt_pts, src_pts) - features = pn2.three_interpolate(features, idx, weight) - features = features.unsqueeze(2) - return features - - def forward(self, features): - batch_size, _, num_points = features.size() - pt1 = features[:, 0:3, :] - - idx1 = [] - for i in range(len(self.k)): - idx = knn(pt1, self.k[i]) - idx1.append(idx) - - pt1 = pt1.transpose(1, 2).contiguous() - - x = features.unsqueeze(2) - x = self.sam_res1(x, idx1) - x1 = self.af(x) - - x, _, _, pt2 = self._edge_pooling(x1, pt1, self.rate, self.pk, self.pts_num[1]) - idx2 = [] - for i in range(len(self.k)): - idx = knn(pt2.transpose(1, 2).contiguous(), self.k[i]) - idx2.append(idx) - - x = self.sam_res2(x, idx2) - x2 = self.af(x) - - x, _, _, pt3 = self._edge_pooling(x2, pt2, self.rate, self.pk, self.pts_num[2]) - idx3 = [] - for i in range(len(self.k)): - idx = knn(pt3.transpose(1, 2).contiguous(), self.k[i]) - idx3.append(idx) - - x = self.sam_res3(x, idx3) - x3 = self.af(x) - - x, _, _, pt4 = self._edge_pooling(x3, pt3, self.rate, self.pk, self.pts_num[3]) - idx4 = [] - for i in range(len(self.k)): - idx = knn(pt4.transpose(1, 2).contiguous(), self.k[i]) - idx4.append(idx) - - x = self.sam_res4(x, idx4) - x4 = self.af(x) - x = self.conv5(x4) - x, _ = torch.max(x, -1) - x = x.view(batch_size, -1) - x = self.dropout(self.af(self.fc2(self.dropout(self.af(self.fc1(x)))))) - - x = x.unsqueeze(2).repeat(1, 1, self.pts_num[3]).unsqueeze(2) - x = self.af(self.conv6(torch.cat([x, x4], 1))) - x = self._edge_unpooling(x, pt4, pt3) - x = self.af(self.conv7(torch.cat([x, x3], 1))) - x = self._edge_unpooling(x, pt3, pt2) - x = self.af(self.conv8(torch.cat([x, x2], 1))) - x = self._edge_unpooling(x, pt2, pt1) - x = self.af(self.conv9(torch.cat([x, x1], 1))) - x = self.conv_out(x) - x = x.squeeze(2) - return x - diff --git a/other_models/pointr/PoinTr.yaml b/other_models/pointr/PoinTr.yaml deleted file mode 100644 index 78db7f4..0000000 --- a/other_models/pointr/PoinTr.yaml +++ /dev/null @@ -1,41 +0,0 @@ -optimizer : { - type: AdamW, - kwargs: { - lr : 0.0005, - weight_decay : 0.0005 -}} - -scheduler: { - type: LambdaLR, - kwargs: { - decay_step: 21, - lr_decay: 0.9, - lowest_decay: 0.02 # min lr = lowest_decay * lr -}} - -bnmscheduler: { - type: Lambda, - kwargs: { - decay_step: 21, - bn_decay: 0.5, - bn_momentum: 0.9, - lowest_decay: 0.01 -}} - -dataset : { - train : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'train'}}, - val : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'test'}}, - test : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'test'}}} -# model : { -# NAME: PoinTr, num_pred: 8192, num_query: 224, knn_layer: 1, trans_dim: 256, refiner: 'foldingnet'} -model : { - NAME: PoinTr, num_pred: 8192, num_query: 224, knn_layer: 1, trans_dim: 384, refiner: 'foldingnet'} - -total_bs : 48 -step_per_update : 1 -max_epoch : 300 - -consider_metric: CDL1 diff --git a/other_models/pointr/Transformer.py b/other_models/pointr/Transformer.py deleted file mode 100644 index 9730545..0000000 --- a/other_models/pointr/Transformer.py +++ /dev/null @@ -1,392 +0,0 @@ -import torch -import torch.nn as nn - -from timm.models.layers import DropPath,trunc_normal_ - -from .dgcnn_group import DGCNN_Grouper -from .logger import * -import numpy as np -from knn_cuda import KNN -knn = KNN(k=8, transpose_mode=False) - -def get_knn_index(coor_q, coor_k=None): - coor_k = coor_k if coor_k is not None else coor_q - # coor: bs, 3, np - batch_size, _, num_points = coor_q.size() - num_points_k = coor_k.size(2) - - with torch.no_grad(): - _, idx = knn(coor_k, coor_q) # bs k np - idx_base = torch.arange(0, batch_size, device=coor_q.device).view(-1, 1, 1) * num_points_k - idx = idx + idx_base - idx = idx.view(-1) - - return idx # bs*k*np - -def get_graph_feature(x, knn_index, x_q=None): - - #x: bs, np, c, knn_index: bs*k*np - k = 8 - batch_size, num_points, num_dims = x.size() - num_query = x_q.size(1) if x_q is not None else num_points - feature = x.view(batch_size * num_points, num_dims)[knn_index, :] - feature = feature.view(batch_size, k, num_query, num_dims) - x = x_q if x_q is not None else x - x = x.view(batch_size, 1, num_query, num_dims).expand(-1, k, -1, -1) - feature = torch.cat((feature - x, x), dim=-1) - return feature # b k np c - -class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = act_layer() - self.fc2 = nn.Linear(hidden_features, out_features) - self.drop = nn.Dropout(drop) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - x = self.fc2(x) - x = self.drop(x) - return x - - -class Attention(nn.Module): - def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): - super().__init__() - self.num_heads = num_heads - head_dim = dim // num_heads - # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights - self.scale = qk_scale or head_dim ** -0.5 - - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.attn_drop = nn.Dropout(attn_drop) - self.proj = nn.Linear(dim, dim) - self.proj_drop = nn.Dropout(proj_drop) - - def forward(self, x): - B, N, C = x.shape - qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) - q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) - - attn = (q @ k.transpose(-2, -1)) * self.scale - attn = attn.softmax(dim=-1) - attn = self.attn_drop(attn) - - x = (attn @ v).transpose(1, 2).reshape(B, N, C) - x = self.proj(x) - x = self.proj_drop(x) - return x - - - -class CrossAttention(nn.Module): - def __init__(self, dim, out_dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): - super().__init__() - self.num_heads = num_heads - self.dim = dim - self.out_dim = out_dim - head_dim = out_dim // num_heads - # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights - self.scale = qk_scale or head_dim ** -0.5 - - self.q_map = nn.Linear(dim, out_dim, bias=qkv_bias) - self.k_map = nn.Linear(dim, out_dim, bias=qkv_bias) - self.v_map = nn.Linear(dim, out_dim, bias=qkv_bias) - self.attn_drop = nn.Dropout(attn_drop) - - self.proj = nn.Linear(out_dim, out_dim) - self.proj_drop = nn.Dropout(proj_drop) - - def forward(self, q, v): - B, N, _ = q.shape - C = self.out_dim - k = v - NK = k.size(1) - - q = self.q_map(q).view(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) - k = self.k_map(k).view(B, NK, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) - v = self.v_map(v).view(B, NK, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) - - attn = (q @ k.transpose(-2, -1)) * self.scale - attn = attn.softmax(dim=-1) - attn = self.attn_drop(attn) - - x = (attn @ v).transpose(1, 2).reshape(B, N, C) - x = self.proj(x) - x = self.proj_drop(x) - return x - - -class DecoderBlock(nn.Module): - def __init__(self, dim, num_heads, dim_q = None, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): - super().__init__() - self.norm1 = norm_layer(dim) - self.self_attn = Attention( - dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - dim_q = dim_q or dim - self.norm_q = norm_layer(dim_q) - self.norm_v = norm_layer(dim) - self.attn = CrossAttention( - dim, dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) - - self.knn_map = nn.Sequential( - nn.Linear(dim * 2, dim), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.merge_map = nn.Linear(dim*2, dim) - - self.knn_map_cross = nn.Sequential( - nn.Linear(dim * 2, dim), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.merge_map_cross = nn.Linear(dim*2, dim) - - def forward(self, q, v, self_knn_index=None, cross_knn_index=None): - # q = q + self.drop_path(self.self_attn(self.norm1(q))) - norm_q = self.norm1(q) - q_1 = self.self_attn(norm_q) - - if self_knn_index is not None: - knn_f = get_graph_feature(norm_q, self_knn_index) - knn_f = self.knn_map(knn_f) - knn_f = knn_f.max(dim=1, keepdim=False)[0] - q_1 = torch.cat([q_1, knn_f], dim=-1) - q_1 = self.merge_map(q_1) - - q = q + self.drop_path(q_1) - - norm_q = self.norm_q(q) - norm_v = self.norm_v(v) - q_2 = self.attn(norm_q, norm_v) - - if cross_knn_index is not None: - knn_f = get_graph_feature(norm_v, cross_knn_index, norm_q) - knn_f = self.knn_map_cross(knn_f) - knn_f = knn_f.max(dim=1, keepdim=False)[0] - q_2 = torch.cat([q_2, knn_f], dim=-1) - q_2 = self.merge_map_cross(q_2) - - q = q + self.drop_path(q_2) - - # q = q + self.drop_path(self.attn(self.norm_q(q), self.norm_v(v))) - q = q + self.drop_path(self.mlp(self.norm2(q))) - return q - - -class Block(nn.Module): - - def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): - super().__init__() - self.norm1 = norm_layer(dim) - self.attn = Attention( - dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - - self.knn_map = nn.Sequential( - nn.Linear(dim * 2, dim), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.merge_map = nn.Linear(dim*2, dim) - - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) - - def forward(self, x, knn_index = None): - # x = x + self.drop_path(self.attn(self.norm1(x))) - norm_x = self.norm1(x) - x_1 = self.attn(norm_x) - - if knn_index is not None: - knn_f = get_graph_feature(norm_x, knn_index) - knn_f = self.knn_map(knn_f) - knn_f = knn_f.max(dim=1, keepdim=False)[0] - x_1 = torch.cat([x_1, knn_f], dim=-1) - x_1 = self.merge_map(x_1) - - x = x + self.drop_path(x_1) - x = x + self.drop_path(self.mlp(self.norm2(x))) - return x - - - -class PCTransformer(nn.Module): - """ Vision Transformer with support for point cloud completion - """ - def __init__(self, in_chans=3, embed_dim=768, depth=[6, 6], num_heads=6, mlp_ratio=2., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., - num_query = 224, knn_layer = -1): - super().__init__() - - self.num_features = self.embed_dim = embed_dim - - self.knn_layer = knn_layer - - print_log(' Transformer with knn_layer %d' % self.knn_layer, logger='MODEL') - - self.grouper = DGCNN_Grouper() # B 3 N to B C(3) N(128) and B C(128) N(128) - - self.pos_embed = nn.Sequential( - nn.Conv1d(in_chans, 128, 1), - nn.BatchNorm1d(128), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(128, embed_dim, 1) - ) - # self.pos_embed_wave = nn.Sequential( - # nn.Conv1d(60, 128, 1), - # nn.BatchNorm1d(128), - # nn.LeakyReLU(negative_slope=0.2), - # nn.Conv1d(128, embed_dim, 1) - # ) - - # self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) - # self.cls_pos = nn.Parameter(torch.zeros(1, 1, embed_dim)) - self.input_proj = nn.Sequential( - nn.Conv1d(128, embed_dim, 1), - nn.BatchNorm1d(embed_dim), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(embed_dim, embed_dim, 1) - ) - - self.encoder = nn.ModuleList([ - Block( - dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate) - for i in range(depth[0])]) - - # self.increase_dim = nn.Sequential( - # nn.Linear(embed_dim,1024), - # nn.ReLU(inplace=True), - # nn.Linear(1024, 1024) - # ) - - self.increase_dim = nn.Sequential( - nn.Conv1d(embed_dim, 1024, 1), - nn.BatchNorm1d(1024), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(1024, 1024, 1) - ) - - self.num_query = num_query - self.coarse_pred = nn.Sequential( - nn.Linear(1024, 1024), - nn.ReLU(inplace=True), - nn.Linear(1024, 3 * num_query) - ) - self.mlp_query = nn.Sequential( - nn.Conv1d(1024 + 3, 1024, 1), - # nn.BatchNorm1d(1024), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(1024, 1024, 1), - # nn.BatchNorm1d(1024), - nn.LeakyReLU(negative_slope=0.2), - nn.Conv1d(1024, embed_dim, 1) - ) - - self.decoder = nn.ModuleList([ - DecoderBlock( - dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate) - for i in range(depth[1])]) - - # trunc_normal_(self.cls_token, std=.02) - # trunc_normal_(self.cls_pos, std=.02) - self.apply(self._init_weights) - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) - if isinstance(m, nn.Linear) and m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.LayerNorm): - nn.init.constant_(m.bias, 0) - nn.init.constant_(m.weight, 1.0) - elif isinstance(m, nn.Conv1d): - nn.init.xavier_normal_(m.weight.data, gain=1) - elif isinstance(m, nn.BatchNorm1d): - nn.init.constant_(m.weight.data, 1) - nn.init.constant_(m.bias.data, 0) - - def pos_encoding_sin_wave(self, coor): - # ref to https://arxiv.org/pdf/2003.08934v2.pdf - D = 64 # - # normal the coor into [-1, 1], batch wise - normal_coor = 2 * ((coor - coor.min()) / (coor.max() - coor.min())) - 1 - - # define sin wave freq - freqs = torch.arange(D, dtype=torch.float).cuda() - freqs = np.pi * (2**freqs) - - freqs = freqs.view(*[1]*len(normal_coor.shape), -1) # 1 x 1 x 1 x D - normal_coor = normal_coor.unsqueeze(-1) # B x 3 x N x 1 - k = normal_coor * freqs # B x 3 x N x D - s = torch.sin(k) # B x 3 x N x D - c = torch.cos(k) # B x 3 x N x D - x = torch.cat([s,c], -1) # B x 3 x N x 2D - pos = x.transpose(-1,-2).reshape(coor.shape[0], -1, coor.shape[-1]) # B 6D N - # zero_pad = torch.zeros(x.size(0), 2, x.size(-1)).cuda() - # pos = torch.cat([x, zero_pad], dim = 1) - # pos = self.pos_embed_wave(x) - return pos - - def forward(self, inpc): - ''' - inpc : input incomplete point cloud with shape B N(2048) C(3) - ''' - # build point proxy - bs = inpc.size(0) - coor, f = self.grouper(inpc.transpose(1,2).contiguous()) - knn_index = get_knn_index(coor) - # NOTE: try to use a sin wave coor B 3 N, change the pos_embed input dim - pos = self.pos_encoding_sin_wave(coor).transpose(1,2) - # pos = self.pos_embed(coor).transpose(1,2) - x = self.input_proj(f).transpose(1,2) - # cls_pos = self.cls_pos.expand(bs, -1, -1) - # cls_token = self.cls_pos.expand(bs, -1, -1) - # x = torch.cat([cls_token, x], dim=1) - # pos = torch.cat([cls_pos, pos], dim=1) - # encoder - for i, blk in enumerate(self.encoder): - if i < self.knn_layer: - x = blk(x + pos, knn_index) # B N C - else: - x = blk(x + pos) - # build the query feature for decoder - # global_feature = x[:, 0] # B C - - global_feature = self.increase_dim(x.transpose(1,2)) # B 1024 N - global_feature = torch.max(global_feature, dim=-1)[0] # B 1024 - - coarse_point_cloud = self.coarse_pred(global_feature).reshape(bs, -1, 3) # B M C(3) - - new_knn_index = get_knn_index(coarse_point_cloud.transpose(1, 2).contiguous()) - cross_knn_index = get_knn_index(coor_k=coor, coor_q=coarse_point_cloud.transpose(1, 2).contiguous()) - - query_feature = torch.cat([ - global_feature.unsqueeze(1).expand(-1, self.num_query, -1), - coarse_point_cloud], dim=-1) # B M C+3 - q = self.mlp_query(query_feature.transpose(1,2)).transpose(1,2) # B M C - # decoder - for i, blk in enumerate(self.decoder): - if i < self.knn_layer: - q = blk(q, x, new_knn_index, cross_knn_index) # B M C - else: - q = blk(q, x) - - return q, coarse_point_cloud - diff --git a/other_models/pointr/__pycache__/PoinTr.cpython-36.pyc b/other_models/pointr/__pycache__/PoinTr.cpython-36.pyc deleted file mode 100644 index 6572c4d..0000000 Binary files a/other_models/pointr/__pycache__/PoinTr.cpython-36.pyc and /dev/null differ diff --git a/other_models/pointr/__pycache__/Transformer.cpython-36.pyc b/other_models/pointr/__pycache__/Transformer.cpython-36.pyc deleted file mode 100644 index 5df556a..0000000 Binary files a/other_models/pointr/__pycache__/Transformer.cpython-36.pyc and /dev/null differ diff --git a/other_models/pointr/__pycache__/build.cpython-36.pyc b/other_models/pointr/__pycache__/build.cpython-36.pyc deleted file mode 100644 index 9784eeb..0000000 Binary files a/other_models/pointr/__pycache__/build.cpython-36.pyc and /dev/null differ diff --git a/other_models/pointr/__pycache__/config.cpython-36.pyc b/other_models/pointr/__pycache__/config.cpython-36.pyc deleted file mode 100644 index 133c4cf..0000000 Binary files a/other_models/pointr/__pycache__/config.cpython-36.pyc and /dev/null differ diff --git a/other_models/pointr/__pycache__/dgcnn_group.cpython-36.pyc b/other_models/pointr/__pycache__/dgcnn_group.cpython-36.pyc deleted file mode 100644 index 12296f8..0000000 Binary files a/other_models/pointr/__pycache__/dgcnn_group.cpython-36.pyc and /dev/null differ diff --git a/other_models/pointr/__pycache__/logger.cpython-36.pyc b/other_models/pointr/__pycache__/logger.cpython-36.pyc deleted file mode 100644 index ac0411d..0000000 Binary files a/other_models/pointr/__pycache__/logger.cpython-36.pyc and /dev/null differ diff --git a/other_models/pointr/__pycache__/registry.cpython-36.pyc b/other_models/pointr/__pycache__/registry.cpython-36.pyc deleted file mode 100644 index 62193c2..0000000 Binary files a/other_models/pointr/__pycache__/registry.cpython-36.pyc and /dev/null differ diff --git a/other_models/pointr/build.py b/other_models/pointr/build.py deleted file mode 100644 index 12216e2..0000000 --- a/other_models/pointr/build.py +++ /dev/null @@ -1,17 +0,0 @@ -from . import registry - - -MODELS = registry.Registry('models') - - -def build_model_from_cfg(cfg, **kwargs): - """ - Build a dataset, defined by `dataset_name`. - Args: - cfg (eDICT): - Returns: - Dataset: a constructed dataset specified by dataset_name. - """ - return MODELS.build(cfg, **kwargs) - - diff --git a/other_models/pointr/config.py b/other_models/pointr/config.py deleted file mode 100644 index 8c86d76..0000000 --- a/other_models/pointr/config.py +++ /dev/null @@ -1,63 +0,0 @@ -import yaml -from easydict import EasyDict -import os -from .logger import print_log - -def log_args_to_file(args, pre='args', logger=None): - for key, val in args.__dict__.items(): - print_log(f'{pre}.{key} : {val}', logger = logger) - -def log_config_to_file(cfg, pre='cfg', logger=None): - for key, val in cfg.items(): - if isinstance(cfg[key], EasyDict): - print_log(f'{pre}.{key} = edict()', logger = logger) - log_config_to_file(cfg[key], pre=pre + '.' + key, logger=logger) - continue - print_log(f'{pre}.{key} : {val}', logger = logger) - -def merge_new_config(config, new_config): - for key, val in new_config.items(): - if not isinstance(val, dict): - if key == '_base_': - with open(new_config['_base_'], 'r') as f: - try: - val = yaml.load(f, Loader=yaml.FullLoader) - except: - val = yaml.load(f) - config[key] = EasyDict() - merge_new_config(config[key], val) - else: - config[key] = val - continue - if key not in config: - config[key] = EasyDict() - merge_new_config(config[key], val) - return config - -def cfg_from_yaml_file(cfg_file): - config = EasyDict() - with open(cfg_file, 'r') as f: - try: - new_config = yaml.load(f, Loader=yaml.FullLoader) - except: - new_config = yaml.load(f) - merge_new_config(config=config, new_config=new_config) - return config - -def get_config(args, logger=None): - if args.resume: - cfg_path = os.path.join(args.experiment_path, 'config.yaml') - if not os.path.exists(cfg_path): - print_log("Failed to resume", logger = logger) - raise FileNotFoundError() - print_log(f'Resume yaml from {cfg_path}', logger = logger) - args.config = cfg_path - config = cfg_from_yaml_file(args.config) - if not args.resume and args.local_rank == 0: - save_experiment_config(args, config, logger) - return config - -def save_experiment_config(args, config, logger = None): - config_path = os.path.join(args.experiment_path, 'config.yaml') - os.system('cp %s %s' % (args.config, config_path)) - print_log(f'Copy the Config file from {args.config} to {config_path}',logger = logger ) \ No newline at end of file diff --git a/other_models/pointr/dgcnn_group.py b/other_models/pointr/dgcnn_group.py deleted file mode 100644 index ebf8e90..0000000 --- a/other_models/pointr/dgcnn_group.py +++ /dev/null @@ -1,111 +0,0 @@ -import torch -from torch import nn -import sys, os -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "VRCNet/utils/pointnet2")) -# from pointnet2_ops import pointnet2_utils -import pointnet2_utils -from knn_cuda import KNN -knn = KNN(k=16, transpose_mode=False) - - -class DGCNN_Grouper(nn.Module): - def __init__(self): - super().__init__() - ''' - K has to be 16 - ''' - self.input_trans = nn.Conv1d(3, 8, 1) - - self.layer1 = nn.Sequential(nn.Conv2d(16, 32, kernel_size=1, bias=False), - nn.GroupNorm(4, 32), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.layer2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=1, bias=False), - nn.GroupNorm(4, 64), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.layer3 = nn.Sequential(nn.Conv2d(128, 64, kernel_size=1, bias=False), - nn.GroupNorm(4, 64), - nn.LeakyReLU(negative_slope=0.2) - ) - - self.layer4 = nn.Sequential(nn.Conv2d(128, 128, kernel_size=1, bias=False), - nn.GroupNorm(4, 128), - nn.LeakyReLU(negative_slope=0.2) - ) - - - @staticmethod - def fps_downsample(coor, x, num_group): - xyz = coor.transpose(1, 2).contiguous() # b, n, 3 - fps_idx = pointnet2_utils.furthest_point_sample(xyz, num_group) - - combined_x = torch.cat([coor, x], dim=1) - - new_combined_x = ( - pointnet2_utils.gather_operation( - combined_x, fps_idx - ) - ) - - new_coor = new_combined_x[:, :3] - new_x = new_combined_x[:, 3:] - - return new_coor, new_x - - @staticmethod - def get_graph_feature(coor_q, x_q, coor_k, x_k): - - # coor: bs, 3, np, x: bs, c, np - - k = 16 - batch_size = x_k.size(0) - num_points_k = x_k.size(2) - num_points_q = x_q.size(2) - - with torch.no_grad(): - _, idx = knn(coor_k, coor_q) # bs k np - assert idx.shape[1] == k - idx_base = torch.arange(0, batch_size, device=x_q.device).view(-1, 1, 1) * num_points_k - idx = idx + idx_base - idx = idx.view(-1) - num_dims = x_k.size(1) - x_k = x_k.transpose(2, 1).contiguous() - feature = x_k.view(batch_size * num_points_k, -1)[idx, :] - feature = feature.view(batch_size, k, num_points_q, num_dims).permute(0, 3, 2, 1).contiguous() - x_q = x_q.view(batch_size, num_dims, num_points_q, 1).expand(-1, -1, -1, k) - feature = torch.cat((feature - x_q, x_q), dim=1) - return feature - - def forward(self, x): - - # x: bs, 3, np - - # bs 3 N(128) bs C(224)128 N(128) - coor = x - f = self.input_trans(x) - - f = self.get_graph_feature(coor, f, coor, f) - f = self.layer1(f) - f = f.max(dim=-1, keepdim=False)[0] - - coor_q, f_q = self.fps_downsample(coor, f, 512) - f = self.get_graph_feature(coor_q, f_q, coor, f) - f = self.layer2(f) - f = f.max(dim=-1, keepdim=False)[0] - coor = coor_q - - f = self.get_graph_feature(coor, f, coor, f) - f = self.layer3(f) - f = f.max(dim=-1, keepdim=False)[0] - - coor_q, f_q = self.fps_downsample(coor, f, 128) - f = self.get_graph_feature(coor_q, f_q, coor, f) - f = self.layer4(f) - f = f.max(dim=-1, keepdim=False)[0] - coor = coor_q - - return coor, f diff --git a/other_models/pointr/ftrans.yaml b/other_models/pointr/ftrans.yaml deleted file mode 100644 index 69565fc..0000000 --- a/other_models/pointr/ftrans.yaml +++ /dev/null @@ -1,39 +0,0 @@ -optimizer : { - type: AdamW, - kwargs: { - lr : 0.0005, - weight_decay : 0.0005 -}} - -scheduler: { - type: LambdaLR, - kwargs: { - decay_step: 21, - lr_decay: 0.9, - lowest_decay: 0.02 # min lr = lowest_decay * lr -}} - -bnmscheduler: { - type: Lambda, - kwargs: { - decay_step: 21, - bn_decay: 0.5, - bn_momentum: 0.9, - lowest_decay: 0.01 -}} - -dataset : { - train : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'train'}}, - val : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'test'}}, - test : { _base_: cfgs/dataset_configs/PCN.yaml, - others: {subset: 'test'}}} -model : { - NAME: PoinTr, num_pred: 2048, num_query: 224, knn_layer: 1, trans_dim: 384, refiner: 'renet'} - -total_bs : 48 -step_per_update : 1 -max_epoch : 300 - -consider_metric: CDL1 diff --git a/other_models/pointr/logger.py b/other_models/pointr/logger.py deleted file mode 100644 index 847c1c7..0000000 --- a/other_models/pointr/logger.py +++ /dev/null @@ -1,127 +0,0 @@ -import logging -import torch.distributed as dist - -logger_initialized = {} - -def get_root_logger(log_file=None, log_level=logging.INFO, name='main'): - """Get root logger and add a keyword filter to it. - The logger will be initialized if it has not been initialized. By default a - StreamHandler will be added. If `log_file` is specified, a FileHandler will - also be added. The name of the root logger is the top-level package name, - e.g., "mmdet3d". - Args: - log_file (str, optional): File path of log. Defaults to None. - log_level (int, optional): The level of logger. - Defaults to logging.INFO. - name (str, optional): The name of the root logger, also used as a - filter keyword. Defaults to 'mmdet3d'. - Returns: - :obj:`logging.Logger`: The obtained logger - """ - logger = get_logger(name=name, log_file=log_file, log_level=log_level) - # add a logging filter - logging_filter = logging.Filter(name) - logging_filter.filter = lambda record: record.find(name) != -1 - - return logger - - -def get_logger(name, log_file=None, log_level=logging.INFO, file_mode='w'): - """Initialize and get a logger by name. - If the logger has not been initialized, this method will initialize the - logger by adding one or two handlers, otherwise the initialized logger will - be directly returned. During initialization, a StreamHandler will always be - added. If `log_file` is specified and the process rank is 0, a FileHandler - will also be added. - Args: - name (str): Logger name. - log_file (str | None): The log filename. If specified, a FileHandler - will be added to the logger. - log_level (int): The logger level. Note that only the process of - rank 0 is affected, and other processes will set the level to - "Error" thus be silent most of the time. - file_mode (str): The file mode used in opening log file. - Defaults to 'w'. - Returns: - logging.Logger: The expected logger. - """ - logger = logging.getLogger(name) - if name in logger_initialized: - return logger - # handle hierarchical names - # e.g., logger "a" is initialized, then logger "a.b" will skip the - # initialization since it is a child of "a". - for logger_name in logger_initialized: - if name.startswith(logger_name): - return logger - - # handle duplicate logs to the console - # Starting in 1.8.0, PyTorch DDP attaches a StreamHandler (NOTSET) - # to the root logger. As logger.propagate is True by default, this root - # level handler causes logging messages from rank>0 processes to - # unexpectedly show up on the console, creating much unwanted clutter. - # To fix this issue, we set the root logger's StreamHandler, if any, to log - # at the ERROR level. - for handler in logger.root.handlers: - if type(handler) is logging.StreamHandler: - handler.setLevel(logging.ERROR) - - stream_handler = logging.StreamHandler() - handlers = [stream_handler] - - if dist.is_available() and dist.is_initialized(): - rank = dist.get_rank() - else: - rank = 0 - - # only rank 0 will add a FileHandler - if rank == 0 and log_file is not None: - # Here, the default behaviour of the official logger is 'a'. Thus, we - # provide an interface to change the file mode to the default - # behaviour. - file_handler = logging.FileHandler(log_file, file_mode) - handlers.append(file_handler) - - formatter = logging.Formatter( - '%(asctime)s - %(name)s - %(levelname)s - %(message)s') - for handler in handlers: - handler.setFormatter(formatter) - handler.setLevel(log_level) - logger.addHandler(handler) - - if rank == 0: - logger.setLevel(log_level) - else: - logger.setLevel(logging.ERROR) - - logger_initialized[name] = True - - - return logger - - -def print_log(msg, logger=None, level=logging.INFO): - """Print a log message. - Args: - msg (str): The message to be logged. - logger (logging.Logger | str | None): The logger to be used. - Some special loggers are: - - "silent": no message will be printed. - - other str: the logger obtained with `get_root_logger(logger)`. - - None: The `print()` method will be used to print log messages. - level (int): Logging level. Only available when `logger` is a Logger - object or "root". - """ - if logger is None: - print(msg) - elif isinstance(logger, logging.Logger): - logger.log(level, msg) - elif logger == 'silent': - pass - elif isinstance(logger, str): - _logger = get_logger(logger) - _logger.log(level, msg) - else: - raise TypeError( - 'logger should be either a logging.Logger object, str, ' - f'"silent" or None, but got {type(logger)}') \ No newline at end of file diff --git a/other_models/pointr/registry.py b/other_models/pointr/registry.py deleted file mode 100644 index f103ec6..0000000 --- a/other_models/pointr/registry.py +++ /dev/null @@ -1,288 +0,0 @@ -import inspect -import warnings -from functools import partial -from . import config - -class Registry: - """A registry to map strings to classes. - Registered object could be built from registry. - Example: - >>> MODELS = Registry('models') - >>> @MODELS.register_module() - >>> class ResNet: - >>> pass - >>> resnet = MODELS.build(dict(NAME='ResNet')) - Please refer to https://mmcv.readthedocs.io/en/latest/registry.html for - advanced useage. - Args: - name (str): Registry name. - build_func(func, optional): Build function to construct instance from - Registry, func:`build_from_cfg` is used if neither ``parent`` or - ``build_func`` is specified. If ``parent`` is specified and - ``build_func`` is not given, ``build_func`` will be inherited - from ``parent``. Default: None. - parent (Registry, optional): Parent registry. The class registered in - children registry could be built from parent. Default: None. - scope (str, optional): The scope of registry. It is the key to search - for children registry. If not specified, scope will be the name of - the package where class is defined, e.g. mmdet, mmcls, mmseg. - Default: None. - """ - - def __init__(self, name, build_func=None, parent=None, scope=None): - self._name = name - self._module_dict = dict() - self._children = dict() - self._scope = self.infer_scope() if scope is None else scope - - # self.build_func will be set with the following priority: - # 1. build_func - # 2. parent.build_func - # 3. build_from_cfg - if build_func is None: - if parent is not None: - self.build_func = parent.build_func - else: - self.build_func = build_from_cfg - else: - self.build_func = build_func - if parent is not None: - assert isinstance(parent, Registry) - parent._add_children(self) - self.parent = parent - else: - self.parent = None - - def __len__(self): - return len(self._module_dict) - - def __contains__(self, key): - return self.get(key) is not None - - def __repr__(self): - format_str = self.__class__.__name__ + \ - f'(name={self._name}, ' \ - f'items={self._module_dict})' - return format_str - - @staticmethod - def infer_scope(): - """Infer the scope of registry. - The name of the package where registry is defined will be returned. - Example: - # in mmdet/models/backbone/resnet.py - >>> MODELS = Registry('models') - >>> @MODELS.register_module() - >>> class ResNet: - >>> pass - The scope of ``ResNet`` will be ``mmdet``. - Returns: - scope (str): The inferred scope name. - """ - # inspect.stack() trace where this function is called, the index-2 - # indicates the frame where `infer_scope()` is called - filename = inspect.getmodule(inspect.stack()[2][0]).__name__ - split_filename = filename.split('.') - return split_filename[0] - - @staticmethod - def split_scope_key(key): - """Split scope and key. - The first scope will be split from key. - Examples: - >>> Registry.split_scope_key('mmdet.ResNet') - 'mmdet', 'ResNet' - >>> Registry.split_scope_key('ResNet') - None, 'ResNet' - Return: - scope (str, None): The first scope. - key (str): The remaining key. - """ - split_index = key.find('.') - if split_index != -1: - return key[:split_index], key[split_index + 1:] - else: - return None, key - - @property - def name(self): - return self._name - - @property - def scope(self): - return self._scope - - @property - def module_dict(self): - return self._module_dict - - @property - def children(self): - return self._children - - def get(self, key): - """Get the registry record. - Args: - key (str): The class name in string format. - Returns: - class: The corresponding class. - """ - scope, real_key = self.split_scope_key(key) - if scope is None or scope == self._scope: - # get from self - if real_key in self._module_dict: - return self._module_dict[real_key] - else: - # get from self._children - if scope in self._children: - return self._children[scope].get(real_key) - else: - # goto root - parent = self.parent - while parent.parent is not None: - parent = parent.parent - return parent.get(key) - - def build(self, *args, **kwargs): - return self.build_func(*args, **kwargs, registry=self) - - def _add_children(self, registry): - """Add children for a registry. - The ``registry`` will be added as children based on its scope. - The parent registry could build objects from children registry. - Example: - >>> models = Registry('models') - >>> mmdet_models = Registry('models', parent=models) - >>> @mmdet_models.register_module() - >>> class ResNet: - >>> pass - >>> resnet = models.build(dict(NAME='mmdet.ResNet')) - """ - - assert isinstance(registry, Registry) - assert registry.scope is not None - assert registry.scope not in self.children, \ - f'scope {registry.scope} exists in {self.name} registry' - self.children[registry.scope] = registry - - def _register_module(self, module_class, module_name=None, force=False): - if not inspect.isclass(module_class): - raise TypeError('module must be a class, ' - f'but got {type(module_class)}') - - if module_name is None: - module_name = module_class.__name__ - if isinstance(module_name, str): - module_name = [module_name] - for name in module_name: - if not force and name in self._module_dict: - raise KeyError(f'{name} is already registered ' - f'in {self.name}') - self._module_dict[name] = module_class - - def deprecated_register_module(self, cls=None, force=False): - warnings.warn( - 'The old API of register_module(module, force=False) ' - 'is deprecated and will be removed, please use the new API ' - 'register_module(name=None, force=False, module=None) instead.') - if cls is None: - return partial(self.deprecated_register_module, force=force) - self._register_module(cls, force=force) - return cls - - def register_module(self, name=None, force=False, module=None): - """Register a module. - A record will be added to `self._module_dict`, whose key is the class - name or the specified name, and value is the class itself. - It can be used as a decorator or a normal function. - Example: - >>> backbones = Registry('backbone') - >>> @backbones.register_module() - >>> class ResNet: - >>> pass - >>> backbones = Registry('backbone') - >>> @backbones.register_module(name='mnet') - >>> class MobileNet: - >>> pass - >>> backbones = Registry('backbone') - >>> class ResNet: - >>> pass - >>> backbones.register_module(ResNet) - Args: - name (str | None): The module name to be registered. If not - specified, the class name will be used. - force (bool, optional): Whether to override an existing class with - the same name. Default: False. - module (type): Module class to be registered. - """ - if not isinstance(force, bool): - raise TypeError(f'force must be a boolean, but got {type(force)}') - # NOTE: This is a walkaround to be compatible with the old api, - # while it may introduce unexpected bugs. - if isinstance(name, type): - return self.deprecated_register_module(name, force=force) - - # raise the error ahead of time - if not (name is None or isinstance(name, str) or misc.is_seq_of(name, str)): - raise TypeError( - 'name must be either of None, an instance of str or a sequence' - f' of str, but got {type(name)}') - - # use it as a normal method: x.register_module(module=SomeClass) - if module is not None: - self._register_module( - module_class=module, module_name=name, force=force) - return module - - # use it as a decorator: @x.register_module() - def _register(cls): - self._register_module( - module_class=cls, module_name=name, force=force) - return cls - - return _register - - -def build_from_cfg(cfg, registry, default_args=None): - """Build a module from config dict. - Args: - cfg (edict): Config dict. It should at least contain the key "NAME". - registry (:obj:`Registry`): The registry to search the type from. - Returns: - object: The constructed object. - """ - if not isinstance(cfg, dict): - raise TypeError(f'cfg must be a dict, but got {type(cfg)}') - if 'NAME' not in cfg: - if default_args is None or 'NAME' not in default_args: - raise KeyError( - '`cfg` or `default_args` must contain the key "NAME", ' - f'but got {cfg}\n{default_args}') - if not isinstance(registry, Registry): - raise TypeError('registry must be an mmcv.Registry object, ' - f'but got {type(registry)}') - - if not (isinstance(default_args, dict) or default_args is None): - raise TypeError('default_args must be a dict or None, ' - f'but got {type(default_args)}') - - if default_args is not None: - cfg = config.merge_new_config(cfg, default_args) - - obj_type = cfg.get('NAME') - - if isinstance(obj_type, str): - obj_cls = registry.get(obj_type) - if obj_cls is None: - raise KeyError( - f'{obj_type} is not in the {registry.name} registry') - elif inspect.isclass(obj_type): - obj_cls = obj_type - else: - raise TypeError( - f'type must be a str or valid type, but got {type(obj_type)}') - try: - return obj_cls(cfg) - except Exception as e: - # Normal TypeError does not print class name. - raise type(e)(f'{obj_cls.__name__}: {e}') diff --git a/other_models/shapegf/models/__init__.py b/other_models/shapegf/models/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/other_models/shapegf/models/__pycache__/__init__.cpython-36.pyc b/other_models/shapegf/models/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index 172e95b..0000000 Binary files a/other_models/shapegf/models/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/shapegf/models/decoders/__pycache__/resnet_add.cpython-36.pyc b/other_models/shapegf/models/decoders/__pycache__/resnet_add.cpython-36.pyc deleted file mode 100644 index 18ab343..0000000 Binary files a/other_models/shapegf/models/decoders/__pycache__/resnet_add.cpython-36.pyc and /dev/null differ diff --git a/other_models/shapegf/models/decoders/resnet_add.py b/other_models/shapegf/models/decoders/resnet_add.py deleted file mode 100644 index 4bc6a75..0000000 --- a/other_models/shapegf/models/decoders/resnet_add.py +++ /dev/null @@ -1,110 +0,0 @@ -import torch -import torch.nn as nn - - -class ResnetBlockConv1d(nn.Module): - """ 1D-Convolutional ResNet block class. - Args: - size_in (int): input dimension - size_out (int): output dimension - size_h (int): hidden dimension - """ - - def __init__(self, c_dim, size_in, size_h=None, size_out=None, - norm_method='batch_norm', legacy=False): - super().__init__() - # Attributes - if size_h is None: - size_h = size_in - if size_out is None: - size_out = size_in - - self.size_in = size_in - self.size_h = size_h - self.size_out = size_out - # Submodules - if norm_method == 'batch_norm': - norm = nn.BatchNorm1d - elif norm_method == 'sync_batch_norm': - norm = nn.SyncBatchNorm - else: - raise Exception("Invalid norm method: %s" % norm_method) - - self.bn_0 = norm(size_in) - self.bn_1 = norm(size_h) - - self.fc_0 = nn.Conv1d(size_in, size_h, 1) - self.fc_1 = nn.Conv1d(size_h, size_out, 1) - self.fc_c = nn.Conv1d(c_dim, size_out, 1) - self.actvn = nn.ReLU() - - if size_in == size_out: - self.shortcut = None - else: - self.shortcut = nn.Conv1d(size_in, size_out, 1, bias=False) - - # Initialization - nn.init.zeros_(self.fc_1.weight) - - def forward(self, x, c): - net = self.fc_0(self.actvn(self.bn_0(x))) - dx = self.fc_1(self.actvn(self.bn_1(net))) - - if self.shortcut is not None: - x_s = self.shortcut(x) - else: - x_s = x - - out = x_s + dx + self.fc_c(c) - - return out - - -class Decoder(nn.Module): - """ Decoder conditioned by adding. - - Example configuration: - z_dim: 128 - hidden_size: 256 - n_blocks: 5 - out_dim: 3 # we are outputting the gradient - sigma_condition: True - xyz_condition: True - """ - def __init__(self, _, cfg): - super().__init__() - self.cfg = cfg - self.z_dim = z_dim = cfg.z_dim - self.dim = dim = cfg.dim - self.out_dim = out_dim = cfg.out_dim - self.hidden_size = hidden_size = cfg.hidden_size - self.n_blocks = n_blocks = cfg.n_blocks - - # Input = Conditional = zdim (shape) + dim (xyz) + 1 (sigma) - c_dim = z_dim + dim + 1 - self.conv_p = nn.Conv1d(c_dim, hidden_size, 1) - self.blocks = nn.ModuleList([ - ResnetBlockConv1d(c_dim, hidden_size) for _ in range(n_blocks) - ]) - self.bn_out = nn.BatchNorm1d(hidden_size) - self.conv_out = nn.Conv1d(hidden_size, out_dim, 1) - self.actvn_out = nn.ReLU() - - # This should have the same signature as the sig condition one - def forward(self, x, c): - """ - :param x: (bs, npoints, self.dim) Input coordinate (xyz) - :param c: (bs, self.zdim + 1) Shape latent code + sigma - :return: (bs, npoints, self.dim) Gradient (self.dim dimension) - """ - p = x.transpose(1, 2) # (bs, dim, n_points) - batch_size, D, num_points = p.size() - - c_expand = c.unsqueeze(2).expand(-1, -1, num_points) - c_xyz = torch.cat([p, c_expand], dim=1) - net = self.conv_p(c_xyz) - for block in self.blocks: - net = block(net, c_xyz) - out = self.conv_out(self.actvn_out(self.bn_out(net))).transpose(1, 2) - return out - diff --git a/other_models/shapegf/models/decoders/resnet_cbn.py b/other_models/shapegf/models/decoders/resnet_cbn.py deleted file mode 100644 index 20dad07..0000000 --- a/other_models/shapegf/models/decoders/resnet_cbn.py +++ /dev/null @@ -1,171 +0,0 @@ -import torch -import torch.nn as nn - - -class CBatchNorm1d(nn.Module): - """ Conditional batch normalization layer class. - Args: - c_dim (int): dimension of latent conditioned code c - f_dim (int): feature dimension - norm_method (str): normalization method - """ - - def __init__(self, c_dim, f_dim, norm_method='batch_norm'): - super().__init__() - self.c_dim = c_dim - self.f_dim = f_dim - self.norm_method = norm_method - # Submodules - self.conv_gamma = nn.Conv1d(c_dim, f_dim, 1) - self.conv_beta = nn.Conv1d(c_dim, f_dim, 1) - if norm_method == 'batch_norm': - self.bn = nn.BatchNorm1d(f_dim, affine=False) - elif norm_method == 'instance_norm': - self.bn = nn.InstanceNorm1d(f_dim, affine=False) - elif norm_method == 'group_norm': - self.bn = nn.GroupNorm1d(f_dim, affine=False) - else: - raise ValueError('Invalid normalization method!') - self.reset_parameters() - - def reset_parameters(self): - nn.init.zeros_(self.conv_gamma.weight) - nn.init.zeros_(self.conv_beta.weight) - nn.init.ones_(self.conv_gamma.bias) - nn.init.zeros_(self.conv_beta.bias) - - def forward(self, x, c): - assert(x.size(0) == c.size(0)) - assert(c.size(1) == self.c_dim), "%d, %d" % (c.size(1), self.c_dim) - - # c is assumed to be of size batch_size x c_dim x T - if len(c.size()) == 2: - c = c.unsqueeze(2) - - # Affine mapping - gamma = self.conv_gamma(c) - beta = self.conv_beta(c) - - # Batchnorm - net = self.bn(x) - out = gamma * net + beta - - return out - - -class CResnetBlockConv1d(nn.Module): - """ Conditional batch normalization-based Resnet block class. - Args: - c_dim (int): dimension of latend conditioned code c - size_in (int): input dimension - size_out (int): output dimension - size_h (int): hidden dimension - norm_method (str): normalization method - legacy (bool): whether to use legacy blocks - """ - - def __init__(self, c_dim, size_in, size_h=None, size_out=None, - norm_method='batch_norm', legacy=False): - super().__init__() - # Attributes - if size_h is None: - size_h = size_in - if size_out is None: - size_out = size_in - self.size_in = size_in - self.size_h = size_h - self.size_out = size_out - - # Submodules - self.bn_0 = CBatchNorm1d( - c_dim, size_in, norm_method=norm_method) - self.bn_1 = CBatchNorm1d( - c_dim, size_h, norm_method=norm_method) - - self.fc_0 = nn.Conv1d(size_in, size_h, 1) - self.fc_1 = nn.Conv1d(size_h, size_out, 1) - self.actvn = nn.ReLU() - - if size_in == size_out: - self.shortcut = None - else: - self.shortcut = nn.Conv1d(size_in, size_out, 1, bias=False) - # Initialization - nn.init.zeros_(self.fc_1.weight) - - def forward(self, x, c): - net = self.fc_0(self.actvn(self.bn_0(x, c))) - dx = self.fc_1(self.actvn(self.bn_1(net, c))) - - if self.shortcut is not None: - x_s = self.shortcut(x) - else: - x_s = x - - return x_s + dx - - -class Decoder(nn.Module): - """ Decoder with CBN class 2. - - It differs from the previous one in that the number of blocks can be - chosen. - - Args: - dim (int): input dimension - z_dim (int): dimension of latent code z - c_dim (int): dimension of latent conditioned code c - hidden_size (int): hidden size of Decoder network - leaky (bool): whether to use leaky ReLUs - n_blocks (int): number of ResNet blocks - - Example configuration: - z_dim: 128 - hidden_size: 256 - n_blocks: 5 - out_dim: 3 # we are outputting the gradient - sigma_condition: True - xyz_condition: True - """ - - def __init__(self, _, cfg): - super().__init__() - self.cfg = cfg - self.z_dim = z_dim = cfg.z_dim - self.dim = dim = cfg.dim - self.out_dim = out_dim = cfg.out_dim - self.hidden_size = hidden_size = cfg.hidden_size - self.n_blocks = n_blocks = cfg.n_blocks - - assert getattr(cfg, "xyz_condition", False) - c_dim = z_dim + dim - self.sigma_condition = getattr(cfg, "sigma_condition", False) - if self.sigma_condition: - c_dim += 1 - - # Input: xyz + latent code (+1 if use sigma) - self.conv_p = nn.Conv1d(c_dim, hidden_size, 1) - self.blocks = nn.ModuleList([ - CResnetBlockConv1d(c_dim, hidden_size) for i in range(n_blocks) - ]) - - self.bn = CBatchNorm1d(c_dim, hidden_size) - self.conv_out = nn.Conv1d(hidden_size, out_dim, 1) - self.actvn = nn.ReLU() - - def forward(self, x, c): - """ - :param x: (bs, npoints, dim) xyz coordinates - :param c: (bs, cdim + 1) conditional - :return: (bs, npoints, dim) gradient - """ - p = x.transpose(1, 2) # (bs, dim, n_points) - batch_size, D, num_points = p.size() - - c_expand = c.unsqueeze(2).expand(-1, -1, num_points) - c_xyz = torch.cat([p, c_expand], dim=1) - net = self.conv_p(c_xyz) - for block in self.blocks: - net = block(net, c_xyz) - out = self.conv_out(self.actvn(self.bn(net, c_xyz))).transpose(1, 2) - return out diff --git a/other_models/shapegf/models/discriminators/mlp_dis.py b/other_models/shapegf/models/discriminators/mlp_dis.py deleted file mode 100644 index 66d8717..0000000 --- a/other_models/shapegf/models/discriminators/mlp_dis.py +++ /dev/null @@ -1,61 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class Discriminator(nn.Module): - - def __init__(self, cfg, cfgmodel): - super().__init__() - self.cfg = cfg - self.cfgmodel = cfgmodel - - self.inp_dim = cfgmodel.inp_dim - self.use_bn = getattr(cfgmodel, "use_bn", False) - self.use_ln = getattr(cfgmodel, "use_ln", False) - self.use_sigmoid = getattr(cfgmodel, "use_sigmoid", False) - self.dims = cfgmodel.dims - - curr_dim = self.inp_dim - self.layers = [] - self.bns = [] - self.lns = [] - for hid in self.dims: - self.layers.append(nn.Linear(curr_dim, hid)) - if self.use_bn: - self.bns.append(nn.BatchNorm1d(hid)) - else: - self.bns.append(None) - if self.use_ln: - self.lns.append(nn.LayerNorm(hid)) - else: - self.lns.append(None) - curr_dim = hid - self.layers = nn.ModuleList(self.layers) - self.bns = nn.ModuleList(self.bns) - self.lns = nn.ModuleList(self.lns) - self.out = nn.Linear(curr_dim, 1) - - def forward(self, z=None, bs=None, return_all=False): - if z is None: - assert bs is not None - z = torch.randn(bs, self.inp_dim).cuda() - - y = z - for layer, bn, ln in zip(self.layers, self.bns, self.lns): - y = layer(y) - if self.use_bn: - y = bn(y) - if self.use_ln: - y = ln(y) - y = F.leaky_relu(y, 0.2) - y = self.out(y) - - if self.use_sigmoid: - y = torch.sigmoid(y) - if return_all: - return { - 'x': y - } - else: - return y diff --git a/other_models/shapegf/models/encoders/__init__.py b/other_models/shapegf/models/encoders/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/other_models/shapegf/models/encoders/__pycache__/__init__.cpython-36.pyc b/other_models/shapegf/models/encoders/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index 79db8e8..0000000 Binary files a/other_models/shapegf/models/encoders/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/shapegf/models/encoders/__pycache__/l3dp_encoder.cpython-36.pyc b/other_models/shapegf/models/encoders/__pycache__/l3dp_encoder.cpython-36.pyc deleted file mode 100644 index e3f063c..0000000 Binary files a/other_models/shapegf/models/encoders/__pycache__/l3dp_encoder.cpython-36.pyc and /dev/null differ diff --git a/other_models/shapegf/models/encoders/constant_encoder.py b/other_models/shapegf/models/encoders/constant_encoder.py deleted file mode 100644 index 395c3b8..0000000 --- a/other_models/shapegf/models/encoders/constant_encoder.py +++ /dev/null @@ -1,14 +0,0 @@ -import torch -from torch import nn - - -class Encoder(nn.Module): - def __init__(self, cfgmodel): - super(Encoder, self).__init__() - self.zdim = cfgmodel.zdim - self.out = nn.Parameter(torch.randn(1, self.zdim), requires_grad=True) - - def forward(self, x): - bs = x.size(0) - m = v = self.out.expand(bs, -1) - return m, v diff --git a/other_models/shapegf/models/encoders/l3dp_encoder.py b/other_models/shapegf/models/encoders/l3dp_encoder.py deleted file mode 100644 index 52f719c..0000000 --- a/other_models/shapegf/models/encoders/l3dp_encoder.py +++ /dev/null @@ -1,64 +0,0 @@ -import torch -from torch import nn -import torch.nn.functional as F - - -class Encoder(nn.Module): - def __init__(self, cfgmodel): - super(Encoder, self).__init__() - self.zdim = cfgmodel.zdim - self.use_deterministic_encoder = cfgmodel.use_deterministic_encoder - self.input_dim = cfgmodel.input_dim - self.conv1 = nn.Conv1d(self.input_dim, 128, 1) - self.conv2 = nn.Conv1d(128, 128, 1) - self.conv3 = nn.Conv1d(128, 256, 1) - self.conv4 = nn.Conv1d(256, 512, 1) - self.bn1 = nn.BatchNorm1d(128) - self.bn2 = nn.BatchNorm1d(128) - self.bn3 = nn.BatchNorm1d(256) - self.bn4 = nn.BatchNorm1d(512) - - if self.use_deterministic_encoder: - self.fc1 = nn.Linear(512, 256) - self.fc2 = nn.Linear(256, 128) - self.fc_bn1 = nn.BatchNorm1d(256) - self.fc_bn2 = nn.BatchNorm1d(128) - self.fc3 = nn.Linear(128, self.zdim) - else: - # Mapping to [c], cmean - self.fc1_m = nn.Linear(512, 256) - self.fc2_m = nn.Linear(256, 128) - self.fc3_m = nn.Linear(128, self.zdim) - self.fc_bn1_m = nn.BatchNorm1d(256) - self.fc_bn2_m = nn.BatchNorm1d(128) - - # Mapping to [c], cmean - self.fc1_v = nn.Linear(512, 256) - self.fc2_v = nn.Linear(256, 128) - self.fc3_v = nn.Linear(128, self.zdim) - self.fc_bn1_v = nn.BatchNorm1d(256) - self.fc_bn2_v = nn.BatchNorm1d(128) - - def forward(self, x): - x = x.transpose(1, 2) - x = F.relu(self.bn1(self.conv1(x))) - x = F.relu(self.bn2(self.conv2(x))) - x = F.relu(self.bn3(self.conv3(x))) - x = self.bn4(self.conv4(x)) - x = torch.max(x, 2, keepdim=True)[0] - x = x.view(-1, 512) - - if self.use_deterministic_encoder: - ms = F.relu(self.fc_bn1(self.fc1(x))) - ms = F.relu(self.fc_bn2(self.fc2(ms))) - ms = self.fc3(ms) - m, v = ms, 0 - else: - m = F.relu(self.fc_bn1_m(self.fc1_m(x))) - m = F.relu(self.fc_bn2_m(self.fc2_m(m))) - m = self.fc3_m(m) - v = F.relu(self.fc_bn1_v(self.fc1_v(x))) - v = F.relu(self.fc_bn2_v(self.fc2_v(v))) - v = self.fc3_v(v) - - return m, v diff --git a/other_models/shapegf/models/generators/mlp_gen.py b/other_models/shapegf/models/generators/mlp_gen.py deleted file mode 100644 index 7670032..0000000 --- a/other_models/shapegf/models/generators/mlp_gen.py +++ /dev/null @@ -1,75 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - - -# Taken from https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/15 -def truncated_normal(tensor, mean=0, std=1, trunc_std=2): - size = tensor.shape - tmp = tensor.new_empty(size + (4,)).normal_() - valid = (tmp < trunc_std) & (tmp > -trunc_std) - ind = valid.max(-1, keepdim=True)[1] - tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1)) - tensor.data.mul_(std).add_(mean) - return tensor - - -class Generator(nn.Module): - - def __init__(self, cfg, cfgmodel): - super().__init__() - self.cfg = cfg - self.cfgmodel = cfgmodel - - self.inp_dim = cfgmodel.inp_dim - self.out_dim = cfgmodel.out_dim - self.use_bn = getattr(cfgmodel, "use_bn", False) - self.output_bn = getattr(cfgmodel, "output_bn", False) - self.dims = cfgmodel.dims - - curr_dim = self.inp_dim - self.layers = [] - self.bns = [] - for hid in self.dims: - self.layers.append(nn.Linear(curr_dim, hid)) - self.bns.append(nn.BatchNorm1d(hid)) - curr_dim = hid - self.layers = nn.ModuleList(self.layers) - self.bns = nn.ModuleList(self.bns) - self.out = nn.Linear(curr_dim, self.out_dim) - self.out_bn = nn.BatchNorm1d(self.out_dim) - self.prior_type = getattr(cfgmodel, "prior", "gaussian") - - def get_prior(self, bs): - if self.prior_type == "truncate_gaussian": - gaussian_scale = getattr(self.cfgmodel, "gaussian_scale", 1.) - truncate_std = getattr(self.cfgmodel, "truncate_std", 2.) - noise = (torch.randn(bs, self.inp_dim) * gaussian_scale).cuda() - noise = truncated_normal( - noise, mean=0, std=gaussian_scale, trunc_std=truncate_std) - return noise - elif self.prior_type == "gaussian": - gaussian_scale = getattr(self.cfgmodel, "gaussian_scale", 1.) - return torch.randn(bs, self.inp_dim).cuda() * gaussian_scale - - else: - raise NotImplementedError( - "Invalid prior type:%s" % self.prior_type) - - def forward(self, z=None, bs=None): - if z is None: - assert bs is not None - z = self.get_prior(bs).cuda() - - y = z - for layer, bn in zip(self.layers, self.bns): - y = layer(y) - if self.use_bn: - y = bn(y) - y = F.relu(y) - y = self.out(y) - - if self.output_bn: - y = self.out_bn(y) - return y - diff --git a/other_models/shapegf/shapenet_recon.yaml b/other_models/shapegf/shapenet_recon.yaml deleted file mode 100644 index 2ebbe60..0000000 --- a/other_models/shapegf/shapenet_recon.yaml +++ /dev/null @@ -1,76 +0,0 @@ -data: - type: datasets.pointflow_datasets - dataset_type: shapenet15k - num_workers: 2 - cates: ['all'] - tr_max_sample_points: 4096 - te_max_sample_points: 4096 - data_dir: data/ShapeNetCore.v2.PC15k - batch_size: 32 - dataset_scale: 1 - normalize_per_shape: False - normalize_std_per_axis: False - recenter_per_shape: True - -viz: - log_freq: 10 - viz_freq: 1000 - save_freq: 25 - val_freq: 500 - -trainer: - type: shapegf.trainers.ae_trainer_3D - epochs: 1200 - seed: 100 - sigma_begin: 1 - sigma_end: 0.01 - sigma_num: 10 - loss_type: "fix_sigma" - scale_gt: True - surface: True - mixture: False - opt_dec: - type: 'adam' - lr: 1e-3 # use bs*1e-5/8 - beta1: 0.9 - beta2: 0.999 - momentum: 0.9 # for SGD - weight_decay: 0. - scheduler: linear - step_epoch: 1000 - start_ratio: 0.35 - duration_ratio: 0.5 - opt_enc: - type: 'adam' - lr: 1e-3 - beta1: 0.9 - beta2: 0.999 - momentum: 0.9 # for SGD - weight_decay: 0. - scheduler: linear - step_epoch: 1000 - start_ratio: 0.35 - duration_ratio: 0.5 - -models: - scorenet: - type: shapegf.models.decoders.resnet_add - dim: 3 - out_dim: 3 - hidden_size: 256 - n_blocks: 24 - sigma_condition: True - param_likelihood: False - xyz_condition: True - z_dim: 256 - encoder: - type: shapegf.models.encoders.l3dp_encoder - use_deterministic_encoder: True - zdim: 256 - input_dim: 3 - -inference: - step_size_ratio: 1 - num_steps: 10 - num_points: 2048 - weight: 1 diff --git a/other_models/shapegf/trainers/__pycache__/ae_trainer_3D.cpython-36.pyc b/other_models/shapegf/trainers/__pycache__/ae_trainer_3D.cpython-36.pyc deleted file mode 100644 index fe5f597..0000000 Binary files a/other_models/shapegf/trainers/__pycache__/ae_trainer_3D.cpython-36.pyc and /dev/null differ diff --git a/other_models/shapegf/trainers/__pycache__/base_trainer.cpython-36.pyc b/other_models/shapegf/trainers/__pycache__/base_trainer.cpython-36.pyc deleted file mode 100644 index c9342e2..0000000 Binary files a/other_models/shapegf/trainers/__pycache__/base_trainer.cpython-36.pyc and /dev/null differ diff --git a/other_models/shapegf/trainers/ae_trainer_3D.py b/other_models/shapegf/trainers/ae_trainer_3D.py deleted file mode 100644 index 3de6a18..0000000 --- a/other_models/shapegf/trainers/ae_trainer_3D.py +++ /dev/null @@ -1,339 +0,0 @@ -import os -import tqdm -import torch -import importlib -import numpy as np -from shapegf.trainers.base_trainer import BaseTrainer -from shapegf.trainers.utils.vis_utils import visualize_point_clouds_3d, \ - visualize_procedure -from shapegf.trainers.utils.utils import get_opt, get_prior, \ - ground_truth_reconstruct_multi, set_random_seed - - -try: - from evaluation.evaluation_metrics import EMD_CD - eval_reconstruciton = True -except: # noqa - # Skip evaluation - eval_reconstruciton = False - - -def score_matching_loss(score_net, shape_latent, tr_pts, sigma): - bs, num_pts = tr_pts.size(0), tr_pts.size(1) - sigma = sigma.view(bs, 1, 1) - # perturbed_points = tr_pts + torch.randn_like(tr_pts) * sigma - perturbed_points = tr_pts + torch.randn_like(tr_pts) * sigma - - # For numerical stability, the network predicts the field in a normalized - # scale (i.e. the norm of the gradient is not scaled by `sigma`) - # As a result, when computing the ground truth for supervision, we are using - # its original scale without scaling by `sigma` - y_pred = score_net(perturbed_points, shape_latent) # field (B, #points, 3) - y_gtr = - (perturbed_points - tr_pts).view(bs, num_pts, -1) - - # The loss for each sigma is weighted - lambda_sigma = 1. / sigma - loss = 0.5 * ((y_gtr - y_pred) ** 2. * lambda_sigma).sum(dim=2).mean() - return { - "loss": loss, - "x": perturbed_points - } - - -class Trainer(BaseTrainer): - - def __init__(self, cfg): - super().__init__(cfg) - self.cfg = cfg - # self.args = args - set_random_seed(getattr(self.cfg.trainer, "seed", 666)) - - # The networks - sn_lib = importlib.import_module(cfg.models.scorenet.type) - self.score_net = sn_lib.Decoder(cfg, cfg.models.scorenet) - self.score_net.cuda() - print("ScoreNet:") - print(self.score_net) - - encoder_lib = importlib.import_module(cfg.models.encoder.type) - self.encoder = encoder_lib.Encoder(cfg.models.encoder) - self.encoder.cuda() - print("Encoder:") - print(self.encoder) - - # The optimizer - if not (hasattr(self.cfg.trainer, "opt_enc") and - hasattr(self.cfg.trainer, "opt_dec")): - self.cfg.trainer.opt_enc = self.cfg.trainer.opt - self.cfg.trainer.opt_dec = self.cfg.trainer.opt - - self.opt_enc, self.scheduler_enc = get_opt( - self.encoder.parameters(), self.cfg.trainer.opt_enc) - self.opt_dec, self.scheduler_dec = get_opt( - self.score_net.parameters(), self.cfg.trainer.opt_dec) - - # Sigmas - if hasattr(cfg.trainer, "sigmas"): - self.sigmas = cfg.trainer.sigmas - else: - self.sigma_begin = float(cfg.trainer.sigma_begin) - self.sigma_end = float(cfg.trainer.sigma_end) - self.num_classes = int(cfg.trainer.sigma_num) - self.sigmas = np.exp( - np.linspace(np.log(self.sigma_begin), - np.log(self.sigma_end), - self.num_classes)) - print("Sigma:, ", self.sigmas) - - # Prepare save directory - """ - os.makedirs(os.path.join(cfg.save_dir, "images"), exist_ok=True) - os.makedirs(os.path.join(cfg.save_dir, "checkpoints"), exist_ok=True) - os.makedirs(os.path.join(cfg.save_dir, "val"), exist_ok=True) - """ - - # Prepare variable for summy - self.oracle_res = None - - def multi_gpu_wrapper(self, wrapper): - self.encoder = wrapper(self.encoder) - self.score_net = wrapper(self.score_net) - - def epoch_end(self, epoch, writer=None, **kwargs): - if self.scheduler_dec is not None: - self.scheduler_dec.step(epoch=epoch) - if writer is not None: - writer.add_scalar( - 'train/opt_dec_lr', self.scheduler_dec.get_lr()[0], epoch) - if self.scheduler_enc is not None: - self.scheduler_enc.step(epoch=epoch) - if writer is not None: - writer.add_scalar( - 'train/opt_enc_lr', self.scheduler_enc.get_lr()[0], epoch) - - def update(self, data, data_gt, **kwargs): - if 'no_update' in kwargs: - no_update = kwargs['no_update'] - else: - no_update = False - if not no_update: - self.encoder.train() - self.score_net.train() - self.opt_enc.zero_grad() - self.opt_dec.zero_grad() - - # tr_pts = data['tr_points'].cuda() # (B, #points, 3)smn_ae_trainer.py - tr_pts = data.cuda() # (B, #points, 3)smn_ae_trainer.py - tr_gt_pts = data_gt.cuda() # (B, #points, 3)smn_ae_trainer.py - batch_size = tr_pts.size(0) - z_mu, z_sigma = self.encoder(tr_pts) - z = z_mu + 0 * z_sigma - - # Randomly sample sigma - labels = torch.randint( - 0, len(self.sigmas), (batch_size,), device=tr_pts.device) - used_sigmas = torch.tensor( - np.array(self.sigmas))[labels].float().view(batch_size, 1).cuda() - z = torch.cat((z, used_sigmas), dim=1) - - res = score_matching_loss(self.score_net, z, tr_gt_pts, used_sigmas) - loss = res['loss'] - """ - if not no_update: - loss.backward() - self.opt_enc.step() - self.opt_dec.step() - - return { - 'loss': loss.detach().cpu().item(), - 'x': res['x'].detach().cpu() # perturbed data - } - """ - return { - 'loss': loss, - 'x': res['x'] # perturbed data - } - - def log_train(self, train_info, train_data, writer=None, - step=None, epoch=None, visualize=False, **kwargs): - if writer is None: - return - - # Log training information to tensorboard - train_info = {k: (v.cpu() if not isinstance(v, float) else v) - for k, v in train_info.items()} - for k, v in train_info.items(): - if not ('loss' in k): - continue - if step is not None: - writer.add_scalar('train/' + k, v, step) - else: - assert epoch is not None - writer.add_scalar('train/' + k, v, epoch) - - if visualize: - with torch.no_grad(): - print("Visualize: %s" % step) - gtr = train_data['te_points'] # ground truth point cloud - inp = train_data['tr_points'] # input for encoder - ptb = train_info['x'] # perturbed data - num_vis = min( - getattr(self.cfg.viz, "num_vis_samples", 5), - gtr.size(0)) - - print("Recon:") - rec, rec_list = self.reconstruct( - inp[:num_vis].cuda(), num_points=inp.size(1)) - print("Ground truth recon:") - rec_gt, rec_gt_list = ground_truth_reconstruct_multi( - inp[:num_vis].cuda(), self.cfg) - # Overview - all_imgs = [] - for idx in range(num_vis): - img = visualize_point_clouds_3d( - [rec_gt[idx], rec[idx], gtr[idx], ptb[idx]], - ["rec_gt", "recon", "shape", "perturbed"]) - all_imgs.append(img) - img = np.concatenate(all_imgs, axis=1) - writer.add_image( - 'tr_vis/overview', torch.as_tensor(img), step) - - # Reconstruction gt procedure - img = visualize_procedure( - self.sigmas, rec_gt_list, gtr, num_vis, self.cfg, "Rec_gt") - writer.add_image( - 'tr_vis/rec_gt_process', torch.as_tensor(img), step) - - # Reconstruction procedure - img = visualize_procedure( - self.sigmas, rec_list, gtr, num_vis, self.cfg, "Rec") - writer.add_image( - 'tr_vis/rec_process', torch.as_tensor(img), step) - - def validate(self, test_loader, epoch, **kwargs): - if not eval_reconstruciton: - return {} - - print("Validation (reconstruction):") - all_ref, all_rec, all_smp, all_ref_denorm = [], [], [], [] - all_rec_gt, all_inp_denorm, all_inp = [], [], [] - for data in tqdm.tqdm(test_loader): - ref_pts = data['te_points'].cuda() - inp_pts = data['tr_points'].cuda() - m = data['mean'].cuda() - std = data['std'].cuda() - rec_pts, _ = self.reconstruct(inp_pts, num_points=inp_pts.size(1)) - - # denormalize - inp_pts_denorm = inp_pts.clone() * std + m - ref_pts_denorm = ref_pts.clone() * std + m - rec_pts = rec_pts * std + m - - all_inp.append(inp_pts) - all_inp_denorm.append(inp_pts_denorm.view(*inp_pts.size())) - all_ref_denorm.append(ref_pts_denorm.view(*ref_pts.size())) - all_rec.append(rec_pts.view(*ref_pts.size())) - all_ref.append(ref_pts) - - inp = torch.cat(all_inp, dim=0) - rec = torch.cat(all_rec, dim=0) - ref = torch.cat(all_ref, dim=0) - ref_denorm = torch.cat(all_ref_denorm, dim=0) - inp_denorm = torch.cat(all_inp_denorm, dim=0) - for name, arr in [ - ('inp', inp), ('rec', rec), ('ref', ref), - ('ref_denorm', ref_denorm), ('inp_denorm', inp_denorm)]: - np.save( - os.path.join( - self.cfg.save_dir, 'val', '%s_ep%d.npy' % (name, epoch)), - arr.detach().cpu().numpy() - ) - all_res = {} - - # Oracle CD/EMD, will compute only once - if self.oracle_res is None: - rec_res = EMD_CD(inp_denorm, ref_denorm, 1) - rec_res = { - ("val/rec/%s" % k): (v if isinstance(v, float) else v.item()) - for k, v in rec_res.items()} - all_res.update(rec_res) - print("Validation oracle (denormalize) Epoch:%d " % epoch, rec_res) - self.oracle_res = rec_res - else: - all_res.update(self.oracle_res) - - # Reconstruction CD/EMD - all_res = {} - rec_res = EMD_CD(rec, ref_denorm, 1) - rec_res = { - ("val/rec/%s" % k): (v if isinstance(v, float) else v.item()) - for k, v in rec_res.items()} - all_res.update(rec_res) - print("Validation Recon (denormalize) Epoch:%d " % epoch, rec_res) - - return all_res - - def save(self, epoch=None, step=None, appendix=None, **kwargs): - d = { - 'opt_enc': self.opt_enc.state_dict(), - 'opt_dec': self.opt_dec.state_dict(), - 'sn': self.score_net.state_dict(), - 'enc': self.encoder.state_dict(), - 'epoch': epoch, - 'step': step - } - if appendix is not None: - d.update(appendix) - save_name = "epoch_%s_iters_%s.pt" % (epoch, step) - path = os.path.join(self.cfg.save_dir, "checkpoints", save_name) - torch.save(d, path) - - def resume(self, path, strict=True, **kwargs): - ckpt = torch.load(path) - self.encoder.load_state_dict(ckpt['enc'], strict=strict) - self.score_net.load_state_dict(ckpt['sn'], strict=strict) - self.opt_enc.load_state_dict(ckpt['opt_enc']) - self.opt_dec.load_state_dict(ckpt['opt_dec']) - start_epoch = ckpt['epoch'] - return start_epoch - - def langevin_dynamics(self, z, num_points=2048): - with torch.no_grad(): - assert hasattr(self.cfg, "inference") - step_size_ratio = float(getattr( - self.cfg.inference, "step_size_ratio", 1)) - num_steps = int(getattr(self.cfg.inference, "num_steps", 5)) - num_points = int(getattr( - self.cfg.inference, "num_points", num_points)) - weight = float(getattr(self.cfg.inference, "weight", 1)) - sigmas = self.sigmas - - x_list = [] - self.score_net.eval() - x = get_prior(z.size(0), num_points, self.cfg.models.scorenet.dim) - x = x.to(z) - x_list.append(x.clone()) - for sigma in sigmas: - sigma = torch.ones((1,)).cuda() * sigma - z_sigma = torch.cat((z, sigma.expand(z.size(0), 1)), dim=1) - step_size = 2 * sigma ** 2 * step_size_ratio - for t in range(num_steps): - z_t = torch.randn_like(x) * weight - x += torch.sqrt(step_size) * z_t - grad = self.score_net(x, z_sigma) - grad = grad / sigma ** 2 - x += 0.5 * step_size * grad - x_list.append(x.clone()) - return x, x_list - - def sample(self, num_shapes=1, num_points=2048): - with torch.no_grad(): - z = torch.randn(num_shapes, self.cfg.models.encoder.zdim).cuda() - return self.langevin_dynamics(z, num_points=num_points) - - def reconstruct(self, inp, num_points=2048): - with torch.no_grad(): - self.encoder.eval() - z, _ = self.encoder(inp) - return self.langevin_dynamics(z, num_points=num_points) - diff --git a/other_models/shapegf/trainers/base_trainer.py b/other_models/shapegf/trainers/base_trainer.py deleted file mode 100644 index a7ab501..0000000 --- a/other_models/shapegf/trainers/base_trainer.py +++ /dev/null @@ -1,37 +0,0 @@ - -class BaseTrainer(): - - def __init__(self, cfg): - pass - - def update(self, data, **kwargs): - raise NotImplementedError("Trainer [update] not implemented.") - - def epoch_end(self, epoch, writer=None, **kwargs): - # Signal now that the epoch ends.... - pass - - def multi_gpu_wrapper(self, wrapper): - raise NotImplementedError("Trainer [multi_gpu_wrapper] not implemented.") - - def log_train(self, train_info, train_data, - writer=None, step=None, epoch=None, visualize=False, - **kwargs): - raise NotImplementedError("Trainer [log_train] not implemented.") - - def validate(self, test_loader, epoch, **kwargs): - raise NotImplementedError("Trainer [validate] not implemented.") - - def log_val(self, val_info, writer=None, step=None, epoch=None, **kwargs): - if writer is not None: - for k, v in val_info.items(): - if step is not None: - writer.add_scalar(k, v, step) - else: - writer.add_scalar(k, v, epoch) - - def save(self, epoch=None, step=None, appendix=None, **kwargs): - raise NotImplementedError("Trainer [save] not implemented.") - - def resume(self, path, strict=True, **kwargs): - raise NotImplementedError("Trainer [resume] not implemented.") diff --git a/other_models/shapegf/trainers/lgan_trainer_3D.py b/other_models/shapegf/trainers/lgan_trainer_3D.py deleted file mode 100644 index f36c3c9..0000000 --- a/other_models/shapegf/trainers/lgan_trainer_3D.py +++ /dev/null @@ -1,310 +0,0 @@ -import os -import tqdm -import torch -import random -import importlib -import numpy as np -from trainers.utils.utils import get_opt -from trainers.ae_trainer_3D import Trainer as BaseTrainer -from trainers.utils.gan_losses import gen_loss, dis_loss, gradient_penalty -from trainers.utils.vis_utils import visualize_procedure, \ - visualize_point_clouds_3d - - -try: - from evaluation.evaluation_metrics import compute_all_metrics - - eval_generation = True -except: # noqa - eval_generation = False - - -class Trainer(BaseTrainer): - - def __init__(self, cfg, args): - super().__init__(cfg, args) - - # Now initialize the GAN part - gen_lib = importlib.import_module(cfg.models.gen.type) - self.gen = gen_lib.Generator(cfg, cfg.models.gen) - self.gen.cuda() - print("Generator:") - print(self.gen) - - dis_lib = importlib.import_module(cfg.models.dis.type) - self.dis = dis_lib.Discriminator(cfg, cfg.models.dis) - self.dis.cuda() - print("Discriminator:") - print(self.dis) - - # Optimizers - if not (hasattr(self.cfg.trainer, "opt_gen") and - hasattr(self.cfg.trainer, "opt_dis")): - self.cfg.trainer.opt_gen = self.cfg.trainer.opt - self.cfg.trainer.opt_dis = self.cfg.trainer.opt - self.opt_gen, self.scheduler_gen = get_opt( - self.gen.parameters(), self.cfg.trainer.opt_gen) - self.opt_dis, self.scheduler_dis = get_opt( - self.dis.parameters(), self.cfg.trainer.opt_dis) - - # book keeping - self.total_iters = 0 - self.total_gan_iters = 0 - self.n_critics = getattr(self.cfg.trainer, "n_critics", 1) - self.gan_only = getattr(self.cfg.trainer, "gan_only", True) - - # If pretrained AE, then load it up - if hasattr(self.cfg.trainer, "ae_pretrained"): - ckpt = torch.load(self.cfg.trainer.ae_pretrained) - print(self.cfg.trainer.ae_pretrained) - strict = getattr(self.cfg.trainer, "resume_strict", True) - self.encoder.load_state_dict(ckpt['enc'], strict=strict) - self.score_net.load_state_dict(ckpt['sn'], strict=strict) - if getattr(self.cfg.trainer, "resume_opt", False): - self.opt_enc.load_state_dict(ckpt['opt_enc']) - self.opt_dec.load_state_dict(ckpt['opt_dec']) - self.gan_pass_update_enc = getattr( - self.cfg.trainer, "gan_pass_update_enc", False) - - def epoch_end(self, epoch, writer=None, **kwargs): - super().epoch_end(epoch, writer=writer) - - if self.scheduler_dis is not None: - self.scheduler_dis.step(epoch=epoch) - if writer is not None: - writer.add_scalar( - 'train/opt_dis_lr', self.scheduler_dis.get_lr()[0], epoch) - - if self.scheduler_gen is not None: - self.scheduler_gen.step(epoch=epoch) - if writer is not None: - writer.add_scalar( - 'train/opt_gen_lr', self.scheduler_gen.get_lr()[0], epoch) - - def _update_gan_(self, data, gen=False): - self.gen.train() - self.dis.train() - self.opt_gen.zero_grad() - self.opt_dis.zero_grad() - if self.gan_pass_update_enc: - self.encoder.train() - self.opt_enc.zero_grad() - else: - self.encoder.eval() - - tr_pts = data['tr_points'].cuda() # (B, #points, 3)smn_ae_trainer.py - x_real, _ = self.encoder(tr_pts) - batch_size = x_real.size(0) - x_fake = self.gen(bs=batch_size) - - x_inp = torch.cat([x_real, x_fake], dim=0) - d_res = self.dis(x_inp, return_all=True) - d_out = d_res['x'] - d_real = d_out[:batch_size, ...] - d_fake = d_out[batch_size:, ...] - - loss, loss_res = None, {} - loss_type = getattr(self.cfg.trainer, "gan_loss_type", "wgan") - if gen: - gen_loss_weight = getattr(self.cfg.trainer, "gen_loss_weight", 1.) - loss_gen, gen_loss_res = gen_loss( - d_real, d_fake, weight=gen_loss_weight, loss_type=loss_type) - loss = loss_gen + (0. if loss is None else loss) - loss.backward() - self.opt_gen.step() - if self.gan_pass_update_enc: - assert self.opt_enc is not None - self.opt_enc.step() - loss_res.update({ - ("train/gan_pass/gen/%s" % k): v - for k, v in gen_loss_res.items() - }) - loss_res['loss'] = loss.detach().cpu().item() - else: - # Get gradient penalty - gp_weight = getattr(self.cfg.trainer, 'gp_weight', 0.) - if gp_weight > 0: - gp_type = getattr(self.cfg.trainer, 'gp_type', "zero_center") - gp, gp_res = gradient_penalty( - x_real, x_fake, d_real, d_fake, - weight=gp_weight, gp_type=gp_type) - loss = gp + (0. if loss is None else loss) - loss_res.update({ - ("train/gan_pass/gp_loss/%s" % k): v - for k, v in gp_res.items() - }) - - dis_loss_weight = getattr(self.cfg.trainer, "dis_loss_weight", 1.) - loss_dis, dis_loss_res = dis_loss( - d_real, d_fake, weight=dis_loss_weight, loss_type=loss_type) - loss = loss_dis + (0. if loss is None else loss) - loss.backward() - self.opt_dis.step() - loss_res.update({ - ("train/gan_pass/dis/%s" % k): v for k, v in dis_loss_res.items() - }) - loss_res['loss'] = loss.detach().cpu().item() - - loss_res['x_real'] = x_real.clone().detach().cpu() - loss_res['x_fake'] = x_fake.clone().detach().cpu() - return loss_res - - def update_lgan(self, data): - self.total_gan_iters += 1 - res = {} - if self.total_gan_iters % self.n_critics == 0: - gen_res = self._update_gan_(data, gen=True) - res.update(gen_res) - dis_res = self._update_gan_(data, gen=False) - res.update(dis_res) - return res - - def update(self, data, *args, **kwargs): - res = {} - if not self.gan_only: - ae_res = super().update(data, *args, **kwargs) - res.update(ae_res) - gan_res = self.update_lgan(data) - res.update(gan_res) - return res - - def log_train(self, train_info, train_data, writer=None, - step=None, epoch=None, visualize=False, **kwargs): - with torch.no_grad(): - train_info.update(super().update(train_data, no_update=True)) - super().log_train(train_info, train_data, writer=writer, step=step, - epoch=epoch, visualize=visualize) - if step is not None: - writer.add_histogram('tr/latent_real', train_info['x_real'], step) - writer.add_histogram('tr/latent_fake', train_info['x_fake'], step) - else: - assert epoch is not None - writer.add_histogram('tr/latent_real', train_info['x_real'], epoch) - writer.add_histogram('tr/latent_fake', train_info['x_fake'], epoch) - - if visualize: - with torch.no_grad(): - print("Visualize generation results: %s" % step) - gtr = train_data['te_points'] # ground truth point cloud - inp = train_data['tr_points'] # input for encoder - num_vis = min( - getattr(self.cfg.viz, "num_vis_samples", 5), - gtr.size(0) - ) - smp, smp_list = self.sample(num_shapes=num_vis, - num_points=inp.size(1)) - - all_imgs = [] - for idx in range(num_vis): - img = visualize_point_clouds_3d( - [smp[idx], gtr[idx]], ["gen", "ref"]) - all_imgs.append(img) - img = np.concatenate(all_imgs, axis=1) - writer.add_image( - 'tr_vis/gen', torch.as_tensor(img), step) - - img = visualize_procedure( - self.sigmas, smp_list, gtr, num_vis, self.cfg, "gen") - writer.add_image( - 'tr_vis/gen_process', torch.as_tensor(img), step) - - def save(self, epoch=None, step=None, appendix=None, **kwargs): - d = { - 'opt_enc': self.opt_enc.state_dict(), - 'opt_dec': self.opt_dec.state_dict(), - 'opt_dis': self.opt_dis.state_dict(), - 'opt_gen': self.opt_gen.state_dict(), - 'sn': self.score_net.state_dict(), - 'enc': self.encoder.state_dict(), - 'dis': self.dis.state_dict(), - 'gen': self.gen.state_dict(), - 'epoch': epoch, - 'step': step - } - if appendix is not None: - d.update(appendix) - save_name = "epoch_%s_iters_%s.pt" % (epoch, step) - path = os.path.join(self.cfg.save_dir, "checkpoints", save_name) - torch.save(d, path) - - def resume(self, path, strict=True, **args): - ckpt = torch.load(path) - self.encoder.load_state_dict(ckpt['enc'], strict=strict) - self.score_net.load_state_dict(ckpt['sn'], strict=strict) - self.opt_enc.load_state_dict(ckpt['opt_enc']) - self.opt_dec.load_state_dict(ckpt['opt_dec']) - start_epoch = ckpt['epoch'] - - if 'gen' in ckpt: - self.gen.load_state_dict(ckpt['gen'], strict=strict) - if 'dis' in ckpt: - self.dis.load_state_dict(ckpt['dis'], strict=strict) - if 'opt_gen' in ckpt: - self.opt_gen.load_state_dict(ckpt['opt_gen']) - if 'opt_dis' in ckpt: - self.opt_dis.load_state_dict(ckpt['opt_dis']) - return start_epoch - - def sample(self, num_shapes=1, num_points=2048): - with torch.no_grad(): - self.gen.eval() - z = self.gen(bs=num_shapes) - return self.langevin_dynamics(z, num_points=num_points) - - def validate(self, test_loader, epoch, *args, **kwargs): - all_res = {} - - if eval_generation: - with torch.no_grad(): - print("l-GAN validation:") - all_ref, all_smp = [], [] - for data in tqdm.tqdm(test_loader): - ref_pts = data['te_points'].cuda() - inp_pts = data['tr_points'].cuda() - smp_pts, _ = self.sample( - num_shapes=inp_pts.size(0), - num_points=inp_pts.size(1), - ) - all_smp.append(smp_pts.view( - ref_pts.size(0), ref_pts.size(1), ref_pts.size(2))) - all_ref.append( - ref_pts.view(ref_pts.size(0), ref_pts.size(1), - ref_pts.size(2))) - - smp = torch.cat(all_smp, dim=0) - np.save( - os.path.join(self.cfg.save_dir, 'val', - 'smp_ep%d.npy' % epoch), - smp.detach().cpu().numpy() - ) - ref = torch.cat(all_ref, dim=0) - - # Sample CD/EMD - # step 1: subsample shapes - max_gen_vali_shape = int(getattr( - self.cfg.trainer, "max_gen_validate_shapes", - int(smp.size(0)))) - sub_sampled = random.sample( - range(smp.size(0)), min(smp.size(0), max_gen_vali_shape)) - smp_sub = smp[sub_sampled, ...].contiguous() - ref_sub = ref[sub_sampled, ...].contiguous() - - gen_res = compute_all_metrics( - smp_sub, ref_sub, - batch_size=int(getattr( - self.cfg.trainer, "val_metrics_batch_size", 100)), - accelerated_cd=True - ) - all_res = { - ("val/gen/%s" % k): - (v if isinstance(v, float) else v.item()) - for k, v in gen_res.items()} - print("Validation Sample (unit) Epoch:%d " % epoch, gen_res) - - - # Call super class validation - if getattr(self.cfg.trainer, "validate_recon", False): - all_res.update(super().validate( - test_loader, epoch, *args, **kwargs)) - - return all_res \ No newline at end of file diff --git a/other_models/shapegf/trainers/utils/__pycache__/utils.cpython-36.pyc b/other_models/shapegf/trainers/utils/__pycache__/utils.cpython-36.pyc deleted file mode 100644 index 666b796..0000000 Binary files a/other_models/shapegf/trainers/utils/__pycache__/utils.cpython-36.pyc and /dev/null differ diff --git a/other_models/shapegf/trainers/utils/__pycache__/vis_utils.cpython-36.pyc b/other_models/shapegf/trainers/utils/__pycache__/vis_utils.cpython-36.pyc deleted file mode 100644 index 86b34d4..0000000 Binary files a/other_models/shapegf/trainers/utils/__pycache__/vis_utils.cpython-36.pyc and /dev/null differ diff --git a/other_models/shapegf/trainers/utils/gan_losses.py b/other_models/shapegf/trainers/utils/gan_losses.py deleted file mode 100644 index 78375dd..0000000 --- a/other_models/shapegf/trainers/utils/gan_losses.py +++ /dev/null @@ -1,89 +0,0 @@ -import torch - - -def gen_loss(d_real, d_fake, loss_type="wgan", weight=1., **kwargs): - if loss_type.lower() == "wgan": - wg_loss_orig = - d_fake.mean() - wg_loss = wg_loss_orig * weight - return wg_loss, { - "wgan_gen_loss": wg_loss.clone().detach().item(), - "wgan_gen_loss_orig": wg_loss_orig.clone().detach().item(), - } - elif loss_type.lower() == "hinge": - g_loss = -d_fake.mean() - d_correct = (d_real >= 0.).float().sum() + (d_fake < 0.).float().sum() - d_acc = d_correct / float(d_real.size(0) + d_fake.size(0)) - - loss = weight * g_loss - return loss, { - 'loss': loss.clone().detach(), - "dis_acc": d_acc.clone().detach(), - "dis_correct": d_correct.clone().detach(), - 'g_loss': g_loss.clone().detach() - } - - else: - raise NotImplementedError("Not implement: %s" % loss_type) - - -def dis_loss(d_real, d_fake, loss_type="wgan", weight=1., **kwargs): - if loss_type.lower() == "wgan": - loss_fake = d_fake.mean() - loss_real = d_real.mean() - wg_loss_orig = loss_fake - loss_real - wg_loss = wg_loss_orig * weight - return wg_loss, { - "wgan_dis_loss": wg_loss.clone().detach().item(), - "wgan_dis_loss_orig": wg_loss_orig.clone().detach().item(), - "wgan_dis_loss_real": loss_real.clone().detach().item(), - "wgan_dis_loss_fake": loss_fake.clone().detach().item() - } - elif loss_type.lower() == "hinge": - d_loss_real = -torch.min(d_real - 1, d_real * 0).mean() - d_loss_fake = -torch.min(-d_fake - 1, d_fake * 0).mean() - d_correct = (d_real >= 0.).float().sum() + (d_fake < 0.).float().sum() - d_acc = d_correct / float(d_real.size(0) + d_fake.size(0)) - - d_loss = d_loss_real + d_loss_fake - loss = d_loss * weight - return loss, { - "loss": loss.clone().detach(), - "d_loss": d_loss.clone().detach(), - "dis_acc": d_acc.clone().detach(), - "dis_correct": d_correct.clone().detach(), - "loss_real": d_loss_real.clone().detach(), - "loss_fake": d_loss_fake.clone().detach(), - } - else: - raise NotImplementedError("Not implement: %s" % loss_type) - - -def dis_acc(d_real, d_fake, loss_type="wgan", **kwargs): - if loss_type.lower() == "wgan": - # No threshold, don't know which one is correct which is not - return {} - elif loss_type.lower() == "hinge": - return {} - else: - raise NotImplementedError("Not implement: %s" % loss_type) - - -def gradient_penalty(x_real, x_fake, d_real, d_fake, - weight=1., gp_type='zero_center', eps=1e-8): - if gp_type == "zero_center": - bs = d_real.size(0) - grad = torch.autograd.grad( - outputs=d_real, inputs=x_real, - grad_outputs=torch.ones_like(d_real).to(d_real), - create_graph=True, retain_graph=True)[0] - # [grad] should be either (B, D) or (B, #points, D) - grad = grad.reshape(bs, -1) - grad_norm = gp_orig = torch.sqrt(torch.sum(grad ** 2, dim=1)).mean() - gp = gp_orig ** 2. * weight - return gp, { - 'gp': gp.clone().detach().cpu(), - 'gp_orig': gp_orig.clone().detach().cpu(), - 'grad_norm': grad_norm.clone().detach().cpu() - } - else: - raise NotImplemented("Invalid gp type:%s" % gp_type) diff --git a/other_models/shapegf/trainers/utils/utils.py b/other_models/shapegf/trainers/utils/utils.py deleted file mode 100644 index 7f76a9f..0000000 --- a/other_models/shapegf/trainers/utils/utils.py +++ /dev/null @@ -1,145 +0,0 @@ -import torch -import random -import numpy as np -from torch import optim - - -def get_opt(params, cfgopt): - if cfgopt.type == 'adam': - optimizer = optim.Adam(params, lr=float(cfgopt.lr), - betas=(cfgopt.beta1, cfgopt.beta2), - weight_decay=cfgopt.weight_decay) - elif cfgopt.type == 'sgd': - optimizer = torch.optim.SGD( - params, lr=float(cfgopt.lr), momentum=cfgopt.momentum) - else: - assert 0, "Optimizer type should be either 'adam' or 'sgd'" - - scheduler = None - scheduler_type = getattr(cfgopt, "scheduler", None) - if scheduler_type is not None: - if scheduler_type == 'exponential': - decay = float(getattr(cfgopt, "step_decay", 0.1)) - scheduler = optim.lr_scheduler.ExponentialLR(optimizer, decay) - elif scheduler_type == 'step': - step_size = int(getattr(cfgopt, "step_epoch", 500)) - decay = float(getattr(cfgopt, "step_decay", 0.1)) - scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=decay) - elif scheduler_type == 'linear': - step_size = int(getattr(cfgopt, "step_epoch", 2000)) - final_ratio = float(getattr(cfgopt, "final_ratio", 0.01)) - start_ratio = float(getattr(cfgopt, "start_ratio", 0.5)) - duration_ratio = float(getattr(cfgopt, "duration_ratio", 0.45)) - - def lambda_rule(ep): - lr_l = 1.0 - min(1, max(0, ep - start_ratio * step_size) / float(duration_ratio * step_size)) * (1 - final_ratio) - return lr_l - - scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) - - elif scheduler_type == 'cosine_anneal_nocycle': - final_lr_ratio = float(getattr(cfgopt, "final_lr_ratio", 0.01)) - eta_min = float(cfgopt.lr) * final_lr_ratio - eta_max = float(cfgopt.lr) - - total_epoch = int(getattr(cfgopt, "step_epoch", 2000)) - start_ratio = float(getattr(cfgopt, "start_ratio", 0.2)) - T_max = total_epoch * (1 - start_ratio) - - def lambda_rule(ep): - curr_ep = max(0., ep - start_ratio * total_epoch) - lr = eta_min + 0.5 * (eta_max - eta_min) * (1 + np.cos(np.pi * curr_ep / T_max)) - lr_l = lr / eta_max - return lr_l - - scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) - - else: - assert 0, "args.schedulers should be either 'exponential' or 'linear' or 'step'" - return optimizer, scheduler - - -def set_random_seed(seed): - """set random seed""" - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -def ground_truth_field(prior_points, tr_pts, sigma): - bs, num_pts = tr_pts.size(0), tr_pts.size(1) - smp_pts = prior_points.size(1) - prior_points = prior_points.view(bs, smp_pts, 1, -1) - tr_pts = tr_pts.view(bs, 1, num_pts, -1) - dist = (prior_points - tr_pts).norm(dim=3, keepdim=True) ** 2. - a = - dist / sigma ** 2. - max_a, _ = torch.max(a, dim=2, keepdim=True) - diff = torch.exp(a - max_a) - w_i = diff / diff.sum(dim=2, keepdim=True) - - # (bs, #pts-prior, 1, dim) - trg_pts = (w_i * tr_pts).sum(dim=2, keepdim=True) - y = - ((prior_points - trg_pts) / sigma ** 2.).view(bs, smp_pts, -1) - return y - - -def ground_truth_reconstruct(inp, sigma, step_size, num_points=2048, - num_steps=100, decay=1, interval=10, weight=1): - with torch.no_grad(): - x = get_prior(inp.size(0), inp.size(1), inp.size(-1)).cuda() - x_list = [] - x_list.append(x.clone()) - - for t in range(num_steps): - z_t = torch.randn_like(x) * weight - x += np.sqrt(step_size) * z_t - grad = ground_truth_field(x, inp, sigma) - x += 0.5 * step_size * grad - if t % (num_steps // interval) == 0: - step_size *= decay - x_list.append(x.clone()) - return x, x_list - - -def ground_truth_reconstruct_multi(inp, cfg): - with torch.no_grad(): - assert hasattr(cfg, "inference") - step_size_ratio = float(getattr(cfg.inference, "step_size_ratio", 1)) - num_steps = int(getattr(cfg.inference, "num_steps", 5)) - num_points = int(getattr(cfg.inference, "num_points", inp.size(1))) - weight = float(getattr(cfg.inference, "weight", 1)) - - x = get_prior( - inp.size(0), num_points, cfg.models.scorenet.dim).cuda() - if hasattr(cfg.trainer, "sigmas"): - sigmas = cfg.trainer.sigmas - else: - sigma_begin = float(cfg.trainer.sigma_begin) - sigma_end = float(cfg.trainer.sigma_end) - num_classes = int(cfg.trainer.sigma_num) - sigmas = np.exp(np.linspace(np.log(sigma_begin), np.log(sigma_end), - num_classes)) - x_list = [] - x_list.append(x.clone()) - bs, num_pts = x.size(0), x.size(1) - - for sigma in sigmas: - sigma = torch.ones((1,)) * sigma - sigma = sigma.cuda() - step_size = 2 * sigma ** 2 * step_size_ratio - for t in range(num_steps): - z_t = torch.randn_like(x) * weight - x += torch.sqrt(step_size) * z_t - grad = ground_truth_field(x, inp, sigma) - x += 0.5 * step_size * grad - x_list.append(x.clone()) - return x, x_list - - -def get_prior(batch_size, num_points, inp_dim): - # -1 to 1, uniform - return (torch.rand(batch_size, num_points, inp_dim) * 2 - 1.) * 1.5 - - diff --git a/other_models/shapegf/trainers/utils/vis_utils.py b/other_models/shapegf/trainers/utils/vis_utils.py deleted file mode 100644 index 228e6b6..0000000 --- a/other_models/shapegf/trainers/utils/vis_utils.py +++ /dev/null @@ -1,180 +0,0 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np -import matplotlib - -matplotlib.use('Agg') -import matplotlib.pyplot as plt -import torch -import matplotlib.cm as cm -from mpl_toolkits.mplot3d import Axes3D # nofa: #401 - - -# Visualization -def visualize_point_clouds_3d(pcl_lst, title_lst=None): - # pts, gtr, inp): - pcl_lst = [pcl.cpu().detach().numpy() for pcl in pcl_lst] - if title_lst is None: - title_lst = [""] * len(pcl_lst) - - fig = plt.figure(figsize=(3 * len(pcl_lst), 3)) - for idx, (pts, title) in enumerate(zip(pcl_lst, title_lst)): - ax1 = fig.add_subplot(1, len(pcl_lst), 1 + idx, projection='3d') - ax1.set_title(title) - ax1.scatter(pts[:, 0], pts[:, 1], pts[:, 2], s=5) - ax1.set_xlim(-1, 1) - ax1.set_ylim(-1, 1) - ax1.set_zlim(-1, 1) - fig.canvas.draw() - - # grab the pixel buffer and dump it into a numpy array - res = np.array(fig.canvas.renderer._renderer) - res = np.transpose(res, (2, 0, 1)) - - plt.close() - return res - - -# Visualization -def visualize_point_clouds_3d_scan(pcl_lst, title_lst=None): - # pts, gtr, inp): - pcl_lst = [pcl.cpu().detach().numpy() for pcl in pcl_lst] - if title_lst is None: - title_lst = [""] * len(pcl_lst) - - fig = plt.figure(figsize=(3 * len(pcl_lst), 3)) - for idx, (pts, title) in enumerate(zip(pcl_lst, title_lst)): - ax1 = fig.add_subplot(1, len(pcl_lst), 1 + idx, projection='3d') - ax1.set_title(title) - ax1.scatter(pts[:, 0], pts[:, 1], pts[:, 2], c=pts[:, 1], s=0.5) - # print(min(pts[:, 0]), max(pts[:, 0]), min(pts[:, 1]), max(pts[:, 1]), min(pts[:, 2]), max(pts[:, 2])) - ax1.set_xlim(-1, 1) - ax1.set_ylim(-1, 1) - ax1.set_zlim(-1, 1) - fig.canvas.draw() - - # grab the pixel buffer and dump it into a numpy array - res = np.array(fig.canvas.renderer._renderer) - res = np.transpose(res, (2, 0, 1)) - - plt.close() - return res - - -# Visualization moving field and likelihood -def get_grid(x, k=10): - # TODO: set the range of field - # x = self.get_prior( - # 1, num_points, self.cfg.models.scorenet.dim).cuda() - # ind_x = np.arange(x[:,:,0].min(),x[:,:,0].max(),3/k) - # ind_y = np.arange(x[:,:,1].min(),x[:,:,0].max(),3/k) - ind_x = np.arange(-1.5, 1.5, 3 / k) - ind_y = np.arange(-1.5, 1.5, 3 / k) - X, Y = np.meshgrid(ind_x, ind_y) - X = torch.tensor(X).view(k * k).to(x) - Y = torch.tensor(Y).view(k * k).to(x) - - point_grid = torch.ones((1, k * k, 2), dtype=torch.double).to(x) - point_grid[0, :, 1] = point_grid[0, :, 1] * X - point_grid[0, :, 0] = point_grid[0, :, 0] * Y - point_grid = point_grid.float() - point_grid = point_grid.expand(x.size(0), -1, -1) - return point_grid - - -def visualize_point_clouds_2d_overlay(pcl_lst, title_lst=None, path=None): - # pts, gtr, inp): - pcl_lst = [pcl.cpu().detach().numpy() for pcl in pcl_lst] - if title_lst is None: - title_lst = [""] * len(pcl_lst) - - fig = plt.figure(figsize=(3, 3)) - ax1 = fig.add_subplot(1, 1, 1) - ax1.set_title(title_lst[0]) - for idx, pts in enumerate(pcl_lst): - ax1.scatter(pts[:, 0], pts[:, 1], s=5) - fig.canvas.draw() - - # grab the pixel buffer and dump it into a numpy array - res = np.array(fig.canvas.renderer._renderer) - res = np.transpose(res, (2, 0, 1)) - if path: - plt.savefig(path) - plt.close() - return res - - -def visualize_field(gtr, grid, field, k, label='field'): - grid_ = np.reshape(grid.cpu().detach().numpy(), (1, k * k, 2)) - if field.size(-1) == 2: - field = np.reshape(field.cpu().detach().numpy(), (1, k * k, 2)) - fig = plt.figure(figsize=(int(k / 100) * 2, int(k / 100))) - plt.title(label) - cs = fig.add_subplot(1, 2, 1) - field_val = np.sqrt(np.reshape((field ** 2).sum(axis=-1), (1, k * k, 1))) - else: - fig = plt.figure(figsize=(int(k / 100), int(k / 100))) - plt.title(label) - cs = fig.add_subplot(1, 1, 1) - field_val = np.reshape(field.cpu().detach().numpy(), (1, k * k, 1)) - - gt = gtr.cpu().detach().numpy() - - for i in range(np.shape(field_val)[0]): - # cs = fig.add_subplot(1, 2, 1) - X = np.reshape(grid_[i, :, 0], (k, k)) - Y = np.reshape(grid_[i, :, 1], (k, k)) - cs.contourf(X, Y, np.reshape(field_val[i, :], (k, k)), 20, - vmin=min(field_val[i, :]), vmax=max(field_val[i, :]), - cmap=cm.coolwarm) - print(min(field_val[i, :]), max(field_val[i, :])) - m = plt.cm.ScalarMappable(cmap=cm.coolwarm) - m.set_array(np.reshape(field_val[i, :], (k, k))) - m.set_clim(min(field_val[i, :]), max(field_val[i, :])) - - if np.shape(field)[-1] == 2: - for i in range(np.shape(field_val)[0]): - ax = fig.add_subplot(1, 2, 2) - scale = 20 - indx = np.array([np.arange(0, k, scale) + t * k for t in range(0, k, scale)]) - X = np.reshape(grid_[i, indx, 0], int(k * k / scale / scale)) - Y = np.reshape(grid_[i, indx, 1], int(k * k / scale / scale)) - u = np.reshape(field[i, indx, 0], int(k * k / scale / scale)) - v = np.reshape(field[i, indx, 1], int(k * k / scale / scale)) - - color = np.sqrt(v ** 2 + u ** 2) - field_norm = field / field_val - u = np.reshape(field_norm[i, indx, 0], int(k * k / scale / scale)) - v = np.reshape(field_norm[i, indx, 1], int(k * k / scale / scale)) - ax.quiver(X, Y, u, v, color, alpha=0.8, cmap=cm.coolwarm) - ax.xaxis.set_ticks([]) - ax.yaxis.set_ticks([]) - ax.set_aspect('equal') - ax.scatter(gt[:, 0], gt[:, 1], s=1, color='r') - - fig.canvas.draw() - # grab the pixel buffer and dump it into a numpy array - res = np.array(fig.canvas.renderer._renderer) - res = np.transpose(res, (2, 0, 1)) - - plt.close() - return res - - -def visualize_procedure(sigmas, fig_list, gtr, num_vis, cfg, name="Rec_gt"): - all_imgs = [] - sigmas = np.append([0], sigmas) - for idx in range(num_vis): - img = visualize_point_clouds_3d( - [fig_list[i][idx] for i in - range(0, len(fig_list), 1)] + [gtr[idx]], - [(name + " step" + - str(i * int(getattr(cfg.inference, "num_steps", 5))) + - " sigma%.3f" % sigmas[i]) - for i in range(0, len(fig_list), 1)] + ["gt shape"]) - all_imgs.append(img) - img = np.concatenate(all_imgs, axis=1) - return img - diff --git a/other_models/snowflake/__init__.py b/other_models/snowflake/__init__.py deleted file mode 100644 index 0b5771c..0000000 --- a/other_models/snowflake/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -import sys -sys.path.append('../pointnet2_ops_lib') -sys.path.append('..') \ No newline at end of file diff --git a/other_models/snowflake/__pycache__/__init__.cpython-36.pyc b/other_models/snowflake/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index 03ab493..0000000 Binary files a/other_models/snowflake/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/other_models/snowflake/__pycache__/model.cpython-36.pyc b/other_models/snowflake/__pycache__/model.cpython-36.pyc deleted file mode 100644 index 6ed617a..0000000 Binary files a/other_models/snowflake/__pycache__/model.cpython-36.pyc and /dev/null differ diff --git a/other_models/snowflake/__pycache__/skip_transformer.cpython-36.pyc b/other_models/snowflake/__pycache__/skip_transformer.cpython-36.pyc deleted file mode 100644 index 9c3e929..0000000 Binary files a/other_models/snowflake/__pycache__/skip_transformer.cpython-36.pyc and /dev/null differ diff --git a/other_models/snowflake/__pycache__/utils.cpython-36.pyc b/other_models/snowflake/__pycache__/utils.cpython-36.pyc deleted file mode 100644 index 0484413..0000000 Binary files a/other_models/snowflake/__pycache__/utils.cpython-36.pyc and /dev/null differ diff --git a/other_models/snowflake/model.py b/other_models/snowflake/model.py deleted file mode 100644 index 1d8c45f..0000000 --- a/other_models/snowflake/model.py +++ /dev/null @@ -1,202 +0,0 @@ -#! /usr/bin/python3 -# -*- coding: utf-8 -*- -# @Author: Peng Xiang - -import torch -import torch.nn as nn -from snowflake.utils import PointNet_SA_Module_KNN, MLP_Res, MLP_CONV, fps_subsample, Transformer -from snowflake.skip_transformer import SkipTransformer - - -class FeatureExtractor(nn.Module): - def __init__(self, out_dim=1024): - """Encoder that encodes information of partial point cloud - """ - super(FeatureExtractor, self).__init__() - self.sa_module_1 = PointNet_SA_Module_KNN(512, 16, 3, [64, 128], group_all=False, if_bn=False, if_idx=True) - self.transformer_1 = Transformer(128, dim=64) - self.sa_module_2 = PointNet_SA_Module_KNN(128, 16, 128, [128, 256], group_all=False, if_bn=False, if_idx=True) - self.transformer_2 = Transformer(256, dim=64) - self.sa_module_3 = PointNet_SA_Module_KNN(None, None, 256, [512, out_dim], group_all=True, if_bn=False) - - def forward(self, point_cloud): - """ - Args: - point_cloud: b, 3, n - - Returns: - l3_points: (B, out_dim, 1) - """ - l0_xyz = point_cloud - l0_points = point_cloud - - l1_xyz, l1_points, idx1 = self.sa_module_1(l0_xyz, l0_points) # (B, 3, 512), (B, 128, 512) - l1_points = self.transformer_1(l1_points, l1_xyz) - l2_xyz, l2_points, idx2 = self.sa_module_2(l1_xyz, l1_points) # (B, 3, 128), (B, 256, 512) - l2_points = self.transformer_2(l2_points, l2_xyz) - l3_xyz, l3_points = self.sa_module_3(l2_xyz, l2_points) # (B, 3, 1), (B, out_dim, 1) - - return l3_points - - -class SeedGenerator(nn.Module): - def __init__(self, dim_feat=512, num_pc=256): - super(SeedGenerator, self).__init__() - self.ps = nn.ConvTranspose1d(dim_feat, 128, num_pc, bias=True) - self.mlp_1 = MLP_Res(in_dim=dim_feat + 128, hidden_dim=128, out_dim=128) - self.mlp_2 = MLP_Res(in_dim=128, hidden_dim=64, out_dim=128) - self.mlp_3 = MLP_Res(in_dim=dim_feat + 128, hidden_dim=128, out_dim=128) - self.mlp_4 = nn.Sequential( - nn.Conv1d(128, 64, 1), - nn.ReLU(), - nn.Conv1d(64, 3, 1) - ) - - def forward(self, feat): - """ - Args: - feat: Tensor (b, dim_feat, 1) - """ - x1 = self.ps(feat) # (b, 128, 256) - x1 = self.mlp_1(torch.cat([x1, feat.repeat((1, 1, x1.size(2)))], 1)) - x2 = self.mlp_2(x1) - x3 = self.mlp_3(torch.cat([x2, feat.repeat((1, 1, x2.size(2)))], 1)) # (b, 128, 256) - completion = self.mlp_4(x3) # (b, 3, 256) - return completion - - -class SPD(nn.Module): - def __init__(self, dim_feat=512, up_factor=2, i=0, radius=1): - """Snowflake Point Deconvolution""" - super(SPD, self).__init__() - self.i = i - self.up_factor = up_factor - self.radius = radius - self.mlp_1 = MLP_CONV(in_channel=3, layer_dims=[64, 128]) - self.mlp_2 = MLP_CONV(in_channel=128 * 2 + dim_feat, layer_dims=[256, 128]) - - self.skip_transformer = SkipTransformer(in_channel=128, dim=64) - - self.mlp_ps = MLP_CONV(in_channel=128, layer_dims=[64, 32]) - self.ps = nn.ConvTranspose1d(32, 128, up_factor, up_factor, bias=False) # point-wise splitting - - self.up_sampler = nn.Upsample(scale_factor=up_factor) - self.mlp_delta_feature = MLP_Res(in_dim=256, hidden_dim=128, out_dim=128) - - self.mlp_delta = MLP_CONV(in_channel=128, layer_dims=[64, 3]) - - def forward(self, pcd_prev, feat_global, K_prev=None): - """ - Args: - pcd_prev: Tensor, (B, 3, N_prev) - feat_global: Tensor, (B, dim_feat, 1) - K_prev: Tensor, (B, 128, N_prev) - - Returns: - pcd_child: Tensor, up sampled point cloud, (B, 3, N_prev * up_factor) - K_curr: Tensor, displacement feature of current step, (B, 128, N_prev * up_factor) - """ - b, _, n_prev = pcd_prev.shape - feat_1 = self.mlp_1(pcd_prev) - feat_1 = torch.cat([feat_1, - torch.max(feat_1, 2, keepdim=True)[0].repeat((1, 1, feat_1.size(2))), - feat_global.repeat(1, 1, feat_1.size(2))], 1) - Q = self.mlp_2(feat_1) - - H = self.skip_transformer(pcd_prev, K_prev if K_prev is not None else Q, Q) - - feat_child = self.mlp_ps(H) - feat_child = self.ps(feat_child) # (B, 128, N_prev * up_factor) - H_up = self.up_sampler(H) - K_curr = self.mlp_delta_feature(torch.cat([feat_child, H_up], 1)) - - delta = torch.tanh(self.mlp_delta(torch.relu(K_curr))) / self.radius**self.i # (B, 3, N_prev * up_factor) - pcd_child = self.up_sampler(pcd_prev) - pcd_child = pcd_child + delta - - return pcd_child, K_curr - - -class Decoder(nn.Module): - def __init__(self, dim_feat=512, num_pc=256, num_p0=512, radius=1, up_factors=None): - super(Decoder, self).__init__() - self.num_p0 = num_p0 - self.decoder_coarse = SeedGenerator(dim_feat=dim_feat, num_pc=num_pc) - if up_factors is None: - up_factors = [1] - else: - up_factors = [1] + up_factors - - uppers = [] - for i, factor in enumerate(up_factors): - uppers.append(SPD(dim_feat=dim_feat, up_factor=factor, i=i, radius=radius)) - - self.uppers = nn.ModuleList(uppers) - - def forward(self, feat, partial=[], return_P0=False): - """ - Args: - feat: Tensor, (b, dim_feat, n) - partial: Tensor, (b, n, 3) - """ - arr_pcd = [] - pcd = self.decoder_coarse(feat).permute(0, 2, 1).contiguous() # (B, num_pc, 3) - arr_pcd.append(pcd) - - if partial != []: - pcd = fps_subsample(torch.cat([pcd, partial], 1), self.num_p0) - else: - pcd = fps_subsample(pcd, self.num_p0) - if return_P0: - arr_pcd.append(pcd) - K_prev = None - pcd = pcd.permute(0, 2, 1).contiguous() - for upper in self.uppers: - pcd, K_prev = upper(pcd, feat, K_prev) - arr_pcd.append(pcd.permute(0, 2, 1).contiguous()) - - return arr_pcd - - -class SnowflakeNet(nn.Module): - def __init__(self, dim_feat=512, num_pc=256, num_p0=512, radius=1, up_factors=None, global_feat=False): - """ - Args: - dim_feat: int, dimension of global feature - num_pc: int - num_p0: int - radius: searching radius - up_factors: list of int - """ - super(SnowflakeNet, self).__init__() - self.feat_extractor = FeatureExtractor(out_dim=dim_feat) - # NOTE: just one type of input - # self.decoder = Decoder(dim_feat=dim_feat, num_pc=num_pc, num_p0=num_p0, radius=radius, up_factors=up_factors) - # NOTE: dim_feat=dim_feat*2 if image and point cloud are fed in the same time - if global_feat == False: - self.decoder = Decoder(dim_feat=dim_feat, num_pc=num_pc, num_p0=num_p0, radius=radius, up_factors=up_factors) - elif global_feat == True: - self.decoder = Decoder(dim_feat=dim_feat*2, num_pc=num_pc, num_p0=num_p0, radius=radius, up_factors=up_factors) - - # NOTE: original codes - # def forward(self, point_cloud, return_P0=False): - def forward(self, point_cloud=[], global_feature=[], return_P0=False): - """ - Args: - point_cloud: (B, N, 3) - """ - if point_cloud != [] and global_feature != []: - pcd_bnc = point_cloud - point_cloud = point_cloud.permute(0, 2, 1).contiguous() - feat = self.feat_extractor(point_cloud) - feat = torch.cat((feat, global_feature), 1) - out = self.decoder(feat, pcd_bnc, return_P0=return_P0) - elif point_cloud != []: - pcd_bnc = point_cloud - point_cloud = point_cloud.permute(0, 2, 1).contiguous() - feat = self.feat_extractor(point_cloud) - out = self.decoder(feat, pcd_bnc, return_P0=return_P0) - elif global_feature != []: - feat = global_feature - out = self.decoder(feat, return_P0=return_P0) - return out diff --git a/other_models/snowflake/skip_transformer.py b/other_models/snowflake/skip_transformer.py deleted file mode 100644 index 9a00674..0000000 --- a/other_models/snowflake/skip_transformer.py +++ /dev/null @@ -1,70 +0,0 @@ -#! /usr/bin/python3 -# -*- coding: utf-8 -*- -# @Author: Peng Xiang - -import torch -from torch import nn, einsum -from snowflake.utils import MLP_Res, grouping_operation, query_knn - - -class SkipTransformer(nn.Module): - def __init__(self, in_channel, dim=256, n_knn=16, pos_hidden_dim=64, attn_hidden_multiplier=4): - super(SkipTransformer, self).__init__() - self.mlp_v = MLP_Res(in_dim=in_channel*2, hidden_dim=in_channel, out_dim=in_channel) - self.n_knn = n_knn - self.conv_key = nn.Conv1d(in_channel, dim, 1) - self.conv_query = nn.Conv1d(in_channel, dim, 1) - self.conv_value = nn.Conv1d(in_channel, dim, 1) - - self.pos_mlp = nn.Sequential( - nn.Conv2d(3, pos_hidden_dim, 1), - nn.BatchNorm2d(pos_hidden_dim), - nn.ReLU(), - nn.Conv2d(pos_hidden_dim, dim, 1) - ) - - self.attn_mlp = nn.Sequential( - nn.Conv2d(dim, dim * attn_hidden_multiplier, 1), - nn.BatchNorm2d(dim * attn_hidden_multiplier), - nn.ReLU(), - nn.Conv2d(dim * attn_hidden_multiplier, dim, 1) - ) - - self.conv_end = nn.Conv1d(dim, in_channel, 1) - - def forward(self, pos, key, query, include_self=True): - """ - Args: - pos: (B, 3, N) - key: (B, in_channel, N) - query: (B, in_channel, N) - include_self: boolean - - Returns: - Tensor: (B, in_channel, N), shape context feature - """ - value = self.mlp_v(torch.cat([key, query], 1)) - identity = value - key = self.conv_key(key) - query = self.conv_query(query) - value = self.conv_value(value) - b, dim, n = value.shape - - pos_flipped = pos.permute(0, 2, 1).contiguous() - idx_knn = query_knn(self.n_knn, pos_flipped, pos_flipped, include_self=include_self) - - key = grouping_operation(key, idx_knn) # b, dim, n, n_knn - qk_rel = query.reshape((b, -1, n, 1)) - key - - pos_rel = pos.reshape((b, -1, n, 1)) - grouping_operation(pos, idx_knn) # b, 3, n, n_knn - pos_embedding = self.pos_mlp(pos_rel) - - attention = self.attn_mlp(qk_rel + pos_embedding) # b, dim, n, n_knn - attention = torch.softmax(attention, -1) - - value = value.reshape((b, -1, n, 1)) + pos_embedding # - - agg = einsum('b c i j, b c i j -> b c i', attention, value) # b, dim, n - y = self.conv_end(agg) - - return y + identity diff --git a/other_models/snowflake/utils.py b/other_models/snowflake/utils.py deleted file mode 100644 index a53fd94..0000000 --- a/other_models/snowflake/utils.py +++ /dev/null @@ -1,484 +0,0 @@ -#! /usr/bin/python3 -# -*- coding: utf-8 -*- -# @Author: Peng Xiang - -import torch -from torch import nn, einsum -# from pointnet2_ops.pointnet2_utils import furthest_point_sample, gather_operation, ball_query, three_nn, three_interpolate, grouping_operation - -import os, sys -proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(proj_dir, "VRCNet/utils/pointnet2")) -from pointnet2_utils import furthest_point_sample, gather_operation, ball_query, three_nn, three_interpolate, grouping_operation - -class Conv1d(nn.Module): - def __init__(self, in_channel, out_channel, kernel_size=1, stride=1, if_bn=True, activation_fn=torch.relu): - super(Conv1d, self).__init__() - self.conv = nn.Conv1d(in_channel, out_channel, kernel_size, stride=stride) - self.if_bn = if_bn - self.bn = nn.BatchNorm1d(out_channel) - self.activation_fn = activation_fn - - def forward(self, input): - out = self.conv(input) - if self.if_bn: - out = self.bn(out) - - if self.activation_fn is not None: - out = self.activation_fn(out) - - return out - -class Conv2d(nn.Module): - def __init__(self, in_channel, out_channel, kernel_size=(1, 1), stride=(1, 1), if_bn=True, activation_fn=torch.relu): - super(Conv2d, self).__init__() - self.conv = nn.Conv2d(in_channel, out_channel, kernel_size, stride=stride) - self.if_bn = if_bn - self.bn = nn.BatchNorm2d(out_channel) - self.activation_fn = activation_fn - - def forward(self, input): - out = self.conv(input) - if self.if_bn: - out = self.bn(out) - - if self.activation_fn is not None: - out = self.activation_fn(out) - - return out - -class MLP(nn.Module): - def __init__(self, in_channel, layer_dims, bn=None): - super(MLP, self).__init__() - layers = [] - last_channel = in_channel - for out_channel in layer_dims[:-1]: - layers.append(nn.Linear(last_channel, out_channel)) - if bn: - layers.append(nn.BatchNorm1d(out_channel)) - layers.append(nn.ReLU()) - last_channel = out_channel - layers.append(nn.Linear(last_channel, layer_dims[-1])) - self.mlp = nn.Sequential(*layers) - - def forward(self, inputs): - return self.mlp(inputs) - -class MLP_CONV(nn.Module): - def __init__(self, in_channel, layer_dims, bn=None): - super(MLP_CONV, self).__init__() - layers = [] - last_channel = in_channel - for out_channel in layer_dims[:-1]: - layers.append(nn.Conv1d(last_channel, out_channel, 1)) - if bn: - layers.append(nn.BatchNorm1d(out_channel)) - layers.append(nn.ReLU()) - last_channel = out_channel - layers.append(nn.Conv1d(last_channel, layer_dims[-1], 1)) - self.mlp = nn.Sequential(*layers) - - def forward(self, inputs): - return self.mlp(inputs) - -class MLP_Res(nn.Module): - def __init__(self, in_dim=128, hidden_dim=None, out_dim=128): - super(MLP_Res, self).__init__() - if hidden_dim is None: - hidden_dim = in_dim - self.conv_1 = nn.Conv1d(in_dim, hidden_dim, 1) - self.conv_2 = nn.Conv1d(hidden_dim, out_dim, 1) - self.conv_shortcut = nn.Conv1d(in_dim, out_dim, 1) - - def forward(self, x): - """ - Args: - x: (B, out_dim, n) - """ - shortcut = self.conv_shortcut(x) - out = self.conv_2(torch.relu(self.conv_1(x))) + shortcut - return out - - -def sample_and_group(xyz, points, npoint, nsample, radius, use_xyz=True): - """ - Args: - xyz: Tensor, (B, 3, N) - points: Tensor, (B, f, N) - npoint: int - nsample: int - radius: float - use_xyz: boolean - - Returns: - new_xyz: Tensor, (B, 3, npoint) - new_points: Tensor, (B, 3 | f+3 | f, npoint, nsample) - idx_local: Tensor, (B, npoint, nsample) - grouped_xyz: Tensor, (B, 3, npoint, nsample) - - """ - xyz_flipped = xyz.permute(0, 2, 1).contiguous() # (B, N, 3) - new_xyz = gather_operation(xyz, furthest_point_sample(xyz_flipped, npoint)) # (B, 3, npoint) - - idx = ball_query(radius, nsample, xyz_flipped, new_xyz.permute(0, 2, 1).contiguous()) # (B, npoint, nsample) - grouped_xyz = grouping_operation(xyz, idx) # (B, 3, npoint, nsample) - grouped_xyz -= new_xyz.unsqueeze(3).repeat(1, 1, 1, nsample) - - if points is not None: - grouped_points = grouping_operation(points, idx) # (B, f, npoint, nsample) - if use_xyz: - new_points = torch.cat([grouped_xyz, grouped_points], 1) - else: - new_points = grouped_points - else: - new_points = grouped_xyz - - return new_xyz, new_points, idx, grouped_xyz - - -def sample_and_group_all(xyz, points, use_xyz=True): - """ - Args: - xyz: Tensor, (B, 3, nsample) - points: Tensor, (B, f, nsample) - use_xyz: boolean - - Returns: - new_xyz: Tensor, (B, 3, 1) - new_points: Tensor, (B, f|f+3|3, 1, nsample) - idx: Tensor, (B, 1, nsample) - grouped_xyz: Tensor, (B, 3, 1, nsample) - """ - b, _, nsample = xyz.shape - device = xyz.device - new_xyz = torch.zeros((1, 3, 1), dtype=torch.float, device=device).repeat(b, 1, 1) - grouped_xyz = xyz.reshape((b, 3, 1, nsample)) - idx = torch.arange(nsample, device=device).reshape(1, 1, nsample).repeat(b, 1, 1) - if points is not None: - if use_xyz: - new_points = torch.cat([xyz, points], 1) - else: - new_points = points - new_points = new_points.unsqueeze(2) - else: - new_points = grouped_xyz - - return new_xyz, new_points, idx, grouped_xyz - - -class PointNet_SA_Module(nn.Module): - def __init__(self, npoint, nsample, radius, in_channel, mlp, if_bn=True, group_all=False, use_xyz=True): - """ - Args: - npoint: int, number of points to sample - nsample: int, number of points in each local region - radius: float - in_channel: int, input channel of features(points) - mlp: list of int, - """ - super(PointNet_SA_Module, self).__init__() - self.npoint = npoint - self.nsample = nsample - self.radius = radius - self.mlp = mlp - self.group_all = group_all - self.use_xyz = use_xyz - if use_xyz: - in_channel += 3 - - last_channel = in_channel - self.mlp_conv = [] - for out_channel in mlp: - self.mlp_conv.append(Conv2d(last_channel, out_channel, if_bn=if_bn)) - last_channel = out_channel - - self.mlp_conv = nn.Sequential(*self.mlp_conv) - - def forward(self, xyz, points): - """ - Args: - xyz: Tensor, (B, 3, N) - points: Tensor, (B, f, N) - - Returns: - new_xyz: Tensor, (B, 3, npoint) - new_points: Tensor, (B, mlp[-1], npoint) - """ - if self.group_all: - new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, self.use_xyz) - else: - new_xyz, new_points, idx, grouped_xyz = sample_and_group(xyz, points, self.npoint, self.nsample, self.radius, self.use_xyz) - - new_points = self.mlp_conv(new_points) - new_points = torch.max(new_points, 3)[0] - - return new_xyz, new_points - - -class PointNet_FP_Module(nn.Module): - def __init__(self, in_channel, mlp, use_points1=False, in_channel_points1=None, if_bn=True): - """ - Args: - in_channel: int, input channel of points2 - mlp: list of int - use_points1: boolean, if use points - in_channel_points1: int, input channel of points1 - """ - super(PointNet_FP_Module, self).__init__() - self.use_points1 = use_points1 - - if use_points1: - in_channel += in_channel_points1 - - last_channel = in_channel - self.mlp_conv = [] - for out_channel in mlp: - self.mlp_conv.append(Conv1d(last_channel, out_channel, if_bn=if_bn)) - last_channel = out_channel - - self.mlp_conv = nn.Sequential(*self.mlp_conv) - - def forward(self, xyz1, xyz2, points1, points2): - """ - Args: - xyz1: Tensor, (B, 3, N) - xyz2: Tensor, (B, 3, M) - points1: Tensor, (B, in_channel, N) - points2: Tensor, (B, in_channel, M) - - Returns:MLP_CONV - new_points: Tensor, (B, mlp[-1], N) - """ - dist, idx = three_nn(xyz1.permute(0, 2, 1).contiguous(), xyz2.permute(0, 2, 1).contiguous()) - dist = torch.clamp_min(dist, 1e-10) # (B, N, 3) - recip_dist = 1.0/dist - norm = torch.sum(recip_dist, 2, keepdim=True).repeat((1, 1, 3)) - weight = recip_dist / norm - interpolated_points = three_interpolate(points2, idx, weight) # B, in_channel, N - - if self.use_points1: - new_points = torch.cat([interpolated_points, points1], 1) - else: - new_points = interpolated_points - - new_points = self.mlp_conv(new_points) - return new_points - - -def square_distance(src, dst): - """ - Calculate Euclid distance between each two points. - - src^T * dst = xn * xm + yn * ym + zn * zm； - sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn; - sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm; - dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2 - = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst - - Input: - src: source points, [B, N, C] - dst: target points, [B, M, C] - Output: - dist: per-point square distance, [B, N, M] - """ - B, N, _ = src.shape - _, M, _ = dst.shape - dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) # B, N, M - dist += torch.sum(src ** 2, -1).view(B, N, 1) - dist += torch.sum(dst ** 2, -1).view(B, 1, M) - return dist - - -def query_knn(nsample, xyz, new_xyz, include_self=True): - """Find k-NN of new_xyz in xyz""" - pad = 0 if include_self else 1 - sqrdists = square_distance(new_xyz, xyz) # B, S, N - idx = torch.argsort(sqrdists, dim=-1, descending=False)[:, :, pad: nsample+pad] - return idx.int() - - -def sample_and_group_knn(xyz, points, npoint, k, use_xyz=True, idx=None): - """ - Args: - xyz: Tensor, (B, 3, N) - points: Tensor, (B, f, N) - npoint: int - nsample: int - radius: float - use_xyz: boolean - - Returns: - new_xyz: Tensor, (B, 3, npoint) - new_points: Tensor, (B, 3 | f+3 | f, npoint, nsample) - idx_local: Tensor, (B, npoint, nsample) - grouped_xyz: Tensor, (B, 3, npoint, nsample) - - """ - xyz_flipped = xyz.permute(0, 2, 1).contiguous() # (B, N, 3) - new_xyz = gather_operation(xyz, furthest_point_sample(xyz_flipped, npoint)) # (B, 3, npoint) - if idx is None: - idx = query_knn(k, xyz_flipped, new_xyz.permute(0, 2, 1).contiguous()) - grouped_xyz = grouping_operation(xyz, idx) # (B, 3, npoint, nsample) - grouped_xyz -= new_xyz.unsqueeze(3).repeat(1, 1, 1, k) - - if points is not None: - grouped_points = grouping_operation(points, idx) # (B, f, npoint, nsample) - if use_xyz: - new_points = torch.cat([grouped_xyz, grouped_points], 1) - else: - new_points = grouped_points - else: - new_points = grouped_xyz - - return new_xyz, new_points, idx, grouped_xyz - - -class PointNet_SA_Module_KNN(nn.Module): - def __init__(self, npoint, nsample, in_channel, mlp, if_bn=True, group_all=False, use_xyz=True, if_idx=False): - """ - Args: - npoint: int, number of points to sample - nsample: int, number of points in each local region - radius: float - in_channel: int, input channel of features(points) - mlp: list of int, - """ - super(PointNet_SA_Module_KNN, self).__init__() - self.npoint = npoint - self.nsample = nsample - self.mlp = mlp - self.group_all = group_all - self.use_xyz = use_xyz - self.if_idx = if_idx - if use_xyz: - in_channel += 3 - - last_channel = in_channel - self.mlp_conv = [] - for out_channel in mlp[:-1]: - self.mlp_conv.append(Conv2d(last_channel, out_channel, if_bn=if_bn)) - last_channel = out_channel - self.mlp_conv.append(Conv2d(last_channel, mlp[-1], if_bn=False, activation_fn=None)) - self.mlp_conv = nn.Sequential(*self.mlp_conv) - - def forward(self, xyz, points, idx=None): - """ - Args: - xyz: Tensor, (B, 3, N) - points: Tensor, (B, f, N) - - Returns: - new_xyz: Tensor, (B, 3, npoint) - new_points: Tensor, (B, mlp[-1], npoint) - """ - if self.group_all: - new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, self.use_xyz) - else: - new_xyz, new_points, idx, grouped_xyz = sample_and_group_knn(xyz, points, self.npoint, self.nsample, self.use_xyz, idx=idx) - - new_points = self.mlp_conv(new_points) - new_points = torch.max(new_points, 3)[0] - - if self.if_idx: - return new_xyz, new_points, idx - else: - return new_xyz, new_points - - -def fps_subsample(pcd, n_points=2048): - """ - Args - pcd: (b, 16384, 3) - - returns - new_pcd: (b, n_points, 3) - """ - new_pcd = gather_operation(pcd.permute(0, 2, 1).contiguous(), furthest_point_sample(pcd, n_points)) - new_pcd = new_pcd.permute(0, 2, 1).contiguous() - return new_pcd - - -class Transformer(nn.Module): - def __init__(self, in_channel, dim=256, n_knn=16, pos_hidden_dim=64, attn_hidden_multiplier=4): - super(Transformer, self).__init__() - self.n_knn = n_knn - self.conv_key = nn.Conv1d(dim, dim, 1) - self.conv_query = nn.Conv1d(dim, dim, 1) - self.conv_value = nn.Conv1d(dim, dim, 1) - - self.pos_mlp = nn.Sequential( - nn.Conv2d(3, pos_hidden_dim, 1), - nn.BatchNorm2d(pos_hidden_dim), - nn.ReLU(), - nn.Conv2d(pos_hidden_dim, dim, 1) - ) - - self.attn_mlp = nn.Sequential( - nn.Conv2d(dim, dim * attn_hidden_multiplier, 1), - nn.BatchNorm2d(dim * attn_hidden_multiplier), - nn.ReLU(), - nn.Conv2d(dim * attn_hidden_multiplier, dim, 1) - ) - - self.linear_start = nn.Conv1d(in_channel, dim, 1) - self.linear_end = nn.Conv1d(dim, in_channel, 1) - - def forward(self, x, pos): - """feed forward of transformer - Args: - x: Tensor of features, (B, in_channel, n) - pos: Tensor of positions, (B, 3, n) - - Returns: - y: Tensor of features with attention, (B, in_channel, n) - """ - - identity = x - - x = self.linear_start(x) - b, dim, n = x.shape - - pos_flipped = pos.permute(0, 2, 1).contiguous() - idx_knn = query_knn(self.n_knn, pos_flipped, pos_flipped) - key = self.conv_key(x) - value = self.conv_value(x) - query = self.conv_query(x) - - key = grouping_operation(key, idx_knn) # b, dim, n, n_knn - qk_rel = query.reshape((b, -1, n, 1)) - key - - pos_rel = pos.reshape((b, -1, n, 1)) - grouping_operation(pos, idx_knn) # b, 3, n, n_knn - pos_embedding = self.pos_mlp(pos_rel) # b, dim, n, n_knn - - attention = self.attn_mlp(qk_rel + pos_embedding) - attention = torch.softmax(attention, -1) - - value = value.reshape((b, -1, n, 1)) + pos_embedding - - agg = einsum('b c i j, b c i j -> b c i', attention, value) # b, dim, n - y = self.linear_end(agg) - - return y+identity - - - - - - - - - - - - - - - - - - - - - - -