使数据加载器(产出函数)缓冲其下一个产出

Question

到目前为止，我只在functions上使用了Multiprocessing和Multi-threading，并在最后return了一个结果。我知道multiprocessing.Queue和multiprocessing.Queue.get()，但我就是不明白如何将其应用于data_loader。

我为以下任务而苦苦挣扎：

def data_loader():
    for _ in range(10**6):
        #calculates for some seconds
        yield result

for data in data_loader():
    train_AI(data)
    #Here an AI is being trained for another some seconds

所以我的问题是:当AI在data_loader上训练时，有没有什么简单的方法可以让我现有的GPU计算(预缓冲)它的下一个yield？

或者我必须完全重构它，使用一个外部迭代器调用一个内部较小的data_loader，每次调用它时都会returns一个批处理？

Answer 1

是的，你可以使用Python的队列：

from multiprocessing import Process, Queue
from time import sleep    
FINISHED_LOADING_DATA = 'LAST ONE'  # just make sure it's not something that can be returned by some_function()

def some_function():
    print('getting data')
    sleep(0.5)
    return 'some_result'


def train_AI(x):
    print('training AI')
    sleep(2)


q = Queue()
final_results = []

def data_loader(q):
    for _ in range(10):
        result = some_function()
        q.put(result)
    q.put(FINISHED_LOADING_DATA)


def train_if_data_available():
    while True:
        data = q.get()
        if data == FINISHED_LOADING_DATA:
            return 'DONE'
        train_AI(data)


t = Process(target=data_loader, args=(q,))
t.daemon = True
t.start()

train_if_data_available()