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,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};
|
||||
}
|
||||
@@ -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};
|
||||
}
|
||||
+63
@@ -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");
|
||||
}
|
||||
Reference in New Issue
Block a user