Repository navigation
[Performance] Accelerate PerceptualLoss with torch.compile #9175
Description
Is your feature request related to a problem? Please describe.
The VGG-based PerceptualLoss had a large positive effect on the image quality of my trained models, but it decreased my training speed quite a lot.
Describe the solution you'd like
My testing showed that we can speed up each iteration after the first one by 30% by adding self.perceptual_function = torch.compile(self.perceptual_function) to the class init
My testing even shows that it speeds things up if you calculate the loss on the CPU, but I am not sure that is correct.
Describe alternatives you've considered
I have already tuned the fake_3d_ratio, which is also a nice feature!
I can clean up the code on my end and open a PR if you would be open to it.
Additional context
Volume size: (2,1, 128,128,128)
network_type: 'vgg'
Iterations: 50
| GPU | total time | first iteration | subsequent iterations |
|---|---|---|---|
| RTX 3090 (not compiled) | 21 seconds | 2.6 seconds | 0.387 seconds |
| RTX 3090 (compiled) | 20 seconds | 6 seconds | 0.292 seconds |
| RTX PRO 5000 (not compiled) | 11.4 seconds | 0.49 seconds | 0.22 seconds |
| RTX PRO 5000 (compiled) | 8.6 seconds | 1.0 seconds | 0.15 seconds |
Script:
from __future__ import annotations
import logging
import torch
import pytest
from monai.losses import PerceptualLoss
import time
@pytest.mark.parametrize("use_compile", [False, True])
@pytest.mark.parametrize("device", ["cpu", "cuda"])
@pytest.mark.parametrize("input_shape", [(2,1, 128,128,128)])
def test_perceptual_loss_speed(use_compile, device, input_shape):
# Create random input tensors
input_tensor = torch.rand(input_shape).to(device)
target_tensor = torch.rand(input_shape).to(device)
# Initialize the PerceptualLoss with the specified use_compile option
loss_fn = PerceptualLoss(spatial_dims=len(input_shape) - 2, network_type="vgg", use_compile=use_compile).to(device)
if device == "cuda":
torch.cuda.synchronize() # Ensure all CUDA operations are complete before timing
r = 50
else:
r = 1
start_time = time.time()
for _ in range(r):
iteration_start = time.time()
loss_value = loss_fn(input_tensor, target_tensor)
iteration_end = time.time()
logging.warning(f"Iteration time {iteration_end - iteration_start:.6f} seconds")
end_time = time.time()
logging.warning(f"PerceptualLoss forward pass time (use_compile={use_compile}, device={device}, input_shape={input_shape}): {end_time - start_time:.6f} seconds")
Activity
Willy1118 commentedon Oct 26, 2026
Hi @StijnvWijn and maintainers,
I'm interested in working on this. I've set up a dev environment and can reproduce the setup (torch.compile works on my machine), so I'd like to check a few things before starting:
@StijnvWijn, do you still plan to open the PR yourself? If so, I'm happy to help with testing or benchmarking instead. If not, I'd like to take this on.
Maintainers, is this a feature MONAI would accept? In particular, would an opt-in
use_compile: bool = Falseargument onPerceptualLoss(off by default, so existing behavior doesn't change) be an acceptable interface?Since the reported numbers show the first iteration gets slower when compiled, I plan to measure the break-even point (how many iterations before compiling pays off) for a few input sizes, and to fall back to the uncompiled function with a warning if
torch.compilefails or is unavailable on the platform.Please let me know if you have any preferences on the design before I start. Thanks!
StijnvWijn commentedon Oct 27, 2026
Hi @Willy1118
I do still plan to open a PR if the maintainers agree. What kind of hardware do you have access to for testing and benchmarking?
Your proposal for the opt-in argument is exactly what I have implemented as well. However, I think that letting the loss error out if compilation fails is a better solution, as it is most likely experimental training code and you don't want to waste GPU-hours on a training with a slow loss.
Let's wait on what the maintainers prefer and then we can start work on the implementation.
Hoyant-Su commentedon Nov 1, 2026
I ran the proposed VGG compilation change on one H100 80 GB, including the backward pass used during training.
The inputs were generated tensors with shape
(2, 1, D, D, D),fake_3d_ratio=0.5, and values in[-1, 1]. The environment was PyTorch 2.6.0+cu124 and torchvision 0.21.0+cu124, using float32 with TF32 disabled and this MONAI source revision. Each measurement used five warmup steps and thirty timed steps, with CUDA synchronization around forward and backward. Compilation was applied toloss.perceptual_function.The break-even estimate divides the extra first-step cost by the median time saved per subsequent step. These measurements favor longer training runs, particularly for the larger volume.
The largest loss-value difference was 1.79e-7. Across three matched slice seeds per volume, the relative L2 difference in input gradients ranged from 0.048% to 0.173%, with maximum absolute differences up to 2.87e-6. Repeated eager runs differed by at most 5.10e-11 per gradient element. The CPU random-generator state after each compared step matched between eager and compiled execution.
The benchmark script and analysis were prepared with OpenAI Codex.
Reproduction script
Install the linked MONAI revision with
lpips==0.1.4in the stated PyTorch environment. Save this script asbenchmark.py, then run it in a fresh process for each volume size:Use an empty TorchInductor and Triton cache for each run when comparing first-step compilation costs. Pretrained VGG weights can be reused.