Основные шаги для совместного программирования с Pytorch и CUDA

искусственный интеллект PyTorch

Ссылка на ссылку - w3school

Ссылка на ссылку - CSDN

Код этой статьи — github

задний план

В настоящее время PyTorch предоставляет множество интерфейсов, которые можно вызывать напрямую, но все еще есть некоторые сильно настраиваемые операции, которые не могут быть эффективно выполнены с помощью PyToch или Python, поэтому PyTorch также предоставляет расширенные интерфейсы для программирования на C++ и CUDA. Существует две основные формы расширений C++: одна создается заранее с помощью setuptools или может быть создана во время выполнения с помощью torch.utils.cpp_extension.load(). Ниже представлен только первый метод, а второй метод будет изучен позже.

Основные шаги

Общие шаги совместного программирования Pytorch, CUDA, C++ следующие:

  1. Сначала вам нужно определить файл C++, который объявляет функции, определенные в файле CUDA, вам также нужно выполнить некоторые проверки и, наконец, перенаправить его вызов на.cuдокумент. Кроме того, в файле необходимо объявить функции, которые будут вызываться в Python и привязываться к python с помощью pybind11. OpenPCDet делит вышеуказанные шаги на следующие шаги:

    1. Сначала определите файл заголовка, файл заголовка.hсодержит.cuфункции, определенные в файле, и.cppфункции, определенные в файле.

    2. затем определите.cppфайл, где роль функции отвечает за выполнение некоторых проверок и вызовов.cuфункции, определенные в файле

    3. .cuФайлы отвечают за выполнение определенных операций программирования CUDA.

    4. апи файл будет.cppФункции, определенные в файле, привязаны к PYBIND11, чтобы Python мог их вызывать.

  2. существуетsetup.pyВ файле объявляется имя компилируемого модуля, путь к исходному файлу и т.д.
  3. Используйте import для импорта объявленного модуля и используйте Python для реализации вычислений прямого и обратного распространения.

Например

.
├── ball_query_src
│   ├── api.cpp
│   ├── ball_query.cpp
│   ├── ball_query_cuda.cu
│   ├── ball_query_cuda.h
│   └── cuda_utils.h
├── setup.py
└── test_ball_query.py

Каталог проекта показан на рисунке выше, гдеapi.cppфайл будетball_query.cppОбъявленная функция использует PYBIND11 для привязки к python. вapi.cppСодержание следующее:

#include <torch/serialize/tensor.h>
#include <torch/extension.h>

#include "ball_query_cuda.h"


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    // 第一个参数表示的是在python中调用的名称，第二个参数是对应的cpp函数，第三个参数对应的是这个函数的说明
    m.def("ball_query_wrapper", &ball_query_wrapper_fast, "ball_query_wrapper_fast");
}

Здесь, чтобы сделать структуру кода более понятной, авторball_query.hВ файле объявлены две функции, одна из которых вызывается на C++, а другая — функция, реализованная в CUDA. Его конкретное содержание состоит в следующем:

#ifndef _BALL_QUERY_GPU_H
#define _BALL_QUERY_GPU_H

#include <torch/serialize/tensor.h>
#include <vector>
#include <cuda.h>
#include <cuda_runtime_api.h>

// 与pybind11绑定的函数，其主要作用是调用下面的cuda函数
int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, 
	at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor);

// CUDA文件中的函数
void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample, 
	const float *xyz, const float *new_xyz, int *idx);

#endif

существует.hПосле того, как два вышеуказанных файла объявлены в файле,ball_query.cppиball_query_cuda.cuКонкретная реализация этих двух функций завершена в файле.

#include <torch/serialize/tensor.h>
#include <vector>
#include <THC/THC.h>
#include <cuda.h>
#include <cuda_runtime_api.h>

#include "ball_query_cuda.h"

extern THCState *state;

// 定义检查数据类型的宏
#define CHECK_CUDA(x) do { \
	  if (!x.type().is_cuda()) { \
		      fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
		      exit(-1); \
		    } \
} while (0)
#define CHECK_CONTIGUOUS(x) do { \
	  if (!x.is_contiguous()) { \
		      fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
		      exit(-1); \
		    } \
} while (0)
#define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)

// 完成输入数据类型的检查，同时调用cu文件中定义的函数
int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, 
    at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) {
    CHECK_INPUT(new_xyz_tensor);
    CHECK_INPUT(xyz_tensor);
    const float *new_xyz = new_xyz_tensor.data<float>();
    const float *xyz = xyz_tensor.data<float>();
    int *idx = idx_tensor.data<int>();
    
    ball_query_kernel_launcher_fast(b, n, m, radius, nsample, new_xyz, xyz, idx);
    return 1;
}

