caffe安装_CenterNet的Caffe-C++/CUDA推理实现(上篇)

import torch
import torch.onnx
import pytorch_to_caffe
 
...
 
model = create_model(opt.arch, opt.heads, opt.head_conv)
model = load_model(model, model_path)
model = model.cpu()
model.eval()
 
input = torch.ones([1, 3, 512, 512])
pytorch_to_caffe.trans_net(model, input, 'res50')
pytorch_to_caffe.save_prototxt('{}.prototxt'.format('res50'))
pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format('res50'))

# src/lib/detectors/ctdet.py
with torch.no_grad():
    output = self.model(images)[-1]
    # 首先对关键点数组做sigmoid
    hm = output['hm'].sigmoid_()
    wh = output['wh']
    reg = output['reg'] if self.opt.reg_offset else None
    ...
    # 再放入后处理步骤中对输出结果进行调整
    dets = ctdet_decode(hm, wh, reg=reg, cat_spec_wh=self.opt.cat_spec_wh, K=self.opt.K)

def _nms(heat, kernel=3):
    pad = (kernel - 1) // 2
    # hmax为最大池化后的特征图，通过padding后与输入的heat特征图大小一致
    hmax = nn.functional.max_pool2d(
        heat, (kernel, kernel), stride=1, padding=pad)
    # keep索引存放hmax于heat相同的部分，也就是判断当前点是否为最大点
    keep = (hmax == heat).float()
 
    return heat * keep

// 经过处理的hm层
  Blob<float>* hm_result_blob = net_->output_blobs()[2];
  const float* hm_result = hm_result_blob->cpu_data();
  vector<int> hm_shape = hm_result_blob->shape();
  //wh层
  Blob<float>* result_blob1 = net_->output_blobs()[0];
  const float* result1 = result_blob1->cpu_data();
  //reg层
  Blob<float>* result_blob0 = net_->output_blobs()[1];
  const float* result0 = result_blob0->cpu_data();
 
  // 这一层是没有max-pooling的heat输出 用于作对比
  boost::shared_ptr<caffe::Blob<float>> layerData = net_->blob_by_name("sigmoid_blob55");  // 获得指定层的输出
  const float* pstart = layerData->cpu_data();
 
  // 这里相当于 之前所说的keep 操作
  int classes = hm_shape[1];
  int feature_size = hm_shape[2]*hm_shape[3];
 
  vector<vector<float>> fscore_max;
  for(int j=0; j < classes; j++)   //class
  {
    for(int k=0; k<feature_size; k++)
      if(pstart[j*feature_size+k]
          == hm_result[j*feature_size+k])
      {
        vector<float> inx;
        inx.push_back(j*feature_size+k);         // 位置
        inx.push_back(pstart[j*feature_size+k]); // 分数
        fscore_max.push_back(inx);
      }
  }
 
  std::sort(fscore_max.begin(), fscore_max.end(),[](const std::vector<float>& a, const std::vector<float>& b){ return a[1] > b[1];});
  // get top 100
  int iters = std::min<int>(fscore_max.size(), 100);
  int only_threshbox=0;
  for(int i=0;i<iters;i++)
  {
    // 这里根据阈值进行筛选
    if(fscore_max[i][1]<thresh)
    {
      break;
    }
    only_threshbox++;
  }
 
  // 这里进行行了一些修改，使得对于任意尺寸输入不会出错
  vector<vector<float>> boxes;
  for(int i = 0; i < only_threshbox;i++)
  {
    vector<float> box;
    int index = ((int)fscore_max[i][0]) / (hm_shape[2] * hm_shape[3]);
    int center_index = ((int)fscore_max[i][0]) % (hm_shape[2]*hm_shape[3]) - hm_shape[3];
    int cls = index;
 
    float xs= center_index % hm_shape[3];
    float ys= int(center_index / hm_shape[3] ) % hm_shape[2];
 
    //reg batch 1
    xs += result0[(int)(((int)ys)*hm_shape[3] + xs)];
    ys += result0[(int)(hm_shape[3]*hm_shape[2]+((int)ys)*hm_shape[3]+xs)];
 
    float w = result1[(int)(((int)ys)*hm_shape[3]+xs)];
    float h = result1[(int)(hm_shape[2]*hm_shape[3]+((int)ys)*hm_shape[3]+xs)];
 
    box.push_back((float)cls);
    box.push_back((float)fscore_max[i][1]);
 
    box.push_back((float)(xs-w/2.0));
    box.push_back((float)(ys-h/2.0));
    box.push_back((float)(xs+w/2.0));
    box.push_back((float)(ys+h/2.0));
    // 输出四个点
    boxes.push_back(box);
  }

