C++およびCUDAによるPytorchの拡張方法
116278 ワード
PytorchはすでにNVIDIA cuDNn、Intel MKL、NNPACKなどの下位層を使用して訓練速度を速めているが、場合によってはPytorchの既存の操作を組み合わせるだけでは足りない特定のアルゴリズムを実現するなど、いくつかの場合がある.これはPytorchが特定の操作の上でとても良い最適化を経たが、Pytorchがすでに書いたこれらの操作に対して、もし私たちが組み合わせて、私たちの新しいアルゴリズムを構成するならば、Pytorchはあなたのアルゴリズムの具体的な実行の流れにかかわらず、一般的にPytorchは設計した操作に従ってGPUの通路を使うだけで、これにより、チャネルが十分に利用されないか、直接過負荷になり、python解釈器も最適化されず、プログラムの実行速度が遅くなります.
準備作業
サーバ上で開発されるため、一般的なユーザーにはroot権限がないため、開発中に依存ソフトウェアが互換性がない場合、依存ソフトウェアを更新する権限がありません.これらの問題を回避するために、まず、必要な依存環境を自分のユーザーディレクトリの下で構成します.最初のステップはanaconda 3環境をインストールする必要があります.公式サイトで相応のバージョンをダウンロードし、公式サイトのコマンドに従ってインストールを実行すればいいです.anaconda 3をインストールしてpythonを入力すると、起動したのはシステムが持っているpython環境なのか、PATHでanaconda 3のpython環境をインポートする必要があります.私が実行するコマンドはexport PATH=/home/guangyuan/anaconda 3/bin:$PATHです.このコマンドはインストールしたanaconda 3パスに応じて調整する必要があります.第2のステップはconda create-n venv pip python=3.6で新しい仮想環境を作成することで、新しい仮想環境で開発することができ、エラーが発生したら、直接仮想環境を削除し、再構成すればよい.仮想環境のインストールが完了するには、source activate venvを使用して仮想化環境をアクティブにし、新しく作成された仮想環境開発に進む必要があります.仮想環境を終了するにはsource deactivateを入力します.
注意:新しく作成した仮想環境をアクティブにするには、PYTHONPATHを使用して仮想環境にインストールされているライブラリをインポートする必要があります.コマンドは次のとおりです:export PYTHONPATH=/home/guangyuan/anaconda 3/envs/venv/lib/python 3.6/site-packages:$PYTHONPATH;具体的なパスは、自分で構成した環境に合わせて調整する必要があります.ステップ3 pytorch 1をインストールします.2(インストールリファレンス公式サイトチュートリアルhttps://pytorch.org/)第4部では、pip install pytest pip install pybind 11が必要な環境構成でほぼ完了し、開発中にライブラリが不足している場合はpip installライブラリ名で直接インストールすればよい.
c++およびCUDAによるプログラミング
本稿では、pytorchにdeptwise convを追加する例として、pytorchtにソースコードを追加する方法を一歩一歩示します.まず、私たちのエンジニアリング組織構造を示し、その後、各ファイルの役割を詳しく説明します. depthwise_conv depthwise_conv.py depthwise_conv_cuda.cpp depthwise_conv_kernel.cu setup.py __init__.py
まずdepthwiseがありますconvフォルダ、フォルダには5つのファイルがあります.c++とcudaにかかわるのはdepthwise_conv_cuda.cppとdepthwise_conv_kernel.Cu、その後setup.pyはコンパイルされpythonにバインドされ、最後にpythonモジュール、すなわちdepthwise_にカプセル化されます.conv.pyの内容.c++とcudaでpytorchを拡張するには、一般的にこれらのファイルが必要です.
まずdepthwise_を見てみましょうconv_kernel.Cuファイル.pytorchに追加されたレイヤの主な実装プロセスを含む.主にcuda kernel実装とcuda kernel呼び出し関数の実装が含まれており、この関数はまず、入力されたパラメータに基づいて、入力されたデータテンソルに基づいて、データのbatch sizeやchannelサイズなどを得ることができ、計算結果を保存するためにいくつかのデータテンソルを再割り当てする必要があります.
pytorchにdepthwise convを追加してdepthwise_を具体的に説明します.conv_kernel.Cuは過成を実現した.c++とcudaによっていくつかの深さ学習フレームワークに層を追加することは,一般的にこの層のforwardとbackwardプロセスを実現する必要があることを知っているが,計算速度を向上させるためにforwardとbackwardのコア計算はcuda kernelで実現されている.depthwise_の場合conv_kernel.Cu我々は実現した
__global__ void ConvForward() __global__ void ConvBackward() __global__ void ConvBackwardWeight() __global__ void ConvBackwardBias()等核関数とDepthWiseConvForwardLaucher()、DepthWiseConvBackwardLaucher()は、cuda kernel関数をカプセル化する.
次にdepwise convのforward実装プロセスをソースコードと組み合わせて具体的に説明する.まず実現します_global__ void ConvForward()関数、コードは以下のように実現されます.
このセグメントコードはdepthwise convのforward cuda kernel実装プロセスであり,pytorchフレームワークとは無関係にcudaプログラミングの基礎が必要である.
その後、forward cuda kernelをpytorchフレームワークにカプセル化する方法を実装する必要があります.一般的にはincludeを含む必要があります.
#includeという2つのヘッダファイル.
コードは、以下のように具体的に実装されます.
この関数は一般的に以下のことをする必要がある:1、入力Tensorによってbatch_を得るsize、channels、input_heightとinput_widthなどの情報は後で使用するために使用される.2、得られた情報と計算式から出力Tensor次元の大きさを計算する.3、計算した出力Tensor次元情報に基づいて出力Tensor申請メモリである;4、pytorchから提供されたAT_DISPATCH_FLOATING_TYPES_AND_HALFインタフェースはcuda kernel関数を呼び出します.
私たちのデータはat::Tensorが定義した変数に格納されているので、at::Tensorの一般的な関数、例えば通過することを理解する必要があります.size()
Tensor次元の情報を得ることができる.Data()はTensorの裸のポインタを得ることができ、通過する.type()はこのTensorのデータ型などを入手できる.
AT_について簡単にご紹介しますDISPATCH_FLOATING_TYPES_AND_HALFの使い方、AT_DISPATCH_FLOATING_TYPES_AND_HALFは3つのパラメータを与える必要があり、1つ目はデータ型でtensor呼び出しを入力することができる.type()
得られ,2番目は1つの文字列であり,特に要求はなく,3番目のパラメータは匿名関数を受け入れ,cuda kernel関数論理を呼び出すのはすべてこの匿名関数で実現される.これでdepthwise conv forwardプロセスが完了しました.
deptwise convのbackward実装手順はforwardと同様であり,ここでは詳細に説明せず,ソースコード実装のみを与える.
これでdepthwise_conv_kernel.Cuがやるべき仕事はもう終わった.
その後、もう一つ作成します.cppファイルは、
本稿では、上記の機能を
このソースコードは、次のことを行います.1、
これで
その後depthwise_に入りますconvディレクトリの下でpython setup.py build_ext--inplaceはソースコードをコンパイルできます.コンパイルが完了したらdepthwiseでconvディレクトリの下には、コンパイルによって生成された
導入するPYTHONPATHパスは、自分の実際のパスに合わせて調整する必要があります.
次に、depthwise_でカスタムレイヤを定義できます.conv.pyでは、次のように実装されます.
このコードは基本的なpytorchカスタムレイヤの実装手順であり、公式サイトのチュートリアルを参照してください.https://pytorch.org/docs/stable/notes/extending.html唯一の違いは、DepthWiseConvFunctionレイヤでforwardとbackwardの具体的な計算です.pytorchのpythonインタフェース関数ではなく、呼び出した自分がソースコードで拡張されています.
準備作業
サーバ上で開発されるため、一般的なユーザーにはroot権限がないため、開発中に依存ソフトウェアが互換性がない場合、依存ソフトウェアを更新する権限がありません.これらの問題を回避するために、まず、必要な依存環境を自分のユーザーディレクトリの下で構成します.最初のステップはanaconda 3環境をインストールする必要があります.公式サイトで相応のバージョンをダウンロードし、公式サイトのコマンドに従ってインストールを実行すればいいです.anaconda 3をインストールしてpythonを入力すると、起動したのはシステムが持っているpython環境なのか、PATHでanaconda 3のpython環境をインポートする必要があります.私が実行するコマンドはexport PATH=/home/guangyuan/anaconda 3/bin:$PATHです.このコマンドはインストールしたanaconda 3パスに応じて調整する必要があります.第2のステップはconda create-n venv pip python=3.6で新しい仮想環境を作成することで、新しい仮想環境で開発することができ、エラーが発生したら、直接仮想環境を削除し、再構成すればよい.仮想環境のインストールが完了するには、source activate venvを使用して仮想化環境をアクティブにし、新しく作成された仮想環境開発に進む必要があります.仮想環境を終了するにはsource deactivateを入力します.
注意:新しく作成した仮想環境をアクティブにするには、PYTHONPATHを使用して仮想環境にインストールされているライブラリをインポートする必要があります.コマンドは次のとおりです:export PYTHONPATH=/home/guangyuan/anaconda 3/envs/venv/lib/python 3.6/site-packages:$PYTHONPATH;具体的なパスは、自分で構成した環境に合わせて調整する必要があります.ステップ3 pytorch 1をインストールします.2(インストールリファレンス公式サイトチュートリアルhttps://pytorch.org/)第4部では、pip install pytest pip install pybind 11が必要な環境構成でほぼ完了し、開発中にライブラリが不足している場合はpip installライブラリ名で直接インストールすればよい.
c++およびCUDAによるプログラミング
本稿では、pytorchにdeptwise convを追加する例として、pytorchtにソースコードを追加する方法を一歩一歩示します.まず、私たちのエンジニアリング組織構造を示し、その後、各ファイルの役割を詳しく説明します.
まずdepthwiseがありますconvフォルダ、フォルダには5つのファイルがあります.c++とcudaにかかわるのはdepthwise_conv_cuda.cppとdepthwise_conv_kernel.Cu、その後setup.pyはコンパイルされpythonにバインドされ、最後にpythonモジュール、すなわちdepthwise_にカプセル化されます.conv.pyの内容.c++とcudaでpytorchを拡張するには、一般的にこれらのファイルが必要です.
まずdepthwise_を見てみましょうconv_kernel.Cuファイル.pytorchに追加されたレイヤの主な実装プロセスを含む.主にcuda kernel実装とcuda kernel呼び出し関数の実装が含まれており、この関数はまず、入力されたパラメータに基づいて、入力されたデータテンソルに基づいて、データのbatch sizeやchannelサイズなどを得ることができ、計算結果を保存するためにいくつかのデータテンソルを再割り当てする必要があります.
pytorchにdepthwise convを追加してdepthwise_を具体的に説明します.conv_kernel.Cuは過成を実現した.c++とcudaによっていくつかの深さ学習フレームワークに層を追加することは,一般的にこの層のforwardとbackwardプロセスを実現する必要があることを知っているが,計算速度を向上させるためにforwardとbackwardのコア計算はcuda kernelで実現されている.depthwise_の場合conv_kernel.Cu我々は実現した
__global__ void ConvForward() __global__ void ConvBackward() __global__ void ConvBackwardWeight() __global__ void ConvBackwardBias()等核関数とDepthWiseConvForwardLaucher()、DepthWiseConvBackwardLaucher()は、cuda kernel関数をカプセル化する.
次にdepwise convのforward実装プロセスをソースコードと組み合わせて具体的に説明する.まず実現します_global__ void ConvForward()関数、コードは以下のように実現されます.
template <typename scalar_t>
__global__ void ConvForward(const int nthreads,
const scalar_t* const bottom_data, const int num, const int channels,
const int height, const int width,const int conved_height,
const int conved_width,const int kernel_h, const int kernel_w,
const int stride_h, const int stride_w, const int pad_h, const int pad_w,
scalar_t* const top_data,const scalar_t* const weight,const scalar_t* const bias,const bool bias_term_) {
CUDA_KERNEL_LOOP(index, nthreads) {
const int pw = index % conved_width;
const int ph = (index / conved_width) % conved_height;
const int c = (index / conved_width / conved_height) % channels;
const int n = index / conved_width / conved_height / channels;
int hstart = ph * stride_h - pad_h;
int wstart = pw * stride_w - pad_w;
int hend = min(hstart + kernel_h, height + pad_h);
int wend = min(wstart + kernel_w, width + pad_w);
hstart = max(hstart, 0);
wstart = max(wstart, 0);
hend = min(hend, height);
wend = min(wend, width);
scalar_t aveval = 0;
const scalar_t* const bottom_slice =
bottom_data + (n * channels + c) * height * width;
const scalar_t* const weight_slice =
weight + c * kernel_h * kernel_w;
int khstart=hend<kernel_h?kernel_h-hend:0;
int kwstart=wend<kernel_w?kernel_w-wend:0;
for (int h = hstart; h < hend; ++h) {
for (int w = wstart; w < wend; ++w) {
aveval += bottom_slice[h * width + w]*weight_slice[(khstart+h-hstart) * kernel_w + (kwstart+w-wstart)];
}
}
if(bias_term_) {
aveval+=bias[c];
}
top_data[index] = aveval;
}
}
このセグメントコードはdepthwise convのforward cuda kernel実装プロセスであり,pytorchフレームワークとは無関係にcudaプログラミングの基礎が必要である.
その後、forward cuda kernelをpytorchフレームワークにカプセル化する方法を実装する必要があります.一般的にはincludeを含む必要があります.
#includeという2つのヘッダファイル.
コードは、以下のように具体的に実装されます.
#include
", cudaGetErrorString(err));
exit(-1);
}
return output;
}
この関数は一般的に以下のことをする必要がある:1、入力Tensorによってbatch_を得るsize、channels、input_heightとinput_widthなどの情報は後で使用するために使用される.2、得られた情報と計算式から出力Tensor次元の大きさを計算する.3、計算した出力Tensor次元情報に基づいて出力Tensor申請メモリである;4、pytorchから提供されたAT_DISPATCH_FLOATING_TYPES_AND_HALFインタフェースはcuda kernel関数を呼び出します.
私たちのデータはat::Tensorが定義した変数に格納されているので、at::Tensorの一般的な関数、例えば通過することを理解する必要があります.size()
Tensor次元の情報を得ることができる.Data()はTensorの裸のポインタを得ることができ、通過する.type()はこのTensorのデータ型などを入手できる.
AT_について簡単にご紹介しますDISPATCH_FLOATING_TYPES_AND_HALFの使い方、AT_DISPATCH_FLOATING_TYPES_AND_HALFは3つのパラメータを与える必要があり、1つ目はデータ型でtensor呼び出しを入力することができる.type()
得られ,2番目は1つの文字列であり,特に要求はなく,3番目のパラメータは匿名関数を受け入れ,cuda kernel関数論理を呼び出すのはすべてこの匿名関数で実現される.これでdepthwise conv forwardプロセスが完了しました.
deptwise convのbackward実装手順はforwardと同様であり,ここでは詳細に説明せず,ソースコード実装のみを与える.
// （Tensor）、 Weight Bias backward 。
template <typename scalar_t>
__global__ void ConvBackward(const int nthreads,
const scalar_t* const top_diff,
const int num, const int channels, const int height,
const int width, const int conved_height, const int conved_width,
const int kernel_h, const int kernel_w, const int stride_h,
const int stride_w, const int pad_h, const int pad_w,
scalar_t* const bottom_diff,
const scalar_t* const weight) {
CUDA_KERNEL_LOOP(index, nthreads) {
const int w = index % width + pad_w;
const int h = (index / width) % height + pad_h;
const int c = (index / width / height) % channels;
const int n = index / width / height / channels;
const int phstart = (h < kernel_h) ? 0 : (h - kernel_h) / stride_h + 1;
const int phend = min(h / stride_h + 1, conved_height);
const int pwstart = (w < kernel_w) ? 0 : (w - kernel_w) / stride_w + 1;
const int pwend = min(w / stride_w + 1, conved_width);
const int khstart=(h >= kernel_h) ? ((h-kernel_h)%stride_h)+(kernel_h-stride_h): h;
const int kwstart=(w >= kernel_w) ? ((w-kernel_w)%stride_w)+(kernel_w-stride_w) : w;
scalar_t gradient = 0;
const scalar_t* const top_diff_slice =
top_diff + (n * channels + c) * conved_height * conved_width;
const scalar_t* const weight_slice =weight + c * kernel_h * kernel_w;
for (int ph = phstart; ph < phend; ++ph) {
for (int pw = pwstart; pw < pwend; ++pw) {
int kh=khstart-(ph-phstart)*stride_h;
int kw=kwstart-(pw-pwstart)*stride_w;
gradient += top_diff_slice[ph * conved_width + pw] *weight_slice[kh*kernel_w+kw];
}
}
bottom_diff[index] = gradient;
}
}
template <typename scalar_t>
__global__ void ConvBackwardWeight(const int nthreads,
const scalar_t* const top_diff,
const int num, const int channels, const int height,
const int width, const int conved_height, const int conved_width,
const int kernel_h, const int kernel_w, const int stride_h,
const int stride_w, const int pad_h, const int pad_w,
scalar_t* const weight_diff,
const scalar_t* const bottom_data) {
CUDA_KERNEL_LOOP(index, nthreads) {
const int kw=index % kernel_w;
const int kh= (index /kernel_w)%kernel_h;
const int c=index /kernel_w/kernel_h;
scalar_t gradient = 0;
for( int n=0;n<num;n++) {
const scalar_t* const top_diff_slice = top_diff + (n * channels + c) * conved_height * conved_width;
const scalar_t* const bottom_data_slice = bottom_data + (n * channels + c) * height * width;
const int phstart=max(DIVIDE_CEIL((pad_h-kh),stride_h),0);
const int phend=min(DIVIDE_CEIL((height+pad_h-kh),stride_h),conved_height);
const int pwstart=max(DIVIDE_CEIL((pad_w-kw),stride_w),0);
const int pwend=min(DIVIDE_CEIL((width+pad_w-kw),stride_w),conved_width);
for(int ph=phstart;ph<phend;ph++){
for (int pw=pwstart;pw<pwend;pw++){
const int h=ph*stride_h+kh-pad_h;
const int w=pw*stride_w+kw-pad_w;
gradient+=top_diff_slice[ph * conved_width + pw]*bottom_data_slice[h*width+w];
}
}
}
weight_diff[c * kernel_h * kernel_w+kh*kernel_w+kw]+=gradient;
}
}
template <typename scalar_t>
__global__ void ConvBackwardBias(const int nthreads,
const scalar_t* const top_diff,
const int num, const int channels, const int height,
const int width, const int conved_height, const int conved_width,
const int kernel_h, const int kernel_w, const int stride_h,
const int stride_w, const int pad_h, const int pad_w,
scalar_t* const bias_diff) {
CUDA_KERNEL_LOOP(index, nthreads) {
const int c = index;
scalar_t gradient=0;
for( int n=0;n<num;n++) {
const scalar_t* const top_diff_slice =
top_diff + (n * channels + c) * conved_height * conved_width;
for(int ph=0;ph<conved_height;ph++) {
for (int pw=0;pw<conved_width;pw++) {
gradient+=top_diff_slice[ph * conved_width + pw];
}
}
}
bias_diff[c]+=gradient;
}
}
// DepthWiseConvBackwarddLaucher cuda kernel 。
std::vector<at::Tensor> DepthWiseConvBackwarddLaucher(const at::Tensor output_grad, const at::Tensor input, const at::Tensor weight, const at::Tensor bias,
const int kernel_h, const int kernel_w, const int stride_h, const int stride_w,
const int pad_h, const int pad_w, const bool bias_term_){
const auto batch_size = input.size(0);
const auto channels = input.size(1);
const auto input_height = input.size(2);
const auto input_width = input.size(3);
const auto kernal_extent_h = /* dilation_h * */ (kernel_h - 1) + 1;
const auto conved_height = (input_height + 2 * pad_h - kernal_extent_h) / stride_h + 1;
const auto kernal_extent_w = /* dilation_w * */ (kernel_w - 1) + 1;
const auto conved_width = (input_width + 2 * pad_w - kernal_extent_w) / stride_w + 1;
const int count_weight = channels * kernel_h * kernel_w;
const int count_input = batch_size * channels * input_height * input_width;
auto weight_diff = at::zeros_like(weight);
auto bottom_diff = at::zeros_like(input);
at::Tensor bias_diff;
int count_bias = 0;
if (bias_term_){
count_bias = channels;
bias_diff = at::zeros_like(bias);
}
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
output_grad.type(), "ConvLaucherBackward",
([&]{
const scalar_t *bottom_data = input.data<scalar_t>();
const scalar_t *depthwise_weight = weight.data<scalar_t>();
const scalar_t *top_diff = output_grad.data<scalar_t>();
scalar_t *depthwise_weight_diff = weight_diff.data<scalar_t>();
scalar_t *depthwise_bottom_diff = bottom_diff.data<scalar_t>();
if (bias_term_){
scalar_t *depthwise_bias_diff = bias_diff.data<scalar_t>();
ConvBackwardBias<scalar_t><<<GET_BLOCKS(count_bias), THREADS_PER_BLOCK>>>(count_bias, top_diff, batch_size,
channels, input_height, input_width, conved_height, conved_width, kernel_h, kernel_w, stride_h,
stride_w, pad_h, pad_w, depthwise_bias_diff);
}
ConvBackwardWeight<scalar_t><<<GET_BLOCKS(count_weight), THREADS_PER_BLOCK>>>(count_weight, top_diff, batch_size,
channels, input_height, input_width, conved_height, conved_width, kernel_h, kernel_w, stride_h,
stride_w, pad_h, pad_w, depthwise_weight_diff, bottom_data);
ConvBackward<scalar_t><<<GET_BLOCKS(count_input), THREADS_PER_BLOCK>>>(count_input, top_diff, batch_size,
channels, input_height, input_width, conved_height, conved_width, kernel_h, kernel_w, stride_h,
stride_w, pad_h, pad_w, depthwise_bottom_diff, depthwise_weight);
}));
cudaError_t err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "cudaCheckError() failed : %s
", cudaGetErrorString(err));
exit(-1);
}
if (bias_term_){
return {bottom_diff, weight_diff, bias_diff};
}
else{
return {bottom_diff, weight_diff};
}
}
これでdepthwise_conv_kernel.Cuがやるべき仕事はもう終わった.
その後、もう一つ作成します.cppファイルは、
DepthWiseConvForwardLaucherとDepthWiseConvBackwarddLaucherの関数をカプセル化し、さっきc++とcudaで実現したレイヤをPythonにバインドします.本稿では、上記の機能を
depthwise_conv_cuda.cppで実現し、まずソースコードを与える.#include このソースコードは、次のことを行います.1、
depthwise_conv_forward_cuda、depthwise_conv_backward_cuda関数を作成します.この2つは主に2つのことをしました.1つはCHECK_を呼び出すことです.INPUTは、入力Tensorsが適合しているかどうかを確認し、その後、DepthWiseConvForwardLaucherおよびDepthWiseConvBackwarddLaucherを呼び出す.2.PYBIND11_MODULEによってdepthwise_conv_forward_cudaとdepthwise_conv_backward_cudaをpythonにバインドする.これで
depthwise_conv_cuda.cppで完了する作業は完了しました.その後、setup.pyファイルを作成してソースコードをコンパイルする必要があります.setup.pyの内容は以下の通りです.from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='depethwise_conv',
ext_modules=[
CUDAExtension('depthwise_conv_cuda', [
'./depthwise_conv_cuda.cpp',
'./depthwise_conv_kernel.cu',
])
],
cmdclass={'build_ext': BuildExtension})
setup.pyではsetuptoolsを使用してC++コードをコンパイルし、CUDAExtensionとBuildExtensionはC++とCudaのコンパイルを容易に実現することができます.このコードは、他のレイヤに移植された追加タスクを簡単に変更することができます.その後depthwise_に入りますconvディレクトリの下でpython setup.py build_ext--inplaceはソースコードをコンパイルできます.コンパイルが完了したらdepthwiseでconvディレクトリの下には、コンパイルによって生成された
depthwise_conv_cuda.cpython-36m-x86_64-linux-gnu.soファイルが表示されます.その後、python環境に入って検証することができます.depthwise_にいることに注意してください.convディレクトリの下でpython環境に入ります.すなわち、python環境の現在のディレクトリはdepthwise_です.convの下、現在のディレクトリの下にdepthwise_conv_cuda.cpython-36m-x86_64-linux-gnu.soがあり、import depthwise_conv_cuda;任意のディレクトリでimport depthwise_conv_Cuda、PYTHONPATHでインポートする必要があります.コマンドは以下のexport PYTHONPATH=/data/nfs_です.share/public/guangyuan/workplace/ops/depthwise_conv:$PYTHONPATH 導入するPYTHONPATHパスは、自分の実際のパスに合わせて調整する必要があります.
In [1]: import torch
In [2]: import depthwise_conv_cuda
In [3]: depthwise_conv_cuda.forward
Out[3]: <built-in method forward of PyCapsule object at 0x7fb70c690d50>
次に、depthwise_でカスタムレイヤを定義できます.conv.pyでは、次のように実装されます.
import math
import collections
from itertools import repeat
import torch
import torch.nn as nn
from torch.autograd import Function
from torch.nn.parameter import Parameter
import depthwise_conv_cuda
def _ntuple(n):
def parse(x):
if isinstance(x, collections.Iterable):
return x
return tuple(repeat(x, n))
return parse
_pair = _ntuple(2)
class DepthWiseConvFunction(Function):
@staticmethod
def forward(ctx, inputs, weight, bias, kernel_size, stride=1, padding=0, use_bias=False):
kernel_size = _pair(kernel_size)
stride = _pair(stride)
padding = _pair(padding)
assert inputs.size(1) == weight.size(0)
if use_bias:
ctx.save_for_backward(*[inputs, weight, bias])
# ctx.save_for_backward(inputs, weight, bias)
else:
bias = torch.zeros(1, dtype=weight.dtype, device=weight.device)
ctx.save_for_backward(*[inputs, weight, bias])
ctx.use_bias = use_bias
ctx.kernel_size = kernel_size
ctx.stride = stride
ctx.padding = padding
output = depthwise_conv_cuda.forward(inputs, weight, bias, kernel_size[0], kernel_size[1],
stride[0], stride[1], padding[0], padding[1], use_bias)
return output
@staticmethod
def backward(ctx, grad_output):
assert grad_output.is_cuda
use_bias = ctx.use_bias
kernel_size = ctx.kernel_size
stride = ctx.stride
padding = ctx.padding
if use_bias:
inputs, weight, bias = ctx.saved_variables
output_grads = depthwise_conv_cuda.backward(grad_output.contiguous(), inputs, weight, bias, kernel_size[0],
kernel_size[1], stride[0], stride[1], padding[0],
padding[1], use_bias)
inputs_grad, weight_grad, bias_grad = output_grads
return inputs_grad, weight_grad, bias_grad, None, None, None, None
else:
inputs, weight, bias = ctx.saved_variables
output_grads = depthwise_conv_cuda.backward(grad_output.contiguous(), inputs, weight, bias, kernel_size[0],
kernel_size[1], stride[0], stride[1], padding[0],
padding[1], use_bias)
inputs_grad, weight_grad = output_grads
return inputs_grad, weight_grad, None, None, None, None, None
depthwise_conv = DepthWiseConvFunction.apply
class DepthWiseConv2d(nn.Module):
def __init__(self, in_channels, kernel_size=3, stride=1, padding=0, use_bias=False):
super(DepthWiseConv2d, self).__init__()
self.in_channels = in_channels
self.kernel_size = _pair(kernel_size)
self.stride = _pair(stride)
self.padding = _pair(padding)
self.use_bias = use_bias
self.weight = Parameter(torch.Tensor(in_channels, 1, *self.kernel_size))
if use_bias:
self.bias = Parameter(torch.Tensor(in_channels))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
n = self.in_channels
for k in self.kernel_size:
n *= k
stdv = 1. / math.sqrt(n)
self.weight.data.uniform_(-stdv, stdv)
if self.bias is not None:
self.bias.data.uniform_(-stdv, stdv)
def forward(self, inputs):
return depthwise_conv(inputs, self.weight, self.bias, self.kernel_size, self.stride, self.padding, self.use_bias)
def testDepthWiseConvFunction():
from torch.autograd import Variable
from torch.nn import functional as F
device = "cuda:0"
inputs = torch.randn(4, 3, 7, 7).to(device)
w = Variable(torch.randn(3, 1, 3, 3), requires_grad=True).to(device)
b = Variable(torch.randn(3), requires_grad=True).to(device)
# opt = F.conv2d(inputs, w, bias=None, stride=1, padding=0, dilation=1, groups=3)
# print(opt.size())
inp = depthwise_conv(inputs, w, b, 3)
loss = inp.sum()
loss.backward()
def testDepthWiseConv2d():
device = "cuda:0"
inputs = torch.randn(4, 3, 7, 7).to(device)
depthwcon = DepthWiseConv2d(3)
depthwcon.to(device)
outp = depthwcon(inputs)
loss = outp.sum()
loss.backward()
print(outp.size())
if __name__ == '__main__':
testDepthWiseConvFunction()
testDepthWiseConv2d()
このコードは基本的なpytorchカスタムレイヤの実装手順であり、公式サイトのチュートリアルを参照してください.https://pytorch.org/docs/stable/notes/extending.html唯一の違いは、DepthWiseConvFunctionレイヤでforwardとbackwardの具体的な計算です.pytorchのpythonインタフェース関数ではなく、呼び出した自分がソースコードで拡張されています.