Сделать буфер загрузки данных (функция выдачи) следующим выходом

Question

До сих пор я использовал Multiprocessing и Multi-threading на functions, что return результат в конце. Я знаю multiprocessing.Queue и multiprocessing.Queue.get(), но я просто не понимаю, как я могу применить это к загрузчику данных ..

Я борюсь со следующей задачей:

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

Так что мой Вопрос в том, есть ли какой-нибудь простой способ заставить мой существующий data_loader вычислить (пребуферировать) его следующий yield, пока ИИ обучается на GPU?

Или мне пришлось бы полностью реструктурировать это с помощью внешнего итератора, который вызывает внутренний меньший загрузчик 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()