loop efficiency w/ cupy

.github/workflows/test_cupy.yml

@@ -0,0 +1,32 @@
+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 ffmpeg python3-soundfile
+      - run: pip install pyre-check pytest torchvision
+      - run: pip install -e .[av]
+      - run: pyre --source-directory "." --noninteractive check || true
+      - run: pytest --durations=10 .

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<version>
+```
+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)

augly/audio/functional.py

@@ -4,13 +4,18 @@
 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
 import torch
 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,