src/caffe/layers/crop_layer.cu

   1 #include <vector>
   2
   3 #include "caffe/layers/crop_layer.hpp"
   4
   5 namespace caffe {
   6
   7 // Copy (one line per thread) from one array to another, with arbitrary
   8 // strides in the last two dimensions.
   9 template <typename Dtype>
  10 __global__ void copy_kernel(const int n, const int height, const int width,
  11     const int src_outer_stride, const int src_inner_stride,
  12     const int dest_outer_stride, const int dest_inner_stride,
  13     const Dtype* src, Dtype* dest) {
  14   CUDA_KERNEL_LOOP(index, n) {
  15     int src_start = index / height * src_outer_stride
  16                   + index % height * src_inner_stride;
  17     int dest_start = index / height * dest_outer_stride
  18                    + index % height * dest_inner_stride;
  19     for (int i = 0; i < width; ++i) {
  20       dest[dest_start + i] = src[src_start + i];
  21     }
  22   }
  23 }
  24
  25 template <typename Dtype>
  26 void CropLayer<Dtype>::crop_copy_gpu(const vector<Blob<Dtype>*>& bottom,
  27              const vector<Blob<Dtype>*>& top,
  28              const vector<int>& offsets,
  29              vector<int> indices,
  30              int cur_dim,
  31              const Dtype* src_data,
  32              Dtype* dest_data,
  33              bool is_forward) {
  34   if (cur_dim + 2 < top[0]->num_axes()) {
  35     // We are not yet at the final dimension, call copy recursivley
  36     for (int i = 0; i < top[0]->shape(cur_dim); ++i) {
  37       indices[cur_dim] = i;
  38       crop_copy_gpu(bottom, top, offsets, indices, cur_dim+1,
  39                 src_data, dest_data, is_forward);
  40     }
  41   } else {
  42     // We are at the last two dimensions, which are stored continuously in
  43     // memory With (N,C,H,W)
  44     //             (0,1,2,3) cur_dim   -> H
  45     //                       cur_dim+1 -> W
  46     const int lines = top[0]->shape(cur_dim);
  47     const int height = top[0]->shape(cur_dim);
  48     const int width = top[0]->shape(cur_dim+1);
  49     std::vector<int> ind_off(cur_dim+2, 0);
  50     for (int j = 0; j < cur_dim; ++j) {
  51         ind_off[j] = indices[j] + offsets[j];
  52     }
  53     ind_off[cur_dim] = offsets[cur_dim];
  54     ind_off[cur_dim+1] = offsets[cur_dim+1];
  55     // Compute copy strides
  56     const int src_outer_stride =
  57         bottom[0]->shape(cur_dim)*bottom[0]->shape(cur_dim+1);
  58     const int src_inner_stride = bottom[0]->shape(cur_dim+1);
  59     const int dest_outer_stride =
  60         top[0]->shape(cur_dim)*top[0]->shape(cur_dim+1);
  61     const int dest_inner_stride = top[0]->shape(cur_dim+1);
  62
  63     if (is_forward) {
  64       const Dtype* bottom_data = bottom[0]->gpu_data() +
  65           bottom[0]->offset(ind_off);
  66       Dtype* top_data = top[0]->mutable_gpu_data() +
  67           top[0]->offset(indices);
  68       // NOLINT_NEXT_LINE(whitespace/operators)
  69       copy_kernel<<<CAFFE_GET_BLOCKS(lines), CAFFE_CUDA_NUM_THREADS>>>(
  70           lines, height, width,
  71           src_outer_stride, src_inner_stride,
  72           dest_outer_stride, dest_inner_stride,
  73           bottom_data, top_data);
  74
  75     } else {
  76       const Dtype* top_diff = top[0]->gpu_diff() +
  77           top[0]->offset(indices);
  78       Dtype* bottom_diff = bottom[0]->mutable_gpu_diff() +
  79           bottom[0]->offset(ind_off);
  80       // NOLINT_NEXT_LINE(whitespace/operators)
  81       copy_kernel<<<CAFFE_GET_BLOCKS(lines), CAFFE_CUDA_NUM_THREADS>>>(
  82           lines, height, width,
  83           dest_outer_stride, dest_inner_stride,
  84           src_outer_stride, src_inner_stride,
  85           top_diff, bottom_diff);
  86     }
  87   }
  88 }
  89
  90 template <typename Dtype>
  91 void CropLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
  92     const vector<Blob<Dtype>*>& top) {
  93   std::vector<int> indices(top[0]->num_axes(), 0);
  94   const Dtype* bottom_data = bottom[0]->gpu_data();
  95   Dtype* top_data = top[0]->mutable_gpu_data();
  96   crop_copy_gpu(bottom, top, offsets, indices, 0, bottom_data, top_data, true);
  97 }
  98
  99 template <typename Dtype>
 100 void CropLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
 101     const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
 102   const Dtype* top_diff = top[0]->gpu_diff();
 103   Dtype* bottom_diff = bottom[0]->mutable_gpu_diff();
 104
 105   if (propagate_down[0]) {
 106     caffe_gpu_set(bottom[0]->count(), static_cast<Dtype>(0), bottom_diff);
 107     std::vector<int> indices(top[0]->num_axes(), 0);
 108     crop_copy_gpu(bottom, top, offsets, indices, 0, top_diff, bottom_diff,
 109                   false);
 110   }
 111 }
 112
 113 INSTANTIATE_LAYER_GPU_FUNCS(CropLayer);
 114
 115 }  // namespace caffe