multithreading - 使用多线程编写 tfrecord 并不像预期的那样快

Question

尝试编写 w/和 w/o 多线程的 tfrecord，发现速度差异不大(w/4 线程:434 秒；w/o 多线程 590 秒)。不确定我是否正确使用它。有没有更好的方法来更快地编写 tfrecord？

import tensorflow as tf 
import numpy as np 
import threading 
import time 


def generate_data(shape=[15,28,60,1]):
    return np.random.uniform(size=shape)


def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def write_instances_to_tfrecord(tfrecord_file, filenames):
    tfrecord_writer = tf.python_io.TFRecordWriter(tfrecord_file)
    for i, filename in enumerate(filenames):
        curr_MFCC = generate_data()
        curr_MFCC_raw = curr_MFCC.tostring()
        curr_filename_raw = str(filename)+'-'+str(i)
        example = tf.train.Example(features=tf.train.Features(
            feature={
            'MFCC': _bytes_feature(curr_MFCC_raw),
            'filename': _bytes_feature(curr_filename_raw)
            })
        )
        tfrecord_writer.write(example.SerializeToString())
    tfrecord_writer.close()


def test():
    threading_start = time.time()
    coord = tf.train.Coordinator()
    threads = []
    for thread_index in xrange(4):
        args = (str(thread_index), range(200000))
        t = threading.Thread(target=write_instances_to_tfrecord, args=args)
        t.start()
        threads.append(t)
    coord.join(threads)
    print 'w/ threading takes', time.time()-threading_start

    start = time.time()
    write_instances_to_tfrecord('5', range(800000))
    print 'w/o threading takes', time.time()-start

if __name__ == '__main__':
    test()

Answer 1

使用 python 线程时，由于 cPython 实现中的 GIL 限制，CPU 利用率将被限制为 1 个核心。无论您添加多少线程，您都不会看到加速。

在您的情况下，一个简单的解决方案是使用 multiprocessing模块。
代码几乎和你的完全一样，只是将线程切换到进程:

from multiprocessing import Process
coord = tf.train.Coordinator()
processes = []
for thread_index in xrange(4):
    args = (str(thread_index), range(200000))
    p = Process(target=write_instances_to_tfrecord, args=args)
    p.start()
    processes.append(p)
coord.join(processes)

我在自己的 tfrecord 编写器代码上对此进行了测试，并获得了线性扩展加速。进程总数受内存限制。