GPU-accelerated motion magnification via pytorch_wavelets + cuFFT

[joeljose]

Problem Statement

DTCWT motion magnification on CPU takes ~2 minutes for a 301-frame 528×592 video (3 channels). Profiling revealed the temporal filter (flat-top convolution) is the dominant bottleneck at 54% of runtime, followed by forward/inverse DTCWT at 26%.

GPU acceleration via pytorch_wavelets (DTCWTForward/DTCWTInverse) + cuFFT (torch.fft) can deliver ~17x speedup per channel based on prototype testing on an RTX 4050 (6 GB VRAM).

Who: Users processing motion magnification videos, especially at higher resolutions or longer durations.

Why now: CPU implementation is stable (v1.1.0). Visual-Mic project validates the GPU DTCWT pattern. Prototype experiments confirm 16.8x speedup is achievable.

User Stories

1. As a user, I want to run --gpu to process videos faster using my NVIDIA GPU.
2. As a user, I want to control GPU memory with --batch-size so I can tune for my GPU's VRAM.
3. As a user, I want a Dockerfile.gpu so I can run GPU processing without installing CUDA/PyTorch on my host.
4. As a user, I want clear error messages when my GPU runs out of memory, with suggestions to reduce batch size.

Research & Profiling Results

CPU Profiling (301 frames, 528×592, single channel)

Step	Original (v1.1.0)	With FFT filter
Temporal filter	42.3s (57%)	9.9s (27%)
Forward DTCWT	11.7s (16%)	7.9s (22%)
Inverse DTCWT	8.2s (11%)	8.2s (22%)
Coeff reconstruct	7.2s (10%)	6.4s (18%)
Phase extraction	4.9s (7%)	4.1s (11%)
Total	74.2s	36.4s (2.0x)

CPU FFT filter optimization (already implemented): Switches scipy.ndimage.convolve1d to scipy.signal.fftconvolve with manual reflect-padding and 10K chunking for cache efficiency. Delivers 4.3x filter speedup, 2.0x total. PSNR 65 dB vs original (visually identical).

GPU Prototype Results (30 frames, RTX 4050 6GB)

Step	CPU	GPU	Speedup
Forward DTCWT	1.1s	0.080s	14x
Phase extraction	0.36s	0.021s	17x
Temporal filter	4.0s	0.207s	19x
Coeff reconstruct	0.7s	0.012s	58x
Inverse DTCWT	0.8s	0.086s	9x
Total	6.81s	0.41s	16.8x

pytorch_wavelets Findings

• DTCWTForward + DTCWTInverse roundtrip: PSNR 131 dB (near-perfect reconstruction)
• Float32 only (CPU dtcwt uses float64) — outputs are self-consistent but not cross-comparable with CPU path
• Yh tensor shape: (N, 1, 6, H, W, 2) where last dim is [real, imag]
• Library is unmaintained (~2020 last commit) but stable and only option for GPU DTCWT in PyTorch
• Filters: biort='near_sym_b', qshift='qshift_b' (same Kingsbury filters as CPU path)

VRAM Measurements

Batch size	Forward+Inverse VRAM	Full pipeline VRAM
5 frames	739 MB	~1.5 GB
10 frames	814 MB	~2.5 GB
15 frames	890 MB	~3.5 GB
30 frames	1,124 MB	~5.0 GB

Proposed Solution: Batched GPU Pipeline

Why Batching is Needed

The motion magnification algorithm requires all frames' phase data simultaneously for temporal filtering (flat-top convolution operates along the time axis across all frames). For a 301-frame video at 528×592 with 8 DTCWT levels, the phase arrays total ~1.4 GB in float32. Combined with DTCWT coefficient storage and FFT buffers, the full pipeline for 301 frames would need ~15+ GB VRAM — far exceeding typical consumer GPUs (6-8 GB).

Batched Pipeline Architecture

The solution splits the pipeline into two GPU-accelerated passes with CPU-resident phase storage:

