pytorch 添加c++实现的自定义op

pytorch已经基本实现了常见的各种op,然而,当想实现一个pytorch中没有的op时,有两种方式。一种方式是这个op可以由pytorch中已有的op进行组合而成,因此只需要使用python接口进行组合就可以了。反之,就必须使用c++或者cuda实现该op,然后添加到pytorch中。本文将介绍添加c++实现的自定义op,因为我还不会cuda : (

本文介绍使用python的setuptools将c++实现的op添加到pytorch中。首先要用c++实现定义的op。比如想实现一个op为 z=3x-y 。头文件为my_op.h

#include <torch/extension.h> //这一句是无论要实现任何op都必须添加的
#include <vector>

//前向传播
torch::Tensor my_op_forward(const torch::Tensor& x, const torch::Tensor& y);
//反向传播
std::vector<torch::Tensor> my_op_backward(const torch::Tensor& gradOutput);

源文件为my_op.cpp

#include "my_op.h"

torch::Tensor my_op_forward(const torch::Tensor& x,                             const torch::Tensor& y) {     
     AT_ASSERTM(x.sizes() == y.sizes(), "x must be the same size as y");
     torch::Tensor z = torch::zeros(x.sizes());
     z = 3 * x - y;
     return z; } 

std::vector<torch::Tensor> my_op_backward(const torch::Tensor& gradOutput) {
     torch::Tensor gradOutputX = 3 * gradOutput * torch::ones(gradOutput.sizes());
     torch::Tensor gradOutputY = -1 * gradOutput * torch::ones(gradOutput.sizes());
     return {gradOutputX, gradOutputY}; } 

// pybind11 绑定 
PYBIND11_MODULE(my_op_api, m) {
     m.def("forward", &my_op_forward, "MY_OP forward");
     m.def("backward", &my_op_backward, "MY_OP backward"); 
} 

其中最后的PYBIND11_MODULE是用来将C++函数绑定到python上的。其中第一个参数my_op_api为要生成的python模块名,以后import my_op_api就可以调用该op了。第二个参数固定为m
函数体中的两个语句分别是绑定前向传播与反向传播到实现的两个函数上。

然后编写setup.py,用来构建pytorch的c++扩展。

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(name='my_op_api',
      version='0.l',
      ext_modules=[CppExtension('my_op_api', sources=['my_op.cpp'], extra_compile_args=['-std=c++11'])],
      cmdclass={'build_ext':BuildExtension})

其中,setup中的name以及CppExtension中第一个参数(也是name)要和PYBIND11_MODULE里设的模块名保持一致,这里都是my_op_api。CppExtension中的extra_compile_args=[‘-std=c++11’]是趟坑发现的,不加的话gcc可能会报n多错(pytorch是用c++11编译的,因此这里用gcc编译的时候也要使用c++11)

然后运行python setup.py install,如果没有问题的话,就生成了所需的python模块。可以从输出的信息看到该模块在所在python环境下的site-packages文件夹下,以.egg结尾。另外,在当前目录下会有3个文件夹生成,build、dist、my_op_api.egg-info,其中dist下也有.egg文件,可以发布到其它python环境。

然后就可以在python中import my_op_api进行调用扩展的op了。这里需要注意一点的是,在import 自定义的op之前,必须先import torch。 但是,这样的op和我们日常使用的还是不太一样,这时需要将它包装为pytorch中的函数和模块,以便我们像使用其它模块一样使用自定义的op。要包装为模块,首先包装成函数。

包装成函数,需要继承torch.autograd.Function。然后包装成模块,需要继承torch.nn.Module

import torch
from torch.autograd import Function
from torch.nn import Module
import my_op_api

class MyOpFunction(Function):
    @staticmethod
    def forward(ctx, x, y):
        #如果有一些信息,需要在梯度反向传播时用到,可以使用ctx.save_for_backward()进行保存
        return my_op_api.forward(x, y)
    @staticmethod
    def backward(ctx, gradOutput):
        #如果在forward中保存了信息,可以使用ctx.saved_tensors取回
        grad_x, grad_y = my_op_api.backward(gradOutput)
        return grad_x, grad_y

class MyOpModule(Module):
    def __init__(self):
        super(MyOpModule, self).__init__()
    def forward(self, input_x, input_y):#只需要定义forward的函数就可以了
        return MyOpFunction.apply(input_x, input_y)

tensorflow 获取所有tensor、op的name

获取所有tensor(每个op的输出张量)的name:

for tensor in tf.contrib.graph_editor.get_tensors(tf.get_default_graph()):
    print(tensor.name)

获取所有op及其输入输出的name:

for node in sess.graph_def.node:
      print(node) 

tensorflow调试的一种方式

由于tensorflow采用构建图,在会话中再运行图的方式,使得调试非常麻烦。尤其是想获取网络中间某个tensor的时候。对此,可以采用以下方法

在构建完成网络所有结构之后,G = tf.get_default_graph()获取图,然后使用tensr = G.get_tensor_by_name(“TensorName:0”)的方式获取想要的tensor,再sess.run该tensor就可以获取它的值了。TensorName可以通过查看网络ckpt、pb文件的方式获取。

windows、ubuntu下校验文件

windows的powershell中自带了校验工具certutil,可以生成文件的MD5值、SHA1值、SHA256值。具体命令如下:

certutil -hashfile xxx MD5
certutil -hashfile xxx SHA1
certutil -hashfile xxx SHA256

xxx表示文件路径

ubuntu下可以使用md5sum/sha256sum生成某待测文件的哈希值

md5sum xxx
或者
sha256sum xxx

tensorflow 导入pb模型进行前向推导

tensorflow一般使用pb文件进行前向推导(在非部署环境使用ckpt也可以)载入pb文件到图的函数

def load_pb_to_graph(sess, pb_file):
with tf.gfile.FastGFile(pb_file, "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
sess.graph.as_default()
tf.import_graph_def(graph_def,name="")

在会话中,调用该函数,并根据名称获取输入和输出的tensor,然后就可以sess.run进行前向推导

load_pb_to_graph(sess, "xxnet.pb")
inputs = tf.get_default_graph().get_tensor_by_name("xxnet/input:0")
outputs = tf.get_default_graph().get_tensor_by_name("xxnet/score:0")
scores = sess.run([outputs], feed_dict={inputs: input_data})

tensorflow 将ckpt文件导出为pb文件

tensorflow训练时将模型保存为ckpt文件,它包含了网络结构、网络权重、训练过程中间变量等等信息。而网络部署一般是使用pb文件,它将变量保存为常量,以及网络前向传播的所有必要结构。如何将ckpt文件导出为pb文件?

首先,使用tfrecord训练的ckpt一般包含读取训练tfrecord文件的结构,而这是pb文件所不需要的。pb文件通常使用placeholder接受输入。因此,要以placeholder为输入重新定义一遍网络结构(通常就是调用一次网络构建函数)。假设为
output = xxnet(input_placeholder)
要获取输出节点的名称
output_nd_name = output.op.name

然后,载入ckpt的权重
saver = tf.train.Saver()
saver.restore(sess, “xxnet.ckpt”)

然后,将其中的变量转化为常量,保存模型

out_graph_def = tf.graph_util.convert_variables_to_constants(
    sess=sess,
    input_graph_def=sess.graph_def,
    output_node_names=[output_nd_name]
)
with tf.gfile.GFile("xxnet.pb","wb") as f:
    f.write(out_graph_def.SerializeToString())

tensorflow 训练网络的一般步骤

本文不针对tensorflow2.0。首先要构建数据的输入,一般是将数据转化为pb格式

然后构建自己的网络,并构建损失函数的节点。构建网络有多种方式,可以用代码构建(利用slim、keras等高级api,或者基础的api,或者已有的代码),也可以从ckpt.meta中载入网络结构(断点继续训练等情况)tf.train.import_meta_graph(“xxx.ckpt.meta”)。这里要注意,一般训练时会同时进行网络在验证集上的测试,比如每训练n步后在训练集上进行测试。因此构建网络需要同时构建一个验证网络,共享训练网络的变量权重。构建验证网络时要在variable_scope中设置reuse=True。

定义优化器,如opt=tf.train.AdamOptimizer()
将优化器应用在损失节点上计算梯度。grads=opt.compute_gradients(L)
梯度下降优化节点 apply_grad_op = opt.apply_gradients(grads)

训练模型需要保存,定义一个saver
saver = tf.train.Saver(max_to_keep=10) 最多保留10个ckpt
在训练时,使用saver.save(sess, “xxx.ckpt”, global_step=step)保存ckpt文件

希望在训练时看到训练过程, 使用tf.summary.scalar 添加想要的变量到训练过程日志中。
如 tf.summary.scalar(“training loss”, L)添加训练损失到训练过程。然后定义summary_op = tf.summary.merge_all()
然后要定义一个summary_writer
summary_writer = tf.summary.FileWriter(logdir, sess.graph)
训练时,每隔n步,使用summary_writer.add_summary(sess.run(summary_op), step)保存训练过程日志
训练开始后,就可以使用tensorboard查看训练过程了

训练过程一般在一个for循环中进行,
sess.run([apply_grad_op])进行网络的训练
在这个循环中,还要进行上面所说的保存ckpt文件、训练日志

tensorflow 初始化新增变量,保持载入的预训练模型权重不变

我们经常会遇到这样一些问题,想要使用一些预训练好的模型,然后在其基础上进行一些增减,以适应新的任务。在开始训练之前,要对所有的新增变量进行初始化,但是要保持预训练模型中已有的权重不变,即只初始化新增变量。而如果使用tensorflow中提供的saver.restore(sess, ckpt_path),会报找不到新增节点的错误!记录一下这种情况要如何处理。

首先,在对预训练模型进行增减之前,先进行saver.restore(sess, ckpt_path)载入预训练权重,然后再进行对网络结构的增减(这里可以使用tf.get_default_graph().get_tensor_by_name(tensor_name)获取到原网络中的tensor,来进行新增节点)

在增加完新增节点之后,要初始化这些新增节点权重变量。接下来就是最关键的一步,获取网络中所有未初始化的权重变量。

def get_uninitialized_variables(sess):
global_vars = tf.global_variables()
is_not_initialized = sess.run([tf.is_variable_initialized(var) for var in global_vars])
not_initialized_vars = [v for (v, f) in zip(global_vars, is_not_initialized) if not f]
print([str(i.name) for i in not_initialized_vars])
return not_initialized_vars

然后在会话中,初始化这些变量

sess.run(tf.variables_initializer(get_uninitialized_variables(sess)))

接下来就可以愉快地进行训练了