Pytorch学习 (二十一) ------自定义C++/ATen扩展

总说

没办法, 出来混总是要还的, 不会写点底层代码没法混啊. 废话不多说, 简单来说, 有时候我们需要写一些自定义的操作, 这些操作如果用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</