[TensorFlow series] [9] Use tf.py_func to customize the operator

This article addresses the following questions:

1. How to define placeholder of list type?

2. How to wrap ordinary python functions into TensorFlow operators and add them to the NN network?

 

See the code for details:

import tensorflow as tf
import numpy as np

def gen_tfrecords():
    with tf.python_io.TFRecordWriter(r"D:\my.tfrecords") as tf_writer:
        features = {}
        features['scale'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[15]))
        xmin = []
        xmax = []
        ymin = []
        ymax = []
        for i in range(2):
            xmin.append(float(i))
            xmax.append(float(i+500))
            ymin.append(float(i))
            ymax.append(float(i+500))
        # 变长数据以list形式存储
        features['xmin'] = tf.train.Feature(float_list=tf.train.FloatList(value=xmin))
        features['xmax'] = tf.train.Feature(float_list=tf.train.FloatList(value=xmax))
        features['ymin'] = tf.train.Feature(float_list=tf.train.FloatList(value=ymin))
        features['ymax'] = tf.train.Feature(float_list=tf.train.FloatList(value=ymax))
        tf_features = tf.train.Features(feature=features)
        tf_example = tf.train.Example(features=tf_features)
        tf_serialized = tf_example.SerializeToString()
        tf_writer.write(tf_serialized)
gen_tfrecords()
def parse_tf(example_proto):
    dics = {}
    #定长数据解析
    dics['scale'] = tf.FixedLenFeature(shape=[], dtype=tf.int64)
    #列表数据解析
    dics['xmin'] = tf.VarLenFeature(tf.float32)
    dics['xmax'] = tf.VarLenFeature(tf.float32)
    dics['ymin'] = tf.VarLenFeature(tf.float32)
    dics['ymax'] = tf.VarLenFeature(tf.float32)
    parse_example = tf.parse_single_example(serialized=example_proto,features=dics)
    xmin = parse_example['xmin']
    xmax = parse_example['xmax']
    ymin = parse_example['ymin']
    ymax = parse_example['ymax']
    scale = parse_example['scale']

    return scale,xmin,xmax,ymin,ymax

def scale_image(scale):
    w = 10
    h = 10
    w = w*scale
    h = h*scale
    return w,h
def scale_image2(scale):
    w = 10
    h = 10
    w = w*scale
    h = h*scale
    lst = [w,h]
    #如果想要返回一个list，需要将其封装为一个ndarray
    return np.array(lst)
def image_s(scale):
    w = 10
    h = 10
    w = w*scale
    h = h*scale
    return w*h

def calc_image_s(xmin,xmax,ymin,ymax):
    ss = []
    for i in range(len(xmin)):
        s = (xmax[i]-xmin[i])*(ymax[i]-ymin[i])
        ss.append(s)
    return np.array(ss)

scale_p = tf.placeholder(dtype=tf.int64)
#如果placeholder的shape不写，则可表示各种类型的数据，这里可用于表示list类型的数据
x_min_p = tf.placeholder(dtype=tf.float32)
x_max_p = tf.placeholder(dtype=tf.float32)
y_min_p = tf.placeholder(dtype=tf.float32)
y_max_p = tf.placeholder(dtype=tf.float32)

#tf.py_func用来将不同python函数包裹成TensorFlow算子，返回值是tensor，Tout表示函数返回值的类型，单个返回值不用[]，多个返回值，要用[]
nw,nh = tf.py_func(scale_image,inp=[scale_p],Tout=[tf.int64,tf.int64])
nw_nh = tf.py_func(scale_image2,inp=[scale_p],Tout=tf.int64)
s = tf.py_func(image_s,inp=[scale_p],Tout=tf.int64)
ss= tf.py_func(calc_image_s,inp=[x_min_p,x_max_p,y_min_p,y_max_p],Tout=tf.float32)

two = tf.constant(value=2,dtype=tf.float32)
s2 = tf.multiply(ss,two)

dataset = tf.data.TFRecordDataset(r"D:\my.tfrecords")
dataset = dataset.map(parse_tf).batch(1).repeat(1)

iterator = dataset.make_one_shot_iterator()

next_element = iterator.get_next()
with tf.Session() as session:
    scale, xmin, xmax, ymin, ymax = session.run(fetches=next_element)
    w,h = session.run(fetches=[nw,nh],feed_dict={scale_p:scale})
    print(w,h)

    w_h = session.run(fetches=[nw_nh], feed_dict={scale_p: scale})
    print(w_h)

    s1 = session.run(fetches=[s], feed_dict={scale_p: scale})
    print(s1)

    s1 = session.run(fetches=[ss], feed_dict={x_min_p:xmin.values,x_max_p:xmax.values,y_min_p:ymin.values,y_max_p:ymax.values})
    print(s1)

    s22 = session.run(fetches=[s2], feed_dict={x_min_p: xmin.values, x_max_p: xmax.values, y_min_p: ymin.values, y_max_p: ymax.values})
    print(s22)

The result is as follows:

[150] [150]
[array([[150],
       [150]], dtype=int64)]
[array([22500], dtype=int64)]
[array([250000., 250000.], dtype=float32)]
[array([500000., 500000.], dtype=float32)]

This method is used when defining a new op. According to the official website, this method currently does not support distribution and model preservation. But for auxiliary ops, it is basically enough, such as the generation of anchors in faster r-cnn and the generation of labels during RPN training.

tf.py_func requires the wrapped function, both input and output are ndarray