如何在处理前一个结果的同时生成生成器?

2024-09-28 03:13:24 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一个生成器,它从磁盘上读取一些数据,在CPU上进行一些计算,然后将其移动到GPU后生成结果。然后在GPU上处理结果。 处理时间比生成时间长,所以我认为在最后一批处理完成的同时生成下一批应该很有用

我试着用python的线程库来实现这一点。为线程指定一个生成器作为目标:

class ThreadGenerator(threading.Thread):
    def __init__(self, group=None, target=None, name=None,
                 args=(), kwargs={}, Verbose=None):
        threading.Thread.__init__(self, group, target, name, args, kwargs)
    def run(self):
        if self._target is not None:
            self._gen = self._target(*self._args, **self._kwargs)
            self.output = next(self._gen)
    def get_next(self):
        self.output = next(self._gen)

问题是只有run与主进程并行发生,而get_next没有

其思想是获取t.output,调用t.get_next,然后在生成下一个生成器输出时对输出进行计算:

# example - in reality the generator takes much longer than just outputting an integer.
t = ThreadGenerator(target=range, args=(10,))
t.run()

for _ in num_steps:
    o = t.output
    t.get_next()
    result = computations(o)

Tags: runselfnonetargetoutputgetgpudef

热门问题