ENH: Adds TeeGenerator for named_parameters

[thomasjpfan]

Fixes #378

[benjamin-work]

I have some comments on this:

1. I'm not sure if LazyGenerator is the right name, because "generator" already implies "lazy". Furthermore, it's not so much about laziness and more about caching.

2. I have a strong feeling that such a solution already exists, though a quick search did not reveal a ready solution.

3. The generator creates the whole list right away when called first. I'm not sure if that's a problem but it is possible to avoid it, e.g.:

from itertools import tee
class BetterName:
    def __init__(self, gen):
        self.gen = gen
        self.gen_ = None

    def __iter__(self):
        if self.gen_ is None:
            self.gen_, it = tee(self.gen())
        else:
            it, *_ = tee(self.gen_, 1)
        yield from it

4. Did you measure how much the overhead is reduced?

[ottonemo]

Suggested change

	"""Stores a generator and calls list on the generator on the first
	"""Stores a function that produces a generator and calls list on the generator on the first

[ottonemo]

1. I'm not sure if LazyGenerator is the right name, because "generator" already implies "lazy". Furthermore, it's not so much about laziness and more about caching.

Well, both actually. The function is evaluated lazily and the generator is cached (and/or multiplexed).
But since the lazy part is not actually the important part, it shouldn't influence the name.

3. The generator creates the whole list right away when called first. I'm not sure if that's a problem but it is possible to avoid it, e.g.:

from itertools import tee
class BetterName:
    def __init__(self, gen):
        self.gen = gen
        self.gen_ = None

    def __iter__(self):
        if self.gen_ is None:
            self.gen_, it = tee(self.gen())
        else:
            it, *_ = tee(self.gen_, 1)
        yield from it

You need to always copy the generator:

it, *_ = tee(self.gen_, 1)

will fail if the generator is evaluated more than twice. Better:

from itertools import tee
class LazyGenerator_:
    def __init__(self, gen):
        self.gen = gen()

    def __iter__(self):
        self.gen, it = tee(self.gen)
        yield from it

