diff --git a/.github/workflows/test_cupy.yml b/.github/workflows/test_cupy.yml new file mode 100644 index 000000000..8c73775ab --- /dev/null +++ b/.github/workflows/test_cupy.yml @@ -0,0 +1,36 @@ +name: test_cupy +on: + pull_request: + branches: + - main + push: + branches: + - main +jobs: + test_cupy: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - uses: actions/setup-python@v2 + with: + python-version: '3.9' + - uses: Jimver/cuda-toolkit@v0.2.4 + id: cuda-toolkit + with: + cuda: '11.2.2' + - run: + echo "Installed cuda version is ${{steps.cuda-toolkit.outputs.cuda}}" + - run: + echo "Cuda install location is ${{steps.cuda-toolkit.outputs + .CUDA_PATH}}" + - run: nvcc --version + - run: sudo apt-get update + - run: sudo apt-get install --fix-missing python3-soundfile + + - run: pip install pyre-check pytest cupy-cuda112 + + - run: pip install -e .[av] + - run: pip freeze + - run: pyre --source-directory "." --noninteractive check || true + - run: pytest --durations=10 augly/tests/audio_tests/ + diff --git a/augly/audio/README.md b/augly/audio/README.md index 502cd9b9d..3ab3cde6a 100644 --- a/augly/audio/README.md +++ b/augly/audio/README.md @@ -9,6 +9,24 @@ pip install augly[av] This ensures that not only the base dependencies, but also the heavier dependencies required for audio & video processing, are installed. + +If you have [CUDA](https://developer.nvidia.com/cuda-gpus), you can also +install AugLy with builtin GPU-accelerated audio augmentation support. + +Check if your system has CUDA by running: +```bash +nvcc --version +``` + +Based on the compiler version, you can manually install `cupy`: +```bash +pip install cupy-cuda +``` +Note: Remove any periods from the version number. + +AugLy will now automatically detect if `cupy` is available and use it for +GPU-accelerated audio augmentation. + ## Augmentations [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/facebookresearch/AugLy/blob/main/examples/AugLy_audio.ipynb) diff --git a/augly/audio/functional.py b/augly/audio/functional.py index 656d5da4a..4920c24d1 100644 --- a/augly/audio/functional.py +++ b/augly/audio/functional.py @@ -4,6 +4,7 @@ import math from copy import deepcopy from typing import Any, Callable, Dict, List, Optional, Tuple, Union +import importlib import augly.audio.utils as audutils import numpy as np @@ -11,6 +12,10 @@ from augly.utils import DEFAULT_SAMPLE_RATE from augly.utils.libsndfile import install_libsndfile +cupy_spec = importlib.util.find_spec("cupy") +cupy_found = cupy_spec is not None +if cupy_found: + import cupy as cp install_libsndfile() import librosa @@ -645,10 +650,14 @@ def loop( """ assert isinstance(n, int) and n >= 0, "Expected 'n' to be a nonnegative integer" audio, sample_rate = audutils.validate_and_load_audio(audio, sample_rate) + if cupy_found: + audio = cp.array(audio) aug_audio = audio for _ in range(n): - aug_audio = np.append(aug_audio, audio, axis=(0 if audio.ndim == 1 else 1)) + aug_audio = (cp if cupy_found else np).append( + aug_audio, audio, axis=(0 if audio.ndim == 1 else 1) + ) audutils.get_metadata( metadata=metadata,