merge to merged branch

.gitignore

@@ -148,3 +148,6 @@ debug/
 
 pytorch/sampler_cuda/bin
 pytorch/sampler_cuda/build
+
+jittor/sampler_cuda/bin
+jittor/sampler_cuda/build

jittor/LICENSE

@@ -0,0 +1,24 @@
+This is free and unencumbered software released into the public domain.
+
+Anyone is free to copy, modify, publish, use, compile, sell, or
+distribute this software, either in source code form or as a compiled
+binary, for any purpose, commercial or non-commercial, and by any
+means.
+
+In jurisdictions that recognize copyright laws, the author or authors
+of this software dedicate any and all copyright interest in the
+software to the public domain. We make this dedication for the benefit
+of the public at large and to the detriment of our heirs and
+successors. We intend this dedication to be an overt act of
+relinquishment in perpetuity of all present and future rights to this
+software under copyright law.
+
+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 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.
+
+For more information, please refer to <https://unlicense.org>

jittor/README.md

@@ -0,0 +1,81 @@
+# DI-Fusion: 基于深度先验的在线隐式三维重建
+本仓库是[此文章](https://github.com/huangjh-pub/di-fusion)提出的用于隐式三维重建的网络部分的[计图](https://cg.cs.tsinghua.edu.cn/jittor/)实现，作者是：[黄家晖](https://cg.cs.tsinghua.edu.cn/people/~huangjh/zh)、[黄石生](https://cg.cs.tsinghua.edu.cn/people/~shisheng/)、宋浩轩和[胡事民](https://cg.cs.tsinghua.edu.cn/shimin.htm)
+
+- 计图是由清华大学计算机系 [图形学实验室](https://cg.cs.tsinghua.edu.cn/) 推出的一个完全基于动态编译、内部使用创新的元算子和统一计算图的深度学习框架。
+
+- DI-Fusion是一个基于RGBD输入的在线三维重建系统。它的相机定位追踪模块以及地图表示完全基于由深度神经网络建模的局部隐式表示。请进一步参考我们的[ArXiv报告](http://arxiv.org/abs/2012.05551)和[视频](https://youtu.be/yxkIQFXQ6rw)。
+- PyTorch version **available [here](https://github.com/huangjh-pub/di-fusion).**
+
+## 网络训练
+
+### 训练数据生成
+
+首先，请编译我们的CUDA点云采样器：
+
+```bash
+cd sampler_cuda
+mkdir build; cd build
+cmake ..
+make -j
+```
+
+编译成功之后，会在`sampler_cuda/bin/`文件夹下，生成名为`PreprocessMeshCUDA`的可执行文件，接着运行：
+
+```bash
+python data_generator.py data-shapenet.yaml --nproc 4
+```
+
+即可生成用于训练的数据。
+
+### 网络训练
+
+当完成了训练数据的生成之后，请运行：
+
+```bash
+python train.py train.yaml
+```
+
+即可开始训练，如果您在上一步更改了数据存放的位置，可能需要修改`train.yaml`中对应的路径，让程序正确的找到数据集。
+
+### 速度对比
+
+计图框架采用了先进的元算子融合以及统一计算图技术，这使得执行效率大大提高，下表对比了使用计图的版本和使用PyTorch的版本的训练速度，训练同一个epoch**计图所用的时间仅有PyTorch的三分之一**，原需要训练1-2天的模型，可以在半天内取得较好的收敛效果。
+
+|                     | PyTorch | PyTorch JIT | 计图 |
+| ------------------- | ------- | ----------- | ---- |
+| 每秒训练步数 (it/s) | 13      | 14          | 39   |
+
+## 运行 (beta)
+
+如果需要运行完整的DI-Fusion系统，首先需要进行简单的权值格式转换：
+
+```bash
+python convert.py
+```
+
+假设本仓库的路径是`<DIR>`，那么上述权值转换程序将输出`<DIR>/model_300.pth.tar`和`<DIR>/encoder_300.pth.tar`两个文件。
+
+接着，执行如下命令，拷贝官方实现的代码：
+
+```bash
+git clone https://github.com/huangjh-pub/di-fusion.git
+cd di-fusion
+cp <DIR>/model_300.pth.tar ./ckpt/default/
+cp <DIR>/encoder_300.pth.tar ./ckpt/default/
+```
+
+执行完成之后，请依照[这里](https://github.com/huangjh-pub/di-fusion#running)的做法继续之后的操作步骤。
+
+## 引用
+
+欢迎您引用我们的工作，蟹蟹：
+
+```
+@inproceedings{huang2021difusion,
+  title={DI-Fusion: Online Implicit 3D Reconstruction with Deep Priors},
+  author={Huang, Jiahui and Huang, Shi-Sheng and Song, Haoxuan and Hu, Shi-Min},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  year={2021}
+}
+```
+

jittor/convert.py

@@ -0,0 +1,33 @@
+import torch
+import pickle
+from pathlib import Path
+
+ENC_MAPPING = {
+    "layer0.conv": 0,
+    "layer0.normlayer.bn": 1,
+    "layer1.conv": 3,
+    "layer1.normlayer.bn": 4,
+    "layer2.conv": 6,
+    "layer2.normlayer.bn": 7,
+    "layer3.conv": 9
+}
+
+enc_jittor_path = Path("../di-checkpoints/default/encoder_300.jt.tar")
+dec_jittor_path = Path("../di-checkpoints/default/model_300.jt.tar")
+
+with enc_jittor_path.open("rb") as f:
+    enc_jt_weight = pickle.load(f)
+pth_dict = {}
+for wkey in list(enc_jt_weight.keys()):
+    wnew_key = None
+    for mkey in ENC_MAPPING.keys():
+        if str(ENC_MAPPING[mkey]) in wkey:
+            wnew_key = wkey.replace(str(ENC_MAPPING[mkey]), mkey)
+    pth_dict[wnew_key] = torch.from_numpy(enc_jt_weight[wkey]).cuda()
+torch.save({"epoch": 300, "model_state": pth_dict}, "./encoder_300.pth.tar")
+
+with dec_jittor_path.open("rb") as f:
+    dec_jt_weight = pickle.load(f)
+for wkey in list(dec_jt_weight.keys()):
+    dec_jt_weight[wkey] = torch.from_numpy(dec_jt_weight[wkey]).cuda()
+torch.save({"epoch": 300, "model_state": dec_jt_weight}, "./model_300.pth.tar")

jittor/criterion.py

@@ -0,0 +1,30 @@
+import math
+import jittor as jt
+
+
+def normal_log_prob(loc, scale, value):
+    var = (scale ** 2)
+    log_scale = scale.log()
+    return -((value - loc) ** 2) / (2 * var) - log_scale - math.log(math.sqrt(2 * math.pi))
+
+
+def neg_log_likelihood(args, info: dict, pd_sdf, pd_sdf_std, gt_sdf, **kwargs):
+    if args.enforce_minmax:
+        gt_sdf = jt.clamp(gt_sdf, -args.clamping_distance, args.clamping_distance)
+        pd_sdf = jt.clamp(pd_sdf, -args.clamping_distance, args.clamping_distance)
+
+    pd_dist_loc = pd_sdf.squeeze(1)
+    pd_dist_std = pd_sdf_std.squeeze(1)
+
+    sdf_loss = -normal_log_prob(pd_dist_loc, pd_dist_std, gt_sdf.squeeze(1)).sum() / info["num_sdf_samples"]
+    return {
+        'll': sdf_loss
+    }
+
+
+def reg_loss(args, info: dict, latent_vecs, **kwargs):
+    l2_size_loss = jt.sum(jt.norm(latent_vecs, k=2, dim=1))
+    reg_loss = min(1, info["epoch"] / 100) * l2_size_loss / info["num_sdf_samples"]
+    return {
+        'reg': reg_loss * args.code_reg_lambda
+    }

jittor/data-shapenet.yaml

@@ -0,0 +1,16 @@
+provider: shapenet_model
+
+provider_kwargs:
+  shapenet_path: "/dataset/ShapeNetCore.v2"
+  categories: ["03001627", "02871439", "03211117", "04256520", "03636649", "04379243"]
+  shapes_per_category: [100, 100, 100, 100, 100, 100]
+  scale: [1.0, 1.7, 0.6, 1.8, 0.5, 1.5]
+output: "../di-datasets/shapenet_plivoxs"
+
+sample_method: 1
+sampler_var: 0.00015
+sampler_count: 800000
+voxel_size: 0.08
+nn_size: 2.0
+
+nproc: 4