tensorflow自定义op,如何在Tensorflow中使用自定义python函数预取数据

╰半橙微兮° 2022-08-30 11:17 224阅读 0赞

![Image 1][]

I am trying to prefetch training data to hide I/O latency. I would like to write custom Python code that loads data from disk and preprocesses the data (e.g. by adding a context window). In other words, one thread does data preprocessing and the other does training. Is this possible in TensorFlow?

Update: I have a working example based on @mrry's example.

import numpy as np

import tensorflow as tf

import threading

BATCH\_SIZE = 5

TRAINING\_ITERS = 4100

feature\_input = tf.placeholder(tf.float32, shape=\[128\])

label\_input = tf.placeholder(tf.float32, shape=\[128\])

q = tf.FIFOQueue(200, \[tf.float32, tf.float32\], shapes=\[\[128\], \[128\]\])

enqueue\_op = q.enqueue(\[label\_input, feature\_input\])

label\_batch, feature\_batch = q.dequeue\_many(BATCH\_SIZE)

c = tf.reshape(feature\_batch, \[BATCH\_SIZE, 128\]) + tf.reshape(label\_batch, \[BATCH\_SIZE, 128\])

sess = tf.Session()

def load\_and\_enqueue(sess, enqueue\_op, coord):

with open('dummy\_data/features.bin') as feature\_file, open('dummy\_data/labels.bin') as label\_file:

while not coord.should\_stop():

feature\_array = np.fromfile(feature\_file, np.float32, 128)

if feature\_array.shape\[0\] == 0:

print('reach end of file, reset using seek(0,0)')

feature\_file.seek(0,0)

label\_file.seek(0,0)

continue

label\_value = np.fromfile(label\_file, np.float32, 128)

sess.run(enqueue\_op, feed\_dict=\{feature\_input: feature\_array,

label\_input: label\_value\})

coord = tf.train.Coordinator()

t = threading.Thread(target=load\_and\_enqueue, args=(sess,enqueue\_op, coord))

t.start()

for i in range(TRAINING\_ITERS):

sum = sess.run(c)

print('train\_iter='+str(i))

print(sum)

coord.request\_stop()

coord.join(\[t\])

解决方案

This is a common use case, and most implementations use TensorFlow's queues to decouple the preprocessing code from the training code. There is a tutorial on how to use queues, but the main steps are as follows:

Define a queue, q, that will buffer the preprocessed data. TensorFlow supports the simple tf.FIFOQueue that produces elements in the order they were enqueued, and the more advanced tf.RandomShuffleQueue that produces elements in a random order. A queue element is a tuple of one or more tensors (which can have different types and shapes). All queues support single-element (enqueue, dequeue) and batch (enqueue\_many, dequeue\_many) operations, but to use the batch operations you must specify the shapes of each tensor in a queue element when constructing the queue.

Build a subgraph that enqueues preprocessed elements into the queue. One way to do this would be to define some tf.placeholder() ops for tensors corresponding to a single input example, then pass them to q.enqueue(). (If your preprocessing produces a batch at once, you should use q.enqueue\_many() instead.) You might also include TensorFlow ops in this subgraph.

Build a subgraph that performs training. This will look like a regular TensorFlow graph, but will get its input by calling q.dequeue\_many(BATCH\_SIZE).

Start your session.

Create one or more threads that execute your preprocessing logic, then execute the enqueue op, feeding in the preprocessed data. You may find the tf.train.Coordinator and tf.train.QueueRunner utility classes useful for this.

Run your training graph (optimizer, etc.) as normal.

EDIT: Here's a simple load\_and\_enqueue() function and code fragment to get you started:

\# Features are length-100 vectors of floats

feature\_input = tf.placeholder(tf.float32, shape=\[100\])

\# Labels are scalar integers.

label\_input = tf.placeholder(tf.int32, shape=\[\])

\# Alternatively, could do:

\# feature\_batch\_input = tf.placeholder(tf.float32, shape=\[None, 100\])

\# label\_batch\_input = tf.placeholder(tf.int32, shape=\[None\])

q = tf.FIFOQueue(100, \[tf.float32, tf.int32\], shapes=\[\[100\], \[\]\])

enqueue\_op = q.enqueue(\[feature\_input, label\_input\])

\# For batch input, do: