Caching disabled on AnchorGenerator

[datumbox]

🐛 Bug

Currently the caching mechanism of AnchorGenerator is completely disabled because the Cache dictionary is wiped every time at the end of forward:

vision/torchvision/models/detection/anchor_utils.py

Lines 156 to 157 in e617474

	# Clear the cache in case that memory leaks.
	self._cache.clear()

The clear statement was introduced at #1657 to solve a memory leak reported at #1689.

Before providing a solution, I think it's worth to:

1. Measure the speed benefit of caching. This should be done by measuring the speed of AnchorGenerator both individually and while being used by a detection model to assess its true impact.
2. Validate that the memory leak is still present by measuring precisely the size of the cache in AnchorGenerator. The best way to do this would be to pass through it a sample of a real-world dataset like coco.

After having the above we could decide how to solve the bug. Here are a few proposed options:

1. Remove the clear statement from forward, if the benchmarks show that the AnchorGenerator has no significant memory issues.
2. Remove the caching completely, if the benchmarks show that is does not offer a material speed improvement (for example improving by over 5% the speed of end-to-end image prediction).
3. Replace the cache with an LRU cache, if the benchmarks show that the memory leak exists and the caching is beneficial. Note that we need to ensure that the LRU cache is compatible with JIT.

[prabhat00155]

I used the script shared in #1689 to verify the memory usage with and without caching and didn't see much difference:

import os
import gc
import random

import torch
import psutil
from torchvision.models.detection import fasterrcnn_resnet50_fpn

process = psutil.Process(os.getpid())
model = fasterrcnn_resnet50_fpn(pretrained=True)
# model.eval() # Gives no issues

for i in range(10):
    size = random.randint(700, 900)
    images = torch.rand([1, 3, size, 800]) 
    # images = torch.rand([1, 3, size, size]) # Gives no issues
    targets = [{'boxes': torch.tensor([[ 10., 20., 30., 40.]]), 'labels': torch.tensor([1])}]
    model(images, targets)
    gc.collect()
    print("Current memory: ", process.memory_info()[0] / float(2**20))

Without caching(current version):

Current memory:  1134.390625
Current memory:  1203.33984375
Current memory:  1223.8828125
Current memory:  1231.9921875
Current memory:  1543.68359375
Current memory:  1546.71875
Current memory:  1571.2265625
Current memory:  1580.12109375
Current memory:  1580.51171875
Current memory:  1580.7890625

With caching(by commenting out clear()):
Try 1:

Current memory:  1171.6171875
Current memory:  1242.02734375
Current memory:  1562.94921875
Current memory:  1592.4140625
Current memory:  1600.921875
Current memory:  1609.3125
Current memory:  1622.9140625
Current memory:  1622.953125
Current memory:  1638.8515625
Current memory:  1649.53125

Try 2:

Current memory:  1119.1328125
Current memory:  1467.59765625
Current memory:  1748.96875
Current memory:  1759.48828125
Current memory:  1772.94140625
Current memory:  1793.87890625
Current memory:  1805.3515625
Current memory:  1810.6953125
Current memory:  1811.0859375
Current memory:  1811.21484375

I also tried comparing the time to run AnchorGenerator forward() using this script(in Juypyter notebook):

size = random.randint(700, 900)
images = torch.rand([10, 3, size, 800])
sizes = [(size, 800)] * 10
    
image_list = ImageList(images, sizes)
features = torch.rand([3, 128, 32, 32])

%%timeit

anchor_gen = AnchorGenerator(
        sizes=((16, 32, 64, 128, 256, 512),),
        aspect_ratios=((0.5, 1.0, 2.0),)
    )
anchors = anchor_gen(image_list, [features])

Without caching(current version):

648 µs ± 92.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

With caching(by commenting out clear()):

642 µs ± 66.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

There doesn't seem to be any significant difference in speed.

Note: I am running these on my laptop, so there could be other processes affecting the results.

[datumbox]

@prabhat00155 Thanks a lot for the detailed benchmarks! Since I spotted this bug, I was inclined to remove caching as it makes the code more complex. Can't say I changed my mind after reading your analysis... I propose to remove the caching mechanism completely.

