python - 如何帮助 tqdm 计算出自定义迭代器中的总数

Question

我正在实现自己的迭代器。 tqdm 不显示进度条，因为它不知道列表中元素的总数。我不想使用“total=”，因为它看起来很丑。相反，我更愿意向我的迭代器添加一些东西，tqdm 可以使用它来计算总数。

class Batches:
    def __init__(self, batches, target_input):
        self.batches = batches
        self.pos = 0
        self.target_input = target_input

    def __iter__(self):
        return self

    def __next__(self):
        if self.pos < len(self.batches):
            minibatch = self.batches[self.pos]
            target = minibatch[:, :, self.target_input]
            self.pos += 1
            return minibatch, target
        else:
            raise StopIteration

    def __len__(self):
        return self.batches.len()

这可能吗？在上面的代码中添加什么...

像下面这样使用 tqdm..

for minibatch, target in tqdm(Batches(test, target_input)):

    output = lstm(minibatch)
    loss = criterion(output, target)
    writer.add_scalar('loss', loss, tensorboard_step)

Answer 1

我知道已经有一段时间了，但我一直在寻找相同的答案，这里是解决方案。而不是像这样用 tqdm 包装你的 iterable

for i in tqdm(my_iterable):
    do_something()

改用“with”结尾，如:

with tqdm(total=len_of_my_iterable) as progress_bar:
    for i in tqdm(my_iterable):
        do_something()
        progress_bar.update(1) # update progress

对于您的批处理，您可以将总数设置为批处理数，并更新为 1(如上所述)。或者，您可以将总数设置为实际的项目总数，并将更新设置为当前处理的批处理的大小。