vector<vector<float> > CenterNet_Detector::apply_nms(vector<vector<float>> &box, float thres)
{
  vector<vector<float> > rlt;
  if (box.empty())
    return vector<vector<float> >();
 
  std::sort(box.begin(), box.end(),[](const std::vector<float>& a, const std::vector<float>& b){ return a[1] > b[1];});
  std::vector<int> pindex;
 
  for(int i=0;i<box.size();i++)
  {
    if(std::find(pindex.begin(),pindex.end(),i)!=pindex.end())
    {
      continue;
    }
    vector<float> truth = box[i];
    for(int j=i+1;j<box.size();j++)
    {
      if(std::find(pindex.begin(),pindex.end(),j)!=pindex.end())
      {
        continue;
      }
 
      vector<float> lbox = box[j];
      float iou = cal_iou(lbox, truth);
      if(iou >= thres)
        pindex.push_back(j);//p[j] = 1
 
    }
  }
  for(int i=0;i<box.size();i++)
  {
    if(std::find(pindex.begin(),pindex.end(),i)==pindex.end())
    {
      rlt.push_back(box[i]);
    }
  }
  return rlt;
}

__device__ float Logist(float data){ return 1./(1. + exp(-data)); }
 
__global__ void CTdetforward_kernel(const float *hm, const float *reg,const float *wh ,
    float *output,const int w,const int h,const int classes,const int kernel_size,const float visthresh  ) {
 
int idx = blockDim.x * blockIdx.x + threadIdx.x;
if (idx >= w * h * classes) return;
int padding = (kernel_size - 1) / 2;
int offset = -padding;
int stride = w * h;
int grid_x = idx % w;
int grid_y = (idx / w) % h;
int cls = idx/w/h ;
int  l, m;
int reg_index = idx - cls*stride;
float c_x, c_y;
float objProb = Logist(hm[idx]);
if (objProb > visthresh) {
    float max = -1;
    int max_index = 0;
    for (l = 0; l < kernel_size; ++l)
        for (m = 0; m < kernel_size; ++m) {
            int cur_x = offset + l + grid_x;
            int cur_y = offset + m + grid_y;
            int cur_index = cur_y * w + cur_x + stride * cls;
            int valid = (cur_x >= 0 && cur_x < w && cur_y >= 0 && cur_y < h);
            float val = (valid != 0) ? Logist(hm[cur_index]) : -1;
            max_index = (val > max) ? cur_index : max_index;
            max = (val > max) ? val : max;
        }
 
    if(idx == max_index){
        // 原子操作，在指针的第一个地址的数据上加1，用于累加挑选后的框的数量
        int resCount = (int) atomicAdd(output, 1);
        // 这里每次将获取到的box数据放入指定的位置 
        char *data = (char *) output + sizeof(float) + resCount * sizeof(Detection);
        Detection *det = (Detection *) (data);
        c_x = grid_x + reg[reg_index];
        c_y = grid_y + reg[reg_index + stride];
        det->bbox.x1 = (c_x - wh[reg_index] / 2) * 4;
        det->bbox.y1 = (c_y - wh[reg_index + stride] / 2) * 4;
        det->bbox.x2 = (c_x + wh[reg_index] / 2) * 4;
        det->bbox.y2 = (c_y + wh[reg_index + stride] / 2) * 4;
        det->classId = cls;
        det->prob = objProb;
    }
}
}
 
