pytorch 學習筆記之編寫 C 擴展,又漲姿勢了

pytorch利用 CFFI 進行 C 語言擴展。包括兩個基本的步驟（docs）：

1. 編寫 C 代碼；

2. python 調用 C 代碼，實現相應的 Function 或 Module。

在之前的文章中，我們已經了解了如何自定義 Module。至於 [py]torch 的 C 代碼庫的結構，我們留待之后討論； 這里，重點關注，如何在 pytorch C 代碼庫高層接口的基礎上，編寫 C 代碼，以及如何調用自己編寫的 C 代碼。

官方示例了如何定義一個加法運算（見 repo）。這里我們定義ReLU函數（見 repo）。

1. C 代碼

pytorch C 的基本數據結構是 THTensor（THFloatTensor、THByteTensor等）。我們以簡單的 ReLU 函數為例，示例編寫 C 。

y=ReLU(x)=max(x,0)

Function 需要定義前向和后向兩個方向的操作，因此，C 代碼要實現相應的功能。

1.1 頭文件聲明

/* ext_lib.h */ int relu_forward(THFloatTensor *input, THFloatTensor *output); int relu_backward(THFloatTensor *grad_output, THFloatTensor *input, THFloatTensor *grad_input); 

1.2 函數實現

TH/TH.h 包括了 pytorch C 代碼數據結構和函數的聲明，這是唯一需要添加的 include 依賴。

/* ext_lib.c */ #include <TH/TH.h> int relu_forward(THFloatTensor *input, THFloatTensor *output) { THFloatTensor_resizeAs(output, input); THFloatTensor_clamp(output, input, 0, INFINITY); return 1; } int relu_backward(THFloatTensor *grad_output, THFloatTensor *input, THFloatTensor *grad_input) { THFloatTensor_resizeAs(grad_input, grad_output); THFloatTensor_zero(grad_input); THLongStorage* size = THFloatTensor_newSizeOf(grad_output); THLongStorage *stride = THFloatTensor_newStrideOf(grad_output); THByteTensor *mask = THByteTensor_newWithSize(size, stride); THFloatTensor_geValue(mask, input, 0); THFloatTensor_maskedCopy(grad_input, mask, grad_output); return 1; } 

2. 編譯代碼

2.1 依賴

由於 pytorch 的代碼是純 C 的，因此沒有過多的依賴，只需要安裝：

• pytorch - 安裝方法見官網

• cffi - pip install cffi

編譯文件非常簡單，主要是添加頭文件和實現文件，以及相關的宏定義; 同時文件還指定了編譯后的調用位置（此外為_ext.ext_lib）：

# build.py import os import torch from torch.utils.ffi import create_extension sources = ['src/ext_lib.c'] headers = ['src/ext_lib.h'] defines = [] with_cuda = False if torch.cuda.is_available(): print('Including CUDA code.') sources += ['src/ext_lib_cuda.c'] headers += ['src/ext_lib_cuda.h'] defines += [('WITH_CUDA', None)] with_cuda = True ffi = create_extension( '_ext.ext_lib', headers=headers, sources=sources, define_macros=defines, relative_to=__file__, with_cuda=with_cuda ) if __name__ == '__main__': ffi.build() python build.py 

3. python 調用

3.1 編寫配置文件

python 的調用非常簡單——pytorch 的 tensor 對象，對應 C 代碼的 THTensor 對象，以此作參數進行調用即可。配置文件如下：

import torch from torch.autograd import Function from _ext import ext_lib class ReLUF(Function): def forward(self, input): self.save_for_backward(input) output = input.new() if not input.is_cuda: ext_lib.relu_forward(input, output) else: raise Exception, "No CUDA Implementation" return output def backward(self, grad_output): input, = self.saved_tensors grad_input = grad_output.new() if not grad_output.is_cuda: ext_lib.relu_backward(grad_output, input, grad_input) else: raise Exception, "No CUDA Implementation" return grad_input 

3.2 測試

此處省略 Module 的定義。下面測試下新定義的基於 C 的 ReLU 函數。

import torch import torch.nn as nn from torch.autograd import Variable from modules.relu import ReLUM torch.manual_seed(1111) class MyNetwork(nn.Module): def __init__(self): super(MyNetwork, self).__init__() self.relu = ReLUM() def forward(self, input): return self.relu(input) model = MyNetwork() x = torch.randn(1, 25).view(5, 5) input = Variable(x, requires_grad=True) output = model(input) print(output) print(input.clamp(min=0)) output.backward(torch.ones(input.size())) print(input.grad.data) 

輸出結果如下：

Variable containing: 0.8749 0.5990 0.6844 0.0000 0.0000 0.6516 0.0000 1.5117 0.5734 0.0072 0.1286 1.4171 0.0796 1.0355 0.0000 0.0000 0.0000 0.0312 0.0999 0.0000 1.0401 1.0599 0.0000 0.0000 0.0000 [torch.FloatTensor of size 5x5] Variable containing: 0.8749 0.5990 0.6844 0.0000 0.0000 0.6516 0.0000 1.5117 0.5734 0.0072 0.1286 1.4171 0.0796 1.0355 0.0000 0.0000 0.0000 0.0312 0.0999 0.0000 1.0401 1.0599 0.0000 0.0000 0.0000 [torch.FloatTensor of size 5x5] 1 1 1 0 0 1 0 1 1 1 1 1 1 1 0 0 0 1 1 0 1 1 0 0 0

原文出自騰訊雲技術社區

原文鏈接https://www.qcloud.com/community/article/314920