Set cudnn.benchmark locally (embedder but not classifier)

vision
1

I am using a ResNet50 as feature extractor and I would like it to run it with cudnn.benchmark=True, as I know the input shape does not change. As the number of created embeddings can differ, the following classifier will have a variable-length input, so I would prefer to set cudnn.benchmark=False for it.

The documentation is rather succinct and I did not find corresponding code in the repo.
Where can I find more information about the benchmarks inner working?

2

You can use the context manager to enable/disable benchmark for a region:

with torch.backends.cudnn.flags(enabled, benchmark, deterministic, allow_tf32):
    some_layer
1 Like