void CTdetforward_gpu(const float *hm, const float *reg,const float *wh ,float *output,
    const int w,const int h,const int classes,const int kernerl_size, const float visthresh ){
uint num = w * h * classes;
int nbrBlocks = ceil((float)num / (float)BLOCK);
CTdetforward_kernel<<<nbrBlocks,BLOCK>>>(hm,reg,wh,output,w,h,classes,kernerl_size,visthresh);
}

Blob<float>* input_layer = net_->input_blobs()[0];
 
input_layer->Reshape(1, num_channels_, img.rows, img.cols);
input_geometry_ = cv::Size(input_layer->width(), input_layer->height());
net_->Reshape();
 
std::vector<cv::Mat> input_channels;
WrapInputLayer(&input_channels);
cv::Mat tm = Preprocess(img, &input_channels);
 
net_->Forward();
 
Blob<float>* hm_result_blob = net_->output_blobs()[2];
const float* hm_result = hm_result_blob->cpu_data();
vector<int> hm_shape = hm_result_blob->shape();
 
Blob<float>* result_blob1 = net_->output_blobs()[0];
const float* result1 = result_blob1->cpu_data();
  vector<int> wh_shape = result_blob1->shape();
 
Blob<float>* result_blob0 = net_->output_blobs()[1];
const float* result0 = result_blob0->cpu_data();
vector<int> reg_shape = result_blob0->shape();
 
boost::shared_ptr<caffe::Blob<float>> layerData1 = net_->blob_by_name("conv_blob55");  
// 得到没有经过任何后处理的原始heat map
const float* hm_raw = layerData1->cpu_data();
 
boost::shared_ptr<caffe::Blob<float>> layerData = net_->blob_by_name("sigmoid_blob55");  // 获得指定层的输出
const float* pstart = layerData->cpu_data();
 
/*cuda part*/
std::vector<void*> mCudaBuffers(3);
std::vector<int64_t> mBindBufferSizes(3);  //  保存size
int64_t totalSize = 0;
int64_t outputBufferSize;
void * cudaOutputBuffer;
cudaEvent_t start1;
cudaEventCreate(&start1);
cudaEvent_t stop1;
cudaEventCreate(&stop1);
cudaEventRecord(start1, NULL);
 
totalSize = hm_shape[1] * hm_shape[2] * hm_shape[3] * sizeof(float);
mCudaBuffers[0] = safeCudaMalloc(totalSize);
mBindBufferSizes[0] = totalSize;
 
totalSize = wh_shape[1] * wh_shape[2] * wh_shape[3] * sizeof(float);
mCudaBuffers[1] = safeCudaMalloc(totalSize);
mBindBufferSizes[1] = totalSize;
 
totalSize = reg_shape[1] * reg_shape[2] * reg_shape[3] * sizeof(float);
mCudaBuffers[2] = safeCudaMalloc(totalSize);
mBindBufferSizes[2] = totalSize;
 
outputBufferSize = mBindBufferSizes[0] * 6 ;
cudaOutputBuffer = safeCudaMalloc(outputBufferSize);
CUDA_CHECK(cudaMemcpy(mCudaBuffers[0], hm_raw, mBindBufferSizes[0], cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(mCudaBuffers[1], result1, mBindBufferSizes[1], cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(mCudaBuffers[2], result0, mBindBufferSizes[2], cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemset(cudaOutputBuffer, 0, sizeof(float))); 
 
std::unique_ptr<float[]> outputData(new float[outputBufferSize]);
 
CTdetforward_gpu(static_cast<const float *>(mCudaBuffers[0]),static_cast<const float *>(mCudaBuffers[2]),
                        static_cast<const float *>(mCudaBuffers[1]),static_cast<float *>(cudaOutputBuffer),
                            hm_shape[3],hm_shape[2] ,2, 3, 0.2);
 
CUDA_CHECK(cudaDeviceSynchronize());
// std::cout << "cuda excute "  << std::endl;
CUDA_CHECK(cudaMemcpy(outputData.get(), cudaOutputBuffer, outputBufferSize, cudaMemcpyDeviceToHost));
int num_det = static_cast<int>(outputData[0]);
std::vector<Detection> result;
result.resize(num_det);
memcpy(result.data(), &outputData[1], num_det * sizeof(Detection));