Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add an API to estimate peak VRAM usage of a model #280

Open
RunDevelopment opened this issue Jul 3, 2024 · 0 comments
Open

Add an API to estimate peak VRAM usage of a model #280

RunDevelopment opened this issue Jul 3, 2024 · 0 comments
Labels
enhancement New feature or request

Comments

@RunDevelopment
Copy link
Member

Add an API for estimating the peak VRAM usage of a model. Ideally, this would be something simple like:

class ImageModelDescriptor:
    ...
    def estimate_peak_memory(self, size: tuple[int, int]) -> int | None:
        """
        Returns an estimate for the peak memory usage (VRAM if on the GPU) of the model for the given image size.

        If no estimate is available, `None` will be returned.
        """

This would allow us to use better VRAM estimates in chaiNNer, which would allow us to pick better tile sizes.


As for actually finding good estimates. I think we could use a script to determine an estimate automatically. The script would be given an architecture, possible hyperparameters, and a series of image size to test. It would then simply run a model on a random image and measure the peak VRAM usage for all combinations of images sizes and hyperparameters. This gives us some data points to work with. The rest is a curve-fitting exercise, and generating python code for what we could out.

With this script in place, we can just run it and automatically get a good function for estimating the VRAM usage of an architecture.

@RunDevelopment RunDevelopment added the enhancement New feature or request label Jul 3, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

1 participant