tensorflow与python交互系列，tf.py_function()、tf.py_func、tf.numpy_function()（一）

前言：前面在介绍使用tensorflow进行data pipeline的时候，遇到了一些问题，特意整理了两篇文章，请参见：

tfrecord文件的map在使用的时候所踩的坑总结（map、py_function、numpy_function）

tensorflow2.x之由dataset.map引发出的关于tf.py_function以及tf.numpy_function问题

本系列文章将深入详解这是哪个函数，即tf.py_function()、tf.py_func、tf.numpy_function()，其中

tf.py_function()、tf.numpy_function()是tensorflow2.x版本的，我使用的版本是tensorflow2.1 

tf.py_func () 是tensorflow1.x版本的，我们使用的是tensorflow1.13.1

理解的核心，将上面三者包装的python函数理解成tensorflow的operation，只不过函数的参数与返回值有所限制。

一、tensorflow2.x的tf.py_function()深入详解

1.1 描述与声明

作用：Wraps a python function into a TensorFlow op that executes it eagerly.

将一个Python函数包装成一个tensorflow的operation操作，然后在eager模式下执行，它的作用是包装自定义的Python函数的，那什么时候回用到呢？

当我们需要进行一些tensorflow没有的操作，需要我们自己编写运算逻辑的时候，我们就需要用到这个函数。

tf.py_function(
    func, inp, Tout, name=None
)

 

参数解析
func	即需要包装的函数，是一个python函数，这个函数接受一个或者是多个Tensor对象作为函数参数，并且返回一个或者是多个Tensor，当然也可以返回None，需要注意的函数的参数定义的Tensor需要与inp所指定的参数个数相符合；返回的Tensor的元素的类型需要与Tout所指定的元素类型相符合。
inp	A list of Tensor objects. 由一个或者是几个Tensor组成的list
Tout	由tensorflow的数据类型组成的list或者是tuple，要与func函数的返回值是匹配的，如果是空的话则代表func函数无返回值，即None。
name	A name for the operation (optional).

返回值：

该函数的返回值就是定义的python函数func运算之后的结果，在tensorflow2.x中，尅直接获取该返回值，返回值可以是一个或者是多个，但是需要注意的是，不管函数func返回的是什么内容，py_function的返回值总是会自动将其包装成与之对应的匹配的Tensor再返回，在tf2.x中，自动包装成EagerTensor，可以直接获取值。

 

1.2 简单示例一

问题描述，现在有一个分段函数，函数定义如下：

当x<m的时候，f(x)=x**2

当x>m的时候，f(x)=m**2(1-2logm)+log(x**2)

当然我们可以直接使用tensorflow的内定义算子来实现上面的运算，但是我们通过下面的方式来实现

import tensorflow as tf

def log_huber(x, m):
    '''
    因为是tensorflow2.1,所以默认使用的是EagerTensor，我们可以在里面直接使用python的条件判断，而不用使用
    tf.while_loop、tf.cond等条件判断
    '''
    if tf.abs(x) <= m:
        return x**2  # 无论这里返回什么，总是会被包装成默认的与之匹配的Tensor返回，这是自动实现的
    else:
        return m**2 * (1 - 2 * tf.math.log(m) + tf.math.log(x**2))

x = tf.Variable(initial_value = 1.5)
m = tf.Variable(initial_value = 2.0)

with tf.GradientTape(persistent=True) as g:
    g.watch(x)
    y = tf.py_function(func=log_huber, inp=[x, m], Tout=tf.float32)
    print(y)  # 应该是1.5**2 = 2.25
    dy_dx = g.gradient(y, x)     # y对x求导，应该是2x，即3.0
    print(dy_dx)

1.3 简单示例二

比如我要实现这样一个算子，来对两个张量tensor进行如下操作，

第一，他们的和减去它们的差

第二，他们的和加上他们的差

并且同时返回这两个结果，代码如下：

def add_sub(x, y):
    '''
    在此函数中使用纯python编程方式
    '''
    x_ = x.numpy()
    y_ = y.numpy()
    result1 = x_ + y_ - (x_ - y_)
    result2 = x_ + y_ + (x_ - y_)  
    # 返回的就是普通的python对象，但是它会自动转化成tensor来作为最终的结果，是自动完成的
    return result1,result2

x = tf.Variable(initial_value = [10,20,30])
y = tf.Variable(initial_value = [100,200,300])

