如何实现数据集混洗的多线程馈送

时间:2019-02-25 19:23:21

标签: python tensorflow tensorflow-datasets

我想实现一个多线程生成器来填充我的DataSet Shuffle对象。这就是我的操作方式,但是如果我想在通过多线程数据混入进行训练的同时实时创建数据,那么我现在想这样做。我没有找到任何有关如何执行此操作的文档,这就是为什么我不确定。

model = models.Model(inputs=input, outputs=y_output)
optimizer = tf.train.AdamOptimizer(LEARNING_RATE)
# model = utils.multi_gpu_model(model, gpus=NUM_GPUS, cpu_relocation=True)
model.compile(loss=lossFunc, optimizer=optimizer)

def generator(n):
    while True:
        try:
            imgBatch = []

            ...

            yield imgBatch

        except ValueError:
            pass

def get_generator(n):
    return partial(generator, n)

def dataset(n):
    return tf.data.Dataset.from_generator(get_generator(n), output_types=[(tf.float32, tf.float32)], output_shapes=[(tf.TensorShape([None,None,1]),tf.TensorShape([None,None,1]))

def input_fn():
    ds = tf.data.Dataset.range(len(dataSets)).apply(tf.data.experimental.parallel_interleave(dataset, cycle_length=len(dataSets), sloppy=True))
    ds = ds.map(map_func=lambda imgBatch: processImage(img,lbl), num_parallel_calls=12)
    ds = ds.shuffle(SHUFFLE_SIZE)
    ds = ds.batch(BATCH_SIZE)
    ds = ds.prefetch(1)
    return ds

strategy = tf.contrib.distribute.MirroredStrategy(num_gpus=NUM_GPUS)
config = tf.estimator.RunConfig(train_distribute=strategy)
estimator = tf.keras.estimator.model_to_estimator(model,
                                                  config=config)

estimator.train(lambda:input_fn())

0 个答案:

没有答案