ball_query_cuda.cuвыполнить

#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include "ball_query_cuda.h"
#include "cuda_utils.h"


__global__ void ball_query_kernel_fast(int b, int n, int m, float radius, int nsample, 
    const float *__restrict__ new_xyz, const float *__restrict__ xyz, int *__restrict__ idx) {
    // new_xyz: (B, M, 3)
    // xyz: (B, N, 3)
    // output:
    //      idx: (B, M, nsample)
    int bs_idx = blockIdx.y;
    int pt_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (bs_idx >= b || pt_idx >= m) return;

    new_xyz += bs_idx * m * 3 + pt_idx * 3;
    xyz += bs_idx * n * 3;
    idx += bs_idx * m * nsample + pt_idx * nsample;

    float radius2 = radius * radius;
    float new_x = new_xyz[0];
    float new_y = new_xyz[1];
    float new_z = new_xyz[2];

    int cnt = 0;
    for (int k = 0; k < n; ++k) {
        float x = xyz[k * 3 + 0];
        float y = xyz[k * 3 + 1];
        float z = xyz[k * 3 + 2];
        float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z);
        if (d2 < radius2){
            if (cnt == 0){
                for (int l = 0; l < nsample; ++l) {
                    idx[l] = k;
                }
            }
            idx[cnt] = k;
            ++cnt;
            if (cnt >= nsample) break;
        }
    }
}


void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample, \
    const float *new_xyz, const float *xyz, int *idx) {
    // new_xyz: (B, M, 3)
    // xyz: (B, N, 3)
    // output:
    //      idx: (B, M, nsample)

    cudaError_t err;

    dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b);  // blockIdx.x(col), blockIdx.y(row)
    dim3 threads(THREADS_PER_BLOCK);

    ball_query_kernel_fast<<<blocks, threads>>>(b, n, m, radius, nsample, new_xyz, xyz, idx);
    // cudaDeviceSynchronize();  // for using printf in kernel function
    err = cudaGetLastError();
    if (cudaSuccess != err) {
        fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
        exit(-1);
    }
}

На этом основная функция ball_query завершена, и теперь нам нужно использоватьsetup.pyфайл для компиляции вышеуказанного файла.setup.pyКонкретная реализация файла выглядит следующим образом:

import os
import subprocess

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


def make_cuda_ext(name, module, sources):
    cuda_ext = CUDAExtension(
        name='%s.%s' % (module, name),
        sources=[os.path.join(module.split('.')[-1], src) for src in sources]
    )
    print([os.path.join(*module.split('.'), src) for src in sources])
    return cuda_ext

if __name__ == '__main__':
    setup(
        name='ballquery',
        packages=find_packages(),
        ext_modules=[
            CUDAExtension('ball_query_cuda',[
                'ball_query_src/api.cpp',
                'ball_query_src/ball_query.cpp',
                'ball_query_src/ball_query_cuda.cu',  
            ])
        ],
        cmdclass={
            'build_ext': BuildExtension
        }
    )

На данный момент мы сгенерировали библиотеку ссылок из приведенного выше кода, но если нам нужно встроить ее в нейронную сеть, нам также необходимо определить ее методы прямого и обратного распространения. мы здесьtest_ball_query.pyОн завершает свое прямое распространение и обратное распространение в файле.

import torch
import torch.nn as nn
from torch.autograd import Function, Variable
import math

import ball_query_cuda

# 定义该方法的前向传播和反向传播方法
class BallQuery(Function):
    
    @staticmethod
    def forward(ctx, radius: float, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor) -> torch.Tensor:
        """
        :param ctx:
        :param radius: float, radius of the balls
        :param nsample: int, maximum number of features in the balls
        :param xyz: (B, N, 3) xyz coordinates of the features
        :param new_xyz: (B, npoint, 3) centers of the ball query
        :return:
            idx: (B, npoint, nsample) tensor with the indicies of the features that form the query balls
        """
        assert new_xyz.is_contiguous()
        assert xyz.is_contiguous()

        B, N, _ = xyz.size()
        npoint = new_xyz.size(1)
        idx = torch.cuda.IntTensor(B, npoint, nsample).zero_()
        
        ball_query_cuda.ball_query_wrapper(B, N, npoint, radius, nsample, new_xyz, xyz, idx)
        return idx

    @staticmethod
    def backward(ctx, a=None):
        return None, None, None, None


ball_query = BallQuery.apply

xyz = torch.randn(2, 128, 3).cuda()
new_xyz = xyz

result = ball_query(0.8, 3, xyz, new_xyz)

print(result.shape)
Категории