SO-Net中的index_max的功能及具体实现

一、背景

  在SO-Net的分类模型的Encoder部分，first_pointnet给出特征向量后，有一个求索引的操作：

		M = node.size()[2]
        with torch.cuda.device(self.first_pn_out.get_device()):
            gather_index = index_max.forward_cuda(self.first_pn_out.detach(),
                                                  min_idx.int(),
                                                  M).detach().long()

这里的index_max在./models/index_max_ext文件下，是pytorch的CUDA扩展，接口为C++接口。因为没有CUDA编程基础，只好找了相关的书籍，简单了解一下基本知识，看了相关的博客和帖子，逐渐明白了这个扩展程序的功能是什么。

二、具体实现

2.1 文件结构

  实现一个pytorch的CUDA扩展，并以C++为接口，至少需要三个部分，setup.py，.cu以及.cpp文件，具体到index_max上就是setup.py，index_max.cpp和index_max_cuda.cu文件，下面一一介绍这三个文件具体的作用。

2.2 setup.py

  这个文件是为了编译后面的文件用的。其具体代码如下：

import setuptools
import torch
from setuptools import setup
from torch.utils.cpp_extension import CppExtension, CUDAExtension, BuildExtension


setup(name='index_max',  # 编译后的链接库名称
      ext_modules=[CUDAExtension('index_max', ['index_max.cpp', 'index_max_cuda.cu'])],  # 待编译文件以及编译函数
      cmdclass={
   'build_ext': BuildExtension})  # 执行编译命令设置

2.3 index_max.cpp

  在原来的文件中有些部分没有使用，就不介绍了，直接解释用到的部分。在.cpp文件中声明了.cu中用到的函数，实现函数具体功能。

#include <torch/extension.h>
#include <iostream>
#include <vector>
#include <thread>

torch::Tensor index_max_forward_cpu(const torch::Tensor data,const torch::Tensor index,const int K) 
{
   
	int B = data.size(0);
	int C = data.size(1);
	int N = data.size(2);
	
	torch::Tensor max_idx = torch::zeros({
   B, C, K}, torch::TensorOptions