Will using Torch improve optimization speed?

[Yekai97]

I tried to optimize using torch, but I found that using CUDA takes longer than using the CPU. Is there something wrong with my code, or is CUDA unable to provide acceleration? Here is my code.

import torch
import time
import matplotlib.pyplot as plt
import optiland.backend as be
from optiland import optic, optimization
from optiland.ml import OpticalSystemModule

be.set_backend("torch") # Set the backend to PyTorch
be.grad_mode.enable() # Enable gradient tracking
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

lens = optic.Optic()

add surfaces

lens.add_surface(index=0, radius=be.inf, thickness=be.inf)
lens.add_surface(index=1, radius=50, thickness=3, material="SK16")
lens.add_surface(index=2, radius=-100, thickness=5)
lens.add_surface(index=3, radius=-100, thickness=3, material=("F2", "schott"))
lens.add_surface(index=4, radius=50, thickness=5, is_stop=True)
lens.add_surface(index=5, radius=50, thickness=3, material="SK16")
lens.add_surface(index=6, radius=-20, thickness=40)
lens.add_surface(index=7)

set aperture

lens.set_aperture(aperture_type="EPD", value=10)

set fields

lens.set_field_type(field_type="angle")
lens.add_field(y=0)

set wavelengths

lens.add_wavelength(value=0.48)
lens.add_wavelength(value=0.55, is_primary=True)
lens.add_wavelength(value=0.65)

lens.draw()

start = time.time()
problem = optimization.OptimizationProblem()
for wave in lens.wavelengths.get_wavelengths():
input_data = {
"optic": lens,
"Hx": 0,
"Hy": 0,
"num_rays": 3,
"wavelength": wave,
"distribution": "gaussian_quad",
}
problem.add_operand(
operand_type="OPD_difference",
target=0,
weight=1,
input_data=input_data,
)

for surface_number in range(1, 7):
problem.add_variable(lens, "radius", surface_number=surface_number)

problem.info()

model = OpticalSystemModule(lens, problem)

optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

Store loss values for plotting

losses = []

Run the optimization for 250 steps

for step in range(100):
optimizer.zero_grad() # Reset gradients from the previous step
loss = model() # Forward pass: calculate the loss (merit function value)
loss.backward(retain_graph=True) # Backward pass: compute gradients
optimizer.step() # Update lens radii based on gradients
model.apply_bounds() # (Optional) Apply any defined parameter constraints
losses.append(loss.item()) # Record the loss for this step

Display optimization problem information

problem.info()

print(f"Optimization completed in {time.time() - start:.2f} seconds with {device}")

[HarrisonKramer]

This kind of question has come up many times in different forms, and there isn’t a single right answer. Whether GPU or CPU is faster depends entirely on the problem size, data movement, precision, and how the operations are structured. Small-scale optimizations often don’t benefit from CUDA because the overhead outweighs the gain.

Without clear context (hardware specs, profiling data, etc.), it’s not really possible to say what’s “wrong” here. In general, it’s best to benchmark both approaches for your specific setup and use whichever performs better.