Skip to content

Commit

Permalink
feat (CI): cache model weights (kornia#2984)
Browse files Browse the repository at this point in the history
  • Loading branch information
johnnv1 authored Aug 18, 2024
1 parent ac99883 commit 93114bf
Show file tree
Hide file tree
Showing 4 changed files with 147 additions and 16 deletions.
21 changes: 21 additions & 0 deletions .github/download-models-weights.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
import argparse

import torch

fonts = {
"sold2_wireframe": "http://cmp.felk.cvut.cz/~mishkdmy/models/sold2_wireframe.pth",
}

if __name__ == "__main__":
parser = argparse.ArgumentParser("WeightsDownloader")
parser.add_argument("--target_directory", "-t", required=False, default="target_directory")

args = parser.parse_args()

torch.hub.set_dir(args.target_directory)

for name, url in fonts.items():
print(f"Downloading weights of `{name}` from `url`. Caching to dir `{args.target_directory}`")
torch.hub.load_state_dict_from_url(url, model_dir=args.target_directory, map_location=torch.device("cpu"))

raise SystemExit(0)
65 changes: 59 additions & 6 deletions .github/workflows/pr_test_cpu.yml
Original file line number Diff line number Diff line change
Expand Up @@ -13,41 +13,88 @@ concurrency:
cancel-in-progress: true

jobs:

pre-tests:
runs-on: ubuntu-latest
outputs:
hash: ${{ steps.hashid.outputs.weights-hash }}
steps:
- uses: kornia/workflows/.github/actions/env@v1.12.1
- uses: actions/cache@v4
id: cache-weights
with:
path: weights/
key: model-weights-${{ hashFiles('.github/download-models-weights.py') }}
enableCrossOsArchive: true

- name: Download models weights...
if: steps.cache-weights.outputs.cache-hit != 'true'
run: python .github/download-models-weights.py -t weights/

- name: write hashid
id: hashid
run: echo "weights-hash=${{ hashFiles('.github/download-models-weights.py') }}" >> "$GITHUB_OUTPUT"

tests-cpu:
needs: [pre-tests]
strategy:
fail-fast: true
matrix:
os: ['Ubuntu-latest', 'Windows-latest']
pytorch-dtype: ['float32', 'float64']

uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
with:
os: ${{ matrix.os }}
python-version: '["3.8", "3.12"]'
pytorch-version: '["1.9.1", "2.4.0"]'
pytorch-dtype: ${{ matrix.pytorch-dtype }}
cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
tests-cpu-macos:
uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
needs: [pre-tests]
uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
with:
os: 'MacOS-latest'
python-version: '["3.8", "3.12"]'
pytorch-dtype: 'float32'
cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
coverage:
uses: kornia/workflows/.github/workflows/coverage.yml@v1.11.0
needs: [pre-tests]
uses: kornia/workflows/.github/workflows/coverage.yml@v1.12.1
with:
cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
typing:
uses: kornia/workflows/.github/workflows/mypy.yml@v1.11.0
uses: kornia/workflows/.github/workflows/mypy.yml@v1.12.1

tutorials:
uses: kornia/workflows/.github/workflows/tutorials.yml@v1.12.1

docs:
uses: kornia/workflows/.github/workflows/docs.yml@v1.11.0
needs: [pre-tests]
uses: kornia/workflows/.github/workflows/docs.yml@v1.12.1
with:
python-version: '["3.11"]'
cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
collector:
needs: [coverage, tests-cpu, tutorials, typing, docs]
Expand All @@ -60,6 +107,7 @@ jobs:

tests-nightly:
if: contains(github.event.pull_request.labels.*.name, 'nightly')
needs: [pre-tests]
name: ${{ matrix.os }} - torch-nightly, ${{ matrix.pytorch-dtype }}

strategy:
Expand All @@ -68,8 +116,13 @@ jobs:
os: ['Ubuntu-latest', 'Windows-latest'] #, 'MacOS-latest'] add it when https://github.com/pytorch/pytorch/pull/89262 be merged
pytorch-dtype: ['float32', 'float64']

uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
with:
os: ${{ matrix.os }}
pytorch-version: '["nightly"]'
pytorch-dtype: ${{ matrix.pytorch-dtype }}
cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
66 changes: 58 additions & 8 deletions .github/workflows/scheduled_test_cpu.yml
Original file line number Diff line number Diff line change
Expand Up @@ -11,57 +11,107 @@ concurrency:
cancel-in-progress: true

jobs:
pre-tests:
runs-on: ubuntu-latest
outputs:
hash: ${{ steps.hashid.outputs.weights-hash }}
steps:
- uses: kornia/workflows/.github/actions/env@v1.12.1
- uses: actions/cache@v4
id: cache-weights
with:
path: weights/
key: model-weights-${{ hashFiles('.github/download-models-weights.py') }}
enableCrossOsArchive: true

- name: Download models weights...
if: steps.cache-weights.outputs.cache-hit != 'true'
run: python .github/download-models-weights.py -t weights/

- name: write hashid
id: hashid
run: echo "weights-hash=${{ hashFiles('.github/download-models-weights.py') }}" >> "$GITHUB_OUTPUT"

tests-cpu-ubuntu:
needs: [pre-tests]
strategy:
fail-fast: false
matrix:
# os: ['Ubuntu-latest', 'Windows-latest', 'MacOS-latest']
pytorch-dtype: ['float32', 'float64']

uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
with:
os: 'Ubuntu-latest'
python-version: '["3.8", "3.9", "3.10", "3.11", "3.12"]'
pytorch-version: '["1.9.1", "1.10.2", "1.11.0", "1.12.1", "1.13.1", "2.0.1", "2.1.2", "2.2.2", "2.3.1", "2.4.0"]'
pytorch-dtype: ${{ matrix.pytorch-dtype }}
pytest-extra: '--runslow'


cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
tests-cpu-windows:
needs: [pre-tests]
strategy:
fail-fast: true
matrix:
pytorch-dtype: ['float32', 'float64']

uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
with:
os: 'Windows-latest'
python-version: '["3.12"]'
pytorch-version: '["1.9.1", "2.4.0"]'
pytorch-dtype: ${{ matrix.pytorch-dtype }}
cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
tests-cpu-mac:
needs: [pre-tests]
strategy:
fail-fast: true
matrix:
pytorch-dtype: ['float32', 'float64']

uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
with:
os: 'MacOS-latest'
pytorch-dtype: ${{ matrix.pytorch-dtype }}
cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
coverage:
uses: kornia/workflows/.github/workflows/coverage.yml@v1.11.0
needs: [pre-tests]
uses: kornia/workflows/.github/workflows/coverage.yml@v1.12.1
with:
cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
typing:
uses: kornia/workflows/.github/workflows/mypy.yml@v1.11.0
uses: kornia/workflows/.github/workflows/mypy.yml@v1.12.1

tutorials:
uses: kornia/workflows/.github/workflows/tutorials.yml@v1.12.1

docs:
uses: kornia/workflows/.github/workflows/docs.yml@v1.11.0
needs: [pre-tests]
uses: kornia/workflows/.github/workflows/docs.yml@v1.12.1
with:
python-version: '["3.11"]'
cache-path: weights/
cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
cache-restore-keys: |
model-weights-${{ needs.pre-tests.outputs.hash }}
model-weights-
11 changes: 9 additions & 2 deletions conftest.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
import sys
from functools import partial
from itertools import product
Expand All @@ -18,6 +19,9 @@
_backends_non_experimental = []


WEIGHTS_CACHE_DIR = "weights/"


def get_test_devices() -> Dict[str, torch.device]:
"""Create a dictionary with the devices to test the source code. CUDA devices will be test only in case the
current hardware supports it.
Expand Down Expand Up @@ -189,7 +193,9 @@ def pytest_sessionstart(session):
ex
) and "Dynamo is not supported on Python 3.12+" not in str(ex):
raise ex
# TODO: cache all torch.load weights/states here to not impact on test suite

os.makedirs(WEIGHTS_CACHE_DIR, exist_ok=True)
torch.hub.set_dir(WEIGHTS_CACHE_DIR)


def _get_env_info() -> Dict[str, Dict[str, str]]:
Expand Down Expand Up @@ -249,7 +255,7 @@ def pytest_report_header(config):
import onnx

env_info = _get_env_info()

CACHED_WEIGTHS = os.listdir(WEIGHTS_CACHE_DIR)
if "cpu" in env_info:
desired_cpu_info = ["Model name", "Architecture", "CPU(s)", "Thread(s) per core", "CPU max MHz", "CPU min MHz"]
cpu_info = "cpu info:\n" + "\n".join(
Expand All @@ -276,6 +282,7 @@ def pytest_report_header(config):
- onnx-{onnx.__version__}
{gcc_info}
available optimizers: {TEST_OPTIMIZER_BACKEND}
model weights cached: {CACHED_WEIGTHS}
"""


Expand Down

0 comments on commit 93114bf

Please sign in to comment.