y1,y2 = tf.py_function(func=add_sub, inp=[x, y], Tout=[tf.int32,tf.int32])
print(y1)  
print(y2)  
'''
tf.Tensor([200 400 600], shape=(3,), dtype=int32)
tf.Tensor([20 40 60], shape=(3,), dtype=int32)
'''

1.4 简单示例三

现在有一个字符串，这个字符串是一个文件名称，我想使用python方法来读取这个文本文件里面的内容，然后返回它的读取的值，比如文件名称如下：file1.txt，里面的内容是1,2,3,4,5

读取里面的内容，然后返回它们求和的结果。

def read_txt(filename):
    filename_ = filename.numpy()  # 通过numpy获取每一条样本的文件名称
    filename_ = filename_.decode("utf-8") # 需要将bytes转化成str，得到 file1.txt、file2.txt
    
    # 现在可以使用纯python操作了，获取文件路径
    filename = "./" + filename_
    f =  open(filename,mode="r")
    s =f.readline()
    x_ =s.split(',')
    result =[]
    for i in x_:
        result.append(int(i))
    
    result_sum = sum(result)
    return result,result_sum

x = tf.Variable(initial_value = "file1.txt")

y1,y2 = tf.py_function(func=read_txt, inp=[x], Tout=[tf.int32,tf.int32])
print(y1)  
print(y2) 
'''
tf.Tensor([1 2 3 4 5], shape=(5,), dtype=int32)
tf.Tensor(15, shape=(), dtype=int32)
'''

1.5 注意事项

（1）自定义的python函数func函数体中的内容不能够被序列化为GraphDef，因此这里面的内容不能够通过序列化保存graph结构，也不能够restore，这是很重要的；

（2）函数的定义必须与调用包装函数py_function在同一个空间中，如果使用tensorflow的分布式训练，必须要进行相关的设置，参考官方文档如下：

https://www.tensorflow.org/api_docs/python/tf/py_function

二、tensorflow1.x的tf.py_func()详解

2.1 描述与声明

作用：依然是将python函数包装成tensorflow的operation

tf.py_func(
    func, inp, Tout, stateful=True, name=None
)

 

Args
func	是一个Python函数，需要格外注意的是，这个函数的参数必须是一个或者是多个numpy.ndarray类型，不再是Tensor，它的返回值也是一个或者是多个array数组，并且函数参数必须与inp的参数相匹配，返回的array的数据类型必须与Tout指定的类型相匹配。 因为它的参数和返回值都必须是numpy数组类型，所以不再像tf2.x的py_function那么灵活了，使用起来不太好用，下面的例子会说明。
inp	A list of Tensor objects.由Tensor组成列表，由于func只接受array，所以我们一般传递给inp的都只是placeholder占位符
Tout	返回的Tensor的类型

2.2 看一个简单的例子

import tensorflow as tf
import os
import numpy as np

def my_func(x):  # 参数X是numpy类型，所以在下面传递的时候不能再传递Tensor给inp参数，只能放一个空的placeholder
    print(type(x))  # <class 'numpy.ndarray'>
    return np.sinh(x)

input = tf.placeholder(tf.float32)
y = tf.py_func(my_func, [input], tf.float32)

with tf.Session() as sess:
    y_ = sess.run(y, feed_dict={input:[1,2,3,4]})
    print(y_)
'''
<class 'numpy.ndarray'>
[ 1.1752012  3.6268604 10.017875  27.289919 ]
'''

2.3 再看一个案例

描述，现在我需要自己定义一个这样的功能函数，根据传递的参数，来对一个数组进行按行求和或者是按列求和，代码如下：

def my_func(x,m):
    '''
    x: 是一个array的二维数组
    m: 是一个数值，但要写成array的形式，如[1]，因为他只能接受array的参数
    '''
    if m[0] == 1.0:
        return np.sum(x,axis=0)
    else:
        return np.sum(x,axis=1)

input_x = tf.placeholder(tf.float32,shape=(2,3))
input_m = tf.placeholder(tf.float32,shape=(1,))

y = tf.py_func(my_func, [input_x,input_m], tf.float32)

x_ = np.array([[1.0,2.0,3.0],[4.0,5.0,6.0]])
m_ =np.array([1.0])

with tf.Session() as sess:
    y_ = sess.run(y, feed_dict={input_x:x_,input_m:m_})
    print(y_)
'''
[5. 7. 9.]
'''

三、tensorflow2.x的tf.numpy_function()

这个函数实际上有点类似于py_func，它所包装的函数只能够接受numpy.array作为函数参数，并且返回的也是numpy.array，这里就不再赘述了。