┌─────────────────────────────────────────────────────────────┐
│  PASS 1: Forward DTCWT + Phase/Amplitude Extraction         │
│  (GPU batched, results stored on CPU)                       │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  For each batch of B frames:                                │
│    1. Transfer B frames to GPU as (B, 1, H, W) float32     │
│    2. DTCWTForward → Yl, Yh                                │
│    3. For each level:                                       │
│       - Extract amplitude: sqrt(real² + imag²)              │
│       - Normalize: real/mag, imag/mag                       │
│       - Compute frame-to-frame phase via conjugate multiply │
│       - Transfer phase angles + amplitudes to CPU (float32) │
│    4. Free GPU memory: del batch_tensor, Yl, Yh            │
│                                                             │
│  CPU storage after Pass 1:                                  │
│    - phase_angles[level]: (num_frames, H_l, W_l, 6) f32    │
│    - amplitudes[level]:   (num_frames, H_l, W_l, 6) f32    │
│    - cumsum phase angles to get cumulative phase            │
│                                                             │
├─────────────────────────────────────────────────────────────┤
│  TEMPORAL FILTERING + PHASE MODIFICATION                    │
│  (GPU per-level, or CPU if VRAM insufficient)               │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  For each level (largest first, bail to CPU if OOM):        │
│    1. Transfer phase array to GPU                           │
│    2. cuFFT forward (torch.fft.rfft along time axis)        │
│    3. Multiply by FFT of flat-top window                    │
│    4. cuFFT inverse → phase0 (base motion)                  │
│    5. detail = phase - phase0                               │
│    6. phase = phase0 + detail * magnification               │
│    7. Repeat steps 2-4 with width=2 window (smoothing)      │
│    8. Transfer modified phase back to CPU                   │
│    9. Free GPU memory                                       │
│                                                             │
│  Fallback: If a level's phase array doesn't fit in VRAM,    │
│  use CPU FFT filter (already optimized, ~2.3x slower)       │
│                                                             │
├─────────────────────────────────────────────────────────────┤
│  PASS 2: Coefficient Reconstruction + Inverse DTCWT         │
│  (GPU batched)                                              │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  For each batch of B frames:                                │
│    1. Reconstruct Yh from amplitude + modified phase:       │
│       - real = amplitude * cos(phase)                       │
│       - imag = amplitude * sin(phase)                       │
│       - Stack to (B, 1, 6, H, W, 2) per level              │
│    2. Transfer Yl (original lowpass) + modified Yh to GPU   │
│    3. DTCWTInverse → reconstructed frames                   │
│    4. Transfer result to CPU, free GPU memory               │
│                                                             │
│  Output: (num_frames, H, W) float32 per channel            │
│                                                             │
└─────────────────────────────────────────────────────────────┘

Memory Budget

For 301 frames at 528×592 with nlevels=8:

CPU RAM (same as current, but float32 instead of float64):

• Phase arrays: ~718 MB (half of current float64)
• Amplitude arrays: ~718 MB (new — not stored in current impl)
• Total: ~1.4 GB per channel (vs ~2.9 GB current)

GPU VRAM (peak, controlled by batch_size):

• Pass 1: batch_size × 13 MB (input + DTCWT coefficients)
• Temporal filter: phase_array_for_one_level (L0 = 358 MB, L1 = 89 MB, ...)
• Pass 2: batch_size × 16 MB (coefficients + reconstruction + output)
• Overhead: ~200 MB (PyTorch, filter weights, FFT buffers)

Default batch_size=10: ~350 MB for passes, ~560 MB for L0 filter = ~1 GB peak VRAM.

Key Design Decisions

1. CPU and GPU produce different outputs — accepted. pytorch_wavelets and dtcwt use different filter implementations. Both are valid DTCWT implementations; outputs are self-consistent but not cross-comparable.

2. Float32 on GPU — pytorch_wavelets only supports float32. Phase arrays are extracted as float32. Temporal filtering on GPU uses float32. This is sufficient for motion magnification (validated by Visual-Mic project).

3. Store Yl (lowpass) on CPU between passes — needed for inverse DTCWT. Small: ~0.1 MB per frame at level 8.

4. Graceful OOM fallback — if GPU temporal filter OOMs on a large level, fall back to CPU FFT filter for that level only. Log which levels ran on GPU vs CPU.

Scope

In (v2.0.0)

• motion_mag_gpu.py — new GPU implementation file
• CLI flags: --gpu, --batch-size (default 10), --device (default 0)
• Dockerfile.gpu based on pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
• docker-build-gpu.sh build script
• requirements-gpu.txt
• GPU unit tests (tests/test_motion_mag_gpu.py)
• VRAM estimation + graceful OOM handling with actionable error messages
• CPU FFT filter optimization (already implemented in motion_mag.py)
• CI job for GPU (lint-only, no GPU runners)

Deferred

• Multi-GPU support
• Mixed-precision (float16) experiments
• Streaming temporal filter for videos that exceed CPU RAM
• GPU-accelerated video I/O (nvdec/nvenc)

Testing Plan

Automated tests (tests/test_motion_mag_gpu.py):

• Tier 1 (strict): VRAM estimation arithmetic (pure math, exact equality)
• Tier 2 (moderate):
  • GPU DTCWT forward/inverse roundtrip (PSNR > 120 dB on synthetic data)
  • GPU phase extraction produces finite values with correct shapes
  • GPU temporal filter output shape and finiteness
  • Batched processing produces same result as single-batch (verify batch boundaries don't introduce artifacts)
• Tier 3 (smoke):
  • Full GPU pipeline on face.mp4 produces valid output (finite, correct dimensions, reasonable value range)
  • OOM fallback triggers and completes successfully with tiny VRAM limit

Hard to test (manual):

• Visual quality comparison of GPU vs CPU output videos
• Actual GPU hardware — CI has no GPU; tests skip with pytest.mark.skipif(not torch.cuda.is_available())
• VRAM edge cases on different GPU models

Dependencies to stub:

• torch and pytorch_wavelets imports guarded with try/except
• GPU tests import-guarded so CPU test suite passes without PyTorch

Open Questions

1. What should --batch-size default be? 10 works for 6 GB GPUs, but 8 GB+ GPUs could use 15-20 for better throughput.
2. Should we auto-detect optimal batch size from available VRAM, or keep it manual?

[joeljose]

Completed in PR #18, released in v2.0.0