总说
没办法, 出来混总是要还的, 不会写点底层代码没法混啊. 废话不多说, 简单来说, 有时候我们需要写一些自定义的操作, 这些操作如果用python写会很慢, 我们需要用CUDA写, 然后这些操作与python绑定, 以供python端调用.
主要是简略拿出 https://pytorch.org/tutorials/advanced/cpp_extension.html 的东西, 根据实践, 补充了一些东西(否则, 直接看官方文档可能会有一些地方需要花点实践), 没毛病. 看这个博客, 可以的.
示例程序
class LLTM(torch.nn.Module):
def __init__(self, input_features, state_size):
super(LLTM, self).__init__()
self.input_features = input_features
self.state_size = state_size
# 3 * state_size for input gate, output gate and candidate cell gate.
# input_features + state_size because we will multiply with [input, h].
self.weights = torch.nn.Parameter(
torch.empty(3 * state_size, input_features + state_size))
self.bias = torch.nn.Parameter(torch.empty(3 * state_size))
self.reset_parameters()
def reset_parameters(self):
stdv = 1.0 / math.sqrt(self.state_size)
for weight in self.parameters():
weight.data.uniform_(-stdv, +stdv)
def forward(self, input, state):
old_h, old_cell = state
X = torch.cat([old_h, input], dim=1)
# 自定义C++扩展, 可以让这些操作, 变成一个fused的版本
# Compute the input, output and candidate cell gates with one MM.
gate_weights = F.linear(X, self.weights, self.bias)
# Split the combined gate weight matrix into its components.
gates = gate_weights.chunk(3, dim=1)
input_gate = F.sigmoid(gates[0])
output_gate = F.sigmoid(gates[1])
# Here we use an ELU instead of the usual tanh.
candidate_cell = F.elu(gates[2])
# Compute the new cell state.
new_cell = old_cell + candidate_cell * input_gate
# Compute the new hidden state and output.
new_h = F.tanh(new_cell) * output_gate
return new_h, new_cell
import torch
X = torch.randn(batch_size, input_features)
h = torch.randn(batch_size, state_size)
C = torch.randn(batch_size, state_size)
rnn = LLTM(input_features, state_size)
new_h, new_C = rnn(X, (h, C))
C++扩展
简单来说, 我们写完C++程序后, python要用这些程序, 可以用pybind11
. 然而, 安装pybind11
需要用到pytest
, 而pytest
貌似只能在python3.5以上才能运行. 所以我们先弄个基于python3的Anaconda. 装好pytorch之后(随便你咋装上的). 然后再进行后续操作.
pybind11安装
git clone https://github.com/pybind/pybind11.git
pip install pytest
注意一下, 这里的pip -V
最好是显示anaconda3
中的pip
, 从而确保下载的pytest
是python3版本.
cd pybind11
mkdir build
cd build
cmake ..
make check -j 4
编译好的动态库是test
目录下的so文件.
pytorch的相关事宜
我们要写pytorch扩展, 得下载pytorch源代码.
git clone --recursive https://github.com/pytorch/pytorch
然后我们在pytorch根目录下, 建立一个文件夹, 比如 lltm-extension.
并在该文件夹下, 建立setup.py</