From becc9fe78e253ea1631641613cd6ed6a21585f25 Mon Sep 17 00:00:00 2001 From: Bharadwaj Date: Sun, 20 Sep 2026 19:03:49 +0200 Subject: [PATCH 1/5] Fix device handling in build_fdn and model_response --- src/pyFDN/auxiliary/flamo.py | 10 +++++++--- src/pyFDN/train/build.py | 2 +- src/pyFDN/train/response.py | 2 +- src/pyFDN/translate/dss_to_flamo.py | 2 +- tests/test_train.py | 6 ++++-- 5 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/pyFDN/auxiliary/flamo.py b/src/pyFDN/auxiliary/flamo.py index f5febc26..722305bf 100644 --- a/src/pyFDN/auxiliary/flamo.py +++ b/src/pyFDN/auxiliary/flamo.py @@ -715,7 +715,7 @@ def core_alias_decay_db(core: Any) -> float: return abs(float(value)) -def wrap_fdn_shell(core: Any, *, nfft: int, dtype: Any = None) -> Any: +def wrap_fdn_shell(core: Any, *, nfft: int, dtype: Any = None, device: Any = None) -> Any: r""" Wrap an FDN core in a FLAMO ``Shell`` that returns the impulse response. @@ -741,6 +741,8 @@ def wrap_fdn_shell(core: Any, *, nfft: int, dtype: Any = None) -> Any: FFT size. dtype : torch.dtype or None Dtype for the FFT/iFFT layers; defaults to float32. + device : torch.device or None + Device for the FFT/iFFT layers; defaults to cuda if available else cpu. Returns ------- @@ -757,10 +759,12 @@ def wrap_fdn_shell(core: Any, *, nfft: int, dtype: Any = None) -> Any: from flamo.processor import dsp, system torch_dtype = torch.float32 if dtype is None else dtype - return system.Shell( + dev = _get_device(device) + shell = system.Shell( core=core, input_layer=dsp.FFT(nfft, dtype=torch_dtype), output_layer=dsp.iFFTAntiAlias( - nfft, alias_decay_db=core_alias_decay_db(core), dtype=torch_dtype + nfft, alias_decay_db=core_alias_decay_db(core), device=dev, dtype=torch_dtype ), ) + return shell.to(dev) diff --git a/src/pyFDN/train/build.py b/src/pyFDN/train/build.py index b3e594a9..4e100961 100644 --- a/src/pyFDN/train/build.py +++ b/src/pyFDN/train/build.py @@ -389,7 +389,7 @@ def _hook(module: Any, baked: np.ndarray | None) -> Any: post_matrix=_hook(post_matrix, build.post_matrix), post_output=_hook(post_output, build.post_output), ) - return wrap_fdn_shell(core, nfft=nfft, dtype=dtype) + return wrap_fdn_shell(core, nfft=nfft, dtype=dtype, device=device) def trainable_from_preset( diff --git a/src/pyFDN/train/response.py b/src/pyFDN/train/response.py index a93e7730..cad29b2b 100644 --- a/src/pyFDN/train/response.py +++ b/src/pyFDN/train/response.py @@ -127,7 +127,7 @@ def model_response(model: Any, excitation: torch.Tensor | None = None) -> Respon n_in, n_out = int(model.input_channels), int(model.output_channels) if excitation is None: excitation = impulse_excitation( - n_in, int(model.nfft), dtype=getattr(model, "dtype", None) + n_in, int(model.nfft), device=model.device, dtype=getattr(model, "dtype", None) ) y = model(excitation) # (n_in, nfft, n_out) if y.shape[0] != n_in or y.shape[2] != n_out: diff --git a/src/pyFDN/translate/dss_to_flamo.py b/src/pyFDN/translate/dss_to_flamo.py index 19532608..2ae69e7b 100644 --- a/src/pyFDN/translate/dss_to_flamo.py +++ b/src/pyFDN/translate/dss_to_flamo.py @@ -148,7 +148,7 @@ def dss_to_flamo( ) if shell: - return wrap_fdn_shell(core, nfft=nfft, dtype=dtype) + return wrap_fdn_shell(core, nfft=nfft, dtype=dtype, device=device) return core diff --git a/tests/test_train.py b/tests/test_train.py index 14747e77..c3969f29 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -1341,6 +1341,7 @@ def test_shelf_decay_pulled_below_zero_stays_at_the_floor(): design="first_order_shelf", nfft=2**10, dtype=torch.float64, + device="cpu", ) sos = module.map(module.param).detach().numpy() assert np.all(np.isfinite(sos)) @@ -1356,6 +1357,7 @@ def test_shelf_decay_pulled_below_zero_stays_at_the_floor(): design="first_order_shelf", nfft=2**10, dtype=torch.float64, + device="cpu", ) np.testing.assert_allclose( sos, deeper.map(deeper.param).detach().numpy(), rtol=1e-6 @@ -1421,10 +1423,10 @@ def test_per_line_rt_floor_is_each_line_s_own_round_trip(): fs, nfft = 48000.0, 2**10 delays = np.array([809.0, 4096.0]) shared = pyFDN.AttenuationFilter( - np.full(10, 1.0), delays, fs, nfft=nfft, dtype=torch.float64 + np.full(10, 1.0), delays, fs, nfft=nfft, dtype=torch.float64, device="cpu" ) per_line = pyFDN.AttenuationFilter( - np.full((10, 2), 1.0), delays, fs, nfft=nfft, dtype=torch.float64 + np.full((10, 2), 1.0), delays, fs, nfft=nfft, dtype=torch.float64, device="cpu" ) assert shared.rt_floor.ndim == 0 np.testing.assert_allclose(float(shared.rt_floor), 4096.0 / fs) From 567bb3001ef7f6928c3fc43e6d6c7c9327f12ff5 Mon Sep 17 00:00:00 2001 From: Bharadwaj Date: Sun, 20 Sep 2026 21:09:52 +0200 Subject: [PATCH 2/5] fix tests --- .../allpass_FDN/example_allpass_FDN_in_FDN.py | 19 ++++++++++--------- .../example_allpass_FDN_schroeder_in_loop.py | 4 +++- examples/example_coupled_rooms.py | 1 + examples/example_delay_matrix_density.py | 1 + examples/example_fdn_to_faust.py | 1 + examples/example_multislope_rir_to_fdn.py | 6 ++++-- examples/example_paraunitary_fdn.py | 1 + examples/example_pole_boundaries.py | 1 + examples/example_process_fdn_vs_flamo.py | 1 + examples/example_rir_to_fdn.py | 2 ++ examples/example_sdn.py | 2 +- examples/example_train_fdn_to_rir.py | 2 -- examples/example_vanilla_FDN.py | 1 + src/pyFDN/auxiliary/coupled_rooms.py | 6 +++--- tests/test_auxiliary_flamo.py | 2 +- tests/test_flamo_to_pr.py | 2 +- 16 files changed, 32 insertions(+), 20 deletions(-) diff --git a/examples/allpass_FDN/example_allpass_FDN_in_FDN.py b/examples/allpass_FDN/example_allpass_FDN_in_FDN.py index b2aefe9f..ca14e608 100644 --- a/examples/allpass_FDN/example_allpass_FDN_in_FDN.py +++ b/examples/allpass_FDN/example_allpass_FDN_in_FDN.py @@ -80,7 +80,8 @@ def _(fs, N, nfft, np, pyFDN): # FLAMO core (N→N), no Shell — for use inside the recursion allpass_fdn_core = pyFDN.dss_to_flamo( - A_sch, B_sch, C_sch, D_sch, delays_sch, fs, nfft=nfft, shell=False + A_sch, B_sch, C_sch, D_sch, delays_sch, fs, nfft=nfft, shell=False, + device="cpu", ) return (allpass_fdn_core,) @@ -104,15 +105,15 @@ def _(fs, N, delay_module, nfft, np, pyFDN, sos_filter_module): ) output_delay_sec = np.linspace(0.01, 0.02, N) + np.random.uniform(0, 0.001, size=N) - main_delays = delay_module(main_delay_sec, nfft, fs=fs) - input_delays = delay_module(input_delay_sec, nfft, fs=fs) - output_delays = delay_module(output_delay_sec, nfft, fs=fs) + main_delays = delay_module(main_delay_sec, nfft, fs=fs, device="cpu") + input_delays = delay_module(input_delay_sec, nfft, fs=fs, device="cpu") + output_delays = delay_module(output_delay_sec, nfft, fs=fs, device="cpu") # Attenuation: first-order absorption, canonical (1, 6, N) SOS bank. main_delay_smp = np.round(main_delay_sec * fs).astype(float) rt_dc, rt_ny = 1.4, 0.3 sos = pyFDN.decay_to_first_order_shelf(rt_dc, rt_ny, None, main_delay_smp, fs=fs) - attenuation = sos_filter_module(sos, nfft) + attenuation = sos_filter_module(sos, nfft, device="cpu") return attenuation, input_delays, main_delays, output_delays @@ -147,8 +148,8 @@ def _( C_out = np.random.randn(2, N) C_out = C_out / np.linalg.norm(C_out, axis=1, keepdims=True) - gain_B_in = gain_module(B_in, nfft) - gain_C_out = gain_module(C_out, nfft) + gain_B_in = gain_module(B_in, nfft, device="cpu") + gain_C_out = gain_module(C_out, nfft, device="cpu") # Recursion: fF = allpass FDN → attenuation, fB = main delays feedforward = system.Series( @@ -176,8 +177,8 @@ def _( model = system.Shell( core=core_chain, - input_layer=dsp.FFT(nfft), - output_layer=dsp.iFFT(nfft), + input_layer=dsp.FFT(nfft).to("cpu"), + output_layer=dsp.iFFT(nfft).to("cpu"), ) ir_stereo = pyFDN.flamo_time_response(model).squeeze() diff --git a/examples/allpass_FDN/example_allpass_FDN_schroeder_in_loop.py b/examples/allpass_FDN/example_allpass_FDN_schroeder_in_loop.py index cdd9843c..3477412a 100644 --- a/examples/allpass_FDN/example_allpass_FDN_schroeder_in_loop.py +++ b/examples/allpass_FDN/example_allpass_FDN_schroeder_in_loop.py @@ -150,7 +150,7 @@ def _(fs, N, nfft, pyFDN): rng=6, ) - model_fdn = pyFDN.build_to_flamo(fdn_build, nfft=nfft) + model_fdn = pyFDN.build_to_flamo(fdn_build, nfft=nfft, device="cpu") ir_fdn = pyFDN.flamo_time_response(model_fdn).squeeze() return fdn_build, ir_fdn @@ -199,11 +199,13 @@ def _( fs, nfft=nfft, shell=False, + device="cpu", ) model_fdn_allpass = pyFDN.build_to_flamo( fdn_build, nfft=nfft, post_delay=schroeder_core, + device="cpu", ) ir_fdn_allpass = pyFDN.flamo_time_response(model_fdn_allpass).squeeze() diff --git a/examples/example_coupled_rooms.py b/examples/example_coupled_rooms.py index 703c9e45..f3e88d1b 100644 --- a/examples/example_coupled_rooms.py +++ b/examples/example_coupled_rooms.py @@ -162,6 +162,7 @@ def _(A, B, C, D, attenuation_sos, delays, fs, nfft, post_eq_sos, pyFDN): nfft=nfft, post_delay=attenuation_sos, post_output=post_eq_sos, + device="cpu", ) return (model,) diff --git a/examples/example_delay_matrix_density.py b/examples/example_delay_matrix_density.py index 8769125c..f925eb48 100644 --- a/examples/example_delay_matrix_density.py +++ b/examples/example_delay_matrix_density.py @@ -107,6 +107,7 @@ def _(fs, gain_per_sample, nfft, np, pyFDN, total_delay): nfft=nfft, post_delay=build.post_delay, post_output=build.post_output, + device="cpu", ) ir_vanilla = pyFDN.flamo_time_response(model).flatten() pyFDN.plot_flamo_graph(model) diff --git a/examples/example_fdn_to_faust.py b/examples/example_fdn_to_faust.py index e56cc879..25e926f2 100644 --- a/examples/example_fdn_to_faust.py +++ b/examples/example_fdn_to_faust.py @@ -123,6 +123,7 @@ def _(pyFDN, torch): nfft=2**18, post_delay=build.post_delay, post_output=build.post_output, + device="cpu", ) # (1 input, n_samples, 2 outputs) -> (n_samples, 2), the shape pyFDN plots take. ir_flamo = pyFDN.flamo_time_response(model)[0] diff --git a/examples/example_multislope_rir_to_fdn.py b/examples/example_multislope_rir_to_fdn.py index ce84fb90..9b24ad81 100644 --- a/examples/example_multislope_rir_to_fdn.py +++ b/examples/example_multislope_rir_to_fdn.py @@ -226,7 +226,8 @@ def geq_grid(band_values): nfft=nfft, post_delay=_absorption, shell=True, - ) + device="cpu", + ) ).squeeze()[: len(rir)] _rt_flat, _ = pyFDN.estimate_rt_bands(_ir_flat, fs) @@ -248,7 +249,8 @@ def geq_grid(band_values): post_delay=_absorption, post_output=_eq[:, :, np.newaxis], shell=True, - ) + device="cpu", + ) ).squeeze()[: len(rir)] slope_fdn_rt.append(_rt_flat) diff --git a/examples/example_paraunitary_fdn.py b/examples/example_paraunitary_fdn.py index 3cfebb4d..d868b0a7 100644 --- a/examples/example_paraunitary_fdn.py +++ b/examples/example_paraunitary_fdn.py @@ -152,6 +152,7 @@ def _( fs=fs, shell=False, dtype=torch.float64, + device="cpu", ) residues, poles, direct_term, is_pair, _meta = pyFDN.flamo_to_pr( model, diff --git a/examples/example_pole_boundaries.py b/examples/example_pole_boundaries.py index 49745f87..be06e20a 100644 --- a/examples/example_pole_boundaries.py +++ b/examples/example_pole_boundaries.py @@ -131,6 +131,7 @@ def _( shell=False, post_delay=sos_loop, dtype=torch.float64, + device="cpu", ) _residues, poles, _direct_term, _is_pair, _meta = pyFDN.flamo_to_pr( model, diff --git a/examples/example_process_fdn_vs_flamo.py b/examples/example_process_fdn_vs_flamo.py index 6a83c83b..49b082b4 100644 --- a/examples/example_process_fdn_vs_flamo.py +++ b/examples/example_process_fdn_vs_flamo.py @@ -137,6 +137,7 @@ def _( post_delay=sos_absorption, # canonical (n_sections, 6, N) bank shell=True, dtype=torch.float64, + device="cpu", ) ir_flamo = pyFDN.flamo_time_response(model).squeeze().astype(np.float64)[:ir_len] diff --git a/examples/example_rir_to_fdn.py b/examples/example_rir_to_fdn.py index 5ea76e4e..5f66f9eb 100644 --- a/examples/example_rir_to_fdn.py +++ b/examples/example_rir_to_fdn.py @@ -172,6 +172,7 @@ def _( nfft=nfft, post_delay=sos_absorption, shell=True, + device="cpu", ) ir_unequalized = pyFDN.flamo_time_response(_model).squeeze()[:rir_len] print(f"Unequalized FDN IR computed: {rir_len} samples") @@ -226,6 +227,7 @@ def _( post_delay=sos_absorption, post_output=equalization_sos[:, :, np.newaxis], shell=True, + device="cpu", ) ir_fdn = pyFDN.flamo_time_response(model_eq).squeeze()[:rir_len] diff --git a/examples/example_sdn.py b/examples/example_sdn.py index b62546fa..1d3f0d0d 100644 --- a/examples/example_sdn.py +++ b/examples/example_sdn.py @@ -120,7 +120,7 @@ def _(mo): @app.cell def _(sdn): nfft = 2**17 - model, sdn_result = sdn.sdn_to_flamo(nfft=nfft) + model, sdn_result = sdn.sdn_to_flamo(nfft=nfft, device="cpu") return model, sdn_result diff --git a/examples/example_train_fdn_to_rir.py b/examples/example_train_fdn_to_rir.py index 5d658281..4f3ceaad 100644 --- a/examples/example_train_fdn_to_rir.py +++ b/examples/example_train_fdn_to_rir.py @@ -265,7 +265,6 @@ def _(mo): @app.cell def _(device, fs, loss, model, pyFDN, render_nfft, rir_len, train_nfft): - model.set_nfft(train_nfft) log = pyFDN.train_fdn( model, loss, @@ -275,7 +274,6 @@ def _(device, fs, loss, model, pyFDN, render_nfft, rir_len, train_nfft): device=device, rng=0, ) - model.set_nfft(render_nfft) ir_trained = pyFDN.flamo_time_response(model, fs=fs).squeeze()[:rir_len] print( f"{log.steps_run} steps: loss " diff --git a/examples/example_vanilla_FDN.py b/examples/example_vanilla_FDN.py index aeda24b8..4350aa77 100644 --- a/examples/example_vanilla_FDN.py +++ b/examples/example_vanilla_FDN.py @@ -73,6 +73,7 @@ def _(pyFDN, torch): nfft=2**18, post_delay=build.post_delay, post_output=build.post_output, + device="cpu", ) ir = pyFDN.flamo_time_response(model).flatten() return build, fs, ir, model diff --git a/src/pyFDN/auxiliary/coupled_rooms.py b/src/pyFDN/auxiliary/coupled_rooms.py index c85a33d0..f52a35f0 100644 --- a/src/pyFDN/auxiliary/coupled_rooms.py +++ b/src/pyFDN/auxiliary/coupled_rooms.py @@ -174,14 +174,14 @@ def matrix_sqrt(A): ) ) - input_layer = dsp.FFT(nfft) - output_layer = dsp.iFFT(nfft) + input_layer = dsp.FFT(nfft).to(device) + output_layer = dsp.iFFT(nfft).to(device) model = system.Shell(core=fdn, input_layer=input_layer, output_layer=output_layer) # Generate impulse response with torch.no_grad(): - impulse = torch.zeros(1, nfft, 1) + impulse = torch.zeros(1, nfft, 1, device=device) impulse[0, 0, 0] = 1.0 ir = model(impulse).squeeze().cpu().numpy() diff --git a/tests/test_auxiliary_flamo.py b/tests/test_auxiliary_flamo.py index 9d482299..82dcd91a 100644 --- a/tests/test_auxiliary_flamo.py +++ b/tests/test_auxiliary_flamo.py @@ -121,6 +121,6 @@ def test_wrap_fdn_shell_returns_the_time_response(): impulse = torch.zeros(1, nfft, 1) impulse[:, 0, :] = 1.0 - time = wrap_fdn_shell(core, nfft=nfft)(impulse) + time = wrap_fdn_shell(core, nfft=nfft, device="cpu")(impulse) assert time.shape == (1, nfft, 1) assert not time.is_complex() diff --git a/tests/test_flamo_to_pr.py b/tests/test_flamo_to_pr.py index 682a9ebb..25f85fe5 100644 --- a/tests/test_flamo_to_pr.py +++ b/tests/test_flamo_to_pr.py @@ -127,7 +127,7 @@ def test_flamo_to_pr_biquad_in_loop_reconstructs_ir(): # Modal reconstruction matches FLAMO's true impulse response (the only IR # reference once there is an IIR filter in the loop). ir_flamo = np.asarray( - model.get_time_response(fs=int(fs)).squeeze(), dtype=np.float64 + model.get_time_response(fs=int(fs)).squeeze().cpu(), dtype=np.float64 ) if ir_flamo.ndim == 3: ir_flamo = ir_flamo[:, 0, 0] From 38cd49aec777472442fd0aefd0278d71b11a790f Mon Sep 17 00:00:00 2001 From: Bharadwaj Date: Sun, 20 Sep 2026 22:52:28 +0200 Subject: [PATCH 3/5] fix: Change device from string to torch object --- .../allpass_FDN/example_allpass_FDN_in_FDN.py | 26 +++++++++++-------- src/pyFDN/auxiliary/coupled_rooms.py | 2 +- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/examples/allpass_FDN/example_allpass_FDN_in_FDN.py b/examples/allpass_FDN/example_allpass_FDN_in_FDN.py index ca14e608..7dad5376 100644 --- a/examples/allpass_FDN/example_allpass_FDN_in_FDN.py +++ b/examples/allpass_FDN/example_allpass_FDN_in_FDN.py @@ -34,6 +34,7 @@ def _(): from collections import OrderedDict import numpy as np + import torch from flamo.processor import dsp, system import pyFDN @@ -43,7 +44,9 @@ def _(): fs = 48000 nfft = 2**17 N = 4 # number of delay lines (1→N input gain, N→2 output gain) + device = torch.device("cpu") return ( + device, fs, N, OrderedDict, @@ -69,7 +72,7 @@ def _(mo): @app.cell -def _(fs, N, nfft, np, pyFDN): +def _(device, fs, N, nfft, np, pyFDN): # Homogeneous MIMO allpass FDN: delays, G = diag(g^delays), A = G @ U, complete_fdn delays_sch = np.random.randint(51, 300, size=N) _g = pyFDN.rt_to_gain_per_sample(0.07, fs) @@ -81,7 +84,7 @@ def _(fs, N, nfft, np, pyFDN): # FLAMO core (N→N), no Shell — for use inside the recursion allpass_fdn_core = pyFDN.dss_to_flamo( A_sch, B_sch, C_sch, D_sch, delays_sch, fs, nfft=nfft, shell=False, - device="cpu", + device=device, ) return (allpass_fdn_core,) @@ -97,7 +100,7 @@ def _(mo): @app.cell -def _(fs, N, delay_module, nfft, np, pyFDN, sos_filter_module): +def _(device, fs, N, delay_module, nfft, np, pyFDN, sos_filter_module): # Main delays (feedback path), input and output delays — in seconds main_delay_sec = np.random.uniform(0.02, 0.04, size=N) input_delay_sec = np.linspace(0.01, 0.02 * N, N) + np.random.uniform( @@ -105,15 +108,15 @@ def _(fs, N, delay_module, nfft, np, pyFDN, sos_filter_module): ) output_delay_sec = np.linspace(0.01, 0.02, N) + np.random.uniform(0, 0.001, size=N) - main_delays = delay_module(main_delay_sec, nfft, fs=fs, device="cpu") - input_delays = delay_module(input_delay_sec, nfft, fs=fs, device="cpu") - output_delays = delay_module(output_delay_sec, nfft, fs=fs, device="cpu") + main_delays = delay_module(main_delay_sec, nfft, fs=fs, device=device) + input_delays = delay_module(input_delay_sec, nfft, fs=fs, device=device) + output_delays = delay_module(output_delay_sec, nfft, fs=fs, device=device) # Attenuation: first-order absorption, canonical (1, 6, N) SOS bank. main_delay_smp = np.round(main_delay_sec * fs).astype(float) rt_dc, rt_ny = 1.4, 0.3 sos = pyFDN.decay_to_first_order_shelf(rt_dc, rt_ny, None, main_delay_smp, fs=fs) - attenuation = sos_filter_module(sos, nfft, device="cpu") + attenuation = sos_filter_module(sos, nfft, device=device) return attenuation, input_delays, main_delays, output_delays @@ -133,6 +136,7 @@ def _( OrderedDict, allpass_fdn_core, attenuation, + device, dsp, gain_module, input_delays, @@ -148,8 +152,8 @@ def _( C_out = np.random.randn(2, N) C_out = C_out / np.linalg.norm(C_out, axis=1, keepdims=True) - gain_B_in = gain_module(B_in, nfft, device="cpu") - gain_C_out = gain_module(C_out, nfft, device="cpu") + gain_B_in = gain_module(B_in, nfft, device=device) + gain_C_out = gain_module(C_out, nfft, device=device) # Recursion: fF = allpass FDN → attenuation, fB = main delays feedforward = system.Series( @@ -177,8 +181,8 @@ def _( model = system.Shell( core=core_chain, - input_layer=dsp.FFT(nfft).to("cpu"), - output_layer=dsp.iFFT(nfft).to("cpu"), + input_layer=dsp.FFT(nfft).to(device), + output_layer=dsp.iFFT(nfft).to(device), ) ir_stereo = pyFDN.flamo_time_response(model).squeeze() diff --git a/src/pyFDN/auxiliary/coupled_rooms.py b/src/pyFDN/auxiliary/coupled_rooms.py index f52a35f0..cecbc72c 100644 --- a/src/pyFDN/auxiliary/coupled_rooms.py +++ b/src/pyFDN/auxiliary/coupled_rooms.py @@ -47,7 +47,7 @@ def create_coupled_rooms_fdn(): num_output = 2 # Device configuration - device = "cuda" if torch.cuda.is_available() else "cpu" + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") alias_decay_db = 0 # No anti-aliasing for exact reproduction # Exact delay values From 13f7e1560756a594015d0b1a7d653fdb2418544a Mon Sep 17 00:00:00 2001 From: Sebastian Jiro Schlecht <45119263+SebastianJiroSchlecht@users.noreply.github.com> Date: Fri, 25 Sep 2026 10:29:40 +0200 Subject: [PATCH 4/5] Require flamo>=0.2.20 and drop the torch.device workaround FLAMO 0.2.20 stores every module's device the way a tensor reports it (gdalsanto/flamo#38), so device="cpu" strings and .to() mix again and torch.device("cuda") no longer mismatches cuda:0. Revert 38cd49a: the allpass example is back on plain "cpu" strings and now doubles as the regression test in test_marimo_examples_run. Co-Authored-By: Claude Opus 5.5 --- .../allpass_FDN/example_allpass_FDN_in_FDN.py | 35 ++++++++++--------- pyproject.toml | 2 +- src/pyFDN/auxiliary/coupled_rooms.py | 2 +- 3 files changed, 21 insertions(+), 18 deletions(-) diff --git a/examples/allpass_FDN/example_allpass_FDN_in_FDN.py b/examples/allpass_FDN/example_allpass_FDN_in_FDN.py index 7dad5376..4f931938 100644 --- a/examples/allpass_FDN/example_allpass_FDN_in_FDN.py +++ b/examples/allpass_FDN/example_allpass_FDN_in_FDN.py @@ -34,7 +34,6 @@ def _(): from collections import OrderedDict import numpy as np - import torch from flamo.processor import dsp, system import pyFDN @@ -44,9 +43,7 @@ def _(): fs = 48000 nfft = 2**17 N = 4 # number of delay lines (1→N input gain, N→2 output gain) - device = torch.device("cpu") return ( - device, fs, N, OrderedDict, @@ -72,7 +69,7 @@ def _(mo): @app.cell -def _(device, fs, N, nfft, np, pyFDN): +def _(fs, N, nfft, np, pyFDN): # Homogeneous MIMO allpass FDN: delays, G = diag(g^delays), A = G @ U, complete_fdn delays_sch = np.random.randint(51, 300, size=N) _g = pyFDN.rt_to_gain_per_sample(0.07, fs) @@ -83,8 +80,15 @@ def _(device, fs, N, nfft, np, pyFDN): # FLAMO core (N→N), no Shell — for use inside the recursion allpass_fdn_core = pyFDN.dss_to_flamo( - A_sch, B_sch, C_sch, D_sch, delays_sch, fs, nfft=nfft, shell=False, - device=device, + A_sch, + B_sch, + C_sch, + D_sch, + delays_sch, + fs, + nfft=nfft, + shell=False, + device="cpu", ) return (allpass_fdn_core,) @@ -100,7 +104,7 @@ def _(mo): @app.cell -def _(device, fs, N, delay_module, nfft, np, pyFDN, sos_filter_module): +def _(fs, N, delay_module, nfft, np, pyFDN, sos_filter_module): # Main delays (feedback path), input and output delays — in seconds main_delay_sec = np.random.uniform(0.02, 0.04, size=N) input_delay_sec = np.linspace(0.01, 0.02 * N, N) + np.random.uniform( @@ -108,15 +112,15 @@ def _(device, fs, N, delay_module, nfft, np, pyFDN, sos_filter_module): ) output_delay_sec = np.linspace(0.01, 0.02, N) + np.random.uniform(0, 0.001, size=N) - main_delays = delay_module(main_delay_sec, nfft, fs=fs, device=device) - input_delays = delay_module(input_delay_sec, nfft, fs=fs, device=device) - output_delays = delay_module(output_delay_sec, nfft, fs=fs, device=device) + main_delays = delay_module(main_delay_sec, nfft, fs=fs, device="cpu") + input_delays = delay_module(input_delay_sec, nfft, fs=fs, device="cpu") + output_delays = delay_module(output_delay_sec, nfft, fs=fs, device="cpu") # Attenuation: first-order absorption, canonical (1, 6, N) SOS bank. main_delay_smp = np.round(main_delay_sec * fs).astype(float) rt_dc, rt_ny = 1.4, 0.3 sos = pyFDN.decay_to_first_order_shelf(rt_dc, rt_ny, None, main_delay_smp, fs=fs) - attenuation = sos_filter_module(sos, nfft, device=device) + attenuation = sos_filter_module(sos, nfft, device="cpu") return attenuation, input_delays, main_delays, output_delays @@ -136,7 +140,6 @@ def _( OrderedDict, allpass_fdn_core, attenuation, - device, dsp, gain_module, input_delays, @@ -152,8 +155,8 @@ def _( C_out = np.random.randn(2, N) C_out = C_out / np.linalg.norm(C_out, axis=1, keepdims=True) - gain_B_in = gain_module(B_in, nfft, device=device) - gain_C_out = gain_module(C_out, nfft, device=device) + gain_B_in = gain_module(B_in, nfft, device="cpu") + gain_C_out = gain_module(C_out, nfft, device="cpu") # Recursion: fF = allpass FDN → attenuation, fB = main delays feedforward = system.Series( @@ -181,8 +184,8 @@ def _( model = system.Shell( core=core_chain, - input_layer=dsp.FFT(nfft).to(device), - output_layer=dsp.iFFT(nfft).to(device), + input_layer=dsp.FFT(nfft).to("cpu"), + output_layer=dsp.iFFT(nfft).to("cpu"), ) ir_stereo = pyFDN.flamo_time_response(model).squeeze() diff --git a/pyproject.toml b/pyproject.toml index 32aa1ec7..46224337 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ "numpy", "soundfile", "scipy", - "flamo>=0.2.18", # Shell.set_nfft (train and render at different FFT sizes) + "flamo>=0.2.20", # canonical module devices (string "cpu"/"cuda" mix with .to()) "flareverb", ] diff --git a/src/pyFDN/auxiliary/coupled_rooms.py b/src/pyFDN/auxiliary/coupled_rooms.py index cecbc72c..f52a35f0 100644 --- a/src/pyFDN/auxiliary/coupled_rooms.py +++ b/src/pyFDN/auxiliary/coupled_rooms.py @@ -47,7 +47,7 @@ def create_coupled_rooms_fdn(): num_output = 2 # Device configuration - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = "cuda" if torch.cuda.is_available() else "cpu" alias_decay_db = 0 # No anti-aliasing for exact reproduction # Exact delay values From 9966d1fa8571ca2ac2113d631a6d1fcec5ce97b2 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 26 Sep 2026 06:59:38 +0000 Subject: [PATCH 5/5] Finish the GPU training fix: device defaults, tests, example restore - train_fdn defaults to the model's device instead of "cpu", so a model built on CUDA is no longer moved back to the CPU for training. - wrap_fdn_shell defaults to the core's own device, so wrapping a CPU core on a GPU machine no longer fails in Shell's device check. - Restore model.set_nfft(train_nfft / render_nfft) in the training example, removed in 567bb30; the example trains and renders on different FFT grids again. - Fix the indentation of the device="cpu" lines in the multislope example. - Tests: Shell assembled from a device="cpu" core and .to("cpu") layers (fails on flamo 0.2.19, passes on 0.2.20); build, render and train with default arguments stay on the model's device, and wrap_fdn_shell keeps the core's device, on CPU and on CUDA/MPS when available. - HISTORY entry. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Q6VyZo1pfcFKwK1Nz6cWcC --- HISTORY.rst | 6 ++ examples/example_multislope_rir_to_fdn.py | 8 +-- examples/example_train_fdn_to_rir.py | 2 + src/pyFDN/auxiliary/flamo.py | 15 +++-- src/pyFDN/train/engine.py | 7 +- src/pyFDN/train/response.py | 5 +- tests/test_train.py | 82 +++++++++++++++++++++++ 7 files changed, 114 insertions(+), 11 deletions(-) diff --git a/HISTORY.rst b/HISTORY.rst index c6e67958..262e3736 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -5,6 +5,12 @@ History Unreleased ---------- +* Fix GPU training (#231): ``model_response`` excites the model on its own + device, ``wrap_fdn_shell`` / ``dss_to_flamo`` / ``trainable_from_build`` + build the FFT layers on the requested device, and ``train_fdn`` now defaults + to the model's device instead of the CPU. Requires ``flamo>=0.2.20``, whose + canonical module devices let ``device="cpu"`` / ``"cuda"`` mix with + ``.to(...)`` (previously ``Shell`` rejected them as different devices). * Add ``Match``, a composable response loss that extracts a feature from the model response and the reference (waveform, magnitude, phase, mel magnitude, multi-resolution spectrogram, octave-band EDC, cumulative energy), diff --git a/examples/example_multislope_rir_to_fdn.py b/examples/example_multislope_rir_to_fdn.py index 9e15e732..d1a6e913 100644 --- a/examples/example_multislope_rir_to_fdn.py +++ b/examples/example_multislope_rir_to_fdn.py @@ -226,8 +226,8 @@ def geq_grid(band_values): nfft=nfft, post_delay=_absorption, shell=True, - device="cpu", - ) + device="cpu", + ) ).squeeze()[: len(rir)] _rt_flat, _ = pyFDN.estimate_rt_bands(_ir_flat, fs) @@ -249,8 +249,8 @@ def geq_grid(band_values): post_delay=_absorption, post_output=_eq[:, :, np.newaxis], shell=True, - device="cpu", - ) + device="cpu", + ) ).squeeze()[: len(rir)] slope_fdn_rt.append(_rt_flat) diff --git a/examples/example_train_fdn_to_rir.py b/examples/example_train_fdn_to_rir.py index 4f3ceaad..5d658281 100644 --- a/examples/example_train_fdn_to_rir.py +++ b/examples/example_train_fdn_to_rir.py @@ -265,6 +265,7 @@ def _(mo): @app.cell def _(device, fs, loss, model, pyFDN, render_nfft, rir_len, train_nfft): + model.set_nfft(train_nfft) log = pyFDN.train_fdn( model, loss, @@ -274,6 +275,7 @@ def _(device, fs, loss, model, pyFDN, render_nfft, rir_len, train_nfft): device=device, rng=0, ) + model.set_nfft(render_nfft) ir_trained = pyFDN.flamo_time_response(model, fs=fs).squeeze()[:rir_len] print( f"{log.steps_run} steps: loss " diff --git a/src/pyFDN/auxiliary/flamo.py b/src/pyFDN/auxiliary/flamo.py index 722305bf..b306e95c 100644 --- a/src/pyFDN/auxiliary/flamo.py +++ b/src/pyFDN/auxiliary/flamo.py @@ -715,7 +715,9 @@ def core_alias_decay_db(core: Any) -> float: return abs(float(value)) -def wrap_fdn_shell(core: Any, *, nfft: int, dtype: Any = None, device: Any = None) -> Any: +def wrap_fdn_shell( + core: Any, *, nfft: int, dtype: Any = None, device: Any = None +) -> Any: r""" Wrap an FDN core in a FLAMO ``Shell`` that returns the impulse response. @@ -742,7 +744,9 @@ def wrap_fdn_shell(core: Any, *, nfft: int, dtype: Any = None, device: Any = Non dtype : torch.dtype or None Dtype for the FFT/iFFT layers; defaults to float32. device : torch.device or None - Device for the FFT/iFFT layers; defaults to cuda if available else cpu. + Device of the returned shell. Defaults to the core's own device, so + wrapping never moves a core; with neither given, CUDA if available, + else CPU. Returns ------- @@ -759,12 +763,15 @@ def wrap_fdn_shell(core: Any, *, nfft: int, dtype: Any = None, device: Any = Non from flamo.processor import dsp, system torch_dtype = torch.float32 if dtype is None else dtype - dev = _get_device(device) + dev = _get_device(device if device is not None else getattr(core, "device", None)) shell = system.Shell( core=core, input_layer=dsp.FFT(nfft, dtype=torch_dtype), output_layer=dsp.iFFTAntiAlias( - nfft, alias_decay_db=core_alias_decay_db(core), device=dev, dtype=torch_dtype + nfft, + alias_decay_db=core_alias_decay_db(core), + device=dev, + dtype=torch_dtype, ), ) return shell.to(dev) diff --git a/src/pyFDN/train/engine.py b/src/pyFDN/train/engine.py index 2d4bafef..2481a110 100644 --- a/src/pyFDN/train/engine.py +++ b/src/pyFDN/train/engine.py @@ -81,7 +81,10 @@ def train_fdn( tol : float Relative-improvement threshold for the plateau early stop. device, dtype : optional - Torch device / dtype (default cpu / float32). + Torch device / dtype. ``device`` defaults to the model's own device + (``model.device``, e.g. CUDA when it was built there), so training + does not silently move a GPU model back to the CPU; ``dtype`` + defaults to float32. rng : int or None Integer seed for ``torch.manual_seed``. log : bool @@ -92,7 +95,7 @@ def train_fdn( import torch from flamo.optimize.trainer import EagerTrainer - dev = "cpu" if device is None else device + dev = getattr(model, "device", "cpu") if device is None else device torch_dtype = torch.float32 if dtype is None else dtype if rng is not None: diff --git a/src/pyFDN/train/response.py b/src/pyFDN/train/response.py index cad29b2b..20ed0c53 100644 --- a/src/pyFDN/train/response.py +++ b/src/pyFDN/train/response.py @@ -127,7 +127,10 @@ def model_response(model: Any, excitation: torch.Tensor | None = None) -> Respon n_in, n_out = int(model.input_channels), int(model.output_channels) if excitation is None: excitation = impulse_excitation( - n_in, int(model.nfft), device=model.device, dtype=getattr(model, "dtype", None) + n_in, + int(model.nfft), + device=model.device, + dtype=getattr(model, "dtype", None), ) y = model(excitation) # (n_in, nfft, n_out) if y.shape[0] != n_in or y.shape[2] != n_out: diff --git a/tests/test_train.py b/tests/test_train.py index bbd4228e..5f51cc18 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -1733,3 +1733,85 @@ def test_a_loss_reused_on_a_new_response_rebuilds_its_reference(make_loss): long_ir = np.concatenate([ir, np.zeros(2**12)]) moved = Response(h=_as_h(long_ir).to(torch.float64), fs=fs) assert float(loss(moved)) == pytest.approx(float(make_loss(reference)(moved))) + + +# --- devices (#231) ------------------------------------------------------------ + + +def test_shell_accepts_parts_built_with_a_string_device_and_moved_with_to(): + """Regression for #231: a core built with ``device="cpu"`` next to layers + moved with ``.to("cpu")`` used to fail inside ``Shell`` (FLAMO < 0.2.20 + compared the string with the ``torch.device`` that ``.to`` stores).""" + from flamo.processor import dsp, system + + nfft = 2**10 + core = pyFDN.dss_to_flamo( + 0.5 * np.eye(2), + np.ones((2, 1)), + np.ones((1, 2)), + np.zeros((1, 1)), + np.array([101, 137]), + 48000, + nfft=nfft, + shell=False, + device="cpu", + ) + model = system.Shell( + core=core, + input_layer=dsp.FFT(nfft).to("cpu"), + output_layer=dsp.iFFT(nfft).to("cpu"), + ) + impulse = torch.zeros(1, nfft, 1) + impulse[0, 0, 0] = 1.0 + assert torch.isfinite(model(impulse)).all() + + +def _devices(): + """CPU always; each accelerator when present, else a skipped CUDA case.""" + accelerators = [] + if torch.cuda.is_available(): + accelerators.append("cuda") + mps = getattr(torch.backends, "mps", None) + if mps is not None and mps.is_available(): + accelerators.append("mps") + if not accelerators: + accelerators.append( + pytest.param("cuda", marks=pytest.mark.skip(reason="no GPU available")) + ) + return ["cpu", *accelerators] + + +@pytest.mark.parametrize("device", _devices()) +def test_training_stays_on_the_device_the_model_was_built_on(device): + """Build on the accelerator, then render and train with default arguments: + the response and the trained parameters stay on that device (#231).""" + model = build_fdn(N=4, rt=None, nfft=2**10, device=device, rng=0) + + response = model_response(model) + assert response.h.device.type == device + + log = train_fdn(model, FlatMagnitude(), max_steps=3, lr=3e-3, rng=0) + assert log.steps_run > 0 + assert all(p.device.type == device for p in model.parameters()) + + +@pytest.mark.parametrize("device", _devices()) +def test_wrap_fdn_shell_keeps_the_core_on_its_device(device): + """Without ``device=``, the shell follows the core instead of the default + accelerator, so wrapping a CPU core on a GPU machine does not fail.""" + from pyFDN.auxiliary.flamo import wrap_fdn_shell + + core = pyFDN.dss_to_flamo( + 0.5 * np.eye(2), + np.ones((2, 1)), + np.ones((1, 2)), + np.zeros((1, 1)), + np.array([101, 137]), + 48000, + nfft=2**10, + shell=False, + device=device, + ) + shell = wrap_fdn_shell(core, nfft=2**10) + assert torch.device(shell.device).type == device + assert model_response(shell).h.device.type == device