From 45f70f6e3f9afdeb5586cc19fff5f3d2aab1a8a7 Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Thu, 9 Apr 2026 09:45:05 -0400 Subject: [PATCH 1/5] modules to check later --- pyproject.toml | 2 +- src/electrai/entrypoints/main.py | 2 + uv.lock | 189 +++++++++++++++++++++---------- 3 files changed, 132 insertions(+), 61 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d8920e88..1536e088 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ requires-python = ">=3.11" dependencies = [ "numpy~=2.3.3", "scikit-learn>=1.7.2", - "torch~=2.9.1", + "torch~=2.10.0", "torchvision>=0.24.0", "lightning~= 2.5.6", "wandb>=0.12.10", diff --git a/src/electrai/entrypoints/main.py b/src/electrai/entrypoints/main.py index 6ccc4725..ad76ca28 100644 --- a/src/electrai/entrypoints/main.py +++ b/src/electrai/entrypoints/main.py @@ -8,6 +8,8 @@ from electrai.entrypoints.train import train torch.backends.cudnn.conv.fp32_precision = "tf32" +torch.backends.cudnn.benchmark = False +torch.backends.cudnn.deterministic = True def main() -> None: diff --git a/uv.lock b/uv.lock index fbe3c234..fa9d0fda 100644 --- a/uv.lock +++ b/uv.lock @@ -170,6 +170,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + [[package]] name = "attrs" version = "25.4.0" @@ -680,6 +686,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0f/e7/aa315e6a749d9b96c2504a1ba0ba031ba2d0517e972ce22682e3fccecb09/cssselect2-0.8.0-py3-none-any.whl", hash = "sha256:46fc70ebc41ced7a32cd42d58b1884d72ade23d21e5a4eaaf022401c13f0e76e", size = 15454, upload-time = "2025-03-05T14:46:06.463Z" }, ] +[[package]] +name = "cuda-bindings" +version = "12.9.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder", marker = "sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/e7/b47792cc2d01c7e1d37c32402182524774dadd2d26339bd224e0e913832e/cuda_bindings-12.9.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c912a3d9e6b6651853eed8eed96d6800d69c08e94052c292fec3f282c5a817c9", size = 12210593, upload-time = "2025-10-21T14:51:36.574Z" }, + { url = "https://files.pythonhosted.org/packages/a9/c1/dabe88f52c3e3760d861401bb994df08f672ec893b8f7592dc91626adcf3/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fda147a344e8eaeca0c6ff113d2851ffca8f7dfc0a6c932374ee5c47caa649c8", size = 12151019, upload-time = "2025-10-21T14:51:43.167Z" }, + { url = "https://files.pythonhosted.org/packages/63/56/e465c31dc9111be3441a9ba7df1941fe98f4aa6e71e8788a3fb4534ce24d/cuda_bindings-12.9.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:32bdc5a76906be4c61eb98f546a6786c5773a881f3b166486449b5d141e4a39f", size = 11906628, upload-time = "2025-10-21T14:51:49.905Z" }, + { url = "https://files.pythonhosted.org/packages/a3/84/1e6be415e37478070aeeee5884c2022713c1ecc735e6d82d744de0252eee/cuda_bindings-12.9.4-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56e0043c457a99ac473ddc926fe0dc4046694d99caef633e92601ab52cbe17eb", size = 11925991, upload-time = "2025-10-21T14:51:56.535Z" }, + { url = "https://files.pythonhosted.org/packages/d1/af/6dfd8f2ed90b1d4719bc053ff8940e494640fe4212dc3dd72f383e4992da/cuda_bindings-12.9.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8b72ee72a9cc1b531db31eebaaee5c69a8ec3500e32c6933f2d3b15297b53686", size = 11922703, upload-time = "2025-10-21T14:52:03.585Z" }, + { url = "https://files.pythonhosted.org/packages/6c/19/90ac264acc00f6df8a49378eedec9fd2db3061bf9263bf9f39fd3d8377c3/cuda_bindings-12.9.4-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d80bffc357df9988dca279734bc9674c3934a654cab10cadeed27ce17d8635ee", size = 11924658, upload-time = "2025-10-21T14:52:10.411Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.5.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/f9/1b9b60a30fc463c14cdea7a77228131a0ccc89572e8df9cb86c9648271ab/cuda_pathfinder-1.5.2-py3-none-any.whl", hash = "sha256:0c5f160a7756c5b072723cbbd6d861e38917ef956c68150b02f0b6e9271c71fa", size = 49988, upload-time = "2026-04-06T23:01:05.17Z" }, +] + [[package]] name = "cycler" version = "0.12.1" @@ -715,6 +745,7 @@ name = "electrai" version = "0.0.1" source = { editable = "." } dependencies = [ + { name = "hydra-core" }, { name = "lightning" }, { name = "numpy" }, { name = "pymatgen" }, @@ -759,6 +790,7 @@ dev = [ requires-dist = [ { name = "cairosvg", marker = "extra == 'docs'", specifier = ">=2.7.1" }, { name = "fire", marker = "extra == 'zarr-conversion'", specifier = ">=0.5.0" }, + { name = "hydra-core", specifier = ">=1.3.2" }, { name = "lightning", specifier = "~=2.5.6" }, { name = "mkdocs-gen-files", marker = "extra == 'docs'", specifier = ">=0.5.0" }, { name = "mkdocs-literate-nav", marker = "extra == 'docs'", specifier = ">=0.6.0" }, @@ -775,7 +807,7 @@ requires-dist = [ { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.0.285" }, { name = "s3fs", marker = "extra == 'zarr-conversion'", specifier = ">=2024.5.0" }, { name = "scikit-learn", specifier = ">=1.7.2" }, - { name = "torch", specifier = "~=2.9.1" }, + { name = "torch", specifier = "~=2.10.0" }, { name = "torchvision", specifier = ">=0.24.0" }, { name = "wandb", specifier = ">=0.12.10" }, { name = "zarr", specifier = ">=3.1.3" }, @@ -1027,6 +1059,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/b1/9ff6578d789a89812ff21e4e0f80ffae20a65d5dd84e7a17873fe3b365be/griffe-1.14.0-py3-none-any.whl", hash = "sha256:0e9d52832cccf0f7188cfe585ba962d2674b241c01916d780925df34873bceb0", size = 144439, upload-time = "2025-09-05T15:02:27.511Z" }, ] +[[package]] +name = "hydra-core" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "omegaconf" }, + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, +] + [[package]] name = "idna" version = "3.11" @@ -1888,10 +1934,10 @@ wheels = [ [[package]] name = "nvidia-nvshmem-cu12" -version = "3.3.20" +version = "3.4.5" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3b/6c/99acb2f9eb85c29fc6f3a7ac4dccfd992e22666dd08a642b303311326a97/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d00f26d3f9b2e3c3065be895e3059d6479ea5c638a3f38c9fec49b1b9dd7c1e5", size = 124657145, upload-time = "2025-08-04T20:25:19.995Z" }, + { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, ] [[package]] @@ -1902,6 +1948,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, ] +[[package]] +name = "omegaconf" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, +] + [[package]] name = "orjson" version = "3.11.4" @@ -3032,9 +3091,10 @@ wheels = [ [[package]] name = "torch" -version = "2.9.1" +version = "2.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "filelock" }, { name = "fsspec" }, { name = "jinja2" }, @@ -3060,30 +3120,39 @@ dependencies = [ { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/15/db/c064112ac0089af3d2f7a2b5bfbabf4aa407a78b74f87889e524b91c5402/torch-2.9.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:62b3fd888277946918cba4478cf849303da5359f0fb4e3bfb86b0533ba2eaf8d", size = 104220430, upload-time = "2025-11-12T15:20:31.705Z" }, - { url = "https://files.pythonhosted.org/packages/56/be/76eaa36c9cd032d3b01b001e2c5a05943df75f26211f68fae79e62f87734/torch-2.9.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:d033ff0ac3f5400df862a51bdde9bad83561f3739ea0046e68f5401ebfa67c1b", size = 899821446, upload-time = "2025-11-12T15:20:15.544Z" }, - { url = "https://files.pythonhosted.org/packages/47/cc/7a2949e38dfe3244c4df21f0e1c27bce8aedd6c604a587dd44fc21017cb4/torch-2.9.1-cp311-cp311-win_amd64.whl", hash = "sha256:0d06b30a9207b7c3516a9e0102114024755a07045f0c1d2f2a56b1819ac06bcb", size = 110973074, upload-time = "2025-11-12T15:21:39.958Z" }, - { url = "https://files.pythonhosted.org/packages/1e/ce/7d251155a783fb2c1bb6837b2b7023c622a2070a0a72726ca1df47e7ea34/torch-2.9.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:52347912d868653e1528b47cafaf79b285b98be3f4f35d5955389b1b95224475", size = 74463887, upload-time = "2025-11-12T15:20:36.611Z" }, - { url = "https://files.pythonhosted.org/packages/0f/27/07c645c7673e73e53ded71705045d6cb5bae94c4b021b03aa8d03eee90ab/torch-2.9.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:da5f6f4d7f4940a173e5572791af238cb0b9e21b1aab592bd8b26da4c99f1cd6", size = 104126592, upload-time = "2025-11-12T15:20:41.62Z" }, - { url = "https://files.pythonhosted.org/packages/19/17/e377a460603132b00760511299fceba4102bd95db1a0ee788da21298ccff/torch-2.9.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:27331cd902fb4322252657f3902adf1c4f6acad9dcad81d8df3ae14c7c4f07c4", size = 899742281, upload-time = "2025-11-12T15:22:17.602Z" }, - { url = "https://files.pythonhosted.org/packages/b1/1a/64f5769025db846a82567fa5b7d21dba4558a7234ee631712ee4771c436c/torch-2.9.1-cp312-cp312-win_amd64.whl", hash = "sha256:81a285002d7b8cfd3fdf1b98aa8df138d41f1a8334fd9ea37511517cedf43083", size = 110940568, upload-time = "2025-11-12T15:21:18.689Z" }, - { url = "https://files.pythonhosted.org/packages/6e/ab/07739fd776618e5882661d04c43f5b5586323e2f6a2d7d84aac20d8f20bd/torch-2.9.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:c0d25d1d8e531b8343bea0ed811d5d528958f1dcbd37e7245bc686273177ad7e", size = 74479191, upload-time = "2025-11-12T15:21:25.816Z" }, - { url = "https://files.pythonhosted.org/packages/20/60/8fc5e828d050bddfab469b3fe78e5ab9a7e53dda9c3bdc6a43d17ce99e63/torch-2.9.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:c29455d2b910b98738131990394da3e50eea8291dfeb4b12de71ecf1fdeb21cb", size = 104135743, upload-time = "2025-11-12T15:21:34.936Z" }, - { url = "https://files.pythonhosted.org/packages/f2/b7/6d3f80e6918213babddb2a37b46dbb14c15b14c5f473e347869a51f40e1f/torch-2.9.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:524de44cd13931208ba2c4bde9ec7741fd4ae6bfd06409a604fc32f6520c2bc9", size = 899749493, upload-time = "2025-11-12T15:24:36.356Z" }, - { url = "https://files.pythonhosted.org/packages/a6/47/c7843d69d6de8938c1cbb1eba426b1d48ddf375f101473d3e31a5fc52b74/torch-2.9.1-cp313-cp313-win_amd64.whl", hash = "sha256:545844cc16b3f91e08ce3b40e9c2d77012dd33a48d505aed34b7740ed627a1b2", size = 110944162, upload-time = "2025-11-12T15:21:53.151Z" }, - { url = "https://files.pythonhosted.org/packages/28/0e/2a37247957e72c12151b33a01e4df651d9d155dd74d8cfcbfad15a79b44a/torch-2.9.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5be4bf7496f1e3ffb1dd44b672adb1ac3f081f204c5ca81eba6442f5f634df8e", size = 74830751, upload-time = "2025-11-12T15:21:43.792Z" }, - { url = "https://files.pythonhosted.org/packages/4b/f7/7a18745edcd7b9ca2381aa03353647bca8aace91683c4975f19ac233809d/torch-2.9.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:30a3e170a84894f3652434b56d59a64a2c11366b0ed5776fab33c2439396bf9a", size = 104142929, upload-time = "2025-11-12T15:21:48.319Z" }, - { url = "https://files.pythonhosted.org/packages/f4/dd/f1c0d879f2863ef209e18823a988dc7a1bf40470750e3ebe927efdb9407f/torch-2.9.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:8301a7b431e51764629208d0edaa4f9e4c33e6df0f2f90b90e261d623df6a4e2", size = 899748978, upload-time = "2025-11-12T15:23:04.568Z" }, - { url = "https://files.pythonhosted.org/packages/1f/9f/6986b83a53b4d043e36f3f898b798ab51f7f20fdf1a9b01a2720f445043d/torch-2.9.1-cp313-cp313t-win_amd64.whl", hash = "sha256:2e1c42c0ae92bf803a4b2409fdfed85e30f9027a66887f5e7dcdbc014c7531db", size = 111176995, upload-time = "2025-11-12T15:22:01.618Z" }, - { url = "https://files.pythonhosted.org/packages/40/60/71c698b466dd01e65d0e9514b5405faae200c52a76901baf6906856f17e4/torch-2.9.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:2c14b3da5df416cf9cb5efab83aa3056f5b8cd8620b8fde81b4987ecab730587", size = 74480347, upload-time = "2025-11-12T15:21:57.648Z" }, - { url = "https://files.pythonhosted.org/packages/48/50/c4b5112546d0d13cc9eaa1c732b823d676a9f49ae8b6f97772f795874a03/torch-2.9.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1edee27a7c9897f4e0b7c14cfc2f3008c571921134522d5b9b5ec4ebbc69041a", size = 74433245, upload-time = "2025-11-12T15:22:39.027Z" }, - { url = "https://files.pythonhosted.org/packages/81/c9/2628f408f0518b3bae49c95f5af3728b6ab498c8624ab1e03a43dd53d650/torch-2.9.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:19d144d6b3e29921f1fc70503e9f2fc572cde6a5115c0c0de2f7ca8b1483e8b6", size = 104134804, upload-time = "2025-11-12T15:22:35.222Z" }, - { url = "https://files.pythonhosted.org/packages/28/fc/5bc91d6d831ae41bf6e9e6da6468f25330522e92347c9156eb3f1cb95956/torch-2.9.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:c432d04376f6d9767a9852ea0def7b47a7bbc8e7af3b16ac9cf9ce02b12851c9", size = 899747132, upload-time = "2025-11-12T15:23:36.068Z" }, - { url = "https://files.pythonhosted.org/packages/63/5d/e8d4e009e52b6b2cf1684bde2a6be157b96fb873732542fb2a9a99e85a83/torch-2.9.1-cp314-cp314-win_amd64.whl", hash = "sha256:d187566a2cdc726fc80138c3cdb260970fab1c27e99f85452721f7759bbd554d", size = 110934845, upload-time = "2025-11-12T15:22:48.367Z" }, - { url = "https://files.pythonhosted.org/packages/bd/b2/2d15a52516b2ea3f414643b8de68fa4cb220d3877ac8b1028c83dc8ca1c4/torch-2.9.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:cb10896a1f7fedaddbccc2017ce6ca9ecaaf990f0973bdfcf405439750118d2c", size = 74823558, upload-time = "2025-11-12T15:22:43.392Z" }, - { url = "https://files.pythonhosted.org/packages/86/5c/5b2e5d84f5b9850cd1e71af07524d8cbb74cba19379800f1f9f7c997fc70/torch-2.9.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:0a2bd769944991c74acf0c4ef23603b9c777fdf7637f115605a4b2d8023110c7", size = 104145788, upload-time = "2025-11-12T15:23:52.109Z" }, - { url = "https://files.pythonhosted.org/packages/a9/8c/3da60787bcf70add986c4ad485993026ac0ca74f2fc21410bc4eb1bb7695/torch-2.9.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:07c8a9660bc9414c39cac530ac83b1fb1b679d7155824144a40a54f4a47bfa73", size = 899735500, upload-time = "2025-11-12T15:24:08.788Z" }, - { url = "https://files.pythonhosted.org/packages/db/2b/f7818f6ec88758dfd21da46b6cd46af9d1b3433e53ddbb19ad1e0da17f9b/torch-2.9.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c88d3299ddeb2b35dcc31753305612db485ab6f1823e37fb29451c8b2732b87e", size = 111163659, upload-time = "2025-11-12T15:23:20.009Z" }, + { url = "https://files.pythonhosted.org/packages/0f/8b/4b61d6e13f7108f36910df9ab4b58fd389cc2520d54d81b88660804aad99/torch-2.10.0-2-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:418997cb02d0a0f1497cf6a09f63166f9f5df9f3e16c8a716ab76a72127c714f", size = 79423467, upload-time = "2026-02-10T21:44:48.711Z" }, + { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" }, + { url = "https://files.pythonhosted.org/packages/ec/23/2c9fe0c9c27f7f6cb865abcea8a4568f29f00acaeadfc6a37f6801f84cb4/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254, upload-time = "2026-02-10T21:44:44.095Z" }, + { url = "https://files.pythonhosted.org/packages/36/ab/7b562f1808d3f65414cd80a4f7d4bb00979d9355616c034c171249e1a303/torch-2.10.0-3-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ac5bdcbb074384c66fa160c15b1ead77839e3fe7ed117d667249afce0acabfac", size = 915518691, upload-time = "2026-03-11T14:15:43.147Z" }, + { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" }, + { url = "https://files.pythonhosted.org/packages/ab/c6/4dfe238342ffdcec5aef1c96c457548762d33c40b45a1ab7033bb26d2ff2/torch-2.10.0-3-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:80b1b5bfe38eb0e9f5ff09f206dcac0a87aadd084230d4a36eea5ec5232c115b", size = 915627275, upload-time = "2026-03-11T14:16:11.325Z" }, + { url = "https://files.pythonhosted.org/packages/d8/f0/72bf18847f58f877a6a8acf60614b14935e2f156d942483af1ffc081aea0/torch-2.10.0-3-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:46b3574d93a2a8134b3f5475cfb98e2eb46771794c57015f6ad1fb795ec25e49", size = 915523474, upload-time = "2026-03-11T14:17:44.422Z" }, + { url = "https://files.pythonhosted.org/packages/f4/39/590742415c3030551944edc2ddc273ea1fdfe8ffb2780992e824f1ebee98/torch-2.10.0-3-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:b1d5e2aba4eb7f8e87fbe04f86442887f9167a35f092afe4c237dfcaaef6e328", size = 915632474, upload-time = "2026-03-11T14:15:13.666Z" }, + { url = "https://files.pythonhosted.org/packages/b6/8e/34949484f764dde5b222b7fe3fede43e4a6f0da9d7f8c370bb617d629ee2/torch-2.10.0-3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:0228d20b06701c05a8f978357f657817a4a63984b0c90745def81c18aedfa591", size = 915523882, upload-time = "2026-03-11T14:14:46.311Z" }, + { url = "https://files.pythonhosted.org/packages/78/89/f5554b13ebd71e05c0b002f95148033e730d3f7067f67423026cc9c69410/torch-2.10.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:3282d9febd1e4e476630a099692b44fdc214ee9bf8ee5377732d9d9dfe5712e4", size = 145992610, upload-time = "2026-01-21T16:25:26.327Z" }, + { url = "https://files.pythonhosted.org/packages/ae/30/a3a2120621bf9c17779b169fc17e3dc29b230c29d0f8222f499f5e159aa8/torch-2.10.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a2f9edd8dbc99f62bc4dfb78af7bf89499bca3d753423ac1b4e06592e467b763", size = 915607863, upload-time = "2026-01-21T16:25:06.696Z" }, + { url = "https://files.pythonhosted.org/packages/6f/3d/c87b33c5f260a2a8ad68da7147e105f05868c281c63d65ed85aa4da98c66/torch-2.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:29b7009dba4b7a1c960260fc8ac85022c784250af43af9fb0ebafc9883782ebd", size = 113723116, upload-time = "2026-01-21T16:25:21.916Z" }, + { url = "https://files.pythonhosted.org/packages/61/d8/15b9d9d3a6b0c01b883787bd056acbe5cc321090d4b216d3ea89a8fcfdf3/torch-2.10.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:b7bd80f3477b830dd166c707c5b0b82a898e7b16f59a7d9d42778dd058272e8b", size = 79423461, upload-time = "2026-01-21T16:24:50.266Z" }, + { url = "https://files.pythonhosted.org/packages/cc/af/758e242e9102e9988969b5e621d41f36b8f258bb4a099109b7a4b4b50ea4/torch-2.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5fd4117d89ffd47e3dcc71e71a22efac24828ad781c7e46aaaf56bf7f2796acf", size = 145996088, upload-time = "2026-01-21T16:24:44.171Z" }, + { url = "https://files.pythonhosted.org/packages/23/8e/3c74db5e53bff7ed9e34c8123e6a8bfef718b2450c35eefab85bb4a7e270/torch-2.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:787124e7db3b379d4f1ed54dd12ae7c741c16a4d29b49c0226a89bea50923ffb", size = 915711952, upload-time = "2026-01-21T16:23:53.503Z" }, + { url = "https://files.pythonhosted.org/packages/6e/01/624c4324ca01f66ae4c7cd1b74eb16fb52596dce66dbe51eff95ef9e7a4c/torch-2.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c66c61f44c5f903046cc696d088e21062644cbe541c7f1c4eaae88b2ad23547", size = 113757972, upload-time = "2026-01-21T16:24:39.516Z" }, + { url = "https://files.pythonhosted.org/packages/c9/5c/dee910b87c4d5c0fcb41b50839ae04df87c1cfc663cf1b5fca7ea565eeaa/torch-2.10.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:6d3707a61863d1c4d6ebba7be4ca320f42b869ee657e9b2c21c736bf17000294", size = 79498198, upload-time = "2026-01-21T16:24:34.704Z" }, + { url = "https://files.pythonhosted.org/packages/c9/6f/f2e91e34e3fcba2e3fc8d8f74e7d6c22e74e480bbd1db7bc8900fdf3e95c/torch-2.10.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5c4d217b14741e40776dd7074d9006fd28b8a97ef5654db959d8635b2fe5f29b", size = 146004247, upload-time = "2026-01-21T16:24:29.335Z" }, + { url = "https://files.pythonhosted.org/packages/98/fb/5160261aeb5e1ee12ee95fe599d0541f7c976c3701d607d8fc29e623229f/torch-2.10.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6b71486353fce0f9714ca0c9ef1c850a2ae766b409808acd58e9678a3edb7738", size = 915716445, upload-time = "2026-01-21T16:22:45.353Z" }, + { url = "https://files.pythonhosted.org/packages/6a/16/502fb1b41e6d868e8deb5b0e3ae926bbb36dab8ceb0d1b769b266ad7b0c3/torch-2.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2ee399c644dc92ef7bc0d4f7e74b5360c37cdbe7c5ba11318dda49ffac2bc57", size = 113757050, upload-time = "2026-01-21T16:24:19.204Z" }, + { url = "https://files.pythonhosted.org/packages/1a/0b/39929b148f4824bc3ad6f9f72a29d4ad865bcf7ebfc2fa67584773e083d2/torch-2.10.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:3202429f58309b9fa96a614885eace4b7995729f44beb54d3e4a47773649d382", size = 79851305, upload-time = "2026-01-21T16:24:09.209Z" }, + { url = "https://files.pythonhosted.org/packages/d8/14/21fbce63bc452381ba5f74a2c0a959fdf5ad5803ccc0c654e752e0dbe91a/torch-2.10.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:aae1b29cd68e50a9397f5ee897b9c24742e9e306f88a807a27d617f07adb3bd8", size = 146005472, upload-time = "2026-01-21T16:22:29.022Z" }, + { url = "https://files.pythonhosted.org/packages/54/fd/b207d1c525cb570ef47f3e9f836b154685011fce11a2f444ba8a4084d042/torch-2.10.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:6021db85958db2f07ec94e1bc77212721ba4920c12a18dc552d2ae36a3eb163f", size = 915612644, upload-time = "2026-01-21T16:21:47.019Z" }, + { url = "https://files.pythonhosted.org/packages/36/53/0197f868c75f1050b199fe58f9bf3bf3aecac9b4e85cc9c964383d745403/torch-2.10.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff43db38af76fda183156153983c9a096fc4c78d0cd1e07b14a2314c7f01c2c8", size = 113997015, upload-time = "2026-01-21T16:23:00.767Z" }, + { url = "https://files.pythonhosted.org/packages/0e/13/e76b4d9c160e89fff48bf16b449ea324bda84745d2ab30294c37c2434c0d/torch-2.10.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:cdf2a523d699b70d613243211ecaac14fe9c5df8a0b0a9c02add60fb2a413e0f", size = 79498248, upload-time = "2026-01-21T16:23:09.315Z" }, + { url = "https://files.pythonhosted.org/packages/4f/93/716b5ac0155f1be70ed81bacc21269c3ece8dba0c249b9994094110bfc51/torch-2.10.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:bf0d9ff448b0218e0433aeb198805192346c4fd659c852370d5cc245f602a06a", size = 79464992, upload-time = "2026-01-21T16:23:05.162Z" }, + { url = "https://files.pythonhosted.org/packages/69/2b/51e663ff190c9d16d4a8271203b71bc73a16aa7619b9f271a69b9d4a936b/torch-2.10.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:233aed0659a2503b831d8a67e9da66a62c996204c0bba4f4c442ccc0c68a3f60", size = 146018567, upload-time = "2026-01-21T16:22:23.393Z" }, + { url = "https://files.pythonhosted.org/packages/5e/cd/4b95ef7f293b927c283db0b136c42be91c8ec6845c44de0238c8c23bdc80/torch-2.10.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:682497e16bdfa6efeec8cde66531bc8d1fbbbb4d8788ec6173c089ed3cc2bfe5", size = 915721646, upload-time = "2026-01-21T16:21:16.983Z" }, + { url = "https://files.pythonhosted.org/packages/56/97/078a007208f8056d88ae43198833469e61a0a355abc0b070edd2c085eb9a/torch-2.10.0-cp314-cp314-win_amd64.whl", hash = "sha256:6528f13d2a8593a1a412ea07a99812495bec07e9224c28b2a25c0a30c7da025c", size = 113752373, upload-time = "2026-01-21T16:22:13.471Z" }, + { url = "https://files.pythonhosted.org/packages/d8/94/71994e7d0d5238393df9732fdab607e37e2b56d26a746cb59fdb415f8966/torch-2.10.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:f5ab4ba32383061be0fb74bda772d470140a12c1c3b58a0cfbf3dae94d164c28", size = 79850324, upload-time = "2026-01-21T16:22:09.494Z" }, + { url = "https://files.pythonhosted.org/packages/e2/65/1a05346b418ea8ccd10360eef4b3e0ce688fba544e76edec26913a8d0ee0/torch-2.10.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:716b01a176c2a5659c98f6b01bf868244abdd896526f1c692712ab36dbaf9b63", size = 146006482, upload-time = "2026-01-21T16:22:18.42Z" }, + { url = "https://files.pythonhosted.org/packages/1d/b9/5f6f9d9e859fc3235f60578fa64f52c9c6e9b4327f0fe0defb6de5c0de31/torch-2.10.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:d8f5912ba938233f86361e891789595ff35ca4b4e2ac8fe3670895e5976731d6", size = 915613050, upload-time = "2026-01-21T16:20:49.035Z" }, + { url = "https://files.pythonhosted.org/packages/66/4d/35352043ee0eaffdeff154fad67cd4a31dbed7ff8e3be1cc4549717d6d51/torch-2.10.0-cp314-cp314t-win_amd64.whl", hash = "sha256:71283a373f0ee2c89e0f0d5f446039bdabe8dbc3c9ccf35f0f784908b0acd185", size = 113995816, upload-time = "2026-01-21T16:22:05.312Z" }, ] [[package]] @@ -3103,7 +3172,7 @@ wheels = [ [[package]] name = "torchvision" -version = "0.24.1" +version = "0.25.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, @@ -3111,30 +3180,30 @@ dependencies = [ { name = "torch" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/69/30f5f03752aa1a7c23931d2519b31e557f3f10af5089d787cddf3b903ecf/torchvision-0.24.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:056c525dc875f18fe8e9c27079ada166a7b2755cea5a2199b0bc7f1f8364e600", size = 1891436, upload-time = "2025-11-12T15:25:04.3Z" }, - { url = "https://files.pythonhosted.org/packages/0c/69/49aae86edb75fe16460b59a191fcc0f568c2378f780bb063850db0fe007a/torchvision-0.24.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:1e39619de698e2821d71976c92c8a9e50cdfd1e993507dfb340f2688bfdd8283", size = 2387757, upload-time = "2025-11-12T15:25:06.795Z" }, - { url = "https://files.pythonhosted.org/packages/11/c9/1dfc3db98797b326f1d0c3f3bb61c83b167a813fc7eab6fcd2edb8c7eb9d/torchvision-0.24.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a0f106663e60332aa4fcb1ca2159ef8c3f2ed266b0e6df88de261048a840e0df", size = 8047682, upload-time = "2025-11-12T15:25:21.125Z" }, - { url = "https://files.pythonhosted.org/packages/fa/bb/cfc6a6f6ccc84a534ed1fdf029ae5716dd6ff04e57ed9dc2dab38bf652d5/torchvision-0.24.1-cp311-cp311-win_amd64.whl", hash = "sha256:a9308cdd37d8a42e14a3e7fd9d271830c7fecb150dd929b642f3c1460514599a", size = 4037588, upload-time = "2025-11-12T15:25:14.402Z" }, - { url = "https://files.pythonhosted.org/packages/f0/af/18e2c6b9538a045f60718a0c5a058908ccb24f88fde8e6f0fc12d5ff7bd3/torchvision-0.24.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e48bf6a8ec95872eb45763f06499f87bd2fb246b9b96cb00aae260fda2f96193", size = 1891433, upload-time = "2025-11-12T15:25:03.232Z" }, - { url = "https://files.pythonhosted.org/packages/9d/43/600e5cfb0643d10d633124f5982d7abc2170dfd7ce985584ff16edab3e76/torchvision-0.24.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:7fb7590c737ebe3e1c077ad60c0e5e2e56bb26e7bccc3b9d04dbfc34fd09f050", size = 2386737, upload-time = "2025-11-12T15:25:08.288Z" }, - { url = "https://files.pythonhosted.org/packages/93/b1/db2941526ecddd84884132e2742a55c9311296a6a38627f9e2627f5ac889/torchvision-0.24.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:66a98471fc18cad9064123106d810a75f57f0838eee20edc56233fd8484b0cc7", size = 8049868, upload-time = "2025-11-12T15:25:13.058Z" }, - { url = "https://files.pythonhosted.org/packages/69/98/16e583f59f86cd59949f59d52bfa8fc286f86341a229a9d15cbe7a694f0c/torchvision-0.24.1-cp312-cp312-win_amd64.whl", hash = "sha256:4aa6cb806eb8541e92c9b313e96192c6b826e9eb0042720e2fa250d021079952", size = 4302006, upload-time = "2025-11-12T15:25:16.184Z" }, - { url = "https://files.pythonhosted.org/packages/e4/97/ab40550f482577f2788304c27220e8ba02c63313bd74cf2f8920526aac20/torchvision-0.24.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:8a6696db7fb71eadb2c6a48602106e136c785642e598eb1533e0b27744f2cce6", size = 1891435, upload-time = "2025-11-12T15:25:28.642Z" }, - { url = "https://files.pythonhosted.org/packages/30/65/ac0a3f9be6abdbe4e1d82c915d7e20de97e7fd0e9a277970508b015309f3/torchvision-0.24.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:db2125c46f9cb25dc740be831ce3ce99303cfe60439249a41b04fd9f373be671", size = 2338718, upload-time = "2025-11-12T15:25:26.19Z" }, - { url = "https://files.pythonhosted.org/packages/10/b5/5bba24ff9d325181508501ed7f0c3de8ed3dd2edca0784d48b144b6c5252/torchvision-0.24.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:f035f0cacd1f44a8ff6cb7ca3627d84c54d685055961d73a1a9fb9827a5414c8", size = 8049661, upload-time = "2025-11-12T15:25:22.558Z" }, - { url = "https://files.pythonhosted.org/packages/5c/ec/54a96ae9ab6a0dd66d4bba27771f892e36478a9c3489fa56e51c70abcc4d/torchvision-0.24.1-cp313-cp313-win_amd64.whl", hash = "sha256:16274823b93048e0a29d83415166a2e9e0bf4e1b432668357b657612a4802864", size = 4319808, upload-time = "2025-11-12T15:25:17.318Z" }, - { url = "https://files.pythonhosted.org/packages/d5/f3/a90a389a7e547f3eb8821b13f96ea7c0563cdefbbbb60a10e08dda9720ff/torchvision-0.24.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e3f96208b4bef54cd60e415545f5200346a65024e04f29a26cd0006dbf9e8e66", size = 2005342, upload-time = "2025-11-12T15:25:11.871Z" }, - { url = "https://files.pythonhosted.org/packages/a9/fe/ff27d2ed1b524078164bea1062f23d2618a5fc3208e247d6153c18c91a76/torchvision-0.24.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:f231f6a4f2aa6522713326d0d2563538fa72d613741ae364f9913027fa52ea35", size = 2341708, upload-time = "2025-11-12T15:25:25.08Z" }, - { url = "https://files.pythonhosted.org/packages/b1/b9/d6c903495cbdfd2533b3ef6f7b5643ff589ea062f8feb5c206ee79b9d9e5/torchvision-0.24.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:1540a9e7f8cf55fe17554482f5a125a7e426347b71de07327d5de6bfd8d17caa", size = 8177239, upload-time = "2025-11-12T15:25:18.554Z" }, - { url = "https://files.pythonhosted.org/packages/4f/2b/ba02e4261369c3798310483028495cf507e6cb3f394f42e4796981ecf3a7/torchvision-0.24.1-cp313-cp313t-win_amd64.whl", hash = "sha256:d83e16d70ea85d2f196d678bfb702c36be7a655b003abed84e465988b6128938", size = 4251604, upload-time = "2025-11-12T15:25:34.069Z" }, - { url = "https://files.pythonhosted.org/packages/42/84/577b2cef8f32094add5f52887867da4c2a3e6b4261538447e9b48eb25812/torchvision-0.24.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:cccf4b4fec7fdfcd3431b9ea75d1588c0a8596d0333245dafebee0462abe3388", size = 2005319, upload-time = "2025-11-12T15:25:23.827Z" }, - { url = "https://files.pythonhosted.org/packages/5f/34/ecb786bffe0159a3b49941a61caaae089853132f3cd1e8f555e3621f7e6f/torchvision-0.24.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:1b495edd3a8f9911292424117544f0b4ab780452e998649425d1f4b2bed6695f", size = 2338844, upload-time = "2025-11-12T15:25:32.625Z" }, - { url = "https://files.pythonhosted.org/packages/51/99/a84623786a6969504c87f2dc3892200f586ee13503f519d282faab0bb4f0/torchvision-0.24.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:ab211e1807dc3e53acf8f6638df9a7444c80c0ad050466e8d652b3e83776987b", size = 8175144, upload-time = "2025-11-12T15:25:31.355Z" }, - { url = "https://files.pythonhosted.org/packages/6d/ba/8fae3525b233e109317ce6a9c1de922ab2881737b029a7e88021f81e068f/torchvision-0.24.1-cp314-cp314-win_amd64.whl", hash = "sha256:18f9cb60e64b37b551cd605a3d62c15730c086362b40682d23e24b616a697d41", size = 4234459, upload-time = "2025-11-12T15:25:19.859Z" }, - { url = "https://files.pythonhosted.org/packages/50/33/481602c1c72d0485d4b3a6b48c9534b71c2957c9d83bf860eb837bf5a620/torchvision-0.24.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ec9d7379c519428395e4ffda4dbb99ec56be64b0a75b95989e00f9ec7ae0b2d7", size = 2005336, upload-time = "2025-11-12T15:25:27.225Z" }, - { url = "https://files.pythonhosted.org/packages/d0/7f/372de60bf3dd8f5593bd0d03f4aecf0d1fd58f5bc6943618d9d913f5e6d5/torchvision-0.24.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:af9201184c2712d808bd4eb656899011afdfce1e83721c7cb08000034df353fe", size = 2341704, upload-time = "2025-11-12T15:25:29.857Z" }, - { url = "https://files.pythonhosted.org/packages/36/9b/0f3b9ff3d0225ee2324ec663de0e7fb3eb855615ca958ac1875f22f1f8e5/torchvision-0.24.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:9ef95d819fd6df81bc7cc97b8f21a15d2c0d3ac5dbfaab5cbc2d2ce57114b19e", size = 8177422, upload-time = "2025-11-12T15:25:37.357Z" }, - { url = "https://files.pythonhosted.org/packages/d6/ab/e2bcc7c2f13d882a58f8b30ff86f794210b075736587ea50f8c545834f8a/torchvision-0.24.1-cp314-cp314t-win_amd64.whl", hash = "sha256:480b271d6edff83ac2e8d69bbb4cf2073f93366516a50d48f140ccfceedb002e", size = 4335190, upload-time = "2025-11-12T15:25:35.745Z" }, + { url = "https://files.pythonhosted.org/packages/3e/be/c704bceaf11c4f6b19d64337a34a877fcdfe3bd68160a8c9ae9bea4a35a3/torchvision-0.25.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:db74a551946b75d19f9996c419a799ffdf6a223ecf17c656f90da011f1d75b20", size = 1874923, upload-time = "2026-01-21T16:27:46.574Z" }, + { url = "https://files.pythonhosted.org/packages/ae/e9/f143cd71232430de1f547ceab840f68c55e127d72558b1061a71d0b193cd/torchvision-0.25.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:f49964f96644dbac2506dffe1a0a7ec0f2bf8cf7a588c3319fed26e6329ffdf3", size = 2344808, upload-time = "2026-01-21T16:27:43.191Z" }, + { url = "https://files.pythonhosted.org/packages/43/ae/ad5d6165797de234c9658752acb4fce65b78a6a18d82efdf8367c940d8da/torchvision-0.25.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:153c0d2cbc34b7cf2da19d73450f24ba36d2b75ec9211b9962b5022fb9e4ecee", size = 8070752, upload-time = "2026-01-21T16:27:33.748Z" }, + { url = "https://files.pythonhosted.org/packages/23/19/55b28aecdc7f38df57b8eb55eb0b14a62b470ed8efeb22cdc74224df1d6a/torchvision-0.25.0-cp311-cp311-win_amd64.whl", hash = "sha256:ea580ffd6094cc01914ad32f8c8118174f18974629af905cea08cb6d5d48c7b7", size = 4038722, upload-time = "2026-01-21T16:27:41.355Z" }, + { url = "https://files.pythonhosted.org/packages/56/3a/6ea0d73f49a9bef38a1b3a92e8dd455cea58470985d25635beab93841748/torchvision-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c2abe430c90b1d5e552680037d68da4eb80a5852ebb1c811b2b89d299b10573b", size = 1874920, upload-time = "2026-01-21T16:27:45.348Z" }, + { url = "https://files.pythonhosted.org/packages/51/f8/c0e1ef27c66e15406fece94930e7d6feee4cb6374bbc02d945a630d6426e/torchvision-0.25.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b75deafa2dfea3e2c2a525559b04783515e3463f6e830cb71de0fb7ea36fe233", size = 2344556, upload-time = "2026-01-21T16:27:40.125Z" }, + { url = "https://files.pythonhosted.org/packages/68/2f/f24b039169db474e8688f649377de082a965fbf85daf4e46c44412f1d15a/torchvision-0.25.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:f25aa9e380865b11ea6e9d99d84df86b9cc959f1a007cd966fc6f1ab2ed0e248", size = 8072351, upload-time = "2026-01-21T16:27:21.074Z" }, + { url = "https://files.pythonhosted.org/packages/ad/16/8f650c2e288977cf0f8f85184b90ee56ed170a4919347fc74ee99286ed6f/torchvision-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:f9c55ae8d673ab493325d1267cbd285bb94d56f99626c00ac4644de32a59ede3", size = 4303059, upload-time = "2026-01-21T16:27:11.08Z" }, + { url = "https://files.pythonhosted.org/packages/f5/5b/1562a04a6a5a4cf8cf40016a0cdeda91ede75d6962cff7f809a85ae966a5/torchvision-0.25.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:24e11199e4d84ba9c5ee7825ebdf1cd37ce8deec225117f10243cae984ced3ec", size = 1874918, upload-time = "2026-01-21T16:27:39.02Z" }, + { url = "https://files.pythonhosted.org/packages/36/b1/3d6c42f62c272ce34fcce609bb8939bdf873dab5f1b798fd4e880255f129/torchvision-0.25.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5f271136d2d2c0b7a24c5671795c6e4fd8da4e0ea98aeb1041f62bc04c4370ef", size = 2309106, upload-time = "2026-01-21T16:27:30.624Z" }, + { url = "https://files.pythonhosted.org/packages/c7/60/59bb9c8b67cce356daeed4cb96a717caa4f69c9822f72e223a0eae7a9bd9/torchvision-0.25.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:855c0dc6d37f462482da7531c6788518baedca1e0847f3df42a911713acdfe52", size = 8071522, upload-time = "2026-01-21T16:27:29.392Z" }, + { url = "https://files.pythonhosted.org/packages/32/a5/9a9b1de0720f884ea50dbf9acb22cbe5312e51d7b8c4ac6ba9b51efd9bba/torchvision-0.25.0-cp313-cp313-win_amd64.whl", hash = "sha256:cef0196be31be421f6f462d1e9da1101be7332d91984caa6f8022e6c78a5877f", size = 4321911, upload-time = "2026-01-21T16:27:35.195Z" }, + { url = "https://files.pythonhosted.org/packages/52/99/dca81ed21ebaeff2b67cc9f815a20fdaa418b69f5f9ea4c6ed71721470db/torchvision-0.25.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a8f8061284395ce31bcd460f2169013382ccf411148ceb2ee38e718e9860f5a7", size = 1896209, upload-time = "2026-01-21T16:27:32.159Z" }, + { url = "https://files.pythonhosted.org/packages/28/cc/2103149761fdb4eaed58a53e8437b2d716d48f05174fab1d9fcf1e2a2244/torchvision-0.25.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:146d02c9876858420adf41f3189fe90e3d6a409cbfa65454c09f25fb33bf7266", size = 2310735, upload-time = "2026-01-21T16:27:22.327Z" }, + { url = "https://files.pythonhosted.org/packages/76/ad/f4c985ad52ddd3b22711c588501be1b330adaeaf6850317f66751711b78c/torchvision-0.25.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:c4d395cb2c4a2712f6eb93a34476cdf7aae74bb6ea2ea1917f858e96344b00aa", size = 8089557, upload-time = "2026-01-21T16:27:27.666Z" }, + { url = "https://files.pythonhosted.org/packages/63/cc/0ea68b5802e5e3c31f44b307e74947bad5a38cc655231d845534ed50ddb8/torchvision-0.25.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5e6b449e9fa7d642142c0e27c41e5a43b508d57ed8e79b7c0a0c28652da8678c", size = 4344260, upload-time = "2026-01-21T16:27:17.018Z" }, + { url = "https://files.pythonhosted.org/packages/9e/1f/fa839532660e2602b7e704d65010787c5bb296258b44fa8b9c1cd6175e7d/torchvision-0.25.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:620a236288d594dcec7634c754484542dc0a5c1b0e0b83a34bda5e91e9b7c3a1", size = 1896193, upload-time = "2026-01-21T16:27:24.785Z" }, + { url = "https://files.pythonhosted.org/packages/80/ed/d51889da7ceaf5ff7a0574fb28f9b6b223df19667265395891f81b364ab3/torchvision-0.25.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:0b5e7f50002a8145a98c5694a018e738c50e2972608310c7e88e1bd4c058f6ce", size = 2309331, upload-time = "2026-01-21T16:27:19.97Z" }, + { url = "https://files.pythonhosted.org/packages/90/a5/f93fcffaddd8f12f9e812256830ec9c9ca65abbf1bc369379f9c364d1ff4/torchvision-0.25.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:632db02300e83793812eee4f61ae6a2686dab10b4cfd628b620dc47747aa9d03", size = 8088713, upload-time = "2026-01-21T16:27:15.281Z" }, + { url = "https://files.pythonhosted.org/packages/1f/eb/d0096eed5690d962853213f2ee00d91478dfcb586b62dbbb449fb8abc3a6/torchvision-0.25.0-cp314-cp314-win_amd64.whl", hash = "sha256:d1abd5ed030c708f5dbf4812ad5f6fbe9384b63c40d6bd79f8df41a4a759a917", size = 4325058, upload-time = "2026-01-21T16:27:26.165Z" }, + { url = "https://files.pythonhosted.org/packages/97/36/96374a4c7ab50dea9787ce987815614ccfe988a42e10ac1a2e3e5b60319a/torchvision-0.25.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ad9a8a5877782944d99186e4502a614770fe906626d76e9cd32446a0ac3075f2", size = 1896207, upload-time = "2026-01-21T16:27:23.383Z" }, + { url = "https://files.pythonhosted.org/packages/b5/e2/7abb10a867db79b226b41da419b63b69c0bd5b82438c4a4ed50e084c552f/torchvision-0.25.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:40a122c3cf4d14b651f095e0f672b688dde78632783fc5cd3d4d5e4f6a828563", size = 2310741, upload-time = "2026-01-21T16:27:18.712Z" }, + { url = "https://files.pythonhosted.org/packages/08/e6/0927784e6ffc340b6676befde1c60260bd51641c9c574b9298d791a9cda4/torchvision-0.25.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:846890161b825b38aa85fc37fb3ba5eea74e7091ff28bab378287111483b6443", size = 8089772, upload-time = "2026-01-21T16:27:14.048Z" }, + { url = "https://files.pythonhosted.org/packages/b6/37/e7ca4ec820d434c0f23f824eb29f0676a0c3e7a118f1514f5b949c3356da/torchvision-0.25.0-cp314-cp314t-win_amd64.whl", hash = "sha256:f07f01d27375ad89d72aa2b3f2180f07da95dd9d2e4c758e015c0acb2da72977", size = 4425879, upload-time = "2026-01-21T16:27:12.579Z" }, ] [[package]] @@ -3151,15 +3220,15 @@ wheels = [ [[package]] name = "triton" -version = "3.5.1" +version = "3.6.0" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b0/72/ec90c3519eaf168f22cb1757ad412f3a2add4782ad3a92861c9ad135d886/triton-3.5.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:61413522a48add32302353fdbaaf92daaaab06f6b5e3229940d21b5207f47579", size = 170425802, upload-time = "2025-11-11T17:40:53.209Z" }, - { url = "https://files.pythonhosted.org/packages/f2/50/9a8358d3ef58162c0a415d173cfb45b67de60176e1024f71fbc4d24c0b6d/triton-3.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d2c6b915a03888ab931a9fd3e55ba36785e1fe70cbea0b40c6ef93b20fc85232", size = 170470207, upload-time = "2025-11-11T17:41:00.253Z" }, - { url = "https://files.pythonhosted.org/packages/27/46/8c3bbb5b0a19313f50edcaa363b599e5a1a5ac9683ead82b9b80fe497c8d/triton-3.5.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f3f4346b6ebbd4fad18773f5ba839114f4826037c9f2f34e0148894cd5dd3dba", size = 170470410, upload-time = "2025-11-11T17:41:06.319Z" }, - { url = "https://files.pythonhosted.org/packages/37/92/e97fcc6b2c27cdb87ce5ee063d77f8f26f19f06916aa680464c8104ef0f6/triton-3.5.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0b4d2c70127fca6a23e247f9348b8adde979d2e7a20391bfbabaac6aebc7e6a8", size = 170579924, upload-time = "2025-11-11T17:41:12.455Z" }, - { url = "https://files.pythonhosted.org/packages/a4/e6/c595c35e5c50c4bc56a7bac96493dad321e9e29b953b526bbbe20f9911d0/triton-3.5.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0637b1efb1db599a8e9dc960d53ab6e4637db7d4ab6630a0974705d77b14b60", size = 170480488, upload-time = "2025-11-11T17:41:18.222Z" }, - { url = "https://files.pythonhosted.org/packages/16/b5/b0d3d8b901b6a04ca38df5e24c27e53afb15b93624d7fd7d658c7cd9352a/triton-3.5.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bac7f7d959ad0f48c0e97d6643a1cc0fd5786fe61cb1f83b537c6b2d54776478", size = 170582192, upload-time = "2025-11-11T17:41:23.963Z" }, + { url = "https://files.pythonhosted.org/packages/e0/12/b05ba554d2c623bffa59922b94b0775673de251f468a9609bc9e45de95e9/triton-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8e323d608e3a9bfcc2d9efcc90ceefb764a82b99dea12a86d643c72539ad5d3", size = 188214640, upload-time = "2026-01-20T16:00:35.869Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, + { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" }, + { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" }, + { url = "https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a17a5d5985f0ac494ed8a8e54568f092f7057ef60e1b0fa09d3fd1512064e803", size = 188273063, upload-time = "2026-01-20T16:01:07.278Z" }, + { url = "https://files.pythonhosted.org/packages/f6/56/6113c23ff46c00aae423333eb58b3e60bdfe9179d542781955a5e1514cb3/triton-3.6.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46bd1c1af4b6704e554cad2eeb3b0a6513a980d470ccfa63189737340c7746a7", size = 188397994, upload-time = "2026-01-20T16:01:14.236Z" }, ] [[package]] From 46d00af03ee51c3bfaf1525114b36fa39cb1eac5 Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Thu, 9 Apr 2026 10:41:12 -0400 Subject: [PATCH 2/5] compare memory usage across torch versions --- scripts/benchmark_conv3d_memory.py | 160 +++++++++++++++++++++++++++++ 1 file changed, 160 insertions(+) create mode 100644 scripts/benchmark_conv3d_memory.py diff --git a/scripts/benchmark_conv3d_memory.py b/scripts/benchmark_conv3d_memory.py new file mode 100644 index 00000000..1381f5e6 --- /dev/null +++ b/scripts/benchmark_conv3d_memory.py @@ -0,0 +1,160 @@ +"""Standalone nn.Conv3d peak-memory benchmark across dtypes. + +Fixed tensor (consistent across all torch versions): + Input: (N=1, C_in=32, D=64, H=64, W=64) + Conv3d: in=32, out=32, kernel=5, padding=2, padding_mode='zeros' + +Collects: + - torch / CUDA / cuDNN versions + - GPU name and total VRAM + - Forward and backward peak GPU memory (MB) for float32 and bfloat16 + - Forward and backward wall-clock time (ms) + +Usage: + python benchmark_conv3d_memory.py --output results.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +import torch +import torch.nn as nn + +# ── Fixed benchmark configuration ───────────────────────────────────────────── +N, CIN, COUT = 1, 32, 32 +D = H = W = 64 +K = 5 +PADDING = K // 2 +WARMUP_ITERS = 3 +BENCH_ITERS = 5 # average over multiple runs for stable timing + + +# ── Memory helpers ──────────────────────────────────────────────────────────── + + +def reset_mem(): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + +def peak_mb() -> float: + return torch.cuda.max_memory_allocated() / 1024**2 + + +# ── Per-dtype measurement ───────────────────────────────────────────────────── + + +def bench_dtype(dtype_str: str, device: torch.device) -> dict: + dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16}[dtype_str] + + torch.manual_seed(42) + conv = nn.Conv3d(CIN, COUT, K, padding=PADDING).to(device=device, dtype=dtype) + x_base = torch.randn(N, CIN, D, H, W, device=device, dtype=dtype) + + # ── Warmup (no grad, no memory tracking) ────────────────────────────────── + for _ in range(WARMUP_ITERS): + with torch.no_grad(): + _ = conv(x_base) + torch.cuda.synchronize() + + # ── Forward pass ────────────────────────────────────────────────────────── + fwd_peaks, fwd_times = [], [] + for _ in range(BENCH_ITERS): + x = x_base.detach().requires_grad_(True) + reset_mem() + torch.cuda.synchronize() + t0 = time.perf_counter() + out = conv(x) + torch.cuda.synchronize() + fwd_times.append(time.perf_counter() - t0) + fwd_peaks.append(peak_mb()) + + # ── Backward pass ───────────────────────────────────────────────────────── + bwd_peaks, bwd_times = [], [] + for _ in range(BENCH_ITERS): + x = x_base.detach().requires_grad_(True) + out = conv(x) # fresh forward to build the computation graph + reset_mem() + torch.cuda.synchronize() + t0 = time.perf_counter() + out.sum().backward() + torch.cuda.synchronize() + bwd_times.append(time.perf_counter() - t0) + bwd_peaks.append(peak_mb()) + + def avg(lst: list) -> float: + return round(sum(lst) / len(lst), 2) + + return { + "fwd_peak_mb": avg(fwd_peaks), + "bwd_peak_mb": avg(bwd_peaks), + "fwd_time_ms": round(avg(fwd_times) * 1e3, 3), + "bwd_time_ms": round(avg(bwd_times) * 1e3, 3), + } + + +# ── Entry point ─────────────────────────────────────────────────────────────── + + +def main(): + parser = argparse.ArgumentParser(description="Conv3d memory benchmark") + parser.add_argument("--output", required=True, help="Path to write JSON results") + args = parser.parse_args() + + if not torch.cuda.is_available(): + result = {"error": "CUDA not available"} + with Path.open(args.output, "w") as f: + json.dump(result, f, indent=2) + sys.exit(1) + + device = torch.device("cuda:0") + gpu_props = torch.cuda.get_device_properties(0) + + # Decode cuDNN version integer. + # cuDNN < 9 : MAJOR*1000 + MINOR*100 + PATCH (e.g. 8904 → "8.9.4") + # cuDNN 9+ : MAJOR*10000 + MINOR*1000 + PATCH*100 + BUILD (e.g. 91002 → "9.1.0.2") + raw_cudnn = torch.backends.cudnn.version() + if not isinstance(raw_cudnn, int): + cudnn_str = str(raw_cudnn) + elif raw_cudnn >= 10000: + major = raw_cudnn // 10000 + minor = (raw_cudnn % 10000) // 1000 + patch = (raw_cudnn % 1000) // 100 + build = raw_cudnn % 100 + cudnn_str = f"{major}.{minor}.{patch}.{build}" + else: + major = raw_cudnn // 1000 + minor = (raw_cudnn % 1000) // 100 + patch = raw_cudnn % 100 + cudnn_str = f"{major}.{minor}.{patch}" + + results = { + "torch_version": torch.__version__, + "cuda_runtime": torch.version.cuda, + "cudnn_version": cudnn_str, + "gpu_name": gpu_props.name, + "gpu_total_memory_gb": round(gpu_props.total_memory / 1024**3, 1), + "gpu_sm_count": gpu_props.multi_processor_count, + "tensor_shape": [N, CIN, D, H, W], + "conv_config": f"Conv3d(in={CIN}, out={COUT}, k={K}, padding={PADDING})", + "bench_iters": BENCH_ITERS, + "measurements": {}, + } + + for dtype in ("float32", "bfloat16"): + try: + results["measurements"][dtype] = bench_dtype(dtype, device) + except Exception as exc: + results["measurements"][dtype] = {"error": str(exc)} + + with Path.open(args.output, "w") as f: + json.dump(results, f, indent=2) + + +if __name__ == "__main__": + main() From b11cc87dcbe488f50de5197b94ca8ea349e7b94b Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Thu, 9 Apr 2026 10:43:18 -0400 Subject: [PATCH 3/5] modified main.py --- src/electrai/entrypoints/main.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/electrai/entrypoints/main.py b/src/electrai/entrypoints/main.py index ad76ca28..6ccc4725 100644 --- a/src/electrai/entrypoints/main.py +++ b/src/electrai/entrypoints/main.py @@ -8,8 +8,6 @@ from electrai.entrypoints.train import train torch.backends.cudnn.conv.fp32_precision = "tf32" -torch.backends.cudnn.benchmark = False -torch.backends.cudnn.deterministic = True def main() -> None: From 9f0f4cfb1b5673252f86fcaaaa44929671e31e40 Mon Sep 17 00:00:00 2001 From: Betsy Cannon Date: Fri, 17 Apr 2026 11:33:45 -0400 Subject: [PATCH 4/5] Polish benchmark_conv3d_memory.py: default output path, clarify peak semantics, drop stale ref - Make `--output` default to `results.json` so the script runs standalone without args. - Add comment explaining that `bwd_peak_mb` reflects peak *during* backward (includes fresh-forward activations held by the autograd graph); relative comparisons across torch versions remain valid. - `del out` after the forward loop so the stale reference doesn't linger into the backward loop. Co-Authored-By: Claude Opus 4.7 (1M context) --- scripts/benchmark_conv3d_memory.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/scripts/benchmark_conv3d_memory.py b/scripts/benchmark_conv3d_memory.py index 1381f5e6..93ec228e 100644 --- a/scripts/benchmark_conv3d_memory.py +++ b/scripts/benchmark_conv3d_memory.py @@ -11,7 +11,7 @@ - Forward and backward wall-clock time (ms) Usage: - python benchmark_conv3d_memory.py --output results.json + python benchmark_conv3d_memory.py [--output results.json] """ from __future__ import annotations @@ -73,8 +73,14 @@ def bench_dtype(dtype_str: str, device: torch.device) -> dict: torch.cuda.synchronize() fwd_times.append(time.perf_counter() - t0) fwd_peaks.append(peak_mb()) + del out # ── Backward pass ───────────────────────────────────────────────────────── + # Note: reset_peak_memory_stats() zeros the peak counter but does not free + # live tensors, so bwd_peak_mb reflects the peak *during* backward, which + # includes the fresh forward activations still held by the autograd graph — + # not purely the backward-specific allocation. The methodology is identical + # across torch versions, so relative comparisons remain valid. bwd_peaks, bwd_times = [], [] for _ in range(BENCH_ITERS): x = x_base.detach().requires_grad_(True) @@ -103,7 +109,11 @@ def avg(lst: list) -> float: def main(): parser = argparse.ArgumentParser(description="Conv3d memory benchmark") - parser.add_argument("--output", required=True, help="Path to write JSON results") + parser.add_argument( + "--output", + default="results.json", + help="Path to write JSON results (default: results.json)", + ) args = parser.parse_args() if not torch.cuda.is_available(): From 80ec65ca3f3523c411fdd6988f888e5b0a80234d Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Wed, 20 May 2026 13:37:20 -0400 Subject: [PATCH 5/5] f32 vs bf16-mixed benchmark (#119) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary This PR adds a benchmark script that trains one sample (task ID) at a time for 3 epochs under both f32 and bf16-mixed precision on a single A100 GPU, recording peak GPU memory, forward/backward times, and OOM status for 10 large-grid samples. ## The problem The 10 task IDs selected for this benchmark are Materials Project entries with relatively large charge-density grids, spanning 3.4 M – 46.7 M voxels across a variety of shapes and aspect ratios. It was not previously known which samples would fit on an A100 (79.3 GB) under f32 vs bf16-mixed, or how training time scales with grid volume. This script establishes those baselines. ## Benchmark results (A100-SXM4-80GB, CUDA 12.8) **Config**: `n_channels=32`, `n_residual_blocks=1`, `kernel_size=5`, `depth=2`, `batch_size=1`, single GPU, 3 epochs per experiment. | Task ID | Grid shape | Voxels | f32 status | f32 peak (GB) | f32 epoch (s) | bf16 status | bf16 peak (GB) | bf16 epoch (s) | Mem ratio | |---------|-----------|:------:|:----------:|:-------------:|:-------------:|:-----------:|:--------------:|:--------------:|:---------:| | mp-1890579 | 56 × 56 × 1080 | 3.4 M | ✅ | 10.5 | 1.65 | ✅ | 5.5 | 1.07 | 1.91× | | mp-1849767 | 60 × 60 × 1120 | 4.0 M | ✅ | 12.4 | 1.95 | ✅ | 6.5 | 1.27 | 1.91× | | mp-1851604 | 60 × 60 × 1120 | 4.0 M | ✅ | 12.4 | 1.93 | ✅ | 6.5 | 1.25 | 1.91× | | mp-1862536 | 80 × 80 × 1024 | 6.6 M | ✅ | 19.9 | 3.17 | ✅ | 10.4 | 1.94 | 1.91× | | mp-1847208 | 1120 × 84 × 84 | 7.9 M | ✅ | 23.9 | 3.87 | ✅ | 12.5 | 2.46 | 1.91× | | mp-1936557 | 80 × 756 × 216 | 13.1 M | ✅ | 39.1 | 6.47 | ✅ | 20.4 | 3.73 | 1.91× | | mp-1850168 | 972 × 240 × 128 | 29.9 M | ❌ OOM | 70.2 | — | ✅ | 46.3 | 8.77 | — | | mp-1887804 | 320 × 320 × 320 | 32.8 M | ❌ OOM | 68.8 | — | ✅ | 50.6 | 169.91 | — | | mp-1889246 | 540 × 144 × 432 | 33.6 M | ❌ OOM | 70.6 | — | ✅ | 51.9 | 208.94 | — | | mp-1871122 | 360 × 360 × 360 | 46.7 M | ❌ OOM | 71.4 | — | ❌ OOM | 71.9 | — | — | ### Summary | Precision | ✅ Completed | ❌ OOM | |-----------|:-----------:|:------:| | f32 | 6 / 10 | 4 / 10 | | bf16-mixed | 9 / 10 | 1 / 10 (mp-1871122, 360³) | ### Key findings - **Memory**: bf16-mixed yields a consistent **1.91× reduction** in peak GPU memory for all grids that fit under both precisions. - **Speed**: bf16-mixed is **1.5–1.7× faster** per epoch on grids ≤ 13 M voxels; the backward pass benefits the most (~2.2×). - **OOM threshold on A100**: ~13 M voxels under f32; ~33 M voxels under bf16-mixed. - **Largest grid (360³, 46.7 M voxels)**: exceeds A100 capacity even with bf16-mixed. - **Very large grids (320³, 540×144×432)**: fit under bf16-mixed but epoch times are 170–209 s, dominated by the backward pass (~144–176 s through large skip connections). ## Files - `scripts/benchmark_precision.py` -- Runs 3-epoch single-GPU training for each (task_id, precision) pair, recording CUDA-event forward/backward times and peak GPU memory --------- Co-authored-by: Hananeh Oliaei Co-authored-by: Betsy Cannon Co-authored-by: Claude Opus 4.7 (1M context) --- scripts/benchmark_gpus.py | 149 +++++++++++++++++++ scripts/benchmark_precision.py | 252 +++++++++++++++++++++++++++++++++ 2 files changed, 401 insertions(+) create mode 100644 scripts/benchmark_gpus.py create mode 100644 scripts/benchmark_precision.py diff --git a/scripts/benchmark_gpus.py b/scripts/benchmark_gpus.py new file mode 100644 index 00000000..a3e9b26d --- /dev/null +++ b/scripts/benchmark_gpus.py @@ -0,0 +1,149 @@ +""" +Compare per-sample benchmark results across two GPUs (e.g. A100 vs H200). + +Reads two JSON files produced by benchmark_precision.py and prints +a markdown report section with a combined table (memory, time, and ratios). + +Usage: + uv run python scripts/benchmark_gpus.py \ + --gpu1 benchmark_results/per_task_precision.json \ + --gpu2 benchmark_results/per_task_precision_h200.json \ + --out benchmark_results/gpu_comparison.md +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def load(path: Path) -> tuple[dict, list[dict]]: + d = json.loads(path.read_text()) + return d["gpu"], d["results"] + + +def by_key(results: list[dict]) -> dict[tuple[str, str], dict]: + return {(r["task_id"], r["precision"]): r for r in results} + + +def ratio(a, b): + """Return a/b ratio string.""" + if a is None or b is None: + return "—" + return f"{a / b:.2f}x" + + +def status(r): + if r is None: + return "—" + return "❌ OOM" if r["oom"] else "✅" + + +def peak_gb(r): + if r is None or r["oom"] or r["peak_mem_mb"] is None: + return None + return round(r["peak_mem_mb"] / 1024, 2) + + +def epoch_s(r): + if r is None or r["oom"]: + return None + return r["avg_epoch_s"] + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--gpu1", type=Path, required=True, help="First GPU results JSON (reference)" + ) + parser.add_argument( + "--gpu2", type=Path, required=True, help="Second GPU results JSON" + ) + parser.add_argument( + "--out", type=Path, default=None, help="Output markdown file (default: stdout)" + ) + args = parser.parse_args() + + gpu1_info, res1 = load(args.gpu1) + gpu2_info, res2 = load(args.gpu2) + + idx1 = by_key(res1) + idx2 = by_key(res2) + + gpu1_name = gpu1_info["name"] + gpu2_name = gpu2_info["name"] + + # grid shape lookup + grid_shapes = {} + for r in res1 + res2: + if r.get("grid_shape"): + grid_shapes[r["task_id"]] = r["grid_shape"] + + def shape_str(tid): + s = grid_shapes.get(tid) + return f"{s[0]} x {s[1]} x {s[2]}" if s else "—" + + def voxels(tid): + s = grid_shapes.get(tid) + return f"{s[0] * s[1] * s[2] / 1e6:.1f} M" if s else "—" + + all_task_ids = sorted({k[0] for k in set(idx1) | set(idx2)}) + + lines = [] + a = lines.append + + a("| GPU | Model | VRAM |") + a("|-----|-------|:----:|") + a(f"| GPU 1 (reference) | {gpu1_name} | {gpu1_info['total_mem_gb']} GB |") + a(f"| GPU 2 | {gpu2_name} | {gpu2_info['total_mem_gb']} GB |") + a("") + + for prec in ["f32", "bf16-mixed"]: + a(f"### {prec}\n") + a( + f"| Task ID | Grid shape | Voxels " + f"| {gpu1_name} status | {gpu1_name} peak (GB) | {gpu1_name} epoch (s) " + f"| {gpu2_name} status | {gpu2_name} peak (GB) | {gpu2_name} epoch (s) " + f"| Peak mem ratio (GPU1/GPU2) | Epoch time ratio (GPU1/GPU2) |" + ) + a( + "|---------|-----------|:------:" + "|:---------:|:-------------------:|:--------------------:" + "|:---------:|:-------------------:|:--------------------:" + "|:-------------------------:|:----------------------------:|" + ) + + for tid in all_task_ids: + key = (tid, prec) + r1 = idx1.get(key) + r2 = idx2.get(key) + + p1, p2 = peak_gb(r1), peak_gb(r2) + e1, e2 = epoch_s(r1), epoch_s(r2) + + p1_str = "—" if p1 is None else f"{p1:.2f}" + p2_str = "—" if p2 is None else f"{p2:.2f}" + e1_str = "—" if e1 is None else f"{e1:.2f}" + e2_str = "—" if e2 is None else f"{e2:.2f}" + + a( + f"| {tid} | {shape_str(tid)} | {voxels(tid)} " + f"| {status(r1)} | {p1_str} | {e1_str} " + f"| {status(r2)} | {p2_str} | {e2_str} " + f"| {ratio(p1, p2)} | {ratio(e1, e2)} |" + ) + a("") + + output = "\n".join(lines) + + if args.out: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(output) + else: + sys.stdout.write(output + "\n") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_precision.py b/scripts/benchmark_precision.py new file mode 100644 index 00000000..9fab8127 --- /dev/null +++ b/scripts/benchmark_precision.py @@ -0,0 +1,252 @@ +""" +Benchmark: f32 vs bf16-mixed training, one task_id at a time. + +For each (task_id, precision) pair: + - Build model and load the single zarr sample + - Run 3 epochs of forward+backward (full training step) + - Record peak GPU memory and per-epoch wall time + - Catch OOM and record it + +Usage: + uv run python scripts/benchmark_precision.py \ + --config path/to/config.yaml \ + --zarr_root path/to/zarr_root \ + --results path/to/results.json +""" + +from __future__ import annotations + +import argparse +import gc +import json +import time +import warnings +from pathlib import Path +from types import SimpleNamespace + +import torch +import yaml +from hydra.utils import instantiate +from torch.utils.data import DataLoader + +from electrai.dataloader.collate import collate_fn +from electrai.dataloader.large_grid_json_gz import LargeGridZarrDataset +from electrai.model.loss.charge import NormMAE + +# --------------------------------------------------------------------------- +# Defaults (override via CLI arguments) +# --------------------------------------------------------------------------- +TASK_IDS = [ + "mp-1862536", + "mp-1936557", + "mp-1847208", + "mp-1850168", + "mp-1890579", + "mp-1871122", + "mp-1889246", + "mp-1849767", + "mp-1851604", + "mp-1887804", +] +PRECISIONS = ["f32", "bf16-mixed"] +EPOCHS = 3 +DEVICE = torch.device("cuda:0") + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Per-sample f32 vs bf16-mixed benchmark" + ) + parser.add_argument( + "--config", type=Path, required=True, help="Path to training config YAML" + ) + parser.add_argument( + "--zarr_root", type=str, required=True, help="Path to zarr root directory" + ) + parser.add_argument( + "--results", + type=Path, + default=Path("benchmark_results/per_sample_precision.json"), + help="Output JSON path (default: benchmark_results/per_sample_precision.json)", + ) + parser.add_argument( + "--task_ids", + nargs="+", + default=None, + help="Override task IDs to benchmark (default: all 10)", + ) + return parser.parse_args() + + +def load_cfg(config_path: Path): + with config_path.open() as f: + d = yaml.safe_load(f) + return SimpleNamespace(**d) + + +def build_model(cfg): + model = instantiate(cfg.model) + return model.to(DEVICE) + + +def get_grid_shape(task_id: str, zarr_root: str) -> tuple[int, ...]: + """Return the charge density grid shape for a task_id.""" + import zarr + + store = zarr.open_group(f"{zarr_root}/{task_id}.zarr", mode="r") + return tuple(store["charge_density_total"].shape) + + +def run_experiment(task_id: str, precision: str, cfg, zarr_root: str) -> dict: + use_bf16 = precision == "bf16-mixed" + + # Initialize so all are in scope for the finally block even on early OOM + model = optimizer = loss_fn = dataset = loader = None + x = y = pred = loss = None + + # ---- build dataloader (single sample, no split needed) ---- + dataset = LargeGridZarrDataset( + task_ids=[task_id], + zarr_root=zarr_root, + precision="f32", # load as f32; autocast handles the cast + ) + loader = DataLoader( + dataset, batch_size=1, shuffle=False, num_workers=0, collate_fn=collate_fn + ) + + # ---- build model & optimizer ---- + model = build_model(cfg) + optimizer = torch.optim.Adam(model.parameters(), lr=float(cfg.lr)) + loss_fn = NormMAE().to(DEVICE) + + gc.collect() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats(DEVICE) + + epoch_times: list[float] = [] + fwd_times: list[float] = [] + bwd_times: list[float] = [] + oom = False + peak_mem_mb: float | None = None + + try: + for _epoch in range(EPOCHS): + epoch_start = time.perf_counter() + + for batch in loader: + x = batch["data"] + y = batch["label"] + + # collate_fn may return a list when shapes vary + if isinstance(x, list): + x = x[0].unsqueeze(0).to(DEVICE) + y = y[0].unsqueeze(0).to(DEVICE) + else: + x = x.to(DEVICE) + y = y.to(DEVICE) + + optimizer.zero_grad() + + # ---- forward ---- + fwd_start = torch.cuda.Event(enable_timing=True) + fwd_end = torch.cuda.Event(enable_timing=True) + fwd_start.record() + + if use_bf16: + with torch.autocast("cuda", dtype=torch.bfloat16): + pred = model(x) + loss = loss_fn(pred, y) + else: + pred = model(x) + loss = loss_fn(pred, y) + + fwd_end.record() + torch.cuda.synchronize() + fwd_ms = fwd_start.elapsed_time(fwd_end) + + # ---- backward ---- + bwd_start = torch.cuda.Event(enable_timing=True) + bwd_end = torch.cuda.Event(enable_timing=True) + bwd_start.record() + loss.backward() + bwd_end.record() + torch.cuda.synchronize() + bwd_ms = bwd_start.elapsed_time(bwd_end) + + optimizer.step() + + fwd_times.append(fwd_ms / 1000.0) + bwd_times.append(bwd_ms / 1000.0) + + epoch_end = time.perf_counter() + epoch_times.append(epoch_end - epoch_start) + + peak_mem_mb = torch.cuda.max_memory_allocated(DEVICE) / 1024**2 + + except torch.cuda.OutOfMemoryError: + oom = True + peak_mem_mb = torch.cuda.max_memory_allocated(DEVICE) / 1024**2 + + finally: + del model, optimizer, loss_fn, dataset, loader, x, y, pred, loss + gc.collect() + torch.cuda.empty_cache() + + return { + "task_id": task_id, + "precision": precision, + "oom": oom, + "peak_mem_mb": round(peak_mem_mb, 1) if peak_mem_mb is not None else None, + "epoch_times_s": [round(t, 2) for t in epoch_times], + "fwd_times_s": [round(t, 3) for t in fwd_times], + "bwd_times_s": [round(t, 3) for t in bwd_times], + "avg_epoch_s": round(sum(epoch_times) / len(epoch_times), 2) + if epoch_times + else None, + "avg_fwd_s": round(sum(fwd_times) / len(fwd_times), 3) if fwd_times else None, + "avg_bwd_s": round(sum(bwd_times) / len(bwd_times), 3) if bwd_times else None, + } + + +def get_gpu_info() -> dict: + props = torch.cuda.get_device_properties(DEVICE) + return { + "name": props.name, + "total_mem_gb": round(props.total_memory / 1024**3, 1), + "major": props.major, + "minor": props.minor, + "multi_processor_count": props.multi_processor_count, + } + + +def main(): + args = parse_args() + args.results.parent.mkdir(parents=True, exist_ok=True) + cfg = load_cfg(args.config) + task_ids = args.task_ids if args.task_ids is not None else TASK_IDS + + gpu_info = get_gpu_info() + + # get grid shapes upfront + grid_shapes = {} + for tid in task_ids: + try: + grid_shapes[tid] = get_grid_shape(tid, args.zarr_root) + except Exception as e: + warnings.warn(f"could not read grid shape for {tid}: {e}", stacklevel=2) + grid_shapes[tid] = None + + results = [] + for tid in task_ids: + for prec in PRECISIONS: + r = run_experiment(tid, prec, cfg, args.zarr_root) + r["grid_shape"] = grid_shapes.get(tid) + results.append(r) + # save incrementally + args.results.write_text( + json.dumps({"gpu": gpu_info, "results": results}, indent=2) + ) + + +if __name__ == "__main__": + main()