231 gpu training still faulty

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).
 * ``import pyFDN`` no longer loads torch or flamo, which makes it about three
   times faster; the training API and ``flamo_to_pr`` load on first use. Every
   public name is still exported from ``pyFDN``.

examples/allpass_FDN/example_allpass_FDN_in_FDN.py

@@ -80,7 +80,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(
-        delays_sch, A_sch, B_sch, C_sch, D_sch, fs, nfft=nfft, shell=False
+        delays_sch, A_sch, B_sch, C_sch, D_sch, fs, nfft=nfft, shell=False, device="cpu"
     )
     return (allpass_fdn_core,)
 
@@ -104,15 +104,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 +147,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 +176,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()

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()
 

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,)
 

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)

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]

examples/example_multislope_rir_to_fdn.py

@@ -202,6 +202,7 @@ def geq_grid(band_values):
                 nfft=nfft,
                 post_delay=_absorption,
                 shell=True,
+                device="cpu",
             )
         ).squeeze()[: len(rir)]
 
@@ -224,6 +225,7 @@ def geq_grid(band_values):
                 post_delay=_absorption,
                 post_output=_eq[:, :, np.newaxis],
                 shell=True,
+                device="cpu",
             )
         ).squeeze()[: len(rir)]
 

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,

examples/example_pole_boundaries.py

@@ -122,6 +122,7 @@ def _(
         shell=False,
         post_delay=absorption,
         dtype=torch.float64,
+        device="cpu",
     )
     _residues, poles, _direct_term, _is_pair, _meta = pyFDN.flamo_to_pr(
         model,

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]
 

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]
 

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
 
 

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

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())
 ]
 
 [project.optional-dependencies]

src/pyFDN/auxiliary/flamo.py

@@ -695,7 +695,9 @@ 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.
 
@@ -721,6 +723,10 @@ 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 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
     -------
@@ -735,10 +741,17 @@ 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 = default_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), dtype=torch_dtype
+            nfft,
+            alias_decay_db=core_alias_decay_db(core),
+            device=dev,
+            dtype=torch_dtype,
         ),
     )
+    return shell.to(dev)

src/pyFDN/train/build.py

@@ -319,7 +319,7 @@ def trainable_from_build(
         post_matrix=build.post_matrix if post_matrix is None else post_matrix,
         post_output=build.post_output if post_output is None else 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(

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:

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), 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:

src/pyFDN/translate/dss_to_flamo.py

@@ -102,7 +102,7 @@ def dss_to_flamo(
     if shell:
         from pyFDN.auxiliary.flamo import wrap_fdn_shell
 
-        return wrap_fdn_shell(core, nfft=nfft, dtype=dtype)
+        return wrap_fdn_shell(core, nfft=nfft, dtype=dtype, device=device)
     return core
 
 

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()

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]