Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,60 @@
|
||||
{
|
||||
"[cpp]": {
|
||||
"editor.formatOnSave": true,
|
||||
"editor.defaultFormatter": "ms-vscode.cpptools"
|
||||
},
|
||||
"[cuda-cpp]": {
|
||||
"editor.formatOnSave": true,
|
||||
"editor.defaultFormatter": "ms-vscode.cpptools"
|
||||
},
|
||||
"[python]": {
|
||||
"editor.defaultFormatter": "charliermarsh.ruff",
|
||||
"editor.formatOnSave": true,
|
||||
"editor.codeActionsOnSave": {
|
||||
"source.organizeImports": "explicit"
|
||||
},
|
||||
},
|
||||
"pylint.enabled": false,
|
||||
"files.associations": {
|
||||
"optional": "cpp",
|
||||
"istream": "cpp",
|
||||
"ostream": "cpp",
|
||||
"ratio": "cpp",
|
||||
"system_error": "cpp",
|
||||
"array": "cpp",
|
||||
"functional": "cpp",
|
||||
"tuple": "cpp",
|
||||
"type_traits": "cpp",
|
||||
"utility": "cpp",
|
||||
"variant": "cpp",
|
||||
"charconv": "cpp",
|
||||
"chrono": "cpp",
|
||||
"format": "cpp",
|
||||
"__bit_reference": "cpp",
|
||||
"__node_handle": "cpp",
|
||||
"bitset": "cpp",
|
||||
"deque": "cpp",
|
||||
"limits": "cpp",
|
||||
"locale": "cpp",
|
||||
"vector": "cpp",
|
||||
"complex": "cpp",
|
||||
"string_view": "cpp",
|
||||
"initializer_list": "cpp",
|
||||
"ranges": "cpp",
|
||||
"span": "cpp",
|
||||
"__hash_table": "cpp",
|
||||
"__split_buffer": "cpp",
|
||||
"__tree": "cpp",
|
||||
"list": "cpp",
|
||||
"map": "cpp",
|
||||
"queue": "cpp",
|
||||
"set": "cpp",
|
||||
"stack": "cpp",
|
||||
"string": "cpp",
|
||||
"unordered_map": "cpp",
|
||||
"unordered_set": "cpp",
|
||||
"iostream": "cpp",
|
||||
"dense": "cpp",
|
||||
"*.tcc": "cpp"
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
# ViPE Contribution Rules
|
||||
|
||||
## Coding Guidelines
|
||||
|
||||
- Please follow the existing conventions in the relevant file, submodule, module, and project when you add new code or when you extend/fix existing functionality.
|
||||
|
||||
- Make sure that you can contribute your work to open source (no license and/or patent conflict is introduced by your code). You will need to [`sign`](#signing-your-work) your commit.
|
||||
|
||||
- Thanks in advance for your patience as we review your contributions; we do appreciate them!
|
||||
|
||||
## Signing Your Work
|
||||
|
||||
* We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
|
||||
|
||||
* Any contribution which contains commits that are not Signed-Off will not be accepted.
|
||||
|
||||
* To sign off on a commit you simply use the `--signoff` (or `-s`) option when committing your changes:
|
||||
```bash
|
||||
$ git commit -s -m "Add cool feature."
|
||||
```
|
||||
This will append the following to your commit message:
|
||||
```
|
||||
Signed-off-by: Your Name <your@email.com>
|
||||
```
|
||||
|
||||
* Full text of the DCO:
|
||||
|
||||
```
|
||||
Developer Certificate of Origin
|
||||
Version 1.1
|
||||
|
||||
Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
|
||||
1 Letterman Drive
|
||||
Suite D4700
|
||||
San Francisco, CA, 94129
|
||||
|
||||
Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed.
|
||||
```
|
||||
|
||||
```
|
||||
Developer's Certificate of Origin 1.1
|
||||
|
||||
By making a contribution to this project, I certify that:
|
||||
|
||||
(a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or
|
||||
|
||||
(b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications, whether created in whole or in part by me, under the same open source license (unless I am permitted to submit under a different license), as indicated in the file; or
|
||||
|
||||
(c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it.
|
||||
|
||||
(d) I understand and agree that this project and the contribution are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved.
|
||||
```
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
Copyright 2025 NVIDIA CORPORATION & AFFILIATES.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
+119
@@ -0,0 +1,119 @@
|
||||
# ViPE: Video Pose Engine for Geometric 3D Perception
|
||||
|
||||
<p align="center">
|
||||
<img src="assets/teaser.gif" alt="teaser"/>
|
||||
</p>
|
||||
|
||||
**TL;DR: ViPE is a useful open-source spatial AI tool for annotating camera poses and dense depth maps from raw videos!**
|
||||
|
||||
**Contributors**: NVIDIA (Spatial Intelligence Lab, Dynamic Vision Lab, NVIDIA Issac, NVIDIA Research).
|
||||
|
||||
**Full Abstract**: Accurate 3D geometric perception is an important prerequisite for a wide range of spatial AI systems. While state-of-the-art methods depend on large-scale training data, acquiring consistent and precise 3D annotations from in-the-wild videos remains a key challenge. In this work, we introduce ViPE, a handy and versatile video processing engine designed to bridge this gap. ViPE efficiently estimates camera intrinsics, camera motion, and dense, near-metric depth maps from unconstrained raw videos. It is robust to diverse scenarios, including dynamic selfie videos, cinematic shots, or dashcams, and supports various camera models such as pinhole, wide-angle, and 360° panoramas.
|
||||
We use ViPE to annotate a large-scale collection of videos. This collection includes around 100K real-world internet videos, 1M high-quality AI-generated videos, and 2K panoramic videos, totaling approximately 96M frames -- all annotated with accurate camera poses and dense depth maps. We open source ViPE and the annotated dataset with the hope to accelerate the development of spatial AI systems.
|
||||
|
||||
**[Technical Whitepaper](https://research.nvidia.com/labs/toronto-ai/vipe/assets/paper.pdf), [Project Page](https://research.nvidia.com/labs/toronto-ai/vipe), [Dataset](https://huggingface.co/) (Coming Soon)**
|
||||
|
||||
## Installation
|
||||
|
||||
To ensure the reproducibility, we recommend creating the runtime environment using [conda](https://docs.conda.io/projects/conda/en/latest/user-guide/getting-started.html).
|
||||
|
||||
```bash
|
||||
# Create a new conda environment and install 3rd-party dependencies
|
||||
conda env create -f envs/base.yml
|
||||
conda activate vipe
|
||||
pip install -r envs/requirements.txt
|
||||
|
||||
# Build the project and install it into the current environment
|
||||
# Omit the -e flag to install the project as a regular package
|
||||
pip install --no-build-isolation -e .
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Using the ViPE CLI
|
||||
|
||||
Once the python package is installed, you can use the `vipe` CLI to process raw videos in mp4 format.
|
||||
|
||||
```bash
|
||||
# Replace YOUR_VIDEO.mp4 with the path to your video. We provide sample videos in assets/examples.
|
||||
vipe infer YOUR_VIDEO.mp4
|
||||
# Additional options:
|
||||
# --output: Output directory (default: vipe_results)
|
||||
# --visualize: Enable visualization of intermediate and final results (default: false)
|
||||
# --pipeline: Pipeline configuration to use (default: default)
|
||||
```
|
||||
|
||||
|
||||
|
||||
One can visualize the results that ViPE produces by running (supported by `viser`):
|
||||
```bash
|
||||
vipe visualize vipe_results/
|
||||
# Please modify the above vipe_results/ path to the output directory of your choice.
|
||||
```
|
||||
|
||||
|
||||
|
||||
> We found that running [video-depth-anything](https://github.com/DepthAnything/Video-Depth-Anything) might eat up too much of GPU memory. To that end we provide a `no_vda` config that produces less temporally-stable depth (but empirically more 3D consistent) maps. This can be triggered by adding `--pipeline no_vda` to the `vipe infer` command.
|
||||
|
||||
### Using the `run.py` script
|
||||
|
||||
The `run.py` script is a more flexible way to run ViPE. Compared to the CLI, the script supports running on multiple videos at once and allows more fine-grained control over the pipeline with `hydra` configs. It also provides an example of using `vipe` as a library in your own project.
|
||||
|
||||
Example usages:
|
||||
|
||||
```bash
|
||||
# Running the full pipeline.
|
||||
python run.py pipeline=default streams=raw_mp4_stream streams.base_path=YOUR_VIDEO_OR_DIR_PATH
|
||||
|
||||
# Running the pose-only pipeline without depth estimation.
|
||||
python run.py pipeline=default streams=raw_mp4_stream streams.base_path=YOUR_VIDEO_OR_DIR_PATH pipeline.post.depth_align_model=null
|
||||
```
|
||||
|
||||
### Converting to COLMAP format
|
||||
|
||||
You can use the following script to convert the ViPE results to COLMAP format. For example:
|
||||
```bash
|
||||
python scripts/vipe_to_colmap.py vipe_results/ --sequence dog_example
|
||||
```
|
||||
This will unproject the dense depth maps to create the 3D point cloud.
|
||||
Alternatively for a more lightweight and 3D consistent point cloud, you can add the `--use_slam_map` flag to the above command. This requires you to run the full pipeline with `pipeline.output.save_slam_map=true` to save the additional information.
|
||||
|
||||
## Acknowledgments
|
||||
|
||||
ViPE is built on top of many great open-source research projects and codebases. Some of these include (not exhaustive):
|
||||
- [DROID-SLAM](https://github.com/princeton-vl/DROID-SLAM)
|
||||
- [Depth Anything V2](https://github.com/DepthAnything/Depth-Anything-V2)
|
||||
- [Metric3Dv2](https://github.com/YvanYin/Metric3D)
|
||||
- [PriorDA](https://github.com/SpatialVision/Prior-Depth-Anything)
|
||||
- [UniDepth](https://github.com/lpiccinelli-eth/UniDepth)
|
||||
- [VideoDepthAnything](https://github.com/DepthAnything/Video-Depth-Anything)
|
||||
- [GeoCalib](https://github.com/cvg/GeoCalib)
|
||||
- [Segment and Track Anything](https://github.com/z-x-yang/Segment-and-Track-Anything)
|
||||
|
||||
Please refer to the [THIRD_PARTY_LICENSES.md](THIRD_PARTY_LICENSES.md) for a full list of projects and their licenses.
|
||||
|
||||
We thank useful discussions from Aigul Dzhumamuratova, Viktor Kuznetsov, Soha Pouya, and Ming-Yu Liu, as well as release support from Vishal Kulkarni.
|
||||
|
||||
## TODO
|
||||
|
||||
- [x] Initial code released under Apache 2.0 license.
|
||||
- [ ] Full dataset uploaded to Hugging Face for download.
|
||||
- [ ] Add instructions to run inference on wide-angle and 360° videos.
|
||||
- [ ] Add instructions for benchmarking.
|
||||
|
||||
## Citation
|
||||
|
||||
If you find ViPE useful in your research or application, please consider citing the following whitepaper:
|
||||
|
||||
```
|
||||
@inproceedings{huang2025vipe,
|
||||
title={ViPE: Video Pose Engine for 3D Geometric Perception},
|
||||
author={Huang, Jiahui and Zhou, Qunjie and Rabeti, Hesam and Korovko, Aleksandr and Ling, Huan and Ren, Xuanchi and Shen, Tianchang and Gao, Jun and Slepichev, Dmitry and Lin, Chen-Hsuan and Ren, Jiawei and Xie, Kevin and Biswas, Joydeep and Leal-Taixe, Laura and Fidler, Sanja},
|
||||
booktitle={NVIDIA Research Whitepapers},
|
||||
year={2025}
|
||||
}
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
This project will download and install additional third-party **models and softwares**. Note that these models or softwares are not distributed by NVIDIA. Review the license terms of these models and projects before use. This source code is released under the [Apache 2 License](https://www.apache.org/licenses/LICENSE-2.0).
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,16 @@
|
||||
defaults:
|
||||
- _self_
|
||||
- streams: raw_mp4_stream
|
||||
- pipeline: ???
|
||||
- override hydra/hydra_logging: none
|
||||
- override hydra/job_logging: disabled
|
||||
|
||||
ray: false
|
||||
prefilter: false
|
||||
|
||||
hydra:
|
||||
output_subdir: null
|
||||
job:
|
||||
chdir: false
|
||||
run:
|
||||
dir: .
|
||||
@@ -0,0 +1,47 @@
|
||||
defaults:
|
||||
- /slam: default
|
||||
|
||||
instance: vipe.pipeline.default.DefaultAnnotationPipeline
|
||||
|
||||
# Init configs
|
||||
init:
|
||||
camera_type: "pinhole"
|
||||
intrinsics: "geocalib"
|
||||
instance:
|
||||
kf_gap_sec: 2.0
|
||||
phrases:
|
||||
- person
|
||||
- animal
|
||||
- vehicle
|
||||
- ball
|
||||
- balloon
|
||||
- gun
|
||||
- pet
|
||||
- car
|
||||
- bus
|
||||
add_sky: true
|
||||
|
||||
slam:
|
||||
keyframe_depth: unidepth-l
|
||||
optimize_intrinsics: ${neq:${..init.intrinsics},"gt"}
|
||||
|
||||
# Post-processing configs
|
||||
post:
|
||||
depth_align_model: "adaptive_unidepth-l_svda"
|
||||
|
||||
# Output configs
|
||||
output:
|
||||
# Path to save results
|
||||
path: vipe_results/
|
||||
skip_exists: false
|
||||
|
||||
# Save artifacts and a info file
|
||||
save_artifacts: false
|
||||
|
||||
# Save reconstruction from SLAM
|
||||
save_slam_map: false
|
||||
|
||||
# Visualization videos to BASE_PATH/vipe/xxx.mp4
|
||||
save_viz: true
|
||||
viz_downsample: 2
|
||||
viz_attributes: [['rgb', 'instance'], ['depth', 'pcd']]
|
||||
@@ -0,0 +1,5 @@
|
||||
defaults:
|
||||
- default
|
||||
|
||||
post:
|
||||
depth_align_model: "adaptive_unidepth-l"
|
||||
@@ -0,0 +1,8 @@
|
||||
defaults:
|
||||
- default
|
||||
|
||||
init:
|
||||
instance: null
|
||||
|
||||
post:
|
||||
depth_align_model: "adaptive_unidepth-l_vda"
|
||||
@@ -0,0 +1,52 @@
|
||||
buffer: 1024
|
||||
# weight for translation / rotation components of flow
|
||||
beta: 0.3
|
||||
# how much motion before considering new keyframe
|
||||
filter_thresh: 2.4
|
||||
# number of warmup frames
|
||||
warmup: 8
|
||||
# threshold to create a new keyframe
|
||||
keyframe_thresh: 4.0
|
||||
# add edges between frames whithin this distance
|
||||
frontend_thresh: 16.0
|
||||
# frontend optimization window
|
||||
frontend_window: 25
|
||||
# force edges between frames within radius
|
||||
frontend_radius: 2
|
||||
# non-maximal supression of edges
|
||||
frontend_nms: 1
|
||||
# Whether to initialize sequentially
|
||||
seq_init: true
|
||||
# Keyframe ids for running backends
|
||||
frontend_backend_iters: [16, 64, 256]
|
||||
backend_thresh: 22.0
|
||||
backend_radius: 2
|
||||
backend_nms: 3
|
||||
backend_iters: 24
|
||||
# Initial values for disparity
|
||||
init_disp: 1.0
|
||||
# Whether to optimize intrinsics
|
||||
optimize_intrinsics: ???
|
||||
# Whether to optimize rig
|
||||
optimize_rig_rotation: false
|
||||
# Whether add reproj factors for cross-views
|
||||
cross_view: true
|
||||
cross_view_idx: null # null means default
|
||||
adaptive_cross_view: false # Backend only, recompute cross_view_idx
|
||||
# Chunk size of the infill
|
||||
infill_chunk_size: 16
|
||||
# Whether or not to infill dense disparity
|
||||
infill_dense_disp: False
|
||||
# Filter threshold for SLAM map (percentage of the mean depth difference allowed)
|
||||
# (This is also used for dense disp extraction)
|
||||
map_filter_thresh: 0.05
|
||||
# Whether to visualize the results using rerun
|
||||
visualize: false
|
||||
# Use keyframe depth for recovering metric depth
|
||||
keyframe_depth: metric3d-small
|
||||
|
||||
ba:
|
||||
dense_disp_alpha: 0.001
|
||||
|
||||
sparse_tracks:
|
||||
name: dummy
|
||||
@@ -0,0 +1,7 @@
|
||||
instance: vipe.streams.raw_mp4_stream.RawMP4StreamList
|
||||
|
||||
base_path: ??
|
||||
frame_start: 0
|
||||
frame_end: 1000
|
||||
frame_skip: 1
|
||||
cached: false
|
||||
@@ -0,0 +1,49 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
void pybind_droid_net_ext(py::module &m);
|
||||
void pybind_grounding_dino_ext(py::module &m);
|
||||
void pybind_utils_ext(py::module &m);
|
||||
void pybind_slam_ext(py::module &m);
|
||||
void pybind_scatter_ext(py::module &m);
|
||||
void pybind_lietorch_ext(py::module &m);
|
||||
void pybind_corr_ext(py::module &m);
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
py::module m_droid_net = m.def_submodule("droid_net_ext");
|
||||
pybind_droid_net_ext(m_droid_net);
|
||||
|
||||
py::module m_gdino = m.def_submodule("grounding_dino_ext");
|
||||
pybind_grounding_dino_ext(m_gdino);
|
||||
|
||||
py::module m_utils = m.def_submodule("utils_ext");
|
||||
pybind_utils_ext(m_utils);
|
||||
|
||||
py::module m_slam = m.def_submodule("slam_ext");
|
||||
pybind_slam_ext(m_slam);
|
||||
|
||||
py::module m_scatter = m.def_submodule("scatter_ext");
|
||||
pybind_scatter_ext(m_scatter);
|
||||
|
||||
py::module m_lietorch = m.def_submodule("lietorch_ext");
|
||||
pybind_lietorch_ext(m_lietorch);
|
||||
|
||||
py::module m_corr = m.def_submodule("corr_ext");
|
||||
pybind_corr_ext(m_corr);
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch correlation repository:
|
||||
* https://github.com/ClementPinard/Pytorch-Correlation-extension、
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
using namespace torch;
|
||||
|
||||
#include <vector>
|
||||
|
||||
#define WITHIN_BOUNDS(x, y, H, W) (x >= 0 && x < H && y >= 0 && y < W)
|
||||
|
||||
template <typename scalar_t>
|
||||
static void correlate_patch(TensorAccessor<scalar_t, 3> input1, TensorAccessor<scalar_t, 3> input2, scalar_t *dst,
|
||||
int kH, int kW, int dilationH, int dilationW, int u, int v, int shiftU, int shiftV) {
|
||||
const int C = input1.size(0);
|
||||
const int iH = input1.size(1);
|
||||
const int iW = input1.size(2);
|
||||
for (int c = 0; c < C; ++c) {
|
||||
for (int i = 0; i < kH; ++i) {
|
||||
int i1 = u + i * dilationH;
|
||||
int i2 = i1 + shiftU;
|
||||
if WITHIN_BOUNDS (i1, i2, iH, iH) {
|
||||
for (int j = 0; j < kW; ++j) {
|
||||
int j1 = v + j * dilationW;
|
||||
int j2 = j1 + shiftV;
|
||||
if WITHIN_BOUNDS (j1, j2, iW, iW) {
|
||||
scalar_t v1 = input1[c][i1][j1];
|
||||
scalar_t v2 = input2[c][i2][j2];
|
||||
*dst += v1 * v2;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
static void correlate_patch_grad(TensorAccessor<scalar_t, 3> input1, TensorAccessor<scalar_t, 3> gradInput1,
|
||||
TensorAccessor<scalar_t, 3> input2, TensorAccessor<scalar_t, 3> gradInput2,
|
||||
scalar_t gradOutput, int kH, int kW, int dilationH, int dilationW, int u, int v,
|
||||
int shiftU, int shiftV) {
|
||||
const int C = input1.size(0);
|
||||
const int iH = input1.size(1);
|
||||
const int iW = input1.size(2);
|
||||
|
||||
for (int c = 0; c < C; ++c) {
|
||||
for (int i = 0; i < kH; ++i) {
|
||||
int i1 = u + i * dilationH;
|
||||
int i2 = i1 + shiftU;
|
||||
if WITHIN_BOUNDS (i1, i2, iH, iH) {
|
||||
for (int j = 0; j < kW; ++j) {
|
||||
int j1 = v + j * dilationW;
|
||||
int j2 = j1 + shiftV;
|
||||
if WITHIN_BOUNDS (j1, j2, iW, iW) {
|
||||
scalar_t v1 = input1[c][i1][j1];
|
||||
scalar_t v2 = input2[c][i2][j2];
|
||||
gradInput2[c][i2][j2] += gradOutput * v1;
|
||||
gradInput1[c][i1][j1] += gradOutput * v2;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor correlation_cpp_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
|
||||
int patchW, int padH, int padW, int dilationH, int dilationW, int dilation_patchH,
|
||||
int dilation_patchW, int dH, int dW) {
|
||||
const auto batch_size = input1.size(0);
|
||||
const auto iH = input1.size(2);
|
||||
const auto iW = input1.size(3);
|
||||
const int patchRadH = (patchH - 1) / 2;
|
||||
const int patchRadW = (patchW - 1) / 2;
|
||||
const int dilatedKH = (kH - 1) * dilationH + 1;
|
||||
const int dilatedKW = (kW - 1) * dilationW + 1;
|
||||
|
||||
const auto oH = (iH + 2 * padH - dilatedKH) / dH + 1;
|
||||
const auto oW = (iW + 2 * padW - dilatedKW) / dW + 1;
|
||||
auto output = at::zeros({batch_size, patchH, patchW, oH, oW}, input1.options());
|
||||
|
||||
int n, ph, pw, h, w;
|
||||
#pragma omp parallel for private(n, ph, pw, h, w) collapse(2)
|
||||
for (n = 0; n < batch_size; ++n) {
|
||||
for (ph = 0; ph < patchH; ++ph) {
|
||||
for (pw = 0; pw < patchW; ++pw) {
|
||||
AT_DISPATCH_FLOATING_TYPES(
|
||||
input1.scalar_type(), "correlation_forward_cpp", ([&] {
|
||||
auto input1_acc = input1.accessor<scalar_t, 4>();
|
||||
auto input2_acc = input2.accessor<scalar_t, 4>();
|
||||
auto output_acc = output.accessor<scalar_t, 5>();
|
||||
for (h = 0; h < oH; ++h) {
|
||||
for (w = 0; w < oW; ++w) {
|
||||
correlate_patch(input1_acc[n], input2_acc[n], &output_acc[n][ph][pw][h][w], kH, kW,
|
||||
dilationH, dilationW, -padH + h * dH, -padW + w * dW,
|
||||
(ph - patchRadH) * dilation_patchH, (pw - patchRadW) * dilation_patchW);
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
return output;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> correlation_cpp_backward(torch::Tensor input1, torch::Tensor input2,
|
||||
torch::Tensor gradOutput, int kH, int kW, int patchH, int patchW,
|
||||
int padH, int padW, int dilationH, int dilationW,
|
||||
int dilation_patchH, int dilation_patchW, int dH, int dW) {
|
||||
const int batch_size = input1.size(0);
|
||||
const int patchRadH = (patchH - 1) / 2;
|
||||
const int patchRadW = (patchW - 1) / 2;
|
||||
const int oH = gradOutput.size(3);
|
||||
const int oW = gradOutput.size(4);
|
||||
|
||||
auto gradInput1 = torch::zeros_like(input1);
|
||||
|
||||
auto gradInput2 = torch::zeros_like(input2);
|
||||
|
||||
int n, ph, pw, h, w;
|
||||
#pragma omp parallel for private(n, ph, pw, h, w)
|
||||
for (n = 0; n < batch_size; ++n) {
|
||||
AT_DISPATCH_FLOATING_TYPES(
|
||||
input1.scalar_type(), "correlation_backward_cpp", ([&] {
|
||||
auto input1_acc = input1.accessor<scalar_t, 4>();
|
||||
auto gradInput1_acc = gradInput1.accessor<scalar_t, 4>();
|
||||
auto input2_acc = input2.accessor<scalar_t, 4>();
|
||||
auto gradInput2_acc = gradInput2.accessor<scalar_t, 4>();
|
||||
auto gradOutput_acc = gradOutput.accessor<scalar_t, 5>();
|
||||
|
||||
for (ph = 0; ph < patchH; ++ph) {
|
||||
for (pw = 0; pw < patchW; ++pw) {
|
||||
for (h = 0; h < oH; ++h) {
|
||||
for (w = 0; w < oW; ++w) {
|
||||
correlate_patch_grad(input1_acc[n], gradInput1_acc[n], input2_acc[n], gradInput2_acc[n],
|
||||
gradOutput_acc[n][ph][pw][h][w], kH, kW, dilationH, dilationW,
|
||||
-padH + h * dH, -padW + w * dW, (ph - patchRadH) * dilation_patchH,
|
||||
(pw - patchRadW) * dilation_patchW);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
return {gradInput1, gradInput2};
|
||||
}
|
||||
+286
@@ -0,0 +1,286 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch correlation repository:
|
||||
* https://github.com/ClementPinard/Pytorch-Correlation-extension、
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/types.h>
|
||||
using namespace torch;
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <vector>
|
||||
#include <iostream>
|
||||
|
||||
// Cuda tensor accessor definitions
|
||||
// restrict pointer traits piroritize speed over memory consumption
|
||||
#define TensorAcc4R PackedTensorAccessor32<scalar_t, 4, RestrictPtrTraits>
|
||||
#define TensorAcc5R PackedTensorAccessor32<scalar_t, 5, RestrictPtrTraits>
|
||||
#define WITHIN_BOUNDS(x, y, H, W) (x >= 0 && x < H && y >= 0 && y < W)
|
||||
|
||||
#define THREADS_FORWARD 32
|
||||
#define THREADS_BACKWARD 5
|
||||
|
||||
namespace corr {
|
||||
template <typename scalar_t>
|
||||
__global__ void correlation_cuda_forward_kernel(const TensorAcc4R rInput1, const TensorAcc4R rInput2,
|
||||
TensorAcc5R output, int kH, int kW, int patchH, int patchW, int padH,
|
||||
int padW, int dilationH, int dilationW, int dilation_patchH,
|
||||
int dilation_patchW, int dH, int dW) {
|
||||
const int iH = rInput1.size(1);
|
||||
const int iW = rInput1.size(2);
|
||||
const int C = rInput1.size(3);
|
||||
|
||||
const int n = blockIdx.x;
|
||||
const int h = blockIdx.y;
|
||||
const int w = blockIdx.z;
|
||||
const int thread = threadIdx.x;
|
||||
|
||||
const int start_i = -padH + h * dH;
|
||||
const int start_j = -padW + w * dW;
|
||||
|
||||
const int patchRadH = dilation_patchH * (patchH - 1) / 2;
|
||||
const int patchRadW = dilation_patchW * (patchW - 1) / 2;
|
||||
|
||||
__shared__ scalar_t prod_sum[THREADS_FORWARD];
|
||||
|
||||
for (int ph = 0; ph < patchH; ++ph) {
|
||||
int ph_dilated = ph * dilation_patchH - patchRadH;
|
||||
for (int pw = 0; pw < patchW; ++pw) {
|
||||
int pw_dilated = pw * dilation_patchW - patchRadW;
|
||||
prod_sum[thread] = 0;
|
||||
for (int i = 0; i < kH; ++i) {
|
||||
int i1 = start_i + i * dilationH;
|
||||
int i2 = i1 + ph_dilated;
|
||||
if WITHIN_BOUNDS (i1, i2, iH, iH) {
|
||||
for (int j = 0; j < kW; ++j) {
|
||||
int j1 = start_j + j * dilationW;
|
||||
int j2 = j1 + pw_dilated;
|
||||
if WITHIN_BOUNDS (j1, j2, iW, iW) {
|
||||
for (int c = thread; c < C; c += THREADS_FORWARD) {
|
||||
scalar_t v1 = rInput1[n][i1][j1][c];
|
||||
scalar_t v2 = rInput2[n][i2][j2][c];
|
||||
prod_sum[thread] += v1 * v2;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// accumulate
|
||||
__syncthreads();
|
||||
if (thread == 0) {
|
||||
scalar_t reduce_sum = 0;
|
||||
for (int index = 0; index < THREADS_FORWARD; ++index) {
|
||||
reduce_sum += prod_sum[index];
|
||||
}
|
||||
output[n][ph][pw][h][w] = reduce_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void correlation_cuda_backward_kernel_input1(const TensorAcc5R gradOutput, const TensorAcc4R input2,
|
||||
TensorAcc4R gradInput1, const int kH, const int kW,
|
||||
const int patchH, const int patchW, const int padH,
|
||||
const int padW, const int dilationH, const int dilationW,
|
||||
const int dilation_patchH, const int dilation_patchW,
|
||||
const int dH, const int dW, const int batch) {
|
||||
const int iH = input2.size(2);
|
||||
const int iW = input2.size(3);
|
||||
|
||||
const int H = gradOutput.size(3);
|
||||
const int W = gradOutput.size(4);
|
||||
|
||||
const int patchRadH = (patchH - 1) / 2;
|
||||
const int patchRadW = (patchW - 1) / 2;
|
||||
|
||||
const int n = batch;
|
||||
const int c = blockIdx.x;
|
||||
const int h = blockIdx.y;
|
||||
const int w = blockIdx.z;
|
||||
const int ph_off = threadIdx.x;
|
||||
const int pw_off = threadIdx.y;
|
||||
|
||||
const int h_2 = h + padH;
|
||||
const int w_2 = w + padW;
|
||||
const int min_h = h_2 - kH * dilationH;
|
||||
const int min_w = w_2 - kW * dilationW;
|
||||
|
||||
__shared__ scalar_t prod_sum[THREADS_BACKWARD][THREADS_BACKWARD];
|
||||
prod_sum[ph_off][pw_off] = 0;
|
||||
|
||||
for (int ph = ph_off; ph < patchH; ph += THREADS_BACKWARD) {
|
||||
int i1 = h + dilation_patchH * (ph - patchRadH);
|
||||
for (int pw = pw_off; pw < patchW; pw += THREADS_BACKWARD) {
|
||||
int j1 = w + dilation_patchW * (pw - patchRadW);
|
||||
if (WITHIN_BOUNDS(i1, j1, iH, iW)) {
|
||||
scalar_t val = input2[n][c][i1][j1];
|
||||
for (int h_3 = h_2; h_3 > min_h; h_3 -= dilationH) {
|
||||
int i2 = (h_3) / dH;
|
||||
if (i2 * dH != h_3) continue;
|
||||
for (int w_3 = w_2; w_3 > min_w; w_3 -= dilationW) {
|
||||
int j2 = (w_3) / dW;
|
||||
if (j2 * dW != w_3) continue;
|
||||
if WITHIN_BOUNDS (i2, j2, H, W) {
|
||||
prod_sum[ph_off][pw_off] += gradOutput[n][ph][pw][i2][j2] * val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (ph_off == 0 && pw_off == 0) {
|
||||
scalar_t reduce_sum = 0;
|
||||
for (int ph = 0; ph < THREADS_BACKWARD; ++ph) {
|
||||
for (int pw = 0; pw < THREADS_BACKWARD; ++pw) {
|
||||
reduce_sum += prod_sum[ph][pw];
|
||||
}
|
||||
}
|
||||
gradInput1[n][c][h][w] = reduce_sum;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void correlation_cuda_backward_kernel_input2(const TensorAcc5R gradOutput, const TensorAcc4R input1,
|
||||
TensorAcc4R gradInput2, int kH, int kW, int patchH, int patchW,
|
||||
int padH, int padW, int dilationH, int dilationW,
|
||||
int dilation_patchH, int dilation_patchW, int dH, int dW,
|
||||
int batch) {
|
||||
const int iH = input1.size(2);
|
||||
const int iW = input1.size(3);
|
||||
|
||||
const int patchRadH = (patchH - 1) / 2;
|
||||
const int patchRadW = (patchW - 1) / 2;
|
||||
|
||||
const int H = gradOutput.size(3);
|
||||
const int W = gradOutput.size(4);
|
||||
|
||||
const int dilatedKH = kH * dilationH;
|
||||
const int dilatedKW = kW * dilationW;
|
||||
|
||||
const int n = batch;
|
||||
const int c = blockIdx.x;
|
||||
const int h = blockIdx.y;
|
||||
const int w = blockIdx.z;
|
||||
const int ph_off = threadIdx.x;
|
||||
const int pw_off = threadIdx.y;
|
||||
|
||||
__shared__ scalar_t prod_sum[THREADS_BACKWARD][THREADS_BACKWARD];
|
||||
prod_sum[ph_off][pw_off] = 0;
|
||||
|
||||
for (int ph = ph_off; ph < patchH; ph += THREADS_BACKWARD) {
|
||||
int i1 = h - dilation_patchH * (ph - patchRadH);
|
||||
for (int pw = pw_off; pw < patchW; pw += THREADS_BACKWARD) {
|
||||
int j1 = w - dilation_patchW * (pw - patchRadW);
|
||||
if WITHIN_BOUNDS (i1, j1, iH, iW) {
|
||||
scalar_t val = input1[n][c][i1][j1];
|
||||
|
||||
const int h_2 = i1 + padH;
|
||||
const int w_2 = j1 + padW;
|
||||
const int min_h = h_2 - dilatedKH;
|
||||
const int min_w = w_2 - dilatedKW;
|
||||
|
||||
for (int h_3 = h_2; h_3 > min_h; h_3 -= dilationH) {
|
||||
int i2 = (h_3) / dH;
|
||||
if (i2 * dH != h_3) continue;
|
||||
for (int w_3 = w_2; w_3 > min_w; w_3 -= dilationW) {
|
||||
int j2 = (w_3) / dW;
|
||||
if (j2 * dW != w_3) continue;
|
||||
if WITHIN_BOUNDS (i2, j2, H, W) {
|
||||
prod_sum[ph_off][pw_off] += gradOutput[n][ph][pw][i2][j2] * val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (ph_off == 0 && pw_off == 0) {
|
||||
scalar_t reduce_sum = 0;
|
||||
for (int ph = 0; ph < THREADS_BACKWARD; ++ph) {
|
||||
for (int pw = 0; pw < THREADS_BACKWARD; ++pw) {
|
||||
reduce_sum += prod_sum[ph][pw];
|
||||
}
|
||||
}
|
||||
gradInput2[n][c][h][w] = reduce_sum;
|
||||
}
|
||||
}
|
||||
} // namespace corr
|
||||
|
||||
torch::Tensor correlation_cuda_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
|
||||
int patchW, int padH, int padW, int dilationH, int dilationW,
|
||||
int dilation_patchH, int dilation_patchW, int dH, int dW) {
|
||||
const int batch_size = input1.size(0);
|
||||
const int iH = input1.size(2);
|
||||
const int iW = input1.size(3);
|
||||
const int dilatedKH = (kH - 1) * dilationH + 1;
|
||||
const int dilatedKW = (kW - 1) * dilationW + 1;
|
||||
|
||||
const auto oH = (iH + 2 * padH - dilatedKH) / dH + 1;
|
||||
const auto oW = (iW + 2 * padW - dilatedKW) / dW + 1;
|
||||
auto output = torch::zeros({batch_size, patchH, patchW, oH, oW}, input1.options());
|
||||
|
||||
auto trInput1 = input1.permute({0, 2, 3, 1}).contiguous();
|
||||
auto trInput2 = input2.permute({0, 2, 3, 1}).contiguous();
|
||||
|
||||
const int threads = THREADS_FORWARD;
|
||||
const dim3 blocks(batch_size, oH, oW);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND2(
|
||||
at::ScalarType::Half, at::ScalarType::BFloat16, input1.scalar_type(), "correlation_forward_cuda", ([&] {
|
||||
TensorAcc4R trInput1_acc = trInput1.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
|
||||
TensorAcc4R trInput2_acc = trInput2.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
|
||||
TensorAcc5R output_acc = output.packed_accessor32<scalar_t, 5, RestrictPtrTraits>();
|
||||
corr::correlation_cuda_forward_kernel<scalar_t>
|
||||
<<<blocks, threads>>>(trInput1_acc, trInput2_acc, output_acc, kH, kW, patchH, patchW, padH, padW,
|
||||
dilationH, dilationW, dilation_patchH, dilation_patchW, dH, dW);
|
||||
}));
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> correlation_cuda_backward(torch::Tensor input1, torch::Tensor input2,
|
||||
torch::Tensor gradOutput, int kH, int kW, int patchH, int patchW,
|
||||
int padH, int padW, int dilationH, int dilationW,
|
||||
int dilation_patchH, int dilation_patchW, int dH, int dW) {
|
||||
auto gradInput1 = torch::zeros_like(input1);
|
||||
auto gradInput2 = torch::zeros_like(input2);
|
||||
|
||||
const int batch_size = input1.size(0);
|
||||
const int iH = input1.size(2);
|
||||
const int iW = input1.size(3);
|
||||
const int C = input1.size(1);
|
||||
|
||||
const dim3 blocks(C, iH, iW);
|
||||
const dim3 threads(THREADS_BACKWARD, THREADS_BACKWARD);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND2(
|
||||
at::ScalarType::Half, at::ScalarType::BFloat16, input1.scalar_type(), "correlation_backward_cuda", ([&] {
|
||||
TensorAcc4R input1_acc = input1.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
|
||||
TensorAcc4R input2_acc = input2.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
|
||||
TensorAcc4R gradInput1_acc = gradInput1.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
|
||||
TensorAcc4R gradInput2_acc = gradInput2.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
|
||||
TensorAcc5R gradOutput_acc = gradOutput.packed_accessor32<scalar_t, 5, RestrictPtrTraits>();
|
||||
|
||||
for (int n = 0; n < batch_size; ++n) {
|
||||
corr::correlation_cuda_backward_kernel_input1<scalar_t>
|
||||
<<<blocks, threads>>>(gradOutput_acc, input2_acc, gradInput1_acc, kH, kW, patchH, patchW, padH,
|
||||
padW, dilationH, dilationW, dilation_patchH, dilation_patchW, dH, dW, n);
|
||||
}
|
||||
|
||||
for (int n = 0; n < batch_size; ++n) {
|
||||
corr::correlation_cuda_backward_kernel_input2<scalar_t>
|
||||
<<<blocks, threads>>>(gradOutput_acc, input1_acc, gradInput2_acc, kH, kW, patchH, patchW, padH,
|
||||
padW, dilationH, dilationW, dilation_patchH, dilation_patchW, dH, dW, n);
|
||||
}
|
||||
}));
|
||||
|
||||
return {gradInput1, gradInput2};
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch correlation repository:
|
||||
* https://github.com/ClementPinard/Pytorch-Correlation-extension、
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <vector>
|
||||
#include <iostream>
|
||||
|
||||
// declarations
|
||||
|
||||
torch::Tensor correlation_cpp_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
|
||||
int patchW, int padH, int padW, int dilationH, int dilationW, int dilation_patchH,
|
||||
int dilation_patchW, int dH, int dW);
|
||||
|
||||
std::vector<torch::Tensor> correlation_cpp_backward(torch::Tensor grad_output, torch::Tensor input1,
|
||||
torch::Tensor input2, int kH, int kW, int patchH, int patchW,
|
||||
int padH, int padW, int dilationH, int dilationW,
|
||||
int dilation_patchH, int dilation_patchW, int dH, int dW);
|
||||
|
||||
#define CHECK_CUDA(x) TORCH_CHECK(x.device().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)
|
||||
#define CHECK_SAME_DEVICE(x, y) TORCH_CHECK(x.device() == y.device(), #x " is not on same device as " #y)
|
||||
|
||||
torch::Tensor correlation_cuda_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
|
||||
int patchW, int padH, int padW, int dilationH, int dilationW,
|
||||
int dilation_patchH, int dilation_patchW, int dH, int dW);
|
||||
|
||||
std::vector<torch::Tensor> correlation_cuda_backward(torch::Tensor grad_output, torch::Tensor input1,
|
||||
torch::Tensor input2, int kH, int kW, int patchH, int patchW,
|
||||
int padH, int padW, int dilationH, int dilationW,
|
||||
int dilation_patchH, int dilation_patchW, int dH, int dW);
|
||||
|
||||
// C++ interface
|
||||
|
||||
torch::Tensor correlation_sample_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
|
||||
int patchW, int padH, int padW, int dilationH, int dilationW,
|
||||
int dilation_patchH, int dilation_patchW, int dH, int dW) {
|
||||
if (input1.device().is_cuda()) {
|
||||
CHECK_INPUT(input1);
|
||||
CHECK_INPUT(input2);
|
||||
|
||||
// set device of input1 as default CUDA device
|
||||
// https://pytorch.org/cppdocs/api/structc10_1_1cuda_1_1_optional_c_u_d_a_guard.html
|
||||
const at::cuda::OptionalCUDAGuard guard_input1(device_of(input1));
|
||||
CHECK_SAME_DEVICE(input1, input2);
|
||||
|
||||
return correlation_cuda_forward(input1, input2, kH, kW, patchH, patchW, padH, padW, dilationH, dilationW,
|
||||
dilation_patchH, dilation_patchW, dH, dW);
|
||||
} else {
|
||||
return correlation_cpp_forward(input1, input2, kH, kW, patchH, patchW, padH, padW, dilationH, dilationW,
|
||||
dilation_patchH, dilation_patchW, dH, dW);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> correlation_sample_backward(torch::Tensor input1, torch::Tensor input2,
|
||||
torch::Tensor grad_output, int kH, int kW, int patchH,
|
||||
int patchW, int padH, int padW, int dilationH, int dilationW,
|
||||
int dilation_patchH, int dilation_patchW, int dH, int dW) {
|
||||
if (grad_output.device().is_cuda()) {
|
||||
CHECK_INPUT(input1);
|
||||
CHECK_INPUT(input2);
|
||||
|
||||
// set device of input1 as default CUDA device
|
||||
const at::cuda::OptionalCUDAGuard guard_input1(device_of(input1));
|
||||
CHECK_SAME_DEVICE(input1, input2);
|
||||
CHECK_SAME_DEVICE(input1, grad_output);
|
||||
|
||||
return correlation_cuda_backward(input1, input2, grad_output, kH, kW, patchH, patchW, padH, padW, dilationH,
|
||||
dilationW, dilation_patchH, dilation_patchW, dH, dW);
|
||||
} else {
|
||||
return correlation_cpp_backward(input1, input2, grad_output, kH, kW, patchH, patchW, padH, padW, dilationH,
|
||||
dilationW, dilation_patchH, dilation_patchW, dH, dW);
|
||||
}
|
||||
}
|
||||
|
||||
void pybind_corr_ext(py::module& m) {
|
||||
m.def("forward", &correlation_sample_forward, "Spatial Correlation Sampler Forward");
|
||||
m.def("backward", &correlation_sample_backward, "Spatial Correlation Sampler backward");
|
||||
}
|
||||
@@ -0,0 +1,321 @@
|
||||
/**
|
||||
* This file includes code originally from the DROID-SLAM repository:
|
||||
* https://github.com/princeton-vl/DROID-SLAM
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/NativeFunctions.h>
|
||||
#include <ATen/cuda/CUDAApplyUtils.cuh>
|
||||
#include <ATen/native/cuda/KernelUtils.cuh>
|
||||
|
||||
#define BLOCK_H 4
|
||||
#define BLOCK_W 8
|
||||
#define BLOCK_HW BLOCK_H *BLOCK_W
|
||||
#define CHANNEL_STRIDE 32
|
||||
|
||||
__forceinline__ __device__ bool within_bounds(int h, int w, int H, int W) { return h >= 0 && h < H && w >= 0 && w < W; }
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void altcorr_forward_kernel(const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap1,
|
||||
const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap2,
|
||||
const torch::PackedTensorAccessor32<float, 5, torch::RestrictPtrTraits> coords,
|
||||
torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr,
|
||||
int r) {
|
||||
const int b = blockIdx.x;
|
||||
const int h0 = blockIdx.y * blockDim.x;
|
||||
const int w0 = blockIdx.z * blockDim.y;
|
||||
const int tid = threadIdx.x * blockDim.y + threadIdx.y;
|
||||
|
||||
const int H1 = fmap1.size(1);
|
||||
const int W1 = fmap1.size(2);
|
||||
const int H2 = fmap2.size(1);
|
||||
const int W2 = fmap2.size(2);
|
||||
const int N = coords.size(1);
|
||||
const int C = fmap1.size(3);
|
||||
|
||||
__shared__ scalar_t f1[CHANNEL_STRIDE][BLOCK_HW];
|
||||
__shared__ scalar_t f2[CHANNEL_STRIDE][BLOCK_HW];
|
||||
|
||||
__shared__ float x2s[BLOCK_HW];
|
||||
__shared__ float y2s[BLOCK_HW];
|
||||
|
||||
for (int c = 0; c < C; c += CHANNEL_STRIDE) {
|
||||
for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
|
||||
int k1 = k + tid / CHANNEL_STRIDE;
|
||||
int h1 = h0 + k1 / BLOCK_W;
|
||||
int w1 = w0 + k1 % BLOCK_W;
|
||||
int c1 = tid % CHANNEL_STRIDE;
|
||||
|
||||
if (within_bounds(h1, w1, H1, W1))
|
||||
f1[c1][k1] = fmap1[b][h1][w1][c + c1];
|
||||
|
||||
else
|
||||
f1[c1][k1] = 0.0;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int n = 0; n < N; n++) {
|
||||
int h1 = h0 + threadIdx.x;
|
||||
int w1 = w0 + threadIdx.y;
|
||||
if (within_bounds(h1, w1, H1, W1)) {
|
||||
x2s[tid] = coords[b][n][h1][w1][0];
|
||||
y2s[tid] = coords[b][n][h1][w1][1];
|
||||
}
|
||||
|
||||
float dx = x2s[tid] - floor(x2s[tid]);
|
||||
float dy = y2s[tid] - floor(y2s[tid]);
|
||||
|
||||
int rd = 2 * r + 1;
|
||||
for (int iy = 0; iy < rd + 1; iy++) {
|
||||
for (int ix = 0; ix < rd + 1; ix++) {
|
||||
for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
|
||||
int k1 = k + tid / CHANNEL_STRIDE;
|
||||
int h2 = static_cast<int>(floor(y2s[k1])) - r + iy;
|
||||
int w2 = static_cast<int>(floor(x2s[k1])) - r + ix;
|
||||
int c2 = tid % CHANNEL_STRIDE;
|
||||
|
||||
if (within_bounds(h2, w2, H2, W2))
|
||||
f2[c2][k1] = fmap2[b][h2][w2][c + c2];
|
||||
|
||||
else
|
||||
f2[c2][k1] = static_cast<scalar_t>(0.0);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
scalar_t s = 0.0;
|
||||
for (int k = 0; k < CHANNEL_STRIDE; k++) s += f1[k][tid] * f2[k][tid];
|
||||
|
||||
int ix_nw = H1 * W1 * ((iy - 1) + rd * (ix - 1));
|
||||
int ix_ne = H1 * W1 * ((iy - 1) + rd * ix);
|
||||
int ix_sw = H1 * W1 * (iy + rd * (ix - 1));
|
||||
int ix_se = H1 * W1 * (iy + rd * ix);
|
||||
|
||||
// int ix_nw = ((iy-1) + rd*(ix-1));
|
||||
// int ix_ne = ((iy-1) + rd*ix);
|
||||
// int ix_sw = (iy + rd*(ix-1));
|
||||
// int ix_se = (iy + rd*ix);
|
||||
|
||||
scalar_t nw = s * static_cast<scalar_t>((dy) * (dx));
|
||||
scalar_t ne = s * static_cast<scalar_t>((dy) * (1 - dx));
|
||||
scalar_t sw = s * static_cast<scalar_t>((1 - dy) * (dx));
|
||||
scalar_t se = s * static_cast<scalar_t>((1 - dy) * (1 - dx));
|
||||
|
||||
// if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1))
|
||||
// corr[b][n][ix_nw][h1][w1] += nw;
|
||||
|
||||
// if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1))
|
||||
// corr[b][n][ix_ne][h1][w1] += ne;
|
||||
|
||||
// if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1))
|
||||
// corr[b][n][ix_sw][h1][w1] += sw;
|
||||
|
||||
// if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1))
|
||||
// corr[b][n][ix_se][h1][w1] += se;
|
||||
|
||||
scalar_t *corr_ptr = &corr[b][n][0][h1][w1];
|
||||
|
||||
if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_nw) += nw;
|
||||
|
||||
if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_ne) += ne;
|
||||
|
||||
if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_sw) += sw;
|
||||
|
||||
if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_se) += se;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void altcorr_backward_kernel(
|
||||
const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap1,
|
||||
const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap2,
|
||||
const torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> coords,
|
||||
const torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr_grad,
|
||||
torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap1_grad,
|
||||
torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap2_grad,
|
||||
torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> coords_grad, int r) {
|
||||
const int b = blockIdx.x;
|
||||
const int h0 = blockIdx.y * blockDim.x;
|
||||
const int w0 = blockIdx.z * blockDim.y;
|
||||
const int tid = threadIdx.x * blockDim.y + threadIdx.y;
|
||||
|
||||
const int H1 = fmap1.size(1);
|
||||
const int W1 = fmap1.size(2);
|
||||
const int H2 = fmap2.size(1);
|
||||
const int W2 = fmap2.size(2);
|
||||
const int N = coords.size(1);
|
||||
const int C = fmap1.size(3);
|
||||
|
||||
__shared__ scalar_t f1[CHANNEL_STRIDE][BLOCK_HW + 1];
|
||||
__shared__ scalar_t f2[CHANNEL_STRIDE][BLOCK_HW + 1];
|
||||
|
||||
__shared__ scalar_t f1_grad[CHANNEL_STRIDE][BLOCK_HW + 1];
|
||||
__shared__ scalar_t f2_grad[CHANNEL_STRIDE][BLOCK_HW + 1];
|
||||
|
||||
__shared__ scalar_t x2s[BLOCK_HW];
|
||||
__shared__ scalar_t y2s[BLOCK_HW];
|
||||
|
||||
for (int c = 0; c < C; c += CHANNEL_STRIDE) {
|
||||
for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
|
||||
int k1 = k + tid / CHANNEL_STRIDE;
|
||||
int h1 = h0 + k1 / BLOCK_W;
|
||||
int w1 = w0 + k1 % BLOCK_W;
|
||||
int c1 = tid % CHANNEL_STRIDE;
|
||||
|
||||
auto fptr = fmap1[b][h1][w1];
|
||||
if (within_bounds(h1, w1, H1, W1))
|
||||
f1[c1][k1] = fptr[c + c1];
|
||||
else
|
||||
f1[c1][k1] = 0.0;
|
||||
|
||||
f1_grad[c1][k1] = 0.0;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
int h1 = h0 + threadIdx.x;
|
||||
int w1 = w0 + threadIdx.y;
|
||||
|
||||
for (int n = 0; n < N; n++) {
|
||||
x2s[tid] = coords[b][n][h1][w1][0];
|
||||
y2s[tid] = coords[b][n][h1][w1][1];
|
||||
|
||||
scalar_t dx = x2s[tid] - floor(x2s[tid]);
|
||||
scalar_t dy = y2s[tid] - floor(y2s[tid]);
|
||||
|
||||
int rd = 2 * r + 1;
|
||||
for (int iy = 0; iy < rd + 1; iy++) {
|
||||
for (int ix = 0; ix < rd + 1; ix++) {
|
||||
for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
|
||||
int k1 = k + tid / CHANNEL_STRIDE;
|
||||
int h2 = static_cast<int>(floor(y2s[k1])) - r + iy;
|
||||
int w2 = static_cast<int>(floor(x2s[k1])) - r + ix;
|
||||
int c2 = tid % CHANNEL_STRIDE;
|
||||
|
||||
auto fptr = fmap2[b][h2][w2];
|
||||
if (within_bounds(h2, w2, H2, W2))
|
||||
f2[c2][k1] = fptr[c + c2];
|
||||
else
|
||||
f2[c2][k1] = 0.0;
|
||||
|
||||
f2_grad[c2][k1] = 0.0;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
const scalar_t *grad_ptr = &corr_grad[b][n][0][h1][w1];
|
||||
scalar_t g = 0.0;
|
||||
|
||||
int ix_nw = H1 * W1 * ((iy - 1) + rd * (ix - 1));
|
||||
int ix_ne = H1 * W1 * ((iy - 1) + rd * ix);
|
||||
int ix_sw = H1 * W1 * (iy + rd * (ix - 1));
|
||||
int ix_se = H1 * W1 * (iy + rd * ix);
|
||||
|
||||
if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1)) g += *(grad_ptr + ix_nw) * dy * dx;
|
||||
|
||||
if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1)) g += *(grad_ptr + ix_ne) * dy * (1 - dx);
|
||||
|
||||
if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1)) g += *(grad_ptr + ix_sw) * (1 - dy) * dx;
|
||||
|
||||
if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1))
|
||||
g += *(grad_ptr + ix_se) * (1 - dy) * (1 - dx);
|
||||
|
||||
for (int k = 0; k < CHANNEL_STRIDE; k++) {
|
||||
f1_grad[k][tid] += g * f2[k][tid];
|
||||
f2_grad[k][tid] += g * f1[k][tid];
|
||||
}
|
||||
|
||||
for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
|
||||
int k1 = k + tid / CHANNEL_STRIDE;
|
||||
int h2 = static_cast<int>(floor(y2s[k1])) - r + iy;
|
||||
int w2 = static_cast<int>(floor(x2s[k1])) - r + ix;
|
||||
int c2 = tid % CHANNEL_STRIDE;
|
||||
|
||||
scalar_t *fptr = &fmap2_grad[b][h2][w2][0];
|
||||
if (within_bounds(h2, w2, H2, W2)) atomicAdd(fptr + c + c2, f2_grad[c2][k1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
|
||||
int k1 = k + tid / CHANNEL_STRIDE;
|
||||
int h1 = h0 + k1 / BLOCK_W;
|
||||
int w1 = w0 + k1 % BLOCK_W;
|
||||
int c1 = tid % CHANNEL_STRIDE;
|
||||
|
||||
scalar_t *fptr = &fmap1_grad[b][h1][w1][0];
|
||||
if (within_bounds(h1, w1, H1, W1)) fptr[c + c1] += f1_grad[c1][k1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> altcorr_cuda_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
|
||||
int radius) {
|
||||
const auto B = coords.size(0);
|
||||
const auto N = coords.size(1);
|
||||
const auto H = coords.size(2);
|
||||
const auto W = coords.size(3);
|
||||
|
||||
const auto rd = 2 * radius + 1;
|
||||
auto opts = fmap1.options();
|
||||
auto corr = torch::zeros({B, N, rd * rd, H, W}, opts);
|
||||
|
||||
const dim3 blocks(B, (H + BLOCK_H - 1) / BLOCK_H, (W + BLOCK_W - 1) / BLOCK_W);
|
||||
const dim3 threads(BLOCK_H, BLOCK_W);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(fmap1.scalar_type(), "altcorr_forward_kernel", ([&] {
|
||||
altcorr_forward_kernel<scalar_t><<<blocks, threads>>>(
|
||||
fmap1.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
|
||||
fmap2.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
|
||||
coords.packed_accessor32<float, 5, torch::RestrictPtrTraits>(),
|
||||
corr.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
|
||||
radius);
|
||||
}));
|
||||
|
||||
return {corr};
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> altcorr_cuda_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
|
||||
torch::Tensor corr_grad, int radius) {
|
||||
const auto B = coords.size(0);
|
||||
const auto N = coords.size(1);
|
||||
|
||||
const auto H1 = fmap1.size(1);
|
||||
const auto W1 = fmap1.size(2);
|
||||
const auto H2 = fmap2.size(1);
|
||||
const auto W2 = fmap2.size(2);
|
||||
const auto C = fmap1.size(3);
|
||||
|
||||
auto opts = fmap1.options();
|
||||
auto fmap1_grad = torch::zeros({B, H1, W1, C}, opts);
|
||||
auto fmap2_grad = torch::zeros({B, H2, W2, C}, opts);
|
||||
auto coords_grad = torch::zeros({B, N, H1, W1, 2}, opts);
|
||||
|
||||
const dim3 blocks(B, (H1 + BLOCK_H - 1) / BLOCK_H, (W1 + BLOCK_W - 1) / BLOCK_W);
|
||||
const dim3 threads(BLOCK_H, BLOCK_W);
|
||||
|
||||
altcorr_backward_kernel<float>
|
||||
<<<blocks, threads>>>(fmap1.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
|
||||
fmap2.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
|
||||
coords.packed_accessor32<float, 5, torch::RestrictPtrTraits>(),
|
||||
corr_grad.packed_accessor32<float, 5, torch::RestrictPtrTraits>(),
|
||||
fmap1_grad.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
|
||||
fmap2_grad.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
|
||||
coords_grad.packed_accessor32<float, 5, torch::RestrictPtrTraits>(), radius);
|
||||
|
||||
return {fmap1_grad, fmap2_grad, coords_grad};
|
||||
}
|
||||
+159
@@ -0,0 +1,159 @@
|
||||
/**
|
||||
* This file includes code originally from the DROID-SLAM repository:
|
||||
* https://github.com/princeton-vl/DROID-SLAM
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/NativeFunctions.h>
|
||||
#include <ATen/Parallel.h>
|
||||
|
||||
#define BLOCK 16
|
||||
|
||||
__forceinline__ __device__ bool within_bounds(int h, int w, int H, int W) { return h >= 0 && h < H && w >= 0 && w < W; }
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void corr_index_forward_kernel(
|
||||
const torch::PackedTensorAccessor64<scalar_t, 5, torch::RestrictPtrTraits> volume,
|
||||
const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> coords,
|
||||
torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr, int r) {
|
||||
// batch index
|
||||
const int x = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int y = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
const int n = blockIdx.z;
|
||||
|
||||
const int h1 = volume.size(1);
|
||||
const int w1 = volume.size(2);
|
||||
const int h2 = volume.size(3);
|
||||
const int w2 = volume.size(4);
|
||||
|
||||
if (!within_bounds(y, x, h1, w1)) {
|
||||
return;
|
||||
}
|
||||
|
||||
float x0 = coords[n][0][y][x];
|
||||
float y0 = coords[n][1][y][x];
|
||||
|
||||
float dx = x0 - floor(x0);
|
||||
float dy = y0 - floor(y0);
|
||||
|
||||
int rd = 2 * r + 1;
|
||||
for (int i = 0; i < rd + 1; i++) {
|
||||
for (int j = 0; j < rd + 1; j++) {
|
||||
int x1 = static_cast<int>(floor(x0)) - r + i;
|
||||
int y1 = static_cast<int>(floor(y0)) - r + j;
|
||||
|
||||
if (within_bounds(y1, x1, h2, w2)) {
|
||||
scalar_t s = volume[n][y][x][y1][x1];
|
||||
|
||||
if (i > 0 && j > 0) corr[n][i - 1][j - 1][y][x] += s * scalar_t(dx * dy);
|
||||
|
||||
if (i > 0 && j < rd) corr[n][i - 1][j][y][x] += s * scalar_t(dx * (1.0f - dy));
|
||||
|
||||
if (i < rd && j > 0) corr[n][i][j - 1][y][x] += s * scalar_t((1.0f - dx) * dy);
|
||||
|
||||
if (i < rd && j < rd) corr[n][i][j][y][x] += s * scalar_t((1.0f - dx) * (1.0f - dy));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void corr_index_backward_kernel(
|
||||
const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> coords,
|
||||
const torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr_grad,
|
||||
torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> volume_grad, int r) {
|
||||
// batch index
|
||||
const int x = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int y = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
const int n = blockIdx.z;
|
||||
|
||||
const int h1 = volume_grad.size(1);
|
||||
const int w1 = volume_grad.size(2);
|
||||
const int h2 = volume_grad.size(3);
|
||||
const int w2 = volume_grad.size(4);
|
||||
|
||||
if (!within_bounds(y, x, h1, w1)) {
|
||||
return;
|
||||
}
|
||||
|
||||
float x0 = coords[n][0][y][x];
|
||||
float y0 = coords[n][1][y][x];
|
||||
|
||||
float dx = x0 - floor(x0);
|
||||
float dy = y0 - floor(y0);
|
||||
|
||||
int rd = 2 * r + 1;
|
||||
for (int i = 0; i < rd + 1; i++) {
|
||||
for (int j = 0; j < rd + 1; j++) {
|
||||
int x1 = static_cast<int>(floor(x0)) - r + i;
|
||||
int y1 = static_cast<int>(floor(y0)) - r + j;
|
||||
|
||||
if (within_bounds(y1, x1, h2, w2)) {
|
||||
scalar_t g = 0.0;
|
||||
if (i > 0 && j > 0) g += corr_grad[n][i - 1][j - 1][y][x] * scalar_t(dx * dy);
|
||||
|
||||
if (i > 0 && j < rd) g += corr_grad[n][i - 1][j][y][x] * scalar_t(dx * (1.0f - dy));
|
||||
|
||||
if (i < rd && j > 0) g += corr_grad[n][i][j - 1][y][x] * scalar_t((1.0f - dx) * dy);
|
||||
|
||||
if (i < rd && j < rd) g += corr_grad[n][i][j][y][x] * scalar_t((1.0f - dx) * (1.0f - dy));
|
||||
|
||||
volume_grad[n][y][x][y1][x1] += g;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> corr_index_cuda_forward(torch::Tensor volume, torch::Tensor coords, int radius) {
|
||||
const auto batch_size = volume.size(0);
|
||||
const auto ht = volume.size(1);
|
||||
const auto wd = volume.size(2);
|
||||
|
||||
const dim3 blocks((wd + BLOCK - 1) / BLOCK, (ht + BLOCK - 1) / BLOCK, batch_size);
|
||||
|
||||
const dim3 threads(BLOCK, BLOCK);
|
||||
|
||||
auto opts = volume.options();
|
||||
torch::Tensor corr = torch::zeros({batch_size, 2 * radius + 1, 2 * radius + 1, ht, wd}, opts);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(volume.scalar_type(), "sampler_forward_kernel", ([&] {
|
||||
corr_index_forward_kernel<scalar_t><<<blocks, threads>>>(
|
||||
volume.packed_accessor64<scalar_t, 5, torch::RestrictPtrTraits>(),
|
||||
coords.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
|
||||
corr.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
|
||||
radius);
|
||||
}));
|
||||
|
||||
return {corr};
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> corr_index_cuda_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
|
||||
int radius) {
|
||||
const auto batch_size = volume.size(0);
|
||||
const auto ht = volume.size(1);
|
||||
const auto wd = volume.size(2);
|
||||
|
||||
auto volume_grad = torch::zeros_like(volume);
|
||||
|
||||
const dim3 blocks((wd + BLOCK - 1) / BLOCK, (ht + BLOCK - 1) / BLOCK, batch_size);
|
||||
|
||||
const dim3 threads(BLOCK, BLOCK);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(volume.scalar_type(), "sampler_backward_kernel", ([&] {
|
||||
corr_index_backward_kernel<scalar_t><<<blocks, threads>>>(
|
||||
coords.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
|
||||
corr_grad.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
|
||||
volume_grad.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
|
||||
radius);
|
||||
}));
|
||||
|
||||
return {volume_grad};
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
/**
|
||||
* This file includes code originally from the DROID-SLAM repository:
|
||||
* https://github.com/princeton-vl/DROID-SLAM
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
||||
#define CHECK_INPUT(x) CHECK_CONTIGUOUS(x)
|
||||
|
||||
std::vector<torch::Tensor> corr_index_cuda_forward(torch::Tensor volume, torch::Tensor coords, int radius);
|
||||
std::vector<torch::Tensor> corr_index_cuda_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
|
||||
int radius);
|
||||
std::vector<torch::Tensor> altcorr_cuda_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
|
||||
int radius);
|
||||
std::vector<torch::Tensor> altcorr_cuda_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
|
||||
torch::Tensor corr_grad, int radius);
|
||||
|
||||
// c++ python binding
|
||||
std::vector<torch::Tensor> corr_index_forward(torch::Tensor volume, torch::Tensor coords, int radius) {
|
||||
CHECK_INPUT(volume);
|
||||
CHECK_INPUT(coords);
|
||||
|
||||
return corr_index_cuda_forward(volume, coords, radius);
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> corr_index_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
|
||||
int radius) {
|
||||
CHECK_INPUT(volume);
|
||||
CHECK_INPUT(coords);
|
||||
CHECK_INPUT(corr_grad);
|
||||
|
||||
auto volume_grad = corr_index_cuda_backward(volume, coords, corr_grad, radius);
|
||||
return {volume_grad};
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> altcorr_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords, int radius) {
|
||||
CHECK_INPUT(fmap1);
|
||||
CHECK_INPUT(fmap2);
|
||||
CHECK_INPUT(coords);
|
||||
|
||||
return altcorr_cuda_forward(fmap1, fmap2, coords, radius);
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> altcorr_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
|
||||
torch::Tensor corr_grad, int radius) {
|
||||
CHECK_INPUT(fmap1);
|
||||
CHECK_INPUT(fmap2);
|
||||
CHECK_INPUT(coords);
|
||||
CHECK_INPUT(corr_grad);
|
||||
|
||||
return altcorr_cuda_backward(fmap1, fmap2, coords, corr_grad, radius);
|
||||
}
|
||||
|
||||
void pybind_droid_net_ext(py::module& m) {
|
||||
// correlation volume kernels
|
||||
m.def("altcorr_forward", &altcorr_forward, "ALTCORR forward");
|
||||
m.def("altcorr_backward", &altcorr_backward, "ALTCORR backward");
|
||||
m.def("corr_index_forward", &corr_index_forward, "INDEX forward");
|
||||
m.def("corr_index_backward", &corr_index_backward, "INDEX backward");
|
||||
}
|
||||
+136
@@ -0,0 +1,136 @@
|
||||
/*!
|
||||
**************************************************************************************************
|
||||
* Deformable DETR
|
||||
* Copyright (c) 2020 SenseTime. All Rights Reserved.
|
||||
* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
**************************************************************************************************
|
||||
* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
|
||||
**************************************************************************************************
|
||||
*/
|
||||
|
||||
#include <vector>
|
||||
#include "ms_deform_im2col_cuda.cuh"
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
namespace groundingdino {
|
||||
|
||||
at::Tensor ms_deform_attn_cuda_forward(const at::Tensor &value, const at::Tensor &spatial_shapes,
|
||||
const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
|
||||
const at::Tensor &attn_weight, const int im2col_step) {
|
||||
AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
|
||||
AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
|
||||
AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
|
||||
AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
|
||||
AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
|
||||
|
||||
AT_ASSERTM(value.is_cuda(), "value must be a CUDA tensor");
|
||||
AT_ASSERTM(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor");
|
||||
AT_ASSERTM(level_start_index.is_cuda(), "level_start_index must be a CUDA tensor");
|
||||
AT_ASSERTM(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor");
|
||||
AT_ASSERTM(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor");
|
||||
|
||||
const int batch = value.size(0);
|
||||
const int spatial_size = value.size(1);
|
||||
const int num_heads = value.size(2);
|
||||
const int channels = value.size(3);
|
||||
|
||||
const int num_levels = spatial_shapes.size(0);
|
||||
|
||||
const int num_query = sampling_loc.size(1);
|
||||
const int num_point = sampling_loc.size(4);
|
||||
|
||||
const int im2col_step_ = std::min(batch, im2col_step);
|
||||
|
||||
AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
|
||||
|
||||
auto output = at::zeros({batch, num_query, num_heads, channels}, value.options());
|
||||
|
||||
const int batch_n = im2col_step_;
|
||||
auto output_n = output.view({batch / im2col_step_, batch_n, num_query, num_heads, channels});
|
||||
auto per_value_size = spatial_size * num_heads * channels;
|
||||
auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
|
||||
auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
|
||||
for (int n = 0; n < batch / im2col_step_; ++n) {
|
||||
auto columns = output_n.select(0, n);
|
||||
AT_DISPATCH_FLOATING_TYPES(
|
||||
value.scalar_type(), "ms_deform_attn_forward_cuda", ([&] {
|
||||
ms_deformable_im2col_cuda(
|
||||
at::cuda::getCurrentCUDAStream(), value.data<scalar_t>() + n * im2col_step_ * per_value_size,
|
||||
spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
|
||||
sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
|
||||
attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size, batch_n, spatial_size,
|
||||
num_heads, channels, num_levels, num_query, num_point, columns.data<scalar_t>());
|
||||
}));
|
||||
}
|
||||
|
||||
output = output.view({batch, num_query, num_heads * channels});
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> ms_deform_attn_cuda_backward(const at::Tensor &value, const at::Tensor &spatial_shapes,
|
||||
const at::Tensor &level_start_index,
|
||||
const at::Tensor &sampling_loc, const at::Tensor &attn_weight,
|
||||
const at::Tensor &grad_output, const int im2col_step) {
|
||||
AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
|
||||
AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
|
||||
AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
|
||||
AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
|
||||
AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
|
||||
AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous");
|
||||
|
||||
AT_ASSERTM(value.is_cuda(), "value must be a CUDA tensor");
|
||||
AT_ASSERTM(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor");
|
||||
AT_ASSERTM(level_start_index.is_cuda(), "level_start_index must be a CUDA tensor");
|
||||
AT_ASSERTM(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor");
|
||||
AT_ASSERTM(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor");
|
||||
AT_ASSERTM(grad_output.is_cuda(), "grad_output must be a CUDA tensor");
|
||||
|
||||
const int batch = value.size(0);
|
||||
const int spatial_size = value.size(1);
|
||||
const int num_heads = value.size(2);
|
||||
const int channels = value.size(3);
|
||||
|
||||
const int num_levels = spatial_shapes.size(0);
|
||||
|
||||
const int num_query = sampling_loc.size(1);
|
||||
const int num_point = sampling_loc.size(4);
|
||||
|
||||
const int im2col_step_ = std::min(batch, im2col_step);
|
||||
|
||||
AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
|
||||
|
||||
auto grad_value = at::zeros_like(value);
|
||||
auto grad_sampling_loc = at::zeros_like(sampling_loc);
|
||||
auto grad_attn_weight = at::zeros_like(attn_weight);
|
||||
|
||||
const int batch_n = im2col_step_;
|
||||
auto per_value_size = spatial_size * num_heads * channels;
|
||||
auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
|
||||
auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
|
||||
auto grad_output_n = grad_output.view({batch / im2col_step_, batch_n, num_query, num_heads, channels});
|
||||
|
||||
for (int n = 0; n < batch / im2col_step_; ++n) {
|
||||
auto grad_output_g = grad_output_n.select(0, n);
|
||||
AT_DISPATCH_FLOATING_TYPES(value.scalar_type(), "ms_deform_attn_backward_cuda", ([&] {
|
||||
ms_deformable_col2im_cuda(
|
||||
at::cuda::getCurrentCUDAStream(), grad_output_g.data<scalar_t>(),
|
||||
value.data<scalar_t>() + n * im2col_step_ * per_value_size,
|
||||
spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
|
||||
sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
|
||||
attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size,
|
||||
batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point,
|
||||
grad_value.data<scalar_t>() + n * im2col_step_ * per_value_size,
|
||||
grad_sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
|
||||
grad_attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size);
|
||||
}));
|
||||
}
|
||||
|
||||
return {grad_value, grad_sampling_loc, grad_attn_weight};
|
||||
}
|
||||
|
||||
} // namespace groundingdino
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
/*!
|
||||
**************************************************************************************************
|
||||
* Deformable DETR
|
||||
* Copyright (c) 2020 SenseTime. All Rights Reserved.
|
||||
* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
**************************************************************************************************
|
||||
* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
|
||||
**************************************************************************************************
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
#include <torch/extension.h>
|
||||
|
||||
namespace groundingdino {
|
||||
|
||||
at::Tensor ms_deform_attn_cuda_forward(const at::Tensor &value, const at::Tensor &spatial_shapes,
|
||||
const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
|
||||
const at::Tensor &attn_weight, const int im2col_step);
|
||||
|
||||
std::vector<at::Tensor> ms_deform_attn_cuda_backward(const at::Tensor &value, const at::Tensor &spatial_shapes,
|
||||
const at::Tensor &level_start_index,
|
||||
const at::Tensor &sampling_loc, const at::Tensor &attn_weight,
|
||||
const at::Tensor &grad_output, const int im2col_step);
|
||||
|
||||
} // namespace groundingdino
|
||||
+914
@@ -0,0 +1,914 @@
|
||||
/*!
|
||||
**************************************************************************
|
||||
* Deformable DETR
|
||||
* Copyright (c) 2020 SenseTime. All Rights Reserved.
|
||||
* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
**************************************************************************
|
||||
* Modified from DCN (https://github.com/msracver/Deformable-ConvNets)
|
||||
* Copyright (c) 2018 Microsoft
|
||||
**************************************************************************
|
||||
*/
|
||||
|
||||
#include <cstdio>
|
||||
#include <algorithm>
|
||||
#include <cstring>
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
|
||||
#include <THC/THCAtomics.cuh>
|
||||
|
||||
#define CUDA_KERNEL_LOOP(i, n) for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x)
|
||||
|
||||
const int CUDA_NUM_THREADS = 1024;
|
||||
inline int GET_BLOCKS(const int N, const int num_threads) { return (N + num_threads - 1) / num_threads; }
|
||||
|
||||
template <typename scalar_t>
|
||||
__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t *&bottom_data, const int &height, const int &width,
|
||||
const int &nheads, const int &channels, const scalar_t &h,
|
||||
const scalar_t &w, const int &m, const int &c) {
|
||||
const int h_low = floor(h);
|
||||
const int w_low = floor(w);
|
||||
const int h_high = h_low + 1;
|
||||
const int w_high = w_low + 1;
|
||||
|
||||
const scalar_t lh = h - h_low;
|
||||
const scalar_t lw = w - w_low;
|
||||
const scalar_t hh = 1 - lh, hw = 1 - lw;
|
||||
|
||||
const int w_stride = nheads * channels;
|
||||
const int h_stride = width * w_stride;
|
||||
const int h_low_ptr_offset = h_low * h_stride;
|
||||
const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
|
||||
const int w_low_ptr_offset = w_low * w_stride;
|
||||
const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
|
||||
const int base_ptr = m * channels + c;
|
||||
|
||||
scalar_t v1 = 0;
|
||||
if (h_low >= 0 && w_low >= 0) {
|
||||
const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
|
||||
v1 = bottom_data[ptr1];
|
||||
}
|
||||
scalar_t v2 = 0;
|
||||
if (h_low >= 0 && w_high <= width - 1) {
|
||||
const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
|
||||
v2 = bottom_data[ptr2];
|
||||
}
|
||||
scalar_t v3 = 0;
|
||||
if (h_high <= height - 1 && w_low >= 0) {
|
||||
const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
|
||||
v3 = bottom_data[ptr3];
|
||||
}
|
||||
scalar_t v4 = 0;
|
||||
if (h_high <= height - 1 && w_high <= width - 1) {
|
||||
const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
|
||||
v4 = bottom_data[ptr4];
|
||||
}
|
||||
|
||||
const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
|
||||
|
||||
const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
|
||||
return val;
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__device__ void ms_deform_attn_col2im_bilinear(const scalar_t *&bottom_data, const int &height, const int &width,
|
||||
const int &nheads, const int &channels, const scalar_t &h,
|
||||
const scalar_t &w, const int &m, const int &c, const scalar_t &top_grad,
|
||||
const scalar_t &attn_weight, scalar_t *&grad_value,
|
||||
scalar_t *grad_sampling_loc, scalar_t *grad_attn_weight) {
|
||||
const int h_low = floor(h);
|
||||
const int w_low = floor(w);
|
||||
const int h_high = h_low + 1;
|
||||
const int w_high = w_low + 1;
|
||||
|
||||
const scalar_t lh = h - h_low;
|
||||
const scalar_t lw = w - w_low;
|
||||
const scalar_t hh = 1 - lh, hw = 1 - lw;
|
||||
|
||||
const int w_stride = nheads * channels;
|
||||
const int h_stride = width * w_stride;
|
||||
const int h_low_ptr_offset = h_low * h_stride;
|
||||
const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
|
||||
const int w_low_ptr_offset = w_low * w_stride;
|
||||
const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
|
||||
const int base_ptr = m * channels + c;
|
||||
|
||||
const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
|
||||
const scalar_t top_grad_value = top_grad * attn_weight;
|
||||
scalar_t grad_h_weight = 0, grad_w_weight = 0;
|
||||
|
||||
scalar_t v1 = 0;
|
||||
if (h_low >= 0 && w_low >= 0) {
|
||||
const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
|
||||
v1 = bottom_data[ptr1];
|
||||
grad_h_weight -= hw * v1;
|
||||
grad_w_weight -= hh * v1;
|
||||
atomicAdd(grad_value + ptr1, w1 * top_grad_value);
|
||||
}
|
||||
scalar_t v2 = 0;
|
||||
if (h_low >= 0 && w_high <= width - 1) {
|
||||
const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
|
||||
v2 = bottom_data[ptr2];
|
||||
grad_h_weight -= lw * v2;
|
||||
grad_w_weight += hh * v2;
|
||||
atomicAdd(grad_value + ptr2, w2 * top_grad_value);
|
||||
}
|
||||
scalar_t v3 = 0;
|
||||
if (h_high <= height - 1 && w_low >= 0) {
|
||||
const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
|
||||
v3 = bottom_data[ptr3];
|
||||
grad_h_weight += hw * v3;
|
||||
grad_w_weight -= lh * v3;
|
||||
atomicAdd(grad_value + ptr3, w3 * top_grad_value);
|
||||
}
|
||||
scalar_t v4 = 0;
|
||||
if (h_high <= height - 1 && w_high <= width - 1) {
|
||||
const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
|
||||
v4 = bottom_data[ptr4];
|
||||
grad_h_weight += lw * v4;
|
||||
grad_w_weight += lh * v4;
|
||||
atomicAdd(grad_value + ptr4, w4 * top_grad_value);
|
||||
}
|
||||
|
||||
const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
|
||||
*grad_attn_weight = top_grad * val;
|
||||
*grad_sampling_loc = width * grad_w_weight * top_grad_value;
|
||||
*(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value;
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t *&bottom_data, const int &height, const int &width,
|
||||
const int &nheads, const int &channels, const scalar_t &h,
|
||||
const scalar_t &w, const int &m, const int &c,
|
||||
const scalar_t &top_grad, const scalar_t &attn_weight,
|
||||
scalar_t *&grad_value, scalar_t *grad_sampling_loc,
|
||||
scalar_t *grad_attn_weight) {
|
||||
const int h_low = floor(h);
|
||||
const int w_low = floor(w);
|
||||
const int h_high = h_low + 1;
|
||||
const int w_high = w_low + 1;
|
||||
|
||||
const scalar_t lh = h - h_low;
|
||||
const scalar_t lw = w - w_low;
|
||||
const scalar_t hh = 1 - lh, hw = 1 - lw;
|
||||
|
||||
const int w_stride = nheads * channels;
|
||||
const int h_stride = width * w_stride;
|
||||
const int h_low_ptr_offset = h_low * h_stride;
|
||||
const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
|
||||
const int w_low_ptr_offset = w_low * w_stride;
|
||||
const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
|
||||
const int base_ptr = m * channels + c;
|
||||
|
||||
const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
|
||||
const scalar_t top_grad_value = top_grad * attn_weight;
|
||||
scalar_t grad_h_weight = 0, grad_w_weight = 0;
|
||||
|
||||
scalar_t v1 = 0;
|
||||
if (h_low >= 0 && w_low >= 0) {
|
||||
const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
|
||||
v1 = bottom_data[ptr1];
|
||||
grad_h_weight -= hw * v1;
|
||||
grad_w_weight -= hh * v1;
|
||||
atomicAdd(grad_value + ptr1, w1 * top_grad_value);
|
||||
}
|
||||
scalar_t v2 = 0;
|
||||
if (h_low >= 0 && w_high <= width - 1) {
|
||||
const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
|
||||
v2 = bottom_data[ptr2];
|
||||
grad_h_weight -= lw * v2;
|
||||
grad_w_weight += hh * v2;
|
||||
atomicAdd(grad_value + ptr2, w2 * top_grad_value);
|
||||
}
|
||||
scalar_t v3 = 0;
|
||||
if (h_high <= height - 1 && w_low >= 0) {
|
||||
const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
|
||||
v3 = bottom_data[ptr3];
|
||||
grad_h_weight += hw * v3;
|
||||
grad_w_weight -= lh * v3;
|
||||
atomicAdd(grad_value + ptr3, w3 * top_grad_value);
|
||||
}
|
||||
scalar_t v4 = 0;
|
||||
if (h_high <= height - 1 && w_high <= width - 1) {
|
||||
const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
|
||||
v4 = bottom_data[ptr4];
|
||||
grad_h_weight += lw * v4;
|
||||
grad_w_weight += lh * v4;
|
||||
atomicAdd(grad_value + ptr4, w4 * top_grad_value);
|
||||
}
|
||||
|
||||
const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
|
||||
atomicAdd(grad_attn_weight, top_grad * val);
|
||||
atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value);
|
||||
atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void ms_deformable_im2col_gpu_kernel(const int n, const scalar_t *data_value,
|
||||
const int64_t *data_spatial_shapes,
|
||||
const int64_t *data_level_start_index,
|
||||
const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
|
||||
const int batch_size, const int spatial_size, const int num_heads,
|
||||
const int channels, const int num_levels, const int num_query,
|
||||
const int num_point, scalar_t *data_col) {
|
||||
CUDA_KERNEL_LOOP(index, n) {
|
||||
int _temp = index;
|
||||
const int c_col = _temp % channels;
|
||||
_temp /= channels;
|
||||
const int sampling_index = _temp;
|
||||
const int m_col = _temp % num_heads;
|
||||
_temp /= num_heads;
|
||||
const int q_col = _temp % num_query;
|
||||
_temp /= num_query;
|
||||
const int b_col = _temp;
|
||||
|
||||
scalar_t *data_col_ptr = data_col + index;
|
||||
int data_weight_ptr = sampling_index * num_levels * num_point;
|
||||
int data_loc_w_ptr = data_weight_ptr << 1;
|
||||
const int qid_stride = num_heads * channels;
|
||||
const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
|
||||
scalar_t col = 0;
|
||||
|
||||
for (int l_col = 0; l_col < num_levels; ++l_col) {
|
||||
const int level_start_id = data_level_start_index[l_col];
|
||||
const int spatial_h_ptr = l_col << 1;
|
||||
const int spatial_h = data_spatial_shapes[spatial_h_ptr];
|
||||
const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
|
||||
const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride);
|
||||
for (int p_col = 0; p_col < num_point; ++p_col) {
|
||||
const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
|
||||
const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
|
||||
const scalar_t weight = data_attn_weight[data_weight_ptr];
|
||||
|
||||
const scalar_t h_im = loc_h * spatial_h - 0.5;
|
||||
const scalar_t w_im = loc_w * spatial_w - 0.5;
|
||||
|
||||
if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
|
||||
col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels,
|
||||
h_im, w_im, m_col, c_col) *
|
||||
weight;
|
||||
}
|
||||
|
||||
data_weight_ptr += 1;
|
||||
data_loc_w_ptr += 2;
|
||||
}
|
||||
}
|
||||
*data_col_ptr = col;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, unsigned int blockSize>
|
||||
__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(
|
||||
const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
|
||||
const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
|
||||
const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
|
||||
const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
|
||||
scalar_t *grad_attn_weight) {
|
||||
CUDA_KERNEL_LOOP(index, n) {
|
||||
__shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
|
||||
__shared__ scalar_t cache_grad_attn_weight[blockSize];
|
||||
unsigned int tid = threadIdx.x;
|
||||
int _temp = index;
|
||||
const int c_col = _temp % channels;
|
||||
_temp /= channels;
|
||||
const int sampling_index = _temp;
|
||||
const int m_col = _temp % num_heads;
|
||||
_temp /= num_heads;
|
||||
const int q_col = _temp % num_query;
|
||||
_temp /= num_query;
|
||||
const int b_col = _temp;
|
||||
|
||||
const scalar_t top_grad = grad_col[index];
|
||||
|
||||
int data_weight_ptr = sampling_index * num_levels * num_point;
|
||||
int data_loc_w_ptr = data_weight_ptr << 1;
|
||||
const int grad_sampling_ptr = data_weight_ptr;
|
||||
grad_sampling_loc += grad_sampling_ptr << 1;
|
||||
grad_attn_weight += grad_sampling_ptr;
|
||||
const int grad_weight_stride = 1;
|
||||
const int grad_loc_stride = 2;
|
||||
const int qid_stride = num_heads * channels;
|
||||
const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
|
||||
|
||||
for (int l_col = 0; l_col < num_levels; ++l_col) {
|
||||
const int level_start_id = data_level_start_index[l_col];
|
||||
const int spatial_h_ptr = l_col << 1;
|
||||
const int spatial_h = data_spatial_shapes[spatial_h_ptr];
|
||||
const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
|
||||
const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
|
||||
const scalar_t *data_value_ptr = data_value + value_ptr_offset;
|
||||
scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
|
||||
|
||||
for (int p_col = 0; p_col < num_point; ++p_col) {
|
||||
const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
|
||||
const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
|
||||
const scalar_t weight = data_attn_weight[data_weight_ptr];
|
||||
|
||||
const scalar_t h_im = loc_h * spatial_h - 0.5;
|
||||
const scalar_t w_im = loc_w * spatial_w - 0.5;
|
||||
*(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
|
||||
*(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
|
||||
*(cache_grad_attn_weight + threadIdx.x) = 0;
|
||||
if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
|
||||
ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
|
||||
w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
|
||||
cache_grad_sampling_loc + (threadIdx.x << 1),
|
||||
cache_grad_attn_weight + threadIdx.x);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
scalar_t _grad_w = cache_grad_sampling_loc[0], _grad_h = cache_grad_sampling_loc[1],
|
||||
_grad_a = cache_grad_attn_weight[0];
|
||||
int sid = 2;
|
||||
for (unsigned int tid = 1; tid < blockSize; ++tid) {
|
||||
_grad_w += cache_grad_sampling_loc[sid];
|
||||
_grad_h += cache_grad_sampling_loc[sid + 1];
|
||||
_grad_a += cache_grad_attn_weight[tid];
|
||||
sid += 2;
|
||||
}
|
||||
|
||||
*grad_sampling_loc = _grad_w;
|
||||
*(grad_sampling_loc + 1) = _grad_h;
|
||||
*grad_attn_weight = _grad_a;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
data_weight_ptr += 1;
|
||||
data_loc_w_ptr += 2;
|
||||
grad_attn_weight += grad_weight_stride;
|
||||
grad_sampling_loc += grad_loc_stride;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, unsigned int blockSize>
|
||||
__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(
|
||||
const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
|
||||
const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
|
||||
const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
|
||||
const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
|
||||
scalar_t *grad_attn_weight) {
|
||||
CUDA_KERNEL_LOOP(index, n) {
|
||||
__shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
|
||||
__shared__ scalar_t cache_grad_attn_weight[blockSize];
|
||||
unsigned int tid = threadIdx.x;
|
||||
int _temp = index;
|
||||
const int c_col = _temp % channels;
|
||||
_temp /= channels;
|
||||
const int sampling_index = _temp;
|
||||
const int m_col = _temp % num_heads;
|
||||
_temp /= num_heads;
|
||||
const int q_col = _temp % num_query;
|
||||
_temp /= num_query;
|
||||
const int b_col = _temp;
|
||||
|
||||
const scalar_t top_grad = grad_col[index];
|
||||
|
||||
int data_weight_ptr = sampling_index * num_levels * num_point;
|
||||
int data_loc_w_ptr = data_weight_ptr << 1;
|
||||
const int grad_sampling_ptr = data_weight_ptr;
|
||||
grad_sampling_loc += grad_sampling_ptr << 1;
|
||||
grad_attn_weight += grad_sampling_ptr;
|
||||
const int grad_weight_stride = 1;
|
||||
const int grad_loc_stride = 2;
|
||||
const int qid_stride = num_heads * channels;
|
||||
const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
|
||||
|
||||
for (int l_col = 0; l_col < num_levels; ++l_col) {
|
||||
const int level_start_id = data_level_start_index[l_col];
|
||||
const int spatial_h_ptr = l_col << 1;
|
||||
const int spatial_h = data_spatial_shapes[spatial_h_ptr];
|
||||
const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
|
||||
const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
|
||||
const scalar_t *data_value_ptr = data_value + value_ptr_offset;
|
||||
scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
|
||||
|
||||
for (int p_col = 0; p_col < num_point; ++p_col) {
|
||||
const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
|
||||
const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
|
||||
const scalar_t weight = data_attn_weight[data_weight_ptr];
|
||||
|
||||
const scalar_t h_im = loc_h * spatial_h - 0.5;
|
||||
const scalar_t w_im = loc_w * spatial_w - 0.5;
|
||||
*(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
|
||||
*(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
|
||||
*(cache_grad_attn_weight + threadIdx.x) = 0;
|
||||
if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
|
||||
ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
|
||||
w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
|
||||
cache_grad_sampling_loc + (threadIdx.x << 1),
|
||||
cache_grad_attn_weight + threadIdx.x);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (unsigned int s = blockSize / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
const unsigned int xid1 = tid << 1;
|
||||
const unsigned int xid2 = (tid + s) << 1;
|
||||
cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
|
||||
cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
|
||||
cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
*grad_sampling_loc = cache_grad_sampling_loc[0];
|
||||
*(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
|
||||
*grad_attn_weight = cache_grad_attn_weight[0];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
data_weight_ptr += 1;
|
||||
data_loc_w_ptr += 2;
|
||||
grad_attn_weight += grad_weight_stride;
|
||||
grad_sampling_loc += grad_loc_stride;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(
|
||||
const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
|
||||
const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
|
||||
const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
|
||||
const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
|
||||
scalar_t *grad_attn_weight) {
|
||||
CUDA_KERNEL_LOOP(index, n) {
|
||||
extern __shared__ int _s[];
|
||||
scalar_t *cache_grad_sampling_loc = (scalar_t *)_s;
|
||||
scalar_t *cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
|
||||
unsigned int tid = threadIdx.x;
|
||||
int _temp = index;
|
||||
const int c_col = _temp % channels;
|
||||
_temp /= channels;
|
||||
const int sampling_index = _temp;
|
||||
const int m_col = _temp % num_heads;
|
||||
_temp /= num_heads;
|
||||
const int q_col = _temp % num_query;
|
||||
_temp /= num_query;
|
||||
const int b_col = _temp;
|
||||
|
||||
const scalar_t top_grad = grad_col[index];
|
||||
|
||||
int data_weight_ptr = sampling_index * num_levels * num_point;
|
||||
int data_loc_w_ptr = data_weight_ptr << 1;
|
||||
const int grad_sampling_ptr = data_weight_ptr;
|
||||
grad_sampling_loc += grad_sampling_ptr << 1;
|
||||
grad_attn_weight += grad_sampling_ptr;
|
||||
const int grad_weight_stride = 1;
|
||||
const int grad_loc_stride = 2;
|
||||
const int qid_stride = num_heads * channels;
|
||||
const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
|
||||
|
||||
for (int l_col = 0; l_col < num_levels; ++l_col) {
|
||||
const int level_start_id = data_level_start_index[l_col];
|
||||
const int spatial_h_ptr = l_col << 1;
|
||||
const int spatial_h = data_spatial_shapes[spatial_h_ptr];
|
||||
const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
|
||||
const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
|
||||
const scalar_t *data_value_ptr = data_value + value_ptr_offset;
|
||||
scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
|
||||
|
||||
for (int p_col = 0; p_col < num_point; ++p_col) {
|
||||
const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
|
||||
const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
|
||||
const scalar_t weight = data_attn_weight[data_weight_ptr];
|
||||
|
||||
const scalar_t h_im = loc_h * spatial_h - 0.5;
|
||||
const scalar_t w_im = loc_w * spatial_w - 0.5;
|
||||
*(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
|
||||
*(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
|
||||
*(cache_grad_attn_weight + threadIdx.x) = 0;
|
||||
if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
|
||||
ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
|
||||
w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
|
||||
cache_grad_sampling_loc + (threadIdx.x << 1),
|
||||
cache_grad_attn_weight + threadIdx.x);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
scalar_t _grad_w = cache_grad_sampling_loc[0], _grad_h = cache_grad_sampling_loc[1],
|
||||
_grad_a = cache_grad_attn_weight[0];
|
||||
int sid = 2;
|
||||
for (unsigned int tid = 1; tid < blockDim.x; ++tid) {
|
||||
_grad_w += cache_grad_sampling_loc[sid];
|
||||
_grad_h += cache_grad_sampling_loc[sid + 1];
|
||||
_grad_a += cache_grad_attn_weight[tid];
|
||||
sid += 2;
|
||||
}
|
||||
|
||||
*grad_sampling_loc = _grad_w;
|
||||
*(grad_sampling_loc + 1) = _grad_h;
|
||||
*grad_attn_weight = _grad_a;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
data_weight_ptr += 1;
|
||||
data_loc_w_ptr += 2;
|
||||
grad_attn_weight += grad_weight_stride;
|
||||
grad_sampling_loc += grad_loc_stride;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(
|
||||
const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
|
||||
const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
|
||||
const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
|
||||
const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
|
||||
scalar_t *grad_attn_weight) {
|
||||
CUDA_KERNEL_LOOP(index, n) {
|
||||
extern __shared__ int _s[];
|
||||
scalar_t *cache_grad_sampling_loc = (scalar_t *)_s;
|
||||
scalar_t *cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
|
||||
unsigned int tid = threadIdx.x;
|
||||
int _temp = index;
|
||||
const int c_col = _temp % channels;
|
||||
_temp /= channels;
|
||||
const int sampling_index = _temp;
|
||||
const int m_col = _temp % num_heads;
|
||||
_temp /= num_heads;
|
||||
const int q_col = _temp % num_query;
|
||||
_temp /= num_query;
|
||||
const int b_col = _temp;
|
||||
|
||||
const scalar_t top_grad = grad_col[index];
|
||||
|
||||
int data_weight_ptr = sampling_index * num_levels * num_point;
|
||||
int data_loc_w_ptr = data_weight_ptr << 1;
|
||||
const int grad_sampling_ptr = data_weight_ptr;
|
||||
grad_sampling_loc += grad_sampling_ptr << 1;
|
||||
grad_attn_weight += grad_sampling_ptr;
|
||||
const int grad_weight_stride = 1;
|
||||
const int grad_loc_stride = 2;
|
||||
const int qid_stride = num_heads * channels;
|
||||
const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
|
||||
|
||||
for (int l_col = 0; l_col < num_levels; ++l_col) {
|
||||
const int level_start_id = data_level_start_index[l_col];
|
||||
const int spatial_h_ptr = l_col << 1;
|
||||
const int spatial_h = data_spatial_shapes[spatial_h_ptr];
|
||||
const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
|
||||
const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
|
||||
const scalar_t *data_value_ptr = data_value + value_ptr_offset;
|
||||
scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
|
||||
|
||||
for (int p_col = 0; p_col < num_point; ++p_col) {
|
||||
const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
|
||||
const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
|
||||
const scalar_t weight = data_attn_weight[data_weight_ptr];
|
||||
|
||||
const scalar_t h_im = loc_h * spatial_h - 0.5;
|
||||
const scalar_t w_im = loc_w * spatial_w - 0.5;
|
||||
*(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
|
||||
*(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
|
||||
*(cache_grad_attn_weight + threadIdx.x) = 0;
|
||||
if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
|
||||
ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
|
||||
w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
|
||||
cache_grad_sampling_loc + (threadIdx.x << 1),
|
||||
cache_grad_attn_weight + threadIdx.x);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (unsigned int s = blockDim.x / 2, spre = blockDim.x; s > 0; s >>= 1, spre >>= 1) {
|
||||
if (tid < s) {
|
||||
const unsigned int xid1 = tid << 1;
|
||||
const unsigned int xid2 = (tid + s) << 1;
|
||||
cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
|
||||
cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
|
||||
cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
|
||||
if (tid + (s << 1) < spre) {
|
||||
cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
|
||||
cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
|
||||
cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
*grad_sampling_loc = cache_grad_sampling_loc[0];
|
||||
*(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
|
||||
*grad_attn_weight = cache_grad_attn_weight[0];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
data_weight_ptr += 1;
|
||||
data_loc_w_ptr += 2;
|
||||
grad_attn_weight += grad_weight_stride;
|
||||
grad_sampling_loc += grad_loc_stride;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(
|
||||
const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
|
||||
const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
|
||||
const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
|
||||
const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
|
||||
scalar_t *grad_attn_weight) {
|
||||
CUDA_KERNEL_LOOP(index, n) {
|
||||
extern __shared__ int _s[];
|
||||
scalar_t *cache_grad_sampling_loc = (scalar_t *)_s;
|
||||
scalar_t *cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
|
||||
unsigned int tid = threadIdx.x;
|
||||
int _temp = index;
|
||||
const int c_col = _temp % channels;
|
||||
_temp /= channels;
|
||||
const int sampling_index = _temp;
|
||||
const int m_col = _temp % num_heads;
|
||||
_temp /= num_heads;
|
||||
const int q_col = _temp % num_query;
|
||||
_temp /= num_query;
|
||||
const int b_col = _temp;
|
||||
|
||||
const scalar_t top_grad = grad_col[index];
|
||||
|
||||
int data_weight_ptr = sampling_index * num_levels * num_point;
|
||||
int data_loc_w_ptr = data_weight_ptr << 1;
|
||||
const int grad_sampling_ptr = data_weight_ptr;
|
||||
grad_sampling_loc += grad_sampling_ptr << 1;
|
||||
grad_attn_weight += grad_sampling_ptr;
|
||||
const int grad_weight_stride = 1;
|
||||
const int grad_loc_stride = 2;
|
||||
const int qid_stride = num_heads * channels;
|
||||
const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
|
||||
|
||||
for (int l_col = 0; l_col < num_levels; ++l_col) {
|
||||
const int level_start_id = data_level_start_index[l_col];
|
||||
const int spatial_h_ptr = l_col << 1;
|
||||
const int spatial_h = data_spatial_shapes[spatial_h_ptr];
|
||||
const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
|
||||
const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
|
||||
const scalar_t *data_value_ptr = data_value + value_ptr_offset;
|
||||
scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
|
||||
|
||||
for (int p_col = 0; p_col < num_point; ++p_col) {
|
||||
const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
|
||||
const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
|
||||
const scalar_t weight = data_attn_weight[data_weight_ptr];
|
||||
|
||||
const scalar_t h_im = loc_h * spatial_h - 0.5;
|
||||
const scalar_t w_im = loc_w * spatial_w - 0.5;
|
||||
*(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
|
||||
*(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
|
||||
*(cache_grad_attn_weight + threadIdx.x) = 0;
|
||||
if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
|
||||
ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
|
||||
w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
|
||||
cache_grad_sampling_loc + (threadIdx.x << 1),
|
||||
cache_grad_attn_weight + threadIdx.x);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (unsigned int s = blockDim.x / 2, spre = blockDim.x; s > 0; s >>= 1, spre >>= 1) {
|
||||
if (tid < s) {
|
||||
const unsigned int xid1 = tid << 1;
|
||||
const unsigned int xid2 = (tid + s) << 1;
|
||||
cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
|
||||
cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
|
||||
cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
|
||||
if (tid + (s << 1) < spre) {
|
||||
cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
|
||||
cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
|
||||
cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]);
|
||||
atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]);
|
||||
atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
data_weight_ptr += 1;
|
||||
data_loc_w_ptr += 2;
|
||||
grad_attn_weight += grad_weight_stride;
|
||||
grad_sampling_loc += grad_loc_stride;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n, const scalar_t *grad_col, const scalar_t *data_value,
|
||||
const int64_t *data_spatial_shapes,
|
||||
const int64_t *data_level_start_index,
|
||||
const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
|
||||
const int batch_size, const int spatial_size, const int num_heads,
|
||||
const int channels, const int num_levels, const int num_query,
|
||||
const int num_point, scalar_t *grad_value,
|
||||
scalar_t *grad_sampling_loc, scalar_t *grad_attn_weight) {
|
||||
CUDA_KERNEL_LOOP(index, n) {
|
||||
int _temp = index;
|
||||
const int c_col = _temp % channels;
|
||||
_temp /= channels;
|
||||
const int sampling_index = _temp;
|
||||
const int m_col = _temp % num_heads;
|
||||
_temp /= num_heads;
|
||||
const int q_col = _temp % num_query;
|
||||
_temp /= num_query;
|
||||
const int b_col = _temp;
|
||||
|
||||
const scalar_t top_grad = grad_col[index];
|
||||
|
||||
int data_weight_ptr = sampling_index * num_levels * num_point;
|
||||
int data_loc_w_ptr = data_weight_ptr << 1;
|
||||
const int grad_sampling_ptr = data_weight_ptr;
|
||||
grad_sampling_loc += grad_sampling_ptr << 1;
|
||||
grad_attn_weight += grad_sampling_ptr;
|
||||
const int grad_weight_stride = 1;
|
||||
const int grad_loc_stride = 2;
|
||||
const int qid_stride = num_heads * channels;
|
||||
const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
|
||||
|
||||
for (int l_col = 0; l_col < num_levels; ++l_col) {
|
||||
const int level_start_id = data_level_start_index[l_col];
|
||||
const int spatial_h_ptr = l_col << 1;
|
||||
const int spatial_h = data_spatial_shapes[spatial_h_ptr];
|
||||
const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
|
||||
const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
|
||||
const scalar_t *data_value_ptr = data_value + value_ptr_offset;
|
||||
scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
|
||||
|
||||
for (int p_col = 0; p_col < num_point; ++p_col) {
|
||||
const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
|
||||
const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
|
||||
const scalar_t weight = data_attn_weight[data_weight_ptr];
|
||||
|
||||
const scalar_t h_im = loc_h * spatial_h - 0.5;
|
||||
const scalar_t w_im = loc_w * spatial_w - 0.5;
|
||||
if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
|
||||
ms_deform_attn_col2im_bilinear_gm(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
|
||||
w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
|
||||
grad_sampling_loc, grad_attn_weight);
|
||||
}
|
||||
data_weight_ptr += 1;
|
||||
data_loc_w_ptr += 2;
|
||||
grad_attn_weight += grad_weight_stride;
|
||||
grad_sampling_loc += grad_loc_stride;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void ms_deformable_im2col_cuda(cudaStream_t stream, const scalar_t *data_value, const int64_t *data_spatial_shapes,
|
||||
const int64_t *data_level_start_index, const scalar_t *data_sampling_loc,
|
||||
const scalar_t *data_attn_weight, const int batch_size, const int spatial_size,
|
||||
const int num_heads, const int channels, const int num_levels, const int num_query,
|
||||
const int num_point, scalar_t *data_col) {
|
||||
const int num_kernels = batch_size * num_query * num_heads * channels;
|
||||
const int num_actual_kernels = batch_size * num_query * num_heads * channels;
|
||||
const int num_threads = CUDA_NUM_THREADS;
|
||||
ms_deformable_im2col_gpu_kernel<scalar_t><<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight,
|
||||
batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col);
|
||||
|
||||
cudaError_t err = cudaGetLastError();
|
||||
if (err != cudaSuccess) {
|
||||
printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void ms_deformable_col2im_cuda(cudaStream_t stream, const scalar_t *grad_col, const scalar_t *data_value,
|
||||
const int64_t *data_spatial_shapes, const int64_t *data_level_start_index,
|
||||
const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
|
||||
const int batch_size, const int spatial_size, const int num_heads, const int channels,
|
||||
const int num_levels, const int num_query, const int num_point, scalar_t *grad_value,
|
||||
scalar_t *grad_sampling_loc, scalar_t *grad_attn_weight) {
|
||||
const int num_threads = (channels > CUDA_NUM_THREADS) ? CUDA_NUM_THREADS : channels;
|
||||
const int num_kernels = batch_size * num_query * num_heads * channels;
|
||||
const int num_actual_kernels = batch_size * num_query * num_heads * channels;
|
||||
if (channels > 1024) {
|
||||
if ((channels & 1023) == 0) {
|
||||
ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks<scalar_t>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, num_threads * 3 * sizeof(scalar_t),
|
||||
stream>>>(num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels,
|
||||
num_levels, num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
} else {
|
||||
ms_deformable_col2im_gpu_kernel_gm<scalar_t>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc,
|
||||
data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point,
|
||||
grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
}
|
||||
} else {
|
||||
switch (channels) {
|
||||
case 1:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 1>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 2:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 2>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 4:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 4>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 8:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 8>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 16:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 16>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 32:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 32>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 64:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 64>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 128:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 128>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 256:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 256>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 512:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 512>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
case 1024:
|
||||
ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 1024>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
|
||||
num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
|
||||
num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
break;
|
||||
default:
|
||||
if (channels < 64) {
|
||||
ms_deformable_col2im_gpu_kernel_shm_reduce_v1<scalar_t>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, num_threads * 3 * sizeof(scalar_t),
|
||||
stream>>>(num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels,
|
||||
num_levels, num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
} else {
|
||||
ms_deformable_col2im_gpu_kernel_shm_reduce_v2<scalar_t>
|
||||
<<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, num_threads * 3 * sizeof(scalar_t),
|
||||
stream>>>(num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
|
||||
data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels,
|
||||
num_levels, num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
|
||||
}
|
||||
}
|
||||
}
|
||||
cudaError_t err = cudaGetLastError();
|
||||
if (err != cudaSuccess) {
|
||||
printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
||||
|
||||
#include "ms_deform_attn_cuda.h"
|
||||
|
||||
namespace groundingdino {
|
||||
|
||||
at::Tensor ms_deform_attn_forward(const at::Tensor &value, const at::Tensor &spatial_shapes,
|
||||
const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
|
||||
const at::Tensor &attn_weight, const int im2col_step) {
|
||||
if (value.is_cuda()) {
|
||||
return ms_deform_attn_cuda_forward(value, spatial_shapes, level_start_index, sampling_loc, attn_weight,
|
||||
im2col_step);
|
||||
}
|
||||
AT_ERROR("Not implemented on the CPU");
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> ms_deform_attn_backward(const at::Tensor &value, const at::Tensor &spatial_shapes,
|
||||
const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
|
||||
const at::Tensor &attn_weight, const at::Tensor &grad_output,
|
||||
const int im2col_step) {
|
||||
if (value.is_cuda()) {
|
||||
return ms_deform_attn_cuda_backward(value, spatial_shapes, level_start_index, sampling_loc, attn_weight,
|
||||
grad_output, im2col_step);
|
||||
}
|
||||
AT_ERROR("Not implemented on the CPU");
|
||||
}
|
||||
|
||||
} // namespace groundingdino
|
||||
|
||||
void pybind_grounding_dino_ext(py::module &m) {
|
||||
m.def("ms_deform_attn_forward", &groundingdino::ms_deform_attn_forward, "ms_deform_attn_forward");
|
||||
m.def("ms_deform_attn_backward", &groundingdino::ms_deform_attn_backward, "ms_deform_attn_backward");
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#ifndef COMMON_H
|
||||
#define COMMON_H
|
||||
|
||||
#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int
|
||||
#define EIGEN_RUNTIME_NO_MALLOC
|
||||
|
||||
#define EPS 1e-6
|
||||
#define PI 3.14159265358979323846
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,52 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#ifndef DISPATCH_H
|
||||
#define DISPATCH_H
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include "so3.h"
|
||||
#include "rxso3.h"
|
||||
#include "se3.h"
|
||||
#include "sim3.h"
|
||||
|
||||
#define PRIVATE_CASE_TYPE(group_index, enum_type, type, ...) \
|
||||
case enum_type: { \
|
||||
using scalar_t = type; \
|
||||
switch (group_index) { \
|
||||
case 1: { \
|
||||
using group_t = SO3<type>; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case 2: { \
|
||||
using group_t = RxSO3<type>; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case 3: { \
|
||||
using group_t = SE3<type>; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case 4: { \
|
||||
using group_t = Sim3<type>; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
} \
|
||||
}
|
||||
|
||||
#define DISPATCH_GROUP_AND_FLOATING_TYPES(GROUP_INDEX, TYPE, NAME, ...) \
|
||||
[&] { \
|
||||
const auto& the_type = TYPE; \
|
||||
/* don't use TYPE again in case it is an expensive or side-effect op */ \
|
||||
switch (the_type) { \
|
||||
PRIVATE_CASE_TYPE(GROUP_INDEX, at::ScalarType::Double, double, __VA_ARGS__) \
|
||||
PRIVATE_CASE_TYPE(GROUP_INDEX, at::ScalarType::Float, float, __VA_ARGS__) \
|
||||
default: \
|
||||
break; \
|
||||
} \
|
||||
}()
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,336 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
#include "lietorch_gpu.h"
|
||||
#include "lietorch_cpu.h"
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
||||
|
||||
/* Interface for cuda and c++ group operations
|
||||
|
||||
enum group_t { SO3=1, SE3=2, Sim3=3 };
|
||||
X, Y, Z: (uppercase) Lie Group Elements
|
||||
a, b, c: (lowercase) Lie Algebra Elements
|
||||
*/
|
||||
|
||||
// Unary operations
|
||||
torch::Tensor expm(int group_index, torch::Tensor a) {
|
||||
CHECK_CONTIGUOUS(a);
|
||||
if (a.device().type() == torch::DeviceType::CPU) {
|
||||
return exp_forward_cpu(group_index, a);
|
||||
|
||||
} else if (a.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
|
||||
return exp_forward_gpu(group_index, a);
|
||||
}
|
||||
|
||||
return a;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> expm_backward(int group_index, torch::Tensor grad, torch::Tensor a) {
|
||||
CHECK_CONTIGUOUS(a);
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
if (a.device().type() == torch::DeviceType::CPU) {
|
||||
return exp_backward_cpu(group_index, grad, a);
|
||||
|
||||
} else if (a.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
|
||||
return exp_backward_gpu(group_index, grad, a);
|
||||
}
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
torch::Tensor logm(int group_index, torch::Tensor X) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return log_forward_cpu(group_index, X);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return log_forward_gpu(group_index, X);
|
||||
}
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> logm_backward(int group_index, torch::Tensor grad, torch::Tensor X) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return log_backward_cpu(group_index, grad, X);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return log_backward_gpu(group_index, grad, X);
|
||||
}
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
torch::Tensor inv(int group_index, torch::Tensor X) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return inv_forward_cpu(group_index, X);
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return inv_forward_gpu(group_index, X);
|
||||
}
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> inv_backward(int group_index, torch::Tensor grad, torch::Tensor X) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return inv_backward_cpu(group_index, grad, X);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return inv_backward_gpu(group_index, grad, X);
|
||||
}
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
// Binary operations
|
||||
|
||||
torch::Tensor mul(int group_index, torch::Tensor X, torch::Tensor Y) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(Y);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return mul_forward_cpu(group_index, X, Y);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return mul_forward_gpu(group_index, X, Y);
|
||||
}
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> mul_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(Y);
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return mul_backward_cpu(group_index, grad, X, Y);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return mul_backward_gpu(group_index, grad, X, Y);
|
||||
}
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
torch::Tensor adj(int group_index, torch::Tensor X, torch::Tensor a) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(a);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return adj_forward_cpu(group_index, X, a);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return adj_forward_gpu(group_index, X, a);
|
||||
}
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> adj_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(a);
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return adj_backward_cpu(group_index, grad, X, a);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return adj_backward_gpu(group_index, grad, X, a);
|
||||
}
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
torch::Tensor adjT(int group_index, torch::Tensor X, torch::Tensor a) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(a);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return adjT_forward_cpu(group_index, X, a);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return adjT_forward_gpu(group_index, X, a);
|
||||
}
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> adjT_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(a);
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return adjT_backward_cpu(group_index, grad, X, a);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return adjT_backward_gpu(group_index, grad, X, a);
|
||||
}
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
torch::Tensor act(int group_index, torch::Tensor X, torch::Tensor p) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(p);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return act_forward_cpu(group_index, X, p);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return act_forward_gpu(group_index, X, p);
|
||||
}
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> act_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(p);
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return act_backward_cpu(group_index, grad, X, p);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return act_backward_gpu(group_index, grad, X, p);
|
||||
}
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
torch::Tensor act4(int group_index, torch::Tensor X, torch::Tensor p) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(p);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return act4_forward_cpu(group_index, X, p);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return act4_forward_gpu(group_index, X, p);
|
||||
}
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> act4_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(p);
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return act4_backward_cpu(group_index, grad, X, p);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return act4_backward_gpu(group_index, grad, X, p);
|
||||
}
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
torch::Tensor projector(int group_index, torch::Tensor X) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return orthogonal_projector_cpu(group_index, X);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return orthogonal_projector_gpu(group_index, X);
|
||||
}
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
torch::Tensor as_matrix(int group_index, torch::Tensor X) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return as_matrix_forward_cpu(group_index, X);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return as_matrix_forward_gpu(group_index, X);
|
||||
}
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
torch::Tensor Jinv(int group_index, torch::Tensor X, torch::Tensor a) {
|
||||
CHECK_CONTIGUOUS(X);
|
||||
CHECK_CONTIGUOUS(a);
|
||||
|
||||
if (X.device().type() == torch::DeviceType::CPU) {
|
||||
return jleft_forward_cpu(group_index, X, a);
|
||||
|
||||
} else if (X.device().type() == torch::DeviceType::CUDA) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
|
||||
return jleft_forward_gpu(group_index, X, a);
|
||||
}
|
||||
|
||||
return a;
|
||||
}
|
||||
|
||||
void pybind_lietorch_ext(py::module &m) {
|
||||
// {exp, log, inv, mul, adj, adjT, act, act4} forward/backward bindings
|
||||
m.def("expm", &expm, "exp map forward");
|
||||
m.def("expm_backward", &expm_backward, "exp map backward");
|
||||
|
||||
m.def("logm", &logm, "log map forward");
|
||||
m.def("logm_backward", &logm_backward, "log map backward");
|
||||
|
||||
m.def("inv", &inv, "inverse operator");
|
||||
m.def("inv_backward", &inv_backward, "inverse operator backward");
|
||||
|
||||
m.def("mul", &mul, "group operator");
|
||||
m.def("mul_backward", &mul_backward, "group operator backward");
|
||||
|
||||
m.def("adj", &adj, "adjoint operator");
|
||||
m.def("adj_backward", &adj_backward, "adjoint operator backward");
|
||||
|
||||
m.def("adjT", &adjT, "transposed adjoint operator");
|
||||
m.def("adjT_backward", &adjT_backward, "transposed adjoint operator backward");
|
||||
|
||||
m.def("act", &act, "action on point");
|
||||
m.def("act_backward", &act_backward, "action on point backward");
|
||||
|
||||
m.def("act4", &act4, "action on homogeneous point");
|
||||
m.def("act4_backward", &act4_backward, "action on homogeneous point backward");
|
||||
|
||||
// functions with no gradient
|
||||
m.def("as_matrix", &as_matrix, "convert to matrix");
|
||||
m.def("projector", &projector, "orthogonal projection matrix");
|
||||
m.def("Jinv", &Jinv, "left inverse jacobian operator");
|
||||
}
|
||||
@@ -0,0 +1,610 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include "lietorch_cpu.h"
|
||||
#include <eigen3/Eigen/Dense>
|
||||
|
||||
#include <iostream>
|
||||
#include "common.h"
|
||||
#include "dispatch.h"
|
||||
|
||||
#include "so3.h"
|
||||
#include "rxso3.h"
|
||||
#include "se3.h"
|
||||
#include "sim3.h"
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void exp_forward_kernel(const scalar_t* a_ptr, scalar_t* X_ptr, int batch_size) {
|
||||
// exponential map forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Eigen::Map<Data>(X_ptr + i * Group::N) = Group::Exp(a).data();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void exp_backward_kernel(const scalar_t* grad, const scalar_t* a_ptr, scalar_t* da, int batch_size) {
|
||||
// exponential map backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Grad dX(grad + i * Group::N);
|
||||
Eigen::Map<Grad>(da + i * Group::K) = dX * Group::left_jacobian(a);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void log_forward_kernel(const scalar_t* X_ptr, scalar_t* a_ptr, int batch_size) {
|
||||
// logarithm map forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Tangent a = Group(X_ptr + i * Group::N).Log();
|
||||
Eigen::Map<Tangent>(a_ptr + i * Group::K) = a;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void log_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
|
||||
// logarithm map backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Tangent a = Group(X_ptr + i * Group::N).Log();
|
||||
Grad da(grad + i * Group::K);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = da * Group::left_jacobian_inverse(a);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void inv_forward_kernel(const scalar_t* X_ptr, scalar_t* Y_ptr, int batch_size) {
|
||||
// group inverse forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Eigen::Map<Data>(Y_ptr + i * Group::N) = X.inv().data();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void inv_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
|
||||
// group inverse backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group Y = Group(X_ptr + i * Group::N).inv();
|
||||
Grad dY(grad + i * Group::N);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = -dY * Y.Adj();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void mul_forward_kernel(const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* Z_ptr, int batch_size) {
|
||||
// group multiplication forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group Z = Group(X_ptr + i * Group::N) * Group(Y_ptr + i * Group::N);
|
||||
Eigen::Map<Data>(Z_ptr + i * Group::N) = Z.data();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <class Group, typename scalar_t>
|
||||
void mul_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* dX, scalar_t* dY,
|
||||
int batch_size) {
|
||||
// group multiplication backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Grad dZ(grad + i * Group::N);
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = dZ;
|
||||
Eigen::Map<Grad>(dY + i * Group::N) = dZ * X.Adj();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void adj_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
|
||||
// adjoint forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.Adj(a);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void adj_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX, scalar_t* da,
|
||||
int batch_size) {
|
||||
// adjoint backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Grad db(grad + i * Group::K);
|
||||
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Tangent b = X.Adj() * a;
|
||||
|
||||
Eigen::Map<Grad>(da + i * Group::K) = db * X.Adj();
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = -db * Group::adj(b);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void adjT_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
|
||||
// adjoint forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.AdjT(a);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void adjT_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX,
|
||||
scalar_t* da, int batch_size) {
|
||||
// adjoint backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Tangent db(grad + i * Group::K);
|
||||
Grad a(a_ptr + i * Group::K);
|
||||
|
||||
Eigen::Map<Tangent>(da + i * Group::K) = X.Adj(db);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = -a * Group::adj(X.Adj(db));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void act_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
|
||||
// action on point forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
using Point = Eigen::Matrix<scalar_t, 3, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Point p(p_ptr + i * 3);
|
||||
Eigen::Map<Point>(q_ptr + i * 3) = X * p;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void act_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX, scalar_t* dp,
|
||||
int batch_size) {
|
||||
// adjoint backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Point = Eigen::Matrix<scalar_t, 3, 1>;
|
||||
using PointGrad = Eigen::Matrix<scalar_t, 1, 3>;
|
||||
using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Point p(p_ptr + i * 3);
|
||||
PointGrad dq(grad + i * 3);
|
||||
|
||||
Eigen::Map<PointGrad>(dp + i * 3) = dq * X.Matrix().template block<3, 3>(0, 0);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act_jacobian(X * p);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// template <typename Group, typename scalar_t>
|
||||
// void tovec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
|
||||
// // group inverse forward kernel
|
||||
// using Data = Eigen::Matrix<scalar_t,Group::N,1>;
|
||||
// using Grad = Eigen::Matrix<scalar_t,1,Group::N>;
|
||||
|
||||
// at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
// for (int64_t i=start; i<end; i++) {
|
||||
// Group X(X_ptr + i*Group::N);
|
||||
// Grad g(grad + i*Group::N);
|
||||
// Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
|
||||
// }
|
||||
// });
|
||||
// }
|
||||
|
||||
// template <typename Group, typename scalar_t>
|
||||
// void fromvec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
|
||||
// // group inverse forward kernel
|
||||
// using Data = Eigen::Matrix<scalar_t,Group::N,1>;
|
||||
// using Grad = Eigen::Matrix<scalar_t,1,Group::N>;
|
||||
|
||||
// at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
// for (int64_t i=start; i<end; i++) {
|
||||
// Group X(X_ptr + i*Group::N);
|
||||
// Grad g(grad + i*Group::N);
|
||||
// Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
|
||||
// }
|
||||
// });
|
||||
// }
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void act4_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
|
||||
// action on homogeneous point forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
using Point = Eigen::Matrix<scalar_t, 4, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Point p(p_ptr + i * 4);
|
||||
Eigen::Map<Point>(q_ptr + i * 4) = X.act4(p);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void act4_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX,
|
||||
scalar_t* dp, int batch_size) {
|
||||
// action on homogeneous point backward kernel
|
||||
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Point = Eigen::Matrix<scalar_t, 4, 1>;
|
||||
using PointGrad = Eigen::Matrix<scalar_t, 1, 4>;
|
||||
using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Point p(p_ptr + i * 4);
|
||||
PointGrad dq(grad + i * 4);
|
||||
|
||||
Eigen::Map<PointGrad>(dp + i * 4) = dq * X.Matrix4x4();
|
||||
const Point q = X.act4(p);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act4_jacobian(q);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void as_matrix_forward_kernel(const scalar_t* X_ptr, scalar_t* T_ptr, int batch_size) {
|
||||
// group inverse forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
using Matrix4 = Eigen::Matrix<scalar_t, 4, 4, Eigen::RowMajor>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Eigen::Map<Matrix4>(T_ptr + i * 16) = X.Matrix4x4();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void orthogonal_projector_kernel(const scalar_t* X_ptr, scalar_t* P_ptr, int batch_size) {
|
||||
// group inverse forward kernel
|
||||
using Proj = Eigen::Matrix<scalar_t, Group::N, Group::N, Eigen::RowMajor>;
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Eigen::Map<Proj>(P_ptr + i * Group::N * Group::N) = X.orthogonal_projector();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
void jleft_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
|
||||
// left-jacobian inverse action
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
|
||||
for (int64_t i = start; i < end; i++) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Tangent b = Group::left_jacobian_inverse(X.Log()) * a;
|
||||
Eigen::Map<Tangent>(b_ptr + i * Group::K) = b;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// unary operations
|
||||
|
||||
torch::Tensor exp_forward_cpu(int group_id, torch::Tensor a) {
|
||||
int batch_size = a.size(0);
|
||||
torch::Tensor X;
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_forward_kernel", ([&] {
|
||||
X = torch::zeros({batch_size, group_t::N}, a.options());
|
||||
exp_forward_kernel<group_t, scalar_t>(a.data_ptr<scalar_t>(),
|
||||
X.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> exp_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor a) {
|
||||
int batch_size = a.size(0);
|
||||
torch::Tensor da = torch::zeros(a.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_backward_kernel", ([&] {
|
||||
exp_backward_kernel<group_t, scalar_t>(grad.data_ptr<scalar_t>(),
|
||||
a.data_ptr<scalar_t>(),
|
||||
da.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {da};
|
||||
}
|
||||
|
||||
torch::Tensor log_forward_cpu(int group_id, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor a;
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_forward_kernel", ([&] {
|
||||
a = torch::zeros({batch_size, group_t::K}, X.options());
|
||||
log_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
|
||||
a.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return a;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> log_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_backward_kernel", ([&] {
|
||||
log_backward_kernel<group_t, scalar_t>(grad.data_ptr<scalar_t>(),
|
||||
X.data_ptr<scalar_t>(),
|
||||
dX.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX};
|
||||
}
|
||||
|
||||
torch::Tensor inv_forward_cpu(int group_id, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor Y = torch::zeros_like(X);
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "inv_forward_kernel",
|
||||
([&] { inv_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), batch_size); }));
|
||||
|
||||
return Y;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> inv_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "inv_backward_kernel", ([&] {
|
||||
inv_backward_kernel<group_t, scalar_t>(grad.data_ptr<scalar_t>(),
|
||||
X.data_ptr<scalar_t>(),
|
||||
dX.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX};
|
||||
}
|
||||
|
||||
// binary operations
|
||||
torch::Tensor mul_forward_cpu(int group_id, torch::Tensor X, torch::Tensor Y) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor Z = torch::zeros_like(X);
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "mul_forward_kernel", ([&] {
|
||||
mul_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
|
||||
Y.data_ptr<scalar_t>(),
|
||||
Z.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return Z;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> mul_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor dY = torch::zeros(Y.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "mul_backward_kernel", ([&] {
|
||||
mul_backward_kernel<group_t, scalar_t>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(),
|
||||
dX.data_ptr<scalar_t>(), dY.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, dY};
|
||||
}
|
||||
|
||||
torch::Tensor adj_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor b = torch::zeros(a.sizes(), a.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adj_forward_kernel", ([&] {
|
||||
adj_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
|
||||
a.data_ptr<scalar_t>(),
|
||||
b.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return b;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> adj_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor da = torch::zeros(a.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adj_backward_kernel", ([&] {
|
||||
adj_backward_kernel<group_t, scalar_t>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(),
|
||||
dX.data_ptr<scalar_t>(), da.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, da};
|
||||
}
|
||||
|
||||
torch::Tensor adjT_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor b = torch::zeros(a.sizes(), a.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adjT_forward_kernel", ([&] {
|
||||
adjT_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
|
||||
a.data_ptr<scalar_t>(),
|
||||
b.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return b;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> adjT_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor da = torch::zeros(a.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adjT_backward_kernel", ([&] {
|
||||
adjT_backward_kernel<group_t, scalar_t>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(),
|
||||
dX.data_ptr<scalar_t>(), da.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, da};
|
||||
}
|
||||
|
||||
torch::Tensor act_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor q = torch::zeros(p.sizes(), p.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act_forward_kernel", ([&] {
|
||||
act_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
|
||||
p.data_ptr<scalar_t>(),
|
||||
q.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return q;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> act_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act_backward_kernel", ([&] {
|
||||
act_backward_kernel<group_t, scalar_t>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(),
|
||||
dX.data_ptr<scalar_t>(), dp.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, dp};
|
||||
}
|
||||
|
||||
torch::Tensor act4_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor q = torch::zeros(p.sizes(), p.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act4_forward_kernel", ([&] {
|
||||
act4_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
|
||||
p.data_ptr<scalar_t>(),
|
||||
q.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return q;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> act4_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act4_backward_kernel", ([&] {
|
||||
act4_backward_kernel<group_t, scalar_t>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(),
|
||||
dX.data_ptr<scalar_t>(), dp.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, dp};
|
||||
}
|
||||
|
||||
torch::Tensor as_matrix_forward_cpu(int group_id, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor T4x4 = torch::zeros({X.size(0), 4, 4}, X.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "as_matrix_forward_kernel", ([&] {
|
||||
as_matrix_forward_kernel<group_t, scalar_t>(
|
||||
X.data_ptr<scalar_t>(), T4x4.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return T4x4;
|
||||
}
|
||||
|
||||
torch::Tensor orthogonal_projector_cpu(int group_id, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor P;
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "orthogonal_projector_kernel", ([&] {
|
||||
P = torch::zeros({X.size(0), group_t::N, group_t::N}, X.options());
|
||||
orthogonal_projector_kernel<group_t, scalar_t>(
|
||||
X.data_ptr<scalar_t>(), P.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return P;
|
||||
}
|
||||
|
||||
torch::Tensor jleft_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor b = torch::zeros(a.sizes(), a.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "jleft_forward_kernel", ([&] {
|
||||
jleft_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
|
||||
a.data_ptr<scalar_t>(),
|
||||
b.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return b;
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#ifndef LIETORCH_CPU_H_
|
||||
#define LIETORCH_CPU_H_
|
||||
|
||||
#include <vector>
|
||||
#include <torch/extension.h>
|
||||
|
||||
// unary operations
|
||||
torch::Tensor exp_forward_cpu(int, torch::Tensor);
|
||||
std::vector<torch::Tensor> exp_backward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor log_forward_cpu(int, torch::Tensor);
|
||||
std::vector<torch::Tensor> log_backward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor inv_forward_cpu(int, torch::Tensor);
|
||||
std::vector<torch::Tensor> inv_backward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
// binary operations
|
||||
torch::Tensor mul_forward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> mul_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor adj_forward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> adj_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor adjT_forward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> adjT_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor act_forward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> act_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor act4_forward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> act4_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
// conversion operations
|
||||
// std::vector<torch::Tensor> to_vec_backward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
// std::vector<torch::Tensor> from_vec_backward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
// utility operations
|
||||
torch::Tensor orthogonal_projector_cpu(int, torch::Tensor);
|
||||
|
||||
torch::Tensor as_matrix_forward_cpu(int, torch::Tensor);
|
||||
|
||||
torch::Tensor jleft_forward_cpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,562 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include "lietorch_gpu.h"
|
||||
#include <eigen3/Eigen/Dense>
|
||||
|
||||
#include "common.h"
|
||||
#include "dispatch.h"
|
||||
|
||||
#include "so3.h"
|
||||
#include "rxso3.h"
|
||||
#include "se3.h"
|
||||
#include "sim3.h"
|
||||
|
||||
#define GPU_1D_KERNEL_LOOP(i, n) \
|
||||
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
|
||||
|
||||
#define NUM_THREADS 256
|
||||
#define NUM_BLOCKS(batch_size) ((batch_size + NUM_THREADS - 1) / NUM_THREADS)
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void exp_forward_kernel(const scalar_t* a_ptr, scalar_t* X_ptr, int num_threads) {
|
||||
// exponential map forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Eigen::Map<Data>(X_ptr + i * Group::N) = Group::Exp(a).data();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void exp_backward_kernel(const scalar_t* grad, const scalar_t* a_ptr, scalar_t* da, int num_threads) {
|
||||
// exponential map backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Grad dX(grad + i * Group::N);
|
||||
Eigen::Map<Grad>(da + i * Group::K) = dX * Group::left_jacobian(a);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void log_forward_kernel(const scalar_t* X_ptr, scalar_t* a_ptr, int num_threads) {
|
||||
// logarithm map forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Tangent a = Group(X_ptr + i * Group::N).Log();
|
||||
Eigen::Map<Tangent>(a_ptr + i * Group::K) = a;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void log_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int num_threads) {
|
||||
// logarithm map backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Tangent a = Group(X_ptr + i * Group::N).Log();
|
||||
Grad da(grad + i * Group::K);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = da * Group::left_jacobian_inverse(a);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void inv_forward_kernel(const scalar_t* X_ptr, scalar_t* Y_ptr, int num_threads) {
|
||||
// group inverse forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Eigen::Map<Data>(Y_ptr + i * Group::N) = X.inv().data();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void inv_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int num_threads) {
|
||||
// group inverse backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group Y = Group(X_ptr + i * Group::N).inv();
|
||||
Grad dY(grad + i * Group::N);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = -dY * Y.Adj();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void mul_forward_kernel(const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* Z_ptr, int num_threads) {
|
||||
// group multiplication forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group Z = Group(X_ptr + i * Group::N) * Group(Y_ptr + i * Group::N);
|
||||
Eigen::Map<Data>(Z_ptr + i * Group::N) = Z.data();
|
||||
}
|
||||
}
|
||||
|
||||
template <class Group, typename scalar_t>
|
||||
__global__ void mul_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* dX,
|
||||
scalar_t* dY, int num_threads) {
|
||||
// group multiplication backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Grad dZ(grad + i * Group::N);
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = dZ;
|
||||
Eigen::Map<Grad>(dY + i * Group::N) = dZ * X.Adj();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void adj_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int num_threads) {
|
||||
// adjoint forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.Adj(a);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void adj_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX,
|
||||
scalar_t* da, int num_threads) {
|
||||
// adjoint backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Grad db(grad + i * Group::K);
|
||||
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Tangent b = X.Adj() * a;
|
||||
|
||||
Eigen::Map<Grad>(da + i * Group::K) = db * X.Adj();
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = -db * Group::adj(b);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void adjT_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int num_threads) {
|
||||
// adjoint forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.AdjT(a);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void adjT_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX,
|
||||
scalar_t* da, int num_threads) {
|
||||
// adjoint backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Tangent db(grad + i * Group::K);
|
||||
Grad a(a_ptr + i * Group::K);
|
||||
|
||||
Eigen::Map<Tangent>(da + i * Group::K) = X.Adj(db);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = -a * Group::adj(X.Adj(db));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void act_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int num_threads) {
|
||||
// action on point forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
using Point = Eigen::Matrix<scalar_t, 3, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Point p(p_ptr + i * 3);
|
||||
Eigen::Map<Point>(q_ptr + i * 3) = X * p;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void act_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX,
|
||||
scalar_t* dp, int num_threads) {
|
||||
// adjoint backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Point = Eigen::Matrix<scalar_t, 3, 1>;
|
||||
using PointGrad = Eigen::Matrix<scalar_t, 1, 3>;
|
||||
using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Point p(p_ptr + i * 3);
|
||||
PointGrad dq(grad + i * 3);
|
||||
|
||||
Eigen::Map<PointGrad>(dp + i * 3) = dq * X.Matrix4x4().block<3, 3>(0, 0);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act_jacobian(X * p);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void act4_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int num_threads) {
|
||||
// action on point forward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
using Point = Eigen::Matrix<scalar_t, 4, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Point p(p_ptr + i * 4);
|
||||
Eigen::Map<Point>(q_ptr + i * 4) = X.act4(p);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void act4_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX,
|
||||
scalar_t* dp, int num_threads) {
|
||||
// adjoint backward kernel
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
|
||||
using Point = Eigen::Matrix<scalar_t, 4, 1>;
|
||||
using PointGrad = Eigen::Matrix<scalar_t, 1, 4>;
|
||||
using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Point p(p_ptr + i * 4);
|
||||
PointGrad dq(grad + i * 4);
|
||||
|
||||
Eigen::Map<PointGrad>(dp + i * 4) = dq * X.Matrix4x4();
|
||||
const Point q = X.act4(p);
|
||||
Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act4_jacobian(q);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void as_matrix_forward_kernel(const scalar_t* X_ptr, scalar_t* T_ptr, int num_threads) {
|
||||
// convert to 4x4 matrix representation
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
using Matrix4 = Eigen::Matrix<scalar_t, 4, 4, Eigen::RowMajor>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Eigen::Map<Matrix4>(T_ptr + i * 16) = X.Matrix4x4();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void orthogonal_projector_kernel(const scalar_t* X_ptr, scalar_t* P_ptr, int num_threads) {
|
||||
// orthogonal projection matrix
|
||||
using Proj = Eigen::Matrix<scalar_t, Group::N, Group::N, Eigen::RowMajor>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Eigen::Map<Proj>(P_ptr + i * Group::N * Group::N) = X.orthogonal_projector();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Group, typename scalar_t>
|
||||
__global__ void jleft_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int num_threads) {
|
||||
// left jacobian inverse action
|
||||
using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
|
||||
using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
|
||||
|
||||
GPU_1D_KERNEL_LOOP(i, num_threads) {
|
||||
Group X(X_ptr + i * Group::N);
|
||||
Tangent a(a_ptr + i * Group::K);
|
||||
Tangent b = Group::left_jacobian_inverse(X.Log()) * a;
|
||||
Eigen::Map<Tangent>(b_ptr + i * Group::K) = b;
|
||||
}
|
||||
}
|
||||
|
||||
// unary operations
|
||||
|
||||
torch::Tensor exp_forward_gpu(int group_id, torch::Tensor a) {
|
||||
int batch_size = a.size(0);
|
||||
torch::Tensor X;
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_forward_kernel", ([&] {
|
||||
X = torch::zeros({batch_size, group_t::N}, a.options());
|
||||
exp_forward_kernel<group_t, scalar_t>
|
||||
<<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
a.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return X;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> exp_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor a) {
|
||||
int batch_size = a.size(0);
|
||||
torch::Tensor da = torch::zeros(a.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, a.scalar_type(), "exp_backward_kernel", ([&] {
|
||||
exp_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
grad.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), da.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {da};
|
||||
}
|
||||
|
||||
torch::Tensor log_forward_gpu(int group_id, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor a;
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_forward_kernel", ([&] {
|
||||
a = torch::zeros({batch_size, group_t::K}, X.options());
|
||||
log_forward_kernel<group_t, scalar_t>
|
||||
<<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return a;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> log_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "log_backward_kernel", ([&] {
|
||||
log_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX};
|
||||
}
|
||||
|
||||
torch::Tensor inv_forward_gpu(int group_id, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor Y = torch::zeros_like(X);
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "inv_forward_kernel", ([&] {
|
||||
inv_forward_kernel<group_t, scalar_t>
|
||||
<<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return Y;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> inv_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "inv_backward_kernel", ([&] {
|
||||
inv_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX};
|
||||
}
|
||||
|
||||
// binary operations
|
||||
torch::Tensor mul_forward_gpu(int group_id, torch::Tensor X, torch::Tensor Y) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor Z = torch::zeros_like(X);
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "mul_forward_kernel", ([&] {
|
||||
mul_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), Z.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return Z;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> mul_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor dY = torch::zeros(Y.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "mul_backward_kernel", ([&] {
|
||||
mul_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
|
||||
dY.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, dY};
|
||||
}
|
||||
|
||||
torch::Tensor adj_forward_gpu(int group_id, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor b = torch::zeros(a.sizes(), a.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "adj_forward_kernel", ([&] {
|
||||
adj_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), b.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return b;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> adj_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor da = torch::zeros(a.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "adj_backward_kernel", ([&] {
|
||||
adj_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
|
||||
da.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, da};
|
||||
}
|
||||
|
||||
torch::Tensor adjT_forward_gpu(int group_id, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor b = torch::zeros(a.sizes(), a.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "adjT_forward_kernel", ([&] {
|
||||
adjT_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), b.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return b;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> adjT_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor da = torch::zeros(a.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "adjT_backward_kernel", ([&] {
|
||||
adjT_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
|
||||
da.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, da};
|
||||
}
|
||||
|
||||
torch::Tensor act_forward_gpu(int group_id, torch::Tensor X, torch::Tensor p) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor q = torch::zeros(p.sizes(), p.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "act_forward_kernel", ([&] {
|
||||
act_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), q.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return q;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> act_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "act_backward_kernel", ([&] {
|
||||
act_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
|
||||
dp.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, dp};
|
||||
}
|
||||
|
||||
torch::Tensor act4_forward_gpu(int group_id, torch::Tensor X, torch::Tensor p) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor q = torch::zeros(p.sizes(), p.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "act4_forward_kernel", ([&] {
|
||||
act4_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), q.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return q;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> act4_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
|
||||
torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "act4_backward_kernel", ([&] {
|
||||
act4_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
|
||||
dp.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return {dX, dp};
|
||||
}
|
||||
|
||||
torch::Tensor as_matrix_forward_gpu(int group_id, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor T4x4 = torch::zeros({X.size(0), 4, 4}, X.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "as_matrix_forward_kernel", ([&] {
|
||||
as_matrix_forward_kernel<group_t, scalar_t>
|
||||
<<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), T4x4.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return T4x4;
|
||||
}
|
||||
|
||||
torch::Tensor orthogonal_projector_gpu(int group_id, torch::Tensor X) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor P;
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "orthogonal_projector_kernel", ([&] {
|
||||
P = torch::zeros({X.size(0), group_t::N, group_t::N}, X.options());
|
||||
orthogonal_projector_kernel<group_t, scalar_t>
|
||||
<<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), P.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return P;
|
||||
}
|
||||
|
||||
torch::Tensor jleft_forward_gpu(int group_id, torch::Tensor X, torch::Tensor a) {
|
||||
int batch_size = X.size(0);
|
||||
torch::Tensor b = torch::zeros(a.sizes(), a.options());
|
||||
|
||||
DISPATCH_GROUP_AND_FLOATING_TYPES(
|
||||
group_id, X.scalar_type(), "jleft_forward_kernel", ([&] {
|
||||
jleft_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
|
||||
X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), b.data_ptr<scalar_t>(), batch_size);
|
||||
}));
|
||||
|
||||
return b;
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#ifndef LIETORCH_GPU_H_
|
||||
#define LIETORCH_GPU_H_
|
||||
|
||||
#include <vector>
|
||||
#include <torch/extension.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// unary operations
|
||||
torch::Tensor exp_forward_gpu(int, torch::Tensor);
|
||||
std::vector<torch::Tensor> exp_backward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor log_forward_gpu(int, torch::Tensor);
|
||||
std::vector<torch::Tensor> log_backward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor inv_forward_gpu(int, torch::Tensor);
|
||||
std::vector<torch::Tensor> inv_backward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
// binary operations
|
||||
torch::Tensor mul_forward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> mul_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor adj_forward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> adj_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor adjT_forward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> adjT_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor act_forward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> act_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
torch::Tensor act4_forward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
std::vector<torch::Tensor> act4_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
|
||||
|
||||
// conversion operations
|
||||
// std::vector<torch::Tensor> to_vec_backward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
// std::vector<torch::Tensor> from_vec_backward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
// utility operators
|
||||
torch::Tensor orthogonal_projector_gpu(int, torch::Tensor);
|
||||
|
||||
torch::Tensor as_matrix_forward_gpu(int, torch::Tensor);
|
||||
|
||||
torch::Tensor jleft_forward_gpu(int, torch::Tensor, torch::Tensor);
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,313 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#ifndef RxSO3_HEADER
|
||||
#define RxSO3_HEADER
|
||||
|
||||
#include <stdio.h>
|
||||
#include <eigen3/Eigen/Dense>
|
||||
#include <eigen3/Eigen/Geometry>
|
||||
|
||||
#include "common.h"
|
||||
|
||||
template <typename Scalar>
|
||||
class RxSO3 {
|
||||
public:
|
||||
const static int constexpr K = 4; // manifold dimension
|
||||
const static int constexpr N = 5; // embedding dimension
|
||||
|
||||
using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
|
||||
using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
|
||||
using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
|
||||
|
||||
using Tangent = Eigen::Matrix<Scalar, K, 1>;
|
||||
using Data = Eigen::Matrix<Scalar, N, 1>;
|
||||
|
||||
using Point = Eigen::Matrix<Scalar, 3, 1>;
|
||||
using Point4 = Eigen::Matrix<Scalar, 4, 1>;
|
||||
|
||||
using Quaternion = Eigen::Quaternion<Scalar>;
|
||||
using Transformation = Eigen::Matrix<Scalar, 3, 3>;
|
||||
using Adjoint = Eigen::Matrix<Scalar, 4, 4>;
|
||||
|
||||
EIGEN_DEVICE_FUNC RxSO3(Quaternion const& q, Scalar const s) : unit_quaternion(q), scale(s) {
|
||||
unit_quaternion.normalize();
|
||||
};
|
||||
|
||||
EIGEN_DEVICE_FUNC RxSO3(const Scalar* data) : unit_quaternion(data), scale(data[4]) {
|
||||
unit_quaternion.normalize();
|
||||
};
|
||||
|
||||
EIGEN_DEVICE_FUNC RxSO3() {
|
||||
unit_quaternion = Quaternion::Identity();
|
||||
scale = Scalar(1.0);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC RxSO3<Scalar> inv() { return RxSO3<Scalar>(unit_quaternion.conjugate(), 1.0 / scale); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Data data() const {
|
||||
Data data_vec;
|
||||
data_vec << unit_quaternion.coeffs(), scale;
|
||||
return data_vec;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC RxSO3<Scalar> operator*(RxSO3<Scalar> const& other) {
|
||||
return RxSO3<Scalar>(unit_quaternion * other.unit_quaternion, scale * other.scale);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Point operator*(Point const& p) const {
|
||||
const Quaternion& q = unit_quaternion;
|
||||
Point uv = q.vec().cross(p);
|
||||
uv += uv;
|
||||
return scale * (p + q.w() * uv + q.vec().cross(uv));
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
|
||||
Point4 p1;
|
||||
p1 << this->operator*(p.template segment<3>(0)), p(3);
|
||||
return p1;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Adjoint Adj() const {
|
||||
Adjoint Ad = Adjoint::Identity();
|
||||
Ad.template block<3, 3>(0, 0) = unit_quaternion.toRotationMatrix();
|
||||
return Ad;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Transformation Matrix() const { return scale * unit_quaternion.toRotationMatrix(); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 4, 4> Matrix4x4() const {
|
||||
Eigen::Matrix<Scalar, 4, 4> T;
|
||||
T = Eigen::Matrix<Scalar, 4, 4>::Identity();
|
||||
T.template block<3, 3>(0, 0) = Matrix();
|
||||
return T;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 5, 5> orthogonal_projector() const {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 5, 5> J = Eigen::Matrix<Scalar, 5, 5>::Zero();
|
||||
|
||||
J.template block<3, 3>(0, 0) =
|
||||
0.5 * (unit_quaternion.w() * Matrix3::Identity() + SO3<Scalar>::hat(-unit_quaternion.vec()));
|
||||
|
||||
J.template block<1, 3>(3, 0) = 0.5 * (-unit_quaternion.vec());
|
||||
|
||||
// scale
|
||||
J(4, 3) = scale;
|
||||
|
||||
return J;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Transformation Rotation() const { return unit_quaternion.toRotationMatrix(); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
|
||||
|
||||
EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& phi_sigma) {
|
||||
Vector3 const phi = phi_sigma.template segment<3>(0);
|
||||
return SO3<Scalar>::hat(phi) + phi(3) * Transformation::Identity();
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& phi_sigma) {
|
||||
Vector3 const phi = phi_sigma.template segment<3>(0);
|
||||
Matrix3 const Phi = SO3<Scalar>::hat(phi);
|
||||
|
||||
Adjoint ad = Adjoint::Zero();
|
||||
ad.template block<3, 3>(0, 0) = Phi;
|
||||
|
||||
return ad;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent Log() const {
|
||||
using std::abs;
|
||||
using std::atan;
|
||||
using std::sqrt;
|
||||
|
||||
Scalar squared_n = unit_quaternion.vec().squaredNorm();
|
||||
Scalar w = unit_quaternion.w();
|
||||
Scalar two_atan_nbyw_by_n;
|
||||
|
||||
/// Atan-based log thanks to
|
||||
///
|
||||
/// C. Hertzberg et al.:
|
||||
/// "Integrating Generic Sensor Fusion Algorithms with Sound State
|
||||
/// Representation through Encapsulation of Manifolds"
|
||||
/// Information Fusion, 2011
|
||||
|
||||
if (squared_n < EPS * EPS) {
|
||||
two_atan_nbyw_by_n = Scalar(2) / w - Scalar(2.0 / 3.0) * (squared_n) / (w * w * w);
|
||||
} else {
|
||||
Scalar n = sqrt(squared_n);
|
||||
if (abs(w) < EPS) {
|
||||
if (w > Scalar(0)) {
|
||||
two_atan_nbyw_by_n = PI / n;
|
||||
} else {
|
||||
two_atan_nbyw_by_n = -PI / n;
|
||||
}
|
||||
} else {
|
||||
two_atan_nbyw_by_n = Scalar(2) * atan(n / w) / n;
|
||||
}
|
||||
}
|
||||
|
||||
Tangent phi_sigma;
|
||||
phi_sigma << two_atan_nbyw_by_n * unit_quaternion.vec(), log(scale);
|
||||
|
||||
return phi_sigma;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static RxSO3<Scalar> Exp(Tangent const& phi_sigma) {
|
||||
Vector3 phi = phi_sigma.template segment<3>(0);
|
||||
Scalar scale = exp(phi_sigma(3));
|
||||
|
||||
Scalar theta2 = phi.squaredNorm();
|
||||
Scalar theta = sqrt(theta2);
|
||||
Scalar imag_factor;
|
||||
Scalar real_factor;
|
||||
|
||||
if (theta < EPS) {
|
||||
Scalar theta4 = theta2 * theta2;
|
||||
imag_factor = Scalar(0.5) - Scalar(1.0 / 48.0) * theta2 + Scalar(1.0 / 3840.0) * theta4;
|
||||
real_factor = Scalar(1) - Scalar(1.0 / 8.0) * theta2 + Scalar(1.0 / 384.0) * theta4;
|
||||
} else {
|
||||
imag_factor = sin(.5 * theta) / theta;
|
||||
real_factor = cos(.5 * theta);
|
||||
}
|
||||
|
||||
Quaternion q(real_factor, imag_factor * phi.x(), imag_factor * phi.y(), imag_factor * phi.z());
|
||||
return RxSO3<Scalar>(q, scale);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Matrix3 calcW(Tangent const& phi_sigma) {
|
||||
// left jacobian
|
||||
using std::abs;
|
||||
Matrix3 const I = Matrix3::Identity();
|
||||
Scalar const one(1);
|
||||
Scalar const half(0.5);
|
||||
|
||||
Vector3 const phi = phi_sigma.template segment<3>(0);
|
||||
Scalar const sigma = phi_sigma(3);
|
||||
Scalar const theta = phi.norm();
|
||||
|
||||
Matrix3 const Phi = SO3<Scalar>::hat(phi);
|
||||
Matrix3 const Phi2 = Phi * Phi;
|
||||
Scalar const scale = exp(sigma);
|
||||
|
||||
Scalar A, B, C;
|
||||
if (abs(sigma) < EPS) {
|
||||
C = one;
|
||||
if (abs(theta) < EPS) {
|
||||
A = half;
|
||||
B = Scalar(1. / 6.);
|
||||
} else {
|
||||
Scalar theta_sq = theta * theta;
|
||||
A = (one - cos(theta)) / theta_sq;
|
||||
B = (theta - sin(theta)) / (theta_sq * theta);
|
||||
}
|
||||
} else {
|
||||
C = (scale - one) / sigma;
|
||||
if (abs(theta) < EPS) {
|
||||
Scalar sigma_sq = sigma * sigma;
|
||||
A = ((sigma - one) * scale + one) / sigma_sq;
|
||||
B = (scale * half * sigma_sq + scale - one - sigma * scale) / (sigma_sq * sigma);
|
||||
} else {
|
||||
Scalar theta_sq = theta * theta;
|
||||
Scalar a = scale * sin(theta);
|
||||
Scalar b = scale * cos(theta);
|
||||
Scalar c = theta_sq + sigma * sigma;
|
||||
A = (a * sigma + (one - b) * theta) / (theta * c);
|
||||
B = (C - ((b - one) * sigma + a * theta) / (c)) * one / (theta_sq);
|
||||
}
|
||||
}
|
||||
return A * Phi + B * Phi2 + C * I;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Matrix3 calcWInv(Tangent const& phi_sigma) {
|
||||
// left jacobian inverse
|
||||
Matrix3 const I = Matrix3::Identity();
|
||||
Scalar const half(0.5);
|
||||
Scalar const one(1);
|
||||
Scalar const two(2);
|
||||
|
||||
Vector3 const phi = phi_sigma.template segment<3>(0);
|
||||
Scalar const sigma = phi_sigma(3);
|
||||
Scalar const theta = phi.norm();
|
||||
Scalar const scale = exp(sigma);
|
||||
|
||||
Matrix3 const Phi = SO3<Scalar>::hat(phi);
|
||||
Matrix3 const Phi2 = Phi * Phi;
|
||||
Scalar const scale_sq = scale * scale;
|
||||
Scalar const theta_sq = theta * theta;
|
||||
Scalar const sin_theta = sin(theta);
|
||||
Scalar const cos_theta = cos(theta);
|
||||
|
||||
Scalar a, b, c;
|
||||
if (abs(sigma * sigma) < EPS) {
|
||||
c = one - half * sigma;
|
||||
a = -half;
|
||||
if (abs(theta_sq) < EPS) {
|
||||
b = Scalar(1. / 12.);
|
||||
} else {
|
||||
b = (theta * sin_theta + two * cos_theta - two) / (two * theta_sq * (cos_theta - one));
|
||||
}
|
||||
} else {
|
||||
Scalar const scale_cu = scale_sq * scale;
|
||||
c = sigma / (scale - one);
|
||||
if (abs(theta_sq) < EPS) {
|
||||
a = (-sigma * scale + scale - one) / ((scale - one) * (scale - one));
|
||||
b = (scale_sq * sigma - two * scale_sq + scale * sigma + two * scale) /
|
||||
(two * scale_cu - Scalar(6) * scale_sq + Scalar(6) * scale - two);
|
||||
} else {
|
||||
Scalar const s_sin_theta = scale * sin_theta;
|
||||
Scalar const s_cos_theta = scale * cos_theta;
|
||||
a = (theta * s_cos_theta - theta - sigma * s_sin_theta) /
|
||||
(theta * (scale_sq - two * s_cos_theta + one));
|
||||
b = -scale *
|
||||
(theta * s_sin_theta - theta * sin_theta + sigma * s_cos_theta - scale * sigma + sigma * cos_theta -
|
||||
sigma) /
|
||||
(theta_sq * (scale_cu - two * scale * s_cos_theta - scale_sq + two * s_cos_theta + scale - one));
|
||||
}
|
||||
}
|
||||
return a * Phi + b * Phi2 + c * I;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& phi_sigma) {
|
||||
// left jacobian
|
||||
Adjoint J = Adjoint::Identity();
|
||||
Vector3 phi = phi_sigma.template segment<3>(0);
|
||||
J.template block<3, 3>(0, 0) = SO3<Scalar>::left_jacobian(phi);
|
||||
return J;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& phi_sigma) {
|
||||
// left jacobian inverse
|
||||
Adjoint Jinv = Adjoint::Identity();
|
||||
Vector3 phi = phi_sigma.template segment<3>(0);
|
||||
Jinv.template block<3, 3>(0, 0) = SO3<Scalar>::left_jacobian_inverse(phi);
|
||||
return Jinv;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 4> act_jacobian(Point const& p) {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 3, 4> Ja;
|
||||
Ja << SO3<Scalar>::hat(-p), p;
|
||||
return Ja;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 4> act4_jacobian(Point4 const& p) {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 4, 4> J = Eigen::Matrix<Scalar, 4, 4>::Zero();
|
||||
J.template block<3, 3>(0, 0) = SO3<Scalar>::hat(-p.template segment<3>(0));
|
||||
J.template block<3, 1>(0, 3) = p.template segment<3>(0);
|
||||
return J;
|
||||
}
|
||||
|
||||
private:
|
||||
Quaternion unit_quaternion;
|
||||
Scalar scale;
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,214 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#ifndef SE3_HEADER
|
||||
#define SE3_HEADER
|
||||
|
||||
#include <stdio.h>
|
||||
#include <eigen3/Eigen/Dense>
|
||||
#include <eigen3/Eigen/Geometry>
|
||||
|
||||
#include "common.h"
|
||||
#include "so3.h"
|
||||
|
||||
template <typename Scalar>
|
||||
class SE3 {
|
||||
public:
|
||||
const static int constexpr K = 6; // manifold dimension
|
||||
const static int constexpr N = 7; // embedding dimension
|
||||
|
||||
using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
|
||||
using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
|
||||
using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
|
||||
|
||||
using Tangent = Eigen::Matrix<Scalar, K, 1>;
|
||||
using Point = Eigen::Matrix<Scalar, 3, 1>;
|
||||
using Point4 = Eigen::Matrix<Scalar, 4, 1>;
|
||||
using Data = Eigen::Matrix<Scalar, N, 1>;
|
||||
using Transformation = Eigen::Matrix<Scalar, 4, 4>;
|
||||
using Adjoint = Eigen::Matrix<Scalar, K, K>;
|
||||
|
||||
EIGEN_DEVICE_FUNC SE3() { translation = Vector3::Zero(); }
|
||||
|
||||
EIGEN_DEVICE_FUNC SE3(SO3<Scalar> const& so3, Vector3 const& t) : so3(so3), translation(t) {};
|
||||
|
||||
EIGEN_DEVICE_FUNC SE3(const Scalar* data) : translation(data), so3(data + 3) {};
|
||||
|
||||
EIGEN_DEVICE_FUNC SE3<Scalar> inv() { return SE3(so3.inv(), -(so3.inv() * translation)); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Data data() const {
|
||||
Data data_vec;
|
||||
data_vec << translation, so3.data();
|
||||
return data_vec;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC SE3<Scalar> operator*(SE3<Scalar> const& other) {
|
||||
return SE3(so3 * other.so3, translation + so3 * other.translation);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Point operator*(Point const& p) const { return so3 * p + translation; }
|
||||
|
||||
EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
|
||||
Point4 p1;
|
||||
p1 << so3 * p.template segment<3>(0) + translation * p(3), p(3);
|
||||
return p1;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Adjoint Adj() const {
|
||||
Matrix3 R = so3.Matrix();
|
||||
Matrix3 tx = SO3<Scalar>::hat(translation);
|
||||
Matrix3 Zer = Matrix3::Zero();
|
||||
|
||||
Adjoint Ad;
|
||||
Ad << R, tx * R, Zer, R;
|
||||
|
||||
return Ad;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Transformation Matrix() const {
|
||||
Transformation T = Transformation::Identity();
|
||||
T.template block<3, 3>(0, 0) = so3.Matrix();
|
||||
T.template block<3, 1>(0, 3) = translation;
|
||||
return T;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Transformation Matrix4x4() const { return Matrix(); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
|
||||
|
||||
EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& tau_phi) {
|
||||
Vector3 tau = tau_phi.template segment<3>(0);
|
||||
Vector3 phi = tau_phi.template segment<3>(3);
|
||||
|
||||
Transformation TauPhi = Transformation::Zero();
|
||||
TauPhi.template block<3, 3>(0, 0) = SO3<Scalar>::hat(phi);
|
||||
TauPhi.template block<3, 1>(0, 3) = tau;
|
||||
|
||||
return TauPhi;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& tau_phi) {
|
||||
Vector3 tau = tau_phi.template segment<3>(0);
|
||||
Vector3 phi = tau_phi.template segment<3>(3);
|
||||
|
||||
Matrix3 Tau = SO3<Scalar>::hat(tau);
|
||||
Matrix3 Phi = SO3<Scalar>::hat(phi);
|
||||
Matrix3 Zer = Matrix3::Zero();
|
||||
|
||||
Adjoint ad;
|
||||
ad << Phi, Tau, Zer, Phi;
|
||||
|
||||
return ad;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 7, 7> orthogonal_projector() const {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 7, 7> J = Eigen::Matrix<Scalar, 7, 7>::Zero();
|
||||
J.template block<3, 3>(0, 0) = Matrix3::Identity();
|
||||
J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-translation);
|
||||
J.template block<4, 4>(3, 3) = so3.orthogonal_projector();
|
||||
|
||||
return J;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent Log() const {
|
||||
Vector3 phi = so3.Log();
|
||||
Matrix3 Vinv = SO3<Scalar>::left_jacobian_inverse(phi);
|
||||
|
||||
Tangent tau_phi;
|
||||
tau_phi << Vinv * translation, phi;
|
||||
|
||||
return tau_phi;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static SE3<Scalar> Exp(Tangent const& tau_phi) {
|
||||
Vector3 tau = tau_phi.template segment<3>(0);
|
||||
Vector3 phi = tau_phi.template segment<3>(3);
|
||||
|
||||
SO3<Scalar> so3 = SO3<Scalar>::Exp(phi);
|
||||
Vector3 t = SO3<Scalar>::left_jacobian(phi) * tau;
|
||||
|
||||
return SE3<Scalar>(so3, t);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Matrix3 calcQ(Tangent const& tau_phi) {
|
||||
// Q matrix
|
||||
Vector3 tau = tau_phi.template segment<3>(0);
|
||||
Vector3 phi = tau_phi.template segment<3>(3);
|
||||
Matrix3 Tau = SO3<Scalar>::hat(tau);
|
||||
Matrix3 Phi = SO3<Scalar>::hat(phi);
|
||||
|
||||
Scalar theta = phi.norm();
|
||||
Scalar theta_pow2 = theta * theta;
|
||||
Scalar theta_pow4 = theta_pow2 * theta_pow2;
|
||||
|
||||
Scalar coef1 = (theta < EPS) ? Scalar(1.0 / 6.0) - Scalar(1.0 / 120.0) * theta_pow2
|
||||
: (theta - sin(theta)) / (theta_pow2 * theta);
|
||||
|
||||
Scalar coef2 = (theta < EPS) ? Scalar(1.0 / 24.0) - Scalar(1.0 / 720.0) * theta_pow2
|
||||
: (theta_pow2 + 2 * cos(theta) - 2) / (2 * theta_pow4);
|
||||
|
||||
Scalar coef3 = (theta < EPS) ? Scalar(1.0 / 120.0) - Scalar(1.0 / 2520.0) * theta_pow2
|
||||
: (2 * theta - 3 * sin(theta) + theta * cos(theta)) / (2 * theta_pow4 * theta);
|
||||
|
||||
Matrix3 Q = Scalar(0.5) * Tau + coef1 * (Phi * Tau + Tau * Phi + Phi * Tau * Phi) +
|
||||
coef2 * (Phi * Phi * Tau + Tau * Phi * Phi - 3 * Phi * Tau * Phi) +
|
||||
coef3 * (Phi * Tau * Phi * Phi + Phi * Phi * Tau * Phi);
|
||||
|
||||
return Q;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& tau_phi) {
|
||||
// left jacobian
|
||||
Vector3 phi = tau_phi.template segment<3>(3);
|
||||
Matrix3 J = SO3<Scalar>::left_jacobian(phi);
|
||||
Matrix3 Q = SE3<Scalar>::calcQ(tau_phi);
|
||||
Matrix3 Zer = Matrix3::Zero();
|
||||
|
||||
Adjoint J6x6;
|
||||
J6x6 << J, Q, Zer, J;
|
||||
|
||||
return J6x6;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& tau_phi) {
|
||||
// left jacobian inverse
|
||||
Vector3 tau = tau_phi.template segment<3>(0);
|
||||
Vector3 phi = tau_phi.template segment<3>(3);
|
||||
Matrix3 Jinv = SO3<Scalar>::left_jacobian_inverse(phi);
|
||||
Matrix3 Q = SE3<Scalar>::calcQ(tau_phi);
|
||||
Matrix3 Zer = Matrix3::Zero();
|
||||
|
||||
Adjoint J6x6;
|
||||
J6x6 << Jinv, -Jinv * Q * Jinv, Zer, Jinv;
|
||||
|
||||
return J6x6;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 6> act_jacobian(Point const& p) {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 3, 6> J;
|
||||
J.template block<3, 3>(0, 0) = Matrix3::Identity();
|
||||
J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p);
|
||||
return J;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 6> act4_jacobian(Point4 const& p) {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 4, 6> J = Eigen::Matrix<Scalar, 4, 6>::Zero();
|
||||
J.template block<3, 3>(0, 0) = p(3) * Matrix3::Identity();
|
||||
J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p.template segment<3>(0));
|
||||
return J;
|
||||
}
|
||||
|
||||
private:
|
||||
SO3<Scalar> so3;
|
||||
Vector3 translation;
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,203 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#ifndef Sim3_HEADER
|
||||
#define Sim3_HEADER
|
||||
|
||||
#include <stdio.h>
|
||||
#include <iostream>
|
||||
|
||||
#include <eigen3/Eigen/Dense>
|
||||
#include <eigen3/Eigen/Geometry>
|
||||
|
||||
#include "common.h"
|
||||
#include "so3.h"
|
||||
#include "rxso3.h"
|
||||
|
||||
template <typename Scalar>
|
||||
class Sim3 {
|
||||
public:
|
||||
const static int constexpr K = 7; // manifold dimension
|
||||
const static int constexpr N = 8; // embedding dimension
|
||||
|
||||
using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
|
||||
using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
|
||||
using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
|
||||
|
||||
using Tangent = Eigen::Matrix<Scalar, K, 1>;
|
||||
using Point = Eigen::Matrix<Scalar, 3, 1>;
|
||||
using Point4 = Eigen::Matrix<Scalar, 4, 1>;
|
||||
using Data = Eigen::Matrix<Scalar, N, 1>;
|
||||
using Transformation = Eigen::Matrix<Scalar, 4, 4>;
|
||||
using Adjoint = Eigen::Matrix<Scalar, K, K>;
|
||||
|
||||
EIGEN_DEVICE_FUNC Sim3() { translation = Vector3::Zero(); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Sim3(RxSO3<Scalar> const& rxso3, Vector3 const& t) : rxso3(rxso3), translation(t) {};
|
||||
|
||||
EIGEN_DEVICE_FUNC Sim3(const Scalar* data) : translation(data), rxso3(data + 3) {};
|
||||
|
||||
EIGEN_DEVICE_FUNC Sim3<Scalar> inv() { return Sim3<Scalar>(rxso3.inv(), -(rxso3.inv() * translation)); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Data data() const {
|
||||
Data data_vec;
|
||||
data_vec << translation, rxso3.data();
|
||||
return data_vec;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Sim3<Scalar> operator*(Sim3<Scalar> const& other) {
|
||||
return Sim3(rxso3 * other.rxso3, translation + rxso3 * other.translation);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Point operator*(Point const& p) const { return (rxso3 * p) + translation; }
|
||||
|
||||
EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
|
||||
Point4 p1;
|
||||
p1 << rxso3 * p.template segment<3>(0) + p(3) * translation, p(3);
|
||||
return p1;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Transformation Matrix() const {
|
||||
Transformation T = Transformation::Identity();
|
||||
T.template block<3, 3>(0, 0) = rxso3.Matrix();
|
||||
T.template block<3, 1>(0, 3) = translation;
|
||||
return T;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Transformation Matrix4x4() const {
|
||||
Transformation T = Transformation::Identity();
|
||||
T.template block<3, 3>(0, 0) = rxso3.Matrix();
|
||||
T.template block<3, 1>(0, 3) = translation;
|
||||
return T;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 8, 8> orthogonal_projector() const {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 8, 8> J = Eigen::Matrix<Scalar, 8, 8>::Zero();
|
||||
J.template block<3, 3>(0, 0) = Matrix3::Identity();
|
||||
J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-translation);
|
||||
J.template block<3, 1>(0, 6) = translation;
|
||||
J.template block<5, 5>(3, 3) = rxso3.orthogonal_projector();
|
||||
return J;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Adjoint Adj() const {
|
||||
Adjoint Ad = Adjoint::Identity();
|
||||
Matrix3 sR = rxso3.Matrix();
|
||||
Matrix3 tx = SO3<Scalar>::hat(translation);
|
||||
Matrix3 R = rxso3.Rotation();
|
||||
|
||||
Ad.template block<3, 3>(0, 0) = sR;
|
||||
Ad.template block<3, 3>(0, 3) = tx * R;
|
||||
Ad.template block<3, 1>(0, 6) = -translation;
|
||||
Ad.template block<3, 3>(3, 3) = R;
|
||||
|
||||
return Ad;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
|
||||
|
||||
EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& tau_phi_sigma) {
|
||||
Vector3 tau = tau_phi_sigma.template segment<3>(0);
|
||||
Vector3 phi = tau_phi_sigma.template segment<3>(3);
|
||||
Scalar sigma = tau_phi_sigma(6);
|
||||
|
||||
Matrix3 Phi = SO3<Scalar>::hat(phi);
|
||||
Matrix3 I = Matrix3::Identity();
|
||||
|
||||
Transformation Omega = Transformation::Zero();
|
||||
Omega.template block<3, 3>(0, 0) = Phi + sigma * I;
|
||||
Omega.template block<3, 1>(0, 3) = tau;
|
||||
|
||||
return Omega;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& tau_phi_sigma) {
|
||||
Adjoint ad = Adjoint::Zero();
|
||||
Vector3 tau = tau_phi_sigma.template segment<3>(0);
|
||||
Vector3 phi = tau_phi_sigma.template segment<3>(3);
|
||||
Scalar sigma = tau_phi_sigma(6);
|
||||
|
||||
Matrix3 Tau = SO3<Scalar>::hat(tau);
|
||||
Matrix3 Phi = SO3<Scalar>::hat(phi);
|
||||
Matrix3 I = Matrix3::Identity();
|
||||
|
||||
ad.template block<3, 3>(0, 0) = Phi + sigma * I;
|
||||
ad.template block<3, 3>(0, 3) = Tau;
|
||||
ad.template block<3, 1>(0, 6) = -tau;
|
||||
ad.template block<3, 3>(3, 3) = Phi;
|
||||
|
||||
return ad;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent Log() const {
|
||||
// logarithm map
|
||||
Vector4 phi_sigma = rxso3.Log();
|
||||
Matrix3 W = RxSO3<Scalar>::calcW(phi_sigma);
|
||||
|
||||
Tangent tau_phi_sigma;
|
||||
tau_phi_sigma << W.inverse() * translation, phi_sigma;
|
||||
|
||||
return tau_phi_sigma;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Sim3<Scalar> Exp(Tangent const& tau_phi_sigma) {
|
||||
// exponential map
|
||||
Vector3 tau = tau_phi_sigma.template segment<3>(0);
|
||||
Vector4 phi_sigma = tau_phi_sigma.template segment<4>(3);
|
||||
|
||||
RxSO3<Scalar> rxso3 = RxSO3<Scalar>::Exp(phi_sigma);
|
||||
Matrix3 W = RxSO3<Scalar>::calcW(phi_sigma);
|
||||
|
||||
return Sim3<Scalar>(rxso3, W * tau);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& tau_phi_sigma) {
|
||||
// left jacobian
|
||||
Adjoint const Xi = adj(tau_phi_sigma);
|
||||
Adjoint const Xi2 = Xi * Xi;
|
||||
Adjoint const Xi4 = Xi2 * Xi2;
|
||||
|
||||
return Adjoint::Identity() + Scalar(1.0 / 2.0) * Xi + Scalar(1.0 / 6.0) * Xi2 + Scalar(1.0 / 24.0) * Xi * Xi2 +
|
||||
Scalar(1.0 / 120.0) * Xi4;
|
||||
+Scalar(1.0 / 720.0) * Xi* Xi4;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& tau_phi_sigma) {
|
||||
// left jacobian inverse
|
||||
Adjoint const Xi = adj(tau_phi_sigma);
|
||||
Adjoint const Xi2 = Xi * Xi;
|
||||
Adjoint const Xi4 = Xi2 * Xi2;
|
||||
|
||||
return Adjoint::Identity() - Scalar(1.0 / 2.0) * Xi + Scalar(1.0 / 12.0) * Xi2 - Scalar(1.0 / 720.0) * Xi4;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 7> act_jacobian(Point const& p) {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 3, 7> J;
|
||||
J.template block<3, 3>(0, 0) = Matrix3::Identity();
|
||||
J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p);
|
||||
J.template block<3, 1>(0, 6) = p;
|
||||
return J;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 7> act4_jacobian(Point4 const& p) {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 4, 7> J = Eigen::Matrix<Scalar, 4, 7>::Zero();
|
||||
J.template block<3, 3>(0, 0) = p(3) * Matrix3::Identity();
|
||||
J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p.template segment<3>(0));
|
||||
J.template block<3, 1>(0, 6) = p.template segment<3>(0);
|
||||
return J;
|
||||
}
|
||||
|
||||
private:
|
||||
Vector3 translation;
|
||||
RxSO3<Scalar> rxso3;
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,202 @@
|
||||
/**
|
||||
* This file includes code originally from the lietorch repository:
|
||||
* https://github.com/princeton-vl/lietorch
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#ifndef SO3_HEADER
|
||||
#define SO3_HEADER
|
||||
|
||||
#include <cuda.h>
|
||||
#include <stdio.h>
|
||||
#include <eigen3/Eigen/Dense>
|
||||
#include <eigen3/Eigen/Geometry>
|
||||
|
||||
#include "common.h"
|
||||
|
||||
template <typename Scalar>
|
||||
class SO3 {
|
||||
public:
|
||||
const static int constexpr K = 3; // manifold dimension
|
||||
const static int constexpr N = 4; // embedding dimension
|
||||
|
||||
using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
|
||||
using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
|
||||
using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
|
||||
|
||||
using Tangent = Eigen::Matrix<Scalar, K, 1>;
|
||||
using Data = Eigen::Matrix<Scalar, N, 1>;
|
||||
|
||||
using Point = Eigen::Matrix<Scalar, 3, 1>;
|
||||
using Point4 = Eigen::Matrix<Scalar, 4, 1>;
|
||||
using Transformation = Eigen::Matrix<Scalar, 3, 3>;
|
||||
using Adjoint = Eigen::Matrix<Scalar, K, K>;
|
||||
using Quaternion = Eigen::Quaternion<Scalar>;
|
||||
|
||||
EIGEN_DEVICE_FUNC SO3(Quaternion const& q) : unit_quaternion(q) { unit_quaternion.normalize(); };
|
||||
|
||||
EIGEN_DEVICE_FUNC SO3(const Scalar* data) : unit_quaternion(data) { unit_quaternion.normalize(); };
|
||||
|
||||
EIGEN_DEVICE_FUNC SO3() { unit_quaternion = Quaternion::Identity(); }
|
||||
|
||||
EIGEN_DEVICE_FUNC SO3<Scalar> inv() { return SO3<Scalar>(unit_quaternion.conjugate()); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Data data() const { return unit_quaternion.coeffs(); }
|
||||
|
||||
EIGEN_DEVICE_FUNC SO3<Scalar> operator*(SO3<Scalar> const& other) {
|
||||
return SO3(unit_quaternion * other.unit_quaternion);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Point operator*(Point const& p) const {
|
||||
const Quaternion& q = unit_quaternion;
|
||||
Point uv = q.vec().cross(p);
|
||||
uv += uv;
|
||||
return p + q.w() * uv + q.vec().cross(uv);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
|
||||
Point4 p1;
|
||||
p1 << this->operator*(p.template segment<3>(0)), p(3);
|
||||
return p1;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Adjoint Adj() const { return unit_quaternion.toRotationMatrix(); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Transformation Matrix() const { return unit_quaternion.toRotationMatrix(); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 4, 4> Matrix4x4() const {
|
||||
Eigen::Matrix<Scalar, 4, 4> T = Eigen::Matrix<Scalar, 4, 4>::Identity();
|
||||
T.template block<3, 3>(0, 0) = Matrix();
|
||||
return T;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 4, 4> orthogonal_projector() const {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 4, 4> J = Eigen::Matrix<Scalar, 4, 4>::Zero();
|
||||
J.template block<3, 3>(0, 0) =
|
||||
0.5 * (unit_quaternion.w() * Matrix3::Identity() + SO3<Scalar>::hat(-unit_quaternion.vec()));
|
||||
|
||||
J.template block<1, 3>(3, 0) = 0.5 * (-unit_quaternion.vec());
|
||||
return J;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
|
||||
|
||||
EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& phi) {
|
||||
Transformation Phi;
|
||||
Phi << 0.0, -phi(2), phi(1), phi(2), 0.0, -phi(0), -phi(1), phi(0), 0.0;
|
||||
|
||||
return Phi;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& phi) { return SO3<Scalar>::hat(phi); }
|
||||
|
||||
EIGEN_DEVICE_FUNC Tangent Log() const {
|
||||
using std::abs;
|
||||
using std::atan;
|
||||
using std::sqrt;
|
||||
Scalar squared_n = unit_quaternion.vec().squaredNorm();
|
||||
Scalar w = unit_quaternion.w();
|
||||
|
||||
Scalar two_atan_nbyw_by_n;
|
||||
|
||||
/// Atan-based log thanks to
|
||||
///
|
||||
/// C. Hertzberg et al.:
|
||||
/// "Integrating Generic Sensor Fusion Algorithms with Sound State
|
||||
/// Representation through Encapsulation of Manifolds"
|
||||
/// Information Fusion, 2011
|
||||
|
||||
if (squared_n < EPS * EPS) {
|
||||
// If quaternion is normalized and n=0, then w should be 1;
|
||||
// w=0 should never happen here!
|
||||
Scalar squared_w = w * w;
|
||||
two_atan_nbyw_by_n = Scalar(2) / w - Scalar(2.0 / 3.0) * (squared_n) / (w * squared_w);
|
||||
} else {
|
||||
Scalar n = sqrt(squared_n);
|
||||
if (abs(w) < EPS) {
|
||||
if (w > Scalar(0)) {
|
||||
two_atan_nbyw_by_n = Scalar(PI) / n;
|
||||
} else {
|
||||
two_atan_nbyw_by_n = -Scalar(PI) / n;
|
||||
}
|
||||
} else {
|
||||
two_atan_nbyw_by_n = Scalar(2) * atan(n / w) / n;
|
||||
}
|
||||
}
|
||||
|
||||
return two_atan_nbyw_by_n * unit_quaternion.vec();
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static SO3<Scalar> Exp(Tangent const& phi) {
|
||||
Scalar theta2 = phi.squaredNorm();
|
||||
Scalar theta = sqrt(theta2);
|
||||
Scalar imag_factor;
|
||||
Scalar real_factor;
|
||||
|
||||
if (theta < EPS) {
|
||||
Scalar theta4 = theta2 * theta2;
|
||||
imag_factor = Scalar(0.5) - Scalar(1.0 / 48.0) * theta2 + Scalar(1.0 / 3840.0) * theta4;
|
||||
real_factor = Scalar(1) - Scalar(1.0 / 8.0) * theta2 + Scalar(1.0 / 384.0) * theta4;
|
||||
} else {
|
||||
imag_factor = sin(.5 * theta) / theta;
|
||||
real_factor = cos(.5 * theta);
|
||||
}
|
||||
|
||||
Quaternion q(real_factor, imag_factor * phi.x(), imag_factor * phi.y(), imag_factor * phi.z());
|
||||
return SO3<Scalar>(q);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& phi) {
|
||||
// left jacobian
|
||||
Matrix3 I = Matrix3::Identity();
|
||||
Matrix3 Phi = SO3<Scalar>::hat(phi);
|
||||
Matrix3 Phi2 = Phi * Phi;
|
||||
|
||||
Scalar theta2 = phi.squaredNorm();
|
||||
Scalar theta = sqrt(theta2);
|
||||
|
||||
Scalar coef1 = (theta < EPS) ? Scalar(1.0 / 2.0) - Scalar(1.0 / 24.0) * theta2 : (1.0 - cos(theta)) / theta2;
|
||||
|
||||
Scalar coef2 =
|
||||
(theta < EPS) ? Scalar(1.0 / 6.0) - Scalar(1.0 / 120.0) * theta2 : (theta - sin(theta)) / (theta2 * theta);
|
||||
|
||||
return I + coef1 * Phi + coef2 * Phi2;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& phi) {
|
||||
// left jacobian inverse
|
||||
Matrix3 I = Matrix3::Identity();
|
||||
Matrix3 Phi = SO3<Scalar>::hat(phi);
|
||||
Matrix3 Phi2 = Phi * Phi;
|
||||
|
||||
Scalar theta2 = phi.squaredNorm();
|
||||
Scalar theta = sqrt(theta2);
|
||||
Scalar half_theta = Scalar(.5) * theta;
|
||||
|
||||
Scalar coef2 = (theta < EPS)
|
||||
? Scalar(1.0 / 12.0)
|
||||
: (Scalar(1) - theta * cos(half_theta) / (Scalar(2) * sin(half_theta))) / (theta * theta);
|
||||
|
||||
return I + Scalar(-0.5) * Phi + coef2 * Phi2;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 3> act_jacobian(Point const& p) {
|
||||
// jacobian action on a point
|
||||
return SO3<Scalar>::hat(-p);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 3> act4_jacobian(Point4 const& p) {
|
||||
// jacobian action on a point
|
||||
Eigen::Matrix<Scalar, 4, 3> J = Eigen::Matrix<Scalar, 4, 3>::Zero();
|
||||
J.template block<3, 3>(0, 0) = SO3<Scalar>::hat(-p.template segment<3>(0));
|
||||
return J;
|
||||
}
|
||||
|
||||
private:
|
||||
Quaternion unit_quaternion;
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,300 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#define ATOMIC(NAME) \
|
||||
template <typename scalar, size_t size> \
|
||||
struct Atomic##NAME##IntegerImpl; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##IntegerImpl<scalar, 1> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
uint32_t *address_as_ui = (uint32_t *)(address - ((size_t)address & 3)); \
|
||||
uint32_t old = *address_as_ui; \
|
||||
uint32_t shift = ((size_t)address & 3) * 8; \
|
||||
uint32_t sum; \
|
||||
uint32_t assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
sum = OP(val, scalar((old >> shift) & 0xff)); \
|
||||
old = (old & ~(0x000000ff << shift)) | (sum << shift); \
|
||||
old = atomicCAS(address_as_ui, assumed, old); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##IntegerImpl<scalar, 2> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
uint32_t *address_as_ui = (uint32_t *)((char *)address - ((size_t)address & 2)); \
|
||||
uint32_t old = *address_as_ui; \
|
||||
uint32_t sum; \
|
||||
uint32_t newval; \
|
||||
uint32_t assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
sum = OP(val, (size_t)address & 2 ? scalar(old >> 16) : scalar(old & 0xffff)); \
|
||||
newval = (size_t)address & 2 ? (old & 0xffff) | (sum << 16) : (old & 0xffff0000) | sum; \
|
||||
old = atomicCAS(address_as_ui, assumed, newval); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##IntegerImpl<scalar, 4> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
uint32_t *address_as_ui = (uint32_t *)address; \
|
||||
uint32_t old = *address_as_ui; \
|
||||
uint32_t assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = atomicCAS(address_as_ui, assumed, OP(val, (scalar)old)); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##IntegerImpl<scalar, 8> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
unsigned long long *address_as_ull = (unsigned long long *)address; \
|
||||
unsigned long long old = *address_as_ull; \
|
||||
unsigned long long assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = atomicCAS(address_as_ull, assumed, OP(val, (scalar)old)); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar, size_t size> \
|
||||
struct Atomic##NAME##DecimalImpl; \
|
||||
\
|
||||
template <> \
|
||||
struct Atomic##NAME##DecimalImpl<at::Half, 2> { \
|
||||
inline __device__ void operator()(at::Half *address, at::Half val) { \
|
||||
unsigned int *address_as_ui = (unsigned int *)((char *)address - ((size_t)address & 2)); \
|
||||
unsigned int old = *address_as_ui; \
|
||||
unsigned int assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
at::Half hsum; \
|
||||
hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff); \
|
||||
hsum = OP(hsum, val); \
|
||||
old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16) : (old & 0xffff0000) | hsum.x; \
|
||||
old = atomicCAS(address_as_ui, assumed, old); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <> \
|
||||
struct Atomic##NAME##DecimalImpl<at::BFloat16, 2> { \
|
||||
inline __device__ void operator()(at::BFloat16 *address, at::BFloat16 val) { \
|
||||
unsigned int *address_as_ui = (unsigned int *)((char *)address - ((size_t)address & 2)); \
|
||||
unsigned int old = *address_as_ui; \
|
||||
unsigned int assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
at::BFloat16 hsum; \
|
||||
hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff); \
|
||||
hsum = OP(hsum, val); \
|
||||
old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16) : (old & 0xffff0000) | hsum.x; \
|
||||
old = atomicCAS(address_as_ui, assumed, old); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##DecimalImpl<scalar, 4> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
int *address_as_i = (int *)address; \
|
||||
int old = *address_as_i; \
|
||||
int assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = atomicCAS(address_as_i, assumed, __float_as_int(OP(val, __int_as_float(assumed)))); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##DecimalImpl<scalar, 8> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
unsigned long long int *address_as_ull = (unsigned long long int *)address; \
|
||||
unsigned long long int old = *address_as_ull; \
|
||||
unsigned long long int assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = \
|
||||
atomicCAS(address_as_ull, assumed, __double_as_longlong(OP(val, __longlong_as_double(assumed)))); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
};
|
||||
|
||||
#define OP(X, Y) Y + X
|
||||
ATOMIC(Add)
|
||||
#undef OP
|
||||
static inline __device__ void atomAdd(uint8_t *address, uint8_t val) {
|
||||
AtomicAddIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomAdd(int8_t *address, int8_t val) {
|
||||
AtomicAddIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomAdd(int16_t *address, int16_t val) {
|
||||
AtomicAddIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomAdd(int32_t *address, int32_t val) { atomicAdd(address, val); }
|
||||
static inline __device__ void atomAdd(int64_t *address, int64_t val) {
|
||||
AtomicAddIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
#if defined(USE_ROCM) || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700 || CUDA_VERSION < 10000))
|
||||
static inline __device__ void atomAdd(at::Half *address, at::Half val) {
|
||||
AtomicAddDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
#else
|
||||
static inline __device__ void atomAdd(at::Half *address, at::Half val) {
|
||||
atomicAdd(reinterpret_cast<__half *>(address), val);
|
||||
}
|
||||
#endif
|
||||
static inline __device__ void atomAdd(float *address, float val) { atomicAdd(address, val); }
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600 || CUDA_VERSION < 8000)
|
||||
static inline __device__ void atomAdd(double *address, double val) {
|
||||
AtomicAddDecimalImpl<double, sizeof(double)>()(address, val);
|
||||
}
|
||||
#else
|
||||
static inline __device__ void atomAdd(double *address, double val) { atomicAdd(address, val); }
|
||||
#endif
|
||||
static inline __device__ void atomAdd(at::BFloat16 *address, at::BFloat16 val) {
|
||||
AtomicAddDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
|
||||
}
|
||||
|
||||
#define OP(X, Y) Y *X
|
||||
ATOMIC(Mul)
|
||||
#undef OP
|
||||
static inline __device__ void atomMul(uint8_t *address, uint8_t val) {
|
||||
AtomicMulIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(int8_t *address, int8_t val) {
|
||||
AtomicMulIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(int16_t *address, int16_t val) {
|
||||
AtomicMulIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(int32_t *address, int32_t val) {
|
||||
AtomicMulIntegerImpl<int32_t, sizeof(int32_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(int64_t *address, int64_t val) {
|
||||
AtomicMulIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(float *address, float val) {
|
||||
AtomicMulDecimalImpl<float, sizeof(float)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(at::Half *address, at::Half val) {
|
||||
AtomicMulDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(double *address, double val) {
|
||||
AtomicMulDecimalImpl<double, sizeof(double)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(at::BFloat16 *address, at::BFloat16 val) {
|
||||
AtomicMulDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
|
||||
}
|
||||
|
||||
#define OP(X, Y) Y / X
|
||||
ATOMIC(Div)
|
||||
#undef OP
|
||||
static inline __device__ void atomDiv(uint8_t *address, uint8_t val) {
|
||||
AtomicDivIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(int8_t *address, int8_t val) {
|
||||
AtomicDivIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(int16_t *address, int16_t val) {
|
||||
AtomicDivIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(int32_t *address, int32_t val) {
|
||||
AtomicDivIntegerImpl<int32_t, sizeof(int32_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(int64_t *address, int64_t val) {
|
||||
AtomicDivIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(at::Half *address, at::Half val) {
|
||||
AtomicDivDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(float *address, float val) {
|
||||
AtomicDivDecimalImpl<float, sizeof(float)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(double *address, double val) {
|
||||
AtomicDivDecimalImpl<double, sizeof(double)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(at::BFloat16 *address, at::BFloat16 val) {
|
||||
AtomicDivDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
|
||||
}
|
||||
|
||||
#define OP(X, Y) max(Y, X)
|
||||
ATOMIC(Max)
|
||||
#undef OP
|
||||
static inline __device__ void atomMax(uint8_t *address, uint8_t val) {
|
||||
AtomicMaxIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(int8_t *address, int8_t val) {
|
||||
AtomicMaxIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(int16_t *address, int16_t val) {
|
||||
AtomicMaxIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(int32_t *address, int32_t val) { atomicMax(address, val); }
|
||||
static inline __device__ void atomMax(int64_t *address, int64_t val) {
|
||||
AtomicMaxIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(at::Half *address, at::Half val) {
|
||||
AtomicMaxDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(float *address, float val) {
|
||||
AtomicMaxDecimalImpl<float, sizeof(float)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(double *address, double val) {
|
||||
AtomicMaxDecimalImpl<double, sizeof(double)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(at::BFloat16 *address, at::BFloat16 val) {
|
||||
AtomicMaxDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
|
||||
}
|
||||
|
||||
#define OP(X, Y) min(Y, X)
|
||||
ATOMIC(Min)
|
||||
#undef OP
|
||||
static inline __device__ void atomMin(uint8_t *address, uint8_t val) {
|
||||
AtomicMinIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(int8_t *address, int8_t val) {
|
||||
AtomicMinIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(int16_t *address, int16_t val) {
|
||||
AtomicMinIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(int32_t *address, int32_t val) { atomicMin(address, val); }
|
||||
static inline __device__ void atomMin(int64_t *address, int64_t val) {
|
||||
AtomicMinIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(at::Half *address, at::Half val) {
|
||||
AtomicMinDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(float *address, float val) {
|
||||
AtomicMinDecimalImpl<float, sizeof(float)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(double *address, double val) {
|
||||
AtomicMinDecimalImpl<double, sizeof(double)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(at::BFloat16 *address, at::BFloat16 val) {
|
||||
AtomicMinDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <limits>
|
||||
#include <map>
|
||||
|
||||
#include "atomics.cuh"
|
||||
|
||||
enum ReductionType { SUM, MEAN, MUL, DIV, MIN, MAX };
|
||||
|
||||
const std::map<std::string, ReductionType> reduce2REDUCE = {
|
||||
{"sum", SUM}, {"mean", MEAN}, {"mul", MUL}, {"div", DIV}, {"min", MIN}, {"max", MAX},
|
||||
};
|
||||
|
||||
#define AT_DISPATCH_REDUCTION_TYPES(reduce, ...) \
|
||||
[&] { \
|
||||
switch (reduce2REDUCE.at(reduce)) { \
|
||||
case SUM: { \
|
||||
static constexpr ReductionType REDUCE = SUM; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case MEAN: { \
|
||||
static constexpr ReductionType REDUCE = MEAN; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case MUL: { \
|
||||
static constexpr ReductionType REDUCE = MUL; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case DIV: { \
|
||||
static constexpr ReductionType REDUCE = DIV; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case MIN: { \
|
||||
static constexpr ReductionType REDUCE = MIN; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case MAX: { \
|
||||
static constexpr ReductionType REDUCE = MAX; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
} \
|
||||
}()
|
||||
|
||||
template <typename scalar_t, ReductionType REDUCE>
|
||||
struct Reducer {
|
||||
static inline __host__ __device__ scalar_t init() {
|
||||
if (REDUCE == MUL || REDUCE == DIV)
|
||||
return (scalar_t)1;
|
||||
else if (REDUCE == MIN)
|
||||
return std::numeric_limits<scalar_t>::max();
|
||||
else if (REDUCE == MAX)
|
||||
return std::numeric_limits<scalar_t>::lowest();
|
||||
else
|
||||
return (scalar_t)0;
|
||||
}
|
||||
|
||||
static inline __host__ __device__ void update(scalar_t *val, scalar_t new_val) {
|
||||
if (REDUCE == SUM || REDUCE == MEAN)
|
||||
*val = *val + new_val;
|
||||
else if (REDUCE == MUL)
|
||||
*val = *val * new_val;
|
||||
else if (REDUCE == DIV)
|
||||
*val = *val / new_val;
|
||||
else if ((REDUCE == MIN && new_val < *val) || (REDUCE == MAX && new_val > *val)) {
|
||||
*val = new_val;
|
||||
}
|
||||
}
|
||||
|
||||
static inline __host__ __device__ void update(scalar_t *val, scalar_t new_val, int64_t *arg, int64_t new_arg) {
|
||||
if (REDUCE == SUM || REDUCE == MEAN)
|
||||
*val = *val + new_val;
|
||||
else if (REDUCE == MUL)
|
||||
*val = *val * new_val;
|
||||
else if (REDUCE == DIV)
|
||||
*val = *val / new_val;
|
||||
else if ((REDUCE == MIN && new_val < *val) || (REDUCE == MAX && new_val > *val)) {
|
||||
*val = new_val;
|
||||
*arg = new_arg;
|
||||
}
|
||||
}
|
||||
|
||||
static inline __host__ __device__ void write(scalar_t *address, scalar_t val, int64_t *arg_address, int64_t arg,
|
||||
int count) {
|
||||
if (REDUCE == SUM || REDUCE == MUL || REDUCE == DIV)
|
||||
*address = val;
|
||||
else if (REDUCE == MEAN)
|
||||
*address = val / (scalar_t)(count > 0 ? count : 1);
|
||||
else if (REDUCE == MIN || REDUCE == MAX) {
|
||||
if (count > 0) {
|
||||
*address = val;
|
||||
*arg_address = arg;
|
||||
} else
|
||||
*address = (scalar_t)0;
|
||||
}
|
||||
}
|
||||
|
||||
static inline __device__ void atomic_write(scalar_t *address, scalar_t val) {
|
||||
if (REDUCE == SUM || REDUCE == MEAN)
|
||||
atomAdd(address, val);
|
||||
else if (REDUCE == MUL)
|
||||
atomMul(address, val);
|
||||
else if (REDUCE == DIV)
|
||||
atomDiv(address, val);
|
||||
else if (REDUCE == MIN)
|
||||
atomMin(address, val);
|
||||
else if (REDUCE == MAX)
|
||||
atomMax(address, val);
|
||||
}
|
||||
};
|
||||
@@ -0,0 +1,128 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include "scatter_cuda.h"
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/Dispatch.h>
|
||||
#include <ATen/cuda/detail/IndexUtils.cuh>
|
||||
#include <ATen/cuda/detail/TensorInfo.cuh>
|
||||
|
||||
#include "reducer.cuh"
|
||||
#include "utils.cuh"
|
||||
|
||||
#define THREADS 256
|
||||
#define BLOCKS(N) (N + THREADS - 1) / THREADS
|
||||
|
||||
template <typename scalar_t, ReductionType REDUCE>
|
||||
__global__ void scatter_kernel(const scalar_t *src_data, const at::cuda::detail::TensorInfo<int64_t, int> index_info,
|
||||
scalar_t *out_data, int E, int K, int N, int numel) {
|
||||
int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
int b = thread_idx / (E * K);
|
||||
int k = thread_idx % K;
|
||||
|
||||
if (thread_idx < numel) {
|
||||
int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(thread_idx, index_info);
|
||||
int64_t idx = index_info.data[offset];
|
||||
|
||||
Reducer<scalar_t, REDUCE>::atomic_write(out_data + b * N * K + idx * K + k, src_data[thread_idx]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void scatter_arg_kernel(const scalar_t *src_data,
|
||||
const at::cuda::detail::TensorInfo<int64_t, int> index_info,
|
||||
const scalar_t *out_data, int64_t *arg_out_data, int E, int K, int N, int numel) {
|
||||
int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
int b = thread_idx / (E * K);
|
||||
int e = (thread_idx / K) % E;
|
||||
int k = thread_idx % K;
|
||||
|
||||
if (thread_idx < numel) {
|
||||
int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(thread_idx, index_info);
|
||||
int64_t idx = index_info.data[offset];
|
||||
|
||||
if (src_data[thread_idx] == out_data[b * N * K + idx * K + k]) {
|
||||
arg_out_data[b * N * K + idx * K + k] = e;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_cuda(torch::Tensor src, torch::Tensor index,
|
||||
int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size,
|
||||
std::string reduce) {
|
||||
CHECK_CUDA(src);
|
||||
CHECK_CUDA(index);
|
||||
if (optional_out.has_value()) CHECK_CUDA(optional_out.value());
|
||||
c10::cuda::MaybeSetDevice(src.get_device());
|
||||
|
||||
CHECK_INPUT(src.dim() == index.dim());
|
||||
for (auto i = 0; i < index.dim() - 1; i++) CHECK_INPUT(src.size(i) >= index.size(i));
|
||||
|
||||
src = src.contiguous();
|
||||
|
||||
torch::Tensor out;
|
||||
if (optional_out.has_value()) {
|
||||
out = optional_out.value().contiguous();
|
||||
for (auto i = 0; i < out.dim(); i++)
|
||||
if (i != dim) CHECK_INPUT(src.size(i) == out.size(i));
|
||||
} else {
|
||||
auto sizes = src.sizes().vec();
|
||||
if (dim_size.has_value())
|
||||
sizes[dim] = dim_size.value();
|
||||
else if (index.numel() == 0)
|
||||
sizes[dim] = 0;
|
||||
else {
|
||||
sizes[dim] = 1 + index.max().cpu().data_ptr<int64_t>()[0];
|
||||
}
|
||||
out = torch::empty(sizes, src.options());
|
||||
}
|
||||
|
||||
std::optional<torch::Tensor> arg_out = std::nullopt;
|
||||
int64_t *arg_out_data = nullptr;
|
||||
if (reduce2REDUCE.at(reduce) == MIN || reduce2REDUCE.at(reduce) == MAX) {
|
||||
arg_out = torch::full_like(out, src.size(dim), index.options());
|
||||
arg_out_data = arg_out.value().data_ptr<int64_t>();
|
||||
}
|
||||
|
||||
if (src.numel() == 0) {
|
||||
if (!optional_out.has_value()) out.fill_(0);
|
||||
return std::make_tuple(out, arg_out);
|
||||
}
|
||||
|
||||
auto B = 1;
|
||||
for (auto i = 0; i < dim; i++) B *= src.size(i);
|
||||
auto E = src.size(dim);
|
||||
auto K = src.numel() / (B * E);
|
||||
auto N = out.size(dim);
|
||||
|
||||
auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
|
||||
auto src_data = src.data_ptr<scalar_t>();
|
||||
auto out_data = out.data_ptr<scalar_t>();
|
||||
|
||||
AT_DISPATCH_REDUCTION_TYPES(reduce, [&] {
|
||||
if (!optional_out.has_value()) out.fill_(Reducer<scalar_t, REDUCE>::init());
|
||||
|
||||
scatter_kernel<scalar_t, REDUCE>
|
||||
<<<BLOCKS(src.numel()), THREADS, 0, stream>>>(src_data, index_info, out_data, E, K, N, src.numel());
|
||||
|
||||
if (!optional_out.has_value() && (REDUCE == MIN || REDUCE == MAX))
|
||||
out.masked_fill_(out == Reducer<scalar_t, REDUCE>::init(), (scalar_t)0);
|
||||
|
||||
if (REDUCE == MIN || REDUCE == MAX)
|
||||
scatter_arg_kernel<scalar_t><<<BLOCKS(src.numel()), THREADS, 0, stream>>>(
|
||||
src_data, index_info, out_data, arg_out_data, E, K, N, src.numel());
|
||||
});
|
||||
});
|
||||
|
||||
return std::make_tuple(out, arg_out);
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_cuda(torch::Tensor src, torch::Tensor index,
|
||||
int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size,
|
||||
std::string reduce);
|
||||
@@ -0,0 +1,29 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#define CHECK_CUDA(x) AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")
|
||||
#define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch")
|
||||
|
||||
__device__ __inline__ at::Half __shfl_up_sync(const unsigned mask, const at::Half var, const unsigned int delta) {
|
||||
return __shfl_up_sync(mask, var.operator __half(), delta);
|
||||
}
|
||||
|
||||
__device__ __inline__ at::Half __shfl_down_sync(const unsigned mask, const at::Half var, const unsigned int delta) {
|
||||
return __shfl_down_sync(mask, var.operator __half(), delta);
|
||||
}
|
||||
|
||||
__device__ __inline__ at::Half __shfl_up(const at::Half var, const unsigned int delta) {
|
||||
return __shfl_up(var.operator __half(), delta);
|
||||
}
|
||||
|
||||
__device__ __inline__ at::Half __shfl_down(const at::Half var, const unsigned int delta) {
|
||||
return __shfl_down(var.operator __half(), delta);
|
||||
}
|
||||
|
||||
#define SHFL_UP_SYNC __shfl_up_sync
|
||||
#define SHFL_DOWN_SYNC __shfl_down_sync
|
||||
@@ -0,0 +1,238 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
#include "cuda/scatter_cuda.h"
|
||||
|
||||
inline std::vector<int64_t> list2vec(const c10::List<int64_t> list) {
|
||||
std::vector<int64_t> result;
|
||||
result.reserve(list.size());
|
||||
for (size_t i = 0; i < list.size(); i++) result.push_back(list[i]);
|
||||
return result;
|
||||
}
|
||||
|
||||
torch::Tensor broadcast(torch::Tensor src, torch::Tensor other, int64_t dim) {
|
||||
if (src.dim() == 1)
|
||||
for (auto i = 0; i < dim; i++) src = src.unsqueeze(0);
|
||||
for (auto i = src.dim(); i < other.dim(); i++) src = src.unsqueeze(-1);
|
||||
src = src.expand(other.sizes().vec());
|
||||
return src;
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_fw(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size,
|
||||
std::string reduce) {
|
||||
return scatter_cuda(src, index, dim, optional_out, dim_size, reduce);
|
||||
}
|
||||
|
||||
using torch::autograd::AutogradContext;
|
||||
using torch::autograd::Variable;
|
||||
using torch::autograd::variable_list;
|
||||
|
||||
class ScatterSum : public torch::autograd::Function<ScatterSum> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");
|
||||
auto out = std::get<0>(result);
|
||||
ctx->save_for_backward({index});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto index = saved[0];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
auto grad_in = torch::gather(grad_out, dim, index, false);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterMul : public torch::autograd::Function<ScatterMul> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "mul");
|
||||
auto out = std::get<0>(result);
|
||||
ctx->save_for_backward({src, index, out});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto src = saved[0];
|
||||
auto index = saved[1];
|
||||
auto out = saved[2];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
auto grad_in = torch::gather(grad_out * out, dim, index, false).div_(src);
|
||||
grad_in.masked_fill_(grad_in.isnan(), 0);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterMean : public torch::autograd::Function<ScatterMean> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
|
||||
auto old_index = index;
|
||||
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");
|
||||
auto out = std::get<0>(result);
|
||||
|
||||
auto ones = torch::ones(old_index.sizes(), src.options());
|
||||
result = scatter_fw(ones, old_index, old_index.dim() <= dim ? old_index.dim() - 1 : dim, std::nullopt,
|
||||
out.size(dim), "sum");
|
||||
auto count = std::get<0>(result);
|
||||
count.masked_fill_(count < 1, 1);
|
||||
count = broadcast(count, out, dim);
|
||||
if (out.is_floating_point())
|
||||
out.true_divide_(count);
|
||||
else
|
||||
out.div_(count, "floor");
|
||||
|
||||
ctx->save_for_backward({index, count});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto index = saved[0];
|
||||
auto count = saved[1];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
count = torch::gather(count, dim, index, false);
|
||||
auto grad_in = torch::gather(grad_out, dim, index, false);
|
||||
grad_in.true_divide_(count);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterMin : public torch::autograd::Function<ScatterMin> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "min");
|
||||
auto out = std::get<0>(result);
|
||||
auto arg_out = std::get<1>(result).value();
|
||||
ctx->save_for_backward({index, arg_out});
|
||||
ctx->mark_non_differentiable({arg_out});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out, arg_out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto index = saved[0];
|
||||
auto arg_out = saved[1];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
src_shape[dim] += 1;
|
||||
auto grad_in = torch::zeros(src_shape, grad_out.options());
|
||||
grad_in.scatter_(dim, arg_out, grad_out);
|
||||
grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterMax : public torch::autograd::Function<ScatterMax> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "max");
|
||||
auto out = std::get<0>(result);
|
||||
auto arg_out = std::get<1>(result).value();
|
||||
ctx->save_for_backward({index, arg_out});
|
||||
ctx->mark_non_differentiable({arg_out});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out, arg_out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto index = saved[0];
|
||||
auto arg_out = saved[1];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
src_shape[dim] += 1;
|
||||
auto grad_in = torch::zeros(src_shape, grad_out.options());
|
||||
grad_in.scatter_(dim, arg_out, grad_out);
|
||||
grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
torch::Tensor scatter_sum(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
|
||||
return ScatterSum::apply(src, index, dim, optional_out, dim_size)[0];
|
||||
}
|
||||
|
||||
torch::Tensor scatter_mul(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
|
||||
return ScatterMul::apply(src, index, dim, optional_out, dim_size)[0];
|
||||
}
|
||||
|
||||
torch::Tensor scatter_mean(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
|
||||
return ScatterMean::apply(src, index, dim, optional_out, dim_size)[0];
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> scatter_min(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size) {
|
||||
auto result = ScatterMin::apply(src, index, dim, optional_out, dim_size);
|
||||
return std::make_tuple(result[0], result[1]);
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> scatter_max(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size) {
|
||||
auto result = ScatterMax::apply(src, index, dim, optional_out, dim_size);
|
||||
return std::make_tuple(result[0], result[1]);
|
||||
}
|
||||
|
||||
void pybind_scatter_ext(py::module &m) {
|
||||
m.def("scatter_sum", &scatter_sum, "Scatter sum operation");
|
||||
m.def("scatter_mul", &scatter_mul, "Scatter multiplication operation");
|
||||
m.def("scatter_mean", &scatter_mean, "Scatter mean operation");
|
||||
m.def("scatter_min", &scatter_min, "Scatter min operation");
|
||||
m.def("scatter_max", &scatter_max, "Scatter max operation");
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,37 @@
|
||||
/**
|
||||
* This file includes code originally from the DROID-SLAM repository:
|
||||
* https://github.com/princeton-vl/DROID-SLAM
|
||||
* Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
namespace slam_ext {
|
||||
|
||||
torch::Tensor depth_filter_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics, torch::Tensor ix,
|
||||
torch::Tensor thresh);
|
||||
|
||||
torch::Tensor frame_distance_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics, torch::Tensor pi,
|
||||
torch::Tensor pj, torch::Tensor qi, torch::Tensor qj, torch::Tensor di,
|
||||
const float beta);
|
||||
|
||||
std::vector<torch::Tensor> projmap_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics,
|
||||
torch::Tensor ii, torch::Tensor jj);
|
||||
|
||||
torch::Tensor iproj_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics);
|
||||
|
||||
std::vector<torch::Tensor> ba_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics,
|
||||
torch::Tensor disps_sens, torch::Tensor targets, torch::Tensor weights,
|
||||
torch::Tensor eta, torch::Tensor ii, torch::Tensor jj, const int t0, const int t1,
|
||||
const int iterations, const float lm, const float ep, const bool motion_only);
|
||||
|
||||
} // namespace slam_ext
|
||||
|
||||
void pybind_slam_ext(py::module &m) {
|
||||
m.def("ba", &slam_ext::ba_cuda, "bundle adjustment");
|
||||
m.def("frame_distance", &slam_ext::frame_distance_cuda, "frame_distance");
|
||||
m.def("projmap", &slam_ext::projmap_cuda, "projmap");
|
||||
m.def("depth_filter", &slam_ext::depth_filter_cuda, "depth_filter");
|
||||
m.def("iproj", &slam_ext::iproj_cuda, "back projection");
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,137 @@
|
||||
/***********************************************************************
|
||||
* The code in this file is modified from the original code in the
|
||||
* https://github.com/jlblancoc/nanoflann repository.
|
||||
*
|
||||
* Software License Agreement (BSD License)
|
||||
*
|
||||
* Copyright 2008-2009 Marius Muja (mariusm@cs.ubc.ca). All rights reserved.
|
||||
* Copyright 2008-2009 David G. Lowe (lowe@cs.ubc.ca). All rights reserved.
|
||||
* Copyright 2011-2025 Jose Luis Blanco (joseluisblancoc@gmail.com).
|
||||
* All rights reserved.
|
||||
*
|
||||
* THE BSD LICENSE
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions
|
||||
* are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright
|
||||
* notice, this list of conditions and the following disclaimer.
|
||||
* 2. Redistributions in binary form must reproduce the above copyright
|
||||
* notice, this list of conditions and the following disclaimer in the
|
||||
* documentation and/or other materials provided with the distribution.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
|
||||
* OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
|
||||
* IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
|
||||
* INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
|
||||
* NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
|
||||
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
|
||||
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
|
||||
* THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*************************************************************************/
|
||||
|
||||
#ifndef FLANN_CUDA_KD_TREE_BUILDER_H_
|
||||
#define FLANN_CUDA_KD_TREE_BUILDER_H_
|
||||
|
||||
#include <thrust/host_vector.h>
|
||||
#include <thrust/device_vector.h>
|
||||
#include <thrust/sort.h>
|
||||
#include <thrust/partition.h>
|
||||
#include <thrust/unique.h>
|
||||
#include <thrust/scan.h>
|
||||
#include "math_util.h"
|
||||
#include <cstdlib>
|
||||
#include <map>
|
||||
|
||||
namespace tinyflann {
|
||||
|
||||
// Distance types.
|
||||
struct CudaL1;
|
||||
struct CudaL2;
|
||||
|
||||
// Parameters
|
||||
struct KDTreeCuda3dIndexParams {
|
||||
int leaf_max_size = 64;
|
||||
};
|
||||
|
||||
struct SearchParams {
|
||||
SearchParams(int checks_ = 32, float eps_ = 0.0, bool sorted_ = true)
|
||||
: checks(checks_), eps(eps_), sorted(sorted_) {
|
||||
max_neighbors = -1;
|
||||
use_heap = true;
|
||||
}
|
||||
|
||||
// how many leafs to visit when searching for neighbours (-1 for unlimited)
|
||||
int checks;
|
||||
// search for eps-approximate neighbours (default: 0)
|
||||
float eps;
|
||||
// only for radius search, require neighbours sorted by distance (default: true)
|
||||
bool sorted;
|
||||
// maximum number of neighbors radius search should return (-1 for unlimited)
|
||||
int max_neighbors;
|
||||
// use a heap to manage the result set (default: FLANN_Undefined)
|
||||
bool use_heap;
|
||||
};
|
||||
|
||||
template <typename Distance>
|
||||
class KDTreeCuda3dIndex {
|
||||
public:
|
||||
int visited_leafs;
|
||||
KDTreeCuda3dIndex(const float* input_data, size_t input_count,
|
||||
const KDTreeCuda3dIndexParams& params = KDTreeCuda3dIndexParams())
|
||||
: dataset_(input_data), leaf_count_(0), visited_leafs(0), node_count_(0), current_node_count_(0) {
|
||||
size_ = input_count;
|
||||
leaf_max_size_ = params.leaf_max_size;
|
||||
gpu_helper_ = 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Standard destructor
|
||||
*/
|
||||
~KDTreeCuda3dIndex() { clearGpuBuffers(); }
|
||||
|
||||
/**
|
||||
* Builds the index
|
||||
*/
|
||||
void buildIndex() {
|
||||
leaf_count_ = 0;
|
||||
node_count_ = 0;
|
||||
uploadTreeToGpu();
|
||||
}
|
||||
|
||||
/**
|
||||
* queries: (N, p) flattened float cuda array where only first 3 elements are used.
|
||||
* n_query: N
|
||||
* n_query_stride: p
|
||||
* indices: (N, knn) int cuda array
|
||||
* dists: (N, knn) float cuda array
|
||||
*/
|
||||
void knnSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, size_t knn,
|
||||
const SearchParams& params = SearchParams()) const;
|
||||
int radiusSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, float radius,
|
||||
const SearchParams& params = SearchParams()) const;
|
||||
|
||||
private:
|
||||
void uploadTreeToGpu();
|
||||
void clearGpuBuffers();
|
||||
|
||||
private:
|
||||
struct GpuHelper;
|
||||
GpuHelper* gpu_helper_;
|
||||
|
||||
const float* dataset_;
|
||||
int leaf_max_size_;
|
||||
int leaf_count_;
|
||||
int node_count_;
|
||||
//! used by convertTreeToGpuFormat
|
||||
int current_node_count_;
|
||||
size_t size_;
|
||||
|
||||
}; // class KDTreeCuda3dIndex
|
||||
|
||||
} // namespace tinyflann
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,66 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
|
||||
#include "cuda_kdtree.cuh"
|
||||
|
||||
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_IS_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, #x " must be a float tensor")
|
||||
|
||||
std::vector<torch::Tensor> nearestNeighbours(torch::Tensor query, torch::Tensor tree, int knn) {
|
||||
CHECK_CUDA(query);
|
||||
CHECK_IS_FLOAT(query);
|
||||
CHECK_CUDA(tree);
|
||||
CHECK_IS_FLOAT(tree);
|
||||
TORCH_CHECK(tree.size(0) >= knn, "knn is too small compared to the size of point cloud!");
|
||||
|
||||
torch::Tensor strided_tree = tree;
|
||||
torch::Tensor strided_query = query;
|
||||
torch::Device device = tree.device();
|
||||
long n_queries = query.size(0);
|
||||
long n_tree = tree.size(0);
|
||||
|
||||
if (strided_tree.stride(0) != 4) {
|
||||
strided_tree = torch::zeros({n_tree, 4}, torch::dtype(torch::kFloat32).device(device));
|
||||
strided_tree.index_put_(
|
||||
{torch::indexing::Ellipsis, torch::indexing::Slice(torch::indexing::None, tree.size(1))}, tree);
|
||||
}
|
||||
|
||||
if (strided_query.stride(0) != 3) {
|
||||
strided_query = torch::zeros({n_queries, 3}, torch::dtype(torch::kFloat32).device(device));
|
||||
strided_query.index_put_(
|
||||
{torch::indexing::Ellipsis, torch::indexing::Slice(torch::indexing::None, query.size(1))}, query);
|
||||
}
|
||||
|
||||
auto *knn_index = new tinyflann::KDTreeCuda3dIndex<tinyflann::CudaL2>(strided_tree.data_ptr<float>(), n_tree);
|
||||
knn_index->buildIndex();
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
|
||||
torch::Tensor dist = torch::empty(n_queries * knn, torch::dtype(torch::kFloat32).device(device));
|
||||
torch::Tensor indices = torch::empty(n_queries * knn, torch::dtype(torch::kInt32).device(device));
|
||||
|
||||
knn_index->knnSearch(strided_query.data_ptr<float>(), n_queries, 3, indices.data_ptr<int>(), dist.data_ptr<float>(),
|
||||
knn);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
|
||||
delete knn_index;
|
||||
|
||||
return {dist.reshape({n_queries, knn}), indices.reshape({n_queries, knn})};
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,50 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
|
||||
#define _CPU_AND_GPU_CODE_ __host__ __device__ // for CUDA device code
|
||||
#else
|
||||
#define _CPU_AND_GPU_CODE_
|
||||
#endif
|
||||
|
||||
#if defined(__CUDACC__)
|
||||
#define _CPU_AND_GPU_CODE_TEMPLATE_ __host__ __device__ // for CUDA device code
|
||||
#else
|
||||
#define _CPU_AND_GPU_CODE_TEMPLATE_
|
||||
#endif
|
||||
|
||||
#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
|
||||
#define _CPU_AND_GPU_CONSTANT_ __constant__ // for CUDA device code
|
||||
#else
|
||||
#define _CPU_AND_GPU_CONSTANT_ const
|
||||
#endif
|
||||
|
||||
#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
|
||||
#include <THC/THCAtomics.cuh>
|
||||
|
||||
template <typename ScalarT>
|
||||
static inline __device__ void atomicAdd(ScalarT* addr, ScalarT value) {
|
||||
gpuAtomicAddNoReturn(addr, value);
|
||||
}
|
||||
#else
|
||||
template <typename ScalarT>
|
||||
static inline void atomicAdd(ScalarT* addr, ScalarT value) {
|
||||
*addr += value;
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,25 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
std::vector<torch::Tensor> nearestNeighbours(torch::Tensor query, torch::Tensor tree, int knn);
|
||||
|
||||
void pybind_utils_ext(py::module &m) {
|
||||
m.def("nearest_neighbours", &nearestNeighbours, "KNN computation (CUDA) distance is squared L2.", py::arg("query"),
|
||||
py::arg("tree"), py::arg("knn"));
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
name: vipe
|
||||
channels:
|
||||
- nvidia/label/cuda-12.8.0
|
||||
- conda-forge
|
||||
- nodefaults
|
||||
dependencies:
|
||||
- python=3.10
|
||||
- pip
|
||||
- cuda-nvcc
|
||||
- eigen
|
||||
- zlib
|
||||
- libcusparse-dev
|
||||
- libcublas-dev
|
||||
- libcusolver-dev
|
||||
@@ -0,0 +1,23 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cu128
|
||||
torch==2.7.0+cu128
|
||||
torchvision
|
||||
|
||||
click
|
||||
numpy
|
||||
einops
|
||||
omegaconf
|
||||
tqdm
|
||||
rerun-sdk
|
||||
imageio[ffmpeg]
|
||||
opencv-python
|
||||
matplotlib
|
||||
python-pycg
|
||||
timm
|
||||
transformers
|
||||
gdown
|
||||
OpenEXR<3.3.0
|
||||
kornia
|
||||
hydra-core
|
||||
ray
|
||||
ninja
|
||||
viser
|
||||
@@ -0,0 +1,375 @@
|
||||
# This file was autogenerated by uv via the following command:
|
||||
# uv pip compile --emit-index-url --prerelease=allow -o envs/requirements.txt envs/requirements.in
|
||||
--index-url https://pypi.org/simple
|
||||
--extra-index-url https://download.pytorch.org/whl/cu128
|
||||
|
||||
aiosignal==1.3.1
|
||||
# via ray
|
||||
antlr4-python3-runtime==4.9.3
|
||||
# via
|
||||
# hydra-core
|
||||
# omegaconf
|
||||
attrs==24.2.0
|
||||
# via
|
||||
# jsonschema
|
||||
# referencing
|
||||
# rerun-sdk
|
||||
beautifulsoup4==4.13.3
|
||||
# via gdown
|
||||
calmsize==0.1.3
|
||||
# via python-pycg
|
||||
certifi==2022.12.7
|
||||
# via requests
|
||||
charset-normalizer==2.1.1
|
||||
# via
|
||||
# requests
|
||||
# trimesh
|
||||
click==8.1.8
|
||||
# via ray
|
||||
colorlog==6.9.0
|
||||
# via trimesh
|
||||
contourpy==1.2.1
|
||||
# via matplotlib
|
||||
cycler==0.12.1
|
||||
# via matplotlib
|
||||
einops==0.8.1
|
||||
# via -r envs/requirements.in
|
||||
embreex==2.17.7.post6
|
||||
# via trimesh
|
||||
filelock==3.13.1
|
||||
# via
|
||||
# gdown
|
||||
# huggingface-hub
|
||||
# ray
|
||||
# torch
|
||||
# transformers
|
||||
fonttools==4.53.0
|
||||
# via matplotlib
|
||||
frozenlist==1.4.1
|
||||
# via
|
||||
# aiosignal
|
||||
# ray
|
||||
fsspec==2024.6.1
|
||||
# via
|
||||
# huggingface-hub
|
||||
# torch
|
||||
gdown==5.2.0
|
||||
# via -r envs/requirements.in
|
||||
h11==0.16.0
|
||||
# via httpx
|
||||
httpx==1.0.dev1
|
||||
# via trimesh
|
||||
huggingface-hub==0.25.1
|
||||
# via
|
||||
# timm
|
||||
# tokenizers
|
||||
# transformers
|
||||
hydra-core==1.3.2
|
||||
# via -r envs/requirements.in
|
||||
idna==3.4
|
||||
# via requests
|
||||
imageio==2.37.0
|
||||
# via
|
||||
# -r envs/requirements.in
|
||||
# python-pycg
|
||||
# scikit-image
|
||||
# viser
|
||||
imageio-ffmpeg==0.6.0
|
||||
# via imageio
|
||||
jinja2==3.1.4
|
||||
# via torch
|
||||
jsonschema==4.23.0
|
||||
# via
|
||||
# ray
|
||||
# trimesh
|
||||
jsonschema-specifications==2024.10.1
|
||||
# via jsonschema
|
||||
kiwisolver==1.4.5
|
||||
# via matplotlib
|
||||
kornia==0.8.0
|
||||
# via -r envs/requirements.in
|
||||
kornia-rs==0.1.8
|
||||
# via kornia
|
||||
lazy-loader==0.4
|
||||
# via scikit-image
|
||||
lxml==6.0.0
|
||||
# via
|
||||
# trimesh
|
||||
# yourdfpy
|
||||
manifold3d==3.2.0
|
||||
# via trimesh
|
||||
mapbox-earcut==1.0.3
|
||||
# via trimesh
|
||||
markdown-it-py==3.0.0
|
||||
# via rich
|
||||
markupsafe==2.1.5
|
||||
# via jinja2
|
||||
matplotlib==3.9.0
|
||||
# via
|
||||
# -r envs/requirements.in
|
||||
# python-pycg
|
||||
mdurl==0.1.2
|
||||
# via markdown-it-py
|
||||
mpmath==1.3.0
|
||||
# via sympy
|
||||
msgpack==1.1.0
|
||||
# via ray
|
||||
msgspec==0.19.0
|
||||
# via viser
|
||||
networkx==3.3
|
||||
# via
|
||||
# scikit-image
|
||||
# torch
|
||||
# trimesh
|
||||
ninja==1.11.1.4
|
||||
# via -r envs/requirements.in
|
||||
nodeenv==1.9.1
|
||||
# via viser
|
||||
numpy==2.1.2
|
||||
# via
|
||||
# -r envs/requirements.in
|
||||
# contourpy
|
||||
# embreex
|
||||
# imageio
|
||||
# manifold3d
|
||||
# mapbox-earcut
|
||||
# matplotlib
|
||||
# opencv-python
|
||||
# plyfile
|
||||
# pyarrow
|
||||
# pycollada
|
||||
# pyquaternion
|
||||
# python-pycg
|
||||
# rerun-sdk
|
||||
# scikit-image
|
||||
# scipy
|
||||
# shapely
|
||||
# tifffile
|
||||
# torchvision
|
||||
# transformers
|
||||
# trimesh
|
||||
# vhacdx
|
||||
# viser
|
||||
# yourdfpy
|
||||
nvidia-cublas-cu12==12.8.3.14
|
||||
# via
|
||||
# nvidia-cudnn-cu12
|
||||
# nvidia-cusolver-cu12
|
||||
# torch
|
||||
nvidia-cuda-cupti-cu12==12.8.57
|
||||
# via torch
|
||||
nvidia-cuda-nvrtc-cu12==12.8.61
|
||||
# via torch
|
||||
nvidia-cuda-runtime-cu12==12.8.57
|
||||
# via torch
|
||||
nvidia-cudnn-cu12==9.7.1.26
|
||||
# via torch
|
||||
nvidia-cufft-cu12==11.3.3.41
|
||||
# via torch
|
||||
nvidia-cufile-cu12==1.13.0.11
|
||||
# via torch
|
||||
nvidia-curand-cu12==10.3.9.55
|
||||
# via torch
|
||||
nvidia-cusolver-cu12==11.7.2.55
|
||||
# via torch
|
||||
nvidia-cusparse-cu12==12.5.7.53
|
||||
# via
|
||||
# nvidia-cusolver-cu12
|
||||
# torch
|
||||
nvidia-cusparselt-cu12==0.6.3
|
||||
# via torch
|
||||
nvidia-ml-py==12.570.86
|
||||
# via pynvml
|
||||
nvidia-nccl-cu12==2.26.2
|
||||
# via torch
|
||||
nvidia-nvjitlink-cu12==12.8.61
|
||||
# via
|
||||
# nvidia-cufft-cu12
|
||||
# nvidia-cusolver-cu12
|
||||
# nvidia-cusparse-cu12
|
||||
# torch
|
||||
nvidia-nvtx-cu12==12.8.55
|
||||
# via torch
|
||||
omegaconf==2.3.0
|
||||
# via
|
||||
# -r envs/requirements.in
|
||||
# hydra-core
|
||||
# python-pycg
|
||||
opencv-python==4.11.0.86
|
||||
# via
|
||||
# -r envs/requirements.in
|
||||
# viser
|
||||
openexr==3.2.4
|
||||
# via -r envs/requirements.in
|
||||
packaging==24.1
|
||||
# via
|
||||
# huggingface-hub
|
||||
# hydra-core
|
||||
# kornia
|
||||
# lazy-loader
|
||||
# matplotlib
|
||||
# ray
|
||||
# scikit-image
|
||||
# transformers
|
||||
pillow==11.0.0
|
||||
# via
|
||||
# imageio
|
||||
# matplotlib
|
||||
# python-pycg
|
||||
# rerun-sdk
|
||||
# scikit-image
|
||||
# torchvision
|
||||
# trimesh
|
||||
plyfile==1.1
|
||||
# via python-pycg
|
||||
protobuf==6.30.0rc1
|
||||
# via ray
|
||||
psutil==6.0.0
|
||||
# via
|
||||
# imageio
|
||||
# viser
|
||||
pyarrow==17.0.0
|
||||
# via rerun-sdk
|
||||
pycollada==0.9.2
|
||||
# via trimesh
|
||||
pygments==2.19.1
|
||||
# via rich
|
||||
pynvml==12.0.0
|
||||
# via python-pycg
|
||||
pyparsing==3.1.2
|
||||
# via matplotlib
|
||||
pyquaternion==0.9.9
|
||||
# via python-pycg
|
||||
pysocks==1.7.1
|
||||
# via requests
|
||||
python-dateutil==2.9.0
|
||||
# via
|
||||
# matplotlib
|
||||
# pycollada
|
||||
python-pycg==0.10.1
|
||||
# via -r envs/requirements.in
|
||||
pyyaml==6.0.2
|
||||
# via
|
||||
# huggingface-hub
|
||||
# omegaconf
|
||||
# python-pycg
|
||||
# ray
|
||||
# timm
|
||||
# transformers
|
||||
ray==2.42.1
|
||||
# via -r envs/requirements.in
|
||||
referencing==0.36.2
|
||||
# via
|
||||
# jsonschema
|
||||
# jsonschema-specifications
|
||||
regex==2024.9.11
|
||||
# via transformers
|
||||
requests==2.28.1
|
||||
# via
|
||||
# gdown
|
||||
# huggingface-hub
|
||||
# ray
|
||||
# transformers
|
||||
# viser
|
||||
rerun-sdk==0.22.0
|
||||
# via -r envs/requirements.in
|
||||
rich==13.9.4
|
||||
# via
|
||||
# python-pycg
|
||||
# viser
|
||||
rpds-py==0.22.3
|
||||
# via
|
||||
# jsonschema
|
||||
# referencing
|
||||
rtree==1.4.0
|
||||
# via trimesh
|
||||
safetensors==0.4.5
|
||||
# via
|
||||
# timm
|
||||
# transformers
|
||||
scikit-image==0.25.2
|
||||
# via viser
|
||||
scipy==1.15.1
|
||||
# via
|
||||
# python-pycg
|
||||
# scikit-image
|
||||
# trimesh
|
||||
# viser
|
||||
screeninfo==0.8.1
|
||||
# via python-pycg
|
||||
setuptools==70.2.0
|
||||
# via
|
||||
# calmsize
|
||||
# trimesh
|
||||
# triton
|
||||
shapely==2.1.1
|
||||
# via trimesh
|
||||
six==1.16.0
|
||||
# via
|
||||
# python-dateutil
|
||||
# yourdfpy
|
||||
soupsieve==2.6
|
||||
# via beautifulsoup4
|
||||
svg-path==7.0
|
||||
# via trimesh
|
||||
sympy==1.13.3
|
||||
# via torch
|
||||
tifffile==2025.5.10
|
||||
# via scikit-image
|
||||
timm==1.0.14
|
||||
# via -r envs/requirements.in
|
||||
tokenizers==0.21.0
|
||||
# via transformers
|
||||
torch==2.7.0+cu128
|
||||
# via
|
||||
# -r envs/requirements.in
|
||||
# kornia
|
||||
# timm
|
||||
# torchvision
|
||||
torchvision==0.22.0+cu128
|
||||
# via
|
||||
# -r envs/requirements.in
|
||||
# timm
|
||||
tqdm==4.66.5
|
||||
# via
|
||||
# -r envs/requirements.in
|
||||
# gdown
|
||||
# huggingface-hub
|
||||
# python-pycg
|
||||
# transformers
|
||||
# viser
|
||||
transformers==4.48.3
|
||||
# via -r envs/requirements.in
|
||||
trimesh==4.6.2
|
||||
# via
|
||||
# python-pycg
|
||||
# viser
|
||||
# yourdfpy
|
||||
triton==3.3.0
|
||||
# via torch
|
||||
truststore==0.10.0
|
||||
# via httpx
|
||||
typing-extensions==4.12.2
|
||||
# via
|
||||
# beautifulsoup4
|
||||
# huggingface-hub
|
||||
# referencing
|
||||
# rerun-sdk
|
||||
# rich
|
||||
# torch
|
||||
urllib3==1.26.13
|
||||
# via requests
|
||||
usd-core==25.2.post1
|
||||
# via python-pycg
|
||||
vhacdx==0.0.8.post2
|
||||
# via trimesh
|
||||
viser==1.0.1
|
||||
# via -r envs/requirements.in
|
||||
websockets==15.0.1
|
||||
# via viser
|
||||
xatlas==0.0.11
|
||||
# via trimesh
|
||||
xxhash==3.5.0
|
||||
# via trimesh
|
||||
yourdfpy==0.0.58
|
||||
# via viser
|
||||
@@ -0,0 +1,41 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=40.8.0", "wheel", "torch"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py311"
|
||||
line-length = 120
|
||||
|
||||
[tool.ruff.lint.isort]
|
||||
lines-between-types = 1
|
||||
lines-after-imports = 2
|
||||
|
||||
[project]
|
||||
name = "vipe"
|
||||
description = "NVIDIA Video Pose Engine"
|
||||
authors = [
|
||||
{name = "The ViPE Authors", email = "jiahuih@nvidia.com"},
|
||||
]
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dynamic = ["version"]
|
||||
dependencies = [
|
||||
"numpy",
|
||||
"einops",
|
||||
"omegaconf",
|
||||
"tqdm",
|
||||
"rerun-sdk",
|
||||
"imageio[ffmpeg]",
|
||||
"opencv-python",
|
||||
"python-pycg",
|
||||
"timm",
|
||||
"transformers",
|
||||
"gdown",
|
||||
"OpenEXR",
|
||||
"kornia",
|
||||
"viser",
|
||||
"click",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
vipe = "vipe.cli.main:main"
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
import hydra
|
||||
from omegaconf import DictConfig
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="configs", config_name="default")
|
||||
def run(args: DictConfig) -> None:
|
||||
from vipe.streams.base import StreamList
|
||||
|
||||
# Gather all video streams
|
||||
stream_list = StreamList.make(args.streams)
|
||||
|
||||
from vipe.pipeline import make_pipeline
|
||||
from vipe.utils.logging import configure_logging
|
||||
|
||||
# Process each video stream
|
||||
logger = configure_logging()
|
||||
for stream_idx in range(len(stream_list)):
|
||||
video_stream = stream_list[stream_idx]
|
||||
logger.info(
|
||||
f"Processing {video_stream.name()} ({stream_idx + 1} / {len(stream_list)})"
|
||||
)
|
||||
pipeline = make_pipeline(args.pipeline)
|
||||
pipeline.run(video_stream)
|
||||
logger.info(f"Finished processing {video_stream.name()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run()
|
||||
@@ -0,0 +1,313 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Tuple
|
||||
|
||||
import cv2
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
from vipe.slam.interface import SLAMMap
|
||||
from vipe.utils.cameras import CameraType
|
||||
from vipe.utils.depth import reliable_depth_mask_range
|
||||
from vipe.utils.io import (
|
||||
ArtifactPath,
|
||||
read_depth_artifacts,
|
||||
read_intrinsics_artifacts,
|
||||
read_pose_artifacts,
|
||||
read_rgb_artifacts,
|
||||
)
|
||||
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def quaternion_from_matrix(matrix: np.ndarray) -> np.ndarray:
|
||||
"""Convert rotation matrix to quaternion (w, x, y, z)."""
|
||||
rotation = Rotation.from_matrix(matrix[:3, :3])
|
||||
quat_xyzw = rotation.as_quat() # Returns [x, y, z, w]
|
||||
return np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]]) # Convert to [w, x, y, z]
|
||||
|
||||
|
||||
def matrix_to_colmap_pose(c2w_matrix: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Convert camera-to-world matrix to COLMAP format.
|
||||
COLMAP uses world-to-camera transformation.
|
||||
"""
|
||||
w2c = np.linalg.inv(c2w_matrix)
|
||||
quaternion = quaternion_from_matrix(w2c)
|
||||
translation = w2c[:3, 3]
|
||||
return quaternion, translation
|
||||
|
||||
|
||||
def write_cameras_txt(output_dir: Path, artifact: ArtifactPath, frame_width: int, frame_height: int):
|
||||
"""Write COLMAP cameras.txt file."""
|
||||
cameras_file = output_dir / "cameras.txt"
|
||||
|
||||
_, intrinsics, camera_types = read_intrinsics_artifacts(artifact.intrinsics_path)
|
||||
|
||||
# Use first frame's intrinsics (assuming constant intrinsics)
|
||||
assert camera_types[0] == CameraType.PINHOLE, "Only PINHOLE camera type is supported"
|
||||
fx, fy, cx, cy = intrinsics[0].cpu().numpy()
|
||||
|
||||
with open(cameras_file, "w") as f:
|
||||
f.write("# Camera list with one line of data per camera:\n")
|
||||
f.write("# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n")
|
||||
f.write("# Number of cameras: 1\n")
|
||||
|
||||
fx, fy, cx, cy = intrinsics[0]
|
||||
|
||||
# COLMAP camera format: CAMERA_ID MODEL WIDTH HEIGHT fx fy cx cy
|
||||
f.write(f"1 PINHOLE {frame_width} {frame_height} {fx:.6f} {fy:.6f} {cx:.6f} {cy:.6f}\n")
|
||||
|
||||
logger.info(f"Written cameras.txt with intrinsics: fx={fx:.2f}, fy={fy:.2f}, cx={cx:.2f}, cy={cy:.2f}")
|
||||
|
||||
|
||||
def write_images_txt(output_dir: Path, artifact: ArtifactPath):
|
||||
"""Write COLMAP images.txt file."""
|
||||
images_file = output_dir / "images.txt"
|
||||
|
||||
# Load pose data
|
||||
pose_data = np.load(artifact.pose_path)
|
||||
poses = pose_data["data"] # Shape: (N, 4, 4)
|
||||
indices = pose_data["inds"] # Frame indices
|
||||
|
||||
with open(images_file, "w") as f:
|
||||
f.write("# Image list with two lines of data per image:\n")
|
||||
f.write("# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n")
|
||||
f.write("# POINTS2D[] as (X, Y, POINT3D_ID)\n")
|
||||
f.write(f"# Number of images: {len(poses)}\n")
|
||||
|
||||
for i, (pose_matrix, frame_idx) in enumerate(zip(poses, indices)):
|
||||
# Convert pose to COLMAP format
|
||||
quaternion, translation = matrix_to_colmap_pose(pose_matrix)
|
||||
qw, qx, qy, qz = quaternion
|
||||
tx, ty, tz = translation
|
||||
|
||||
# Image filename
|
||||
image_name = f"images/frame_{frame_idx:06d}.jpg"
|
||||
|
||||
# Write image line
|
||||
f.write(f"{i + 1} {qw:.9f} {qx:.9f} {qy:.9f} {qz:.9f} {tx:.9f} {ty:.9f} {tz:.9f} 1 {image_name}\n")
|
||||
# Empty points2D line (no 2D-3D correspondences)
|
||||
f.write("\n")
|
||||
|
||||
logger.info(f"Written images.txt with {len(poses)} images")
|
||||
|
||||
|
||||
def write_points3d_txt_from_slam_map(output_dir: Path, artifact: ArtifactPath):
|
||||
"""Write points3D.txt from SLAM map (placeholder implementation)."""
|
||||
assert artifact.slam_map_path.exists(), "SLAM map not found, please refer to README.md for more details."
|
||||
|
||||
slam_map = SLAMMap.load(artifact.slam_map_path, device=torch.device("cpu"))
|
||||
|
||||
points3d_file = output_dir / "points3D.txt"
|
||||
with open(points3d_file, "w") as f:
|
||||
f.write("# 3D point list with one line of data per point:\n")
|
||||
f.write("# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n")
|
||||
f.write(f"# Number of points: {slam_map.dense_disp_xyz.shape[0]}\n")
|
||||
|
||||
point_id = 1
|
||||
for keyframe_idx, frame_idx in enumerate(slam_map.dense_disp_frame_inds):
|
||||
xyz, rgb = slam_map.get_dense_disp_pcd(keyframe_idx)
|
||||
xyz = xyz.cpu().numpy()
|
||||
rgb = rgb.cpu().numpy()
|
||||
|
||||
for xyz, rgb in zip(xyz, rgb):
|
||||
x, y, z = xyz
|
||||
r, g, b = (rgb * 255).astype(np.uint8)
|
||||
f.write(f"{point_id} {x:.6f} {y:.6f} {z:.6f} {r} {g} {b} 0.0 {frame_idx} {point_id} 0 0 0 0\n")
|
||||
point_id += 1
|
||||
|
||||
|
||||
def write_points3d_txt_from_depth(
|
||||
output_dir: Path, artifact: ArtifactPath, depth_step: int, spatial_subsample: int = 4
|
||||
):
|
||||
"""Write empty COLMAP points3D.txt file."""
|
||||
_, pose_data = read_pose_artifacts(artifact.pose_path)
|
||||
_, intrinsics, camera_types = read_intrinsics_artifacts(artifact.intrinsics_path)
|
||||
camera_type = camera_types[0]
|
||||
points3d_file = output_dir / "points3D.txt"
|
||||
|
||||
image_dir = output_dir / "images"
|
||||
images = sorted(list(image_dir.glob("*.jpg")))
|
||||
# Collect all 3D points first
|
||||
all_points = []
|
||||
point_id = 1
|
||||
|
||||
rays: np.ndarray | None = None
|
||||
|
||||
for idx, (_, depth) in enumerate(read_depth_artifacts(artifact.depth_path)):
|
||||
if idx % 30 == 0:
|
||||
logger.info(f"Processed {idx} depth maps")
|
||||
|
||||
if idx % depth_step != 0:
|
||||
continue
|
||||
|
||||
rgb = cv2.cvtColor(cv2.imread(str(images[idx]), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
|
||||
frame_height, frame_width = rgb.shape[:2]
|
||||
rgb = rgb[::spatial_subsample, ::spatial_subsample]
|
||||
|
||||
if rays is None:
|
||||
camera_model = camera_type.build_camera_model(intrinsics[idx])
|
||||
disp_v, disp_u = torch.meshgrid(
|
||||
torch.arange(frame_height).float()[::spatial_subsample],
|
||||
torch.arange(frame_width).float()[::spatial_subsample],
|
||||
indexing="ij",
|
||||
)
|
||||
if camera_type == CameraType.PANORAMA:
|
||||
disp_v = disp_v / (frame_height - 1)
|
||||
disp_u = disp_u / (frame_width - 1)
|
||||
disp = torch.ones_like(disp_v)
|
||||
pts, _, _ = camera_model.iproj_disp(disp, disp_u, disp_v)
|
||||
rays = pts[..., :3].numpy()
|
||||
if camera_type != CameraType.PANORAMA:
|
||||
rays /= rays[..., 2:3]
|
||||
|
||||
if depth is not None:
|
||||
pcd = rays * depth.numpy()[::spatial_subsample, ::spatial_subsample, None]
|
||||
depth_mask = reliable_depth_mask_range(depth)[::spatial_subsample, ::spatial_subsample].numpy()
|
||||
rgb, pcd = rgb[depth_mask], pcd[depth_mask]
|
||||
c2w_matrix = pose_data[idx].matrix().numpy()
|
||||
pcd = pcd @ c2w_matrix[:3, :3].T + c2w_matrix[:3, 3][None]
|
||||
|
||||
for pts_rgb, pts_xyz in zip(rgb, pcd):
|
||||
all_points.append(
|
||||
(
|
||||
point_id,
|
||||
pts_xyz[0],
|
||||
pts_xyz[1],
|
||||
pts_xyz[2],
|
||||
int(pts_rgb[0]),
|
||||
int(pts_rgb[1]),
|
||||
int(pts_rgb[2]),
|
||||
0.0,
|
||||
idx + 1,
|
||||
)
|
||||
)
|
||||
point_id += 1
|
||||
|
||||
# Write points to file
|
||||
with open(points3d_file, "w") as f:
|
||||
f.write("# 3D point list with one line of data per point:\n")
|
||||
f.write("# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n")
|
||||
f.write(f"# Number of points: {len(all_points)}\n")
|
||||
|
||||
for point_id, point_data in enumerate(all_points):
|
||||
point_id, x, y, z, r, g, b, error, image_id = point_data
|
||||
# The last 4 values are for visualization purposes.
|
||||
f.write(f"{point_id} {x:.6f} {y:.6f} {z:.6f} {r} {g} {b} {error:.6f} {image_id} {point_id} 0 0 0 0\n")
|
||||
|
||||
logger.info(f"Written points3D.txt with {len(all_points)} points")
|
||||
|
||||
|
||||
def extract_frames(artifact: ArtifactPath, output_dir: Path) -> Tuple[int, int]:
|
||||
"""Extract frames from video to individual image files."""
|
||||
video_path = artifact.rgb_path
|
||||
images_dir = output_dir / "images"
|
||||
images_dir.mkdir(exist_ok=True)
|
||||
|
||||
logger.info(f"Extracting frames from {video_path}")
|
||||
|
||||
for frame_idx, rgb in read_rgb_artifacts(video_path):
|
||||
frame_path = images_dir / f"frame_{frame_idx:06d}.jpg"
|
||||
frame_height, frame_width = rgb.shape[:2]
|
||||
imageio.imwrite(str(frame_path), (rgb.cpu().numpy() * 255).astype(np.uint8))
|
||||
if frame_idx % 30 == 0:
|
||||
logger.info(f"Extracted {frame_idx} frames")
|
||||
|
||||
logger.info(f"Extracted {frame_idx} frames to {images_dir}")
|
||||
|
||||
return frame_width, frame_height
|
||||
|
||||
|
||||
def convert_vipe_to_colmap(artifact: ArtifactPath, output_path: Path, depth_step: int, use_slam_map: bool):
|
||||
"""Convert ViPE reconstruction results to COLMAP format."""
|
||||
|
||||
logger.info(
|
||||
f"Converting ViPE results from {artifact.base_path} ({artifact.artifact_name}) to COLMAP format at {output_path}"
|
||||
)
|
||||
|
||||
# Verify required files exist
|
||||
required_files = [artifact.rgb_path, artifact.pose_path, artifact.intrinsics_path, artifact.depth_path]
|
||||
for file_path in required_files:
|
||||
if not file_path.exists():
|
||||
raise FileNotFoundError(f"Required file not found: {file_path}")
|
||||
|
||||
# Create output directory
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Extract frames and get video dimensions
|
||||
frame_width, frame_height = extract_frames(artifact, output_path)
|
||||
|
||||
# Write COLMAP files
|
||||
write_cameras_txt(output_path, artifact, frame_width, frame_height)
|
||||
write_images_txt(output_path, artifact)
|
||||
if use_slam_map:
|
||||
write_points3d_txt_from_slam_map(output_path, artifact)
|
||||
else:
|
||||
write_points3d_txt_from_depth(output_path, artifact, depth_step)
|
||||
|
||||
logger.info("COLMAP conversion completed successfully!")
|
||||
logger.info(f"Output directory: {output_path}")
|
||||
logger.info("Files created:")
|
||||
logger.info(" - cameras.txt: Camera intrinsics")
|
||||
logger.info(" - images.txt: Camera poses")
|
||||
logger.info(" - points3D.txt: 3D points")
|
||||
logger.info(" - images/: Individual frame images")
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function for ViPE to COLMAP conversion script."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert ViPE reconstruction results to COLMAP format",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
parser.add_argument("vipe_path", type=Path, help="Path to ViPE results directory")
|
||||
parser.add_argument(
|
||||
"--sequence",
|
||||
"-s",
|
||||
type=str,
|
||||
help="Sequence name (if not provided will convert all sequences in the directory)",
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument("--use_slam_map", action="store_true", help="Use SLAM map to unproject depth maps")
|
||||
parser.add_argument("--depth_step", type=int, default=16, help="Step size for depth extraction (default: 16)")
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
"-o",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Output directory for COLMAP format (default: <vipe_path>_colmap)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.vipe_path.exists():
|
||||
print(f"Error: ViPE path '{args.vipe_path}' does not exist.")
|
||||
return 1
|
||||
|
||||
# Find artifacts
|
||||
artifacts = list(ArtifactPath.glob_artifacts(args.vipe_path, use_video=True))
|
||||
if args.sequence is not None:
|
||||
artifacts = [artifact for artifact in artifacts if artifact.artifact_name == args.sequence]
|
||||
|
||||
# Set default output path
|
||||
if args.output is None:
|
||||
args.output = args.vipe_path.parent / f"{args.vipe_path.name}_colmap"
|
||||
|
||||
for artifact in artifacts:
|
||||
convert_vipe_to_colmap(artifact, args.output / artifact.artifact_name, args.depth_step, args.use_slam_map)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main())
|
||||
+56
@@ -0,0 +1,56 @@
|
||||
import os
|
||||
import re
|
||||
|
||||
from setuptools import find_packages, setup
|
||||
|
||||
|
||||
try:
|
||||
import torch
|
||||
import torch.version
|
||||
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
|
||||
torch_version = torch.version.__version__.split(".")[:2]
|
||||
cuda_version = torch.version.cuda
|
||||
|
||||
# This will be e.g. "+pt23cu121"
|
||||
assert cuda_version is not None, "Pytorch CUDA is required for this installation."
|
||||
version_suffix = f"+pt{torch_version[0]}{torch_version[1]}cu{cuda_version.replace('.', '')}"
|
||||
|
||||
except ImportError:
|
||||
raise ValueError("Pytorch not found, please install it first.")
|
||||
|
||||
PACKAGE_NAME = "vipe"
|
||||
|
||||
# Avoid directly importing the package
|
||||
with open(f"{PACKAGE_NAME}/__init__.py", "r") as fh:
|
||||
__version__ = re.findall(r"__version__ = \"(.*?)\"", fh.read())[0]
|
||||
__version__ += version_suffix
|
||||
|
||||
coder_finder_path = f"{PACKAGE_NAME}/ext/specs.py"
|
||||
code_finder_namespace = {"__file__": coder_finder_path}
|
||||
with open(coder_finder_path, "r") as fh:
|
||||
exec(fh.read(), code_finder_namespace)
|
||||
get_sources = code_finder_namespace["get_sources"]
|
||||
get_cpp_flags = code_finder_namespace["get_cpp_flags"]
|
||||
get_cuda_flags = code_finder_namespace["get_cuda_flags"]
|
||||
|
||||
# Setup CUDA_HOME for conda environment for consistency
|
||||
if "CONDA_PREFIX" in os.environ:
|
||||
conda_nvcc_path = os.path.join(os.environ["CONDA_PREFIX"], "bin", "nvcc")
|
||||
if os.path.exists(conda_nvcc_path):
|
||||
os.environ["PYTORCH_NVCC"] = conda_nvcc_path
|
||||
|
||||
packages = find_packages()
|
||||
setup(
|
||||
packages=packages,
|
||||
version=__version__,
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
f"{PACKAGE_NAME}_ext",
|
||||
sources=get_sources(), # type: ignore
|
||||
extra_compile_args={"cxx": get_cpp_flags(), "nvcc": get_cuda_flags()}, # type: ignore
|
||||
)
|
||||
],
|
||||
cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
|
||||
)
|
||||
@@ -0,0 +1,33 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from vipe.pipeline import make_pipeline
|
||||
|
||||
|
||||
__version__ = "0.1.1"
|
||||
__version_info__ = (0, 1, 1)
|
||||
|
||||
if not OmegaConf.has_resolver("eq"):
|
||||
OmegaConf.register_new_resolver("eq", lambda a, b: a == b)
|
||||
if not OmegaConf.has_resolver("neq"):
|
||||
OmegaConf.register_new_resolver("neq", lambda a, b: a != b)
|
||||
|
||||
|
||||
def get_config_path() -> Path:
|
||||
return Path(__file__).parent.parent / "configs"
|
||||
@@ -0,0 +1,14 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
@@ -0,0 +1,84 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
import hydra
|
||||
|
||||
from vipe import get_config_path, make_pipeline
|
||||
from vipe.streams.base import ProcessedVideoStream
|
||||
from vipe.streams.raw_mp4_stream import RawMp4Stream
|
||||
from vipe.utils.logging import configure_logging
|
||||
from vipe.utils.viser import run_viser
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.argument("video", type=click.Path(exists=True, path_type=Path))
|
||||
@click.option(
|
||||
"--output",
|
||||
"-o",
|
||||
type=click.Path(path_type=Path),
|
||||
help="Output directory (default: current directory)",
|
||||
default=Path.cwd() / "vipe_results",
|
||||
)
|
||||
@click.option("--pipeline", "-p", default="default", help="Pipeline configuration to use (default: 'default')")
|
||||
@click.option("--visualize", "-v", is_flag=True, help="Enable visualization of intermediate results")
|
||||
def infer(video: Path, output: Path, pipeline: str, visualize: bool):
|
||||
"""Run inference on a video file."""
|
||||
|
||||
logger = configure_logging()
|
||||
|
||||
overrides = [f"pipeline={pipeline}", f"pipeline.output.path={output}", "pipeline.output.save_artifacts=true"]
|
||||
if visualize:
|
||||
overrides.append("pipeline.output.save_viz=true")
|
||||
overrides.append("pipeline.slam.visualize=true")
|
||||
else:
|
||||
overrides.append("pipeline.output.save_viz=false")
|
||||
|
||||
with hydra.initialize_config_dir(config_dir=str(get_config_path()), version_base=None):
|
||||
args = hydra.compose("default", overrides=overrides)
|
||||
|
||||
logger.info(f"Processing {video}...")
|
||||
vipe_pipeline = make_pipeline(args.pipeline)
|
||||
|
||||
# Some input videos can be malformed, so we need to cache the videos to obtain correct number of frames.
|
||||
video_stream = ProcessedVideoStream(RawMp4Stream(video), []).cache(desc="Reading video stream")
|
||||
|
||||
vipe_pipeline.run(video_stream)
|
||||
logger.info("Finished")
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.argument("data_path", type=click.Path(exists=True, path_type=Path), default=Path.cwd() / "vipe_results")
|
||||
@click.option("--port", "-p", default=20540, type=int, help="Port for the visualization server (default: 20540)")
|
||||
def visualize(data_path: Path, port: int):
|
||||
run_viser(data_path, port)
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.version_option()
|
||||
def main():
|
||||
"""NVIDIA Video Pose Engine (ViPE) CLI"""
|
||||
pass
|
||||
|
||||
|
||||
# Add subcommands
|
||||
main.add_command(infer)
|
||||
main.add_command(visualize)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,46 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
from vipe.ext.specs import get_cpp_flags, get_cuda_flags, get_sources
|
||||
|
||||
|
||||
try:
|
||||
import vipe_ext as _C
|
||||
|
||||
vipe_ext_not_found = False
|
||||
except ImportError:
|
||||
vipe_ext_not_found = True
|
||||
|
||||
if vipe_ext_not_found or os.environ.get("VIPE_EXT_JIT", "0") == "1":
|
||||
_C = load(
|
||||
name="vipe_ext_jit",
|
||||
sources=get_sources(),
|
||||
extra_cflags=get_cpp_flags(),
|
||||
extra_cuda_cflags=get_cuda_flags(),
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Reference to submodules
|
||||
droid_net_ext = _C.droid_net_ext
|
||||
grounding_dino_ext = _C.grounding_dino_ext
|
||||
utils_ext = _C.utils_ext
|
||||
slam_ext = _C.slam_ext
|
||||
scatter_ext = _C.scatter_ext
|
||||
lietorch_ext = _C.lietorch_ext
|
||||
corr_ext = _C.corr_ext
|
||||
@@ -0,0 +1,5 @@
|
||||
# This file includes code originally from the Pytorch Correlation repository:
|
||||
# https://github.com/ClementPinard/Pytorch-Correlation-extension
|
||||
# Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
from .spatial_correlation_sampler import SpatialCorrelationSampler, spatial_correlation_sample
|
||||
+126
@@ -0,0 +1,126 @@
|
||||
# This file includes code originally from the Pytorch Correlation repository:
|
||||
# https://github.com/ClementPinard/Pytorch-Correlation-extension
|
||||
# Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
from torch import nn
|
||||
from torch.autograd import Function
|
||||
from torch.autograd.function import once_differentiable
|
||||
from torch.nn.modules.utils import _pair
|
||||
|
||||
from vipe.ext import corr_ext as correlation
|
||||
|
||||
|
||||
def spatial_correlation_sample(
|
||||
input1, input2, kernel_size=1, patch_size=1, stride=1, padding=0, dilation=1, dilation_patch=1
|
||||
):
|
||||
"""Apply spatial correlation sampling on from input1 to input2,
|
||||
|
||||
Every parameter except input1 and input2 can be either single int
|
||||
or a pair of int. For more information about Spatial Correlation
|
||||
Sampling, see this page.
|
||||
https://lmb.informatik.uni-freiburg.de/Publications/2015/DFIB15/
|
||||
|
||||
Args:
|
||||
input1 : The first parameter.
|
||||
input2 : The second parameter.
|
||||
kernel_size : total size of your correlation kernel, in pixels
|
||||
patch_size : total size of your patch, determining how many
|
||||
different shifts will be applied
|
||||
stride : stride of the spatial sampler, will modify output
|
||||
height and width
|
||||
padding : padding applied to input1 and input2 before applying
|
||||
the correlation sampling, will modify output height and width
|
||||
dilation_patch : step for every shift in patch
|
||||
|
||||
Returns:
|
||||
Tensor: Result of correlation sampling
|
||||
|
||||
"""
|
||||
return SpatialCorrelationSamplerFunction.apply(
|
||||
input1, input2, kernel_size, patch_size, stride, padding, dilation, dilation_patch
|
||||
)
|
||||
|
||||
|
||||
class SpatialCorrelationSamplerFunction(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, input1, input2, kernel_size=1, patch_size=1, stride=1, padding=0, dilation=1, dilation_patch=1):
|
||||
ctx.save_for_backward(input1, input2)
|
||||
kH, kW = ctx.kernel_size = _pair(kernel_size)
|
||||
patchH, patchW = ctx.patch_size = _pair(patch_size)
|
||||
padH, padW = ctx.padding = _pair(padding)
|
||||
dilationH, dilationW = ctx.dilation = _pair(dilation)
|
||||
dilation_patchH, dilation_patchW = ctx.dilation_patch = _pair(dilation_patch)
|
||||
dH, dW = ctx.stride = _pair(stride)
|
||||
|
||||
output = correlation.forward(
|
||||
input1,
|
||||
input2,
|
||||
kH,
|
||||
kW,
|
||||
patchH,
|
||||
patchW,
|
||||
padH,
|
||||
padW,
|
||||
dilationH,
|
||||
dilationW,
|
||||
dilation_patchH,
|
||||
dilation_patchW,
|
||||
dH,
|
||||
dW,
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
@once_differentiable
|
||||
def backward(ctx, grad_output):
|
||||
input1, input2 = ctx.saved_tensors
|
||||
|
||||
kH, kW = ctx.kernel_size
|
||||
patchH, patchW = ctx.patch_size
|
||||
padH, padW = ctx.padding
|
||||
dilationH, dilationW = ctx.dilation
|
||||
dilation_patchH, dilation_patchW = ctx.dilation_patch
|
||||
dH, dW = ctx.stride
|
||||
|
||||
grad_input1, grad_input2 = correlation.backward(
|
||||
input1,
|
||||
input2,
|
||||
grad_output,
|
||||
kH,
|
||||
kW,
|
||||
patchH,
|
||||
patchW,
|
||||
padH,
|
||||
padW,
|
||||
dilationH,
|
||||
dilationW,
|
||||
dilation_patchH,
|
||||
dilation_patchW,
|
||||
dH,
|
||||
dW,
|
||||
)
|
||||
return grad_input1, grad_input2, None, None, None, None, None, None
|
||||
|
||||
|
||||
class SpatialCorrelationSampler(nn.Module):
|
||||
def __init__(self, kernel_size=1, patch_size=1, stride=1, padding=0, dilation=1, dilation_patch=1):
|
||||
super(SpatialCorrelationSampler, self).__init__()
|
||||
self.kernel_size = kernel_size
|
||||
self.patch_size = patch_size
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
self.dilation = dilation
|
||||
self.dilation_patch = dilation_patch
|
||||
|
||||
def forward(self, input1, input2):
|
||||
return SpatialCorrelationSamplerFunction.apply(
|
||||
input1,
|
||||
input2,
|
||||
self.kernel_size,
|
||||
self.patch_size,
|
||||
self.stride,
|
||||
self.padding,
|
||||
self.dilation,
|
||||
self.dilation_patch,
|
||||
)
|
||||
@@ -0,0 +1,6 @@
|
||||
# This file includes code originally from the lietorch repository:
|
||||
# https://github.com/princeton-vl/lietorch
|
||||
# Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
__all__ = ["groups"]
|
||||
from .groups import SE3, SO3, LieGroupParameter, RxSO3, Sim3, cat, stack
|
||||
@@ -0,0 +1,37 @@
|
||||
# This file includes code originally from the lietorch repository:
|
||||
# https://github.com/princeton-vl/lietorch
|
||||
# Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def check_broadcastable(x, y):
|
||||
assert len(x.shape) == len(y.shape)
|
||||
for n, m in zip(x.shape[:-1], y.shape[:-1]):
|
||||
assert n == m or n == 1 or m == 1
|
||||
|
||||
|
||||
def broadcast_inputs(x, y):
|
||||
"""Automatic broadcasting of missing dimensions"""
|
||||
if y is None:
|
||||
xs, xd = x.shape[:-1], x.shape[-1]
|
||||
return (x.view(-1, xd).contiguous(),), x.shape[:-1]
|
||||
|
||||
check_broadcastable(x, y)
|
||||
|
||||
xs, xd = x.shape[:-1], x.shape[-1]
|
||||
ys, yd = y.shape[:-1], y.shape[-1]
|
||||
out_shape = [max(n, m) for (n, m) in zip(xs, ys)]
|
||||
|
||||
if x.shape[:-1] == y.shape[-1]:
|
||||
x1 = x.view(-1, xd)
|
||||
y1 = y.view(-1, yd)
|
||||
|
||||
else:
|
||||
x_expand = [m if n == 1 else 1 for (n, m) in zip(xs, ys)]
|
||||
y_expand = [n if m == 1 else 1 for (n, m) in zip(xs, ys)]
|
||||
x1 = x.repeat(x_expand + [1]).reshape(-1, xd).contiguous()
|
||||
y1 = y.repeat(y_expand + [1]).reshape(-1, yd).contiguous()
|
||||
|
||||
return (x1, y1), tuple(out_shape)
|
||||
@@ -0,0 +1,124 @@
|
||||
# This file includes code originally from the lietorch repository:
|
||||
# https://github.com/princeton-vl/lietorch
|
||||
# Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vipe.ext import lietorch_ext as lietorch_backends
|
||||
|
||||
|
||||
class GroupOp(torch.autograd.Function):
|
||||
"""group operation base class"""
|
||||
|
||||
@classmethod
|
||||
def forward(cls, ctx, group_id, *inputs):
|
||||
ctx.group_id = group_id
|
||||
ctx.save_for_backward(*inputs)
|
||||
out = cls.forward_op(ctx.group_id, *inputs)
|
||||
return out
|
||||
|
||||
@classmethod
|
||||
def backward(cls, ctx, grad):
|
||||
error_str = "Backward operation not implemented for {}".format(cls)
|
||||
assert cls.backward_op is not None, error_str
|
||||
|
||||
inputs = ctx.saved_tensors
|
||||
grad = grad.contiguous()
|
||||
grad_inputs = cls.backward_op(ctx.group_id, grad, *inputs)
|
||||
return (None,) + tuple(grad_inputs)
|
||||
|
||||
|
||||
class Exp(GroupOp):
|
||||
"""exponential map"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.expm, lietorch_backends.expm_backward
|
||||
|
||||
|
||||
class Log(GroupOp):
|
||||
"""logarithm map"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.logm, lietorch_backends.logm_backward
|
||||
|
||||
|
||||
class Inv(GroupOp):
|
||||
"""group inverse"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.inv, lietorch_backends.inv_backward
|
||||
|
||||
|
||||
class Mul(GroupOp):
|
||||
"""group multiplication"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.mul, lietorch_backends.mul_backward
|
||||
|
||||
|
||||
class Adj(GroupOp):
|
||||
"""adjoint operator"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.adj, lietorch_backends.adj_backward
|
||||
|
||||
|
||||
class AdjT(GroupOp):
|
||||
"""adjoint operator"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.adjT, lietorch_backends.adjT_backward
|
||||
|
||||
|
||||
class Act3(GroupOp):
|
||||
"""action on point"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.act, lietorch_backends.act_backward
|
||||
|
||||
|
||||
class Act4(GroupOp):
|
||||
"""action on point"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.act4, lietorch_backends.act4_backward
|
||||
|
||||
|
||||
class Jinv(GroupOp):
|
||||
"""adjoint operator"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.Jinv, None
|
||||
|
||||
|
||||
class ToMatrix(GroupOp):
|
||||
"""convert to matrix representation"""
|
||||
|
||||
forward_op, backward_op = lietorch_backends.as_matrix, None
|
||||
|
||||
|
||||
### conversion operations to/from Euclidean embeddings ###
|
||||
|
||||
|
||||
class FromVec(torch.autograd.Function):
|
||||
"""convert vector into group object"""
|
||||
|
||||
@classmethod
|
||||
def forward(cls, ctx, group_id, *inputs):
|
||||
ctx.group_id = group_id
|
||||
ctx.save_for_backward(*inputs)
|
||||
return inputs[0]
|
||||
|
||||
@classmethod
|
||||
def backward(cls, ctx, grad):
|
||||
inputs = ctx.saved_tensors
|
||||
J = lietorch_backends.projector(ctx.group_id, *inputs)
|
||||
return None, torch.matmul(grad.unsqueeze(-2), torch.linalg.pinv(J)).squeeze(-2)
|
||||
|
||||
|
||||
class ToVec(torch.autograd.Function):
|
||||
"""convert group object to vector"""
|
||||
|
||||
@classmethod
|
||||
def forward(cls, ctx, group_id, *inputs):
|
||||
ctx.group_id = group_id
|
||||
ctx.save_for_backward(*inputs)
|
||||
return inputs[0]
|
||||
|
||||
@classmethod
|
||||
def backward(cls, ctx, grad):
|
||||
inputs = ctx.saved_tensors
|
||||
J = lietorch_backends.projector(ctx.group_id, *inputs)
|
||||
return None, torch.matmul(grad.unsqueeze(-2), J).squeeze(-2)
|
||||
@@ -0,0 +1,328 @@
|
||||
# This file includes code originally from the lietorch repository:
|
||||
# https://github.com/princeton-vl/lietorch
|
||||
# Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .broadcasting import broadcast_inputs
|
||||
|
||||
# group operations implemented in cuda
|
||||
from .group_ops import Act3, Act4, Adj, AdjT, Exp, FromVec, Inv, Jinv, Log, Mul, ToMatrix, ToVec
|
||||
|
||||
|
||||
class LieGroupParameter(torch.Tensor):
|
||||
"""Wrapper class for LieGroup"""
|
||||
|
||||
from torch._C import _disabled_torch_function_impl
|
||||
|
||||
__torch_function__ = _disabled_torch_function_impl
|
||||
|
||||
def __new__(cls, group, requires_grad=True):
|
||||
data = torch.zeros(group.tangent_shape, device=group.data.device, dtype=group.data.dtype, requires_grad=True)
|
||||
|
||||
return torch.Tensor._make_subclass(cls, data, requires_grad)
|
||||
|
||||
def __init__(self, group):
|
||||
self.group = group
|
||||
|
||||
def retr(self):
|
||||
return self.group.retr(self)
|
||||
|
||||
def log(self):
|
||||
return self.retr().log()
|
||||
|
||||
def inv(self):
|
||||
return self.retr().inv()
|
||||
|
||||
def adj(self, a):
|
||||
return self.retr().adj(a)
|
||||
|
||||
def __mul__(self, other):
|
||||
if isinstance(other, LieGroupParameter):
|
||||
return self.retr() * other.retr()
|
||||
else:
|
||||
return self.retr() * other
|
||||
|
||||
def add_(self, update, alpha):
|
||||
self.group = self.group.exp(alpha * update) * self.group
|
||||
|
||||
def __getitem__(self, index):
|
||||
return self.retr().__getitem__(index)
|
||||
|
||||
|
||||
class LieGroup:
|
||||
"""Base class for Lie Group"""
|
||||
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
def __repr__(self):
|
||||
return "{}: size={}, device={}, dtype={}".format(self.group_name, self.shape, self.device, self.dtype)
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self.data.shape[:-1]
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return self.data.device
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self.data.dtype
|
||||
|
||||
def vec(self):
|
||||
return self.apply_op(ToVec, self.data)
|
||||
|
||||
@property
|
||||
def tangent_shape(self):
|
||||
return self.data.shape[:-1] + (self.manifold_dim,)
|
||||
|
||||
@classmethod
|
||||
def Identity(cls, *batch_shape, **kwargs):
|
||||
"""Construct identity element with batch shape"""
|
||||
|
||||
if isinstance(batch_shape[0], tuple):
|
||||
batch_shape = batch_shape[0]
|
||||
|
||||
elif isinstance(batch_shape[0], list):
|
||||
batch_shape = tuple(batch_shape[0])
|
||||
|
||||
numel = np.prod(batch_shape)
|
||||
data = cls.id_elem.reshape(1, -1)
|
||||
|
||||
if "device" in kwargs:
|
||||
data = data.to(kwargs["device"])
|
||||
|
||||
if "dtype" in kwargs:
|
||||
data = data.type(kwargs["dtype"])
|
||||
|
||||
data = data.repeat(numel, 1)
|
||||
return cls(data).view(batch_shape)
|
||||
|
||||
@classmethod
|
||||
def IdentityLike(cls, G):
|
||||
return cls.Identity(G.shape, device=G.data.device, dtype=G.data.dtype)
|
||||
|
||||
@classmethod
|
||||
def InitFromVec(cls, data):
|
||||
return cls(cls.apply_op(FromVec, data))
|
||||
|
||||
@classmethod
|
||||
def Random(cls, *batch_shape, sigma=1.0, **kwargs):
|
||||
"""Construct random element with batch_shape by random sampling in tangent space"""
|
||||
|
||||
if isinstance(batch_shape[0], tuple):
|
||||
batch_shape = batch_shape[0]
|
||||
|
||||
elif isinstance(batch_shape[0], list):
|
||||
batch_shape = tuple(batch_shape[0])
|
||||
|
||||
tangent_shape = batch_shape + (cls.manifold_dim,)
|
||||
xi = torch.randn(tangent_shape, **kwargs)
|
||||
return cls.exp(sigma * xi)
|
||||
|
||||
@classmethod
|
||||
def apply_op(cls, op, x, y=None):
|
||||
"""Apply group operator"""
|
||||
inputs, out_shape = broadcast_inputs(x, y)
|
||||
|
||||
data = op.apply(cls.group_id, *inputs)
|
||||
return data.view(out_shape + (-1,))
|
||||
|
||||
@classmethod
|
||||
def exp(cls, x):
|
||||
"""exponential map: x -> X"""
|
||||
return cls(cls.apply_op(Exp, x))
|
||||
|
||||
def quaternion(self):
|
||||
"""extract quaternion"""
|
||||
return self.apply_op(Quat, self.data)
|
||||
|
||||
def log(self):
|
||||
"""logarithm map"""
|
||||
return self.apply_op(Log, self.data)
|
||||
|
||||
def inv(self):
|
||||
"""group inverse"""
|
||||
return self.__class__(self.apply_op(Inv, self.data))
|
||||
|
||||
def mul(self, other):
|
||||
"""group multiplication"""
|
||||
return self.__class__(self.apply_op(Mul, self.data, other.data))
|
||||
|
||||
def retr(self, a):
|
||||
"""retraction: Exp(a) * X"""
|
||||
dX = self.__class__.apply_op(Exp, a)
|
||||
return self.__class__(self.apply_op(Mul, dX, self.data))
|
||||
|
||||
def adj(self, a):
|
||||
"""adjoint operator: b = A(X) * a"""
|
||||
return self.apply_op(Adj, self.data, a)
|
||||
|
||||
def adjT(self, a):
|
||||
"""transposed adjoint operator: b = a * A(X)"""
|
||||
return self.apply_op(AdjT, self.data, a)
|
||||
|
||||
def Jinv(self, a):
|
||||
return self.apply_op(Jinv, self.data, a)
|
||||
|
||||
def act(self, p):
|
||||
"""action on a point cloud"""
|
||||
|
||||
# action on point
|
||||
if p.shape[-1] == 3:
|
||||
return self.apply_op(Act3, self.data, p)
|
||||
|
||||
# action on homogeneous point
|
||||
elif p.shape[-1] == 4:
|
||||
return self.apply_op(Act4, self.data, p)
|
||||
|
||||
def matrix(self):
|
||||
"""convert element to 4x4 matrix"""
|
||||
I = torch.eye(4, dtype=self.dtype, device=self.device)
|
||||
I = I.view([1] * (len(self.data.shape) - 1) + [4, 4])
|
||||
return self.__class__(self.data[..., None, :]).act(I).transpose(-1, -2)
|
||||
|
||||
def translation(self):
|
||||
"""extract translation component"""
|
||||
p = torch.as_tensor([0.0, 0.0, 0.0, 1.0], dtype=self.dtype, device=self.device)
|
||||
p = p.view(
|
||||
[1] * (len(self.data.shape) - 1)
|
||||
+ [
|
||||
4,
|
||||
]
|
||||
)
|
||||
return self.apply_op(Act4, self.data, p)
|
||||
|
||||
def detach(self):
|
||||
return self.__class__(self.data.detach())
|
||||
|
||||
def view(self, dims):
|
||||
data_reshaped = self.data.view(dims + (self.embedded_dim,))
|
||||
return self.__class__(data_reshaped)
|
||||
|
||||
def __mul__(self, other):
|
||||
# group multiplication
|
||||
if isinstance(other, LieGroup):
|
||||
return self.mul(other)
|
||||
|
||||
# action on point
|
||||
elif isinstance(other, torch.Tensor):
|
||||
return self.act(other)
|
||||
|
||||
def __getitem__(self, index):
|
||||
return self.__class__(self.data[index])
|
||||
|
||||
def __setitem__(self, index, item):
|
||||
self.data[index] = item.data
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
return self.__class__(self.data.to(*args, **kwargs))
|
||||
|
||||
def cpu(self):
|
||||
return self.__class__(self.data.cpu())
|
||||
|
||||
def cuda(self):
|
||||
return self.__class__(self.data.cuda())
|
||||
|
||||
def float(self, device):
|
||||
return self.__class__(self.data.float())
|
||||
|
||||
def double(self, device):
|
||||
return self.__class__(self.data.double())
|
||||
|
||||
def unbind(self, dim=0):
|
||||
return [self.__class__(x) for x in self.data.unbind(dim=dim)]
|
||||
|
||||
|
||||
class SO3(LieGroup):
|
||||
group_name = "SO3"
|
||||
group_id = 1
|
||||
manifold_dim = 3
|
||||
embedded_dim = 4
|
||||
|
||||
# unit quaternion
|
||||
id_elem = torch.as_tensor([0.0, 0.0, 0.0, 1.0])
|
||||
|
||||
def __init__(self, data):
|
||||
if isinstance(data, SE3):
|
||||
data = data.data[..., 3:7]
|
||||
|
||||
super(SO3, self).__init__(data)
|
||||
|
||||
|
||||
class RxSO3(LieGroup):
|
||||
group_name = "RxSO3"
|
||||
group_id = 2
|
||||
manifold_dim = 4
|
||||
embedded_dim = 5
|
||||
|
||||
# unit quaternion
|
||||
id_elem = torch.as_tensor([0.0, 0.0, 0.0, 1.0, 1.0])
|
||||
|
||||
def __init__(self, data):
|
||||
if isinstance(data, Sim3):
|
||||
data = data.data[..., 3:8]
|
||||
|
||||
super(RxSO3, self).__init__(data)
|
||||
|
||||
|
||||
class SE3(LieGroup):
|
||||
group_name = "SE3"
|
||||
group_id = 3
|
||||
manifold_dim = 6
|
||||
embedded_dim = 7
|
||||
|
||||
# translation, unit quaternion
|
||||
id_elem = torch.as_tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
|
||||
|
||||
def __init__(self, data):
|
||||
if isinstance(data, SO3):
|
||||
translation = torch.zeros_like(data.data[..., :3])
|
||||
data = torch.cat([translation, data.data], -1)
|
||||
|
||||
super(SE3, self).__init__(data)
|
||||
|
||||
def scale(self, s):
|
||||
t, q = self.data.split([3, 4], -1)
|
||||
t = t * s.unsqueeze(-1)
|
||||
return SE3(torch.cat([t, q], dim=-1))
|
||||
|
||||
|
||||
class Sim3(LieGroup):
|
||||
group_name = "Sim3"
|
||||
group_id = 4
|
||||
manifold_dim = 7
|
||||
embedded_dim = 8
|
||||
|
||||
# translation, unit quaternion, scale
|
||||
id_elem = torch.as_tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0])
|
||||
|
||||
def __init__(self, data):
|
||||
if isinstance(data, SO3):
|
||||
scale = torch.ones_like(SO3.data[..., :1])
|
||||
translation = torch.zeros_like(SO3.data[..., :3])
|
||||
data = torch.cat([translation, SO3.data, scale], -1)
|
||||
|
||||
elif isinstance(data, SE3):
|
||||
scale = torch.ones_like(data.data[..., :1])
|
||||
data = torch.cat([data.data, scale], -1)
|
||||
|
||||
elif isinstance(data, Sim3):
|
||||
data = data.data
|
||||
|
||||
super(Sim3, self).__init__(data)
|
||||
|
||||
|
||||
def cat(group_list, dim):
|
||||
"""Concatenate groups along dimension"""
|
||||
data = torch.cat([X.data for X in group_list], dim=dim)
|
||||
return group_list[0].__class__(data)
|
||||
|
||||
|
||||
def stack(group_list, dim):
|
||||
"""Concatenate groups along dimension"""
|
||||
data = torch.stack([X.data for X in group_list], dim=dim)
|
||||
return group_list[0].__class__(data)
|
||||
@@ -0,0 +1,205 @@
|
||||
# This file includes code originally from the Pytorch Scatter repository:
|
||||
# https://github.com/rusty1s/pytorch_scatter
|
||||
# Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from . import scatter_ext
|
||||
|
||||
|
||||
def broadcast(src: torch.Tensor, other: torch.Tensor, dim: int):
|
||||
if dim < 0:
|
||||
dim = other.dim() + dim
|
||||
if src.dim() == 1:
|
||||
for _ in range(0, dim):
|
||||
src = src.unsqueeze(0)
|
||||
for _ in range(src.dim(), other.dim()):
|
||||
src = src.unsqueeze(-1)
|
||||
src = src.expand(other.size())
|
||||
return src
|
||||
|
||||
|
||||
def scatter_sum(
|
||||
src: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
dim: int = -1,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
dim_size: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
index = broadcast(index, src, dim)
|
||||
if out is None:
|
||||
size = list(src.size())
|
||||
if dim_size is not None:
|
||||
size[dim] = dim_size
|
||||
elif index.numel() == 0:
|
||||
size[dim] = 0
|
||||
else:
|
||||
size[dim] = int(index.max()) + 1
|
||||
out = torch.zeros(size, dtype=src.dtype, device=src.device)
|
||||
return out.scatter_add_(dim, index, src)
|
||||
else:
|
||||
return out.scatter_add_(dim, index, src)
|
||||
|
||||
|
||||
def scatter_add(
|
||||
src: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
dim: int = -1,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
dim_size: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
return scatter_sum(src, index, dim, out, dim_size)
|
||||
|
||||
|
||||
def scatter_mul(
|
||||
src: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
dim: int = -1,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
dim_size: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
return scatter_ext.scatter_mul(src, index, dim, out, dim_size)
|
||||
|
||||
|
||||
def scatter_mean(
|
||||
src: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
dim: int = -1,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
dim_size: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
out = scatter_sum(src, index, dim, out, dim_size)
|
||||
dim_size = out.size(dim)
|
||||
|
||||
index_dim = dim
|
||||
if index_dim < 0:
|
||||
index_dim = index_dim + src.dim()
|
||||
if index.dim() <= index_dim:
|
||||
index_dim = index.dim() - 1
|
||||
|
||||
ones = torch.ones(index.size(), dtype=src.dtype, device=src.device)
|
||||
count = scatter_sum(ones, index, index_dim, None, dim_size)
|
||||
count[count < 1] = 1
|
||||
count = broadcast(count, out, dim)
|
||||
if out.is_floating_point():
|
||||
out.true_divide_(count)
|
||||
else:
|
||||
out.div_(count, rounding_mode="floor")
|
||||
return out
|
||||
|
||||
|
||||
def scatter_min(
|
||||
src: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
dim: int = -1,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
dim_size: Optional[int] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
return scatter_ext.scatter_min(src, index, dim, out, dim_size)
|
||||
|
||||
|
||||
def scatter_max(
|
||||
src: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
dim: int = -1,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
dim_size: Optional[int] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
return scatter_ext.scatter_max(src, index, dim, out, dim_size)
|
||||
|
||||
|
||||
def scatter(
|
||||
src: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
dim: int = -1,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
dim_size: Optional[int] = None,
|
||||
reduce: str = "sum",
|
||||
) -> torch.Tensor:
|
||||
r"""
|
||||
|
|
||||
|
||||
.. image:: https://raw.githubusercontent.com/rusty1s/pytorch_scatter/
|
||||
master/docs/source/_figures/add.svg?sanitize=true
|
||||
:align: center
|
||||
:width: 400px
|
||||
|
||||
|
|
||||
|
||||
Reduces all values from the :attr:`src` tensor into :attr:`out` at the
|
||||
indices specified in the :attr:`index` tensor along a given axis
|
||||
:attr:`dim`.
|
||||
For each value in :attr:`src`, its output index is specified by its index
|
||||
in :attr:`src` for dimensions outside of :attr:`dim` and by the
|
||||
corresponding value in :attr:`index` for dimension :attr:`dim`.
|
||||
The applied reduction is defined via the :attr:`reduce` argument.
|
||||
|
||||
Formally, if :attr:`src` and :attr:`index` are :math:`n`-dimensional
|
||||
tensors with size :math:`(x_0, ..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})`
|
||||
and :attr:`dim` = `i`, then :attr:`out` must be an :math:`n`-dimensional
|
||||
tensor with size :math:`(x_0, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})`.
|
||||
Moreover, the values of :attr:`index` must be between :math:`0` and
|
||||
:math:`y - 1`, although no specific ordering of indices is required.
|
||||
The :attr:`index` tensor supports broadcasting in case its dimensions do
|
||||
not match with :attr:`src`.
|
||||
|
||||
For one-dimensional tensors with :obj:`reduce="sum"`, the operation
|
||||
computes
|
||||
|
||||
.. math::
|
||||
\mathrm{out}_i = \mathrm{out}_i + \sum_j~\mathrm{src}_j
|
||||
|
||||
where :math:`\sum_j` is over :math:`j` such that
|
||||
:math:`\mathrm{index}_j = i`.
|
||||
|
||||
.. note::
|
||||
|
||||
This operation is implemented via atomic operations on the GPU and is
|
||||
therefore **non-deterministic** since the order of parallel operations
|
||||
to the same value is undetermined.
|
||||
For floating-point variables, this results in a source of variance in
|
||||
the result.
|
||||
|
||||
:param src: The source tensor.
|
||||
:param index: The indices of elements to scatter.
|
||||
:param dim: The axis along which to index. (default: :obj:`-1`)
|
||||
:param out: The destination tensor.
|
||||
:param dim_size: If :attr:`out` is not given, automatically create output
|
||||
with size :attr:`dim_size` at dimension :attr:`dim`.
|
||||
If :attr:`dim_size` is not given, a minimal sized output tensor
|
||||
according to :obj:`index.max() + 1` is returned.
|
||||
:param reduce: The reduce operation (:obj:`"sum"`, :obj:`"mul"`,
|
||||
:obj:`"mean"`, :obj:`"min"` or :obj:`"max"`). (default: :obj:`"sum"`)
|
||||
|
||||
:rtype: :class:`Tensor`
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from torch_scatter import scatter
|
||||
|
||||
src = torch.randn(10, 6, 64)
|
||||
index = torch.tensor([0, 1, 0, 1, 2, 1])
|
||||
|
||||
# Broadcasting in the first and last dim.
|
||||
out = scatter(src, index, dim=1, reduce="sum")
|
||||
|
||||
print(out.size())
|
||||
|
||||
.. code-block::
|
||||
|
||||
torch.Size([10, 3, 64])
|
||||
"""
|
||||
if reduce == "sum" or reduce == "add":
|
||||
return scatter_sum(src, index, dim, out, dim_size)
|
||||
if reduce == "mul":
|
||||
return scatter_mul(src, index, dim, out, dim_size)
|
||||
elif reduce == "mean":
|
||||
return scatter_mean(src, index, dim, out, dim_size)
|
||||
elif reduce == "min":
|
||||
return scatter_min(src, index, dim, out, dim_size)[0]
|
||||
elif reduce == "max":
|
||||
return scatter_max(src, index, dim, out, dim_size)[0]
|
||||
else:
|
||||
raise ValueError
|
||||
@@ -0,0 +1,39 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def get_sources() -> list[str]:
|
||||
csrc_path = Path(__file__).parent.parent.parent / "csrc"
|
||||
return [str(p) for p in csrc_path.glob("**/*") if p.suffix in [".cpp", ".cu"]]
|
||||
|
||||
|
||||
def _additional_include_flags() -> list[str]:
|
||||
if "CONDA_PREFIX" in os.environ:
|
||||
conda_include_path = os.path.join(os.environ["CONDA_PREFIX"], "include")
|
||||
if os.path.exists(conda_include_path):
|
||||
return ["-isystem", conda_include_path]
|
||||
return []
|
||||
|
||||
|
||||
def get_cpp_flags() -> list[str]:
|
||||
return ["-O3", "-DWITH_CUDA"] + _additional_include_flags()
|
||||
|
||||
|
||||
def get_cuda_flags() -> list[str]:
|
||||
return ["-O3", "-DWITH_CUDA", "--use_fast_math"] + _additional_include_flags()
|
||||
@@ -0,0 +1,49 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import Optional, Sequence
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def memory_efficient_attention(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
# (B, S, H, D) -> (B, H, S, D)
|
||||
query = query.transpose(1, 2)
|
||||
key = key.transpose(1, 2)
|
||||
value = value.transpose(1, 2)
|
||||
|
||||
out = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_bias)
|
||||
|
||||
return out.transpose(1, 2)
|
||||
|
||||
|
||||
def scaled_index_add(
|
||||
input: torch.Tensor, # [B, M, D]
|
||||
index: torch.Tensor, # [Bi] - int64
|
||||
source: torch.Tensor, # [Bi, M, D]
|
||||
scaling: Optional[torch.Tensor] = None, # [D]
|
||||
alpha: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
return torch.index_add(input, dim=0, source=scaling * source, index=index, alpha=alpha)
|
||||
|
||||
|
||||
def index_select_cat(sources: Sequence[torch.Tensor], indices: Sequence[torch.Tensor]) -> torch.Tensor:
|
||||
return torch.cat([s[i.long()].flatten() for s, i in zip(sources, indices)], dim=0)
|
||||
@@ -0,0 +1,80 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import copy
|
||||
import importlib
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Sequence
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from vipe.streams.base import MultiviewVideoList, VideoStream
|
||||
|
||||
|
||||
@dataclass(kw_only=True, slots=True)
|
||||
class AnnotationPipelineOutput:
|
||||
# Eager return of the payload values that comes from early stages of the pipeline.
|
||||
payload: Any | None = None
|
||||
output_streams: Sequence[VideoStream] | None = None
|
||||
|
||||
|
||||
class Pipeline(ABC):
|
||||
def __init__(self) -> None:
|
||||
self._return_payload = False
|
||||
self._return_output_streams = False
|
||||
|
||||
@property
|
||||
def return_payload(self) -> bool:
|
||||
return self._return_payload
|
||||
|
||||
@return_payload.setter
|
||||
def return_payload(self, value: bool) -> None:
|
||||
assert isinstance(value, bool), "return_payload must be a boolean"
|
||||
self._return_payload = value
|
||||
if value:
|
||||
self._return_output_streams = False
|
||||
|
||||
@property
|
||||
def return_output_streams(self) -> bool:
|
||||
return self._return_output_streams
|
||||
|
||||
@return_output_streams.setter
|
||||
def return_output_streams(self, value: bool) -> None:
|
||||
assert isinstance(value, bool), "return_output_streams must be a boolean"
|
||||
self._return_output_streams = value
|
||||
if value:
|
||||
self._return_payload = False
|
||||
|
||||
def should_filter(self, stream_name: str) -> bool:
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
def run(self, video_data: VideoStream | MultiviewVideoList) -> AnnotationPipelineOutput: ...
|
||||
|
||||
|
||||
def make_pipeline_cls(config: DictConfig) -> type[Pipeline]:
|
||||
module_path, class_name = config.instance.rsplit(".", 1)
|
||||
module = importlib.import_module(module_path)
|
||||
return getattr(module, class_name)
|
||||
|
||||
|
||||
def make_pipeline(config: DictConfig) -> Pipeline:
|
||||
config = copy.deepcopy(config)
|
||||
pipeline_cls = make_pipeline_cls(config)
|
||||
del config.instance
|
||||
return pipeline_cls(**config)
|
||||
@@ -0,0 +1,152 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import logging
|
||||
import pickle
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from vipe.slam.system import SLAMOutput, SLAMSystem
|
||||
from vipe.streams.base import (
|
||||
AssignAttributesProcessor,
|
||||
FrameAttribute,
|
||||
MultiviewVideoList,
|
||||
ProcessedVideoStream,
|
||||
StreamProcessor,
|
||||
VideoStream,
|
||||
)
|
||||
from vipe.utils import io
|
||||
from vipe.utils.cameras import CameraType
|
||||
from vipe.utils.visualization import save_projection_video
|
||||
|
||||
from . import AnnotationPipelineOutput, Pipeline
|
||||
from .processors import AdaptiveDepthProcessor, GeoCalibIntrinsicsProcessor, TrackAnythingProcessor
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DefaultAnnotationPipeline(Pipeline):
|
||||
def __init__(self, init: DictConfig, slam: DictConfig, post: DictConfig, output: DictConfig) -> None:
|
||||
super().__init__()
|
||||
self.init_cfg = init
|
||||
self.slam_cfg = slam
|
||||
self.post_cfg = post
|
||||
self.out_cfg = output
|
||||
self.out_path = Path(self.out_cfg.path)
|
||||
self.out_path.mkdir(exist_ok=True, parents=True)
|
||||
self.camera_type = CameraType(self.init_cfg.camera_type)
|
||||
|
||||
def _add_init_processors(self, video_stream: VideoStream) -> ProcessedVideoStream:
|
||||
init_processors: list[StreamProcessor] = []
|
||||
|
||||
# The assertions make sure that the attributes are not estimated previously.
|
||||
# Otherwise it will be overwritten by the processors.
|
||||
assert FrameAttribute.INTRINSICS not in video_stream.attributes()
|
||||
assert FrameAttribute.CAMERA_TYPE not in video_stream.attributes()
|
||||
assert FrameAttribute.METRIC_DEPTH not in video_stream.attributes()
|
||||
assert FrameAttribute.INSTANCE not in video_stream.attributes()
|
||||
|
||||
init_processors.append(GeoCalibIntrinsicsProcessor(video_stream, camera_type=self.camera_type))
|
||||
if self.init_cfg.instance is not None:
|
||||
init_processors.append(
|
||||
TrackAnythingProcessor(
|
||||
self.init_cfg.instance.phrases,
|
||||
add_sky=self.init_cfg.instance.add_sky,
|
||||
sam_run_gap=int(video_stream.fps() * self.init_cfg.instance.kf_gap_sec),
|
||||
)
|
||||
)
|
||||
return ProcessedVideoStream(video_stream, init_processors)
|
||||
|
||||
def _add_post_processors(
|
||||
self, view_idx: int, video_stream: VideoStream, slam_output: SLAMOutput
|
||||
) -> ProcessedVideoStream:
|
||||
post_processors: list[StreamProcessor] = [
|
||||
AssignAttributesProcessor(
|
||||
{
|
||||
FrameAttribute.POSE: slam_output.get_view_trajectory(view_idx), # type: ignore
|
||||
FrameAttribute.INTRINSICS: [slam_output.intrinsics[view_idx]] * len(video_stream),
|
||||
}
|
||||
)
|
||||
]
|
||||
if (depth_align_model := self.post_cfg.depth_align_model) is not None:
|
||||
post_processors.append(AdaptiveDepthProcessor(slam_output, view_idx, depth_align_model))
|
||||
return ProcessedVideoStream(video_stream, post_processors)
|
||||
|
||||
def run(self, video_data: VideoStream | MultiviewVideoList) -> AnnotationPipelineOutput:
|
||||
if isinstance(video_data, MultiviewVideoList):
|
||||
video_streams = [video_data[view_idx] for view_idx in range(len(video_data))]
|
||||
artifact_paths = [io.ArtifactPath(self.out_path, video_stream.name()) for video_stream in video_streams]
|
||||
slam_rig = video_data.rig()
|
||||
|
||||
else:
|
||||
assert isinstance(video_data, VideoStream)
|
||||
video_streams = [video_data]
|
||||
artifact_paths = [io.ArtifactPath(self.out_path, video_data.name())]
|
||||
slam_rig = None
|
||||
|
||||
annotate_output = AnnotationPipelineOutput()
|
||||
|
||||
if all([self.should_filter(video_stream.name()) for video_stream in video_streams]):
|
||||
logger.info(f"{video_data.name()} has been proccessed already, skip it!!")
|
||||
return annotate_output
|
||||
|
||||
slam_streams: list[VideoStream] = [
|
||||
self._add_init_processors(video_stream).cache("process", online=True) for video_stream in video_streams
|
||||
]
|
||||
|
||||
slam_pipeline = SLAMSystem(device=torch.device("cuda"), config=self.slam_cfg)
|
||||
slam_output = slam_pipeline.run(slam_streams, rig=slam_rig, camera_type=self.camera_type)
|
||||
|
||||
if self.return_payload:
|
||||
annotate_output.payload = slam_output
|
||||
return annotate_output
|
||||
|
||||
output_streams = [
|
||||
self._add_post_processors(view_idx, slam_stream, slam_output).cache("depth", online=True)
|
||||
for view_idx, slam_stream in enumerate(slam_streams)
|
||||
]
|
||||
|
||||
# Dumping artifacts for all views in the streams
|
||||
for output_stream, artifact_path in zip(output_streams, artifact_paths):
|
||||
artifact_path.meta_info_path.parent.mkdir(exist_ok=True, parents=True)
|
||||
if self.out_cfg.save_artifacts:
|
||||
logger.info(f"Saving artifacts to {artifact_path}")
|
||||
io.save_artifacts(artifact_path, output_stream)
|
||||
with artifact_path.meta_info_path.open("wb") as f:
|
||||
pickle.dump({"ba_residual": slam_output.ba_residual}, f)
|
||||
|
||||
if self.out_cfg.save_viz:
|
||||
save_projection_video(
|
||||
artifact_path.meta_vis_path,
|
||||
output_stream,
|
||||
slam_output,
|
||||
self.out_cfg.viz_downsample,
|
||||
self.out_cfg.viz_attributes,
|
||||
)
|
||||
|
||||
if self.out_cfg.save_slam_map and slam_output.slam_map is not None:
|
||||
logger.info(f"Saving SLAM map to {artifact_path.slam_map_path}")
|
||||
slam_output.slam_map.save(artifact_path.slam_map_path)
|
||||
|
||||
if self.return_output_streams:
|
||||
annotate_output.output_streams = output_streams
|
||||
|
||||
return annotate_output
|
||||
@@ -0,0 +1,301 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import logging
|
||||
|
||||
from typing import Iterator
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vipe.priors.depth import DepthEstimationInput, make_depth_model
|
||||
from vipe.priors.depth.alignment import align_inv_depth_to_depth
|
||||
from vipe.priors.depth.priorda import PriorDAModel
|
||||
from vipe.priors.depth.videodepthanything import VdieoDepthAnythingDepthModel
|
||||
from vipe.priors.geocalib import GeoCalib
|
||||
from vipe.priors.track_anything import TrackAnythingPipeline
|
||||
from vipe.slam.interface import SLAMOutput
|
||||
from vipe.streams.base import CachedVideoStream, FrameAttribute, StreamProcessor, VideoFrame, VideoStream
|
||||
from vipe.utils.cameras import CameraType
|
||||
from vipe.utils.logging import pbar
|
||||
from vipe.utils.misc import unpack_optional
|
||||
from vipe.utils.morph import erode
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IntrinsicEstimationProcessor(StreamProcessor):
|
||||
"""Override existing intrinsics with estimated intrinsics."""
|
||||
|
||||
def __init__(self, video_stream: VideoStream, gap_sec: float = 1.0) -> None:
|
||||
super().__init__()
|
||||
gap_frame = int(gap_sec * video_stream.fps())
|
||||
gap_frame = min(gap_frame, (len(video_stream) - 1) // 2)
|
||||
self.sample_frame_inds = [0, gap_frame, gap_frame * 2]
|
||||
self.fov_y = -1.0
|
||||
self.camera_type = CameraType.PINHOLE
|
||||
self.distortion: list[float] = []
|
||||
|
||||
def update_attributes(self, previous_attributes: set[FrameAttribute]) -> set[FrameAttribute]:
|
||||
return previous_attributes | {FrameAttribute.INTRINSICS}
|
||||
|
||||
def __call__(self, frame_idx: int, frame: VideoFrame) -> VideoFrame:
|
||||
assert self.fov_y > 0, "FOV not set"
|
||||
frame_height, frame_width = frame.size()
|
||||
fx = fy = frame_height / (2 * np.tan(self.fov_y / 2))
|
||||
frame.intrinsics = torch.as_tensor(
|
||||
[fx, fy, frame_width / 2, frame_height / 2] + self.distortion,
|
||||
).float()
|
||||
frame.camera_type = self.camera_type
|
||||
return frame
|
||||
|
||||
|
||||
class GeoCalibIntrinsicsProcessor(IntrinsicEstimationProcessor):
|
||||
def __init__(
|
||||
self,
|
||||
video_stream: VideoStream,
|
||||
gap_sec: float = 1.0,
|
||||
camera_type: CameraType = CameraType.PINHOLE,
|
||||
) -> None:
|
||||
super().__init__(video_stream, gap_sec)
|
||||
|
||||
is_pinhole = camera_type == CameraType.PINHOLE
|
||||
weights = "pinhole" if is_pinhole else "distorted"
|
||||
|
||||
model = GeoCalib(weights=weights).cuda()
|
||||
indexable_stream = CachedVideoStream(video_stream)
|
||||
|
||||
if is_pinhole:
|
||||
sample_frames = torch.stack([indexable_stream[i].rgb.moveaxis(-1, 0) for i in self.sample_frame_inds])
|
||||
res = model.calibrate(
|
||||
sample_frames,
|
||||
shared_intrinsics=True,
|
||||
)
|
||||
else:
|
||||
# Use first frame for calibration
|
||||
camera_model = {
|
||||
CameraType.PINHOLE: "pinhole",
|
||||
CameraType.MEI: "simple_mei",
|
||||
}[camera_type]
|
||||
res = model.calibrate(
|
||||
indexable_stream[self.sample_frame_inds[0]].rgb.moveaxis(-1, 0)[None],
|
||||
camera_model=camera_model,
|
||||
)
|
||||
|
||||
self.fov_y = res["camera"].vfov[0].item()
|
||||
self.camera_type = camera_type
|
||||
|
||||
if not is_pinhole:
|
||||
# Assign distortion parameter
|
||||
self.distortion = [res["camera"].dist[0, 0].item()]
|
||||
|
||||
|
||||
class TrackAnythingProcessor(StreamProcessor):
|
||||
"""
|
||||
A processor that tracks a mask caption in the video.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mask_phrases: list[str],
|
||||
add_sky: bool,
|
||||
sam_run_gap: int = 30,
|
||||
mask_expand: int = 5,
|
||||
) -> None:
|
||||
self.mask_phrases = mask_phrases
|
||||
self.sam_run_gap = sam_run_gap
|
||||
self.add_sky = add_sky
|
||||
|
||||
if self.add_sky:
|
||||
self.mask_phrases.append(VideoFrame.SKY_PROMPT)
|
||||
|
||||
self.tracker = TrackAnythingPipeline(self.mask_phrases, sam_points_per_side=50, sam_run_gap=self.sam_run_gap)
|
||||
self.mask_expand = mask_expand
|
||||
|
||||
def update_attributes(self, previous_attributes: set[FrameAttribute]) -> set[FrameAttribute]:
|
||||
return previous_attributes | {FrameAttribute.INSTANCE, FrameAttribute.MASK}
|
||||
|
||||
def __call__(self, frame_idx: int, frame: VideoFrame) -> VideoFrame:
|
||||
frame.instance, frame.instance_phrases = self.tracker.track(frame)
|
||||
self.last_track_frame = frame.raw_frame_idx
|
||||
|
||||
frame_instance_mask = frame.instance == 0
|
||||
if self.add_sky:
|
||||
# We won't mask out the sky.
|
||||
frame_instance_mask |= frame.sky_mask
|
||||
|
||||
frame.mask = erode(frame_instance_mask, self.mask_expand)
|
||||
return frame
|
||||
|
||||
|
||||
class AdaptiveDepthProcessor(StreamProcessor):
|
||||
"""
|
||||
Compute projection of the SLAM map onto the current frames.
|
||||
If it's well-distributed, then use the fast map-prompted video depth model.
|
||||
If not, then use the slow metric depth + video depth alignment model.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
slam_output: SLAMOutput,
|
||||
view_idx: int = 0,
|
||||
model: str = "adaptive_unidepth-l_svda",
|
||||
share_depth_model: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.slam_output = slam_output
|
||||
self.infill_target_pose = self.slam_output.get_view_trajectory(view_idx)
|
||||
assert view_idx == 0, "Adaptive depth processor only supports view_idx=0"
|
||||
assert not share_depth_model, "Adaptive depth processor does not support shared depth model"
|
||||
self.require_cache = True
|
||||
self.model = model
|
||||
|
||||
try:
|
||||
prefix, metric_model, video_model = model.split("_")
|
||||
assert video_model in ["svda", "vda"]
|
||||
self.video_depth_model = VdieoDepthAnythingDepthModel(model="vits" if video_model == "svda" else "vitl")
|
||||
|
||||
except ValueError:
|
||||
prefix, metric_model = model.split("_")
|
||||
video_model = None
|
||||
self.video_depth_model = None
|
||||
|
||||
assert prefix == "adaptive", "Model name should start with 'adaptive_'"
|
||||
|
||||
self.depth_model = make_depth_model(metric_model)
|
||||
self.prompt_model = PriorDAModel()
|
||||
self.update_momentum = 0.99
|
||||
|
||||
def __call__(self, frame_idx: int, frame: VideoFrame) -> VideoFrame:
|
||||
raise NotImplementedError("AdaptiveDepthProcessor should not be called directly.")
|
||||
|
||||
def update_attributes(self, previous_attributes: set[FrameAttribute]) -> set[FrameAttribute]:
|
||||
return previous_attributes | {FrameAttribute.METRIC_DEPTH}
|
||||
|
||||
def _compute_uv_score(self, depth: torch.Tensor, patch_count: int = 10) -> float:
|
||||
h_shape = depth.size(0) // patch_count
|
||||
w_shape = depth.size(1) // patch_count
|
||||
depth_crop = (depth > 0)[: h_shape * patch_count, : w_shape * patch_count]
|
||||
depth_crop = depth_crop.reshape(patch_count, h_shape, patch_count, w_shape)
|
||||
depth_exist = depth_crop.any(dim=(1, 3))
|
||||
return depth_exist.float().mean().item()
|
||||
|
||||
def _compute_video_da(self, frame_iterator: Iterator[VideoFrame]) -> tuple[torch.Tensor, list[VideoFrame]]:
|
||||
frame_list: list[np.ndarray] = []
|
||||
frame_data_list: list[VideoFrame] = []
|
||||
for frame in frame_iterator:
|
||||
frame_data_list.append(frame.cpu())
|
||||
frame_list.append(frame.rgb.cpu().numpy())
|
||||
|
||||
video_depth_result: torch.Tensor = unpack_optional(
|
||||
self.video_depth_model.estimate(DepthEstimationInput(video_frame_list=frame_list)).relative_inv_depth
|
||||
)
|
||||
return video_depth_result, frame_data_list
|
||||
|
||||
def update_iterator(self, previous_iterator: Iterator[VideoFrame]) -> Iterator[VideoFrame]:
|
||||
# Determine the percentage score of the SLAM map.
|
||||
|
||||
self.cache_scale_bias = None
|
||||
min_uv_score: float = 1.0
|
||||
|
||||
if self.video_depth_model is not None:
|
||||
video_depth_result, data_iterator = self._compute_video_da(previous_iterator)
|
||||
else:
|
||||
video_depth_result = None
|
||||
data_iterator = previous_iterator
|
||||
|
||||
for frame_idx, frame in pbar(enumerate(data_iterator), desc="Aligning depth"):
|
||||
# Convert back to GPU if not already.
|
||||
frame = frame.cuda()
|
||||
|
||||
# Compute the minimum UV score only once at the 0-th frame.
|
||||
if frame_idx == 0:
|
||||
for test_frame_idx in range(self.slam_output.trajectory.shape[0]):
|
||||
if test_frame_idx % 10 != 0:
|
||||
continue
|
||||
depth_infilled = self.slam_output.slam_map.project_map(
|
||||
test_frame_idx,
|
||||
0,
|
||||
frame.size(),
|
||||
unpack_optional(frame.intrinsics),
|
||||
self.infill_target_pose[test_frame_idx],
|
||||
unpack_optional(frame.camera_type),
|
||||
infill=False,
|
||||
)
|
||||
uv_score = self._compute_uv_score(depth_infilled)
|
||||
if uv_score < min_uv_score:
|
||||
min_uv_score = uv_score
|
||||
|
||||
logger.info(f"Minimum UV score: {min_uv_score:.4f}")
|
||||
|
||||
if min_uv_score < 0.3:
|
||||
focal_length = frame.intrinsics[0].item()
|
||||
prompt_result = self.depth_model.estimate(
|
||||
DepthEstimationInput(rgb=frame.rgb.float().cuda(), focal_length=focal_length)
|
||||
).metric_depth
|
||||
frame.information = f"uv={min_uv_score:.2f}(Metric)"
|
||||
else:
|
||||
depth_map = self.slam_output.slam_map.project_map(
|
||||
frame_idx,
|
||||
0,
|
||||
frame.size(),
|
||||
unpack_optional(frame.intrinsics),
|
||||
self.infill_target_pose[frame_idx],
|
||||
unpack_optional(frame.camera_type),
|
||||
infill=False,
|
||||
)
|
||||
if frame.mask is not None:
|
||||
depth_map = depth_map * frame.mask.float()
|
||||
prompt_result = self.prompt_model.estimate(
|
||||
DepthEstimationInput(
|
||||
rgb=frame.rgb.float().cuda(),
|
||||
prompt_metric_depth=depth_map,
|
||||
)
|
||||
).metric_depth
|
||||
frame.information = f"uv={min_uv_score:.2f}(SLAM)"
|
||||
|
||||
if video_depth_result is not None:
|
||||
video_depth_inv_depth = video_depth_result[frame_idx]
|
||||
|
||||
align_mask = video_depth_inv_depth > 1e-3
|
||||
if frame.mask is not None:
|
||||
align_mask = align_mask & frame.mask & (~frame.sky_mask)
|
||||
|
||||
try:
|
||||
_, scale, bias = align_inv_depth_to_depth(
|
||||
unpack_optional(video_depth_inv_depth),
|
||||
prompt_result,
|
||||
align_mask,
|
||||
)
|
||||
except RuntimeError:
|
||||
scale, bias = self.cache_scale_bias
|
||||
|
||||
# momentum update
|
||||
if self.cache_scale_bias is None:
|
||||
self.cache_scale_bias = (scale, bias)
|
||||
scale = self.cache_scale_bias[0] * self.update_momentum + scale * (1 - self.update_momentum)
|
||||
bias = self.cache_scale_bias[1] * self.update_momentum + bias * (1 - self.update_momentum)
|
||||
self.cache_scale_bias = (scale, bias)
|
||||
|
||||
video_inv_depth = video_depth_inv_depth * scale + bias
|
||||
video_inv_depth[video_inv_depth < 1e-3] = 1e-3
|
||||
frame.metric_depth = video_inv_depth.reciprocal()
|
||||
|
||||
else:
|
||||
frame.metric_depth = prompt_result
|
||||
|
||||
yield frame
|
||||
@@ -0,0 +1,36 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .base import DepthEstimationInput, DepthEstimationModel, DepthEstimationResult, DepthType
|
||||
|
||||
|
||||
def make_depth_model(model: str):
|
||||
if "-" not in model:
|
||||
model_name, model_sub = model, ""
|
||||
else:
|
||||
model_name, model_sub = model.split("-")
|
||||
|
||||
if model_name == "metric3d":
|
||||
from .metric3d import Metric3DDepthModel
|
||||
|
||||
return Metric3DDepthModel(version=2, model=model_sub)
|
||||
|
||||
elif model_name == "unidepth":
|
||||
from .unidepth import UniDepth2Model
|
||||
|
||||
return UniDepth2Model(type=model_sub)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown depth model: {model}")
|
||||
@@ -0,0 +1,125 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def align_inv_depth_to_depth(
|
||||
source_inv_depth: torch.Tensor,
|
||||
target_depth: torch.Tensor,
|
||||
target_mask: torch.Tensor | None = None,
|
||||
quantile_masking: bool = True,
|
||||
) -> tuple[torch.Tensor, float, float]:
|
||||
"""
|
||||
Apply affine transformation to align source inverse depth to target depth.
|
||||
|
||||
Args:
|
||||
source_inv_depth: Inverse depth map to be aligned. Shape: (H, W).
|
||||
target_depth: Target depth map. Shape: (H, W).
|
||||
target_mask: Mask of valid target pixels. Shape: (H, W).
|
||||
|
||||
Returns:
|
||||
Aligned Depth map. Shape: (H, W).
|
||||
scale: Scaling factor.
|
||||
bias: Bias term.
|
||||
"""
|
||||
target_inv_depth = 1.0 / target_depth
|
||||
source_mask = source_inv_depth > 0
|
||||
target_depth_mask = target_depth > 0
|
||||
|
||||
if target_mask is None:
|
||||
target_mask = target_depth_mask
|
||||
else:
|
||||
target_mask = torch.logical_and(target_mask > 0, target_depth_mask)
|
||||
|
||||
# Remove outliers
|
||||
if quantile_masking:
|
||||
outlier_quantiles = torch.tensor([0.1, 0.9], device=source_inv_depth.device)
|
||||
source_data_low, source_data_high = torch.quantile(source_inv_depth[source_mask], outlier_quantiles)
|
||||
target_data_low, target_data_high = torch.quantile(target_inv_depth[target_mask], outlier_quantiles)
|
||||
source_mask = (source_inv_depth > source_data_low) & (source_inv_depth < source_data_high)
|
||||
target_mask = (target_inv_depth > target_data_low) & (target_inv_depth < target_data_high)
|
||||
|
||||
mask = torch.logical_and(source_mask, target_mask)
|
||||
|
||||
source_data = source_inv_depth[mask].view(-1, 1)
|
||||
target_data = target_inv_depth[mask].view(-1, 1)
|
||||
|
||||
ones = torch.ones((source_data.shape[0], 1), device=source_data.device)
|
||||
source_data_h = torch.cat([source_data, ones], dim=1)
|
||||
transform_matrix = torch.linalg.lstsq(source_data_h, target_data).solution
|
||||
|
||||
scale, bias = transform_matrix[0, 0], transform_matrix[1, 0]
|
||||
aligned_inv_depth = source_inv_depth * scale + bias
|
||||
aligned_depth = torch.clamp(aligned_inv_depth.reciprocal(), min=1e-4)
|
||||
|
||||
return aligned_depth, scale, bias
|
||||
|
||||
|
||||
def align_depth_to_depth(
|
||||
source_depth: torch.Tensor,
|
||||
target_depth: torch.Tensor,
|
||||
target_mask: torch.Tensor | None = None,
|
||||
quantile_masking: bool = True,
|
||||
bias: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Apply affine transformation to align source depth to target depth.
|
||||
|
||||
Args:
|
||||
source_inv_depth: Depth map to be aligned. Shape: (H, W).
|
||||
target_depth: Target depth map. Shape: (H, W).
|
||||
target_mask: Mask of valid target pixels. Shape: (H, W).
|
||||
|
||||
Returns:
|
||||
Aligned Depth map. Shape: (H, W).
|
||||
"""
|
||||
source_mask = source_depth > 0
|
||||
target_depth_mask = target_depth > 0
|
||||
|
||||
if target_mask is None:
|
||||
target_mask = target_depth_mask
|
||||
else:
|
||||
target_mask = torch.logical_and(target_mask > 0, target_depth_mask)
|
||||
|
||||
# Remove outliers
|
||||
if quantile_masking:
|
||||
outlier_quantiles = torch.tensor([0.1, 0.9], device=source_depth.device)
|
||||
|
||||
source_data_low, source_data_high = torch.quantile(source_depth[source_mask], outlier_quantiles)
|
||||
target_data_low, target_data_high = torch.quantile(target_depth[target_mask], outlier_quantiles)
|
||||
source_mask = (source_depth > source_data_low) & (source_depth < source_data_high)
|
||||
target_mask = (target_depth > target_data_low) & (target_depth < target_data_high)
|
||||
|
||||
mask = torch.logical_and(source_mask, target_mask)
|
||||
|
||||
source_data = source_depth[mask].view(-1, 1)
|
||||
target_data = target_depth[mask].view(-1, 1)
|
||||
|
||||
if not bias:
|
||||
ones = torch.ones((source_data.shape[0], 1), device=source_data.device)
|
||||
source_data_h = torch.cat([source_data, ones], dim=1)
|
||||
transform_matrix = torch.linalg.lstsq(source_data_h, target_data).solution
|
||||
|
||||
scale, bias = transform_matrix[0, 0], transform_matrix[1, 0]
|
||||
aligned_depth = source_depth * scale + bias
|
||||
aligned_depth = torch.clamp(aligned_depth, min=1e-4)
|
||||
|
||||
else:
|
||||
scale = torch.median(target_data / source_data)
|
||||
aligned_depth = source_depth * scale
|
||||
aligned_depth = torch.clamp(aligned_depth, min=1e-4)
|
||||
|
||||
return aligned_depth
|
||||
@@ -0,0 +1,104 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class DepthType(Enum):
|
||||
"""
|
||||
Type of depth estimated.
|
||||
"""
|
||||
|
||||
# Inverse metric depth (DepthPro, Metric3D), scale is determined by focal.
|
||||
METRIC_DEPTH = "metric_depth"
|
||||
|
||||
# Metric depth (UniDepth), scale is determined by focal but require re-running the model.
|
||||
MODEL_METRIC_DEPTH = "model_metric_depth"
|
||||
|
||||
# Metric distance, scale is determined by focal but require re-running the model.
|
||||
MODEL_METRIC_DISTANCE = "model_metric_distance"
|
||||
|
||||
# Affine-invariant inverse depth (DepthAnything, MoGE), affine needs to be solved per-estimation.
|
||||
AFFINE_DISP = "affine_disp"
|
||||
|
||||
# Scale-invariant depth (DUSt3R), scale needs to be solved per-estimation.
|
||||
SCALE_DISP = "scale_disp"
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class DepthEstimationResult:
|
||||
"""
|
||||
Dataclass for depth estimation results.
|
||||
|
||||
- relative_inv_depth: The estimated inverse depth map ([B,], H, W) in relative scale.
|
||||
One has to estimate scale and offset of such estimation
|
||||
- metric_depth: The estimated depth map ([B,], H, W) in metric scale.
|
||||
- confidence: The confidence map ([B,], H, W).
|
||||
"""
|
||||
|
||||
relative_inv_depth: torch.Tensor | None = None
|
||||
metric_depth: torch.Tensor | None = None
|
||||
confidence: torch.Tensor | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class DepthEstimationInput:
|
||||
"""
|
||||
Dataclass for depth estimation inputs.
|
||||
|
||||
- rgb: The source image ([B,], H, W, 3), should be within 0-1 float.
|
||||
- video_frame_list: The list of video frames (H, W, 3), should be within 0-1 float, length is T.
|
||||
we use numpy here mainly to enforce CPU tensor.
|
||||
- focal_length: The focal length of the camera.
|
||||
with a simple pinhole assumption where fx = fy and cx = w/2, cy = h/2.
|
||||
"""
|
||||
|
||||
rgb: torch.Tensor | None = None
|
||||
video_frame_list: list[np.ndarray] | None = None
|
||||
prompt_metric_depth: torch.Tensor | None = None
|
||||
focal_length: float | None = None
|
||||
|
||||
|
||||
class DepthEstimationModel(ABC):
|
||||
"""
|
||||
Unified interface for depth prediction models.
|
||||
"""
|
||||
|
||||
@property
|
||||
def depth_type(self) -> DepthType:
|
||||
"""
|
||||
Type of depth estimated.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
|
||||
"""
|
||||
Estimate a single optical flow result from two images.
|
||||
"""
|
||||
|
||||
|
||||
class DummyDepthModel(DepthEstimationModel):
|
||||
"""
|
||||
Dummy model that would not predict anything.
|
||||
"""
|
||||
|
||||
def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
|
||||
return DepthEstimationResult()
|
||||
@@ -0,0 +1,105 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
from vipe.utils.misc import unpack_optional
|
||||
|
||||
from ..base import DepthEstimationInput, DepthEstimationModel, DepthEstimationResult, DepthType
|
||||
from .dpt import DepthAnythingV2
|
||||
|
||||
|
||||
class DepthAnythingDepthModel(DepthEstimationModel):
|
||||
"""
|
||||
https://github.com/DepthAnything/Depth-Anything-V2
|
||||
"""
|
||||
|
||||
def __init__(self, model: str = "vitl", ckpt: str = "default") -> None:
|
||||
super().__init__()
|
||||
|
||||
self.model_config = {
|
||||
"vits": {
|
||||
"encoder": "vits",
|
||||
"features": 64,
|
||||
"out_channels": [48, 96, 192, 384],
|
||||
},
|
||||
"vitb": {
|
||||
"encoder": "vitb",
|
||||
"features": 128,
|
||||
"out_channels": [96, 192, 384, 768],
|
||||
},
|
||||
"vitl": {
|
||||
"encoder": "vitl",
|
||||
"features": 256,
|
||||
"out_channels": [256, 512, 1024, 1024],
|
||||
},
|
||||
"vitg": {
|
||||
"encoder": "vitg",
|
||||
"features": 384,
|
||||
"out_channels": [1536, 1536, 1536, 1536],
|
||||
},
|
||||
}[model]
|
||||
url_name = {"vits": "Small", "vitb": "Base", "vitl": "Large", "vitg": "Giant"}[model]
|
||||
|
||||
if ckpt == "default":
|
||||
self.is_metric = False
|
||||
self.max_depth = None
|
||||
self.ckpt_url = f"https://huggingface.co/depth-anything/Depth-Anything-V2-{url_name}/resolve/main/depth_anything_v2_{model}.pth?download=true"
|
||||
|
||||
elif ckpt == "metric-indoor":
|
||||
self.is_metric = True
|
||||
self.max_depth = 20.0
|
||||
self.ckpt_url = f"https://huggingface.co/depth-anything/Depth-Anything-V2-Metric-Hypersim-{url_name}/resolve/main/depth_anything_v2_metric_hypersim_{model}.pth?download=true"
|
||||
|
||||
elif ckpt == "metric-outdoor":
|
||||
self.is_metric = True
|
||||
self.max_depth = 80.0
|
||||
self.ckpt_url = f"hhttps://huggingface.co/depth-anything/Depth-Anything-V2-Metric-VKITTI-{url_name}/resolve/main/depth_anything_v2_metric_vkitti_{model}.pth?download=true"
|
||||
|
||||
else:
|
||||
raise ValueError("Invalid checkpoint name.")
|
||||
|
||||
self.model = DepthAnythingV2(**self.model_config, max_depth=self.max_depth)
|
||||
self.model.load_state_dict(torch.hub.load_state_dict_from_url(self.ckpt_url, map_location="cpu"))
|
||||
self.model.cuda().eval()
|
||||
|
||||
@property
|
||||
def depth_type(self) -> DepthType:
|
||||
return DepthType.AFFINE_DISP
|
||||
|
||||
def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
|
||||
rgb: torch.Tensor = unpack_optional(src.rgb)
|
||||
assert rgb.dtype == torch.float32, "Input image should be float32"
|
||||
|
||||
if rgb.dim() == 3:
|
||||
rgb, batch_dim = rgb[None], False
|
||||
else:
|
||||
batch_dim = True
|
||||
|
||||
assert rgb.size(0) == 1, "Only batch size 1 is supported."
|
||||
|
||||
# TODO: Support pytorch pre-processing
|
||||
depth = self.model.infer_image(rgb[0].cpu().numpy())
|
||||
|
||||
if batch_dim:
|
||||
depth = depth[None]
|
||||
|
||||
if self.is_metric:
|
||||
return DepthEstimationResult(metric_depth=depth)
|
||||
|
||||
else:
|
||||
return DepthEstimationResult(relative_inv_depth=depth)
|
||||
@@ -0,0 +1,429 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
#
|
||||
# This source code is licensed under the Apache License, Version 2.0
|
||||
# found in the LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
|
||||
|
||||
import logging
|
||||
import math
|
||||
|
||||
from functools import partial
|
||||
from typing import Callable, Sequence, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint
|
||||
|
||||
from torch.nn.init import trunc_normal_
|
||||
|
||||
from .dinov2_layers import MemEffAttention, Mlp, PatchEmbed
|
||||
from .dinov2_layers import NestedTensorBlock as Block
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
|
||||
if not depth_first and include_root:
|
||||
fn(module=module, name=name)
|
||||
for child_name, child_module in module.named_children():
|
||||
child_name = ".".join((name, child_name)) if name else child_name
|
||||
named_apply(
|
||||
fn=fn,
|
||||
module=child_module,
|
||||
name=child_name,
|
||||
depth_first=depth_first,
|
||||
include_root=True,
|
||||
)
|
||||
if depth_first and include_root:
|
||||
fn(module=module, name=name)
|
||||
return module
|
||||
|
||||
|
||||
class BlockChunk(nn.ModuleList):
|
||||
def forward(self, x):
|
||||
for b in self:
|
||||
x = b(x)
|
||||
return x
|
||||
|
||||
|
||||
class DinoVisionTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
img_size=224,
|
||||
patch_size=16,
|
||||
in_chans=3,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4.0,
|
||||
qkv_bias=True,
|
||||
ffn_bias=True,
|
||||
proj_bias=True,
|
||||
drop_path_rate=0.0,
|
||||
drop_path_uniform=False,
|
||||
init_values=None, # for layerscale: None or 0 => no layerscale
|
||||
embed_layer=PatchEmbed,
|
||||
act_layer=nn.GELU,
|
||||
block_fn=Block,
|
||||
ffn_layer="mlp",
|
||||
block_chunks=1,
|
||||
num_register_tokens=0,
|
||||
interpolate_antialias=False,
|
||||
interpolate_offset=0.1,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
img_size (int, tuple): input image size
|
||||
patch_size (int, tuple): patch size
|
||||
in_chans (int): number of input channels
|
||||
embed_dim (int): embedding dimension
|
||||
depth (int): depth of transformer
|
||||
num_heads (int): number of attention heads
|
||||
mlp_ratio (int): ratio of mlp hidden dim to embedding dim
|
||||
qkv_bias (bool): enable bias for qkv if True
|
||||
proj_bias (bool): enable bias for proj in attn if True
|
||||
ffn_bias (bool): enable bias for ffn if True
|
||||
drop_path_rate (float): stochastic depth rate
|
||||
drop_path_uniform (bool): apply uniform drop rate across blocks
|
||||
weight_init (str): weight init scheme
|
||||
init_values (float): layer-scale init values
|
||||
embed_layer (nn.Module): patch embedding layer
|
||||
act_layer (nn.Module): MLP activation layer
|
||||
block_fn (nn.Module): transformer block class
|
||||
ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
|
||||
block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
|
||||
num_register_tokens: (int) number of extra cls tokens (so-called "registers")
|
||||
interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
|
||||
interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
|
||||
"""
|
||||
super().__init__()
|
||||
norm_layer = partial(nn.LayerNorm, eps=1e-6)
|
||||
|
||||
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
|
||||
self.num_tokens = 1
|
||||
self.n_blocks = depth
|
||||
self.num_heads = num_heads
|
||||
self.patch_size = patch_size
|
||||
self.num_register_tokens = num_register_tokens
|
||||
self.interpolate_antialias = interpolate_antialias
|
||||
self.interpolate_offset = interpolate_offset
|
||||
|
||||
self.patch_embed = embed_layer(
|
||||
img_size=img_size,
|
||||
patch_size=patch_size,
|
||||
in_chans=in_chans,
|
||||
embed_dim=embed_dim,
|
||||
)
|
||||
num_patches = self.patch_embed.num_patches
|
||||
|
||||
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
|
||||
assert num_register_tokens >= 0
|
||||
self.register_tokens = (
|
||||
nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
|
||||
)
|
||||
|
||||
if drop_path_uniform is True:
|
||||
dpr = [drop_path_rate] * depth
|
||||
else:
|
||||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
|
||||
|
||||
if ffn_layer == "mlp":
|
||||
# logger.info("using MLP layer as FFN")
|
||||
ffn_layer = Mlp
|
||||
elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
|
||||
# logger.info("using SwiGLU layer as FFN")
|
||||
ffn_layer = SwiGLUFFNFused
|
||||
elif ffn_layer == "identity":
|
||||
# logger.info("using Identity layer as FFN")
|
||||
|
||||
def f(*args, **kwargs):
|
||||
return nn.Identity()
|
||||
|
||||
ffn_layer = f
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
blocks_list = [
|
||||
block_fn(
|
||||
dim=embed_dim,
|
||||
num_heads=num_heads,
|
||||
mlp_ratio=mlp_ratio,
|
||||
qkv_bias=qkv_bias,
|
||||
proj_bias=proj_bias,
|
||||
ffn_bias=ffn_bias,
|
||||
drop_path=dpr[i],
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
ffn_layer=ffn_layer,
|
||||
init_values=init_values,
|
||||
)
|
||||
for i in range(depth)
|
||||
]
|
||||
if block_chunks > 0:
|
||||
self.chunked_blocks = True
|
||||
chunked_blocks = []
|
||||
chunksize = depth // block_chunks
|
||||
for i in range(0, depth, chunksize):
|
||||
# this is to keep the block index consistent if we chunk the block list
|
||||
chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
|
||||
self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
|
||||
else:
|
||||
self.chunked_blocks = False
|
||||
self.blocks = nn.ModuleList(blocks_list)
|
||||
|
||||
self.norm = norm_layer(embed_dim)
|
||||
self.head = nn.Identity()
|
||||
|
||||
self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
|
||||
|
||||
self.init_weights()
|
||||
|
||||
def init_weights(self):
|
||||
trunc_normal_(self.pos_embed, std=0.02)
|
||||
nn.init.normal_(self.cls_token, std=1e-6)
|
||||
if self.register_tokens is not None:
|
||||
nn.init.normal_(self.register_tokens, std=1e-6)
|
||||
named_apply(init_weights_vit_timm, self)
|
||||
|
||||
def interpolate_pos_encoding(self, x, w, h):
|
||||
previous_dtype = x.dtype
|
||||
npatch = x.shape[1] - 1
|
||||
N = self.pos_embed.shape[1] - 1
|
||||
if npatch == N and w == h:
|
||||
return self.pos_embed
|
||||
pos_embed = self.pos_embed.float()
|
||||
class_pos_embed = pos_embed[:, 0]
|
||||
patch_pos_embed = pos_embed[:, 1:]
|
||||
dim = x.shape[-1]
|
||||
w0 = w // self.patch_size
|
||||
h0 = h // self.patch_size
|
||||
# we add a small number to avoid floating point error in the interpolation
|
||||
# see discussion at https://github.com/facebookresearch/dino/issues/8
|
||||
# DINOv2 with register modify the interpolate_offset from 0.1 to 0.0
|
||||
w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
|
||||
# w0, h0 = w0 + 0.1, h0 + 0.1
|
||||
|
||||
sqrt_N = math.sqrt(N)
|
||||
sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
|
||||
patch_pos_embed = nn.functional.interpolate(
|
||||
patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
|
||||
scale_factor=(sx, sy),
|
||||
# (int(w0), int(h0)), # to solve the upsampling shape issue
|
||||
mode="bicubic",
|
||||
antialias=self.interpolate_antialias,
|
||||
)
|
||||
|
||||
assert int(w0) == patch_pos_embed.shape[-2]
|
||||
assert int(h0) == patch_pos_embed.shape[-1]
|
||||
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
|
||||
return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
|
||||
|
||||
def prepare_tokens_with_masks(self, x, masks=None):
|
||||
B, nc, w, h = x.shape
|
||||
x = self.patch_embed(x)
|
||||
if masks is not None:
|
||||
x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
|
||||
|
||||
x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
|
||||
x = x + self.interpolate_pos_encoding(x, w, h)
|
||||
|
||||
if self.register_tokens is not None:
|
||||
x = torch.cat(
|
||||
(
|
||||
x[:, :1],
|
||||
self.register_tokens.expand(x.shape[0], -1, -1),
|
||||
x[:, 1:],
|
||||
),
|
||||
dim=1,
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
def forward_features_list(self, x_list, masks_list):
|
||||
x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
|
||||
for blk in self.blocks:
|
||||
x = blk(x)
|
||||
|
||||
all_x = x
|
||||
output = []
|
||||
for x, masks in zip(all_x, masks_list):
|
||||
x_norm = self.norm(x)
|
||||
output.append(
|
||||
{
|
||||
"x_norm_clstoken": x_norm[:, 0],
|
||||
"x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
|
||||
"x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
|
||||
"x_prenorm": x,
|
||||
"masks": masks,
|
||||
}
|
||||
)
|
||||
return output
|
||||
|
||||
def forward_features(self, x, masks=None):
|
||||
if isinstance(x, list):
|
||||
return self.forward_features_list(x, masks)
|
||||
|
||||
x = self.prepare_tokens_with_masks(x, masks)
|
||||
|
||||
for blk in self.blocks:
|
||||
x = blk(x)
|
||||
|
||||
x_norm = self.norm(x)
|
||||
return {
|
||||
"x_norm_clstoken": x_norm[:, 0],
|
||||
"x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
|
||||
"x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
|
||||
"x_prenorm": x,
|
||||
"masks": masks,
|
||||
}
|
||||
|
||||
def _get_intermediate_layers_not_chunked(self, x, n=1):
|
||||
x = self.prepare_tokens_with_masks(x)
|
||||
# If n is an int, take the n last blocks. If it's a list, take them
|
||||
output, total_block_len = [], len(self.blocks)
|
||||
blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
|
||||
for i, blk in enumerate(self.blocks):
|
||||
x = blk(x)
|
||||
if i in blocks_to_take:
|
||||
output.append(x)
|
||||
assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
|
||||
return output
|
||||
|
||||
def _get_intermediate_layers_chunked(self, x, n=1):
|
||||
x = self.prepare_tokens_with_masks(x)
|
||||
output, i, total_block_len = [], 0, len(self.blocks[-1])
|
||||
# If n is an int, take the n last blocks. If it's a list, take them
|
||||
blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
|
||||
for block_chunk in self.blocks:
|
||||
for blk in block_chunk[i:]: # Passing the nn.Identity()
|
||||
x = blk(x)
|
||||
if i in blocks_to_take:
|
||||
output.append(x)
|
||||
i += 1
|
||||
assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
|
||||
return output
|
||||
|
||||
def get_intermediate_layers(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
n: Union[int, Sequence] = 1, # Layers or n last layers to take
|
||||
reshape: bool = False,
|
||||
return_class_token: bool = False,
|
||||
norm=True,
|
||||
) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
|
||||
if self.chunked_blocks:
|
||||
outputs = self._get_intermediate_layers_chunked(x, n)
|
||||
else:
|
||||
outputs = self._get_intermediate_layers_not_chunked(x, n)
|
||||
if norm:
|
||||
outputs = [self.norm(out) for out in outputs]
|
||||
class_tokens = [out[:, 0] for out in outputs]
|
||||
outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs]
|
||||
if reshape:
|
||||
B, _, w, h = x.shape
|
||||
outputs = [
|
||||
out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
|
||||
for out in outputs
|
||||
]
|
||||
if return_class_token:
|
||||
return tuple(zip(outputs, class_tokens))
|
||||
return tuple(outputs)
|
||||
|
||||
def forward(self, *args, is_training=False, **kwargs):
|
||||
ret = self.forward_features(*args, **kwargs)
|
||||
if is_training:
|
||||
return ret
|
||||
else:
|
||||
return self.head(ret["x_norm_clstoken"])
|
||||
|
||||
|
||||
def init_weights_vit_timm(module: nn.Module, name: str = ""):
|
||||
"""ViT weight initialization, original timm impl (for reproducibility)"""
|
||||
if isinstance(module, nn.Linear):
|
||||
trunc_normal_(module.weight, std=0.02)
|
||||
if module.bias is not None:
|
||||
nn.init.zeros_(module.bias)
|
||||
|
||||
|
||||
def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=384,
|
||||
depth=12,
|
||||
num_heads=6,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=1024,
|
||||
depth=24,
|
||||
num_heads=16,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
"""
|
||||
Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
|
||||
"""
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=1536,
|
||||
depth=40,
|
||||
num_heads=24,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def DINOv2(model_name):
|
||||
model_zoo = {
|
||||
"vits": vit_small,
|
||||
"vitb": vit_base,
|
||||
"vitl": vit_large,
|
||||
"vitg": vit_giant2,
|
||||
}
|
||||
|
||||
return model_zoo[model_name](
|
||||
img_size=518,
|
||||
patch_size=14,
|
||||
init_values=1.0,
|
||||
ffn_layer="mlp" if model_name != "vitg" else "swiglufused",
|
||||
block_chunks=0,
|
||||
num_register_tokens=0,
|
||||
interpolate_antialias=False,
|
||||
interpolate_offset=0.1,
|
||||
)
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .attention import MemEffAttention
|
||||
from .block import NestedTensorBlock
|
||||
from .mlp import Mlp
|
||||
from .patch_embed import PatchEmbed
|
||||
from .swiglu_ffn import SwiGLUFFN
|
||||
+71
@@ -0,0 +1,71 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
|
||||
from torch import Tensor, nn
|
||||
|
||||
from vipe.ext.xformers import memory_efficient_attention
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int = 8,
|
||||
qkv_bias: bool = False,
|
||||
proj_bias: bool = True,
|
||||
attn_drop: float = 0.0,
|
||||
proj_drop: float = 0.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
self.scale = 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, bias=proj_bias)
|
||||
self.proj_drop = nn.Dropout(proj_drop)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
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] * self.scale, qkv[1], qkv[2]
|
||||
attn = q @ k.transpose(-2, -1)
|
||||
|
||||
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 MemEffAttention(Attention):
|
||||
def forward(self, x: Tensor, attn_bias=None) -> Tensor:
|
||||
B, N, C = x.shape
|
||||
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
|
||||
|
||||
q, k, v = torch.unbind(qkv, dim=2)
|
||||
|
||||
x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
|
||||
x = x.reshape([B, N, C])
|
||||
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
return x
|
||||
+253
@@ -0,0 +1,253 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
|
||||
|
||||
import logging
|
||||
|
||||
from typing import Any, Callable, Dict, List, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from torch import Tensor, nn
|
||||
|
||||
from vipe.ext.xformers import index_select_cat, scaled_index_add
|
||||
|
||||
from .attention import Attention, MemEffAttention
|
||||
from .drop_path import DropPath
|
||||
from .layer_scale import LayerScale
|
||||
from .mlp import Mlp
|
||||
|
||||
|
||||
# from xformers.ops import fmha
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qkv_bias: bool = False,
|
||||
proj_bias: bool = True,
|
||||
ffn_bias: bool = True,
|
||||
drop: float = 0.0,
|
||||
attn_drop: float = 0.0,
|
||||
init_values=None,
|
||||
drop_path: float = 0.0,
|
||||
act_layer: Callable[..., nn.Module] = nn.GELU,
|
||||
norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
|
||||
attn_class: Callable[..., nn.Module] = Attention,
|
||||
ffn_layer: Callable[..., nn.Module] = Mlp,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
# print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
|
||||
self.norm1 = norm_layer(dim)
|
||||
self.attn = attn_class(
|
||||
dim,
|
||||
num_heads=num_heads,
|
||||
qkv_bias=qkv_bias,
|
||||
proj_bias=proj_bias,
|
||||
attn_drop=attn_drop,
|
||||
proj_drop=drop,
|
||||
)
|
||||
self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
||||
self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
|
||||
self.norm2 = norm_layer(dim)
|
||||
mlp_hidden_dim = int(dim * mlp_ratio)
|
||||
self.mlp = ffn_layer(
|
||||
in_features=dim,
|
||||
hidden_features=mlp_hidden_dim,
|
||||
act_layer=act_layer,
|
||||
drop=drop,
|
||||
bias=ffn_bias,
|
||||
)
|
||||
self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
||||
self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
|
||||
self.sample_drop_ratio = drop_path
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
def attn_residual_func(x: Tensor) -> Tensor:
|
||||
return self.ls1(self.attn(self.norm1(x)))
|
||||
|
||||
def ffn_residual_func(x: Tensor) -> Tensor:
|
||||
return self.ls2(self.mlp(self.norm2(x)))
|
||||
|
||||
if self.training and self.sample_drop_ratio > 0.1:
|
||||
# the overhead is compensated only for a drop path rate larger than 0.1
|
||||
x = drop_add_residual_stochastic_depth(
|
||||
x,
|
||||
residual_func=attn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
)
|
||||
x = drop_add_residual_stochastic_depth(
|
||||
x,
|
||||
residual_func=ffn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
)
|
||||
elif self.training and self.sample_drop_ratio > 0.0:
|
||||
x = x + self.drop_path1(attn_residual_func(x))
|
||||
x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2
|
||||
else:
|
||||
x = x + attn_residual_func(x)
|
||||
x = x + ffn_residual_func(x)
|
||||
return x
|
||||
|
||||
|
||||
def drop_add_residual_stochastic_depth(
|
||||
x: Tensor,
|
||||
residual_func: Callable[[Tensor], Tensor],
|
||||
sample_drop_ratio: float = 0.0,
|
||||
) -> Tensor:
|
||||
# 1) extract subset using permutation
|
||||
b, n, d = x.shape
|
||||
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
|
||||
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
|
||||
x_subset = x[brange]
|
||||
|
||||
# 2) apply residual_func to get residual
|
||||
residual = residual_func(x_subset)
|
||||
|
||||
x_flat = x.flatten(1)
|
||||
residual = residual.flatten(1)
|
||||
|
||||
residual_scale_factor = b / sample_subset_size
|
||||
|
||||
# 3) add the residual
|
||||
x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
|
||||
return x_plus_residual.view_as(x)
|
||||
|
||||
|
||||
def get_branges_scales(x, sample_drop_ratio=0.0):
|
||||
b, n, d = x.shape
|
||||
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
|
||||
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
|
||||
residual_scale_factor = b / sample_subset_size
|
||||
return brange, residual_scale_factor
|
||||
|
||||
|
||||
def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
|
||||
if scaling_vector is None:
|
||||
x_flat = x.flatten(1)
|
||||
residual = residual.flatten(1)
|
||||
x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
|
||||
else:
|
||||
x_plus_residual = scaled_index_add(
|
||||
x,
|
||||
brange,
|
||||
residual.to(dtype=x.dtype),
|
||||
scaling=scaling_vector,
|
||||
alpha=residual_scale_factor,
|
||||
)
|
||||
return x_plus_residual
|
||||
|
||||
|
||||
attn_bias_cache: Dict[Tuple, Any] = {}
|
||||
|
||||
|
||||
def get_attn_bias_and_cat(x_list, branges=None):
|
||||
"""
|
||||
this will perform the index select, cat the tensors, and provide the attn_bias from cache
|
||||
"""
|
||||
batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
|
||||
all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
|
||||
if all_shapes not in attn_bias_cache.keys():
|
||||
seqlens = []
|
||||
for b, x in zip(batch_sizes, x_list):
|
||||
for _ in range(b):
|
||||
seqlens.append(x.shape[1])
|
||||
attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
|
||||
attn_bias._batch_sizes = batch_sizes
|
||||
attn_bias_cache[all_shapes] = attn_bias
|
||||
|
||||
if branges is not None:
|
||||
cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
|
||||
else:
|
||||
tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
|
||||
cat_tensors = torch.cat(tensors_bs1, dim=1)
|
||||
|
||||
return attn_bias_cache[all_shapes], cat_tensors
|
||||
|
||||
|
||||
def drop_add_residual_stochastic_depth_list(
|
||||
x_list: List[Tensor],
|
||||
residual_func: Callable[[Tensor, Any], Tensor],
|
||||
sample_drop_ratio: float = 0.0,
|
||||
scaling_vector=None,
|
||||
) -> Tensor:
|
||||
# 1) generate random set of indices for dropping samples in the batch
|
||||
branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
|
||||
branges = [s[0] for s in branges_scales]
|
||||
residual_scale_factors = [s[1] for s in branges_scales]
|
||||
|
||||
# 2) get attention bias and index+concat the tensors
|
||||
attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
|
||||
|
||||
# 3) apply residual_func to get residual, and split the result
|
||||
residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore
|
||||
|
||||
outputs = []
|
||||
for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
|
||||
outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
|
||||
return outputs
|
||||
|
||||
|
||||
class NestedTensorBlock(Block):
|
||||
def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
|
||||
"""
|
||||
x_list contains a list of tensors to nest together and run
|
||||
"""
|
||||
assert isinstance(self.attn, MemEffAttention)
|
||||
|
||||
if self.training and self.sample_drop_ratio > 0.0:
|
||||
|
||||
def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.attn(self.norm1(x), attn_bias=attn_bias)
|
||||
|
||||
def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.mlp(self.norm2(x))
|
||||
|
||||
x_list = drop_add_residual_stochastic_depth_list(
|
||||
x_list,
|
||||
residual_func=attn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
scaling_vector=(self.ls1.gamma if isinstance(self.ls1, LayerScale) else None),
|
||||
)
|
||||
x_list = drop_add_residual_stochastic_depth_list(
|
||||
x_list,
|
||||
residual_func=ffn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
scaling_vector=(self.ls2.gamma if isinstance(self.ls1, LayerScale) else None),
|
||||
)
|
||||
return x_list
|
||||
else:
|
||||
|
||||
def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
|
||||
|
||||
def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.ls2(self.mlp(self.norm2(x)))
|
||||
|
||||
attn_bias, x = get_attn_bias_and_cat(x_list)
|
||||
x = x + attn_residual_func(x, attn_bias=attn_bias)
|
||||
x = x + ffn_residual_func(x)
|
||||
return attn_bias.split(x)
|
||||
|
||||
def forward(self, x_or_x_list):
|
||||
if isinstance(x_or_x_list, Tensor):
|
||||
return super().forward(x_or_x_list)
|
||||
elif isinstance(x_or_x_list, list):
|
||||
assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
|
||||
return self.forward_nested(x_or_x_list)
|
||||
else:
|
||||
raise AssertionError
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
|
||||
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
def drop_path(x, drop_prob: float = 0.0, training: bool = False):
|
||||
if drop_prob == 0.0 or not training:
|
||||
return x
|
||||
keep_prob = 1 - drop_prob
|
||||
shape = (x.shape[0],) + (1,) * (
|
||||
x.ndim - 1
|
||||
) # work with diff dim tensors, not just 2D ConvNets
|
||||
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
|
||||
if keep_prob > 0.0:
|
||||
random_tensor.div_(keep_prob)
|
||||
output = x * random_tensor
|
||||
return output
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
|
||||
|
||||
def __init__(self, drop_prob=None):
|
||||
super(DropPath, self).__init__()
|
||||
self.drop_prob = drop_prob
|
||||
|
||||
def forward(self, x):
|
||||
return drop_path(x, self.drop_prob, self.training)
|
||||
+27
@@ -0,0 +1,27 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
|
||||
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
class LayerScale(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
init_values: Union[float, Tensor] = 1e-5,
|
||||
inplace: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.inplace = inplace
|
||||
self.gamma = nn.Parameter(init_values * torch.ones(dim))
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return x.mul_(self.gamma) if self.inplace else x * self.gamma
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
|
||||
|
||||
|
||||
from typing import Callable, Optional
|
||||
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
class Mlp(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: Optional[int] = None,
|
||||
out_features: Optional[int] = None,
|
||||
act_layer: Callable[..., nn.Module] = nn.GELU,
|
||||
drop: float = 0.0,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
|
||||
self.act = act_layer()
|
||||
self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
|
||||
self.drop = nn.Dropout(drop)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
x = self.drop(x)
|
||||
x = self.fc2(x)
|
||||
x = self.drop(x)
|
||||
return x
|
||||
+114
@@ -0,0 +1,114 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
|
||||
|
||||
from typing import Callable, Optional, Tuple, Union
|
||||
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
def make_2tuple(x):
|
||||
if isinstance(x, tuple):
|
||||
assert len(x) == 2
|
||||
return x
|
||||
|
||||
assert isinstance(x, int)
|
||||
return (x, x)
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
"""
|
||||
2D image to patch embedding: (B,C,H,W) -> (B,N,D)
|
||||
|
||||
Args:
|
||||
img_size: Image size.
|
||||
patch_size: Patch token size.
|
||||
in_chans: Number of input image channels.
|
||||
embed_dim: Number of linear projection output channels.
|
||||
norm_layer: Normalization layer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
img_size: Union[int, Tuple[int, int]] = 224,
|
||||
patch_size: Union[int, Tuple[int, int]] = 16,
|
||||
in_chans: int = 3,
|
||||
embed_dim: int = 768,
|
||||
norm_layer: Optional[Callable] = None,
|
||||
flatten_embedding: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
image_HW = make_2tuple(img_size)
|
||||
patch_HW = make_2tuple(patch_size)
|
||||
patch_grid_size = (
|
||||
image_HW[0] // patch_HW[0],
|
||||
image_HW[1] // patch_HW[1],
|
||||
)
|
||||
|
||||
self.img_size = image_HW
|
||||
self.patch_size = patch_HW
|
||||
self.patches_resolution = patch_grid_size
|
||||
self.num_patches = patch_grid_size[0] * patch_grid_size[1]
|
||||
|
||||
self.in_chans = in_chans
|
||||
self.embed_dim = embed_dim
|
||||
|
||||
self.flatten_embedding = flatten_embedding
|
||||
|
||||
self.proj = nn.Conv2d(
|
||||
in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW
|
||||
)
|
||||
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
|
||||
|
||||
def init_alpha_conv(self, cond_channels=2):
|
||||
self.alpha_proj = nn.Conv2d(
|
||||
cond_channels,
|
||||
self.embed_dim,
|
||||
kernel_size=self.patch_size,
|
||||
stride=self.patch_size,
|
||||
)
|
||||
nn.init.constant_(self.alpha_proj.weight, 0)
|
||||
nn.init.constant_(self.alpha_proj.bias, 0)
|
||||
|
||||
def forward(self, x: Tensor, condition=None) -> Tensor:
|
||||
_, _, H, W = x.shape
|
||||
patch_H, patch_W = self.patch_size
|
||||
|
||||
assert (
|
||||
H % patch_H == 0
|
||||
), f"Input image height {H} is not a multiple of patch height {patch_H}"
|
||||
assert (
|
||||
W % patch_W == 0
|
||||
), f"Input image width {W} is not a multiple of patch width: {patch_W}"
|
||||
|
||||
x = self.proj(x) # B C H W
|
||||
if hasattr(self, "alpha_proj"):
|
||||
x += self.alpha_proj(condition)
|
||||
|
||||
H, W = x.size(2), x.size(3)
|
||||
x = x.flatten(2).transpose(1, 2) # B HW C
|
||||
x = self.norm(x)
|
||||
if not self.flatten_embedding:
|
||||
x = x.reshape(-1, H, W, self.embed_dim) # B H W C
|
||||
return x
|
||||
|
||||
def flops(self) -> float:
|
||||
Ho, Wo = self.patches_resolution
|
||||
flops = (
|
||||
Ho
|
||||
* Wo
|
||||
* self.embed_dim
|
||||
* self.in_chans
|
||||
* (self.patch_size[0] * self.patch_size[1])
|
||||
)
|
||||
if self.norm is not None:
|
||||
flops += Ho * Wo * self.embed_dim
|
||||
return flops
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from typing import Callable, Optional
|
||||
|
||||
import torch.nn.functional as F
|
||||
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
class SwiGLUFFN(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: Optional[int] = None,
|
||||
out_features: Optional[int] = None,
|
||||
act_layer: Callable[..., nn.Module] = None,
|
||||
drop: float = 0.0,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
|
||||
self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x12 = self.w12(x)
|
||||
x1, x2 = x12.chunk(2, dim=-1)
|
||||
hidden = F.silu(x1) * x2
|
||||
return self.w3(hidden)
|
||||
@@ -0,0 +1,281 @@
|
||||
# This file includes code originally from the Depth Anything V2 repository:
|
||||
# https://github.com/DepthAnything/Depth-Anything-V2
|
||||
# Licensed under the Apache License, Version 2.0. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from torchvision.transforms import Compose
|
||||
|
||||
from .dinov2 import DINOv2
|
||||
from .util.blocks import FeatureFusionBlock, _make_scratch
|
||||
from .util.transform import NormalizeImage, PrepareForNet, Resize
|
||||
|
||||
|
||||
def _make_fusion_block(features, use_bn, size=None):
|
||||
return FeatureFusionBlock(
|
||||
features,
|
||||
nn.ReLU(False),
|
||||
deconv=False,
|
||||
bn=use_bn,
|
||||
expand=False,
|
||||
align_corners=True,
|
||||
size=size,
|
||||
)
|
||||
|
||||
|
||||
class ConvBlock(nn.Module):
|
||||
def __init__(self, in_feature, out_feature):
|
||||
super().__init__()
|
||||
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.Conv2d(in_feature, out_feature, kernel_size=3, stride=1, padding=1),
|
||||
nn.BatchNorm2d(out_feature),
|
||||
nn.ReLU(True),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv_block(x)
|
||||
|
||||
|
||||
class DPTHead(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
features=256,
|
||||
use_bn=False,
|
||||
out_channels=[256, 512, 1024, 1024],
|
||||
use_clstoken=False,
|
||||
is_metric=False,
|
||||
):
|
||||
super(DPTHead, self).__init__()
|
||||
|
||||
self.use_clstoken = use_clstoken
|
||||
|
||||
self.projects = nn.ModuleList(
|
||||
[
|
||||
nn.Conv2d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channel,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
)
|
||||
for out_channel in out_channels
|
||||
]
|
||||
)
|
||||
|
||||
self.resize_layers = nn.ModuleList(
|
||||
[
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=out_channels[0],
|
||||
out_channels=out_channels[0],
|
||||
kernel_size=4,
|
||||
stride=4,
|
||||
padding=0,
|
||||
),
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=out_channels[1],
|
||||
out_channels=out_channels[1],
|
||||
kernel_size=2,
|
||||
stride=2,
|
||||
padding=0,
|
||||
),
|
||||
nn.Identity(),
|
||||
nn.Conv2d(
|
||||
in_channels=out_channels[3],
|
||||
out_channels=out_channels[3],
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
if use_clstoken:
|
||||
self.readout_projects = nn.ModuleList()
|
||||
for _ in range(len(self.projects)):
|
||||
self.readout_projects.append(nn.Sequential(nn.Linear(2 * in_channels, in_channels), nn.GELU()))
|
||||
|
||||
self.scratch = _make_scratch(
|
||||
out_channels,
|
||||
features,
|
||||
groups=1,
|
||||
expand=False,
|
||||
)
|
||||
|
||||
self.scratch.stem_transpose = None
|
||||
|
||||
self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
|
||||
|
||||
head_features_1 = features
|
||||
head_features_2 = 32
|
||||
|
||||
self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
if is_metric:
|
||||
self.scratch.output_conv2 = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
head_features_1 // 2,
|
||||
head_features_2,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
else:
|
||||
self.scratch.output_conv2 = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
head_features_1 // 2,
|
||||
head_features_2,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
|
||||
nn.ReLU(True),
|
||||
nn.Identity(),
|
||||
)
|
||||
|
||||
def forward(self, out_features, patch_h, patch_w):
|
||||
out = []
|
||||
for i, x in enumerate(out_features):
|
||||
if self.use_clstoken:
|
||||
x, cls_token = x[0], x[1]
|
||||
readout = cls_token.unsqueeze(1).expand_as(x)
|
||||
x = self.readout_projects[i](torch.cat((x, readout), -1))
|
||||
else:
|
||||
x = x[0]
|
||||
|
||||
x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
|
||||
|
||||
x = self.projects[i](x)
|
||||
x = self.resize_layers[i](x)
|
||||
|
||||
out.append(x)
|
||||
|
||||
layer_1, layer_2, layer_3, layer_4 = out
|
||||
|
||||
layer_1_rn = self.scratch.layer1_rn(layer_1)
|
||||
layer_2_rn = self.scratch.layer2_rn(layer_2)
|
||||
layer_3_rn = self.scratch.layer3_rn(layer_3)
|
||||
layer_4_rn = self.scratch.layer4_rn(layer_4)
|
||||
|
||||
path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
|
||||
path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
|
||||
path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
|
||||
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
|
||||
|
||||
out = self.scratch.output_conv1(path_1)
|
||||
out = F.interpolate(
|
||||
out,
|
||||
(int(patch_h * 14), int(patch_w * 14)),
|
||||
mode="bilinear",
|
||||
align_corners=True,
|
||||
)
|
||||
out = self.scratch.output_conv2(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class DepthAnythingV2(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
encoder="vitl",
|
||||
features=256,
|
||||
out_channels=[256, 512, 1024, 1024],
|
||||
use_bn=False,
|
||||
use_clstoken=False,
|
||||
max_depth: float | None = 20.0,
|
||||
):
|
||||
super(DepthAnythingV2, self).__init__()
|
||||
|
||||
self.intermediate_layer_idx = {
|
||||
"vits": [2, 5, 8, 11],
|
||||
"vitb": [2, 5, 8, 11],
|
||||
"vitl": [4, 11, 17, 23],
|
||||
"vitg": [9, 19, 29, 39],
|
||||
}
|
||||
|
||||
self.max_depth = max_depth
|
||||
|
||||
self.encoder = encoder
|
||||
self.pretrained = DINOv2(model_name=encoder)
|
||||
|
||||
self.depth_head = DPTHead(
|
||||
self.pretrained.embed_dim,
|
||||
features,
|
||||
use_bn,
|
||||
out_channels=out_channels,
|
||||
use_clstoken=use_clstoken,
|
||||
is_metric=self.max_depth is not None,
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14
|
||||
|
||||
features = self.pretrained.get_intermediate_layers(
|
||||
x, self.intermediate_layer_idx[self.encoder], return_class_token=True
|
||||
)
|
||||
|
||||
depth = self.depth_head(features, patch_h, patch_w)
|
||||
|
||||
if self.max_depth is not None:
|
||||
depth = self.max_depth * depth
|
||||
|
||||
else:
|
||||
depth = F.relu(depth)
|
||||
|
||||
return depth.squeeze(1)
|
||||
|
||||
@torch.no_grad()
|
||||
def infer_image(self, raw_image, input_size=518):
|
||||
image, (h, w) = self.image2tensor(raw_image, input_size)
|
||||
|
||||
depth = self.forward(image)
|
||||
|
||||
depth = F.interpolate(depth[:, None], (h, w), mode="bilinear", align_corners=True)[0, 0]
|
||||
|
||||
# return depth.cpu().numpy()
|
||||
return depth
|
||||
|
||||
def image2tensor(self, raw_image, input_size=518):
|
||||
transform = Compose(
|
||||
[
|
||||
Resize(
|
||||
width=input_size,
|
||||
height=input_size,
|
||||
resize_target=False,
|
||||
keep_aspect_ratio=True,
|
||||
ensure_multiple_of=14,
|
||||
resize_method="lower_bound",
|
||||
image_interpolation_method=cv2.INTER_CUBIC,
|
||||
),
|
||||
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
PrepareForNet(),
|
||||
]
|
||||
)
|
||||
|
||||
h, w = raw_image.shape[:2]
|
||||
|
||||
# image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0
|
||||
image = raw_image
|
||||
|
||||
image = transform({"image": image})["image"]
|
||||
image = torch.from_numpy(image).unsqueeze(0)
|
||||
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
image = image.to(DEVICE)
|
||||
|
||||
return image, (h, w)
|
||||
@@ -0,0 +1,3 @@
|
||||
# This file includes code originally from the Depth Anything V2 repository:
|
||||
# https://github.com/DepthAnything/Depth-Anything-V2
|
||||
# Licensed under the Apache License, Version 2.0. See THIRD_PARTY_LICENSES.md for details.
|
||||
@@ -0,0 +1,207 @@
|
||||
# This file includes code originally from the Depth Anything V2 repository:
|
||||
# https://github.com/DepthAnything/Depth-Anything-V2
|
||||
# Licensed under the Apache License, Version 2.0. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def _make_scratch(in_shape, out_shape, groups=1, expand=False):
|
||||
scratch = nn.Module()
|
||||
|
||||
out_shape1 = out_shape
|
||||
out_shape2 = out_shape
|
||||
out_shape3 = out_shape
|
||||
if len(in_shape) >= 4:
|
||||
out_shape4 = out_shape
|
||||
|
||||
if expand:
|
||||
out_shape1 = out_shape
|
||||
out_shape2 = out_shape * 2
|
||||
out_shape3 = out_shape * 4
|
||||
if len(in_shape) >= 4:
|
||||
out_shape4 = out_shape * 8
|
||||
|
||||
scratch.layer1_rn = nn.Conv2d(
|
||||
in_shape[0],
|
||||
out_shape1,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=False,
|
||||
groups=groups,
|
||||
)
|
||||
scratch.layer2_rn = nn.Conv2d(
|
||||
in_shape[1],
|
||||
out_shape2,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=False,
|
||||
groups=groups,
|
||||
)
|
||||
scratch.layer3_rn = nn.Conv2d(
|
||||
in_shape[2],
|
||||
out_shape3,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=False,
|
||||
groups=groups,
|
||||
)
|
||||
if len(in_shape) >= 4:
|
||||
scratch.layer4_rn = nn.Conv2d(
|
||||
in_shape[3],
|
||||
out_shape4,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=False,
|
||||
groups=groups,
|
||||
)
|
||||
|
||||
return scratch
|
||||
|
||||
|
||||
class ResidualConvUnit(nn.Module):
|
||||
"""Residual convolution module."""
|
||||
|
||||
def __init__(self, features, activation, bn):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.bn = bn
|
||||
|
||||
self.groups = 1
|
||||
|
||||
self.conv1 = nn.Conv2d(
|
||||
features,
|
||||
features,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=True,
|
||||
groups=self.groups,
|
||||
)
|
||||
|
||||
self.conv2 = nn.Conv2d(
|
||||
features,
|
||||
features,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=True,
|
||||
groups=self.groups,
|
||||
)
|
||||
|
||||
if self.bn == True:
|
||||
self.bn1 = nn.BatchNorm2d(features)
|
||||
self.bn2 = nn.BatchNorm2d(features)
|
||||
|
||||
self.activation = activation
|
||||
|
||||
self.skip_add = nn.quantized.FloatFunctional()
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass.
|
||||
|
||||
Args:
|
||||
x (tensor): input
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
|
||||
out = self.activation(x)
|
||||
out = self.conv1(out)
|
||||
if self.bn == True:
|
||||
out = self.bn1(out)
|
||||
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
if self.bn == True:
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.groups > 1:
|
||||
out = self.conv_merge(out)
|
||||
|
||||
return self.skip_add.add(out, x)
|
||||
|
||||
|
||||
class FeatureFusionBlock(nn.Module):
|
||||
"""Feature fusion block."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
features,
|
||||
activation,
|
||||
deconv=False,
|
||||
bn=False,
|
||||
expand=False,
|
||||
align_corners=True,
|
||||
size=None,
|
||||
):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super(FeatureFusionBlock, self).__init__()
|
||||
|
||||
self.deconv = deconv
|
||||
self.align_corners = align_corners
|
||||
|
||||
self.groups = 1
|
||||
|
||||
self.expand = expand
|
||||
out_features = features
|
||||
if self.expand == True:
|
||||
out_features = features // 2
|
||||
|
||||
self.out_conv = nn.Conv2d(
|
||||
features,
|
||||
out_features,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
bias=True,
|
||||
groups=1,
|
||||
)
|
||||
|
||||
self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
|
||||
self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
|
||||
|
||||
self.skip_add = nn.quantized.FloatFunctional()
|
||||
|
||||
self.size = size
|
||||
|
||||
def forward(self, *xs, size=None):
|
||||
"""Forward pass.
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
output = xs[0]
|
||||
|
||||
if len(xs) == 2:
|
||||
res = self.resConfUnit1(xs[1])
|
||||
output = self.skip_add.add(output, res)
|
||||
|
||||
output = self.resConfUnit2(output)
|
||||
|
||||
if (size is None) and (self.size is None):
|
||||
modifier = {"scale_factor": 2}
|
||||
elif size is None:
|
||||
modifier = {"size": self.size}
|
||||
else:
|
||||
modifier = {"size": size}
|
||||
|
||||
output = nn.functional.interpolate(output, **modifier, mode="bilinear", align_corners=self.align_corners)
|
||||
|
||||
output = self.out_conv(output)
|
||||
|
||||
return output
|
||||
+167
@@ -0,0 +1,167 @@
|
||||
# This file includes code originally from the Depth Anything V2 repository:
|
||||
# https://github.com/DepthAnything/Depth-Anything-V2
|
||||
# Licensed under the Apache License, Version 2.0. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Resize(object):
|
||||
"""Resize sample to given size (width, height)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
width,
|
||||
height,
|
||||
resize_target=True,
|
||||
keep_aspect_ratio=False,
|
||||
ensure_multiple_of=1,
|
||||
resize_method="lower_bound",
|
||||
image_interpolation_method=cv2.INTER_AREA,
|
||||
):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
width (int): desired output width
|
||||
height (int): desired output height
|
||||
resize_target (bool, optional):
|
||||
True: Resize the full sample (image, mask, target).
|
||||
False: Resize image only.
|
||||
Defaults to True.
|
||||
keep_aspect_ratio (bool, optional):
|
||||
True: Keep the aspect ratio of the input sample.
|
||||
Output sample might not have the given width and height, and
|
||||
resize behaviour depends on the parameter 'resize_method'.
|
||||
Defaults to False.
|
||||
ensure_multiple_of (int, optional):
|
||||
Output width and height is constrained to be multiple of this parameter.
|
||||
Defaults to 1.
|
||||
resize_method (str, optional):
|
||||
"lower_bound": Output will be at least as large as the given size.
|
||||
"upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
|
||||
"minimal": Scale as least as possible. (Output size might be smaller than given size.)
|
||||
Defaults to "lower_bound".
|
||||
"""
|
||||
self.__width = width
|
||||
self.__height = height
|
||||
|
||||
self.__resize_target = resize_target
|
||||
self.__keep_aspect_ratio = keep_aspect_ratio
|
||||
self.__multiple_of = ensure_multiple_of
|
||||
self.__resize_method = resize_method
|
||||
self.__image_interpolation_method = image_interpolation_method
|
||||
|
||||
def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
|
||||
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
if max_val is not None and y > max_val:
|
||||
y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
if y < min_val:
|
||||
y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
return y
|
||||
|
||||
def get_size(self, width, height):
|
||||
# determine new height and width
|
||||
scale_height = self.__height / height
|
||||
scale_width = self.__width / width
|
||||
|
||||
if self.__keep_aspect_ratio:
|
||||
if self.__resize_method == "lower_bound":
|
||||
# scale such that output size is lower bound
|
||||
if scale_width > scale_height:
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
elif self.__resize_method == "upper_bound":
|
||||
# scale such that output size is upper bound
|
||||
if scale_width < scale_height:
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
elif self.__resize_method == "minimal":
|
||||
# scale as least as possbile
|
||||
if abs(1 - scale_width) < abs(1 - scale_height):
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
else:
|
||||
raise ValueError(f"resize_method {self.__resize_method} not implemented")
|
||||
|
||||
if self.__resize_method == "lower_bound":
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
|
||||
elif self.__resize_method == "upper_bound":
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
|
||||
elif self.__resize_method == "minimal":
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width)
|
||||
else:
|
||||
raise ValueError(f"resize_method {self.__resize_method} not implemented")
|
||||
|
||||
return (new_width, new_height)
|
||||
|
||||
def __call__(self, sample):
|
||||
width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
|
||||
|
||||
# resize sample
|
||||
sample["image"] = cv2.resize(
|
||||
sample["image"],
|
||||
(width, height),
|
||||
interpolation=self.__image_interpolation_method,
|
||||
)
|
||||
|
||||
if self.__resize_target:
|
||||
if "depth" in sample:
|
||||
sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
if "mask" in sample:
|
||||
sample["mask"] = cv2.resize(
|
||||
sample["mask"].astype(np.float32),
|
||||
(width, height),
|
||||
interpolation=cv2.INTER_NEAREST,
|
||||
)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class NormalizeImage(object):
|
||||
"""Normlize image by given mean and std."""
|
||||
|
||||
def __init__(self, mean, std):
|
||||
self.__mean = mean
|
||||
self.__std = std
|
||||
|
||||
def __call__(self, sample):
|
||||
sample["image"] = (sample["image"] - self.__mean) / self.__std
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class PrepareForNet(object):
|
||||
"""Prepare sample for usage as network input."""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, sample):
|
||||
image = np.transpose(sample["image"], (2, 0, 1))
|
||||
sample["image"] = np.ascontiguousarray(image).astype(np.float32)
|
||||
|
||||
if "depth" in sample:
|
||||
depth = sample["depth"].astype(np.float32)
|
||||
sample["depth"] = np.ascontiguousarray(depth)
|
||||
|
||||
if "mask" in sample:
|
||||
sample["mask"] = sample["mask"].astype(np.float32)
|
||||
sample["mask"] = np.ascontiguousarray(sample["mask"])
|
||||
|
||||
return sample
|
||||
+165
@@ -0,0 +1,165 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
from vipe.utils.misc import unpack_optional
|
||||
|
||||
from ..base import DepthEstimationInput, DepthEstimationModel, DepthEstimationResult, DepthType
|
||||
from .model_fn import (
|
||||
metric3d_convnext_large,
|
||||
metric3d_convnext_tiny,
|
||||
metric3d_vit_giant2,
|
||||
metric3d_vit_large,
|
||||
metric3d_vit_small,
|
||||
)
|
||||
|
||||
|
||||
class Metric3DDepthModel(DepthEstimationModel):
|
||||
"""
|
||||
https://github.com/YvanYin/Metric3D
|
||||
"""
|
||||
|
||||
@dataclass
|
||||
class PadInfo:
|
||||
pad_up: int
|
||||
pad_down: int
|
||||
pad_left: int
|
||||
pad_right: int
|
||||
size_origin: tuple[int, int]
|
||||
focal_length: float
|
||||
|
||||
def __init__(self, version: int = 2, model: str = "giant2") -> None:
|
||||
super().__init__()
|
||||
|
||||
if version == 1:
|
||||
assert model in ["tiny", "large"]
|
||||
if model == "tiny":
|
||||
self.model = metric3d_convnext_tiny(pretrain=True, skip_validation=True)
|
||||
else:
|
||||
self.model = metric3d_convnext_large(pretrain=True, skip_validation=True)
|
||||
|
||||
self.input_size = (544, 1216)
|
||||
|
||||
elif version == 2:
|
||||
assert model in ["small", "large", "giant2"]
|
||||
if model == "small":
|
||||
self.model = metric3d_vit_small(pretrain=True, skip_validation=True)
|
||||
elif model == "large":
|
||||
self.model = metric3d_vit_large(pretrain=True, skip_validation=True)
|
||||
else:
|
||||
self.model = metric3d_vit_giant2(pretrain=True, skip_validation=True)
|
||||
|
||||
self.input_size = (616, 1064)
|
||||
|
||||
else:
|
||||
raise ValueError("Invalid version number.")
|
||||
|
||||
self.model.cuda().eval()
|
||||
|
||||
@property
|
||||
def depth_type(self) -> DepthType:
|
||||
return DepthType.METRIC_DEPTH
|
||||
|
||||
@torch.no_grad()
|
||||
def _prepare_input(self, img: torch.Tensor, focal_length: float) -> tuple[torch.Tensor, PadInfo]:
|
||||
"""
|
||||
Convert a 0-1 RGB image to normalized and padded image required by the model.
|
||||
input: (B, H, W, 3) RGB image
|
||||
return: (B, 3, H', W') tensor and pad_info
|
||||
"""
|
||||
|
||||
h_origin, w_origin = img.shape[1:3]
|
||||
scale = min(self.input_size[0] / h_origin, self.input_size[1] / w_origin)
|
||||
h_scaled, w_scaled = int(h_origin * scale), int(w_origin * scale)
|
||||
|
||||
img = img.permute(0, 3, 1, 2)
|
||||
img = torch.nn.functional.interpolate(img, (h_scaled, w_scaled), mode="bilinear")
|
||||
|
||||
# normalize
|
||||
mean = (
|
||||
torch.tensor([123.675, 116.28, 103.53], device=img.device, dtype=torch.float32)[None, :, None, None] / 255
|
||||
)
|
||||
std = torch.tensor([58.395, 57.12, 57.375], device=img.device, dtype=torch.float32)[None, :, None, None] / 255
|
||||
img = torch.div((img - mean), std)
|
||||
|
||||
# padding to input_size
|
||||
pad_h = self.input_size[0] - h_scaled
|
||||
pad_w = self.input_size[1] - w_scaled
|
||||
|
||||
pad_up, pad_left = pad_h // 2, pad_w // 2
|
||||
pad_down, pad_right = pad_h - pad_up, pad_w - pad_left
|
||||
|
||||
img = torch.nn.functional.pad(
|
||||
img,
|
||||
(pad_left, pad_right, pad_up, pad_down),
|
||||
value=0.0,
|
||||
)
|
||||
|
||||
return img, self.PadInfo(
|
||||
pad_up,
|
||||
pad_down,
|
||||
pad_left,
|
||||
pad_right,
|
||||
(h_origin, w_origin),
|
||||
focal_length * scale,
|
||||
)
|
||||
|
||||
def _post_process(self, feature: torch.Tensor, pad_info: PadInfo, is_depth: bool) -> torch.Tensor:
|
||||
"""
|
||||
Post-process the feature map to remove padding, and de-canonicalize the depth.
|
||||
input: (B, C, H', W') tensor
|
||||
return: (B, C, H, W) tensor
|
||||
"""
|
||||
|
||||
feature = feature[
|
||||
...,
|
||||
pad_info.pad_up : feature.size(-2) - pad_info.pad_down,
|
||||
pad_info.pad_left : feature.size(-1) - pad_info.pad_right,
|
||||
]
|
||||
feature = torch.nn.functional.interpolate(feature, pad_info.size_origin, mode="bilinear")
|
||||
|
||||
if is_depth:
|
||||
canonical_to_real_scale = pad_info.focal_length / 1000.0
|
||||
feature = feature * canonical_to_real_scale
|
||||
|
||||
return feature
|
||||
|
||||
def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
|
||||
rgb: torch.Tensor = unpack_optional(src.rgb)
|
||||
assert rgb.dtype == torch.float32, "Input image should be float32"
|
||||
|
||||
focal_length: float = unpack_optional(src.focal_length)
|
||||
|
||||
if rgb.dim() == 3:
|
||||
rgb, batch_dim = rgb[None], False
|
||||
else:
|
||||
batch_dim = True
|
||||
|
||||
rgb, pad_info = self._prepare_input(rgb, focal_length)
|
||||
pred_depth, confidence, output_dict = self.model.inference({"input": rgb})
|
||||
|
||||
pred_depth = self._post_process(pred_depth, pad_info, is_depth=True)
|
||||
confidence = self._post_process(confidence, pad_info, is_depth=False)
|
||||
|
||||
if not batch_dim:
|
||||
pred_depth, confidence = pred_depth[0], confidence[0]
|
||||
|
||||
return DepthEstimationResult(
|
||||
metric_depth=pred_depth[0],
|
||||
confidence=confidence[0],
|
||||
)
|
||||
+3
@@ -0,0 +1,3 @@
|
||||
# This file includes code originally from the Metric3D repository:
|
||||
# https://github.com/YvanYin/Metric3D
|
||||
# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
Vendored
+310
@@ -0,0 +1,310 @@
|
||||
# This file includes code originally from the Metric3D repository:
|
||||
# https://github.com/YvanYin/Metric3D
|
||||
# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from timm.models.layers import DropPath, trunc_normal_
|
||||
from timm.models.registry import register_model
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
r"""ConvNeXt Block. There are two equivalent implementations:
|
||||
(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
|
||||
(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
|
||||
We use (2) as we find it slightly faster in PyTorch
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
drop_path (float): Stochastic depth rate. Default: 0.0
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
"""
|
||||
|
||||
def __init__(self, dim, drop_path=0.0, layer_scale_init_value=1e-6):
|
||||
super().__init__()
|
||||
self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
|
||||
self.norm = LayerNorm(dim, eps=1e-6)
|
||||
self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
|
||||
self.act = nn.GELU()
|
||||
self.pwconv2 = nn.Linear(4 * dim, dim)
|
||||
self.gamma = (
|
||||
nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
|
||||
if layer_scale_init_value > 0
|
||||
else None
|
||||
)
|
||||
self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
input = x
|
||||
x = self.dwconv(x)
|
||||
x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
|
||||
x = self.norm(x)
|
||||
x = self.pwconv1(x)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x)
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
|
||||
|
||||
x = input + self.drop_path(x)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXt(nn.Module):
|
||||
r"""ConvNeXt
|
||||
A PyTorch impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
Args:
|
||||
in_chans (int): Number of input image channels. Default: 3
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_chans=3,
|
||||
num_classes=1000,
|
||||
depths=[3, 3, 9, 3],
|
||||
dims=[96, 192, 384, 768],
|
||||
drop_path_rate=0.0,
|
||||
layer_scale_init_value=1e-6,
|
||||
head_init_scale=1.0,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
|
||||
stem = nn.Sequential(
|
||||
nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
|
||||
LayerNorm(dims[0], eps=1e-6, data_format="channels_first"),
|
||||
)
|
||||
self.downsample_layers.append(stem)
|
||||
for i in range(3):
|
||||
downsample_layer = nn.Sequential(
|
||||
LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
|
||||
nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
|
||||
)
|
||||
self.downsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(4):
|
||||
stage = nn.Sequential(
|
||||
*[
|
||||
Block(
|
||||
dim=dims[i],
|
||||
drop_path=dp_rates[cur + j],
|
||||
layer_scale_init_value=layer_scale_init_value,
|
||||
)
|
||||
for j in range(depths[i])
|
||||
]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
# self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer
|
||||
# self.head = nn.Linear(dims[-1], num_classes)
|
||||
|
||||
self.apply(self._init_weights)
|
||||
# self.head.weight.data.mul_(head_init_scale)
|
||||
# self.head.bias.data.mul_(head_init_scale)
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
||||
trunc_normal_(m.weight, std=0.02)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward_features(self, x):
|
||||
features = []
|
||||
for i in range(4):
|
||||
x = self.downsample_layers[i](x)
|
||||
x = self.stages[i](x)
|
||||
features.append(x)
|
||||
return features # global average pooling, (N, C, H, W) -> (N, C)
|
||||
|
||||
def forward(self, x):
|
||||
# x = self.forward_features(x)
|
||||
# x = self.head(x)
|
||||
features = self.forward_features(x)
|
||||
return features
|
||||
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
|
||||
The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
|
||||
shape (batch_size, height, width, channels) while channels_first corresponds to inputs
|
||||
with shape (batch_size, channels, height, width).
|
||||
"""
|
||||
|
||||
def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(normalized_shape))
|
||||
self.bias = nn.Parameter(torch.zeros(normalized_shape))
|
||||
self.eps = eps
|
||||
self.data_format = data_format
|
||||
if self.data_format not in ["channels_last", "channels_first"]:
|
||||
raise NotImplementedError
|
||||
self.normalized_shape = (normalized_shape,)
|
||||
|
||||
def forward(self, x):
|
||||
if self.data_format == "channels_last":
|
||||
return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
|
||||
elif self.data_format == "channels_first":
|
||||
u = x.mean(1, keepdim=True)
|
||||
s = (x - u).pow(2).mean(1, keepdim=True)
|
||||
x = (x - u) / torch.sqrt(s + self.eps)
|
||||
x = self.weight[:, None, None] * x + self.bias[:, None, None]
|
||||
return x
|
||||
|
||||
|
||||
model_urls = {
|
||||
"convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
|
||||
"convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
|
||||
"convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
|
||||
"convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
|
||||
"convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
|
||||
"convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
|
||||
"convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
|
||||
"convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
|
||||
"convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
|
||||
}
|
||||
|
||||
|
||||
def convnext_tiny(pretrained=True, in_22k=False, **kwargs):
|
||||
model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
|
||||
if pretrained:
|
||||
checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
|
||||
# url = model_urls['convnext_tiny_22k'] if in_22k else model_urls['convnext_tiny_1k']
|
||||
# checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
|
||||
model_dict = model.state_dict()
|
||||
pretrained_dict = {}
|
||||
unmatched_pretrained_dict = {}
|
||||
for k, v in checkpoint["model"].items():
|
||||
if k in model_dict:
|
||||
pretrained_dict[k] = v
|
||||
else:
|
||||
unmatched_pretrained_dict[k] = v
|
||||
model_dict.update(pretrained_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
print(
|
||||
"Successfully loaded pretrained %d params, and %d paras are unmatched."
|
||||
% (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
|
||||
)
|
||||
print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
|
||||
return model
|
||||
|
||||
|
||||
def convnext_small(pretrained=True, in_22k=False, **kwargs):
|
||||
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
|
||||
if pretrained:
|
||||
checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
|
||||
# url = model_urls['convnext_small_22k'] if in_22k else model_urls['convnext_small_1k']
|
||||
# checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
|
||||
model_dict = model.state_dict()
|
||||
pretrained_dict = {}
|
||||
unmatched_pretrained_dict = {}
|
||||
for k, v in checkpoint["model"].items():
|
||||
if k in model_dict:
|
||||
pretrained_dict[k] = v
|
||||
else:
|
||||
unmatched_pretrained_dict[k] = v
|
||||
model_dict.update(pretrained_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
print(
|
||||
"Successfully loaded pretrained %d params, and %d paras are unmatched."
|
||||
% (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
|
||||
)
|
||||
print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
|
||||
return model
|
||||
|
||||
|
||||
def convnext_base(pretrained=True, in_22k=False, **kwargs):
|
||||
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
|
||||
if pretrained:
|
||||
checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
|
||||
# url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
|
||||
# checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
|
||||
model_dict = model.state_dict()
|
||||
pretrained_dict = {}
|
||||
unmatched_pretrained_dict = {}
|
||||
for k, v in checkpoint["model"].items():
|
||||
if k in model_dict:
|
||||
pretrained_dict[k] = v
|
||||
else:
|
||||
unmatched_pretrained_dict[k] = v
|
||||
model_dict.update(pretrained_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
print(
|
||||
"Successfully loaded pretrained %d params, and %d paras are unmatched."
|
||||
% (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
|
||||
)
|
||||
print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
|
||||
return model
|
||||
|
||||
|
||||
def convnext_large(pretrained=True, in_22k=False, **kwargs):
|
||||
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
|
||||
if pretrained:
|
||||
checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
|
||||
# url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']
|
||||
# checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
|
||||
model_dict = model.state_dict()
|
||||
pretrained_dict = {}
|
||||
unmatched_pretrained_dict = {}
|
||||
for k, v in checkpoint["model"].items():
|
||||
if k in model_dict:
|
||||
pretrained_dict[k] = v
|
||||
else:
|
||||
unmatched_pretrained_dict[k] = v
|
||||
model_dict.update(pretrained_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
print(
|
||||
"Successfully loaded pretrained %d params, and %d paras are unmatched."
|
||||
% (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
|
||||
)
|
||||
print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
|
||||
return model
|
||||
|
||||
|
||||
def convnext_xlarge(pretrained=True, in_22k=False, **kwargs):
|
||||
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
|
||||
if pretrained:
|
||||
assert in_22k, "only ImageNet-22K pre-trained ConvNeXt-XL is available; please set in_22k=True"
|
||||
checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
|
||||
# url = model_urls['convnext_xlarge_22k']
|
||||
# checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
|
||||
model_dict = model.state_dict()
|
||||
pretrained_dict = {}
|
||||
unmatched_pretrained_dict = {}
|
||||
for k, v in checkpoint["model"].items():
|
||||
if k in model_dict:
|
||||
pretrained_dict[k] = v
|
||||
else:
|
||||
unmatched_pretrained_dict[k] = v
|
||||
model_dict.update(pretrained_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
print(
|
||||
"Successfully loaded pretrained %d params, and %d paras are unmatched."
|
||||
% (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
|
||||
)
|
||||
print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
|
||||
return model
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import torch
|
||||
|
||||
model = convnext_base(True, in_22k=False).cuda()
|
||||
|
||||
rgb = torch.rand((2, 3, 256, 256)).cuda()
|
||||
out = model(rgb)
|
||||
print(len(out))
|
||||
for i, ft in enumerate(out):
|
||||
print(i, ft.shape)
|
||||
Vendored
+1362
File diff suppressed because it is too large
Load Diff
Vendored
+18
@@ -0,0 +1,18 @@
|
||||
# This file includes code originally from the Metric3D repository:
|
||||
# https://github.com/YvanYin/Metric3D
|
||||
# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
from .ConvNeXt import convnext_base, convnext_large, convnext_small, convnext_tiny, convnext_xlarge
|
||||
from .ViT_DINO_reg import vit_giant2_reg, vit_large_reg, vit_small_reg
|
||||
|
||||
|
||||
__all__ = [
|
||||
"convnext_xlarge",
|
||||
"convnext_small",
|
||||
"convnext_base",
|
||||
"convnext_large",
|
||||
"convnext_tiny",
|
||||
"vit_small_reg",
|
||||
"vit_large_reg",
|
||||
"vit_giant2_reg",
|
||||
]
|
||||
+1
@@ -0,0 +1 @@
|
||||
# Configuration files for metric3d models
|
||||
Vendored
+53
@@ -0,0 +1,53 @@
|
||||
# This file includes code originally from the Metric3D repository:
|
||||
# https://github.com/YvanYin/Metric3D
|
||||
# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
|
||||
config = DictConfig(
|
||||
{
|
||||
"model": {
|
||||
"backbone": {
|
||||
"type": "convnext_large",
|
||||
"pretrained": False,
|
||||
"in_22k": True,
|
||||
"out_indices": [0, 1, 2, 3],
|
||||
"drop_path_rate": 0.4,
|
||||
"layer_scale_init_value": 1.0,
|
||||
"checkpoint": "data/pretrained_weight_repo/convnext/convnext_large_22k_1k_384.pth",
|
||||
"prefix": "backbones.",
|
||||
"out_channels": [192, 384, 768, 1536],
|
||||
},
|
||||
"type": "DensePredModel",
|
||||
"decode_head": {
|
||||
"type": "HourglassDecoder",
|
||||
"in_channels": [192, 384, 768, 1536],
|
||||
"decoder_channel": [128, 128, 256, 512],
|
||||
"prefix": "decode_heads.",
|
||||
},
|
||||
},
|
||||
"data_basic": {
|
||||
"canonical_space": {"img_size": (512, 960), "focal_length": 1000.0},
|
||||
"depth_range": (0, 1),
|
||||
"depth_normalize": (0.3, 150),
|
||||
"crop_size": (544, 1216),
|
||||
"clip_depth_range": (0.9, 150),
|
||||
},
|
||||
"load_from": None,
|
||||
"cudnn_benchmark": True,
|
||||
"test_metrics": [
|
||||
"abs_rel",
|
||||
"rmse",
|
||||
"silog",
|
||||
"delta1",
|
||||
"delta2",
|
||||
"delta3",
|
||||
"rmse_log",
|
||||
"log10",
|
||||
"sq_rel",
|
||||
],
|
||||
"batchsize_per_gpu": 2,
|
||||
"thread_per_gpu": 4,
|
||||
}
|
||||
)
|
||||
Vendored
+53
@@ -0,0 +1,53 @@
|
||||
# This file includes code originally from the Metric3D repository:
|
||||
# https://github.com/YvanYin/Metric3D
|
||||
# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
|
||||
config = DictConfig(
|
||||
{
|
||||
"model": {
|
||||
"backbone": {
|
||||
"type": "convnext_tiny",
|
||||
"pretrained": False,
|
||||
"in_22k": True,
|
||||
"out_indices": [0, 1, 2, 3],
|
||||
"drop_path_rate": 0.4,
|
||||
"layer_scale_init_value": 1.0,
|
||||
"checkpoint": "",
|
||||
"prefix": "backbones.",
|
||||
"out_channels": [96, 192, 384, 768],
|
||||
},
|
||||
"type": "DensePredModel",
|
||||
"decode_head": {
|
||||
"type": "HourglassDecoder",
|
||||
"in_channels": [96, 192, 384, 768],
|
||||
"decoder_channel": [64, 64, 128, 256],
|
||||
"prefix": "decode_heads.",
|
||||
},
|
||||
},
|
||||
"data_basic": {
|
||||
"canonical_space": {"img_size": (512, 960), "focal_length": 1000.0},
|
||||
"depth_range": (0, 1),
|
||||
"depth_normalize": (0.3, 150),
|
||||
"crop_size": (544, 1216),
|
||||
"clip_depth_range": (0.9, 150),
|
||||
},
|
||||
"load_from": None,
|
||||
"cudnn_benchmark": True,
|
||||
"test_metrics": [
|
||||
"abs_rel",
|
||||
"rmse",
|
||||
"silog",
|
||||
"delta1",
|
||||
"delta2",
|
||||
"delta3",
|
||||
"rmse_log",
|
||||
"log10",
|
||||
"sq_rel",
|
||||
],
|
||||
"batchsize_per_gpu": 2,
|
||||
"thread_per_gpu": 4,
|
||||
}
|
||||
)
|
||||
Vendored
+60
@@ -0,0 +1,60 @@
|
||||
# This file includes code originally from the Metric3D repository:
|
||||
# https://github.com/YvanYin/Metric3D
|
||||
# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
|
||||
config = DictConfig(
|
||||
{
|
||||
"model": {
|
||||
"backbone": {
|
||||
"type": "vit_giant2_reg",
|
||||
"prefix": "backbones.",
|
||||
"out_channels": [1536, 1536, 1536, 1536],
|
||||
"drop_path_rate": 0.0,
|
||||
},
|
||||
"type": "DensePredModel",
|
||||
"decode_head": {
|
||||
"type": "RAFTDepthNormalDPT5",
|
||||
"in_channels": [1536, 1536, 1536, 1536],
|
||||
"use_cls_token": True,
|
||||
"feature_channels": [384, 768, 1536, 1536],
|
||||
"decoder_channels": [192, 384, 768, 1536, 1536],
|
||||
"up_scale": 7,
|
||||
"hidden_channels": [192, 192, 192, 192],
|
||||
"n_gru_layers": 3,
|
||||
"n_downsample": 2,
|
||||
"iters": 8,
|
||||
"slow_fast_gru": True,
|
||||
"num_register_tokens": 4,
|
||||
"prefix": "decode_heads.",
|
||||
"detach": False,
|
||||
},
|
||||
},
|
||||
"data_basic": {
|
||||
"canonical_space": {"img_size": (540, 960), "focal_length": 1000.0},
|
||||
"depth_range": (0, 1),
|
||||
"depth_normalize": (0.1, 200),
|
||||
"crop_size": (616, 1064),
|
||||
"clip_depth_range": (0.1, 200),
|
||||
"vit_size": (616, 1064),
|
||||
},
|
||||
"load_from": None,
|
||||
"cudnn_benchmark": True,
|
||||
"test_metrics": [
|
||||
"abs_rel",
|
||||
"rmse",
|
||||
"silog",
|
||||
"delta1",
|
||||
"delta2",
|
||||
"delta3",
|
||||
"rmse_log",
|
||||
"log10",
|
||||
"sq_rel",
|
||||
],
|
||||
"max_value": 200,
|
||||
"batchsize_per_gpu": 1,
|
||||
"thread_per_gpu": 1,
|
||||
}
|
||||
)
|
||||
Vendored
+60
@@ -0,0 +1,60 @@
|
||||
# This file includes code originally from the Metric3D repository:
|
||||
# https://github.com/YvanYin/Metric3D
|
||||
# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
|
||||
config = DictConfig(
|
||||
{
|
||||
"model": {
|
||||
"backbone": {
|
||||
"type": "vit_large_reg",
|
||||
"prefix": "backbones.",
|
||||
"out_channels": [1024, 1024, 1024, 1024],
|
||||
"drop_path_rate": 0.0,
|
||||
},
|
||||
"type": "DensePredModel",
|
||||
"decode_head": {
|
||||
"type": "RAFTDepthNormalDPT5",
|
||||
"in_channels": [1024, 1024, 1024, 1024],
|
||||
"use_cls_token": True,
|
||||
"feature_channels": [256, 512, 1024, 1024],
|
||||
"decoder_channels": [128, 256, 512, 1024, 1024],
|
||||
"up_scale": 7,
|
||||
"hidden_channels": [128, 128, 128, 128],
|
||||
"n_gru_layers": 3,
|
||||
"n_downsample": 2,
|
||||
"iters": 8,
|
||||
"slow_fast_gru": True,
|
||||
"num_register_tokens": 4,
|
||||
"prefix": "decode_heads.",
|
||||
"detach": False,
|
||||
},
|
||||
},
|
||||
"data_basic": {
|
||||
"canonical_space": {"img_size": (540, 960), "focal_length": 1000.0},
|
||||
"depth_range": (0, 1),
|
||||
"depth_normalize": (0.1, 200),
|
||||
"crop_size": (616, 1064),
|
||||
"clip_depth_range": (0.1, 200),
|
||||
"vit_size": (616, 1064),
|
||||
},
|
||||
"load_from": None,
|
||||
"cudnn_benchmark": True,
|
||||
"test_metrics": [
|
||||
"abs_rel",
|
||||
"rmse",
|
||||
"silog",
|
||||
"delta1",
|
||||
"delta2",
|
||||
"delta3",
|
||||
"rmse_log",
|
||||
"log10",
|
||||
"sq_rel",
|
||||
],
|
||||
"max_value": 200,
|
||||
"batchsize_per_gpu": 1,
|
||||
"thread_per_gpu": 1,
|
||||
}
|
||||
)
|
||||
Vendored
+60
@@ -0,0 +1,60 @@
|
||||
# This file includes code originally from the Metric3D repository:
|
||||
# https://github.com/YvanYin/Metric3D
|
||||
# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
|
||||
config = DictConfig(
|
||||
{
|
||||
"model": {
|
||||
"backbone": {
|
||||
"type": "vit_small_reg",
|
||||
"prefix": "backbones.",
|
||||
"out_channels": [384, 384, 384, 384],
|
||||
"drop_path_rate": 0.0,
|
||||
},
|
||||
"type": "DensePredModel",
|
||||
"decode_head": {
|
||||
"type": "RAFTDepthNormalDPT5",
|
||||
"in_channels": [384, 384, 384, 384],
|
||||
"use_cls_token": True,
|
||||
"feature_channels": [96, 192, 384, 768],
|
||||
"decoder_channels": [48, 96, 192, 384, 384],
|
||||
"up_scale": 7,
|
||||
"hidden_channels": [48, 48, 48, 48],
|
||||
"n_gru_layers": 3,
|
||||
"n_downsample": 2,
|
||||
"iters": 4,
|
||||
"slow_fast_gru": True,
|
||||
"num_register_tokens": 4,
|
||||
"prefix": "decode_heads.",
|
||||
"detach": False,
|
||||
},
|
||||
},
|
||||
"data_basic": {
|
||||
"canonical_space": {"img_size": (540, 960), "focal_length": 1000.0},
|
||||
"depth_range": (0, 1),
|
||||
"depth_normalize": (0.1, 200),
|
||||
"crop_size": (616, 1064),
|
||||
"clip_depth_range": (0.1, 200),
|
||||
"vit_size": (616, 1064),
|
||||
},
|
||||
"load_from": None,
|
||||
"cudnn_benchmark": True,
|
||||
"test_metrics": [
|
||||
"abs_rel",
|
||||
"rmse",
|
||||
"silog",
|
||||
"delta1",
|
||||
"delta2",
|
||||
"delta3",
|
||||
"rmse_log",
|
||||
"log10",
|
||||
"sq_rel",
|
||||
],
|
||||
"max_value": 200,
|
||||
"batchsize_per_gpu": 1,
|
||||
"thread_per_gpu": 1,
|
||||
}
|
||||
)
|
||||
+306
@@ -0,0 +1,306 @@
|
||||
# This file includes code originally from the Metric3D repository:
|
||||
# https://github.com/YvanYin/Metric3D
|
||||
# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def compute_depth_expectation(prob, depth_values):
|
||||
depth_values = depth_values.view(*depth_values.shape, 1, 1)
|
||||
depth = torch.sum(prob * depth_values, 1)
|
||||
return depth
|
||||
|
||||
|
||||
class ConvBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size=3):
|
||||
super(ConvBlock, self).__init__()
|
||||
|
||||
if kernel_size == 3:
|
||||
self.conv = nn.Sequential(
|
||||
nn.ReflectionPad2d(1),
|
||||
nn.Conv2d(in_channels, out_channels, 3, padding=0, stride=1),
|
||||
)
|
||||
elif kernel_size == 1:
|
||||
self.conv = nn.Conv2d(int(in_channels), int(out_channels), 1, padding=0, stride=1)
|
||||
|
||||
self.nonlin = nn.ELU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv(x)
|
||||
out = self.nonlin(out)
|
||||
return out
|
||||
|
||||
|
||||
class ConvBlock_double(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size=3):
|
||||
super(ConvBlock_double, self).__init__()
|
||||
|
||||
if kernel_size == 3:
|
||||
self.conv = nn.Sequential(
|
||||
nn.ReflectionPad2d(1),
|
||||
nn.Conv2d(in_channels, out_channels, 3, padding=0, stride=1),
|
||||
)
|
||||
elif kernel_size == 1:
|
||||
self.conv = nn.Conv2d(int(in_channels), int(out_channels), 1, padding=0, stride=1)
|
||||
|
||||
self.nonlin = nn.ELU(inplace=True)
|
||||
self.conv_2 = nn.Conv2d(out_channels, out_channels, 1, padding=0, stride=1)
|
||||
self.nonlin_2 = nn.ELU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv(x)
|
||||
out = self.nonlin(out)
|
||||
out = self.conv_2(out)
|
||||
out = self.nonlin_2(out)
|
||||
return out
|
||||
|
||||
|
||||
class DecoderFeature(nn.Module):
|
||||
def __init__(self, feat_channels, num_ch_dec=[64, 64, 128, 256]):
|
||||
super(DecoderFeature, self).__init__()
|
||||
self.num_ch_dec = num_ch_dec
|
||||
self.feat_channels = feat_channels
|
||||
|
||||
self.upconv_3_0 = ConvBlock(self.feat_channels[3], self.num_ch_dec[3], kernel_size=1)
|
||||
self.upconv_3_1 = ConvBlock_double(
|
||||
self.feat_channels[2] + self.num_ch_dec[3],
|
||||
self.num_ch_dec[3],
|
||||
kernel_size=1,
|
||||
)
|
||||
|
||||
self.upconv_2_0 = ConvBlock(self.num_ch_dec[3], self.num_ch_dec[2], kernel_size=3)
|
||||
self.upconv_2_1 = ConvBlock_double(
|
||||
self.feat_channels[1] + self.num_ch_dec[2],
|
||||
self.num_ch_dec[2],
|
||||
kernel_size=3,
|
||||
)
|
||||
|
||||
self.upconv_1_0 = ConvBlock(self.num_ch_dec[2], self.num_ch_dec[1], kernel_size=3)
|
||||
self.upconv_1_1 = ConvBlock_double(
|
||||
self.feat_channels[0] + self.num_ch_dec[1],
|
||||
self.num_ch_dec[1],
|
||||
kernel_size=3,
|
||||
)
|
||||
self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
|
||||
|
||||
def forward(self, ref_feature):
|
||||
x = ref_feature[3]
|
||||
|
||||
x = self.upconv_3_0(x)
|
||||
x = torch.cat((self.upsample(x), ref_feature[2]), 1)
|
||||
x = self.upconv_3_1(x)
|
||||
|
||||
x = self.upconv_2_0(x)
|
||||
x = torch.cat((self.upsample(x), ref_feature[1]), 1)
|
||||
x = self.upconv_2_1(x)
|
||||
|
||||
x = self.upconv_1_0(x)
|
||||
x = torch.cat((self.upsample(x), ref_feature[0]), 1)
|
||||
x = self.upconv_1_1(x)
|
||||
return x
|
||||
|
||||
|
||||
class UNet(nn.Module):
|
||||
def __init__(self, inp_ch=32, output_chal=1, down_sample_times=3, channel_mode="v0"):
|
||||
super(UNet, self).__init__()
|
||||
basic_block = ConvBnReLU
|
||||
num_depth = 128
|
||||
|
||||
self.conv0 = basic_block(inp_ch, num_depth)
|
||||
if channel_mode == "v0":
|
||||
channels = [
|
||||
num_depth,
|
||||
num_depth // 2,
|
||||
num_depth // 4,
|
||||
num_depth // 8,
|
||||
num_depth // 8,
|
||||
]
|
||||
elif channel_mode == "v1":
|
||||
channels = [
|
||||
num_depth,
|
||||
num_depth,
|
||||
num_depth,
|
||||
num_depth,
|
||||
num_depth,
|
||||
num_depth,
|
||||
]
|
||||
self.down_sample_times = down_sample_times
|
||||
for i in range(down_sample_times):
|
||||
setattr(
|
||||
self,
|
||||
"conv_%d" % i,
|
||||
nn.Sequential(
|
||||
basic_block(channels[i], channels[i + 1], stride=2),
|
||||
basic_block(channels[i + 1], channels[i + 1]),
|
||||
),
|
||||
)
|
||||
for i in range(down_sample_times - 1, -1, -1):
|
||||
setattr(
|
||||
self,
|
||||
"deconv_%d" % i,
|
||||
nn.Sequential(
|
||||
nn.ConvTranspose2d(
|
||||
channels[i + 1],
|
||||
channels[i],
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
output_padding=1,
|
||||
stride=2,
|
||||
bias=False,
|
||||
),
|
||||
nn.BatchNorm2d(channels[i]),
|
||||
nn.ReLU(inplace=True),
|
||||
),
|
||||
)
|
||||
self.prob = nn.Conv2d(num_depth, output_chal, 1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
features = {}
|
||||
conv0 = self.conv0(x)
|
||||
x = conv0
|
||||
features[0] = conv0
|
||||
for i in range(self.down_sample_times):
|
||||
x = getattr(self, "conv_%d" % i)(x)
|
||||
features[i + 1] = x
|
||||
for i in range(self.down_sample_times - 1, -1, -1):
|
||||
x = features[i] + getattr(self, "deconv_%d" % i)(x)
|
||||
x = self.prob(x)
|
||||
return x
|
||||
|
||||
|
||||
class ConvBnReLU(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
|
||||
super(ConvBnReLU, self).__init__()
|
||||
self.conv = nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
bias=False,
|
||||
)
|
||||
self.bn = nn.BatchNorm2d(out_channels)
|
||||
|
||||
def forward(self, x):
|
||||
return F.relu(self.bn(self.conv(x)), inplace=True)
|
||||
|
||||
|
||||
class HourglassDecoder(nn.Module):
|
||||
def __init__(self, cfg):
|
||||
super(HourglassDecoder, self).__init__()
|
||||
self.inchannels = cfg.model.decode_head.in_channels # [256, 512, 1024, 2048]
|
||||
self.decoder_channels = cfg.model.decode_head.decoder_channel # [64, 64, 128, 256]
|
||||
self.min_val = cfg.data_basic.depth_normalize[0]
|
||||
self.max_val = cfg.data_basic.depth_normalize[1]
|
||||
|
||||
self.num_ch_dec = self.decoder_channels # [64, 64, 128, 256]
|
||||
self.num_depth_regressor_anchor = 512
|
||||
self.feat_channels = self.inchannels
|
||||
unet_in_channel = self.num_ch_dec[1]
|
||||
unet_out_channel = 256
|
||||
|
||||
self.decoder_mono = DecoderFeature(self.feat_channels, self.num_ch_dec)
|
||||
self.conv_out_2 = UNet(
|
||||
inp_ch=unet_in_channel,
|
||||
output_chal=unet_out_channel + 1,
|
||||
down_sample_times=3,
|
||||
channel_mode="v0",
|
||||
)
|
||||
|
||||
self.depth_regressor_2 = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
unet_out_channel,
|
||||
self.num_depth_regressor_anchor,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
),
|
||||
nn.BatchNorm2d(self.num_depth_regressor_anchor),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(
|
||||
self.num_depth_regressor_anchor,
|
||||
self.num_depth_regressor_anchor,
|
||||
kernel_size=1,
|
||||
),
|
||||
)
|
||||
self.residual_channel = 16
|
||||
self.conv_up_2 = nn.Sequential(
|
||||
nn.Conv2d(1 + 2 + unet_out_channel, self.residual_channel, 3, padding=1),
|
||||
nn.BatchNorm2d(self.residual_channel),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(self.residual_channel, self.residual_channel, 3, padding=1),
|
||||
nn.Upsample(scale_factor=4),
|
||||
nn.Conv2d(self.residual_channel, self.residual_channel, 3, padding=1),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(self.residual_channel, 1, 1, padding=0),
|
||||
)
|
||||
|
||||
def get_bins(self, bins_num):
|
||||
depth_bins_vec = torch.linspace(math.log(self.min_val), math.log(self.max_val), bins_num, device="cuda")
|
||||
depth_bins_vec = torch.exp(depth_bins_vec)
|
||||
return depth_bins_vec
|
||||
|
||||
def register_depth_expectation_anchor(self, bins_num, B):
|
||||
depth_bins_vec = self.get_bins(bins_num)
|
||||
depth_bins_vec = depth_bins_vec.unsqueeze(0).repeat(B, 1)
|
||||
self.register_buffer("depth_expectation_anchor", depth_bins_vec, persistent=False)
|
||||
|
||||
def upsample(self, x, scale_factor=2):
|
||||
return F.interpolate(x, scale_factor=scale_factor, mode="nearest")
|
||||
|
||||
def regress_depth_2(self, feature_map_d):
|
||||
prob = self.depth_regressor_2(feature_map_d).softmax(dim=1)
|
||||
B = prob.shape[0]
|
||||
if "depth_expectation_anchor" not in self._buffers:
|
||||
self.register_depth_expectation_anchor(self.num_depth_regressor_anchor, B)
|
||||
d = compute_depth_expectation(prob, self.depth_expectation_anchor[:B, ...]).unsqueeze(1)
|
||||
return d
|
||||
|
||||
def create_mesh_grid(self, height, width, batch, device="cuda", set_buffer=True):
|
||||
y, x = torch.meshgrid(
|
||||
[
|
||||
torch.arange(0, height, dtype=torch.float32, device=device),
|
||||
torch.arange(0, width, dtype=torch.float32, device=device),
|
||||
],
|
||||
indexing="ij",
|
||||
)
|
||||
meshgrid = torch.stack((x, y))
|
||||
meshgrid = meshgrid.unsqueeze(0).repeat(batch, 1, 1, 1)
|
||||
return meshgrid
|
||||
|
||||
def forward(self, features_mono, **kwargs):
|
||||
"""
|
||||
trans_ref2src: list of transformation matrix from the reference view to source view. [B, 4, 4]
|
||||
inv_intrinsic_pool: list of inverse intrinsic matrix.
|
||||
features_mono: features of reference and source views. [[ref_f1, ref_f2, ref_f3, ref_f4],[src1_f1, src1_f2, src1_f3, src1_f4], ...].
|
||||
"""
|
||||
outputs = {}
|
||||
# get encoder feature of the reference view
|
||||
ref_feat = features_mono
|
||||
|
||||
feature_map_mono = self.decoder_mono(ref_feat)
|
||||
feature_map_mono_pred = self.conv_out_2(feature_map_mono)
|
||||
confidence_map_2 = feature_map_mono_pred[:, -1:, :, :]
|
||||
feature_map_d_2 = feature_map_mono_pred[:, :-1, :, :]
|
||||
|
||||
depth_pred_2 = self.regress_depth_2(feature_map_d_2)
|
||||
|
||||
B, _, H, W = depth_pred_2.shape
|
||||
|
||||
meshgrid = self.create_mesh_grid(H, W, B)
|
||||
|
||||
depth_pred_mono = self.upsample(depth_pred_2, scale_factor=4) + 1e-1 * self.conv_up_2(
|
||||
torch.cat((depth_pred_2, meshgrid[:B, ...], feature_map_d_2), 1)
|
||||
)
|
||||
confidence_map_mono = self.upsample(confidence_map_2, scale_factor=4)
|
||||
|
||||
outputs = dict(
|
||||
prediction=depth_pred_mono,
|
||||
confidence=confidence_map_mono,
|
||||
pred_logit=None,
|
||||
)
|
||||
return outputs
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user