vol2col.h 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109
  1. #pragma once
  2. #include <cstring>
  3. namespace at::native {
  4. template <typename T>
  5. void vol2col(
  6. const T* data_vol,
  7. const int64_t channels,
  8. const int64_t depth,
  9. const int64_t height,
  10. const int64_t width,
  11. const int64_t depth_col,
  12. const int64_t height_col,
  13. const int64_t width_col,
  14. const int64_t kT,
  15. const int64_t kernel_height,
  16. const int64_t kernel_width,
  17. const int64_t pT,
  18. const int64_t pH,
  19. const int64_t pW,
  20. const int64_t dT,
  21. const int64_t dH,
  22. const int64_t dW,
  23. const int64_t dilationT,
  24. const int64_t dilationH,
  25. const int64_t dilationW,
  26. T* data_col) {
  27. int64_t c, t, h, w;
  28. int64_t channels_col = channels * kT * kernel_height * kernel_width;
  29. for (c = 0; c < channels_col; ++c) {
  30. int64_t w_offset = c % kernel_width;
  31. int64_t h_offset = (c / kernel_width) % kernel_height;
  32. int64_t t_offset = (c / kernel_width / kernel_height) % kT;
  33. int64_t c_vol = c / kT / kernel_height / kernel_width;
  34. for (t = 0; t < depth_col; ++t) {
  35. int64_t t_pad = t * dT - pT + t_offset * dilationT;
  36. for (h = 0; h < height_col; ++h) {
  37. int64_t h_pad = h * dH - pH + h_offset * dilationH;
  38. for (w = 0; w < width_col; ++w) {
  39. int64_t w_pad = w * dW - pW + w_offset * dilationW;
  40. if (t_pad >= 0 && t_pad < depth && h_pad >= 0 && h_pad < height &&
  41. w_pad >= 0 && w_pad < width)
  42. data_col[((c * depth_col + t) * height_col + h) * width_col + w] =
  43. data_vol
  44. [((c_vol * depth + t_pad) * height + h_pad) * width +
  45. w_pad];
  46. else
  47. data_col[((c * depth_col + t) * height_col + h) * width_col + w] =
  48. 0;
  49. }
  50. }
  51. }
  52. }
  53. }
  54. template <typename T>
  55. void col2vol(
  56. const T* data_col,
  57. const int64_t channels,
  58. const int64_t depth,
  59. const int64_t height,
  60. const int64_t width,
  61. const int64_t out_depth,
  62. const int64_t out_height,
  63. const int64_t out_width,
  64. const int64_t kT,
  65. const int64_t kernel_height,
  66. const int64_t kernel_width,
  67. const int64_t pT,
  68. const int64_t pH,
  69. const int64_t pW,
  70. const int64_t dT,
  71. const int64_t dH,
  72. const int64_t dW,
  73. const int64_t dilationT,
  74. const int64_t dilationH,
  75. const int64_t dilationW,
  76. T* data_vol) {
  77. memset(data_vol, 0, sizeof(T) * depth * height * width * channels);
  78. int64_t depth_col = out_depth;
  79. int64_t height_col = out_height;
  80. int64_t width_col = out_width;
  81. int64_t channels_col = channels * kT * kernel_height * kernel_width;
  82. for (int64_t c = 0; c < channels_col; ++c) {
  83. int64_t w_offset = c % kernel_width;
  84. int64_t h_offset = (c / kernel_width) % kernel_height;
  85. int64_t t_offset = (c / kernel_width / kernel_height) % kT;
  86. int64_t c_vol = c / kT / kernel_height / kernel_width;
  87. for (int64_t t = 0; t < depth_col; ++t) {
  88. int64_t t_pad = t * dT - pT + t_offset * dilationT;
  89. for (int64_t h = 0; h < height_col; ++h) {
  90. int64_t h_pad = h * dH - pH + h_offset * dilationH;
  91. for (int64_t w = 0; w < width_col; ++w) {
  92. int64_t w_pad = w * dW - pW + w_offset * dilationW;
  93. if (t_pad >= 0 && t_pad < depth && h_pad >= 0 && h_pad < height &&
  94. w_pad >= 0 && w_pad < width)
  95. data_vol
  96. [((c_vol * depth + t_pad) * height + h_pad) * width + w_pad] +=
  97. data_col
  98. [((c * depth_col + t) * height_col + h) * width_col + w];
  99. }
  100. }
  101. }
  102. }
  103. }
  104. } // namespace at::native