Caffe2 - (十三) 基于 Python 创建 Operator

Caffe2 - Python Operator

类似于 Caffe 基于 Python 定制 CaffeLayers， Caffe2 也提供了使用 Python 来自定义 Caffe2 Operators.

Forward Python Operator

Caffe2 提供了 high-level 接口，用于 Python ops 创建 —— Net.Python() 接口.

from caffe2.python import core, workspace
import numpy as np

def f(inputs, outputs):
    outputs[0].feed(2 * inputs[0].data)

workspace.ResetWorkspace()
net = core.Net("tutorial")
net.Python(f)(["x"], ["y"]) ##
workspace.FeedBlob("x", np.array([3.]))
workspace.RunNetOnce(net)
print(workspace.FetchBlob("y"))
# [6.]

Caffe2 的 net.Python() 函数是可调用的，类似与其它 operators.

如 net.Python(f)(["x"], ["y"]) 是一个新添加到网络的 Python operator，其输入是 x，输出是 y. 之后，即可保存 net.Python() 的输出；也可多次调用以添加多个 Python operators (可能分别是不同的输入和输出).

Python operators 的函数 f包括连个参数：输入 inputs 列表和输出 outputs 列表. 当执行 operator 时，Caffe2 blobs 被转换为列表元素. 对于 CPU tensor blobs，会被转换为 TensorCPU object，类似与 Numpy arrays.

Caffe2 CPU tensor、Python TensorCPU object 和 Numpy array 的关系：

• C++ tensor objects 和 Numpy objects 间是自动转换的，由 PyBind library 来控制管理.
• 当创建一个 TensorCPU wrapper 时，伴随创建一个新的 Numpy array object，其与对应的 Caffe2 CPU tensor共享相同的内存. 该 Numpy array 可以在 Python 通过 TensorCPU object 的 .data 操作来读取和查看.
• 尽管 Numpy array 和 Caffe2 tensor 共享相同的内存，但是其它的 Caffe2 tensor 数据(如，shape) 是与 Numpy array 分开存储的. 而且，在 operator 函数执行时，Numpy 可能复制并重新分配其 array 到内存的不同位置(如，array 的 resize). 因此，需谨记，在创建 Python operator 代码时，需要保证 Caffe2 和 Numpy 同步输出 tensors.
• TensorCPU feed 数据方式是 Numpy tensor，将 Caffe2 tensor 进行 resize，并将 Numpy 的 tensor 数据复制到 Caffe2 tensor.
• 另一种正确设置 Caffe2 的输出 tensor 方式是调用 reshape 函数，来对相应的 TensorCPU output 处理，然后将 Python 的数据复制到输出的 .data tensor， 如：

def f_reshape(inputs, outputs):
   outputs[0].reshape(inputs[0].shape)
   outputs[0].data[...] = 2 * inputs[0].data

workspace.ResetWorkspace()
net = core.Net("tutorial")
net.Python(f_reshape)(["x"], ["z"])
workspace.FeedBlob("x", np.array([3.]))
workspace.RunNetOnce(net)
print(workspace.FetchBlob("z"))

• reshape 函数将 Caffe2 tensor 进行更新，然后调用 .data 属性返回与 Caffe2 tensor 共享内存的 Numpy array. f_reshape 即是将输出复制到共享内存的位置.

net.Python()函数也可以有其它参数. 当 pass_workspace=True使， workspace 会被传递到 operator 的 Python 函数，如：

def f_workspace(inputs, outputs, workspace):
    outputs[0].feed(2 * workspace.blobs["x"].fetch())

workspace.ResetWorkspace()
net = core.Net("tutorial")
net.Python(f_workspace, pass_workspace=True)([], ["y"])
workspace.FeedBlob("x", np.array([3.]))
workspace.RunNetOnce(net)
print(workspace.FetchBlob("y"))
# [6.]

Backward - Gradient Python Operator

Caffe2 的 net.Python() 的另一个重要参数是 grad_f，其是对应的 gradient operator 的Python 函数.

def f(inputs, outputs):
            outputs[0].reshape(inputs[0].shape)
            outputs[0].data[...] = inputs[0].data * 2

def grad_f(inputs, outputs):
    # Ordering of inputs is [fwd inputs, outputs, grad_outputs]
    grad_output = inputs[2]

    grad_input = outputs[0]
    grad_input.reshape(grad_output.shape)
    grad_input.data[...] = grad_output.data * 2

workspace.ResetWorkspace()
net = core.Net("tutorial")
net.Python(f, grad_f)(["x"], ["y"])
workspace.FeedBlob("x", np.array([3.]))
net.AddGradientOperators(["y"])
workspace.RunNetOnce(net)
print(workspace.FetchBlob("x_grad"))
# [2.]

当指定 gradient 函数，并用 net.Python 进行调用时，会再注册一个序列化 gradient 函数，其可以被对应的 gradient Python operator(PythonGradient) 来使用.

gradient operator 有两个参数： 输入 input 列表和输出 output 列表.

输入 input 列表包含所有的 forward 函数输入，以及其 outputs，其 forward 函数输出 outputs 的 gradients.

输出 output 列表 包含 forward 函数输入 inputs 的 gradient.

net.Python 的 grad_output_indices/grad_input_indices 指定了梯度 gradient/input blobs 的 gradient 函数的 reads/writes 索引.

GPU tensors：

PythonOp 的实现是 CPU 的，采用 Numpy arrays，在 CPU 内存中.

如果需要 PythonOp 是 GPU tensors，Caffe2 定义了 PythonOp 的 CUDA 版本，基于 GPUFallbackOp. 该 operator 对 CPU-operator 进行wrap，并添加 GPU-to-CPU(和反向) copy operations.

因此，在使用 CUDA PythonOp 时，所有的输入 input CUDA tensors 被自动的复制到 CPU 内存，然后所有的 CPU 输出 output 再被复制回 GPU. （GPU-CPU-GPU）