From ea35db7b8c4b1224be5464970ce3c64268c952a4 Mon Sep 17 00:00:00 2001 From: Joey Ballentine Date: Tue, 16 Jan 2024 00:08:58 -0500 Subject: [PATCH 1/6] add pytorch memory budget setting --- backend/src/packages/chaiNNer_pytorch/settings.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/backend/src/packages/chaiNNer_pytorch/settings.py b/backend/src/packages/chaiNNer_pytorch/settings.py index bc1b130a1..d42529950 100644 --- a/backend/src/packages/chaiNNer_pytorch/settings.py +++ b/backend/src/packages/chaiNNer_pytorch/settings.py @@ -3,7 +3,7 @@ import torch from sanic.log import logger -from api import DropdownSetting, NodeContext, ToggleSetting +from api import DropdownSetting, NodeContext, NumberSetting, ToggleSetting from gpu import get_nvidia_helper from system import is_arm_mac @@ -66,12 +66,24 @@ ), ) +package.add_setting( + NumberSetting( + label="Memory Budget Limit (GiB)", + key="budget_limit", + description="Maximum memory to use for PyTorch inference. 0 means no limit. Memory usage measurement is not completely accurate yet; you may need to significantly adjust this budget limit via trial-and-error if it's not having the effect you want.", + default=0, + min=0, + max=1024**2, + ) +) + @dataclass(frozen=True) class PyTorchSettings: use_cpu: bool use_fp16: bool gpu_index: int + budget_limit: int # PyTorch 2.0 does not support FP16 when using CPU def __post_init__(self): @@ -111,4 +123,5 @@ def get_settings(context: NodeContext) -> PyTorchSettings: use_cpu=settings.get_bool("use_cpu", False), use_fp16=settings.get_bool("use_fp16", False), gpu_index=settings.get_int("gpu_index", 0, parse_str=True), + budget_limit=settings.get_int("budget_limit", 0, parse_str=True), ) From 4c92348295598f87ab8c20c2b654ee8df6e6cd8d Mon Sep 17 00:00:00 2001 From: Joey Ballentine Date: Tue, 16 Jan 2024 00:22:09 -0500 Subject: [PATCH 2/6] Make memory budget setting actually do something --- .../chaiNNer_pytorch/pytorch/processing/upscale_image.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py index e8a10ff8e..560bd554d 100644 --- a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py +++ b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py @@ -55,6 +55,8 @@ def estimate(): model_bytes = sum( p.numel() * element_size for p in model.model.parameters() ) + if options.budget_limit > 0: + free = min(int(options.budget_limit * 1024**3), free) budget = int(free * 0.8) return MaxTileSize( From f525881e13b9da12970f110acd9518c50ddc4395 Mon Sep 17 00:00:00 2001 From: Joey Ballentine Date: Tue, 16 Jan 2024 00:25:53 -0500 Subject: [PATCH 3/6] specify VRAM --- backend/src/packages/chaiNNer_pytorch/settings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/src/packages/chaiNNer_pytorch/settings.py b/backend/src/packages/chaiNNer_pytorch/settings.py index d42529950..9218d89d5 100644 --- a/backend/src/packages/chaiNNer_pytorch/settings.py +++ b/backend/src/packages/chaiNNer_pytorch/settings.py @@ -70,7 +70,7 @@ NumberSetting( label="Memory Budget Limit (GiB)", key="budget_limit", - description="Maximum memory to use for PyTorch inference. 0 means no limit. Memory usage measurement is not completely accurate yet; you may need to significantly adjust this budget limit via trial-and-error if it's not having the effect you want.", + description="Maximum memory (VRAM) to use for PyTorch inference. 0 means no limit. Memory usage measurement is not completely accurate yet; you may need to significantly adjust this budget limit via trial-and-error if it's not having the effect you want.", default=0, min=0, max=1024**2, From 328dc7349d50f723d5ac7595e2b0574162fad2a5 Mon Sep 17 00:00:00 2001 From: Joey Ballentine Date: Sat, 20 Jan 2024 19:57:52 -0500 Subject: [PATCH 4/6] also use limit for CPU --- .../pytorch/processing/upscale_image.py | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py index cfe842ca1..cbf12a630 100644 --- a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py +++ b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py @@ -1,6 +1,7 @@ from __future__ import annotations import numpy as np +import psutil import torch from sanic.log import logger from spandrel import ImageModelDescriptor, ModelTiling @@ -48,17 +49,31 @@ def upscale( tile_size = NO_TILING def estimate(): + element_size = 2 if use_fp16 else 4 + model_bytes = sum( + p.numel() * element_size for p in model.model.parameters() + ) + if "cuda" in device.type: mem_info: tuple[int, int] = torch.cuda.mem_get_info(device) # type: ignore free, _total = mem_info element_size = 2 if use_fp16 else 4 - model_bytes = sum( - p.numel() * element_size for p in model.model.parameters() - ) if options.budget_limit > 0: - free = min(int(options.budget_limit * 1024**3), free) + free = min(options.budget_limit * 1024**3, free) budget = int(free * 0.8) + return MaxTileSize( + estimate_tile_size( + budget, + model_bytes, + img, + element_size, + ) + ) + elif options.budget_limit > 0: + free = psutil.virtual_memory().available + free = min(options.budget_limit * 1024**3, free) + budget = int(free * 0.8) return MaxTileSize( estimate_tile_size( budget, From f2577167c16c3f895e12a33d66146d49eb717248 Mon Sep 17 00:00:00 2001 From: Joey Ballentine Date: Sat, 20 Jan 2024 19:58:52 -0500 Subject: [PATCH 5/6] change condition --- .../chaiNNer_pytorch/pytorch/processing/upscale_image.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py index cbf12a630..6cc284a79 100644 --- a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py +++ b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py @@ -70,9 +70,10 @@ def estimate(): element_size, ) ) - elif options.budget_limit > 0: + elif device.type == "cpu": free = psutil.virtual_memory().available - free = min(options.budget_limit * 1024**3, free) + if options.budget_limit > 0: + free = min(options.budget_limit * 1024**3, free) budget = int(free * 0.8) return MaxTileSize( estimate_tile_size( From a99ec85edd8c010bfdb7fa865f32e49ad82683c6 Mon Sep 17 00:00:00 2001 From: Joey Ballentine Date: Sat, 20 Jan 2024 19:59:52 -0500 Subject: [PATCH 6/6] change desc --- backend/src/packages/chaiNNer_pytorch/settings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/src/packages/chaiNNer_pytorch/settings.py b/backend/src/packages/chaiNNer_pytorch/settings.py index 9218d89d5..6dbb8b1c3 100644 --- a/backend/src/packages/chaiNNer_pytorch/settings.py +++ b/backend/src/packages/chaiNNer_pytorch/settings.py @@ -70,7 +70,7 @@ NumberSetting( label="Memory Budget Limit (GiB)", key="budget_limit", - description="Maximum memory (VRAM) to use for PyTorch inference. 0 means no limit. Memory usage measurement is not completely accurate yet; you may need to significantly adjust this budget limit via trial-and-error if it's not having the effect you want.", + description="Maximum memory (VRAM if GPU, RAM if CPU) to use for PyTorch inference. 0 means no limit. Memory usage measurement is not completely accurate yet; you may need to significantly adjust this budget limit via trial-and-error if it's not having the effect you want.", default=0, min=0, max=1024**2,