(this way you don't need the callback at all)

[benjamin-work]

will fail if the generator is evaluated more than twice

What do you mean by that?

Better:

Yes, this also works. I wanted to avoid calling the generator during __init__ because that is what the original code does, but I don't know whether it's a requirement.

[ottonemo]

will fail if the generator is evaluated more than twice

What do you mean by that?

Sorry, should have written an example:

g = BetterName(lambda: [1,2,3])

g1 = list(g) # == [1,2,3]
g2 = list(g) # == [1,2,3]
g3 = list(g) # == [] <- bad

Better:

Yes, this also works. I wanted to avoid calling the generator during __init__ because that is what the original code does, but I don't know whether it's a requirement.

I think it is no requirement. While there might be an overhead by simply creating the generator yielded by named_parameters it is probably less than evaluating the generator fully (as the old code did). Of course this should be evaluated in a benchmark.

[benjamin-work]

Interesting, I didn't see this.

If initialization of the generator during __init__ is not desired, this should still work:

class LazyGenerator:
    def __init__(self, gen):
        self.gen = gen
        self.gen_ = None

    def __iter__(self):
        if self.gen_ is None:
            self.gen_ = self.gen()
        self.gen_, it = tee(self.gen_)
        yield from it

[thomasjpfan]

Here are the results of the microbenchmark comparing the list version and the tee version:

from itertools import tee

class ListLazyGenerator:
    def __init__(self, gen): 
        self.gen = gen
        self.gen_ = None
    def __iter__(self): 
        if self.gen_ is None:
            self.gen_ = list(self.gen())
        yield from self.gen_

class TeeLazyGenerator:
    def __init__(self, gen):
        self.gen = gen
        self.gen_ = None

    def __iter__(self):
        if self.gen_ is None:
            self.gen_ = self.gen()
        self.gen_, it = tee(self.gen_)
        yield from it

class TeeCallInInitLazyGenerator:
    def __init__(self, gen):
        self.gen = gen()

    def __iter__(self):
        self.gen, it = tee(self.gen)
        yield from it

Small Module

import torchvision
m = torchvision.models.squeezenet1_0()

# No generator
%timeit list(m.named_parameters())
# 214 µs ± 10.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

list_lazy = ListLazyGenerator(m.named_parameters)
%timeit list(list_lazy)
# 3.04 µs ± 61.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

tee_lazy = TeeLazyGenerator(m.named_parameters)
%timeit list(tee_lazy)
# 3.51 µs ± 87.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

tee_call_in_init = TeeCallInInitLazyGenerator(m.named_parameters)
%timeit list(tee_call_in_init)
# 3.29 µs ± 47 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Big Module

import torchvision
m = torchvision.models.resnet152()

# No generator
%timeit list(m.named_parameters())
# 1.67 ms ± 75.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

list_lazy = ListLazyGenerator(m.named_parameters)
%timeit list(list_lazy)
# 22.5 µs ± 405 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

tee_lazy = TeeLazyGenerator(m.named_parameters)
%timeit list(tee_lazy)
# 24.8 µs ± 1.36 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

tee_call_in_init = TeeCallInInitLazyGenerator(m.named_parameters)
%timeit list(tee_call_in_init)
# 23.8 µs ± 253 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

I think list is fastest with a tradeoff in memory usage. Given the microbenchmark, tee with the generator called in init would be my preferred solution.

[benjamin-work]

Thanks for the measurements.

I believe the time differences between the 3 implementations is negligible. The only downside of TeeCallInInitLazyGenerator would be if very frequently, the generator is instantiated but not used, so that the overhead for calling named_parameters is paid unnecessarily:

%timeit ListLazyGenerator(m.named_parameters)
# 596 ns ± 14.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%timeit TeeLazyGenerator(m.named_parameters)
# 530 ns ± 55.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%timeit TeeCallInInitLazyGenerator(m.named_parameters)
# 814 ns ± 67 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

This overhead is relevant when named_parameters are never used. However, this could all be premature optimization when looking at those small numbers, in which case simple beats fast.

[thomasjpfan]

This PR was updated to use the tee version of the generator, and renames the class to TeeGenerator.

[benjamin-work]

The test would be a little clearer if the call count would also be tested immediately after the first call (i.e. here).

[benjamin-work]

The docstring is no longer up-to-date. Also, could you add 2 newlines after the first sentence?

[benjamin-work]

A couple of small comments.

[ottonemo]

Suggested change

	`named_parameters` only when it is used.
	`named_parameters` only when it is used to reduce the run-time overhead of the call (#379)

[thomasjpfan]

I am starting to see how you like the CHANGELOG to be worded. 😅

[ottonemo]

I think the TeeGenerator should use a generator, not a callable. The overhead is negligible but it allows for easy passing of partially consumed generators (e.g. when using models with two parameter sets) and is more intuitive.

[ottonemo]

LGTM

[ottonemo]

Suggested change

	when ``TeeGenerator`` is iterated over.
	when ``TeeGenerator`` is iterated over to let you iterate over the given generator more than once.

[benjamin-work]

@thomasjpfan Would you mind looking at the open comments?

[benjamin-work]

@thomasjpfan Some small TODOs left:

• merge conflict in CHANGES.md
• I believe you accidentally deleted a line in the CHANGES.md, currently there is only half a sentence
• What do you think about my suggestion to adjust the unit test?

[thomasjpfan]

What do you think about my suggestion to adjust the unit test?

TeeGenerator no longer calls the function, it now accepts a generator. The test was refactored to not use mocks, thus removing the call_count assertion.

[benjamin-work]

TeeGenerator no longer calls the function, it now accepts a generator. The test was refactored to not use mocks, thus removing the call_count assertion.

Ah okay, I was looking at the outdated state shown by github. With the given change in the test, does the function name still make sense, though?

[thomasjpfan]

The test was renamed to test_returns_copies_of_generator.

[benjamin-work]

Thanks.