@fmassa @NicolasHug thoughts?

[NicolasHug]

Regarding the timing benchmark,

%%timeit

anchor_gen = AnchorGenerator(
        sizes=((16, 32, 64, 128, 256, 512),),
        aspect_ratios=((0.5, 1.0, 2.0),)
    )
anchors = anchor_gen(image_list, [features])

I'm wondering if the above properly assesses the caching mechanism. From what I understand the cache is useful across calls to forward but the snippet above will only call forward once per AnchorGenerator instance.

Should this be instead, e.g.:

%%timeit
anchor_gen = AnchorGenerator(
        sizes=((16, 32, 64, 128, 256, 512),),
        aspect_ratios=((0.5, 1.0, 2.0),)
    )

for _ in range(1000):
	anchors = anchor_gen(image_list, [features])  # cache hits for the next 999 calls

?

[prabhat00155]

Regarding the timing benchmark,

%%timeit

anchor_gen = AnchorGenerator(
        sizes=((16, 32, 64, 128, 256, 512),),
        aspect_ratios=((0.5, 1.0, 2.0),)
    )
anchors = anchor_gen(image_list, [features])

I'm wondering if the above properly assesses the caching mechanism. From what I understand the cache is useful across calls to forward but the snippet above will only call forward once per AnchorGenerator instance.

Should this be instead, e.g.:

%%timeit
anchor_gen = AnchorGenerator(
        sizes=((16, 32, 64, 128, 256, 512),),
        aspect_ratios=((0.5, 1.0, 2.0),)
    )

for _ in range(1000):
	anchors = anchor_gen(image_list, [features])  # cache hits for the next 999 calls

?

Doesn't timeit call the cell multiple times(1000 by default)?

[NicolasHug]

Yes but it will create 1000 AnchorGenerator instances and call forward once on each of them. There won't be any cache hit since the cache hits can only happen when forward is called (at least) twice on the same instance, as far as I understand

[prabhat00155]

Yes but it will create 1000 AnchorGenerator instances and call forward once on each of them. There won't be any cache hit since the cache hits can only happen when forward is called (at least) twice on the same instance, as far as I understand

Yeah, that's true. I should move instantiating AnchorGenerator out of this cell.

[NicolasHug]

I should move instantiating AnchorGenerator out of this cell.

I'm not sure this is still what we want: this will give you an average over 1000 runs where the cache will be hit 999 times except for the first time. The stddev will be off.

In fact if the caching is super effective %timeit will actually detect that the first run took much longer than the rest and give you a warning.

[prabhat00155]

I do see an improvement with caching(after removing self._cache.clear()).
With caching:

337 µs ± 14.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Without caching:

526 µs ± 20.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

[NicolasHug]

I get a similar 2X speed up with caching:

%%timeit
anchor_gen = AnchorGenerator(
        sizes=((16, 32, 64, 128, 256, 512),),
        aspect_ratios=((0.5, 1.0, 2.0),)
    )

for _ in range(1000):
	anchors = anchor_gen(image_list, [features])

357 ms ± 21.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

vs

658 ms ± 54.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[prabhat00155]

Similar results with the script @NicolasHug suggested:
With caching:

313 ms ± 14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Without caching:

567 ms ± 41.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[datumbox]

Thanks both.

Good spot for finding the benchmark issue @NicolasHug.

It would be useful to finally measure how much the performance changes during inference for end-to-end predictions. In other words, assume we are making predictions on a mobile device (cpu). How much time does it take on average to produce one prediction with and without the caching (after some burn in period)?

The reason I'm proposing to measure also this is to assess whether this double speed up makes actually a difference for a model such as the fasterrcnn_mobilenet_v3_large_320_fpn.

@prabhat00155 I wonder if you could run a modified version of the following so that we could get comparable numbers on the same machine?

model = fasterrcnn_mobilenet_v3_large_320_fpn(pretrained=True)
model.eval()

images = torch.rand([1, 3, 700, 800]) 
model(images) # populate cache once

# TODO: measure avg time with and without caching for the following loop
for i in range(1000):
    model(images)