Repository: Rikorose/DeepFilterNet (language: python)

Issue:
WIP ms



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index cd09dcd6d..e5014b968 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,8 +2,11 @@
 *.wav
 *.png
 *.pdf
+assets/hostbs.ini
+.cache*
 out/
 export/
+assets/valid_cache*
 DeepFilterNet/poetry.lock
 
 ### Rust gitignore ###
diff --git a/DeepFilterNet/df/checkpoint.py b/DeepFilterNet/df/checkpoint.py
index 616b4999d..80327fb1f 100644
--- a/DeepFilterNet/df/checkpoint.py
+++ b/DeepFilterNet/df/checkpoint.py
@@ -3,6 +3,7 @@
 import re
 from typing import List, Optional, Tuple, Union
 
+import numpy as np
 import torch
 from loguru import logger
 from torch import nn
@@ -14,7 +15,7 @@
 
 
 def get_epoch(cp) -> int:
-    return int(os.path.splitext(os.path.basename(cp))[0].split("_")[-1])
+    return int(os.path.basename(cp).split(".")[0].split("_")[-1])
 
 
 def load_model(
@@ -23,6 +24,8 @@ def load_model(
     jit: bool = False,
     mask_only: bool = False,
     train_df_only: bool = False,
+    extension: str = "ckpt",
+    epoch: Union[str, int] = "latest",
 ) -> Tuple[nn.Module, int]:
     if mask_only and train_df_only:
         raise ValueError("Only one of `mask_only` `train_df_only` can be enabled")
@@ -30,10 +33,13 @@ def load_model(
     if jit:
         model = torch.jit.script(model)
     blacklist: List[str] = config("CP_BLACKLIST", [], Csv(), save=False, section="train")  # type: ignore
-    epoch = 0
     if cp_dir is not None:
-        epoch = read_cp(model, "model", cp_dir, blacklist=blacklist)
+        epoch = read_cp(
+            model, "model", cp_dir, blacklist=blacklist, extension=extension, epoch=epoch
+        )
         epoch = 0 if epoch is None else epoch
+    else:
+        epoch = 0
     return model, epoch
 
 
@@ -45,11 +51,26 @@ def read_cp(
     extension="ckpt",
     blacklist=[],
 ):
-    checkpoints = glob.glob(os.path.join(dirname, f"{name}*.{extension}"))
+    checkpoints = []
+    if isinstance(epoch, str):
+        assert epoch in ("best", "latest")
+    if epoch == "best":
+        checkpoints = glob.glob(os.path.join(dirname, f"{name}*.{extension}.best"))
+        if len(checkpoints) == 0:
+            logger.warning("Could not find `best` checkpoint. Checking for default...")
+    if len(checkpoints) == 0:
+        checkpoints = glob.glob(os.path.join(dirname, f"{name}*.{extension}"))
+        checkpoints += glob.glob(os.path.join(dirname, f"{name}*.{extension}.best"))
     if len(checkpoints) == 0:
         return None
-    latest = max(checkpoints, key=get_epoch)
-    epoch = get_epoch(latest)
+    if isinstance(epoch, int):
+        latest = next((x for x in checkpoints if get_epoch(x) == epoch), None)
+        if latest is None:
+            logger.error(f"Could not find checkpoint of epoch {epoch}")
+            exit(1)
+    else:
+        latest = max(checkpoints, key=get_epoch)
+        epoch = get_epoch(latest)
     logger.info("Found checkpoint {} with epoch {}".format(latest, epoch))
     latest = torch.load(latest, map_location="cpu")
     latest = {k.replace("clc", "df"): v for k, v in latest.items()}
@@ -87,13 +108,35 @@ def write_cp(
     dirname: str,
     epoch: int,
     extension="ckpt",
+    metric: Optional[float] = None,
+    cmp="min",
 ):
     check_finite_module(obj)
+    n_keep = config("n_checkpoint_history", default=3, cast=int, section="train")
+    n_keep_best = config("n_best_checkpoint_history", default=5, cast=int, section="train")
+    if metric is not None:
+        assert cmp in ("min", "max")
+        metric = float(metric)  # Make sure it is not an integer
+        # Each line contains a previous best with entries: (epoch, metric)
+        with open(os.path.join(dirname, ".best"), "a+") as prev_best_f:
+            prev_best_f.seek(0)  # "a+" creates a file in read/write mode without truncating
+            lines = prev_best_f.readlines()
+            if len(lines) == 0:
+                prev_best = float("inf" if cmp == "min" else "-inf")
+            else:
+                prev_best = float(lines[-1].strip().split(" ")[1])
+            cmp = "__lt__" if cmp == "min" else "__gt__"
+            if getattr(metric, cmp)(prev_best):
+                logger.info(f"Saving new best checkpoint at epoch {epoch} with metric: {metric}")
+                prev_best_f.seek(0, os.SEEK_END)
+                np.savetxt(prev_best_f, np.array([[float(epoch), metric]]))
+                cp_name = os.path.join(dirname, f"{name}_{epoch}.{extension}.best")
+                torch.save(obj.state_dict(), cp_name)
+                cleanup(name, dirname, extension + ".best", nkeep=n_keep_best)
     cp_name = os.path.join(dirname, f"{name}_{epoch}.{extension}")
     logger.info(f"Writing checkpoint {cp_name} with epoch {epoch}")
     torch.save(obj.state_dict(), cp_name)
-    n_old_cps = config("n_checkpoint_history", default=-1, cast=int, save=False, section="train")
-    cleanup(name, dirname, extension, nkeep=n_old_cps)
+    cleanup(name, dirname, extension, nkeep=n_keep)
 
 
 def cleanup(name: str, dirname: str, extension: str, nkeep=5):
@@ -106,3 +149,63 @@ def cleanup(name: str, dirname: str, extension: str, nkeep=5):
     for cp in checkpoints[nkeep:]:
         logger.debug("Removing old checkpoint: {}".format(cp))
         os.remove(cp)
+
+
+def check_patience(
+    dirname: str, max_patience: int, new_metric: float, cmp: str = "min", raise_: bool = True
+):
+    cmp = "__lt__" if cmp == "min" else "__gt__"
+    new_metric = float(new_metric)  # Make sure it is not an integer
+    prev_patience, prev_metric = read_patience(dirname)
+    if prev_patience is None or getattr(new_metric, cmp)(prev_metric):
+        # We have a better new_metric, reset patience
+        write_patience(dirname, 0, new_metric)
+    else:
+        # We don't have a better metric, decrement patience
+        new_patience = prev_patience + 1
+        write_patience(dirname, new_patience, prev_metric)
+        if new_patience >= max_patience:
+            if raise_:
+                raise ValueError(
+                    f"No improvements on validation metric ({new_metric}) for {max_patience} epochs. "
+                    "Stopping."
+                )
+            else:
+                return False
+    return True
+
+
+def read_patience(dirname: str) -> Tuple[Optional[int], float]:
+    fn = os.path.join(dirname, ".patience")
+    if not os.path.isfile(fn):
+        return None, 0.0
+    patience, metric = np.loadtxt(fn)
+    return int(patience), float(metric)
+
+
+def write_patience(dirname: str, new_patience: int, metric: float):
+    return np.savetxt(os.path.join(dirname, ".patience"), [new_patience, metric])
+
+
+def test_check_patience():
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as d:
+        check_patience(d, 3, 1.0)
+        check_patience(d, 3, 1.0)
+        check_patience(d, 3, 1.0)
+        assert check_patience(d, 3, 1.0, raise_=False) is False
+
+    with tempfile.TemporaryDirectory() as d:
+        check_patience(d, 3, 1.0)
+        check_patience(d, 3, 0.9)
+        check_patience(d, 3, 1.0)
+        check_patience(d, 3, 1.0)
+        assert check_patience(d, 3, 1.0, raise_=False) is False
+
+    with tempfile.TemporaryDirectory() as d:
+        check_patience(d, 3, 1.0, cmp="max")
+        check_patience(d, 3, 1.9, cmp="max")
+        check_patience(d, 3, 1.0, cmp="max")
+        check_patience(d, 3, 1.0, cmp="max")
+        assert check_patience(d, 3, 1.0, cmp="max", raise_=False) is False
diff --git a/DeepFilterNet/df/config.py b/DeepFilterNet/df/config.py
index 61bfd162f..574e856ff 100644
--- a/DeepFilterNet/df/config.py
+++ b/DeepFilterNet/df/config.py
@@ -12,14 +12,14 @@
 class DfParams:
     def __init__(self):
         self.sr: int = config("SR", cast=int, default=48_000, section="DF")
-        self.fft_size: int = config("FFT_SIZE", cast=int, default=384, section="DF")
-        self.hop_size: int = config("HOP_SIZE", cast=int, default=192, section="DF")
-        self.nb_erb: int = config("NB_ERB", cast=int, default=16, section="DF")
-        self.nb_df: int = config("NB_DF", cast=int, default=24, section="DF")
+        self.fft_size: int = config("FFT_SIZE", cast=int, default=960, section="DF")
+        self.hop_size: int = config("HOP_SIZE", cast=int, default=480, section="DF")
+        self.nb_erb: int = config("NB_ERB", cast=int, default=32, section="DF")
+        self.nb_df: int = config("NB_DF", cast=int, default=96, section="DF")
         self.norm_tau: float = config("NORM_TAU", 1, float, section="DF")
         self.lsnr_max: int = config("LSNR_MAX", 35, int, section="DF")
         self.lsnr_min: int = config("LSNR_MIN", -15, int, section="DF")
-        self.min_nb_freqs = config("MIN_NB_ERB_FREQS", 1, int, section="DF")
+        self.min_nb_freqs = config("MIN_NB_ERB_FREQS", 2, int, section="DF")
 
 
 class Config:
@@ -33,9 +33,11 @@ def __init__(self):
         self.modified = False
         self.allow_defaults = True
 
-    def load(self, path: Optional[str], config_must_exist=False, allow_defaults=True):
+    def load(
+        self, path: Optional[str], config_must_exist=False, allow_defaults=True, allow_reload=False
+    ):
         self.allow_defaults = allow_defaults
-        if self.parser is not None:
+        if self.parser is not None and not allow_reload:
             raise ValueError("Config already loaded")
         self.parser = ConfigParser()
         self.path = path
@@ -69,15 +71,17 @@ def tostr(self, value, cast):
             return "".join(str(v) + cast.delimiter for v in value)[:-1]
         return str(value)
 
-    def set(self, section: str, option: str, value: T, cast: Type[T]):
+    def set(self, option: str, value: T, cast: Type[T], section: Optional[str] = None) -> T:
+        section = self.DEFAULT_SECTION if section is None else section
         section = section.lower()
         if not self.parser.has_section(section):
-            raise ValueError(f"Section not found: {section}")
+            self.parser.add_section(section)
         if self.parser.has_option(section, option):
             if value == self.cast(self.parser.get(section, option), cast):
-                return
+                return value
         self.modified = True
         self.parser.set(section, option, self.tostr(value, cast))
+        return value
 
     def __call__(
         self,
@@ -115,7 +119,7 @@ def __call__(
         else:
             value = default
             if save:
-                self.set(section, option, value, cast)
+                self.set(option, value, cast, section)
         return self.cast(value, cast)
 
     def cast(self, value, cast):
@@ -129,7 +133,8 @@ def cast(self, value, cast):
             raise ValueError("Parse error")
         return cast(value)
 
-    def get(self, option: str, section: str, cast: Type[T] = str) -> T:
+    def get(self, option: str, cast: Type[T] = str, section: Optional[str] = None) -> T:
+        section = self.DEFAULT_SECTION if section is None else section
         assert self.parser.has_section(section)
         assert self.parser.has_option(section, option)
         return cast(self.parser.get(section, option))
@@ -152,7 +157,7 @@ def read_from_section(
             self.modified = True
         return value
 
-    def overwrite(self, section: str, option: str, value: str):
+    def overwrite(self, section: str, option: str, value: Any):
         if not self.parser.has_section(section):
             return ValueError(f"Section not found: '{section}'")
         if not self.parser.has_option(section, option):
diff --git a/DeepFilterNet/df/deepfilternet2.py b/DeepFilterNet/df/deepfilternet2.py
new file mode 100644
index 000000000..2cc6af0e2
--- /dev/null
+++ b/DeepFilterNet/df/deepfilternet2.py
@@ -0,0 +1,447 @@
+from functools import partial
+from typing import Final, List, Optional, Tuple
+
+import torch
+from loguru import logger
+from torch import Tensor, nn
+
+from df.config import Csv, DfParams, config
+from df.logger import log_deprecated
+from df.modules import (
+    Conv2dNormAct,
+    ConvTranspose2dNormAct,
+    DfOp,
+    GroupedGRU,
+    GroupedLinear,
+    GroupedLinearEinsum,
+    Mask,
+    SqueezedGRU,
+    erb_fb,
+    get_device,
+)
+from df.multistagenet import MagCompression
+from libdf import DF
+
+
+class ModelParams(DfParams):
+    section = "deepfilternet"
+
+    def __init__(self):
+        super().__init__()
+        self.df_order: int = config("DF_ORDER", cast=int, default=5, section=self.section)
+        self.df_lookahead: int = config("DF_LOOKAHEAD", cast=int, default=0, section=self.section)
+        self.conv_lookahead: int = config(
+            "CONV_LOOKAHEAD", cast=int, default=0, section=self.section
+        )
+        self.conv_ch: int = config("CONV_CH", cast=int, default=16, section=self.section)
+        self.conv_depthwise: bool = config(
+            "CONV_DEPTHWISE", cast=bool, default=True, section=self.section
+        )
+        self.convt_depthwise: bool = config(
+            "CONVT_DEPTHWISE", cast=bool, default=True, section=self.section
+        )
+        self.conv_kernel: List[int] = config(
+            "CONV_KERNEL", cast=Csv(int), default=(1, 3), section=self.section  # type: ignore
+        )
+        self.emb_hidden_dim: int = config(
+            "EMB_HIDDEN_DIM", cast=int, default=256, section=self.section
+        )
+        self.emb_num_layers: int = config(
+            "EMB_NUM_LAYERS", cast=int, default=2, section=self.section
+        )
+        self.df_hidden_dim: int = config(
+            "DF_HIDDEN_DIM", cast=int, default=256, section=self.section
+        )
+        self.df_gru_skip: str = config("DF_GRU_SKIP", default="none", section=self.section)
+        self.df_output_layer: str = config(
+            "DF_OUTPUT_LAYER", default="linear", section=self.section
+        )
+        self.df_pathway_kernel_size_t: int = config(
+            "DF_PATHWAY_KERNEL_SIZE_T", cast=int, default=1, section=self.section
+        )
+        self.enc_concat: bool = config("ENC_CONCAT", cast=bool, default=False, section=self.section)
+        self.df_num_layers: int = config("DF_NUM_LAYERS", cast=int, default=3, section=self.section)
+        self.df_n_iter: int = config("DF_N_ITER", cast=int, default=2, section=self.section)
+        self.gru_type: str = config("GRU_TYPE", default="grouped", section=self.section)
+        self.gru_groups: int = config("GRU_GROUPS", cast=int, default=1, section=self.section)
+        self.lin_groups: int = config("LINEAR_GROUPS", cast=int, default=1, section=self.section)
+        self.group_shuffle: bool = config(
+            "GROUP_SHUFFLE", cast=bool, default=True, section=self.section
+        )
+        self.dfop_method: str = config(
+            "DFOP_METHOD", cast=str, default="real_unfold", section=self.section
+        )
+        self.mask_pf: bool = config("MASK_PF", cast=bool, default=False, section=self.section)
+
+
+def init_model(df_state: Optional[DF] = None, run_df: bool = True, train_mask: bool = True):
+    p = ModelParams()
+    if df_state is None:
+        df_state = DF(sr=p.sr, fft_size=p.fft_size, hop_size=p.hop_size, nb_bands=p.nb_erb)
+    erb = erb_fb(df_state.erb_widths(), p.sr, inverse=False)
+    erb_inverse = erb_fb(df_state.erb_widths(), p.sr, inverse=True)
+    model = DfNet(erb, erb_inverse, run_df, train_mask)
+    return model.to(device=get_device())
+
+
+class Add(nn.Module):
+    def forward(self, a, b):
+        return a + b
+
+
+class Concat(nn.Module):
+    def forward(self, a, b):
+        return torch.cat((a, b), dim=-1)
+
+
+class Encoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        p = ModelParams()
+        assert p.nb_erb % 4 == 0, "erb_bins should be divisible by 4"
+
+        self.erb_conv0 = Conv2dNormAct(1, p.conv_ch, kernel_size=(3, 3), bias=False, separable=True)
+        conv_layer = partial(
+            Conv2dNormAct,
+            in_ch=p.conv_ch,
+            out_ch=p.conv_ch,
+            kernel_size=p.conv_kernel,
+            bias=False,
+            separable=True,
+        )
+        self.erb_conv1 = conv_layer(fstride=2)
+        self.erb_conv2 = conv_layer(fstride=2)
+        self.erb_conv3 = conv_layer(fstride=1)
+        self.df_conv0 = Conv2dNormAct(2, p.conv_ch, kernel_size=(3, 3), bias=False, separable=True)
+        self.df_conv1 = conv_layer(fstride=2)
+        self.erb_bins = p.nb_erb
+        self.emb_in_dim = p.conv_ch * p.nb_erb // 4
+        self.emb_out_dim = p.emb_hidden_dim
+        if p.gru_type == "grouped":
+            self.df_fc_emb = GroupedLinear(
+                p.conv_ch * p.nb_df // 2, self.emb_in_dim, groups=p.lin_groups
+            )
+        else:
+            df_fc_emb = GroupedLinearEinsum(
+                p.conv_ch * p.nb_df // 2, self.emb_in_dim, groups=p.lin_groups
+            )
+            self.df_fc_emb = nn.Sequential(df_fc_emb, nn.ReLU(inplace=True))
+        if p.enc_concat:
+            self.emb_in_dim *= 2
+            self.combine = Concat()
+        else:
+            self.combine = Add()
+        self.emb_out_dim = p.emb_hidden_dim
+        self.emb_n_layers = p.emb_num_layers
+        assert p.gru_type in ("grouped", "squeeze"), f"But got {p.gru_type}"
+        if p.gru_type == "grouped":
+            self.emb_gru = GroupedGRU(
+                self.emb_in_dim,
+                self.emb_out_dim,
+                num_layers=1,
+                batch_first=True,
+                groups=p.gru_groups,
+                shuffle=p.group_shuffle,
+                add_outputs=True,
+            )
+        else:
+            self.emb_gru = SqueezedGRU(
+                self.emb_in_dim,
+                self.emb_out_dim,
+                num_layers=1,
+                batch_first=True,
+                linear_groups=p.lin_groups,
+                linear_act_layer=partial(nn.ReLU, inplace=True),
+            )
+        self.lsnr_fc = nn.Sequential(nn.Linear(self.emb_out_dim, 1), nn.Sigmoid())
+        self.lsnr_scale = p.lsnr_max - p.lsnr_min
+        self.lsnr_offset = p.lsnr_min
+
+    def forward(
+        self, feat_erb: Tensor, feat_spec: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+        # Encodes erb; erb should be in dB scale + normalized; Fe are number of erb bands.
+        # erb: [B, 1, T, Fe]
+        # spec: [B, 2, T, Fc]
+        # b, _, t, _ = feat_erb.shape
+        e0 = self.erb_conv0(feat_erb)  # [B, C, T, F]
+        e1 = self.erb_conv1(e0)  # [B, C*2, T, F/2]
+        e2 = self.erb_conv2(e1)  # [B, C*4, T, F/4]
+        e3 = self.erb_conv3(e2)  # [B, C*4, T, F/4]
+        c0 = self.df_conv0(feat_spec)  # [B, C, T, Fc]
+        c1 = self.df_conv1(c0)  # [B, C*2, T, Fc]
+        cemb = c1.permute(0, 2, 3, 1).flatten(2)  # [B, T, -1]
+        cemb = self.df_fc_emb(cemb)  # [T, B, C * F/4]
+        emb = e3.permute(0, 2, 3, 1).flatten(2)  # [B, T, C * F/4]
+        emb = self.combine(emb, cemb)
+        emb, _ = self.emb_gru(emb)  # [B, T, -1]
+        lsnr = self.lsnr_fc(emb) * self.lsnr_scale + self.lsnr_offset
+        return e0, e1, e2, e3, emb, c0, lsnr
+
+
+class ErbDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        p = ModelParams()
+        assert p.nb_erb % 8 == 0, "erb_bins should be divisible by 8"
+
+        self.emb_out_dim = p.emb_hidden_dim
+
+        if p.gru_type == "grouped":
+            self.emb_gru = GroupedGRU(
+                p.conv_ch * p.nb_erb // 4,  # For compat
+                self.emb_out_dim,
+                num_layers=p.emb_num_layers - 1,
+                batch_first=True,
+                groups=p.gru_groups,
+                shuffle=p.group_shuffle,
+                add_outputs=True,
+            )
+            # SqueezedGRU uses GroupedLinearEinsum, so let's use it here as well
+            fc_emb = GroupedLinear(
+                p.emb_hidden_dim,
+                p.conv_ch * p.nb_erb // 4,
+                groups=p.lin_groups,
+                shuffle=p.group_shuffle,
+            )
+            self.fc_emb = nn.Sequential(fc_emb, nn.ReLU(inplace=True))
+        else:
+            self.emb_gru = SqueezedGRU(
+                self.emb_out_dim,
+                self.emb_out_dim,
+                output_size=p.conv_ch * p.nb_erb // 4,
+                num_layers=p.emb_num_layers - 1,
+                batch_first=True,
+                gru_skip_op=nn.Identity,
+                linear_groups=p.lin_groups,
+                linear_act_layer=partial(nn.ReLU, inplace=True),
+            )
+            self.fc_emb = nn.Identity()
+        tconv_layer = partial(
+            ConvTranspose2dNormAct,
+            kernel_size=p.conv_kernel,
+            bias=False,
+            separable=True,
+        )
+        conv_layer = partial(
+            Conv2dNormAct,
+            bias=False,
+            separable=True,
+        )
+        # convt: TransposedConvolution, convp: Pathway (encoder to decoder) convolutions
+        self.conv3p = conv_layer(p.conv_ch, p.conv_ch, kernel_size=1)
+        self.convt3 = conv_layer(p.conv_ch, p.conv_ch, kernel_size=p.conv_kernel)
+        self.conv2p = conv_layer(p.conv_ch, p.conv_ch, kernel_size=1)
+        self.convt2 = tconv_layer(p.conv_ch, p.conv_ch, fstride=2)
+        self.conv1p = conv_layer(p.conv_ch, p.conv_ch, kernel_size=1)
+        self.convt1 = tconv_layer(p.conv_ch, p.conv_ch, fstride=2)
+        self.conv0p = conv_layer(p.conv_ch, p.conv_ch, kernel_size=1)
+        self.conv0_out = conv_layer(
+            p.conv_ch, 1, kernel_size=p.conv_kernel, activation_layer=nn.Sigmoid
+        )
+
+    def forward(self, emb, e3, e2, e1, e0) -> Tensor:
+        # Estimates erb mask
+        b, _, t, f8 = e3.shape
+        emb, _ = self.emb_gru(emb)
+        emb = self.fc_emb(emb)
+        emb = emb.view(b, t, f8, -1).permute(0, 3, 1, 2)  # [B, C*8, T, F/8]
+        e3 = self.convt3(self.conv3p(e3) + emb)  # [B, C*4, T, F/4]
+        e2 = self.convt2(self.conv2p(e2) + e3)  # [B, C*2, T, F/2]
+        e1 = self.convt1(self.conv1p(e1) + e2)  # [B, C, T, F]
+        m = self.conv0_out(self.conv0p(e0) + e1)  # [B, 1, T, F]
+        return m
+
+
+class DfDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        p = ModelParams()
+        layer_width = p.conv_ch
+        self.emb_dim = p.emb_hidden_dim
+
+        self.df_n_hidden = p.df_hidden_dim
+        self.df_n_layers = p.df_num_layers
+        self.df_order = p.df_order
+        self.df_bins = p.nb_df
+        self.df_lookahead = p.df_lookahead
+        self.gru_groups = p.gru_groups
+
+        conv_layer = partial(Conv2dNormAct, separable=True, bias=False)
+        kt = p.df_pathway_kernel_size_t
+        self.df_convp = conv_layer(layer_width, self.df_order * 2, fstride=1, kernel_size=(kt, 1))
+        if p.gru_type == "grouped":
+            self.df_gru = GroupedGRU(
+                p.emb_hidden_dim,
+                p.df_hidden_dim,
+                num_layers=self.df_n_layers,
+                batch_first=True,
+                groups=p.gru_groups,
+                shuffle=p.group_shuffle,
+                add_outputs=True,
+            )
+        else:
+            self.df_gru = SqueezedGRU(
+                p.emb_hidden_dim,
+                p.df_hidden_dim,
+                num_layers=self.df_n_layers,
+                batch_first=True,
+                gru_skip_op=nn.Identity,
+                linear_act_layer=partial(nn.ReLU, inplace=True),
+            )
+        p.df_gru_skip = p.df_gru_skip.lower()
+        assert p.df_gru_skip in ("none", "identity", "groupedlinear")
+        self.df_skip: Optional[nn.Module]
+        if p.df_gru_skip == "none":
+            self.df_skip = None
+        elif p.df_gru_skip == "identity":
+            assert p.emb_hidden_dim == p.df_hidden_dim, "Dimensions do not match"
+            self.df_skip = nn.Identity()
+        elif p.df_gru_skip == "groupedlinear":
+            self.df_skip = GroupedLinearEinsum(
+                p.emb_hidden_dim, p.df_hidden_dim, groups=p.lin_groups
+            )
+        else:
+            raise NotImplementedError()
+        assert p.df_output_layer in ("linear", "groupedlinear")
+        self.df_out: nn.Module
+        out_dim = self.df_bins * self.df_order * 2
+        if p.df_output_layer == "linear":
+            df_out = nn.Linear(self.df_n_hidden, out_dim)
+        elif p.df_output_layer == "groupedlinear":
+            df_out = GroupedLinearEinsum(self.df_n_hidden, out_dim, groups=p.lin_groups)
+        else:
+            raise NotImplementedError
+        self.df_out = nn.Sequential(df_out, nn.Tanh())
+        self.df_fc_a = nn.Sequential(nn.Linear(self.df_n_hidden, 1), nn.Sigmoid())
+
+    def forward(self, emb: Tensor, c0: Tensor) -> Tuple[Tensor, Tensor]:
+        b, t, _ = emb.shape
+        c, _ = self.df_gru(emb)  # [B, T, H], H: df_n_hidden
+        if self.df_skip is not None:
+            c += self.df_skip(emb)
+        c0 = self.df_convp(c0).permute(0, 2, 3, 1)  # [B, T, F, O*2], channels_last
+        alpha = self.df_fc_a(c)  # [B, T, 1]
+        c = self.df_out(c)  # [B, T, F*O*2], O: df_order
+        c = c.view(b, t, self.df_bins, self.df_order * 2)  # [B, T, F, O*2]
+        # TODO: use c.add(c0).view(b, t, self.df_order, self.df_bins, 2)  # [B, T, O, F, 2]
+        c = c.add(c0).view(b, t, self.df_order, 2, self.df_bins).transpose(3, 4)  # [B, T, O, F, 2]
+        return c, alpha
+
+
+class DfDecoderLinear(nn.Module):
+    # For compat
+    def __init__(self):
+        super().__init__()
+        log_deprecated("Use of `linear` for `df_ouput_layer` is marked as deprecated.")
+        p = ModelParams()
+        layer_width = p.conv_ch
+        self.emb_dim = p.emb_hidden_dim
+
+        self.df_n_hidden = p.df_hidden_dim
+        self.df_n_layers = p.df_num_layers
+        self.df_order = p.df_order
+        self.df_bins = p.nb_df
+        self.df_lookahead = p.df_lookahead
+        self.gru_groups = p.gru_groups
+
+        conv_layer = partial(Conv2dNormAct, separable=True, bias=False)
+        kt = p.df_pathway_kernel_size_t
+        self.df_convp = conv_layer(layer_width, self.df_order * 2, fstride=1, kernel_size=(kt, 1))
+        self.df_gru = GroupedGRU(
+            p.emb_hidden_dim,
+            self.df_n_hidden,
+            num_layers=self.df_n_layers,
+            batch_first=True,
+            groups=p.gru_groups,
+            shuffle=p.group_shuffle,
+            add_outputs=True,
+        )
+        assert p.df_output_layer == "linear"
+        self.df_fc_out = nn.Sequential(
+            nn.Linear(self.df_n_hidden, self.df_bins * self.df_order * 2), nn.Tanh()
+        )
+        self.df_fc_a = nn.Sequential(nn.Linear(self.df_n_hidden, 1), nn.Sigmoid())
+
+    def forward(self, emb: Tensor, c0: Tensor) -> Tuple[Tensor, Tensor]:
+        b, t, _ = emb.shape
+        c, _ = self.df_gru(emb)  # [B, T, H], H: df_n_hidden
+        c0 = self.df_convp(c0).transpose(1, 2)  # [B, T, O*2, F], channels_last
+        alpha = self.df_fc_a(c)  # [B, T, 1]
+        c = self.df_fc_out(c)  # [B, T, F*O*2], O: df_order
+        c = c.view(b, t, self.df_order * 2, self.df_bins)  # [B, T, O*2, F]
+        c = c.add(c0).view(b, t, self.df_order, 2, self.df_bins).transpose(3, 4)  # [B, T, O, F, 2]
+        return c, alpha
+
+
+class DfNet(nn.Module):
+    run_df: Final[bool]
+
+    def __init__(
+        self,
+        erb_fb: Tensor,
+        erb_inv_fb: Tensor,
+        run_df: bool = True,
+        train_mask: bool = True,
+    ):
+        super().__init__()
+        p = ModelParams()
+        layer_width = p.conv_ch
+        assert p.nb_erb % 8 == 0, "erb_bins should be divisible by 8"
+        self.lookahead: int = p.conv_lookahead
+        self.freq_bins: int = p.fft_size // 2 + 1
+        self.emb_dim: int = layer_width * p.nb_erb
+        self.erb_bins: int = p.nb_erb
+        self.erb_comp = MagCompression(p.nb_erb)
+        if p.conv_lookahead > 0:
+            pad = (0, 0, -p.conv_lookahead, p.conv_lookahead)
+            self.pad = nn.ConstantPad2d(pad, 0.0)
+        else:
+            self.pad = nn.Identity()
+        self.register_buffer("erb_fb", erb_fb)
+        self.enc = Encoder()
+        self.erb_dec = ErbDecoder()
+        self.mask = Mask(erb_inv_fb, post_filter=p.mask_pf)
+
+        self.df_order = p.df_order
+        self.df_bins = p.nb_df
+        self.df_lookahead = p.df_lookahead
+        if p.df_output_layer == "linear":
+            self.df_dec = DfDecoderLinear()
+        else:
+            self.df_dec = DfDecoder()
+        self.df_op = torch.jit.script(
+            DfOp(
+                p.nb_df,
+                p.df_order,
+                p.df_lookahead,
+                freq_bins=self.freq_bins,
+                method=p.dfop_method,
+            )
+        )
+
+        self.run_df = run_df
+        if not run_df:
+            logger.warning("Runing without DF")
+        self.train_mask = train_mask
+        self.df_iter = p.df_n_iter
+
+    def forward(
+        self,
+        spec: Tensor,
+        feat_erb: Tensor,
+        feat_spec: Tensor,  # Not used, take spec modified by mask instead
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+        feat_spec = feat_spec.squeeze(1).permute(0, 3, 1, 2)
+
+        feat_erb = self.pad(feat_erb)
+        feat_spec = self.pad(feat_spec)
+        e0, e1, e2, e3, emb, c0, lsnr = self.enc(feat_erb, feat_spec)
+        m = self.erb_dec(emb, e3, e2, e1, e0)
+
+        spec = self.mask(spec, m)
+        df_coefs, df_alpha = self.df_dec(emb, c0)
+
+        for _ in range(self.df_iter):
+            spec = self.df_op(spec, df_coefs, df_alpha)
+        return spec, m, lsnr, df_alpha
diff --git a/DeepFilterNet/df/enhance.py b/DeepFilterNet/df/enhance.py
index d73356663..7a357e62d 100644
--- a/DeepFilterNet/df/enhance.py
+++ b/DeepFilterNet/df/enhance.py
@@ -15,7 +15,7 @@
 import df
 from df import config
 from df.checkpoint import load_model as load_model_cp
-from df.logger import init_logger
+from df.logger import init_logger, warn_once
 from df.model import ModelParams
 from df.modules import get_device
 from df.utils import as_complex, as_real, get_norm_alpha, resample
@@ -24,7 +24,10 @@
 
 def main(args):
     model, df_state, suffix = init_df(
-        args.model_base_dir, post_filter=args.pf, log_level=args.log_level
+        args.model_base_dir,
+        post_filter=args.pf,
+        log_level=args.log_level,
+        config_allow_defaults=True,
     )
     if args.output_dir is None:
         args.output_dir = "."
@@ -34,7 +37,9 @@ def main(args):
     for file in args.noisy_audio_files:
         audio, meta = load_audio(file, df_sr)
         t0 = time.time()
-        audio = enhance(model, df_state, audio, pad=args.compensate_delay)
+        audio = enhance(
+            model, df_state, audio, pad=args.compensate_delay, atten_lim_db=args.atten_lim
+        )
         t1 = time.time()
         t_audio = audio.shape[-1] / df_sr
         t = t1 - t0
@@ -50,7 +55,9 @@ def init_df(
     model_base_dir: Optional[str] = None,
     post_filter: bool = False,
     log_level: str = "INFO",
+    log_file: str = "enhance.log",
     config_allow_defaults: bool = False,
+    epoch: str = "best",
 ) -> Tuple[nn.Module, DF, str]:
     """Initializes and loads config, model and deep filtering state.
 
@@ -66,6 +73,13 @@ def init_df(
         suffix (str): Suffix based on the model name. This can be used for saving the enhanced
             audio.
     """
+    try:
+        from icecream import ic, install
+
+        ic.configureOutput(includeContext=True)
+        install()
+    except ImportError:
+        pass
     use_default_model = False
     if model_base_dir is None:
         use_default_model = True
@@ -74,16 +88,17 @@ def init_df(
         )
     if not os.path.isdir(model_base_dir):
         raise NotADirectoryError("Base directory not found at {}".format(model_base_dir))
-    init_logger(file=os.path.join(model_base_dir, "enhance.log"), level=log_level)
+    init_logger(file=os.path.join(model_base_dir, log_file), level=log_level, model=model_base_dir)
     if use_default_model:
         logger.info(f"Using default model at {model_base_dir}")
     config.load(
         os.path.join(model_base_dir, "config.ini"),
         config_must_exist=True,
         allow_defaults=config_allow_defaults,
+        allow_reload=True,
     )
     if post_filter:
-        config.set(ModelParams().section, "mask_pf", True, bool)
+        config.set("mask_pf", True, bool, ModelParams().section)
     p = ModelParams()
     df_state = DF(
         sr=p.sr,
@@ -93,7 +108,11 @@ def init_df(
         min_nb_erb_freqs=p.min_nb_freqs,
     )
     checkpoint_dir = os.path.join(model_base_dir, "checkpoints")
-    model, _ = load_model_cp(checkpoint_dir, df_state)
+    model, epoch = load_model_cp(checkpoint_dir, df_state, epoch=epoch)
+    if epoch is None:
+        logger.error("Could not find a checkpoint")
+        exit(1)
+    logger.debug(f"Loaded checkpoint from epoch {epoch}")
     model = model.to(get_device())
     # Set suffix to model name
     suffix = os.path.basename(os.path.abspath(model_base_dir))
@@ -138,8 +157,9 @@ def load_audio(file: str, sr: int, **kwargs) -> Tuple[Tensor, AudioMetaData]:
     info = ta.info(file, **ikwargs)
     audio, orig_sr = ta.load(file, **kwargs)
     if orig_sr != sr:
-        logger.warning(
-            f"Audio sampling rate does not match model sampling rate ({orig_sr}, {sr}). Resampling..."
+        warn_once(
+            f"Audio sampling rate does not match model sampling rate ({orig_sr}, {sr}). "
+            "Resampling..."
         )
         audio = resample(audio, orig_sr, sr)
     return audio, info
@@ -150,8 +170,9 @@ def save_audio(
     audio: Union[Tensor, ndarray],
     sr: int,
     output_dir: Optional[str] = None,
-    suffix: str = None,
+    suffix: Optional[str] = None,
     log: bool = False,
+    dtype=torch.int16,
 ):
     outpath = file
     if suffix is not None:
@@ -164,13 +185,15 @@ def save_audio(
     audio = torch.as_tensor(audio)
     if audio.ndim == 1:
         audio.unsqueeze_(0)
-    if audio.dtype != torch.int16:
+    if dtype == torch.int16 and audio.dtype != torch.int16:
         audio = (audio * (1 << 15)).to(torch.int16)
     ta.save(outpath, audio, sr)
 
 
 @torch.no_grad()
-def enhance(model: nn.Module, df_state: DF, audio: Tensor, pad=False):
+def enhance(
+    model: nn.Module, df_state: DF, audio: Tensor, pad=False, atten_lim_db: Optional[float] = None
+):
     p = ModelParams()
     model.eval()
     bs = audio.shape[0]
@@ -181,8 +204,12 @@ def enhance(model: nn.Module, df_state: DF, audio: Tensor, pad=False):
         # Pad audio to compensate for the delay due to the real-time STFT implementation
         audio = F.pad(audio, (0, p.fft_size))
     spec, erb_feat, spec_feat = df_features(audio, df_state, device=get_device())
-    spec = model(spec, erb_feat, spec_feat)[0].cpu()
-    audio = torch.as_tensor(df_state.synthesis(as_complex(spec.squeeze(1)).numpy()))
+    enhanced = model(spec, erb_feat, spec_feat)[0].cpu()
+    enhanced = as_complex(enhanced.squeeze(1))
+    if atten_lim_db is not None and abs(atten_lim_db) > 0:
+        lim = 10 ** (-abs(atten_lim_db) / 20)
+        enhanced = as_complex(spec.squeeze(1)) * lim + enhanced * (1 - lim)
+    audio = torch.as_tensor(df_state.synthesis(enhanced.numpy()))
     if pad:
         # The frame size is equal to p.hop_size. Given a new frame, the STFT loop requires e.g.
         # ceil((p.fft_size-p.hop_size)/p.hop_size). I.e. for 50% overlap, then p.hop_size=p.fft_size//2
@@ -194,6 +221,14 @@ def enhance(model: nn.Module, df_state: DF, audio: Tensor, pad=False):
     return audio
 
 
+def parse_epoch_type(value: str) -> Union[int, str]:
+    try:
+        return int(value)
+    except ValueError:
+        assert value in ("best", "latest")
+        return value
+
+
 def setup_df_argument_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -221,6 +256,21 @@ def setup_df_argument_parser() -> argparse.ArgumentParser:
         default="info",
         help="Logger verbosity. Can be one of (debug, info, error, none)",
     )
+    parser.add_argument("--debug", "-d", action="store_const", const="DEBUG", dest="log_level")
+    parser.add_argument(
+        "--atten-lim",
+        "-a",
+        type=int,
+        default=None,
+        help="Attenuation limit in dB by mixing the enhanced signal with the noisy signal.",
+    )
+    parser.add_argument(
+        "--epoch",
+        "-e",
+        default="best",
+        type=parse_epoch_type,
+        help="Epoch for checkpoint loading. Can be one of ['best', 'latest', <int>].",
+    )
     return parser
 
 
@@ -228,7 +278,7 @@ def setup_df_argument_parser() -> argparse.ArgumentParser:
     parser = setup_df_argument_parser()
     parser.add_argument(
         "--compensate-delay",
-        "-d",
+        "-D",
         action="store_true",
         help="Add some paddig to compensate the delay introduced by the real-time STFT/ISTFT implementation.",
     )
diff --git a/DeepFilterNet/df/evaluation_utils.py b/DeepFilterNet/df/evaluation_utils.py
new file mode 100644
index 000000000..05d932d06
--- /dev/null
+++ b/DeepFilterNet/df/evaluation_utils.py
@@ -0,0 +1,323 @@
+import os
+from abc import ABC, abstractmethod
+from collections import deque
+from functools import partial
+from multiprocessing.dummy import Pool as DummyPool
+from typing import Callable, Dict, List, Optional, Union
+
+import numpy as np
+import pystoi
+import torch
+import torch.multiprocessing as mp
+from loguru import logger
+from pesq import pesq
+from torch import Tensor
+from torch.multiprocessing.pool import Pool
+from torchaudio.functional import highpass_biquad
+from torchaudio.transforms import Resample
+
+from df.enhance import df_features, load_audio, save_audio
+from df.sepm import composite as composite_py
+from df.utils import as_complex, get_device, get_resample_params, resample
+from libdf import DF
+
+RESAMPLE_METHOD = "sinc_fast"
+
+try:
+    from tqdm import tqdm
+except ImportError:
+
+    def tqdm(iterable, total: Optional[int] = None, log_freq_percent=25, desc="Progress"):
+        assert 0 < log_freq_percent < 100
+        # tqdm not available using fallback
+        logged = set()
+        try:
+            L = iterable.__len__()
+        except AttributeError:
+            assert total is not None
+            L = total
+
+        for k, i in enumerate(iterable):
+            yield i
+            p = (k + 1) / L
+            progress = int(100 * p)
+            if progress % log_freq_percent == 0 and progress > 0:
+                if progress not in logged:
+                    logger.info("{}: {: >2d}%".format(desc, progress))
+                    logged.add(progress)
+
+
+@torch.no_grad()
+def enhance(model, df_state, audio, f_hp_cutoff: Optional[int] = None):
+    model.eval()
+    if hasattr(model, "reset_h0"):
+        model.reset_h0(batch_size=1, device=get_device())
+    spec, erb_feat, spec_feat = df_features(audio, df_state, get_device())
+    spec = model(spec, erb_feat, spec_feat)[0].squeeze(0)  # [C, T, F, 2]
+    audio = df_state.synthesis(as_complex(spec).cpu().numpy())
+    if f_hp_cutoff is not None:
+        audio = highpass_biquad(
+            torch.from_numpy(audio), df_state.sr(), cutoff_freq=f_hp_cutoff
+        ).numpy()
+    return audio
+
+
+def evaluation_loop(
+    df_state: DF,
+    model,
+    clean_files: List[str],
+    noisy_files: List[str],
+    metrics: List[str] = ["stoi", "composite", "sisdr"],
+    save_audio_callback: Optional[Callable[[str, Tensor], None]] = None,
+    n_workers: int = 4,
+) -> Dict[str, float]:
+    sr = df_state.sr()
+    metrics_dict = {
+        "stoi": partial(StoiMetric, sr=sr),
+        "sisdr": SiSDRMetric,
+        "composite": partial(CompositeMetric, sr=sr),
+        "composite-octave": partial(CompositeMetric, sr=sr, use_octave=True),
+        "pesq": partial(PesqMetric, sr=sr),
+    }
+    if n_workers >= 1:
+        pool_fn = mp.Pool
+    else:
+        pool_fn = DummyPool
+    with pool_fn(processes=max(1, n_workers)) as pool:
+        metrics: List[Metric] = [metrics_dict[m.lower()](pool=pool) for m in metrics]
+        for noisyfn, cleanfn in tqdm(zip(noisy_files, clean_files), total=len(noisy_files)):
+            noisy, _ = load_audio(noisyfn, sr)
+            clean, _ = load_audio(cleanfn, sr)
+            logger.debug(f"Processing {os.path.basename(noisyfn)}, {os.path.basename(cleanfn)}")
+            enh = enhance(model, df_state, noisy, 40)[0]
+            clean = df_state.synthesis(df_state.analysis(clean.numpy()))[0]
+            noisy = df_state.synthesis(df_state.analysis(noisy.numpy()))[0]
+            for m in metrics:
+                m.add(clean=clean, enhanced=enh, noisy=noisy)
+            if save_audio_callback is not None:
+                enh = torch.as_tensor(enh).to(torch.float32).view(1, -1)
+                save_audio_callback(cleanfn, enh)
+        logger.info("Waiting for metrics computation completion. This could take a few minutes.")
+        out_dict = {}
+        for m in metrics:
+            for k, v in m.mean().items():
+                out_dict[k] = v
+        return out_dict
+
+
+def stoi(clean, degraded, sr, extended=False):
+    assert len(clean.shape) == 1
+    if sr != 10000:
+        clean = resample(torch.as_tensor(clean), sr, 10000, method=RESAMPLE_METHOD).numpy()
+        degraded = resample(torch.as_tensor(degraded), sr, 10000, method=RESAMPLE_METHOD).numpy()
+        sr = 10000
+    stoi = pystoi.stoi(x=clean, y=degraded, fs_sig=sr, extended=extended)
+    return stoi
+
+
+def composite(
+    clean: Union[np.ndarray, Tensor], degraded: Union[np.ndarray, Tensor], sr: int, use_octave=False
+) -> np.ndarray:
+    """Compute pesq, csig, cbak, covl, ssnr"""
+    assert len(clean.shape) == 1, f"Input must be 1D array, but got input shape {clean.shape}"
+    if sr != 16000:
+        clean = resample(torch.as_tensor(clean), sr, 16000, method=RESAMPLE_METHOD).numpy()
+        degraded = resample(torch.as_tensor(degraded), sr, 16000, method=RESAMPLE_METHOD).numpy()
+        sr = 16000
+    if use_octave:
+        from tempfile import NamedTemporaryFile
+
+        import semetrics
+
+        with NamedTemporaryFile(suffix=".wav") as cf, NamedTemporaryFile(suffix=".wav") as nf:
+            # Note: Quantizing to int16 results in slightly modified metrics. Thus, keep f32 prec.
+            save_audio(cf.name, clean, sr, dtype=torch.float32)
+            save_audio(nf.name, degraded, sr, dtype=torch.float32)
+            c = semetrics.composite(cf.name, nf.name)
+    else:
+        c = composite_py(as_numpy(clean), as_numpy(degraded), sr)
+    return np.asarray(c)
+
+
+def si_sdr_speechmetrics(reference: np.ndarray, estimate: np.ndarray):
+    """This implementation is adopted from https://github.com/aliutkus/speechmetrics/blob/dde303e/speechmetrics/relative/sisdr.py"""
+    # as provided by @Jonathan-LeRoux and slightly adapted for the case of just one reference
+    # and one estimate.
+    # see original code here: https://github.com/sigsep/bsseval/issues/3#issuecomment-494995846
+    reference = reference.reshape(-1, 1)
+    estimate = estimate.reshape(-1, 1)
+    eps = np.finfo(reference.dtype).eps
+    Rss = np.dot(reference.T, reference)
+
+    # get the scaling factor for clean sources
+    a = (eps + np.dot(reference.T, estimate)) / (Rss + eps)
+
+    e_true = a * reference
+    e_res = estimate - e_true
+
+    Sss = (e_true**2).sum()
+    Snn = (e_res**2).sum()
+
+    sisdr = 10 * np.log10((eps + Sss) / (eps + Snn))
+    return sisdr
+
+
+class Metric(ABC):
+    def __init__(
+        self,
+        name: Union[str, List[str]],
+        source_sr: Optional[int] = None,
+        target_sr: Optional[int] = None,
+        device="cpu",
+    ):
+        self.name = name
+        self.sr = target_sr
+        self.resampler = None
+        if source_sr is not None and target_sr is not None and source_sr != target_sr:
+            params = get_resample_params(RESAMPLE_METHOD)
+            self.resampler = Resample(source_sr, target_sr, **params).to(device)
+        self.enh_values: Dict[str, List[float]] = (
+            {name: []} if isinstance(name, str) else {n: [] for n in name}
+        )
+        self.noisy_values: Dict[str, List[float]] = (
+            {name: []} if isinstance(name, str) else {n: [] for n in name}
+        )
+
+    @abstractmethod
+    def compute_metric(self, clean, degraded) -> Union[float, np.ndarray]:
+        pass
+
+    def _add_values_enh(self, values_enh: Union[float, np.ndarray]):
+        if isinstance(values_enh, float):
+            values_enh = np.asarray([values_enh])
+        for k, v in zip(self.enh_values.keys(), values_enh):
+            self.enh_values[k].append(v)
+
+    def _add_values_noisy(self, values_noisy: Union[float, np.ndarray]):
+        if isinstance(values_noisy, float):
+            values_noisy = np.asarray([values_noisy])
+        for k, v in zip(self.noisy_values.keys(), values_noisy):
+            self.noisy_values[k].append(v)
+
+    def maybe_resample(self, x) -> Tensor:
+        if self.resampler is not None:
+            x = self.resampler.forward(torch.as_tensor(x))
+        return x
+
+    def add(self, clean, enhanced, noisy):
+        clean = self.maybe_resample(clean).squeeze(0)
+        enhanced = self.maybe_resample(enhanced).squeeze(0)
+        values_enh = self.compute_metric(clean=clean, degraded=enhanced)
+        self._add_values_enh(values_enh)
+        if noisy is not None:
+            noisy = self.maybe_resample(noisy).squeeze(0)
+            values_noisy = self.compute_metric(clean=clean, degraded=enhanced)
+            self._add_values_noisy(values_noisy)
+
+    def mean(self) -> Dict[str, float]:
+        out = {}
+        for k in self.enh_values.keys():
+            if k in self.noisy_values and len(self.noisy_values[k]) > 0:
+                out[f"Noisy    {k}"] = np.mean(self.noisy_values[k])
+            out[f"Enhanced {k}"] = np.mean(self.enh_values[k])
+        return out
+
+
+# Multiprocessing Metric
+class MPMetric(Metric):
+    def __init__(
+        self,
+        name,
+        pool: Pool,
+        source_sr: Optional[int] = None,
+        target_sr: Optional[int] = None,
+    ):
+        super().__init__(name, source_sr=source_sr, target_sr=target_sr)
+        self.pool = pool
+        self.worker_results = deque()
+
+    def add(self, clean, enhanced, noisy):
+        clean = self.maybe_resample(torch.as_tensor(clean)).squeeze(0)
+        enhanced = self.maybe_resample(torch.as_tensor(enhanced)).squeeze(0)
+        h = self.pool.apply_async(
+            self.compute_metric,
+            (clean, enhanced),
+            callback=self._add_values_enh,
+            error_callback=logger.error,
+        )
+        h.get()
+        if noisy is not None:
+            noisy = self.maybe_resample(torch.as_tensor(noisy)).squeeze(0)
+            h = self.pool.apply_async(
+                self.compute_metric,
+                (clean, noisy),
+                callback=self._add_values_noisy,
+                error_callback=logger.error,
+            )
+            self.worker_results.append(h)
+
+    def mean(self) -> Dict[str, float]:
+        while len(self.worker_results) > 0:
+            h = self.worker_results.popleft()
+            h.get()
+        self.pool.close()
+        self.pool.join()
+        return super().mean()
+
+    def __getstate__(self):
+        self_dict = self.__dict__.copy()
+        del self_dict["pool"]
+        del self_dict["worker_results"]
+        return self_dict
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+
+class SiSDRMetric(MPMetric):
+    def __init__(self, pool: Pool):
+        super().__init__(name="SISDR", pool=pool)
+
+    def compute_metric(self, clean, degraded) -> float:
+        return si_sdr_speechmetrics(reference=as_numpy(clean), estimate=as_numpy(degraded))
+
+
+class StoiMetric(MPMetric):
+    def __init__(self, sr: int, pool: Pool):
+        super().__init__(name="STOI", pool=pool, source_sr=sr, target_sr=10000)
+
+    def compute_metric(self, clean, degraded) -> float:
+        assert self.sr is not None
+        return stoi(clean=as_numpy(clean), degraded=as_numpy(degraded), sr=self.sr)
+
+
+class PesqMetric(MPMetric):
+    def __init__(self, sr: int, pool: Pool):
+        super().__init__(name="PESQ", pool=pool, source_sr=sr, target_sr=16000)
+
+    def compute_metric(self, clean, degraded) -> Union[float, np.ndarray]:
+        assert self.sr is not None
+        return pesq(self.sr, as_numpy(clean), as_numpy(degraded), "wb")
+
+
+class CompositeMetric(MPMetric):
+    def __init__(self, sr: int, pool: Pool, use_octave: bool = False):
+        names = ["PESQ", "CSIG", "CBAK", "COVL", "SSNR"]
+        super().__init__(names, pool=pool, source_sr=sr, target_sr=16000)
+        self.use_octave = use_octave
+
+    def compute_metric(self, clean, degraded) -> Union[float, np.ndarray]:
+        assert self.sr is not None
+        c = composite(
+            clean=clean.squeeze(0),
+            degraded=degraded.squeeze(0),
+            sr=self.sr,
+            use_octave=self.use_octave,
+        )
+        return c
+
+
+def as_numpy(x) -> np.ndarray:
+    if isinstance(x, torch.Tensor):
+        return x.cpu().detach().numpy()
+    return x
diff --git a/DeepFilterNet/df/logger.py b/DeepFilterNet/df/logger.py
index 1d8319cb2..d1ea1951d 100644
--- a/DeepFilterNet/df/logger.py
+++ b/DeepFilterNet/df/logger.py
@@ -1,61 +1,166 @@
 import os
 import sys
-from typing import Dict, Optional
+from collections import defaultdict
+from copy import deepcopy
+from typing import Dict, Optional, Tuple
 
+import numpy as np
 import torch
 from loguru import logger
 from torch.types import Number
 
+from df.modules import GroupedLinearEinsum
 from df.utils import get_branch_name, get_commit_hash, get_device, get_host
 
 _logger_initialized = False
+WARN_ONCE_NO = logger.level("WARNING").no + 1
+DEPRECATED_NO = logger.level("WARNING").no + 2
 
 
-def init_logger(file: Optional[str] = None, level: str = "INFO"):
-    global _logger_initialized
+def init_logger(file: Optional[str] = None, level: str = "INFO", model: Optional[str] = None):
+    global _logger_initialized, _duplicate_filter
     if _logger_initialized:
         logger.debug("Logger already initialized.")
-        return
-    logger.remove()
-    level = level.upper()
-    if level != "NONE":
-        log_format = get_log_format(debug=level == "DEBUG")
-        logger.add(sys.stdout, level=level, format=log_format)
-        if file is not None:
-            logger.add(file, level=level, format=log_format)
-
-        logger.info(f"Running on torch {torch.__version__}")
-        logger.info(f"Running on host {get_host()}")
-        commit = get_commit_hash()
-        if commit is not None:
-            logger.info(f"Git commit: {commit}, branch: {get_branch_name()}")
-        if (jobid := os.getenv("SLURM_JOB_ID")) is not None:
-            logger.info(f"Slurm jobid: {jobid}")
+    else:
+        logger.remove()
+        level = level.upper()
+        if level.lower() != "none":
+            log_format = Formatter(debug=level == "DEBUG").format
+            logger.add(
+                sys.stdout,
+                level=level,
+                format=log_format,
+                filter=lambda r: r["level"].no not in {WARN_ONCE_NO, DEPRECATED_NO},
+            )
+            if file is not None:
+                logger.add(
+                    file,
+                    level=level,
+                    format=log_format,
+                    filter=lambda r: r["level"].no != WARN_ONCE_NO,
+                )
+
+            logger.info(f"Running on torch {torch.__version__}")
+            logger.info(f"Running on host {get_host()}")
+            commit = get_commit_hash()
+            if commit is not None:
+                logger.info(f"Git commit: {commit}, branch: {get_branch_name()}")
+            if (jobid := os.getenv("SLURM_JOB_ID")) is not None:
+                logger.info(f"Slurm jobid: {jobid}")
+            logger.level("WARNONCE", no=WARN_ONCE_NO, color="<yellow><bold>")
+            logger.add(
+                sys.stderr,
+                level=max(logger.level(level).no, WARN_ONCE_NO),
+                format=log_format,
+                filter=lambda r: r["level"].no == WARN_ONCE_NO and _duplicate_filter(r),
+            )
+            logger.level("DEPRECATED", no=DEPRECATED_NO, color="<yellow><bold>")
+            logger.add(
+                sys.stderr,
+                level=max(logger.level(level).no, DEPRECATED_NO),
+                format=log_format,
+                filter=lambda r: r["level"].no == DEPRECATED_NO and _duplicate_filter(r),
+            )
+    if model is not None:
+        logger.info("Loading model settings of {}", os.path.basename(model.rstrip("/")))
     _logger_initialized = True
 
 
-def get_log_format(debug=False):
-    if debug:
-        return (
-            "<green>{time:YYYY-MM-DD HH:mm:ss}</green>"
-            " | <level>{level: <8}</level>"
-            " | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan>"
-            " | <level>{message}</level>"
-        )
+def warn_once(message, *args, **kwargs):
+    logger.log("WARNONCE", message, *args, **kwargs)
+
+
+def log_deprecated(message, *args, **kwargs):
+    logger.log("DEPRECATED", message, *args, **kwargs)
+
+
+class Formatter:
+    def __init__(self, debug=False):
+        if debug:
+            self.fmt = (
+                "<green>{time:YYYY-MM-DD HH:mm:ss}</green>"
+                " | <level>{level: <8}</level>"
+                " | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan>"
+                " | <level>{message}</level>"
+            )
+        else:
+            self.fmt = (
+                "<green>{time:YYYY-MM-DD HH:mm:ss}</green>"
+                " | <level>{level: <8}</level>"
+                " | <cyan>DF</cyan>"
+                " | <level>{message}</level>"
+            )
+        self.fmt += "\n{exception}"
+
+    def format(self, record):
+        if record["level"].no == WARN_ONCE_NO:
+            return self.fmt.replace("{level: <8}", "WARNING ")
+        return self.fmt
+
+
+def _metrics_key(k_: Tuple[str, float]):
+    k0 = k_[0]
+    ks = k0.split("_")
+    if len(ks) > 2:
+        try:
+            return int(ks[-1])
+        except ValueError:
+            return 1000
+    elif k0 == "loss":
+        return -999
+    elif "loss" in k0.lower():
+        return -998
+    elif k0 == "lr":
+        return 998
+    elif k0 == "wd":
+        return 999
     else:
-        return (
-            "<green>{time:YYYY-MM-DD HH:mm:ss}</green>"
-            " | <level>{level: <8}</level>"
-            " | <cyan>DF</cyan>"
-            " | <level>{message}</level>"
-        )
+        return -101
 
 
 def log_metrics(prefix: str, metrics: Dict[str, Number]):
-    msg = prefix
-    for n, v in sorted(metrics.items()):
-        msg += f" | {n}: {v:.5g}"
-    logger.info(msg)
+    msg = ""
+    stages = defaultdict(str)
+    loss_msg = ""
+    for n, v in sorted(metrics.items(), key=_metrics_key):
+        if abs(v) > 1e-3:
+            m = f" | {n}: {v:.5f}"
+        else:
+            m = f" | {n}: {v:.3E}"
+        if "stage" in n:
+            s = n.split("stage_")[1].split("_snr")[0]
+            stages[s] += m.replace(f"stage_{s}_", "")
+        elif ("valid" in prefix or "test" in prefix) and "loss" in n.lower():
+            loss_msg += m
+        else:
+            msg += m
+    for s, msg_s in stages.items():
+        logger.info(f"{prefix} | stage {s}" + msg_s)
+    if len(stages) == 0:
+        logger.info(prefix + msg)
+    if len(loss_msg) > 0:
+        logger.info(prefix + loss_msg)
+
+
+class DuplicateFilter:
+    """
+    Filters away duplicate log messages.
+    Modified version of: https://stackoverflow.com/a/60462619
+    """
+
+    def __init__(self):
+        self.msgs = set()
+
+    def __call__(self, record) -> bool:
+        k = f"{record['level']}{record['message']}"
+        if k in self.msgs:
+            return False
+        else:
+            self.msgs.add(k)
+            return True
+
+
+_duplicate_filter = DuplicateFilter()
 
 
 def log_model_summary(model: torch.nn.Module, verbose=False):
@@ -77,11 +182,25 @@ def log_model_summary(model: torch.nn.Module, verbose=False):
     feat_spec = torch.randn([b, 1, t, p.nb_df, 2]).to(device)
 
     macs, params = ptflops.get_model_complexity_info(
-        model,
+        deepcopy(model),
         (t,),
         input_constructor=lambda _: {"spec": spec, "feat_erb": feat_erb, "feat_spec": feat_spec},
         as_strings=False,
         print_per_layer_stat=verbose,
         verbose=verbose,
+        custom_modules_hooks={
+            GroupedLinearEinsum: grouped_linear_flops_counter_hook,
+        },
     )
     logger.info(f"Model complexity: {params/1e6:.3f}M #Params, {macs/1e6:.1f}M MACS")
+
+
+def grouped_linear_flops_counter_hook(module: GroupedLinearEinsum, input, output):
+    # input: ([B, T, I],)
+    # output: [B, T, H]
+    input = input[0]  # [B, T, I]
+    output_last_dim = module.weight.shape[-1]
+    input = input.unflatten(-1, (module.groups, module.ws))  # [B, T, G, I/G]
+    # GroupedLinear calculates "...gi,...gih->...gh"
+    weight_flops = np.prod(input.shape) * output_last_dim
+    module.__flops__ += int(weight_flops)  # type: ignore
diff --git a/DeepFilterNet/df/loss.py b/DeepFilterNet/df/loss.py
index 7a2ec5867..90402c9bd 100644
--- a/DeepFilterNet/df/loss.py
+++ b/DeepFilterNet/df/loss.py
@@ -5,13 +5,12 @@
 import torch
 import torch.nn.functional as F
 from torch import Tensor, nn
-from torch.autograd import Function
 
 from df.config import Csv, config
 from df.model import ModelParams
 from df.modules import LocalSnrTarget, erb_fb
 from df.stoi import stoi
-from df.utils import as_complex
+from df.utils import angle, as_complex, get_device
 from libdf import DF
 
 
@@ -35,21 +34,6 @@ def iam(S: Tensor, X: Tensor, eps: float = 1e-10) -> Tensor:
     return (SS_mag / (XX_mag + eps)).clamp(0, 1)
 
 
-class angle(Function):
-    """Similar to torch.angle but robustify the gradient for zero magnitude."""
-
-    @staticmethod
-    def forward(ctx, x: Tensor):
-        ctx.save_for_backward(x)
-        return torch.atan2(x.imag, x.real)
-
-    @staticmethod
-    def backward(ctx, grad: Tensor):
-        (x,) = ctx.saved_tensors
-        grad_inv = grad / (x.real.square() + x.imag.square()).clamp_min_(1e-12)
-        return torch.view_as_complex(torch.stack((-x.imag * grad_inv, x.real * grad_inv), dim=-1))
-
-
 class Stft(nn.Module):
     def __init__(self, n_fft: int, hop: Optional[int] = None, window: Optional[Tensor] = None):
         super().__init__()
@@ -63,16 +47,19 @@ def __init__(self, n_fft: int, hop: Optional[int] = None, window: Optional[Tenso
         self.register_buffer("w", window)
 
     def forward(self, input: Tensor):
-        # Time-domain input shape: [B, T]
+        # Time-domain input shape: [B, *, T]
         t = input.shape[-1]
-        return torch.stft(
-            input.view(-1, t),
+        sh = input.shape[:-1]
+        out = torch.stft(
+            input.reshape(-1, t),
             n_fft=self.n_fft,
             hop_length=self.hop,
             window=self.w,
             normalized=True,
             return_complex=True,
         )
+        out = out.view(*sh, *out.shape[-2:])
+        return out
 
 
 class Istft(nn.Module):
@@ -86,18 +73,22 @@ def __init__(self, n_fft_inv: int, hop_inv: int, window_inv: Tensor):
         self.register_buffer("w_inv", window_inv)
 
     def forward(self, input: Tensor):
-        # Input shape: [B (C) T, F, (2)]
+        # Input shape: [B, * T, F, (2)]
         input = as_complex(input)
         t, f = input.shape[-2:]
+        sh = input.shape[:-2]
         # Even though this is not the DF implementation, it numerical sufficiently close.
         # Pad one extra step at the end to get original signal length
-        return torch.istft(
-            F.pad(input.view(-1, t, f).transpose(1, 2), (0, 1)),
+        out = torch.istft(
+            F.pad(input.reshape(-1, t, f).transpose(1, 2), (0, 1)),
             n_fft=self.n_fft_inv,
             hop_length=self.hop_inv,
             window=self.w_inv,
             normalized=True,
         )
+        if input.ndim > 2:
+            out = out.view(*sh, out.shape[-1])
+        return out
 
 
 class MultiResSpecLoss(nn.Module):
@@ -147,12 +138,7 @@ class SpectralLoss(nn.Module):
     f_m: Final[float]
     f_c: Final[float]
 
-    def __init__(
-        self,
-        gamma: float = 1,
-        factor_magnitude: float = 1,
-        factor_complex: float = 1,
-    ):
+    def __init__(self, gamma: float = 1, factor_magnitude: float = 1, factor_complex: float = 1):
         super().__init__()
         self.gamma = gamma
         self.f_m = factor_magnitude
@@ -314,6 +300,9 @@ def __init__(self):
     def forward(self, input: Tensor, target: Tensor):
         # Input shape: [B, T]
         eps = torch.finfo(input.dtype).eps
+        t = input.shape[-1]
+        target = target.reshape(-1, t)
+        input = input.reshape(-1, t)
         # Einsum for batch vector dot product
         Rss: Tensor = torch.einsum("bi,bi->b", target, target).unsqueeze(-1)
         a: Tensor = torch.einsum("bi,bi->b", target, input).add(eps).unsqueeze(-1) / Rss.add(eps)
@@ -341,9 +330,8 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor:
 
 
 class SegSdrLoss(nn.Module):
-    def __init__(
-        self, window_sizes: Optional[List[int]] = None, factor: float = 0.2, overlap: float = 0
-    ):
+    def __init__(self, window_sizes: List[int], factor: float = 0.2, overlap: float = 0):
+        # Window size in samples
         super().__init__()
         self.window_sizes = window_sizes
         self.factor = factor
@@ -355,21 +343,30 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor:
         if self.factor == 0:
             return torch.zeros((), device=input.device)
         loss = torch.zeros((), device=input.device)
-        if self.window_sizes is None:
-            # Non-segmental SdrLoss
-            loss = self.sdr(input, target).mean()
-        else:
-            for ws in self.window_sizes:
-                if ws > input.size(1):
-                    warnings.warn("Input size smaller than window size. Adjusting window size.")
-                    ws = input.size(1)
-                loss += self.sdr(
-                    input=input.unfold(1, ws, int(self.hop * ws)).reshape(-1, ws),
-                    target=target.unfold(1, ws, int(self.hop * ws)).reshape(-1, ws),
-                ).mean()
+        for ws in self.window_sizes:
+            if ws > input.size(-1):
+                warnings.warn(
+                    f"Input size {input.size(-1)} smaller than window size. Adjusting window size."
+                )
+                ws = input.size(1)
+            loss += self.sdr(
+                input=input.unfold(-1, ws, int(self.hop * ws)).reshape(-1, ws),
+                target=target.unfold(-1, ws, int(self.hop * ws)).reshape(-1, ws),
+            ).mean()
         return -loss * self.factor
 
 
+class LocalSnrLoss(nn.Module):
+    def __init__(self, factor: float = 1):
+        super().__init__()
+        self.factor = factor
+
+    def forward(self, input: Tensor, target_lsnr: Tensor):
+        # input (freq-domain): [B, T, 1]
+        input = input.squeeze(-1)
+        return F.mse_loss(input, target_lsnr) * self.factor
+
+
 class Loss(nn.Module):
     ml_f: Final[float]
     cal_f: Final[float]
@@ -387,7 +384,7 @@ def __init__(self, state: DF, istft: Optional[Istft] = None):
         self.store_losses = False
         self.summaries: Dict[str, List[Tensor]] = self.reset_summaries()
         # Mask Loss
-        self.ml_f = config("factor", 1, float, section="MaskLoss")
+        self.ml_f = config("factor", 0, float, section="MaskLoss")  # e.g. 1
         self.ml_gamma = config("gamma", 0.6, float, section="MaskLoss")
         self.ml_gamma_pred = config("gamma_pred", 0.6, float, section="MaskLoss")
         self.ml_f_under = config("f_under", 2, float, section="MaskLoss")
@@ -402,11 +399,11 @@ def __init__(self, state: DF, istft: Optional[Istft] = None):
             powers=[2, 4],
         )
         # DfAlphaLoss
-        self.cal_f = config("factor", 1e3, float, section="DfAlphaLoss")
+        self.cal_f = config("factor", 0, float, section="DfAlphaLoss")  # e.g. 1e3
         self.cal = DfAlphaLoss(self.cal_f) if self.cal_f > 0 else None
         # SpectralLoss
-        self.sl_fm = config("factor_magnitude", 1e6, float, section="SpectralLoss")
-        self.sl_fc = config("factor_complex", 1e6, float, section="SpectralLoss")
+        self.sl_fm = config("factor_magnitude", 0, float, section="SpectralLoss")  # e.g. 1e4
+        self.sl_fc = config("factor_complex", 0, float, section="SpectralLoss")
         self.sl_gamma = config("gamma", 1, float, section="SpectralLoss")
         self.sl_f = self.sl_fm + self.sl_fc
         if self.sl_f > 0:
@@ -425,6 +422,17 @@ def __init__(self, state: DF, istft: Optional[Istft] = None):
             self.mrsl = MultiResSpecLoss(self.mrsl_ffts, self.mrsl_gamma, self.mrsl_f, self.mrsl_fc)
         else:
             self.mrsl = None
+        self.sdrl_f = config("factor", 0, float, section="SdrLoss")
+        self.sdrl = None
+        if self.sdrl_f > 0:
+            sdr_sgemental_ws = config("segmental_ws", [], Csv(int), section="SdrLoss")
+            if len(sdr_sgemental_ws) > 0 and any(ws > 0 for ws in sdr_sgemental_ws):
+                self.sdrl = SegSdrLoss(sdr_sgemental_ws, factor=self.sdrl_f)
+            else:
+                self.sdrl = SdrLoss(self.sdrl_f)
+        self.lsnr_f = config("factor", 0.0005, float, section="LocalSnrLoss")
+        self.lsnrl = LocalSnrLoss(self.lsnr_f) if self.lsnr_f > 0 else None
+        self.dev_str = get_device().type
 
     def forward(
         self,
@@ -436,6 +444,7 @@ def forward(
         df_alpha: Optional[Tensor],
         snrs: Tensor,
         max_freq: Optional[Tensor] = None,
+        multi_stage_specs: List[Tensor] = [],
     ):
         max_bin: Optional[Tensor] = None
         if max_freq is not None:
@@ -446,24 +455,62 @@ def forward(
             ).long()
         enhanced_td = None
         clean_td = None
+        multi_stage = None
+        multi_stage_td = None
+        if multi_stage_specs:
+            # Stack spectrograms in a channel dimension
+            multi_stage = as_complex(torch.stack(multi_stage_specs, dim=1))
+        lsnr_gt = self.lsnr(clean, noise=noisy - clean)
         if self.istft is not None:
-            if self.store_losses or self.mrsl is not None:
+            if self.store_losses or self.mrsl is not None or self.sdrl is not None:
                 enhanced_td = self.istft(enhanced)
                 clean_td = self.istft(clean)
+                if multi_stage is not None:
+                    # leave out erb enhanced
+                    multi_stage_td = self.istft(multi_stage)
 
-        ml, sl, mrsl, cal = [torch.zeros((), device=clean.device)] * 4
+        ml, sl, mrsl, cal, sdrl, lsnrl = [torch.zeros((), device=clean.device)] * 6
         if self.ml_f != 0 and self.ml is not None:
             ml = self.ml(input=mask, clean=clean, noisy=noisy, max_bin=max_bin)
         if self.sl_f != 0 and self.sl is not None:
-            sl = self.sl(input=enhanced, target=clean)
+            sl = torch.zeros((), device=clean.device)
+            if multi_stage is not None:
+                sl += self.sl(input=multi_stage, target=clean.expand_as(multi_stage))
+            else:
+                sl = self.sl(input=enhanced, target=clean)
         if self.mrsl_f > 0 and self.mrsl is not None:
-            mrsl = self.mrsl(enhanced_td, clean_td)
+            if multi_stage_td is not None:
+                ms = multi_stage_td[:, 1:]
+                mrsl = self.mrsl(ms, clean_td.expand_as(ms))
+            else:
+                mrsl = self.mrsl(enhanced_td, clean_td)
+        if self.lsnr_f != 0:
+            lsnrl = self.lsnrl(input=lsnr, target_lsnr=lsnr_gt)
         if self.cal_f != 0 and self.cal is not None and df_alpha is not None:
             lsnr_gt = self.lsnr(clean, noise=noisy - clean, max_bin=self.nb_df)
             cal = self.cal(df_alpha, target_lsnr=lsnr_gt)
+        if self.sdrl_f != 0:
+            if multi_stage_td is not None:
+                ms = multi_stage_td[:, 1:]
+                sdrl = self.sdrl(ms, clean_td.expand_as(ms))
+            else:
+                sdrl = self.sdrl(enhanced_td, clean_td)
         if self.store_losses and self.istft is not None:
-            self.store_summaries(enhanced_td, clean_td, snrs, ml, sl, mrsl, cal)  # type: ignore
-        return ml + sl + mrsl + cal
+            assert enhanced_td is not None
+            assert clean_td is not None
+            self.store_summaries(
+                enhanced_td,
+                clean_td,
+                snrs,
+                ml,
+                sl,
+                mrsl,
+                sdrl,
+                lsnrl,
+                cal,
+                multi_stage_td=multi_stage_td,
+            )
+        return ml + sl + mrsl + sdrl + lsnrl + cal
 
     def reset_summaries(self):
         self.summaries = defaultdict(list)
@@ -483,26 +530,49 @@ def store_summaries(
         ml: Tensor,
         sl: Tensor,
         mrsl: Tensor,
+        sdrl: Tensor,
+        lsnrl: Tensor,
         cal: Tensor,
+        multi_stage_td: Optional[Tensor] = None,
     ):
         if ml != 0:
             self.summaries["MaskLoss"].append(ml.detach())
         if sl != 0:
             self.summaries["SpectralLoss"].append(sl.detach())
-        if cal != 0:
-            self.summaries["DfAlphaLoss"].append(cal.detach())
         if mrsl != 0:
             self.summaries["MultiResSpecLoss"].append(mrsl.detach())
+        if sdrl != 0:
+            self.summaries["SdrLoss"].append(sdrl.detach())
+        if cal != 0:
+            self.summaries["DfAlphaLoss"].append(cal.detach())
+        if lsnrl != 0:
+            self.summaries["LocalSnrLoss"].append(lsnrl.detach())
         sdr = SiSdr()
-        sdr_vals: Tensor = sdr(enh_td.detach(), target=clean_td.detach())
-        stoi_vals: Tensor = stoi(y=enh_td.detach(), x=clean_td.detach(), fs_source=self.sr)
+        enh_td = enh_td.squeeze(1).detach()
+        clean_td = clean_td.squeeze(1).detach()
+        sdr_vals: Tensor = sdr(enh_td, target=clean_td)
+        stoi_vals: Tensor = stoi(y=enh_td, x=clean_td, fs_source=self.sr)
+        sdr_vals_ms, stoi_vals_ms = [], []
+        if multi_stage_td is not None:
+            for i in range(multi_stage_td.shape[1]):
+                sdr_vals_ms.append(sdr(multi_stage_td[:, i].detach(), clean_td))
+                stoi_vals_ms.append(
+                    stoi(y=multi_stage_td[:, i].detach(), x=clean_td, fs_source=self.sr)
+                )
         for snr in torch.unique(snrs, sorted=False):
             self.summaries[f"sdr_snr_{snr.item()}"].extend(
-                sdr_vals.masked_select(snr == snrs).split(1)
+                sdr_vals.masked_select(snr == snrs).detach().split(1)
             )
             self.summaries[f"stoi_snr_{snr.item()}"].extend(
-                stoi_vals.masked_select(snr == snrs).split(1)
+                stoi_vals.masked_select(snr == snrs).detach().split(1)
             )
+            for i, (sdr_i, stoi_i) in enumerate(zip(sdr_vals_ms, stoi_vals_ms)):
+                self.summaries[f"sdr_stage_{i}_snr_{snr.item()}"].extend(
+                    sdr_i.masked_select(snr == snrs).detach().split(1)
+                )
+                self.summaries[f"stoi_stage_{i}_snr_{snr.item()}"].extend(
+                    stoi_i.masked_select(snr == snrs).detach().split(1)
+                )
 
 
 def test_local_snr():
diff --git a/DeepFilterNet/df/lr.py b/DeepFilterNet/df/lr.py
index 6d9c0fb36..eea69146f 100644
--- a/DeepFilterNet/df/lr.py
+++ b/DeepFilterNet/df/lr.py
@@ -1,65 +1,63 @@
 import math
 
-
-def cosine_lr_scheduler(
-    lr: float,
-    n_warmup: int = 0,
-    warmup_init_lr: float = -1,
-    min_lr: float = 0.0,
-    t_mult: float = 1.0,
-    lr_period_updates: float = -1,
-    lr_shrink: float = 0.1,
-    max_update: int = -1,
+import numpy as np
+
+
+def cosine_scheduler(
+    base_value: float,
+    final_value: float,
+    epochs: int,
+    niter_per_ep: int,
+    warmup_epochs: int = 0,
+    start_warmup_value: float = 0,
+    warmup_steps: int = -1,
+    initial_ep_per_cycle: float = -1,
+    cycle_decay: float = 1,
+    cycle_mul: float = 1,
 ):
-    """Cosine annealing learning rate scheduler with warmup and step decay for LambdaLR.
-
-    Based on fairseq.optim.lr_scheduler.cosine_lr_scheduler.
-
-    Args:
-        lr (float): (Maximum) learning rate.
-        n_warmup (int): Number of warmup steps with a linear lr increase. Default is 0.
-        warmup_init_lr (float): Initial learning rate during warmup phase. Default is `lr`.
-        min_lr (float): Minimum learning rate during cosine annealing. Default is 0.
-        t_mult (float): Factor to grow the length of each period. Default is 1.
-        lr_period_updates (float): Initial number of updates per period.
-        lr_shrink (float): Shrink factor for each period. Default 0.1.
-        max_update (int): Number of maximum updates (epochs). If specified, will result in 1 period
-            over all updates.
-    """
-    max_lr_base = lr
-    min_lr_base = min_lr
-    warmup_end_lr = max_lr_base
-    warmup_init_lr = min_lr if warmup_init_lr < 0 else warmup_init_lr
-    period = lr_period_updates
-    if period <= 0:
-        assert max_update > 0, "Either lr_period_updates or max_update must be set."
-        period = max_update - n_warmup
-    if n_warmup > 0:
-        step_lr = (warmup_end_lr - warmup_init_lr) / n_warmup
+    """Adopted from official ConvNeXt repo."""
+    warmup_schedule = np.array([])
+    warmup_iters = warmup_epochs * niter_per_ep
+    if warmup_steps > 0:
+        warmup_iters = warmup_steps
+    if warmup_epochs > 0:
+        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
+
+    iters_after_warmup = epochs * niter_per_ep - warmup_iters
+    if initial_ep_per_cycle == -1:
+        initial_ep_per_cycle = iters_after_warmup
+        num_cycles = 1
+        cycle_lengths = [iters_after_warmup]
     else:
-        step_lr = 1
-    lr_shrink_base = lr_shrink
-
-    def step(epoch: int) -> float:
-        if epoch < n_warmup:
-            return (warmup_init_lr + epoch * step_lr) / max_lr_base
-        cur_updates = epoch - n_warmup
-
-        if t_mult != 1:
-            i = math.floor(math.log(1 - cur_updates / period * (1 - t_mult), t_mult))
-            t_i = t_mult**i * period
-            t_cur = cur_updates - (1 - t_mult**i) / (1 - t_mult) * period
+        initial_cycle_iter = int(round(initial_ep_per_cycle * niter_per_ep))
+        if cycle_mul == 1:
+            num_cycles = int(math.ceil(iters_after_warmup / (initial_ep_per_cycle * niter_per_ep)))
+            cycle_lengths = [initial_cycle_iter] * num_cycles
         else:
-            i = math.floor(cur_updates / period)
-            t_i = period
-            t_cur = cur_updates - (period * i)
-
-        lr_shrink = lr_shrink_base**i
-        min_lr = min_lr_base * lr_shrink
-        max_lr = max_lr_base * lr_shrink
-
-        return (
-            min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * t_cur / t_i))
-        ) / max_lr_base
-
-    return step
+            num_cycles = 0
+            cycle_lengths = []
+            i = 0
+            while sum(cycle_lengths) < iters_after_warmup:
+                num_cycles += 1
+                cycle_lengths.append(initial_cycle_iter * cycle_mul**i)
+                i += 1
+    schedule_cycles = []
+    for i in range(num_cycles):
+        cycle_base_value = base_value * cycle_decay**i
+        iters = np.arange(cycle_lengths[i])
+        schedule = np.array(
+            [
+                final_value
+                + 0.5
+                * (cycle_base_value - final_value)
+                * (1 + math.cos(math.pi * i / (len(iters))))
+                for i in iters
+            ]
+        )
+        schedule_cycles.append(schedule)
+
+    schedule = np.concatenate((warmup_schedule, *schedule_cycles))
+    schedule = schedule[: epochs * niter_per_ep]
+
+    assert len(schedule) == epochs * niter_per_ep
+    return schedule
diff --git a/DeepFilterNet/df/model.py b/DeepFilterNet/df/model.py
index 8fddacbc4..00328a21a 100644
--- a/DeepFilterNet/df/model.py
+++ b/DeepFilterNet/df/model.py
@@ -1,5 +1,6 @@
 from importlib import import_module
 
+import torch
 from loguru import logger
 
 from df.config import DfParams, config
@@ -18,4 +19,6 @@ def init_model(*args, **kwargs):
     """Initialize the model specified in the config."""
     model = config("MODEL", default="deepfilternet", section="train")
     logger.info(f"Initializing model `{model}`")
-    return getattr(import_module("df." + model), "init_model")(*args, **kwargs)
+    model = getattr(import_module("df." + model), "init_model")(*args, **kwargs)
+    model.to(memory_format=torch.channels_last)
+    return model
diff --git a/DeepFilterNet/df/modules.py b/DeepFilterNet/df/modules.py
index 205e815f8..11fb8b07d 100644
--- a/DeepFilterNet/df/modules.py
+++ b/DeepFilterNet/df/modules.py
@@ -1,10 +1,13 @@
+import math
 from collections import OrderedDict
-from typing import List, Optional, Tuple
+from typing import Callable, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
 from torch import Tensor, nn
 from torch.nn import functional as F
+from torch.nn import init
+from torch.nn.parameter import Parameter
 from typing_extensions import Final
 
 from df.model import ModelParams
@@ -12,6 +15,117 @@
 from libdf import unit_norm_init
 
 
+class Conv2dNormAct(nn.Sequential):
+    def __init__(
+        self,
+        in_ch: int,
+        out_ch: int,
+        kernel_size: Union[int, Iterable[int]],
+        fstride: int = 1,
+        dilation: int = 1,
+        fpad: bool = True,
+        bias: bool = True,
+        separable: bool = False,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+    ):
+        """Causal Conv2d by delaying the signal for any lookahead.
+
+        Expected input format: [B, C, T, F]
+        """
+        lookahead = 0  # This needs to be handled on the input feature side
+        # Padding on time axis
+        kernel_size = (
+            (kernel_size, kernel_size) if isinstance(kernel_size, int) else tuple(kernel_size)
+        )
+        if fpad:
+            fpad_ = kernel_size[1] // 2 + dilation - 1
+        else:
+            fpad_ = 0
+        pad = (0, 0, kernel_size[0] - 1 - lookahead, lookahead)
+        layers = []
+        if any(x > 0 for x in pad):
+            layers.append(nn.ConstantPad2d(pad, 0.0))
+        groups = math.gcd(in_ch, out_ch) if separable else 1
+        if groups == 1:
+            separable = False
+        if max(kernel_size) == 1:
+            separable = False
+        layers.append(
+            nn.Conv2d(
+                in_ch,
+                out_ch,
+                kernel_size=kernel_size,
+                padding=(0, fpad_),
+                stride=(1, fstride),  # Stride over time is always 1
+                dilation=(1, dilation),  # Same for dilation
+                groups=groups,
+                bias=bias,
+            )
+        )
+        if separable:
+            layers.append(nn.Conv2d(out_ch, out_ch, kernel_size=1, bias=False))
+        if norm_layer is not None:
+            layers.append(norm_layer(out_ch))
+        if activation_layer is not None:
+            layers.append(activation_layer())
+        super().__init__(*layers)
+
+
+class ConvTranspose2dNormAct(nn.Sequential):
+    def __init__(
+        self,
+        in_ch: int,
+        out_ch: int,
+        kernel_size: Union[int, Tuple[int, int]],
+        fstride: int = 1,
+        dilation: int = 1,
+        fpad: bool = True,
+        bias: bool = True,
+        separable: bool = False,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+    ):
+        """Causal ConvTranspose2d.
+
+        Expected input format: [B, C, T, F]
+        """
+        # Padding on time axis, with lookahead = 0
+        lookahead = 0  # This needs to be handled on the input feature side
+        kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
+        if fpad:
+            fpad_ = kernel_size[1] // 2
+        else:
+            fpad_ = 0
+        pad = (0, 0, kernel_size[0] - 1 - lookahead, lookahead)
+        layers = []
+        if any(x > 0 for x in pad):
+            layers.append(nn.ConstantPad2d(pad, 0.0))
+        groups = math.gcd(in_ch, out_ch) if separable else 1
+        if groups == 1:
+            separable = False
+        layers.append(
+            nn.ConvTranspose2d(
+                in_ch,
+                out_ch,
+                kernel_size=kernel_size,
+                padding=(kernel_size[0] - 1, fpad_ + dilation - 1),
+                output_padding=(0, fpad_),
+                stride=(1, fstride),  # Stride over time is always 1
+                dilation=(1, dilation),
+                groups=groups,
+                bias=bias,
+            )
+        )
+        if separable:
+            layers.append(nn.Conv2d(out_ch, out_ch, kernel_size=1, bias=False))
+        if norm_layer is not None:
+            layers.append(norm_layer(out_ch))
+        if activation_layer is not None:
+            layers.append(activation_layer())
+        super().__init__(*layers)
+
+
 def convkxf(
     in_ch: int,
     out_ch: Optional[int] = None,
@@ -414,6 +528,10 @@ def __init__(
             (nn.GRU(self.input_size, self.hidden_size, **kwargs) for _ in range(groups))
         )
 
+    def flatten_parameters(self):
+        for layer in self.layers:
+            layer.flatten_parameters()
+
     def get_h0(self, batch_size: int = 1, device: torch.device = torch.device("cpu")):
         return torch.zeros(
             self.groups * self.num_directions,
@@ -477,6 +595,8 @@ def __init__(
         }
         assert input_size % groups == 0
         assert hidden_size % groups == 0
+        assert num_layers > 0
+        self.input_size = input_size
         self.groups = groups
         self.num_layers = num_layers
         self.batch_first = batch_first
@@ -491,6 +611,11 @@ def __init__(
         self.grus.append(GroupedGRULayer(input_size, hidden_size, **kwargs))
         for _ in range(1, num_layers):
             self.grus.append(GroupedGRULayer(hidden_size, hidden_size, **kwargs))
+        self.flatten_parameters()
+
+    def flatten_parameters(self):
+        for gru in self.grus:
+            gru.flatten_parameters()
 
     def get_h0(self, batch_size: int, device: torch.device = torch.device("cpu")) -> Tensor:
         return torch.zeros(
@@ -523,6 +648,77 @@ def forward(self, input: Tensor, state: Optional[Tensor] = None) -> Tuple[Tensor
         return output, outstate
 
 
+class SqueezedGRU(nn.Module):
+    input_size: Final[int]
+    hidden_size: Final[int]
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        output_size: Optional[int] = None,
+        num_layers: int = 1,
+        linear_groups: int = 8,
+        batch_first: bool = True,
+        gru_skip_op: Optional[Callable[..., torch.nn.Module]] = None,
+        linear_act_layer: Callable[..., torch.nn.Module] = nn.Identity,
+    ):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.linear_in = nn.Sequential(
+            GroupedLinearEinsum(input_size, hidden_size, linear_groups), linear_act_layer()
+        )
+        self.gru = nn.GRU(hidden_size, hidden_size, num_layers=num_layers, batch_first=batch_first)
+        self.gru_skip = gru_skip_op() if gru_skip_op is not None else None
+        if output_size is not None:
+            self.linear_out = nn.Sequential(
+                GroupedLinearEinsum(hidden_size, output_size, linear_groups), linear_act_layer()
+            )
+        else:
+            self.linear_out = nn.Identity()
+
+    def forward(self, input: Tensor, h=None) -> Tuple[Tensor, Tensor]:
+        input = self.linear_in(input)
+        x, h = self.gru(input, h)
+        if self.gru_skip is not None:
+            x = x + self.gru_skip(input)
+        x = self.linear_out(x)
+        return x, h
+
+
+class GroupedLinearEinsum(nn.Module):
+    input_size: Final[int]
+    hidden_size: Final[int]
+    groups: Final[int]
+
+    def __init__(self, input_size: int, hidden_size: int, groups: int = 1):
+        super().__init__()
+        # self.weight: Tensor
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.groups = groups
+        assert input_size % groups == 0
+        self.ws = input_size // groups
+        self.register_parameter(
+            "weight",
+            Parameter(
+                torch.zeros(groups, input_size // groups, hidden_size // groups), requires_grad=True
+            ),
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))  # type: ignore
+
+    def forward(self, x: Tensor) -> Tensor:
+        # x: [..., I]
+        x = x.unflatten(-1, (self.groups, self.ws))  # [..., G, I/G]
+        x = torch.einsum("...gi,...gih->...gh", x, self.weight)  # [..., G, H/G]
+        x = x.flatten(2, 3)  # [B, T, H]
+        return x
+
+
 class GroupedLinear(nn.Module):
     input_size: Final[int]
     hidden_size: Final[int]
diff --git a/DeepFilterNet/df/multistagenet.py b/DeepFilterNet/df/multistagenet.py
new file mode 100644
index 000000000..113cf02ae
--- /dev/null
+++ b/DeepFilterNet/df/multistagenet.py
@@ -0,0 +1,577 @@
+import math
+from functools import partial, reduce
+from typing import Callable, Final, List, Optional, Tuple
+
+import torch
+from loguru import logger
+from torch import Tensor, nn
+from torch.nn import init
+from torch.nn.parameter import Parameter
+
+from df.config import Csv, DfParams, config
+from df.modules import Conv2dNormAct, ConvTranspose2dNormAct, GroupedGRU, Mask, erb_fb
+from df.utils import angle_re_im, get_device
+from libdf import DF
+
+
+class ModelParams(DfParams):
+    section = "multistagenet"
+
+    def __init__(self):
+        super().__init__()
+        self.stages: List[int] = config(
+            "STAGES", cast=Csv(int), default=(3, 3, 9, 3), section=self.section  # type: ignore
+        )
+        self.conv_lookahead: int = config(
+            "CONV_LOOKAHEAD", cast=int, default=0, section=self.section
+        )
+        self.erb_widths: List[int] = config(
+            "ERB_WIDTHS", cast=Csv(int), default=[16, 16, 16], section=self.section
+        )
+        self.erb_hidden_dim: int = config(
+            "ERB_HIDDEN_DIM", cast=int, default=64, section=self.section
+        )
+        self.refinement_widths: List[int] = config(
+            "REFINEMENT_WIDTHS", cast=Csv(int), default=[32, 32, 32, 32], section=self.section
+        )
+        self.refinement_hidden_dim: int = config(
+            "REFINEMENT_HIDDEN_DIM", cast=int, default=96, section=self.section
+        )
+        self.refinement_act: str = (
+            config("REFINEMENT_OUTPUT_ACT", default="identity", section=self.section)
+            .lower()
+            .replace("none", "identity")
+        )
+        self.refinement_op: str = config(
+            "REFINEMENT_OP", default="mul", section=self.section
+        ).lower()
+        self.refinement_out_layer: str = config(
+            "REFINEMENT_OUTPUT_LAYER", default="LocallyConnected", section=self.section
+        )
+        self.mask_pf: bool = config("MASK_PF", cast=bool, default=False, section=self.section)
+
+
+# class GruSE(nn.Module):
+#     """GRU with previous adaptive avg pooling like SqueezeExcitation"""
+
+#     avg_dim: Final[int]
+
+#     def __init__(
+#         self,
+#         input_dim: int,
+#         hidden_dim: int,
+#         groups: int = 1,
+#         reduce_: str = "frequencies",
+#         skip: Optional[Callable[..., torch.nn.Module]] = nn.Identity,
+#         scale_activation: Optional[Callable[..., torch.nn.Module]] = None,
+#     ):
+#         super().__init__()
+#         assert reduce_ in ("channels", "frequencies", "none")
+#         if reduce_ == "channels":
+#             self.avg_dim = 1
+#         else:
+#             self.avg_dim = 3
+#         if groups == 1:
+#             self.gru = nn.GRU(input_dim, hidden_dim)
+#         else:
+#             self.gru = GroupedGRU(input_dim, hidden_dim, groups=groups)
+#         assert (
+#             skip or scale_activation is None
+#         ), "Can only either use a skip connection or SqueezeExcitation with `scale_activation`"
+#         self.fc = nn.Linear(hidden_dim, input_dim)
+#         self.skip = skip() if skip is not None else None
+#         self.scale = scale_activation() if scale_activation is not None else None
+
+#     def forward(self, input: Tensor, h=None) -> Tuple[Tensor, Tensor]:
+#         # x: [B, C, T, F]
+#         if self.avg_dim == 1:
+#             x = input.mean(dim=self.avg_dim)  # [B, T, C]
+#         else:
+#             x = input.mean(dim=self.avg_dim).transpose(1, 2)  # [B, T, C]
+#         x, h = self.gru(x, h)
+#         x = self.fc(x)
+#         if self.avg_dim == 1:
+#             x = x.unsqueeze(1)
+#         else:
+#             x = x.transpose(1, 2).unsqueeze(-1)
+#         if self.skip is not None:
+#             x = self.skip(input) + x  # a regular skip connection
+#         elif self.scale is not None:
+#             x = input * self.scale(x)  # like in SqueezeExcitation
+#         return x, h
+
+
+class LSNRNet(nn.Module):
+    lsnr_scale: Final[int]
+    lsnr_offset: Final[int]
+
+    def __init__(
+        self, in_ch: int, hidden_dim: int = 16, fstride=2, lsnr_min: int = -15, lsnr_max: int = 40
+    ):
+        super().__init__()
+        self.conv = Conv2dNormAct(in_ch, in_ch, kernel_size=(1, 3), fstride=fstride, separable=True)
+        self.gru_snr = nn.GRU(in_ch, hidden_dim)
+        self.fc_snr = nn.Sequential(nn.Linear(hidden_dim, 1), nn.Sigmoid())
+        self.lsnr_scale = lsnr_max - lsnr_min
+        self.lsnr_offset = lsnr_min
+
+    def forward(self, x: Tensor, h=None) -> Tuple[Tensor, Tensor]:
+        x = self.conv(x)
+        x, h = self.gru_snr(x.mean(-1).transpose(1, 2), h)
+        x = self.fc_snr(x) * self.lsnr_scale + self.lsnr_offset
+        return x, h
+
+
+class LocalLinearCF(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, n_freqs: int, bias: bool = True):
+        super().__init__()
+        self.n_freqs = n_freqs
+        self.register_parameter(
+            "weight", Parameter(torch.zeros(in_ch, out_ch, n_freqs), requires_grad=True)
+        )
+        if bias:
+            self.bias: Optional[Tensor]
+            self.register_parameter(
+                "bias", Parameter(torch.zeros(out_ch, 1, n_freqs), requires_grad=True)
+            )
+        else:
+            self.bias = None
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))  # type: ignore
+        if self.bias is not None:
+            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # x: [B, Ci, T, F]
+        x = torch.einsum("bctf,cof->botf", x, self.weight)  # [B, Co, T, F]
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+
+class GroupedLinearCF(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, n_freqs: int, n_groups: int, bias: bool = True):
+        super().__init__()
+        # self.weight: Tensor
+        self.n_freqs = n_freqs
+        n_groups = n_groups if n_groups > 0 else n_freqs
+        if n_groups == n_freqs:
+            logger.warning("Use more performant LocallyConnected since they are equivalent now.")
+        assert (
+            n_freqs % n_groups == 0
+        ), "Number of frequencies must be dividable by the number of groups"
+        self.n_groups = n_groups
+        self.n_unfold = n_freqs // n_groups
+        self.register_parameter(
+            "weight",
+            Parameter(
+                torch.zeros(n_groups, n_freqs // n_groups, in_ch, out_ch), requires_grad=True
+            ),
+        )
+        if bias:
+            self.bias: Optional[Tensor]
+            self.register_parameter(
+                "bias", Parameter(torch.zeros(out_ch, 1, n_freqs), requires_grad=True)
+            )
+        else:
+            self.bias = None
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))  # type: ignore
+        if self.bias is not None:
+            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # x: [B, Ci, T, F]
+        x = x.permute(0, 2, 3, 1)  # [B, T, F, Ci]
+        x = x.unflatten(2, (self.n_groups, self.n_unfold))  # [B, T, G, F/G, Ci]
+        x = torch.einsum("btfgi,fgio->btfgo", x, self.weight)  # [B, T, G, F/G, Co]
+        x = x.flatten(2, 3)  # [B, T, F, Co]
+        x = x.permute(0, 3, 1, 2)  # [B, Co, T, F]
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+
+class GroupedGRULayerMS(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, n_freqs: int, n_groups: int, bias: bool = True):
+        super().__init__()
+        assert n_freqs % n_groups == 0
+        self.n_freqs = n_freqs
+        self.g_freqs = n_freqs // n_groups
+        self.n_groups = n_groups
+        self.out_ch = self.g_freqs * out_ch
+        self._in_ch = in_ch
+        self.input_size = self.g_freqs * in_ch
+        # self.weight_ih_l: Tensor
+        self.register_parameter(
+            "weight_ih_l",
+            Parameter(torch.zeros(n_groups, 3 * self.out_ch, self.input_size), requires_grad=True),
+        )
+        # self.weight_hh_l: Tensor
+        self.register_parameter(
+            "weight_hh_l",
+            Parameter(torch.zeros(n_groups, 3 * self.out_ch, self.out_ch), requires_grad=True),
+        )
+        if bias:
+            # self.bias_ih_l: Tensor
+            self.register_parameter(
+                "bias_ih_l", Parameter(torch.zeros(n_groups, 3 * self.out_ch), requires_grad=True)
+            )
+            # self.bias_hh_l: Tensor
+            self.register_parameter(
+                "bias_hh_l", Parameter(torch.zeros(n_groups, 3 * self.out_ch), requires_grad=True)
+            )
+        else:
+            self.bias_ih_l = None  # type: ignore
+            self.bias_hh_l = None  # type: ignore
+
+    def init_hidden(self, batch_size: int, device: torch.device = torch.device("cpu")) -> Tensor:
+        return torch.zeros(batch_size, self.n_groups, self.out_ch, device=device)
+
+    def forward(self, input: Tensor, h=None) -> Tuple[Tensor, Tensor]:
+        # input: [B, Ci, T, F]
+        assert self.n_freqs == input.shape[-1]
+        assert self._in_ch == input.shape[1]
+        if h is None:
+            h = self.init_hidden(input.shape[0])
+        input = input.permute(0, 2, 3, 1).unflatten(2, (self.n_groups, self.g_freqs)).flatten(3)
+        input = torch.einsum("btgi,goi->btgo", input, self.weight_ih_l)
+        if self.bias_ih_l is not None:
+            input = input + self.bias_ih_l
+        h_out: List[Tensor] = []
+        for t in range(input.shape[1]):
+            hh = torch.einsum("bgo,gpo->bgp", h, self.weight_hh_l)
+            if self.bias_hh_l is not None:
+                hh = hh + self.bias_hh_l
+            ri, zi, ni = input[:, t].split(self.out_ch, dim=2)
+            rh, zh, nh = hh.split(self.out_ch, dim=2)
+            r = torch.sigmoid(ri + rh)
+            z = torch.sigmoid(zi + zh)
+            n = torch.tanh(ni + r * nh)
+            h = (1 - z) * n + z * h
+            h_out.append(h)
+        out = torch.stack(h_out, dim=1)  # [B, T, G, F/G*Co]
+        out = out.unflatten(3, (self.g_freqs, -1)).flatten(2, 3)  # [B, T, F, Co]
+        out = out.permute(0, 3, 1, 2)  # [B, Co, T, F]
+        return out, h
+
+
+class GroupedGRUMS(nn.Module):
+    def __init__(
+        self,
+        in_ch: int,
+        out_ch: int,
+        n_freqs: int,
+        n_groups: int,
+        n_layers: int = 1,
+        bias: bool = True,
+        add_outputs: bool = False,
+    ):
+        super().__init__()
+        self.n_layers = n_layers
+        self.grus: List[GroupedGRULayerMS] = nn.ModuleList()  # type: ignore
+        gru_layer = partial(
+            GroupedGRULayerMS, out_ch=out_ch, n_freqs=n_freqs, n_groups=n_groups, bias=bias
+        )
+        self.gru0 = gru_layer(in_ch=in_ch)
+        for _ in range(1, n_layers):
+            self.grus.append(gru_layer(in_ch=out_ch))
+        self.add_outputs = add_outputs
+
+    def init_hidden(self, batch_size: int, device: torch.device = torch.device("cpu")) -> Tensor:
+        return torch.stack(
+            tuple(self.gru0.init_hidden(batch_size, device) for _ in range(self.n_layers))
+        )
+
+    def forward(self, input: Tensor, h=None) -> Tuple[Tensor, Tensor]:
+        if h is None:
+            h = self.init_hidden(input.shape[0], input.device)
+        h_out = []
+        # First layer
+        input, hl = self.gru0(input, h[0])
+        h_out.append(hl)
+        output = input
+        for i, gru in enumerate(self.grus, 1):
+            input, hl = gru(input, h[i])
+            h_out.append(hl)
+            if self.add_outputs:
+                output = output + input
+        if not self.add_outputs:
+            output = input
+        return output, torch.stack(h_out)  # type: ignore
+
+
+class FreqStage(nn.Module):
+    in_ch: Final[int]
+    out_ch: Final[int]
+    depth: Final[int]
+
+    def __init__(
+        self,
+        in_ch: int,
+        out_ch: int,
+        out_act: Optional[Callable[..., torch.nn.Module]],
+        num_freqs: int,
+        gru_dim: int,
+        widths: List[int],
+        fstrides: Optional[List[int]] = None,
+        initial_kernel: Tuple[int, int] = (3, 3),
+        kernel: Tuple[int, int] = (1, 3),
+        separable_conv: bool = False,
+        num_gru_layers: int = 3,
+        num_gru_groups: int = 8,
+        global_pathway: bool = False,
+        decoder_out_layer: Optional[Callable[[int, int], torch.nn.Module]] = None,
+    ):
+        super().__init__()
+        self.in_ch = in_ch
+        self.out_ch = out_ch
+        self.depth = len(widths) - 1
+        if fstrides is not None:
+            assert len(fstrides) >= self.depth
+            overall_stride = reduce(lambda x, y: x * y, fstrides)
+        else:
+            fstrides = [2] * self.depth
+            overall_stride = 2 ** (len(widths) - 1)
+        assert num_freqs % overall_stride == 0, f"num_freqs ({num_freqs}) must correspond to depth"
+        norm_layer = nn.BatchNorm2d
+
+        # Encoder
+        self.enc0 = Conv2dNormAct(
+            in_ch,
+            widths[0],
+            initial_kernel,
+            fstride=1,
+            norm_layer=norm_layer,
+            separable=separable_conv,
+        )
+        self.enc = nn.ModuleList()
+
+        fstrides = fstrides or [2] * self.depth
+        freqs = num_freqs
+        for i in range(self.depth):
+            in_ch = widths[i]
+            out_ch = widths[i + 1]
+            fstride = fstrides[i]
+            freqs = freqs // fstride
+            self.enc.append(
+                Conv2dNormAct(
+                    in_ch, out_ch, kernel_size=kernel, fstride=fstride, separable=separable_conv
+                )
+            )
+        self.inner_freqs = freqs
+        self.lin_emb_in = LocalLinearCF(out_ch, gru_dim // freqs, n_freqs=freqs)
+        if num_gru_groups == 1:
+            self.gru = CFGRU(
+                freqs=self.inner_freqs,
+                input_size=gru_dim // freqs * freqs,
+                hidden_size=gru_dim,
+                num_layers=num_gru_layers,
+            )
+        else:
+            self.gru = GroupedGRU(
+                in_ch=gru_dim // freqs,
+                out_ch=gru_dim // freqs,
+                n_freqs=freqs,
+                n_groups=min(freqs, num_gru_groups),
+                n_layers=num_gru_layers,
+            )
+        self.lin_emb_out = LocalLinearCF(gru_dim // freqs, out_ch, n_freqs=freqs)
+        self.gru_skip = nn.Conv2d(out_ch, out_ch, 1)
+
+        self.dec = nn.ModuleList()
+        for i in range(self.depth - 1, -1, -1):
+            in_ch = widths[i + 1]
+            out_ch = widths[i]
+            fstride = fstrides[i]
+            dec_layer = ConvTranspose2dNormAct if fstride > 1 else Conv2dNormAct
+            self.dec.append(
+                dec_layer(
+                    in_ch, out_ch, kernel_size=kernel, fstride=fstride, separable=separable_conv
+                )
+            )
+        self.global_pathway = (
+            Conv2dNormAct(widths[0], widths[0], 1) if global_pathway else nn.Identity()
+        )
+        if decoder_out_layer is None:
+            self.dec0 = Conv2dNormAct(
+                widths[0],
+                self.out_ch,
+                kernel,
+                fstride=1,
+                norm_layer=None,
+                activation_layer=out_act,
+                separable=separable_conv,
+            )
+        else:
+            self.dec0 = decoder_out_layer(widths[0], self.out_ch)
+
+    def encode(self, x: Tensor) -> Tuple[Tensor, List[Tensor]]:
+        intermediate = []
+        x = self.enc0(x)
+        for enc_layer in self.enc:
+            intermediate.append(x)
+            x = enc_layer(x)
+        return x, intermediate
+
+    def embed(self, input: Tensor, h=None) -> Tuple[Tensor, Tensor]:
+        x = self.lin_emb_in(input)
+        x_gru, h = self.gru(x, h)
+        x_gru = self.lin_emb_out(x_gru)
+        x = self.gru_skip(input) + x_gru
+        return x, h
+
+    def decode(self, x: Tensor, intermediate: List[Tensor]) -> Tensor:
+        intermediate[0] = self.global_pathway(intermediate[0])
+        for dec_layer, x_enc in zip(self.dec, reversed(intermediate)):
+            x = dec_layer(x) + x_enc
+        x = self.dec0(x)
+        return x
+
+    def forward(self, x: Tensor, h=None) -> Tuple[Tensor, Tensor, Tensor]:
+        # input shape: [B, C, T, F]
+        x_enc, intermediate = self.encode(x)
+        x_inner, h = self.embed(x_enc, h)
+        x = self.decode(x_inner, intermediate)
+        return x, x_enc, h
+
+
+class CFGRU(nn.Module):
+    def __init__(self, freqs: int, *args, **kwargs):
+        super().__init__()
+        self.inner_freqs = freqs
+        self.gru = nn.GRU(*args, **kwargs)
+
+    def forward(self, x, h):
+        x = x.permute(0, 2, 3, 1).flatten(2)
+        x_gru, h = self.gru(x, h)
+        x_gru = x.unflatten(2, (self.inner_freqs, -1)).permute(0, 3, 1, 2)
+        return x_gru, h
+
+
+class ComplexCompression(nn.Module):
+    def __init__(self, n_freqs: int, init_value: float = 0.3):
+        super().__init__()
+        self.register_parameter(
+            "c", Parameter(torch.full((n_freqs,), init_value), requires_grad=True)
+        )
+
+    def forward(self, x: Tensor):
+        # x has shape x [B, 2, T, F]
+        x_abs = (x[:, 0].square() + x[:, 1].square()).clamp_min(1e-10).pow(self.c)
+        x_ang = angle_re_im.apply(x[:, 0], x[:, 1])
+        x = torch.stack((x_abs * torch.cos(x_ang), x_abs * torch.sin(x_ang)), dim=1)
+        # x_c = x_abs * torch.exp(1j * x_ang)
+        # x_ = torch.view_as_complex(x.permute(0,2,3,1).contiguous())
+        # torch.allclose(x_, x_c)
+        return x
+
+
+class MagCompression(nn.Module):
+    def __init__(self, n_freqs: int, init_value: float = 0.3):
+        super().__init__()
+        self.c: Tensor
+        self.register_parameter(
+            "c", Parameter(torch.full((n_freqs,), init_value), requires_grad=True)
+        )
+        self.mn: Tensor
+        self.register_parameter("mn", Parameter(torch.full((n_freqs,), -0.2), requires_grad=True))
+
+    def forward(self, x: Tensor):
+        # x has shape x [B, T, F, 2]
+        x = x.pow(self.c) + self.mn
+        return x
+
+
+class ComplexAdd(nn.Module):
+    def forward(self, a, b):
+        return a + b
+
+
+class ComplexMul(nn.Module):
+    def forward(self, a, b):
+        # [B, 2, *]
+        re = a[:, :1] * b[:, :1] - a[:, 1:] * b[:, 1:]
+        im = a[:, :1] * b[:, 1:] + a[:, :1] * b[:, 1:]
+        return torch.cat((re, im), dim=1)
+
+
+class MSNet(nn.Module):
+    def __init__(self, erb_fb: Tensor, erb_inv_fb: Tensor):
+        super().__init__()
+        p = ModelParams()
+        self.stages = p.stages
+        self.freq_bins = p.fft_size // 2 + 1
+        self.erb_bins = p.nb_erb
+        self.df_bins = p.nb_df
+        self.erb_fb: Tensor
+        self.erb_comp = MagCompression(self.erb_bins)
+        self.cplx_comp = ComplexCompression(self.df_bins)
+        self.register_buffer("erb_fb", erb_fb, persistent=False)
+        self.erb_stage = FreqStage(
+            in_ch=1,
+            out_ch=1,
+            widths=p.erb_widths,
+            out_act=nn.Sigmoid,
+            num_freqs=p.nb_erb,
+            gru_dim=p.erb_hidden_dim,
+        )
+        strides = [2, 2, 2, 2, 1, 1, 1]
+        self.mask = Mask(erb_inv_fb, post_filter=p.mask_pf)
+        refinement_act = {"tanh": nn.Tanh, "identity": nn.Identity}[p.refinement_act.lower()]
+        assert p.refinement_out_layer.lower() in ("locallyconnected", "conv2d")
+        refinement_out_layer = (
+            partial(Conv2dNormAct, kernel_size=(3, 1), norm_layer=None, activation_layer=None)
+            if p.refinement_out_layer.lower() == "conv2d"
+            else partial(GroupedLinearCF, n_freqs=p.nb_df, n_groups=8)
+        )
+        self.refinement_stage = FreqStage(
+            in_ch=2,
+            out_ch=2,
+            out_act=refinement_act,
+            widths=p.refinement_widths,
+            num_freqs=p.nb_df,
+            gru_dim=p.refinement_hidden_dim,
+            fstrides=strides,
+            decoder_out_layer=refinement_out_layer,
+        )
+        self.refinement_op = ComplexMul() if p.refinement_op == "mul" else ComplexAdd()
+        self.lsnr_net = LSNRNet(p.erb_widths[-1], lsnr_min=p.lsnr_min, lsnr_max=p.lsnr_max)
+
+    def forward(
+        self, spec: Tensor, atten_lim: Optional[Tensor] = None
+    ) -> Tuple[Tensor, Tensor, Tensor, None]:
+        # Spec shape: [B, 1, T, F, 2]
+        feat_erb = torch.view_as_complex(spec).abs().matmul(self.erb_fb)
+        feat_erb = self.erb_comp(feat_erb)
+        m, x_rnn, _ = self.erb_stage(feat_erb)
+        spec = self.mask(spec, m, atten_lim)  # [B, 1, T, F, 2]
+        lsnr, _ = self.lsnr_net(x_rnn)
+        # re/im into channel axis
+        spec_f = (
+            spec.squeeze(1)[:, :, : self.df_bins].permute(0, 3, 1, 2).clone()
+        )  # [B, 2, T, F_df]
+        r, _, _ = self.refinement_stage(self.cplx_comp(spec_f))
+        spec_f = self.refinement_op(spec_f, r)
+        spec[..., : self.df_bins, :] = spec_f.unsqueeze(-1).transpose(1, -1)
+        return spec, m, lsnr, None
+
+
+def init_model(df_state: Optional[DF] = None, run_df: bool = True, train_mask: bool = True):
+    assert run_df and train_mask
+    p = ModelParams()
+    if df_state is None:
+        df_state = DF(sr=p.sr, fft_size=p.fft_size, hop_size=p.hop_size, nb_bands=p.nb_erb)
+    erb_width = df_state.erb_widths()
+    erb = erb_fb(erb_width, p.sr)
+    erb_inverse = erb_fb(erb_width, p.sr, inverse=True)
+    model = MSNet(erb, erb_inverse)
+    return model.to(device=get_device())
diff --git a/DeepFilterNet/df/scripts/list_attrs_in_hdf5.py b/DeepFilterNet/df/scripts/list_attrs_in_hdf5.py
new file mode 100644
index 000000000..c7bfcdbe6
--- /dev/null
+++ b/DeepFilterNet/df/scripts/list_attrs_in_hdf5.py
@@ -0,0 +1,9 @@
+import sys
+
+import h5py
+
+with h5py.File(sys.argv[1], "r", libver="latest", swmr=True) as h5f:
+    for n, k in h5f.attrs.items():
+        print(f"Found attr {n} '{k}' in {sys.argv[1]}")
+    for group, samples in h5f.items():
+        print(f"Found {len(samples)} samples in {group}")
diff --git a/DeepFilterNet/df/scripts/plot_lrs.py b/DeepFilterNet/df/scripts/plot_lrs.py
new file mode 100644
index 000000000..64d323a69
--- /dev/null
+++ b/DeepFilterNet/df/scripts/plot_lrs.py
@@ -0,0 +1,43 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from icecream import install  # noqa
+
+from df.lr import cosine_scheduler
+
+install()
+
+LR = 1e-2
+WARMUP = 3
+LR_WARMUP_INIT = 1e-5
+LR_MIN = 0
+CYCLE_DECAY = 1
+CYCLE_LIM = 1
+CYCLE_MUL = 1
+
+NUM_EPOCHS = 50
+STEPS_PER_EPOCH = 20
+
+
+def get_lr_updates():
+    lrs = cosine_scheduler(
+        LR,
+        LR_MIN,
+        NUM_EPOCHS,
+        STEPS_PER_EPOCH,
+        WARMUP,
+        LR_WARMUP_INIT,
+        initial_ep_per_cycle=5,
+        cycle_decay=0.9,
+        cycle_mul=1.5,
+    )
+    return lrs
+
+
+if __name__ == "__main__":
+    lr_data_per_step = get_lr_updates()
+    t = np.linspace(0, NUM_EPOCHS, NUM_EPOCHS * STEPS_PER_EPOCH)
+    plt.plot(t, lr_data_per_step, label="Update per step")
+    plt.grid()
+    plt.xticks(range(NUM_EPOCHS))
+    plt.legend()
+    plt.show()
diff --git a/DeepFilterNet/df/scripts/prepare_data.py b/DeepFilterNet/df/scripts/prepare_data.py
index f0af0990c..f18c6b79f 100755
--- a/DeepFilterNet/df/scripts/prepare_data.py
+++ b/DeepFilterNet/df/scripts/prepare_data.py
@@ -135,7 +135,7 @@ def __getitem__(self, index):
         fn = self.file_names[index]
         logger.debug(f"Reading audio file {fn}")
         x = self.read(fn)
-        assert x.dim() == 2 and x.shape[0] <= 16, x.shape
+        assert x.dim() == 2 and x.shape[0] <= 16, f"Got sample {fn} with unexpected shape {x.shape}"
         n_samples = x.shape[1]
         if self.codec == "vorbis":
             with NamedTemporaryFile(suffix=".ogg") as tf:
diff --git a/DeepFilterNet/df/scripts/split_hdf5.py b/DeepFilterNet/df/scripts/split_hdf5.py
new file mode 100644
index 000000000..98ebe6a01
--- /dev/null
+++ b/DeepFilterNet/df/scripts/split_hdf5.py
@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+
+import os
+import shutil
+import sys
+
+import h5py
+import numpy as np
+from icecream import ic
+
+splits = (0.7, 0.15, 0.15)
+assert np.sum(splits) == 1
+
+
+def main(hdf5: str, force: bool = False):
+    assert hdf5.endswith("TRAIN.hdf5")
+    f_train = h5py.File(hdf5, "r+")
+    if os.path.basename(hdf5).lower().startswith("vocalset"):
+        # Correct wrong sampling rate
+        f_train.attrs["sr"] = 16000
+        f_train.attrs["max_freq"] = 8000
+    hdf5_train_new = hdf5.replace("TRAIN", "TRAIN_NEW")
+    hdf5_valid = hdf5.replace("TRAIN", "VALID")
+    hdf5_test = hdf5.replace("TRAIN", "TEST")
+    if (os.path.exists(hdf5_test) or os.path.exists(hdf5_valid)) and not force:
+        raise FileExistsError(f"Dataset {hdf5_test} already exists.")
+    f_valid = h5py.File(hdf5_valid, "w")
+    f_test = h5py.File(hdf5_test, "w")
+    f_train_new = h5py.File(hdf5_train_new, "w")
+    for attr, v in f_train.attrs.items():
+        ic(attr, v)
+        f_train_new.attrs[attr] = v
+        f_valid.attrs[attr] = v
+        f_test.attrs[attr] = v
+    for key in f_train:
+        grp_train = f_train[key]
+        grp_train_new = f_train_new.create_group(key)
+        grp_valid = f_valid.create_group(key)
+        grp_test = f_test.create_group(key)
+        keys = list(grp_train.keys())
+        n = len(keys)
+        # sections is a list of sorted integers; I.e. the cumulated sum
+        len_train, len_valid = (int(s * n) for s in splits[:2])
+        sections = [len_train, len_train + len_valid]
+        # We don't need a section for test, it will be the rest anyways
+        len_test = n - sections[1]
+        keys_train, keys_valid, keys_test = np.split(np.random.permutation(keys), sections)
+        ic(len(keys_train), len(keys_valid), len(keys_test))
+        for k in keys_train:
+            print("train", k)
+            grp_train_new.copy(grp_train[k], dest=k)
+        assert len(grp_train_new) == len_train
+        for k in keys_valid:
+            print("valid", k)
+            grp_valid.copy(grp_train[k], dest=k)
+        assert len(grp_valid) == len_valid
+        for k in keys_test:
+            print("test", k)
+            grp_test.copy(grp_train[k], dest=k)
+        assert len(grp_test) == len_test
+    ic(shutil.move(hdf5_train_new, hdf5))
+
+
+if __name__ == "__main__":
+    force = sys.argv[2] == "-f" if len(sys.argv) > 2 else False
+    main(sys.argv[1], force=force)
diff --git a/DeepFilterNet/df/sepm.py b/DeepFilterNet/df/sepm.py
new file mode 100644
index 000000000..55ea7760a
--- /dev/null
+++ b/DeepFilterNet/df/sepm.py
@@ -0,0 +1,511 @@
+import numpy as np
+from pesq import pesq
+from scipy.linalg import toeplitz
+from scipy.signal import stft
+
+#################################################
+#
+#  SPEECH ENHANCEMENT PERFORMANCE METRICS
+#
+#  DENGFENG.KE @ 2020-03-14 BEIJING.CHINA
+#
+#  Modified version of https://raw.githubusercontent.com/IMLHF/PHASEN-PyTorch/9e68773/phasen_torch/sepm.py
+#
+#  This version is verified to be numerically close to the original octave/matlab version.
+
+
+def extractOverlappedWindows(x, nperseg, noverlap, window=None):
+    # source: https://github.com/scipy/scipy/blob/v1.2.1/scipy/signal/spectral.py
+    step = nperseg - noverlap
+    shape = x.shape[:-1] + ((x.shape[-1] - noverlap) // step, nperseg)
+    strides = x.strides[:-1] + (step * x.strides[-1], x.strides[-1])
+    result = np.lib.stride_tricks.as_strided(x, shape=shape, strides=strides)
+    if window is not None:
+        result = window * result
+    return result
+
+
+def SNRseg(clean_speech, processed_speech, fs, frameLen=0.03, overlap=0.75):
+    eps = np.finfo(np.float64).eps
+
+    winlength = round(frameLen * fs)  # window length in samples
+    skiprate = int(np.floor((1 - overlap) * frameLen * fs))  # window skip in samples
+    MIN_SNR = -10  # minimum SNR in dB
+    MAX_SNR = 35  # maximum SNR in dB
+
+    hannWin = 0.5 * (1 - np.cos(2 * np.pi * np.arange(1, winlength + 1) / (winlength + 1)))
+    clean_speech_framed = extractOverlappedWindows(
+        clean_speech, winlength, winlength - skiprate, hannWin
+    )
+    processed_speech_framed = extractOverlappedWindows(
+        processed_speech, winlength, winlength - skiprate, hannWin
+    )
+
+    signal_energy = np.power(clean_speech_framed, 2).sum(-1)
+    noise_energy = np.power(clean_speech_framed - processed_speech_framed, 2).sum(-1)
+
+    segmental_snr = 10 * np.log10(signal_energy / (noise_energy + eps) + eps)
+    segmental_snr[segmental_snr < MIN_SNR] = MIN_SNR
+    segmental_snr[segmental_snr > MAX_SNR] = MAX_SNR
+    segmental_snr = segmental_snr[:-1]  # remove last frame -> not valid
+    return np.mean(segmental_snr)
+
+
+def fwSNRseg(cleanSig, enhancedSig, fs, frameLen=0.03, overlap=0.75):
+    if cleanSig.shape != enhancedSig.shape:
+        raise ValueError("The two signals do not match!")
+    eps = np.finfo(np.float64).eps
+    cleanSig = cleanSig.astype(np.float64) + eps
+    enhancedSig = enhancedSig.astype(np.float64) + eps
+    winlength = round(frameLen * fs)  # window length in samples
+    skiprate = int(np.floor((1 - overlap) * frameLen * fs))  # window skip in samples
+    max_freq = fs / 2  # maximum bandwidth
+    num_crit = 25  # number of critical bands
+    n_fft = 2 ** np.ceil(np.log2(2 * winlength))
+    n_fftby2 = int(n_fft / 2)
+    gamma = 0.2
+
+    cent_freq = np.zeros((num_crit,))
+    bandwidth = np.zeros((num_crit,))
+
+    cent_freq[0] = 50.0000
+    bandwidth[0] = 70.0000
+    cent_freq[1] = 120.000
+    bandwidth[1] = 70.0000
+    cent_freq[2] = 190.000
+    bandwidth[2] = 70.0000
+    cent_freq[3] = 260.000
+    bandwidth[3] = 70.0000
+    cent_freq[4] = 330.000
+    bandwidth[4] = 70.0000
+    cent_freq[5] = 400.000
+    bandwidth[5] = 70.0000
+    cent_freq[6] = 470.000
+    bandwidth[6] = 70.0000
+    cent_freq[7] = 540.000
+    bandwidth[7] = 77.3724
+    cent_freq[8] = 617.372
+    bandwidth[8] = 86.0056
+    cent_freq[9] = 703.378
+    bandwidth[9] = 95.3398
+    cent_freq[10] = 798.717
+    bandwidth[10] = 105.411
+    cent_freq[11] = 904.128
+    bandwidth[11] = 116.256
+    cent_freq[12] = 1020.38
+    bandwidth[12] = 127.914
+    cent_freq[13] = 1148.30
+    bandwidth[13] = 140.423
+    cent_freq[14] = 1288.72
+    bandwidth[14] = 153.823
+    cent_freq[15] = 1442.54
+    bandwidth[15] = 168.154
+    cent_freq[16] = 1610.70
+    bandwidth[16] = 183.457
+    cent_freq[17] = 1794.16
+    bandwidth[17] = 199.776
+    cent_freq[18] = 1993.93
+    bandwidth[18] = 217.153
+    cent_freq[19] = 2211.08
+    bandwidth[19] = 235.631
+    cent_freq[20] = 2446.71
+    bandwidth[20] = 255.255
+    cent_freq[21] = 2701.97
+    bandwidth[21] = 276.072
+    cent_freq[22] = 2978.04
+    bandwidth[22] = 298.126
+    cent_freq[23] = 3276.17
+    bandwidth[23] = 321.465
+    cent_freq[24] = 3597.63
+    bandwidth[24] = 346.136
+
+    bw_min = bandwidth[0]
+    min_factor = np.exp(-30.0 / (2.0 * 2.303))  # % -30 dB point of filter
+
+    all_f0 = np.zeros((num_crit,))
+    crit_filter = np.zeros((num_crit, int(n_fftby2)))
+    j = np.arange(0, n_fftby2)
+
+    for i in range(num_crit):
+        f0 = (cent_freq[i] / max_freq) * (n_fftby2)
+        all_f0[i] = np.floor(f0)
+        bw = (bandwidth[i] / max_freq) * (n_fftby2)
+        norm_factor = np.log(bw_min) - np.log(bandwidth[i])
+        crit_filter[i, :] = np.exp(-11 * (((j - np.floor(f0)) / bw) ** 2) + norm_factor)
+        crit_filter[i, :] = crit_filter[i, :] * (crit_filter[i, :] > min_factor)
+
+    num_frames = len(cleanSig) / skiprate - (winlength / skiprate)  # number of frames
+
+    hannWin = 0.5 * (1 - np.cos(2 * np.pi * np.arange(1, winlength + 1) / (winlength + 1)))
+    f, t, Zxx = stft(
+        cleanSig[0 : int(num_frames) * skiprate + int(winlength - skiprate)],
+        fs=fs,
+        window=hannWin,
+        nperseg=winlength,
+        noverlap=winlength - skiprate,
+        nfft=n_fft,
+        detrend=False,
+        return_onesided=True,
+        boundary=None,
+        padded=False,
+    )
+    clean_spec = np.abs(Zxx)
+    clean_spec = clean_spec[:-1, :]
+    clean_spec = clean_spec / clean_spec.sum(0)
+    f, t, Zxx = stft(
+        enhancedSig[0 : int(num_frames) * skiprate + int(winlength - skiprate)],
+        fs=fs,
+        window=hannWin,
+        nperseg=winlength,
+        noverlap=winlength - skiprate,
+        nfft=n_fft,
+        detrend=False,
+        return_onesided=True,
+        boundary=None,
+        padded=False,
+    )
+    enh_spec = np.abs(Zxx)
+    enh_spec = enh_spec[:-1, :]
+    enh_spec = enh_spec / enh_spec.sum(0)
+
+    clean_energy = crit_filter.dot(clean_spec)
+    processed_energy = crit_filter.dot(enh_spec)
+    error_energy = np.power(clean_energy - processed_energy, 2)
+    error_energy[error_energy < eps] = eps
+    W_freq = np.power(clean_energy, gamma)
+    SNRlog = 10 * np.log10((clean_energy**2) / error_energy)
+    fwSNR = np.sum(W_freq * SNRlog, 0) / np.sum(W_freq, 0)
+    distortion = fwSNR.copy()
+    distortion[distortion < -10] = -10
+    distortion[distortion > 35] = 35
+
+    return np.mean(distortion)
+
+
+# def lpcoeff(speech_frame, model_order):
+#    # ----------------------------------------------------------
+#    # (1) Compute Autocorrelation Lags
+#    # ----------------------------------------------------------
+
+#     R=correlate(speech_frame,speech_frame)
+#     R=R[len(speech_frame)-1:len(speech_frame)+model_order]
+#    # ----------------------------------------------------------
+#    # (2) Levinson-Durbin
+#    # ----------------------------------------------------------
+#     lpparams=np.ones((model_order+1))
+#     lpparams[1:]=solve_toeplitz(R[0:-1],-R[1:])
+#     return(lpparams,R)
+
+
+def lpcoeff(speech_frame, model_order):
+    # (1) Compute Autocor lags
+    eps = np.finfo(np.float64).eps
+    winlength = speech_frame.shape[0]
+    R = []
+    for k in range(model_order + 1):
+        first = speech_frame[: (winlength - k)]
+        second = speech_frame[k:winlength]
+        R.append(np.sum(first * second))
+
+    # (2) Lev-Durbin
+    a = np.ones((model_order,))
+    E = np.zeros((model_order + 1,))
+    rcoeff = np.zeros((model_order,))
+    E[0] = R[0]
+    for i in range(model_order):
+        if i == 0:
+            sum_term = 0
+        else:
+            a_past = a[:i]
+            sum_term = np.sum(a_past * np.array(R[i:0:-1]))
+        # rcoeff[i] = (R[i+1] - sum_term)/E[i] # fixed by LiHongfeng
+        rcoeff[i] = (R[i + 1] - sum_term) / max(E[i], eps)
+        # if E[i] == 0:
+        #   print(233333, i, eps==0)
+        a[i] = rcoeff[i]
+        if i > 0:
+            a[:i] = a_past[:i] - rcoeff[i] * a_past[::-1]
+        E[i + 1] = (1 - rcoeff[i] * rcoeff[i]) * E[i]
+    acorr = np.array(R, dtype=np.float32)
+    refcoeff = np.array(rcoeff, dtype=np.float32)
+    a = a * -1
+    lpparams = np.array([1] + list(a), dtype=np.float32)
+    acorr = np.array(acorr, dtype=np.float32)
+    refcoeff = np.array(refcoeff, dtype=np.float32)
+    lpparams = np.array(lpparams, dtype=np.float32)
+
+    # return acorr, refcoeff, lpparams
+    return lpparams, acorr
+
+
+def llr(clean_speech, processed_speech, fs, frameLen=0.03, overlap=0.75):
+    eps = np.finfo(np.float64).eps
+    alpha = 0.95
+    winlength = round(frameLen * fs)  # window length in samples
+    skiprate = int(np.floor((1 - overlap) * frameLen * fs))  # window skip in samples
+    if fs < 10000:
+        P = 10  # LPC Analysis Order
+    else:
+        P = 16  # this could vary depending on sampling frequency.
+
+    hannWin = 0.5 * (1 - np.cos(2 * np.pi * np.arange(1, winlength + 1) / (winlength + 1)))
+    clean_speech_framed = extractOverlappedWindows(
+        clean_speech, winlength, winlength - skiprate, hannWin
+    )
+    processed_speech_framed = extractOverlappedWindows(
+        processed_speech, winlength, winlength - skiprate, hannWin
+    )
+    numFrames = clean_speech_framed.shape[0]
+    if numFrames == 0:
+        return None
+    numerators = np.zeros((numFrames - 1,))
+    denominators = np.zeros((numFrames - 1,)) + eps
+
+    for ii in range(numFrames - 1):
+        A_clean, R_clean = lpcoeff(clean_speech_framed[ii, :], P)
+        A_proc, R_proc = lpcoeff(processed_speech_framed[ii, :], P)
+
+        numerators[ii] = A_proc.dot(toeplitz(R_clean).dot(A_proc.T))
+        denominators[ii] = A_clean.dot(toeplitz(R_clean).dot(A_clean.T)) + eps
+
+    frac = numerators / denominators
+    frac[frac <= 0] = 1000
+    distortion = np.log(frac)
+    # distortion[distortion>2]=2 # this line is not in composite measure matlab implementation of loizou
+    distortion = np.sort(distortion)
+    distortion = distortion[: int(round(len(distortion) * alpha))]
+    return np.mean(distortion)
+
+
+def findLocPeaks(slope, energy):
+    num_crit = len(energy)
+
+    loc_peaks = np.zeros_like(slope)
+
+    for ii in range(len(slope)):
+        n = ii
+        if slope[ii] > 0:
+            while (n < num_crit - 1) and (slope[n] > 0):
+                n = n + 1
+            loc_peaks[ii] = energy[n - 1]
+        else:
+            while (n >= 0) and (slope[n] <= 0):
+                n = n - 1
+            loc_peaks[ii] = energy[n + 1]
+
+    return loc_peaks
+
+
+def wss(clean_speech, processed_speech, fs, frameLen=0.03, overlap=0.75):
+
+    Kmax = 20  # value suggested by Klatt, pg 1280
+    Klocmax = 1  # value suggested by Klatt, pg 1280
+    alpha = 0.95
+    if clean_speech.shape != processed_speech.shape:
+        raise ValueError("The two signals do not match!")
+    eps = np.finfo(np.float64).eps
+    clean_speech = clean_speech.astype(np.float64) + eps
+    processed_speech = processed_speech.astype(np.float64) + eps
+    winlength = round(frameLen * fs)  # window length in samples
+    skiprate = int(np.floor((1 - overlap) * frameLen * fs))  # window skip in samples
+    max_freq = fs / 2  # maximum bandwidth
+    num_crit = 25  # number of critical bands
+    n_fft = 2 ** np.ceil(np.log2(2 * winlength))
+    n_fftby2 = int(n_fft / 2)
+
+    cent_freq = np.zeros((num_crit,))
+    bandwidth = np.zeros((num_crit,))
+
+    cent_freq[0] = 50.0000
+    bandwidth[0] = 70.0000
+    cent_freq[1] = 120.000
+    bandwidth[1] = 70.0000
+    cent_freq[2] = 190.000
+    bandwidth[2] = 70.0000
+    cent_freq[3] = 260.000
+    bandwidth[3] = 70.0000
+    cent_freq[4] = 330.000
+    bandwidth[4] = 70.0000
+    cent_freq[5] = 400.000
+    bandwidth[5] = 70.0000
+    cent_freq[6] = 470.000
+    bandwidth[6] = 70.0000
+    cent_freq[7] = 540.000
+    bandwidth[7] = 77.3724
+    cent_freq[8] = 617.372
+    bandwidth[8] = 86.0056
+    cent_freq[9] = 703.378
+    bandwidth[9] = 95.3398
+    cent_freq[10] = 798.717
+    bandwidth[10] = 105.411
+    cent_freq[11] = 904.128
+    bandwidth[11] = 116.256
+    cent_freq[12] = 1020.38
+    bandwidth[12] = 127.914
+    cent_freq[13] = 1148.30
+    bandwidth[13] = 140.423
+    cent_freq[14] = 1288.72
+    bandwidth[14] = 153.823
+    cent_freq[15] = 1442.54
+    bandwidth[15] = 168.154
+    cent_freq[16] = 1610.70
+    bandwidth[16] = 183.457
+    cent_freq[17] = 1794.16
+    bandwidth[17] = 199.776
+    cent_freq[18] = 1993.93
+    bandwidth[18] = 217.153
+    cent_freq[19] = 2211.08
+    bandwidth[19] = 235.631
+    cent_freq[20] = 2446.71
+    bandwidth[20] = 255.255
+    cent_freq[21] = 2701.97
+    bandwidth[21] = 276.072
+    cent_freq[22] = 2978.04
+    bandwidth[22] = 298.126
+    cent_freq[23] = 3276.17
+    bandwidth[23] = 321.465
+    cent_freq[24] = 3597.63
+    bandwidth[24] = 346.136
+
+    W = np.array(
+        [
+            0.003,
+            0.003,
+            0.003,
+            0.007,
+            0.010,
+            0.016,
+            0.016,
+            0.017,
+            0.017,
+            0.022,
+            0.027,
+            0.028,
+            0.030,
+            0.032,
+            0.034,
+            0.035,
+            0.037,
+            0.036,
+            0.036,
+            0.033,
+            0.030,
+            0.029,
+            0.027,
+            0.026,
+            0.026,
+        ]
+    )
+
+    bw_min = bandwidth[0]
+    min_factor = np.exp(-30.0 / (2.0 * 2.303))
+    #      % -30 dB point of filter
+
+    all_f0 = np.zeros((num_crit,))
+    crit_filter = np.zeros((num_crit, int(n_fftby2)))
+    j = np.arange(0, n_fftby2)
+
+    for i in range(num_crit):
+        f0 = (cent_freq[i] / max_freq) * (n_fftby2)
+        all_f0[i] = np.floor(f0)
+        bw = (bandwidth[i] / max_freq) * (n_fftby2)
+        norm_factor = np.log(bw_min) - np.log(bandwidth[i])
+        crit_filter[i, :] = np.exp(-11 * (((j - np.floor(f0)) / bw) ** 2) + norm_factor)
+        crit_filter[i, :] = crit_filter[i, :] * (crit_filter[i, :] > min_factor)
+
+    num_frames = len(clean_speech) / skiprate - (winlength / skiprate)  # number of frames
+
+    hannWin = 0.5 * (1 - np.cos(2 * np.pi * np.arange(1, winlength + 1) / (winlength + 1)))
+    scale = np.sqrt(1.0 / hannWin.sum() ** 2)
+
+    f, t, Zxx = stft(
+        clean_speech[0 : int(num_frames) * skiprate + int(winlength - skiprate)],
+        fs=fs,
+        window=hannWin,
+        nperseg=winlength,
+        noverlap=winlength - skiprate,
+        nfft=n_fft,
+        detrend=False,
+        return_onesided=True,
+        boundary=None,
+        padded=False,
+    )
+    clean_spec = np.power(np.abs(Zxx) / scale, 2)
+    clean_spec = clean_spec[:-1, :]
+
+    f, t, Zxx = stft(
+        processed_speech[0 : int(num_frames) * skiprate + int(winlength - skiprate)],
+        fs=fs,
+        window=hannWin,
+        nperseg=winlength,
+        noverlap=winlength - skiprate,
+        nfft=n_fft,
+        detrend=False,
+        return_onesided=True,
+        boundary=None,
+        padded=False,
+    )
+    proc_spec = np.power(np.abs(Zxx) / scale, 2)
+    proc_spec = proc_spec[:-1, :]
+
+    clean_energy = crit_filter.dot(clean_spec)
+    log_clean_energy = 10 * np.log10(clean_energy)
+    log_clean_energy[log_clean_energy < -100] = -100
+    proc_energy = crit_filter.dot(proc_spec)
+    log_proc_energy = 10 * np.log10(proc_energy)
+    log_proc_energy[log_proc_energy < -100] = -100
+
+    log_clean_energy_slope = np.diff(log_clean_energy, axis=0)
+    log_proc_energy_slope = np.diff(log_proc_energy, axis=0)
+
+    dBMax_clean = np.max(log_clean_energy, axis=0)
+    dBMax_processed = np.max(log_proc_energy, axis=0)
+
+    numFrames = log_clean_energy_slope.shape[-1]
+
+    clean_loc_peaks = np.zeros_like(log_clean_energy_slope)
+    proc_loc_peaks = np.zeros_like(log_proc_energy_slope)
+    for ii in range(numFrames):
+        clean_loc_peaks[:, ii] = findLocPeaks(
+            log_clean_energy_slope[:, ii], log_clean_energy[:, ii]
+        )
+        proc_loc_peaks[:, ii] = findLocPeaks(log_proc_energy_slope[:, ii], log_proc_energy[:, ii])
+
+    Wmax_clean = Kmax / (Kmax + dBMax_clean - log_clean_energy[:-1, :])
+    Wlocmax_clean = Klocmax / (Klocmax + clean_loc_peaks - log_clean_energy[:-1, :])
+    W_clean = Wmax_clean * Wlocmax_clean
+
+    Wmax_proc = Kmax / (Kmax + dBMax_processed - log_proc_energy[:-1])
+    Wlocmax_proc = Klocmax / (Klocmax + proc_loc_peaks - log_proc_energy[:-1, :])
+    W_proc = Wmax_proc * Wlocmax_proc
+
+    W = (W_clean + W_proc) / 2.0
+
+    distortion = np.sum(W * (log_clean_energy_slope - log_proc_energy_slope) ** 2, axis=0)
+    distortion = distortion / np.sum(W, axis=0)
+    distortion = np.sort(distortion)
+    distortion = distortion[: int(round(len(distortion) * alpha))]
+    return np.mean(distortion)
+
+
+def composite(clean_speech, processed_speech, fs):
+    """Computes composite measures: PESQ, CSIG, CBAK, COVL, SSNR."""
+    assert fs == 16000
+    wss_dist = wss(clean_speech, processed_speech, fs)
+    llr_mean = llr(clean_speech, processed_speech, fs)
+    if llr_mean is None:
+        print("Error computing llr")
+        llr_mean = 0
+    segSNR = SNRseg(clean_speech, processed_speech, fs)
+    pesq_mos = pesq(fs, clean_speech, processed_speech, "wb")
+
+    Csig = 3.093 - 1.029 * llr_mean + 0.603 * pesq_mos - 0.009 * wss_dist
+    Csig = np.max((1, Csig))
+    Csig = np.min((5, Csig))  # limit values to [1, 5]
+    Cbak = 1.634 + 0.478 * pesq_mos - 0.007 * wss_dist + 0.063 * segSNR
+    Cbak = np.max((1, Cbak))
+    Cbak = np.min((5, Cbak))  # limit values to [1, 5]
+    Covl = 1.594 + 0.805 * pesq_mos - 0.512 * llr_mean - 0.007 * wss_dist
+    Covl = np.max((1, Covl))
+    Covl = np.min((5, Covl))  # limit values to [1, 5]
+    return pesq_mos, Csig, Cbak, Covl, segSNR
diff --git a/DeepFilterNet/df/stoi.py b/DeepFilterNet/df/stoi.py
index 3e561ba86..2bf906e4b 100644
--- a/DeepFilterNet/df/stoi.py
+++ b/DeepFilterNet/df/stoi.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 import torch
+from loguru import logger
 from torch import Tensor
 from torch.nn import functional as F
 
@@ -191,6 +192,9 @@ def stoi(x, y, fs_source):
 
     for i in range(B):
         # To spectral domain
+        if x_[i].numel() < N_fft:
+            logger.warning("Could not calculate STOI (not enough frames left). Skipping.")
+            continue
         x = _stft(x_[i], win_size=N_frame, fft_size=N_fft, hop_size=N_frame // 2)
         y = _stft(y_[i], win_size=N_frame, fft_size=N_fft, hop_size=N_frame // 2)
         # Power spectrogram
diff --git a/DeepFilterNet/df/train.py b/DeepFilterNet/df/train.py
index 851498af0..825ada687 100644
--- a/DeepFilterNet/df/train.py
+++ b/DeepFilterNet/df/train.py
@@ -1,20 +1,25 @@
 import argparse
 import os
+import random
 import signal
-from typing import Dict, Optional
+import sys
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
 import torchaudio
 from loguru import logger
 from torch import Tensor, nn, optim
+from torch.autograd.anomaly_mode import set_detect_anomaly
+from torch.autograd.grad_mode import set_grad_enabled
+from torch.nn.utils.clip_grad import clip_grad_norm_
 from torch.types import Number
 
-from df.checkpoint import load_model, read_cp, write_cp
+from df.checkpoint import check_patience, load_model, read_cp, write_cp
 from df.config import Csv, config
 from df.logger import init_logger, log_metrics, log_model_summary
 from df.loss import Istft, Loss, MaskLoss
-from df.lr import cosine_lr_scheduler
+from df.lr import cosine_scheduler
 from df.model import ModelParams
 from df.modules import get_device
 from df.utils import (
@@ -22,8 +27,8 @@
     as_real,
     check_finite_module,
     check_manual_seed,
-    clip_grad_norm_,
     detach_hidden,
+    get_host,
     get_norm_alpha,
     make_np,
 )
@@ -32,13 +37,14 @@
 
 should_stop = False
 debug = False
+log_timings = False
 state: Optional[DF] = None
 istft: Optional[nn.Module]
 
 
 @logger.catch
 def main():
-    global should_stop, debug, state
+    global should_stop, debug, state, log_timings
 
     parser = argparse.ArgumentParser()
     parser.add_argument("data_config_file", type=str, help="Path to a dataset config file.")
@@ -48,8 +54,16 @@ def main():
     parser.add_argument(
         "base_dir", type=str, help="Directory to store logs, summaries, checkpoints, etc."
     )
+    parser.add_argument(
+        "--host-batchsize-config",
+        "-b",
+        type=str,
+        default=None,
+        help="Path to a host specific batch size config.",
+    )
     parser.add_argument("--no-resume", action="store_false", dest="resume")
     parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--no-debug", action="store_false", dest="debug")
     args = parser.parse_args()
     if not os.path.isfile(args.data_config_file):
         raise FileNotFoundError("Dataset config not found at {}".format(args.data_config_file))
@@ -60,12 +74,28 @@ def main():
     os.makedirs(summary_dir, exist_ok=True)
     debug = args.debug
     log_level = "DEBUG" if debug else "INFO"
-    init_logger(file=os.path.join(args.base_dir, "train.log"), level=log_level)
-    config.load(os.path.join(args.base_dir, "config.ini"))
+    init_logger(file=os.path.join(args.base_dir, "train.log"), level=log_level, model=args.base_dir)
+    config_file = os.path.join(args.base_dir, "config.ini")
+    config.load(config_file)
     seed = config("SEED", 42, int, section="train")
     check_manual_seed(seed)
     logger.info("Running on device {}".format(get_device()))
 
+    # Maybe update batch size
+    if args.host_batchsize_config is not None:
+        try:
+            sys.path.append(
+                os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+            )
+            from scripts.set_batch_size import main as set_batch_size  # type: ignore
+
+            key = get_host() + "_" + config.get("model", section="train")
+            key += "_" + config.get("fft_size", section="df")
+            set_batch_size(config_file, args.host_batchsize_config, host_key=key)
+            config.load(config_file, allow_reload=True)  # Load again
+        except Exception as e:
+            logger.error(f"Could not apply host specific batch size config: {str(e)}")
+
     signal.signal(signal.SIGUSR1, get_sigusr1_handler(args.base_dir))
 
     p = ModelParams()
@@ -88,24 +118,12 @@ def main():
         mask_only=mask_only,
         train_df_only=train_df_only,
     )
-    opt = load_opt(
-        checkpoint_dir if args.resume else None,
-        model,
-        mask_only,
-        train_df_only,
-    )
-    try:
-        log_model_summary(model, verbose=args.debug)
-    except Exception as e:
-        logger.warning(f"Failed to print model summary: {e}")
-    if jit:
-        # Load as jit after log_model_summary
-        model = torch.jit.script(model)
 
     bs: int = config("BATCH_SIZE", 1, int, section="train")
     bs_eval: int = config("BATCH_SIZE_EVAL", 0, int, section="train")
     bs_eval = bs_eval if bs_eval > 0 else bs
     overfit = config("OVERFIT", False, bool, section="train")
+    log_timings = config("LOG_TIMINGS", False, bool, section="train", save=False)
     dataloader = DataLoader(
         ds_dir=args.data_dir,
         ds_config=args.data_config_file,
@@ -118,33 +136,60 @@ def main():
         hop_size=p.hop_size,
         nb_erb=p.nb_erb,
         nb_spec=p.nb_df,
-        norm_alpha=get_norm_alpha(),
-        p_atten_lim=config("p_atten_lim", 0.2, float, section="train"),
+        norm_alpha=get_norm_alpha(log=False),
         p_reverb=config("p_reverb", 0.2, float, section="train"),
-        prefetch=10,
+        prefetch=config("NUM_PREFETCH_BATCHES", 32, int, section="train"),
         overfit=overfit,
         seed=seed,
         min_nb_erb_freqs=p.min_nb_freqs,
+        log_timings=log_timings,
+        global_sampling_factor=config("GLOBAL_DS_SAMPLING_F", 1.0, float, section="train"),
+        snrs=config("DATALOADER_SNRS", [-5, 0, 5, 10, 20, 40], Csv(int), section="train"),  # type: ignore
+        cache_valid=config("VALIDATION_SET_CACHING", False, bool, section="train"),
     )
 
+    # Batch size scheduling limits the batch size for the first epochs. It will increase the batch
+    # size during training as specified. Used format is a comma separated list containing
+    # epoch/batch size tuples where each tuple is separated via '/':
+    # '<epoch>/<batch_size>,<epoch>/<batch_size>,<epoch>/<batch_size>'
+    # The first epoch has to be 0, later epoch may modify the batch size as specified.
+    # This only applies to training batch size.
+    batch_size_scheduling: List[str] = config("BATCH_SIZE_SCHEDULING", [], Csv(str), section="train")  # type: ignore
+    scheduling_bs = bs
+    prev_scheduling_bs = bs
+    if len(batch_size_scheduling) > 0:
+        batch_size_scheduling = [
+            (int(bs[0]), int(bs[1])) for bs in (bs.split("/") for bs in batch_size_scheduling)
+        ]
+        assert batch_size_scheduling[0][0] == 0  # First epoch must be 0
+        logger.info("Running with learning rate scheduling")
+
     max_epochs = config("MAX_EPOCHS", 10, int, section="train")
     assert epoch >= 0
-    n_warmup = 3
-    last_epoch = max(0, epoch) - 1
-    lr = config.get("lr", "train", float)
-    lrs = optim.lr_scheduler.LambdaLR(
-        opt,
-        cosine_lr_scheduler(
-            lr,
-            n_warmup=n_warmup,
-            warmup_init_lr=0.1 * lr,
-            min_lr=1e-6,
-            t_mult=1.5,
-            lr_period_updates=10,
-            lr_shrink=0.5,
-        ),
-        last_epoch=last_epoch,
+    opt = load_opt(
+        checkpoint_dir if args.resume else None,
+        model,
+        mask_only,
+        train_df_only,
     )
+    lrs = setup_lrs(len(dataloader))
+    wds = setup_wds(len(dataloader))
+    if not args.resume and os.path.isfile(os.path.join(checkpoint_dir, ".patience")):
+        os.remove(os.path.join(checkpoint_dir, ".patience"))
+    try:
+        log_model_summary(model, verbose=args.debug)
+    except Exception as e:
+        logger.warning(f"Failed to print model summary: {e}")
+    if jit:
+        # Load as jit after log_model_summary
+        model = torch.jit.script(model)
+
+    # Validation optimization target. Used for early stopping and selecting best checkpoint
+    val_criteria = []
+    val_criteria_type = config("VALIDATION_CRITERIA", "loss", section="train")  # must be in metrics
+    val_criteria_rule = config("VALIDATION_CRITERIA_RULE", "min", section="train")
+    val_criteria_rule = val_criteria_rule.replace("less", "min").replace("more", "max")
+    patience = config("EARLY_STOPPING_PATIENCE", 5, int, section="train")
 
     losses = setup_losses()
 
@@ -167,6 +212,19 @@ def main():
     # Save default values to disk
     config.save(os.path.join(args.base_dir, "config.ini"))
     for epoch in range(epoch, max_epochs):
+        if len(batch_size_scheduling) > 0:
+            # Get current batch size
+            for (e, b) in batch_size_scheduling:
+                if e <= epoch:
+                    # Update bs, but don't go higher than the batch size specified in the config
+                    scheduling_bs = min(b, bs)
+            if prev_scheduling_bs != scheduling_bs:
+                logger.info(f"Batch scheduling | Setting batch size to {scheduling_bs}")
+                dataloader.set_batch_size(scheduling_bs, "train")
+                # Update lr/wd scheduling since dataloader len changed
+                lrs = setup_lrs(len(dataloader))
+                wds = setup_wds(len(dataloader))
+                prev_scheduling_bs = scheduling_bs
         train_loss = run_epoch(
             model=model,
             epoch=epoch,
@@ -175,10 +233,12 @@ def main():
             opt=opt,
             losses=losses,
             summary_dir=summary_dir,
+            lr_scheduler_values=lrs,
+            wd_scheduler_values=wds,
         )
         metrics = {"loss": train_loss}
         try:
-            metrics["lr"] = lrs.get_last_lr()[0]
+            metrics["lr"] = opt.param_groups[0]["lr"]
         except AttributeError:
             pass
         if debug:
@@ -202,13 +262,26 @@ def main():
         metrics.update(
             {n: torch.mean(torch.stack(vals)).item() for n, vals in losses.get_summaries()}
         )
+        val_criteria = metrics[val_criteria_type]
+        write_cp(
+            model, "model", checkpoint_dir, epoch + 1, metric=val_criteria, cmp=val_criteria_rule
+        )
         log_metrics(f"[{epoch}] [valid]", metrics)
+        check_patience(
+            checkpoint_dir, max_patience=patience, new_metric=val_criteria, cmp=val_criteria_rule
+        )
         if should_stop:
-            logger.info("Stopping training")
+            logger.info("Stopping training due to timeout")
             exit(0)
         losses.reset_summaries()
-        lrs.step()
-        ic([group["lr"] for group in opt.param_groups])
+    model, epoch = load_model(
+        checkpoint_dir,
+        state,
+        jit=jit,
+        epoch="best",
+        mask_only=mask_only,
+        train_df_only=train_df_only,
+    )
     test_loss = run_epoch(
         model=model,
         epoch=epoch,
@@ -232,6 +305,8 @@ def run_epoch(
     opt: optim.Optimizer,
     losses: Loss,
     summary_dir: str,
+    lr_scheduler_values: Optional[np.ndarray] = None,
+    wd_scheduler_values: Optional[np.ndarray] = None,
 ) -> float:
     global debug
 
@@ -248,32 +323,40 @@ def run_epoch(
     dev = get_device()
     l_mem = []
     is_train = split == "train"
-    summary_fn = summary_write  # or summary_noop
     model.train(mode=is_train)
     losses.store_losses = debug or not is_train
     max_steps = loader.len(split) - 1
     seed = epoch if is_train else 42
     n_nans = 0
-    logger.info("Dataloader len: {}".format(loader.len(split)))
+    start_steps = epoch * loader.len(split)
 
     for i, batch in enumerate(loader.iter_epoch(split, seed)):
         opt.zero_grad()
+        it = start_steps + i  # global training iteration
+        if lr_scheduler_values is not None or wd_scheduler_values is not None:
+            for param_group in opt.param_groups:
+                if lr_scheduler_values is not None:
+                    param_group["lr"] = lr_scheduler_values[it] * param_group.get("lr_scale", 1)
+                if wd_scheduler_values is not None and param_group["weight_decay"] > 0:
+                    param_group["weight_decay"] = wd_scheduler_values[it]
         assert batch.feat_spec is not None
         assert batch.feat_erb is not None
         feat_erb = batch.feat_erb.to(dev, non_blocking=True)
         feat_spec = as_real(batch.feat_spec.to(dev, non_blocking=True))
         noisy = batch.noisy.to(dev, non_blocking=True)
         clean = batch.speech.to(dev, non_blocking=True)
-        atten = batch.atten.to(dev, non_blocking=True)
         snrs = batch.snr.to(dev, non_blocking=True)
-        with torch.autograd.set_detect_anomaly(detect_anomaly):
-            with torch.set_grad_enabled(is_train):
-                enh, m, lsnr, df_alpha = model.forward(
-                    spec=as_real(noisy),
-                    feat_erb=feat_erb,
-                    feat_spec=feat_spec,
-                    atten_lim=atten,
-                )
+        with set_detect_anomaly(detect_anomaly and is_train), set_grad_enabled(is_train):
+            enh, m, lsnr, other = model.forward(
+                spec=as_real(noisy),
+                feat_erb=feat_erb,
+                feat_spec=feat_spec,
+            )
+            df_alpha, multi_stage_specs = None, []
+            if isinstance(other, Tensor):
+                df_alpha = other
+            elif isinstance(other, (list, tuple)):
+                multi_stage_specs = other
             try:
                 err = losses.forward(
                     clean,
@@ -284,6 +367,7 @@ def run_epoch(
                     df_alpha=df_alpha,
                     max_freq=batch.max_freq,
                     snrs=snrs,
+                    multi_stage_specs=multi_stage_specs,
                 )
             except Exception as e:
                 if "nan" in str(e).lower() or "finite" in str(e).lower():
@@ -318,6 +402,13 @@ def run_epoch(
             if torch.isnan(l_mean):
                 check_finite_module(model)
             l_dict = {"loss": l_mean.item()}
+            if lr_scheduler_values is not None:
+                l_dict["lr"] = opt.param_groups[0]["lr"]
+            if wd_scheduler_values is not None:
+                l_dict["wd"] = opt.param_groups[0]["weight_decay"]
+            if log_timings:
+                l_dict["t_sample"] = batch.timings[:-1].sum()
+                l_dict["t_batch"] = batch.timings[-1].mean()  # last if for whole batch
             if debug:
                 l_dict.update(
                     {
@@ -325,13 +416,16 @@ def run_epoch(
                         for n, vals in losses.get_summaries()
                     }
                 )
-            log_metrics(f"[{epoch}] [{i}/{max_steps}]", l_dict)
-            summary_fn(
-                clean,
-                noisy,
-                enh,
-                batch.snr,
-                lsnr,
+            step = str(i).rjust(len(str(max_steps)))
+            log_metrics(f"[{epoch}] [{step}/{max_steps}]", l_dict)
+            if df_alpha is not None:
+                df_alpha.detach().float()
+            summary_write(
+                clean.detach(),
+                noisy.detach(),
+                enh.detach(),
+                batch.snr.detach(),
+                lsnr.detach().float(),
                 df_alpha,
                 summary_dir,
                 mask_loss=losses.ml,
@@ -361,11 +455,13 @@ def setup_losses() -> Loss:
 def load_opt(
     cp_dir: Optional[str], model: nn.Module, mask_only: bool = False, df_only: bool = False
 ) -> optim.Optimizer:
-    lr = config("LR", 1e-4, float, section="train")
-    decay = config("WEIGHT_DECAY", 1e-3, float, section="train")
-    optimizer = config("OPTIMIZER", "adam", str, section="train").lower()
-    momentum = config("MOMENTUM", 0, float, section="train", save=False)  # For sgd, rmsprop
-    betas = config("BETAS", [0.9, 0.98], Csv(float), section="train", save=False)  # For adam(w)
+    lr = config("LR", 5e-4, float, section="optim")
+    momentum = config("momentum", 0, float, section="optim")  # For sgd, rmsprop
+    decay = config("weight_decay", 0.05, float, section="optim")
+    optimizer = config("optimizer", "adamw", str, section="optim").lower()
+    betas: Tuple[int, int] = config(
+        "opt_betas", [0.9, 0.999], Csv(float), section="optim", save=False  # type: ignore
+    )
     if mask_only:
         params = []
         for n, p in model.named_parameters():
@@ -376,10 +472,10 @@ def load_opt(
     else:
         params = model.parameters()
     supported = {
-        "adam": lambda p: optim.AdamW(p, lr=lr, weight_decay=decay, betas=betas),
+        "adam": lambda p: optim.Adam(p, lr=lr, weight_decay=decay, betas=betas),
         "adamw": lambda p: optim.AdamW(p, lr=lr, weight_decay=decay, betas=betas),
         "sgd": lambda p: optim.SGD(p, lr=lr, momentum=momentum, nesterov=True, weight_decay=decay),
-        "rmsprop": lambda p: optim.RMSprop(params, lr=lr, momentum=momentum, weight_decay=decay),
+        "rmsprop": lambda p: optim.RMSprop(p, lr=lr, momentum=momentum, weight_decay=decay),
     }
     if optimizer not in supported:
         raise ValueError(
@@ -397,6 +493,47 @@ def load_opt(
     return opt
 
 
+def setup_lrs(steps_per_epoch: int) -> np.ndarray:
+    lr = config.get("lr", float, "optim")
+    num_epochs = config.get("max_epochs", int, "train")
+    lr_min = config("lr_min", 1e-6, float, section="optim")
+    lr_warmup = config("lr_warmup", 1e-4, float, section="optim")
+    assert lr_warmup < lr
+    warmup_epochs = config("warmup_epochs", 3, int, section="optim")
+    lr_cycle_mul = config("lr_cycle_mul", 1.0, float, section="optim")
+    lr_cycle_decay = config("lr_cycle_decay", 0.5, float, section="optim")
+    lr_cycle_epochs = config("lr_cycle_epochs", -1, int, section="optim")
+    lr_values = cosine_scheduler(
+        lr,
+        lr_min,
+        epochs=num_epochs,
+        niter_per_ep=steps_per_epoch,
+        warmup_epochs=warmup_epochs,
+        start_warmup_value=lr_warmup,
+        initial_ep_per_cycle=lr_cycle_epochs,
+        cycle_decay=lr_cycle_decay,
+        cycle_mul=lr_cycle_mul,
+    )
+    return lr_values
+
+
+def setup_wds(steps_per_epoch: int) -> Optional[np.ndarray]:
+    decay = config("weight_decay", 0.05, float, section="optim")
+    decay_end = config("weight_decay_end", -1, float, section="optim")
+    if decay_end == -1:
+        return None
+    if decay == 0.0:
+        decay = 1e-12
+        logger.warning("Got 'weight_decay_end' value > 0, but weight_decay is disabled.")
+        logger.warning(f"Setting initial weight decay to {decay}.")
+        config.overwrite("optim", "weight_decay", decay)
+    num_epochs = config.get("max_epochs", int, "train")
+    decay_values = cosine_scheduler(
+        decay, decay_end, niter_per_ep=steps_per_epoch, epochs=num_epochs
+    )
+    return decay_values
+
+
 @torch.no_grad()
 def summary_write(
     clean: Tensor,
@@ -404,7 +541,7 @@ def summary_write(
     enh: Tensor,
     snrs: Tensor,
     lsnr: Tensor,
-    df_alpha: Tensor,
+    df_alpha: Optional[Tensor],
     summary_dir: str,
     mask_loss: Optional[MaskLoss] = None,
     split="train",
@@ -413,33 +550,38 @@ def summary_write(
     assert state is not None
 
     p = ModelParams()
-    snr = snrs[0].detach().cpu().item()
+    bs = snrs.shape[0]
+    idx = random.randrange(bs)
+    snr = snrs[idx].detach().cpu().item()
 
     def synthesis(x: Tensor) -> Tensor:
         return torch.as_tensor(state.synthesis(make_np(as_complex(x.detach()))))
 
     if mask_loss is not None:
-        ideal = mask_loss.erb_mask_compr(clean[0], noisy[0], compressed=False)
-        ideal = noisy[0] * mask_loss.erb_inv(ideal)
+        ideal = mask_loss.erb_mask_compr(clean[idx], noisy[idx], compressed=False)
+        ideal = noisy[idx] * mask_loss.erb_inv(ideal)
         torchaudio.save(
             os.path.join(summary_dir, f"{split}_idealmask_snr{snr}.wav"), synthesis(ideal), p.sr
         )
     torchaudio.save(
-        os.path.join(summary_dir, f"{split}_clean_snr{snr}.wav"), synthesis(clean[0]), p.sr
+        os.path.join(summary_dir, f"{split}_clean_snr{snr}.wav"), synthesis(clean[idx]), p.sr
+    )
+    torchaudio.save(
+        os.path.join(summary_dir, f"{split}_noisy_snr{snr}.wav"), synthesis(noisy[idx]), p.sr
     )
     torchaudio.save(
-        os.path.join(summary_dir, f"{split}_noisy_snr{snr}.wav"), synthesis(noisy[0]), p.sr
+        os.path.join(summary_dir, f"{split}_enh_snr{snr}.wav"), synthesis(enh[idx]), p.sr
     )
-    torchaudio.save(os.path.join(summary_dir, f"{split}_enh_snr{snr}.wav"), synthesis(enh[0]), p.sr)
     np.savetxt(
         os.path.join(summary_dir, f"{split}_lsnr_snr{snr}.txt"),
-        lsnr[0].detach().cpu().numpy(),
+        lsnr[idx].detach().cpu().numpy(),
         fmt="%.3f",
     )
-    np.savetxt(
-        os.path.join(summary_dir, f"{split}_df_alpha_snr{snr}.txt"),
-        df_alpha[0].detach().cpu().numpy(),
-    )
+    if df_alpha is not None:
+        np.savetxt(
+            os.path.join(summary_dir, f"{split}_df_alpha_snr{snr}.txt"),
+            df_alpha[idx].detach().cpu().numpy(),
+        )
 
 
 def summary_noop(*__args, **__kwargs):  # type: ignore
diff --git a/DeepFilterNet/df/utils.py b/DeepFilterNet/df/utils.py
index 4dbe2c1ad..f733a3ece 100644
--- a/DeepFilterNet/df/utils.py
+++ b/DeepFilterNet/df/utils.py
@@ -3,15 +3,15 @@
 import os
 import random
 import subprocess
-import warnings
 from socket import gethostname
-from typing import Any, Set, Union
+from typing import Any, Dict, Set, Tuple, Union
 
 import numpy as np
 import torch
 from loguru import logger
 from torch import Tensor
 from torch._six import string_classes
+from torch.autograd import Function
 from torch.types import Number
 
 from df.config import config
@@ -23,7 +23,7 @@
     from torchaudio.compliance.kaldi import resample_waveform as ta_resample  # type: ignore
 
 
-def resample(audio: Tensor, orig_sr: int, new_sr: int, method="sinc_fast"):
+def get_resample_params(method: str) -> Dict[str, Any]:
     params = {
         "sinc_fast": {"resampling_method": "sinc_interpolation", "lowpass_filter_width": 16},
         "sinc_best": {"resampling_method": "sinc_interpolation", "lowpass_filter_width": 64},
@@ -40,8 +40,13 @@ def resample(audio: Tensor, orig_sr: int, new_sr: int, method="sinc_fast"):
             "beta": 14.769656459379492,
         },
     }
-    assert method in params.keys()
-    return ta_resample(audio, orig_sr, new_sr, **params[method])
+    assert method in params.keys(), f"method must be one of {list(params.keys())}"
+    return params[method]
+
+
+def resample(audio: Tensor, orig_sr: int, new_sr: int, method="sinc_fast"):
+    params = get_resample_params(method)
+    return ta_resample(audio, orig_sr, new_sr, **params)
 
 
 def get_device():
@@ -72,6 +77,36 @@ def as_real(x: Tensor):
     return x
 
 
+class angle_re_im(Function):
+    """Similar to torch.angle but robustify the gradient for zero magnitude."""
+
+    @staticmethod
+    def forward(ctx, re: Tensor, im: Tensor):
+        ctx.save_for_backward(re, im)
+        return torch.atan2(im, re)
+
+    @staticmethod
+    def backward(ctx, grad: Tensor) -> Tuple[Tensor, Tensor]:
+        re, im = ctx.saved_tensors
+        grad_inv = grad / (re.square() + im.square()).clamp_min_(1e-10)
+        return -im * grad_inv, re * grad_inv
+
+
+class angle(Function):
+    """Similar to torch.angle but robustify the gradient for zero magnitude."""
+
+    @staticmethod
+    def forward(ctx, x: Tensor):
+        ctx.save_for_backward(x)
+        return torch.atan2(x.imag, x.real)
+
+    @staticmethod
+    def backward(ctx, grad: Tensor):
+        (x,) = ctx.saved_tensors
+        grad_inv = grad / (x.real.square() + x.imag.square()).clamp_min_(1e-10)
+        return torch.view_as_complex(torch.stack((-x.imag * grad_inv, x.real * grad_inv), dim=-1))
+
+
 def check_finite_module(obj, name="Module", _raise=True) -> Set[str]:
     out: Set[str] = set()
     if isinstance(obj, torch.nn.Module):
@@ -166,70 +201,6 @@ def get_branch_name():
     return branch
 
 
-def clip_grad_norm_(
-    parameters,
-    max_norm: float,
-    norm_type: float = 2.0,
-    error_if_nonfinite: bool = False,
-) -> torch.Tensor:
-    r"""Pytorch 1.9 backport: Clips gradient norm of an iterable of parameters.
-
-    The norm is computed over all gradients together, as if they were
-    concatenated into a single vector. Gradients are modified in-place.
-
-    Args:
-        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
-            single Tensor that will have gradients normalized
-        max_norm (float or int): max norm of the gradients
-        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
-            infinity norm.
-        error_if_nonfinite (bool): if True, an error is thrown if the total
-            norm of the gradients from :attr:``parameters`` is ``nan``,
-            ``inf``, or ``-inf``. Default: False (will switch to True in the future)
-
-    Returns:
-        Total norm of the parameters (viewed as a single vector).
-    """
-    if isinstance(parameters, torch.Tensor):
-        parameters = [parameters]
-    parameters = [p for p in parameters if p.grad is not None]
-    max_norm = float(max_norm)
-    norm_type = float(norm_type)
-    if len(parameters) == 0:
-        return torch.tensor(0.0)
-    device = parameters[0].grad.device
-    if norm_type == torch._six.inf:
-        norms = [p.grad.detach().abs().max().to(device) for p in parameters]
-        total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
-    else:
-        total_norm = torch.norm(
-            torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]),
-            norm_type,
-        )
-    if total_norm.isnan() or total_norm.isinf():
-        if error_if_nonfinite:
-            raise RuntimeError(
-                f"The total norm of order {norm_type} for gradients from "
-                "`parameters` is non-finite, so it cannot be clipped. To disable "
-                "this error and scale the gradients by the non-finite norm anyway, "
-                "set `error_if_nonfinite=False`"
-            )
-        else:
-            warnings.warn(
-                "Non-finite norm encountered in torch.nn.utils.clip_grad_norm_; continuing anyway. "
-                "Note that the default behavior will change in a future release to error out "
-                "if a non-finite total norm is encountered. At that point, setting "
-                "error_if_nonfinite=false will be required to retain the old behavior.",
-                FutureWarning,
-                stacklevel=2,
-            )
-    clip_coef = max_norm / (total_norm + 1e-6)
-    if clip_coef < 1:
-        for p in parameters:
-            p.grad.detach().mul_(clip_coef.to(p.grad.device))
-    return total_norm
-
-
 # from pytorch/ignite:
 def apply_to_tensor(input_, func):
     """Apply a function on a tensor or mapping, or sequence of tensors."""
diff --git a/DeepFilterNet/pretrained_models/DeepFilterNet2a/checkpoints/model_50.ckpt b/DeepFilterNet/pretrained_models/DeepFilterNet2a/checkpoints/model_50.ckpt
new file mode 100644
index 000000000..13a0bc007
Binary files /dev/null and b/DeepFilterNet/pretrained_models/DeepFilterNet2a/checkpoints/model_50.ckpt differ
diff --git a/DeepFilterNet/pretrained_models/DeepFilterNet2a/config.ini b/DeepFilterNet/pretrained_models/DeepFilterNet2a/config.ini
new file mode 100644
index 000000000..a757161b3
--- /dev/null
+++ b/DeepFilterNet/pretrained_models/DeepFilterNet2a/config.ini
@@ -0,0 +1,102 @@
+[train]
+seed = 43
+device = 
+model = deepfilternet2
+jit = true
+mask_only = false
+df_only = false
+batch_size = 96
+batch_size_eval = 128
+num_workers = 16
+max_sample_len_s = 3.0
+p_atten_lim = 0.0
+p_reverb = 0.2
+overfit = false
+start_eval = true
+max_epochs = 100
+log_freq = 50
+detect_anomaly = false
+validation_criteria = loss
+validation_criteria_rule = min
+early_stopping_patience = 10
+global_ds_sampling_f = 1
+num_prefetch_batches = 8
+dataloader_snrs = -5,0,5,10,20,40
+train_autocast = False
+n_checkpoint_history = 5
+n_best_checkpoint_history = 10
+batch_size_scheduling = 
+
+[df]
+sr = 48000
+fft_size = 960
+hop_size = 480
+nb_erb = 32
+nb_df = 96
+norm_tau = 1
+lsnr_max = 35
+lsnr_min = -15
+min_nb_erb_freqs = 2
+
+[deepfilternet]
+conv_lookahead = 2
+conv_ch = 64
+conv_depthwise = True
+emb_hidden_dim = 512
+emb_num_layers = 3
+gru_groups = 8
+linear_groups = 8
+conv_dec_mode = transposed
+convt_depthwise = True
+mask_pf = False
+df_order = 5
+df_lookahead = 2
+df_hidden_dim = 512
+df_num_layers = 2
+dfop_method = real_unfold
+group_shuffle = False
+conv_kernel = 2,3
+
+[localsnrloss]
+factor = 1e-3
+
+[maskloss]
+factor = 0
+mask = iam
+gamma = 0.6
+gamma_pred = 0.6
+f_under = 1
+
+[spectralloss]
+factor_magnitude = 1000
+factor_complex = 1000
+gamma = 0.3
+
+[dfalphaloss]
+factor = 0.0
+
+[multiresspecloss]
+factor = 500
+factor_complex = 500
+gamma = 0.3
+fft_sizes = 256,512,1024
+
+[optim]
+lr = 0.001
+momentum = 0
+weight_decay = 1e-12
+weight_decay_end = 0.05
+optimizer = adamw
+lr_min = 1e-06
+lr_warmup = 0.0001
+warmup_epochs = 3
+lr_cycle_mul = 1.0
+lr_cycle_decay = 0.5
+lr_cycle_limit = 1
+lr_update_per_epoch = False
+lr_cycle_epochs = -1
+
+[sdrloss]
+factor = 0.0
+segmental_ws = 0
+
diff --git a/DeepFilterNet/pretrained_models/DeepFilterNet2b/checkpoints/model_27.ckpt.best b/DeepFilterNet/pretrained_models/DeepFilterNet2b/checkpoints/model_27.ckpt.best
new file mode 100644
index 000000000..63edef63d
Binary files /dev/null and b/DeepFilterNet/pretrained_models/DeepFilterNet2b/checkpoints/model_27.ckpt.best differ
diff --git a/DeepFilterNet/pretrained_models/DeepFilterNet2b/config.ini b/DeepFilterNet/pretrained_models/DeepFilterNet2b/config.ini
new file mode 100644
index 000000000..f8e2a7894
--- /dev/null
+++ b/DeepFilterNet/pretrained_models/DeepFilterNet2b/config.ini
@@ -0,0 +1,107 @@
+[train]
+seed = 43
+device = 
+model = deepfilternet2
+jit = false
+mask_only = false
+df_only = false
+batch_size = 96
+batch_size_eval = 128
+num_workers = 16
+max_sample_len_s = 3.0
+p_atten_lim = 0.0
+p_reverb = 0.0
+overfit = false
+start_eval = true
+max_epochs = 100
+log_freq = 100
+log_timings = true
+validation_criteria = loss
+validation_criteria_rule = min
+early_stopping_patience = 10
+global_ds_sampling_f = 1
+num_prefetch_batches = 8
+dataloader_snrs = -5,0,5,10,20,40
+detect_anomaly = false
+batch_size_scheduling = 0/8,1/16,2/24,5/32,10/64,20/128,40/9999
+validation_set_caching = False
+
+[df]
+sr = 48000
+fft_size = 960
+hop_size = 480
+nb_erb = 32
+nb_df = 96
+norm_tau = 1
+lsnr_max = 35
+lsnr_min = -15
+min_nb_erb_freqs = 2
+
+[deepfilternet]
+conv_lookahead = 2
+conv_ch = 64
+conv_depthwise = True
+emb_hidden_dim = 256
+emb_num_layers = 3
+gru_groups = 8
+linear_groups = 8
+conv_dec_mode = transposed
+convt_depthwise = True
+mask_pf = False
+df_order = 5
+df_lookahead = 2
+df_hidden_dim = 256
+df_num_layers = 2
+dfop_method = real_unfold
+group_shuffle = False
+conv_kernel = 1,3
+df_gru_skip = none
+df_output_layer = groupedlinear
+gru_type = squeeze
+df_pathway_kernel_size_t = 5
+df_n_iter = 1
+enc_concat = False
+
+[localsnrloss]
+factor = 1e-3
+
+[maskloss]
+factor = 0
+mask = iam
+gamma = 0.6
+gamma_pred = 0.6
+f_under = 1
+
+[spectralloss]
+factor_magnitude = 1000
+factor_complex = 1000
+gamma = 0.3
+
+[dfalphaloss]
+factor = 0.0
+
+[multiresspecloss]
+factor = 500
+factor_complex = 500
+gamma = 0.3
+fft_sizes = 256,512,1024
+
+[optim]
+lr = 0.001
+momentum = 0
+weight_decay = 1e-12
+weight_decay_end = 0.05
+optimizer = adamw
+lr_min = 1e-06
+lr_warmup = 0.0001
+warmup_epochs = 3
+lr_cycle_mul = 1.0
+lr_cycle_decay = 0.5
+lr_cycle_limit = 1
+lr_update_per_epoch = False
+lr_cycle_epochs = -1
+
+[sdrloss]
+factor = 0.0
+segmental_ws = 0
+
diff --git a/DeepFilterNet/pyproject.toml b/DeepFilterNet/pyproject.toml
index a364aec91..54fe76068 100644
--- a/DeepFilterNet/pyproject.toml
+++ b/DeepFilterNet/pyproject.toml
@@ -30,22 +30,26 @@ loguru = ">=0.5"
 soundfile = { version = "^0.10", optional = true }
 icecream = { version = ">=2,<3", optional = true }
 pystoi = { version = "^0.3", optional = true }
+pesq = { version = "^0.0.3", optional = true }
+scipy = { version = "^1", optional = true }
 semetrics = {git = "https://github.com/usimarit/semetrics.git", optional = true}
 
 [tool.poetry.extras]
 train = ["deepfilterdataloader", "icecream"]
-eval = ["pystoi", "semetrics"]
+eval = ["pystoi", "pesq", "scipy"]
+semetrics = ["semetrics"]
 soundfile = ["soundfile"]
 
 [tool.poetry.scripts]
 deepFilter = "df.enhance:main"
 
 [tool.poetry.dev-dependencies]
-poethepoet = "^0.11.0"
+poethepoet = "^0.11"
+pre-commit = "^2.17"
 
 [tool.poe.tasks]
-install-torch-cuda111 = "python -m pip install torch==1.10.0 torchaudio==0.10.0 --extra-index-url https://download.pytorch.org/whl/cu111/"
-install-torch-cpu = "python -m pip install torch==1.10.0 torchaudio==0.10.0 --extra-index-url https://download.pytorch.org/whl/cpu/"
+install-torch-cuda111 = "python -m pip install torch==1.11.0 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113/"
+install-torch-cpu = "python -m pip install torch==1.11.0 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cpu/"
 
 [build-system]
 requires = ["poetry>=1.0"]
@@ -53,7 +57,7 @@ build-backend = "poetry.masonry.api"
 
 [tool.black]
 line-length = 100
-target-version = ["py37", "py38", "py39"]
+target-version = ["py37", "py38", "py39", "py310"]
 include = '\.pyi?$'
 
 [tool.isort]
diff --git a/DeepFilterNet/requirements.txt b/DeepFilterNet/requirements.txt
index bdbbfd85a..cc4ef474b 100644
--- a/DeepFilterNet/requirements.txt
+++ b/DeepFilterNet/requirements.txt
@@ -5,3 +5,4 @@ deepfilterlib
 icecream>=2
 loguru>=0.5
 numpy>=1.20
+ptflops
diff --git a/assets/dataset.cfg b/assets/dataset.cfg
index 70c61f41d..e2be96705 100644
--- a/assets/dataset.cfg
+++ b/assets/dataset.cfg
@@ -12,7 +12,7 @@
     "train": [
         [
             "clean.hdf5",
-            10000
+            100
         ],
         [
             "noise.hdf5",
diff --git a/libDF/Cargo.toml b/libDF/Cargo.toml
index 7cc321b31..f916dd4f0 100644
--- a/libDF/Cargo.toml
+++ b/libDF/Cargo.toml
@@ -12,6 +12,16 @@ readme = "../README.md"
 name = "df"
 path = "src/lib.rs"
 
+[[bin]]
+name = "sample-hdf5"
+path = "src/bin/sample-hdf5.rs"
+required-features = ["dataset", "anyhow", "rand"]
+
+[[bin]]
+name = "sample-dataset"
+path = "src/bin/sample-dataset.rs"
+required-features = ["dataset", "anyhow"]
+
 [features]
 default = []
 
@@ -27,8 +37,9 @@ rand = { version="0.8", optional=true }
 rubato = { version="0.11", optional=true }
 rand_xoshiro = { version="0.6", optional=true }
 thiserror = { version="1.0", optional=true }
+anyhow = { version="1.0", optional=true, features=["backtrace"] }
 hound = { version="3.4", optional=true }
-hdf5 = { version="^0.8", optional=true, git="https://github.com/aldanor/hdf5-rust/", rev="refs/pull/198/head" }
+hdf5 = { version="^0.8", optional=true, features=["blosc"], git="https://github.com/aldanor/hdf5-rust/", rev="refs/pull/198/head" }
 bincode = { version="^1.3", optional=true }
 ndarray = { version="^0.15", optional=true, features=["serde"] }
 rayon = { version="1.5", optional=true }
diff --git a/libDF/src/bin/sample-dataset.rs b/libDF/src/bin/sample-dataset.rs
new file mode 100644
index 000000000..448baf2f7
--- /dev/null
+++ b/libDF/src/bin/sample-dataset.rs
@@ -0,0 +1,26 @@
+use std::env::args;
+
+use anyhow::Result;
+use df::{
+    augmentations::seed_from_u64,
+    dataset::{Dataset, DatasetBuilder, DatasetConfigJson, Split},
+};
+
+fn main() -> Result<()> {
+    let args = args().collect::<Vec<String>>();
+    let cfg = args.get(1).expect("HDF5 config file expected");
+    let dir = args.get(2).expect("HDF5 dataset dir expected");
+    // let out = args.nth(3).unwrap_or_else(|| "out".to_string());
+    seed_from_u64(42);
+    let cfg = DatasetConfigJson::open(cfg).unwrap();
+    let mut ds_builder = DatasetBuilder::new(dir, 48000).df_params(960, None, None, None, None);
+    ds_builder = ds_builder.p_sample_full_speech(1.0);
+    let mut train_ds = ds_builder.dataset(cfg.split_config(Split::Train)).build_fft_dataset()?;
+    train_ds.generate_keys()?;
+
+    for idx in 0..train_ds.len() {
+        let sample = train_ds.get_sample(idx, Some(0)).unwrap();
+        dbg!(sample);
+    }
+    Ok(())
+}
diff --git a/libDF/src/bin/sample-hdf5.rs b/libDF/src/bin/sample-hdf5.rs
new file mode 100644
index 000000000..f46b7d3cc
--- /dev/null
+++ b/libDF/src/bin/sample-hdf5.rs
@@ -0,0 +1,22 @@
+use std::env::args;
+
+use anyhow::Result;
+use df::dataset::Hdf5Dataset;
+use df::wav_utils::write_wav_arr2;
+use rand::seq::SliceRandom;
+
+fn main() -> Result<()> {
+    let args = args().collect::<Vec<String>>();
+    let p = args.get(1).expect("HDF5 dataset path expected");
+    let ds = Hdf5Dataset::new(p)?;
+    let k = match args.get(2) {
+        Some(k) => k.to_string(),
+        None => ds.keys()?.choose(&mut rand::thread_rng()).unwrap().to_string(),
+    };
+    let data = ds.read(&k)?;
+    dbg!(&k, data.shape());
+    let name = "out/".to_owned() + &k;
+    println!("Sampled {:?}", k);
+    write_wav_arr2(&name, data.view(), ds.sr.unwrap() as u32).unwrap();
+    Ok(())
+}
diff --git a/libDF/src/dataloader.rs b/libDF/src/dataloader.rs
index 4f166b27b..65e60616a 100644
--- a/libDF/src/dataloader.rs
+++ b/libDF/src/dataloader.rs
@@ -5,6 +5,7 @@ use std::sync::mpsc::{sync_channel, Receiver};
 use std::sync::{Arc, Mutex};
 use std::thread;
 use std::time::Duration;
+use std::time::Instant;
 
 use crossbeam_channel::unbounded;
 use ndarray::prelude::*;
@@ -83,8 +84,8 @@ pub struct DataLoaderBuilder {
     _batch_size_eval: Option<usize>,
     _prefetch: Option<usize>,
     _num_threads: Option<usize>,
-    _drop_last: Option<bool>,
-    _overfit: Option<bool>,
+    _drop_last: bool,
+    _overfit: bool,
 }
 
 impl DataLoaderBuilder {
@@ -95,8 +96,8 @@ impl DataLoaderBuilder {
             _batch_size_eval: None,
             _prefetch: None,
             _num_threads: None,
-            _drop_last: None,
-            _overfit: None,
+            _drop_last: false,
+            _overfit: false,
         }
     }
     pub fn batch_size(mut self, batch_size: usize) -> Self {
@@ -115,44 +116,26 @@ impl DataLoaderBuilder {
         self._num_threads = Some(num_threads);
         self
     }
-    pub fn overfit(mut self, overfit: bool) -> Self {
-        self._overfit = Some(overfit);
+    pub fn overfit(mut self) -> Self {
+        self._overfit = true;
         self
     }
-    pub fn drop_last(mut self, drop_last: bool) -> Self {
-        self._drop_last = Some(drop_last);
+    pub fn drop_last(mut self) -> Self {
+        self._drop_last = true;
         self
     }
-    fn check_dataset_size(&self, bs_train: usize) -> Result<()> {
-        for split in [Split::Train, Split::Valid, Split::Test] {
-            let batch_size = match split {
-                Split::Train => bs_train,
-                _ => self._batch_size_eval.unwrap_or(bs_train),
-            };
-            let dataset_size = self._ds.as_ref().unwrap().get(split).len();
-            if dataset_size < batch_size {
-                return Err(DfDataloaderError::DatasetTooSmall {
-                    split,
-                    dataset_size,
-                    batch_size,
-                });
-            }
-        }
-        Ok(())
-    }
     pub fn build(self) -> Result<DataLoader> {
         let bs_train = self._batch_size.unwrap_or(1);
-        self.check_dataset_size(bs_train)?;
-        let prefetch = self._prefetch.unwrap_or(bs_train * self._num_threads.unwrap_or(4) * 2);
+        let prefetch = self._prefetch.unwrap_or(bs_train * self._num_threads.unwrap_or(4));
         let mut loader = DataLoader::new(
             self._ds.unwrap(),
             bs_train,
             self._batch_size_eval,
             prefetch,
             self._num_threads,
-            self._drop_last.unwrap_or(false),
+            self._drop_last,
         )?;
-        loader.overfit = self._overfit.unwrap_or(false);
+        loader.overfit = self._overfit;
         Ok(loader)
     }
 }
@@ -170,21 +153,15 @@ impl DataLoader {
         drop_last: bool,
     ) -> Result<Self> {
         // Register global rayon threadpool. It will only be used for data loader workers.
-        let mut poolbuilder = rayon::ThreadPoolBuilder::new();
+        hdf5::sync::sync(|| {});
         let num_workers = num_threads.unwrap_or_else(current_num_threads);
-        poolbuilder = poolbuilder.num_threads(num_workers);
-        match poolbuilder
+        hdf5::sync::sync(|| {});
+        rayon::ThreadPoolBuilder::new()
+            .num_threads(num_workers)
             .thread_name(|idx| format!("DataLoader Worker {}", idx))
+            .start_handler(|_| hdf5::sync::sync(|| {}))
             .build_global()
-        {
-            Ok(()) => (),
-            Err(e) => {
-                if e.to_string() != "The global thread pool has already been initialized." {
-                    return Err(e.into());
-                }
-                // else: already initialized, do not complain.
-            }
-        };
+            .unwrap_or(());
         let batch_size_eval = batch_size_eval.unwrap_or(batch_size_train);
         Ok(DataLoader {
             ds_train: Some(Arc::new(datasets.train)),
@@ -243,14 +220,21 @@ impl DataLoader {
     }
 
     pub fn batch_size<S: Into<Split>>(&self, split: S) -> usize {
-        let split = split.into();
-        if split == Split::Train {
+        if split.into() == Split::Train {
             self.batch_size_train
         } else {
             self.batch_size_eval
         }
     }
 
+    pub fn set_batch_size<S: Into<Split>>(&mut self, batch_size: usize, split: S) {
+        if split.into() == Split::Train {
+            self.batch_size_train = batch_size;
+        } else {
+            self.batch_size_eval = batch_size;
+        }
+    }
+
     pub fn start_idx_worker(
         &mut self,
         split: Split,
@@ -273,6 +257,7 @@ impl DataLoader {
         in_sender.send((0, -1)).expect("Could not send index");
 
         let worker_recievers: Vec<_> = (0..self.num_workers).map(|_| in_receiver.clone()).collect();
+        let overfit = self.overfit;
         let handle = thread::spawn(move || -> Result<()> {
             worker_recievers.par_iter().try_for_each(|r| {
                 while let Ok((sample_idx, ordering_idx)) = r.recv() {
@@ -281,7 +266,12 @@ impl DataLoader {
                         return Ok(());
                     }
                     assert!(ordering_idx >= 0);
-                    let sample = ds.get_sample(sample_idx, Some(epoch_seed));
+                    let seed = Some(if overfit {
+                        epoch_seed
+                    } else {
+                        epoch_seed + sample_idx as u64
+                    });
+                    let sample = ds.get_sample(sample_idx, seed);
                     out_sender.send((ordering_idx as usize, sample.map_err(|e| e.into())))?;
                 }
                 Ok(())
@@ -349,7 +339,9 @@ impl DataLoader {
     where
         C: Collate<Complex32>,
     {
+        let mut t0 = Instant::now();
         let bs = self.batch_size(self.current_split);
+        let mut timings = Vec::with_capacity(bs);
         let mut samples = Vec::with_capacity(bs);
         let target_idx = self.dataset_len(self.current_split).min(self.cur_out_idx + bs);
         if self.cur_out_idx >= self.dataset_len(self.current_split) {
@@ -368,6 +360,9 @@ impl DataLoader {
             if let Some(s) = self.out_buf.remove(&self.cur_out_idx) {
                 ids.push(self.cur_out_idx);
                 samples.push(s);
+                let t1 = Instant::now();
+                timings.push((t1 - t0).as_secs_f32());
+                t0 = t1;
                 self.cur_out_idx += 1;
             } else {
                 // Or check worker threads
@@ -388,6 +383,9 @@ impl DataLoader {
                     Ok((o_idx, Ok(s))) => {
                         if o_idx == self.cur_out_idx {
                             samples.push(s);
+                            let t1 = Instant::now();
+                            timings.push((t1 - t0).as_secs_f32());
+                            t0 = t1;
                             ids.push(o_idx);
                             self.cur_out_idx += 1;
                         } else {
@@ -414,6 +412,7 @@ impl DataLoader {
             if !self.drained && self.cur_out_idx < target_idx {
                 debug_assert_eq!(batch.batch_size(), self.batch_size(self.current_split));
             }
+            batch.timings = timings;
             Some(batch)
         };
         Ok(out)
@@ -453,7 +452,6 @@ impl Collate<f32> for f32 {
         let max_freq = samples.iter().map(|s| s.max_freq).collect();
         let snr = samples.iter().map(|s| s.snr).collect();
         let gain = samples.iter().map(|s| s.gain).collect();
-        let atten = samples.iter().map(|s| s.attenuation.unwrap_or(0)).collect();
         Ok(DsBatch {
             speech,
             noise,
@@ -464,8 +462,8 @@ impl Collate<f32> for f32 {
             max_freq,
             snr,
             gain,
-            atten,
             ids: Vec::new(),
+            timings: Vec::new(),
         })
     }
 }
@@ -496,7 +494,6 @@ impl Collate<Complex32> for Complex32 {
         let max_freq = samples.iter().map(|s| s.max_freq).collect();
         let snr = samples.iter().map(|s| s.snr).collect();
         let gain = samples.iter().map(|s| s.gain).collect();
-        let atten = samples.iter().map(|s| s.attenuation.unwrap_or(0)).collect();
         Ok(DsBatch {
             speech,
             noise,
@@ -507,8 +504,8 @@ impl Collate<Complex32> for Complex32 {
             max_freq,
             snr,
             gain,
-            atten,
             ids: Vec::new(),
+            timings: Vec::new(),
         })
     }
 }
@@ -516,6 +513,21 @@ impl Collate<Complex32> for Complex32 {
 impl Drop for DataLoader {
     fn drop(&mut self) {
         self.join_fill_thread().unwrap(); // Stop out_receiver and join fill thread
+        for split in Split::iter() {
+            let mut ds = match Arc::try_unwrap(
+                match split {
+                    Split::Train => self.ds_train.take(),
+                    Split::Valid => self.ds_valid.take(),
+                    Split::Test => self.ds_test.take(),
+                }
+                .unwrap(),
+            ) {
+                Ok(ds) => ds,
+                Err(_) => panic!("Could not regain ownership over dataset"),
+            };
+            ds.close_cache().unwrap();
+            self.set_ds(split, ds);
+        }
     }
 }
 
@@ -532,8 +544,8 @@ where
     pub max_freq: Array1<usize>,
     pub snr: Vec<i8>,
     pub gain: Vec<i8>,
-    pub atten: Vec<u8>, // attenuation limit in dB; 0 stands for no limit
     pub ids: Vec<usize>,
+    pub timings: Vec<f32>,
 }
 impl<T> DsBatch<T>
 where
@@ -608,7 +620,6 @@ mod tests {
         let sr = 48000;
         let ds_dir = "../assets/";
         let mut cfg = DatasetConfigJson::open("../assets/dataset.cfg")?;
-        let split = Split::Train;
         let builder = DatasetBuilder::new(ds_dir, sr)
             .df_params(fft_size, hop_size, nb_erb, nb_spec, norm_alpha)
             .max_len(1.);
@@ -617,6 +628,14 @@ mod tests {
                 c.1 = dataset_size as f32; // Set sampling factor
                 assert_eq!(c.sampling_factor(), dataset_size as f32);
             }
+            for c in cfg.valid.iter_mut() {
+                c.1 = dataset_size as f32; // Set sampling factor
+                assert_eq!(c.sampling_factor(), dataset_size as f32);
+            }
+            for c in cfg.test.iter_mut() {
+                c.1 = dataset_size as f32; // Set sampling factor
+                assert_eq!(c.sampling_factor(), dataset_size as f32);
+            }
             'inner: for batch_size in [1, 2, 16] {
                 let ds = Datasets {
                     train: builder
@@ -632,7 +651,6 @@ mod tests {
                         .dataset(cfg.split_config(Split::Valid))
                         .build_fft_dataset()?,
                 };
-                dbg!(dataset_size, batch_size);
                 let mut loader = match DataLoader::builder(ds)
                     .num_threads(1)
                     .batch_size(batch_size)
@@ -658,19 +676,22 @@ mod tests {
                         e => return Err(e),
                     },
                 };
-                for epoch in 0..2 {
-                    loader.start_epoch(split, epoch)?;
-                    let mut n_samples = 0;
-                    loop {
-                        let batch = loader.get_batch::<Complex32>()?;
-                        if let Some(batch) = batch {
+                for split in Split::iter() {
+                    for epoch in 0..2 {
+                        println!(
+                        "***** Test: Loader with dataset_size {}, batch_size {}, epoch {} ******",
+                        dataset_size, batch_size, epoch
+                    );
+                        loader.start_epoch(split, epoch)?;
+                        let mut n_samples = 0;
+                        dbg!(dataset_size);
+                        while let Some(batch) = loader.get_batch::<Complex32>()? {
                             n_samples += batch.batch_size();
                             dbg!(n_samples);
-                        } else {
-                            break;
+                            assert!(n_samples <= dataset_size);
                         }
+                        dbg!(n_samples, dataset_size);
                     }
-                    assert_eq!(n_samples, dataset_size);
                 }
             }
         }
diff --git a/libDF/src/dataset.rs b/libDF/src/dataset.rs
index fbf1b9301..3cb7c26c2 100644
--- a/libDF/src/dataset.rs
+++ b/libDF/src/dataset.rs
@@ -1,17 +1,32 @@
+use std::collections::hash_map::DefaultHasher;
 use std::fmt;
 use std::fs;
+use std::hash::Hash;
+use std::hash::Hasher;
 use std::io::BufReader;
+use std::io::BufWriter;
+use std::ops::Deref;
 use std::ops::Range;
 use std::path::Path;
+use std::sync::{
+    mpsc::{sync_channel, SyncSender},
+    Arc,
+};
+use std::thread;
+use std::time::SystemTime;
 
 #[cfg(feature = "flac")]
 use claxon;
-use hdf5::{types::VarLenUnicode, File};
+use crossbeam_channel::Sender;
+use hdf5::Location;
+use hdf5::{filters, types::VarLenUnicode, File, Group, H5Type};
 use ndarray::{prelude::*, Slice};
 use rand::prelude::{IteratorRandom, SliceRandom};
 use rand::Rng;
+use rayon::prelude::*;
 use realfft::num_traits::Zero;
 use serde::Deserialize;
+use serde::Serialize;
 use thiserror::Error;
 #[cfg(feature = "vorbis")]
 use {lewton::inside_ogg::OggStreamReader, ogg::reading::PacketReader as OggPacketReader};
@@ -31,7 +46,7 @@ pub enum DfDatasetError {
     #[error("Unsupported during PCM decode: {0}")]
     PcmUnspportedDimension(usize),
     #[error("Wav Reader Error")]
-    WarReadError(#[from] crate::wav_utils::WavUtilsError),
+    WavReadError(#[from] crate::wav_utils::WavUtilsError),
     #[error("Input Range ({range:?}) larger than dataset size ({size:?})")]
     PcmRangeToLarge {
         range: Range<usize>,
@@ -64,6 +79,32 @@ pub enum DfDatasetError {
     #[cfg(feature = "flac")]
     #[error("Flac Decode Error")]
     FlacError(#[from] claxon::Error),
+    #[error("Multithreading Send Error: {0:?}")]
+    SendError(String),
+    #[error("Multithreading Recv Error: {0:?}")]
+    RecvError(String),
+    #[error("Thread Join Error: {0:?}")]
+    ThreadJoinError(String),
+    #[error("Crossbeam Multithreading Send Error: {0:?}")]
+    CrossbeamSendError(String),
+}
+
+impl<T> From<crossbeam_channel::SendError<T>> for DfDatasetError {
+    fn from(error: crossbeam_channel::SendError<T>) -> Self {
+        DfDatasetError::CrossbeamSendError(error.to_string())
+    }
+}
+
+impl From<std::sync::mpsc::RecvError> for DfDatasetError {
+    fn from(error: std::sync::mpsc::RecvError) -> Self {
+        DfDatasetError::RecvError(error.to_string())
+    }
+}
+
+impl<T> From<std::sync::mpsc::SendError<T>> for DfDatasetError {
+    fn from(error: std::sync::mpsc::SendError<T>) -> Self {
+        DfDatasetError::SendError(error.to_string())
+    }
 }
 
 type Signal = Array2<f32>;
@@ -71,12 +112,22 @@ type Signal = Array2<f32>;
 fn one() -> f32 {
     1.
 }
+#[derive(Hash, Debug, Clone)]
+pub struct DatasetModified(SystemTime, u64); // Modified time and file size
+impl DatasetModified {
+    fn new(file_name: &str) -> Result<Self> {
+        let meta_data = fs::metadata(file_name)?;
+        Ok(DatasetModified(meta_data.modified()?, meta_data.len()))
+    }
+}
 #[derive(Deserialize, Debug, Clone)]
 pub struct Hdf5Cfg(
-    pub String,                                              // file name
-    #[serde(default = "one")] pub f32,                       // dataset sampling factor
-    #[serde(default = "Option::default")] pub Option<usize>, // fallback sampling rate
-    #[serde(default = "Option::default")] pub Option<usize>, // fallback max freq
+    pub String,                                                 // file name
+    #[serde(default = "one")] pub f32,                          // dataset sampling factor
+    #[serde(default = "Option::default")] pub Option<usize>,    // fallback sampling rate
+    #[serde(default = "Option::default")] pub Option<usize>,    // fallback max freq
+    #[serde(default = "Option::default")] pub Option<Hdf5Keys>, // cached key list
+    #[serde(default = "Option::default")] pub Option<u64>,      // modified hash
 );
 impl Hdf5Cfg {
     pub fn filename(&self) -> &str {
@@ -85,12 +136,59 @@ impl Hdf5Cfg {
     pub fn sampling_factor(&self) -> f32 {
         self.1
     }
+    pub fn set_sampling_factor(&mut self, f: f32) {
+        self.1 = f
+    }
     pub fn fallback_sr(&self) -> Option<usize> {
         self.2
     }
     pub fn fallback_max_freq(&self) -> Option<usize> {
         self.3
     }
+    pub fn keys_unchecked(&self) -> Option<&Hdf5Keys> {
+        self.4.as_ref()
+    }
+    pub fn hash(&self) -> Option<u64> {
+        self.5
+    }
+    pub fn store_modified_hash(&mut self, hash: u64) {
+        self.5 = Some(hash);
+    }
+    pub fn hash_from_ds_path(&self, ds_path: &str) -> Result<u64> {
+        Ok(calculate_hash(&DatasetModified::new(ds_path)?))
+    }
+    pub fn load_keys(&self, hash: u64) -> Result<Option<&Hdf5Keys>> {
+        if let Some(keys) = self.keys_unchecked() {
+            if keys.hash == hash {
+                return Ok(Some(keys));
+            }
+            println!("Hash does not match for {}", self.filename());
+            return Ok(None);
+        }
+        Ok(None)
+    }
+    pub fn set_keys_new(&mut self, hash: u64, keys: Vec<String>) -> Result<()> {
+        self.set_keys(Hdf5Keys {
+            filename: self.filename().to_string(),
+            hash,
+            keys,
+        })
+    }
+    pub fn set_keys(&mut self, keys: Hdf5Keys) -> Result<()> {
+        self.4.replace(keys);
+        Ok(())
+    }
+}
+fn calculate_hash<T: Hash>(t: &T) -> u64 {
+    let mut s = DefaultHasher::new();
+    t.hash(&mut s);
+    s.finish()
+}
+#[derive(Deserialize, Serialize, Debug, Clone)]
+pub struct Hdf5Keys {
+    pub filename: String,
+    hash: u64,
+    keys: Vec<String>,
 }
 #[derive(Deserialize, Debug)]
 pub struct DatasetConfigJson {
@@ -99,12 +197,27 @@ pub struct DatasetConfigJson {
     pub test: Vec<Hdf5Cfg>,
 }
 impl DatasetConfigJson {
-    pub fn open(path: &str) -> Result<Self> {
-        let file = fs::File::open(path)?;
+    pub fn open(cfg_path: &str) -> Result<Self> {
+        let file = fs::File::open(cfg_path)?;
         let reader = BufReader::new(file);
         let cfg = serde_json::from_reader(reader)?;
         Ok(cfg)
     }
+    pub fn set_keys<S: Into<Split>>(&mut self, split: S, keys: &[Hdf5Keys]) -> Result<()> {
+        let s = match split.into() {
+            Split::Train => &mut self.train,
+            Split::Valid => &mut self.valid,
+            Split::Test => &mut self.test,
+        };
+        for cfg in s.iter_mut() {
+            if let Some(key_cache) = keys.iter().find(|c| c.filename == cfg.filename()) {
+                cfg.set_keys(key_cache.clone())?;
+            } else {
+                eprintln!("Could not find cached keys for {}", cfg.filename());
+            }
+        }
+        Ok(())
+    }
     pub fn split_config(&self, split: Split) -> DatasetSplitConfig {
         match split {
             Split::Train => DatasetSplitConfig {
@@ -122,6 +235,28 @@ impl DatasetConfigJson {
         }
     }
 }
+#[derive(Deserialize, Serialize, Debug, Clone)]
+pub struct DatasetConfigCacheJson(Vec<Hdf5Keys>);
+impl DatasetConfigCacheJson {
+    pub fn new(keys: Vec<Hdf5Keys>) -> Self {
+        DatasetConfigCacheJson(keys)
+    }
+    pub fn keys(&self) -> &Vec<Hdf5Keys> {
+        &self.0
+    }
+    pub fn open(cache_path: &str) -> Result<Self> {
+        let file = fs::File::open(cache_path)?;
+        let reader = BufReader::new(file);
+        let cfg = serde_json::from_reader(reader)?;
+        Ok(cfg)
+    }
+    pub fn write(&self, cache_path: &str) -> Result<()> {
+        let file = fs::OpenOptions::new().create(true).write(true).open(cache_path)?;
+        let writer = BufWriter::new(file);
+        serde_json::to_writer(writer, &self)?;
+        Ok(())
+    }
+}
 
 #[derive(Debug, Clone)]
 pub struct DatasetSplitConfig {
@@ -200,6 +335,7 @@ pub enum SampleType {
     TimeDomain,
     FreqDomain,
 }
+#[derive(Clone)]
 pub struct Sample<T>
 where
     T: Data,
@@ -212,7 +348,6 @@ where
     pub max_freq: usize,
     pub snr: i8,
     pub gain: i8,
-    pub attenuation: Option<u8>,
     pub idx: usize,
 }
 impl Sample<f32> {
@@ -282,8 +417,8 @@ where
 }
 
 #[derive(Clone)]
-pub struct DatasetBuilder<'a> {
-    ds_dir: &'a str,
+pub struct DatasetBuilder {
+    ds_dir: String,
     sr: usize,
     fft_size: Option<usize>,
     datasets: Option<DatasetSplitConfig>,
@@ -292,16 +427,21 @@ pub struct DatasetBuilder<'a> {
     nb_erb: Option<usize>,
     nb_spec: Option<usize>,
     norm_alpha: Option<f32>,
-    p_atten_lim: Option<f32>,
     p_reverb: Option<f32>,
     p_fill_speech: Option<f32>,
     seed: Option<u64>,
     min_nb_freqs: Option<usize>,
+    global_sampling_f: Option<f32>,
+    snrs: Option<Vec<i8>>,
+    gains: Option<Vec<i8>>,
+    cache_valid: bool,
+    num_threads: Option<usize>,
+    logger: Option<Sender<(LogLevel, String)>>,
 }
-impl<'a> DatasetBuilder<'a> {
-    pub fn new(ds_dir: &'a str, sr: usize) -> Self {
+impl DatasetBuilder {
+    pub fn new(ds_dir: &str, sr: usize) -> Self {
         DatasetBuilder {
-            ds_dir,
+            ds_dir: ds_dir.to_string(),
             sr,
             datasets: None,
             max_len_s: None,
@@ -310,11 +450,16 @@ impl<'a> DatasetBuilder<'a> {
             nb_erb: None,
             nb_spec: None,
             norm_alpha: None,
-            p_atten_lim: None,
             p_reverb: None,
             p_fill_speech: None,
             seed: None,
             min_nb_freqs: None,
+            global_sampling_f: None,
+            snrs: None,
+            gains: None,
+            cache_valid: false,
+            num_threads: None,
+            logger: None,
         }
     }
     pub fn build_fft_dataset(self) -> Result<FftDataset> {
@@ -335,6 +480,21 @@ impl<'a> DatasetBuilder<'a> {
                 return Err(DfDatasetError::DataProcessingError(msg));
             }
         }
+        let split = self.datasets.unwrap().split;
+        let cache = if self.cache_valid && split == Split::Valid {
+            let ds_path = Path::new(&self.ds_dir);
+            let hash = {
+                let mut hash_vec: Vec<u64> = ds.config.iter().map(|c| c.hash().unwrap()).collect();
+                hash_vec.push(fft_size as u64);
+                hash_vec.push(hop_size as u64);
+                hash_vec.push(nb_erb as u64);
+                calculate_hash(&hash_vec)
+            };
+            let cache_path = ds_path.join(format!("{}_cache_{}.hdf5", split, hash));
+            Some(Hdf5Cache::new(&cache_path, hash, ds.len())?)
+        } else {
+            None
+        };
         Ok(FftDataset {
             ds,
             fft_size,
@@ -343,42 +503,89 @@ impl<'a> DatasetBuilder<'a> {
             nb_spec: self.nb_spec,
             norm_alpha: self.norm_alpha,
             min_nb_freqs: self.min_nb_freqs,
+            cache,
         })
     }
     pub fn build_td_dataset(self) -> Result<TdDataset> {
-        let mut datasets = match self.datasets {
+        let datasets = match self.datasets {
             None => panic!("No datasets provided"),
             Some(ds) => ds,
         };
         // TODO: Return all sample by default and not only 10 seconds
         let max_samples: usize = (self.max_len_s.unwrap_or(10.) * self.sr as f32).round() as usize;
         // Get dataset handles and keys. Each key is a unique String.
+        let ds_path = Path::new(&self.ds_dir);
+        let (sender, receiver) = sync_channel(datasets.hdf5s.len() + 1);
+        if let Some(n) = self.num_threads {
+            hdf5::sync::sync(|| {});
+            rayon::ThreadPoolBuilder::new()
+                .num_threads(n)
+                .thread_name(|idx| format!("DataLoader Worker {}", idx))
+                .start_handler(|_| hdf5::sync::sync(|| {}))
+                .build_global()
+                .unwrap_or(());
+        }
+        datasets.hdf5s.par_iter().try_for_each(|cfg| -> Result<()> {
+            let path = ds_path.join(cfg.filename());
+            if (!path.is_file()) && path.read_link().is_err() {
+                if let Some(logger) = self.logger.as_ref() {
+                    logger.send((
+                        LogLevel::Warning,
+                        format!("Dataset {:?} not found. Skipping.", path),
+                    ))?;
+                }
+                return Ok(());
+            }
+            let mut cfg = cfg.clone();
+            let ds = Hdf5Dataset::new(path.to_str().unwrap())?;
+            let modified_hash = cfg.hash_from_ds_path(path.to_str().unwrap())?;
+            cfg.store_modified_hash(modified_hash);
+            let keys = cfg.load_keys(modified_hash)?.cloned();
+            match keys {
+                Some(keys) => cfg.set_keys(keys)?,
+                None => cfg.set_keys_new(modified_hash, ds.keys()?)?,
+            };
+            if let Some(f) = self.global_sampling_f {
+                cfg.set_sampling_factor(cfg.sampling_factor() * f)
+            }
+            sender.send(Some((cfg, ds))).unwrap();
+            Ok(())
+        })?;
+        sender.send(None).unwrap();
+        let mut config = Vec::new();
         let mut hdf5_handles = Vec::new();
         let mut ds_keys = Vec::new();
-        let mut config: Vec<Hdf5Cfg> = Vec::new();
         let mut has_rirs = false;
-        for (i, cfg) in datasets.hdf5s.drain(..).enumerate() {
-            let name = cfg.filename();
-            let path = Path::new(self.ds_dir).join(name);
-            if (!path.is_file()) && path.read_link().is_err() {
-                eprintln!("Dataset {:?} not found. Skipping.", path);
-                continue;
+        let mut i = 0;
+        let mut ds_len: usize = 0;
+        while let Some((cfg, ds)) = receiver.try_recv().unwrap() {
+            has_rirs = has_rirs || ds.dstype == DsType::RIR;
+            let keys = cfg.keys_unchecked().unwrap().keys.clone();
+            if ds.dstype == DsType::Speech {
+                ds_len += (keys.len() as f32 * cfg.sampling_factor()).round() as usize;
             }
-            let ds = Hdf5Dataset::new(path.to_str().unwrap())?;
-            if ds.dstype == DsType::RIR {
-                has_rirs = true
+            ds_keys.push((ds.dstype, i, keys));
+            config.push(cfg);
+            if let Some(logger) = self.logger.as_ref() {
+                logger.send((
+                    LogLevel::Debug,
+                    format!(
+                        "Found {} {} dataset {} with {} samples",
+                        &ds.dstype,
+                        datasets.split,
+                        ds.name(),
+                        ds.len()
+                    ),
+                ))?;
             }
-            ds_keys.push((ds.dstype, i, ds.keys()?));
             hdf5_handles.push(ds);
-            config.push(cfg);
+            i += 1;
         }
         if hdf5_handles.is_empty() {
             return Err(DfDatasetError::NoDatasetFoundError);
         }
-        let snrs = vec![-5, 0, 5, 10, 20, 40];
-        let gains = vec![-6, 0, 6];
-        let attenuation_range = (6, 40);
-        let p_atten_lim = self.p_atten_lim.unwrap_or(0.);
+        let snrs = self.snrs.unwrap_or_else(|| vec![-5, 0, 5, 10, 20, 40]);
+        let gains = self.gains.unwrap_or_else(|| vec![-6, 0, 6]);
         let p_fill_speech = self.p_fill_speech.unwrap_or(0.);
         let sp_transforms = Compose::new(vec![
             Box::new(RandRemoveDc::default_with_prob(0.25)),
@@ -389,11 +596,16 @@ impl<'a> DatasetBuilder<'a> {
         let ns_transforms = sp_transforms.clone();
         let p_reverb = self.p_reverb.unwrap_or(0.);
         if p_reverb > 0. && !has_rirs {
-            eprintln!("Warning: Reverb augmentation enabled but no RIRs provided!");
+            if let Some(logger) = self.logger.as_ref() {
+                logger.send((
+                    LogLevel::Warning,
+                    "Warning: Reverb augmentation enabled but no RIRs provided!".to_string(),
+                ))?;
+            }
         }
         let reverb = RandReverbSim::new(p_reverb, self.sr);
         let seed = self.seed.unwrap_or(0);
-        let mut ds = TdDataset {
+        Ok(TdDataset {
             config,
             hdf5_handles,
             max_samples,
@@ -405,17 +617,14 @@ impl<'a> DatasetBuilder<'a> {
             rir_keys: Vec::new(),
             snrs,
             gains,
-            attenuation_range,
             p_fill_speech,
-            p_atten_lim,
             sp_transforms,
             ns_transforms,
             reverb,
             seed,
-        };
-        // Generate inital speech/noise/rir dataset keys. May be changed at the start of each epoch.
-        ds.generate_keys()?;
-        Ok(ds)
+            ds_len,
+            logger: self.logger,
+        })
     }
     pub fn dataset(mut self, datasets: DatasetSplitConfig) -> Self {
         let has_ds = self.datasets.is_some();
@@ -445,9 +654,8 @@ impl<'a> DatasetBuilder<'a> {
         self.norm_alpha = norm_alpha;
         self
     }
-    pub fn prob_atten_lim(mut self, p_atten_lim: f32) -> Self {
-        assert!((0. ..=1.).contains(&p_atten_lim));
-        self.p_atten_lim = Some(p_atten_lim);
+    pub fn global_sample_factor(mut self, f: f32) -> Self {
+        self.global_sampling_f = Some(f);
         self
     }
     pub fn prob_reverberation(mut self, p_reverb: f32) -> Self {
@@ -467,6 +675,26 @@ impl<'a> DatasetBuilder<'a> {
         self.min_nb_freqs = Some(n);
         self
     }
+    pub fn snrs(mut self, snrs: Vec<i8>) -> Self {
+        self.snrs = Some(snrs);
+        self
+    }
+    pub fn gains(mut self, gains: Vec<i8>) -> Self {
+        self.gains = Some(gains);
+        self
+    }
+    pub fn num_threads(mut self, n: usize) -> Self {
+        self.num_threads = Some(n);
+        self
+    }
+    pub fn cache_valid_dataset(mut self) -> Self {
+        self.cache_valid = true;
+        self
+    }
+    pub fn add_logger(mut self, logger: Sender<(LogLevel, String)>) -> Self {
+        self.logger = Some(logger);
+        self
+    }
 }
 
 pub struct FftDataset {
@@ -477,9 +705,35 @@ pub struct FftDataset {
     nb_spec: Option<usize>,
     norm_alpha: Option<f32>,
     min_nb_freqs: Option<usize>,
+    cache: Option<Hdf5Cache>,
+}
+impl FftDataset {
+    pub fn get_hdf5cfg(&self, filename: &str) -> Option<&Hdf5Cfg> {
+        for cfg in &self.ds.config {
+            if cfg.filename() == filename {
+                return Some(cfg);
+            }
+        }
+        None
+    }
+    pub fn close_cache(&mut self) -> Result<()> {
+        if let Some(mut c) = self.cache.take() {
+            c.close()?;
+        }
+        Ok(())
+    }
 }
 impl Dataset<Complex32> for FftDataset {
     fn get_sample(&self, idx: usize, seed: Option<u64>) -> Result<Sample<Complex32>> {
+        let hash = if let Some(cache) = self.cache.as_ref() {
+            let hash = calculate_hash(&(idx, seed));
+            if let Some(s) = cache.load_sample(hash)? {
+                return Ok(s);
+            }
+            hash
+        } else {
+            0
+        };
         let sample: Sample<f32> = self.ds.get_sample(idx, seed)?;
         let nb_erb = self.nb_erb.unwrap_or(1);
         let mut state = DFState::new(self.sr(), self.fft_size, self.hop_size, nb_erb, 1);
@@ -504,7 +758,7 @@ impl Dataset<Complex32> for FftDataset {
         } else {
             None
         };
-        Ok(Sample {
+        let sample = Sample {
             speech: speech.into_dyn(),
             noise: noise.into_dyn(),
             noisy: noisy.into_dyn(),
@@ -513,13 +767,16 @@ impl Dataset<Complex32> for FftDataset {
             max_freq: sample.max_freq,
             gain: sample.gain,
             snr: sample.snr,
-            attenuation: sample.attenuation,
             idx: sample.idx,
-        })
+        };
+        if let Some(cache) = self.cache.as_ref() {
+            cache.cache_sample(hash, &sample, true)?;
+        }
+        Ok(sample)
     }
 
     fn len(&self) -> usize {
-        self.ds.sp_keys.len()
+        self.ds.len()
     }
 
     fn sr(&self) -> usize {
@@ -553,15 +810,15 @@ pub struct TdDataset {
     sp_keys: Vec<(usize, String)>, // Pair of hdf5 index and dataset keys. Will be generated at each epoch start
     ns_keys: Vec<(usize, String)>,
     rir_keys: Vec<(usize, String)>,
-    snrs: Vec<i8>,               // in dB; SNR to sample from
-    gains: Vec<i8>,              // in dB; Speech (loudness) to sample from
-    attenuation_range: (u8, u8), // in dB; Return a target sample containing noise for attenuation limited algorithm
-    p_atten_lim: f32,            // Probability for containing noise in target
+    snrs: Vec<i8>,          // in dB; SNR to sample from
+    gains: Vec<i8>,         // in dB; Speech (loudness) to sample from
     p_fill_speech: f32, // Probability to completely fill the speech signal to `max_samples` with a different speech sample
     sp_transforms: Compose, // Transforms to augment speech samples
     ns_transforms: Compose, // Transforms to augment noise samples
     reverb: RandReverbSim, // Separate reverb transform that may be applied to both speech and noise
     seed: u64,
+    ds_len: usize,
+    logger: Option<Sender<(LogLevel, String)>>,
 }
 
 impl TdDataset {
@@ -612,15 +869,31 @@ impl TdDataset {
     fn read_max_len(&self, idx: usize, key: &str) -> Result<Array2<f32>> {
         let x = match self._read_from_hdf5(key, idx, Some(self.max_samples)) {
             Err(e) => {
-                eprintln!("Error during speech reading get_data(): {:?}", e);
-                if e.to_string().contains("inflate") {
+                let msg = format!(
+                    "Error during {} read_max_len() for key '{}' from dataset {}: {:?}",
+                    self.ds_type(idx),
+                    key,
+                    self.ds_name(idx),
+                    e
+                );
+                eprintln!("{}", &msg);
+                if let Some(logger) = self.logger.as_ref() {
+                    logger.send((LogLevel::Warning, msg)).unwrap_or(());
+                }
+                let e_str = e.to_string();
+                if e_str.contains("inflate") || e_str.contains("Flac") {
                     // Get a different speech then
                     let idx = thread_rng()?.gen_range(0..self.len());
                     let (sp_idx, sp_key) = &self.sp_keys[idx];
-                    eprintln!(
-                        "Returning a different speech sample from {}",
-                        self.ds_name(*sp_idx)
+                    let msg = format!(
+                        "Returning a different speech sample from {} due to {}",
+                        self.ds_name(*sp_idx),
+                        e_str
                     );
+                    eprintln!("{}", &msg);
+                    if let Some(logger) = self.logger.as_ref() {
+                        logger.send((LogLevel::Warning, msg)).unwrap_or(());
+                    }
                     self.read_max_len(*sp_idx, sp_key)?
                 } else {
                     return Err(e);
@@ -646,11 +919,15 @@ impl TdDataset {
     fn ds_name(&self, idx: usize) -> String {
         self.hdf5_handles[idx].name()
     }
+
+    fn ds_type(&self, idx: usize) -> String {
+        self.hdf5_handles[idx].ds_type()
+    }
 }
 
 impl Dataset<f32> for TdDataset {
     fn get_sample(&self, idx: usize, seed: Option<u64>) -> Result<Sample<f32>> {
-        seed_from_u64(idx as u64 + self.seed + seed.unwrap_or(0));
+        seed_from_u64(self.seed + seed.unwrap_or(idx as u64));
         let mut rng = thread_rng()?;
         let (sp_idx, sp_key) = &self.sp_keys[idx];
         let mut speech = self.read_max_len(*sp_idx, sp_key)?;
@@ -697,7 +974,11 @@ impl Dataset<f32> for TdDataset {
         for (ns_idx, ns_key) in &ns_ids {
             let mut ns = match self.read_max_len(*ns_idx, ns_key) {
                 Err(e) => {
-                    eprintln!("Error during noise reading get_data(): {}", e);
+                    let msg = format!("Error during noise reading get_sample(): {}", e);
+                    eprintln!("{}", &msg);
+                    if let Some(logger) = self.logger.as_ref() {
+                        logger.send((LogLevel::Warning, msg)).unwrap_or(());
+                    }
                     continue;
                 }
                 Ok(n) => n,
@@ -716,12 +997,6 @@ impl Dataset<f32> for TdDataset {
         // Sample SNR and gain
         let &snr = self.snrs.choose(&mut rng).unwrap();
         let &gain = self.gains.choose(&mut rng).unwrap();
-        // Sample attenuation limiting during training
-        let atten = if self.p_atten_lim > 0. && self.p_atten_lim > rng.gen_range(0f32..1f32) {
-            Some(rng.gen_range(self.attenuation_range.0..self.attenuation_range.1))
-        } else {
-            None
-        };
         // Truncate to speech len, combine noises and mix to noisy
         let mut noise = combine_noises(ch, len, &mut noises, Some(noise_gains_f32.as_slice()))?;
         // Apply reverberation using a randomly sampled RIR
@@ -740,7 +1015,6 @@ impl Dataset<f32> for TdDataset {
             noise,
             snr as f32,
             gain as f32,
-            atten.map(|a| a as f32),
             noise_low_pass,
         )?;
         Ok(Sample {
@@ -752,13 +1026,12 @@ impl Dataset<f32> for TdDataset {
             max_freq,
             snr,
             gain,
-            attenuation: atten,
             idx,
         })
     }
 
     fn len(&self) -> usize {
-        self.sp_keys.len()
+        self.ds_len
     }
 
     fn sr(&self) -> usize {
@@ -795,12 +1068,13 @@ impl Dataset<f32> for TdDataset {
         self.rir_keys.clear();
 
         for (dstype, hdf5_idx, keys) in self.ds_keys.iter() {
-            let len = self.hdf5_handles[*hdf5_idx].len();
+            debug_assert_eq!(&self.hdf5_handles[*hdf5_idx].keys().unwrap(), keys);
+            let len = keys.len();
             let n_samples =
                 (self.config[*hdf5_idx].sampling_factor() * len as f32).round() as usize;
             let mut keys = keys.clone();
             if self.ds_split == Split::Train {
-                keys.shuffle(&mut thread_rng()?)
+                keys.shuffle(&mut thread_rng()?);
             }
             let keys: Vec<(usize, String)> =
                 keys.iter().cycle().take(n_samples).map(|k| (*hdf5_idx, k.clone())).collect();
@@ -847,16 +1121,250 @@ pub enum DType {
     F32 = 1,
 }
 
+/// Cache for validation dataset
+#[derive(Debug)]
+struct Hdf5Cache {
+    file: Option<Arc<File>>,
+    hash: u64,
+    filters: Arc<Vec<filters::Filter>>,
+    cacher: Option<thread::JoinHandle<Result<()>>>,
+    sender: SyncSender<Option<(u64, Sample<Complex32>)>>,
+    n_samples: usize,
+    max_size: u64, // in bytes
+}
+
+impl Drop for Hdf5Cache {
+    fn drop(&mut self) {
+        self.close().unwrap();
+    }
+}
+impl Hdf5Cache {
+    fn new(path: &impl AsRef<Path>, hash: u64, n_total_samples: usize) -> Result<Self> {
+        let path = path.as_ref();
+        let file = if path.is_file() {
+            println!("Opening validation cache: {:?}", path);
+            File::open_rw(path)?
+        } else {
+            println!("Validation cache not found. Creating: {:?}", path);
+            File::create(path)?
+        };
+        match file.attr("hash") {
+            Err(_e) => {
+                let attr = file.new_attr::<u64>().create("hash")?;
+                attr.write_scalar::<u64>(&hash)?;
+            }
+            Ok(attr) => assert_eq!(attr.read_scalar::<u64>().unwrap(), hash),
+        };
+        let file = Arc::new(file);
+        let filters = if filters::blosc_available() {
+            filters::blosc_set_nthreads(8);
+            Arc::new(vec![filters::Filter::blosc_lz4hc(5, true)])
+        } else if filters::deflate_available() {
+            Arc::new(vec![filters::Filter::Deflate(4)])
+        } else {
+            Arc::new(Vec::new())
+        };
+        let max_gb = std::env::var("HDF5CACHE_MAX_GB")
+            .map(|s| s.parse::<f32>().unwrap_or_default())
+            .unwrap_or_default();
+        let max_size: u64 = (max_gb * 1024. * 1024. * 1024.) as u64;
+        Hdf5Cache::set_attr(file.deref(), "max_size", &max_size)?;
+        Hdf5Cache::set_attr(&file.as_location()?, "total_samples", &n_total_samples)?;
+        let (sender, receiver) = sync_channel(256);
+        let filters_c = filters.clone();
+        let file_c = file.clone();
+        let cacher = Some(thread::spawn(move || -> Result<()> {
+            seed_from_u64(0); // Only for selecting which samples to store;
+            while let Some((key, sample)) = receiver.recv()? {
+                Hdf5Cache::cache_sample_impl(&file_c, key, &sample, &filters_c)?;
+            }
+            Ok(())
+        }));
+        Ok(Hdf5Cache {
+            file: Some(file),
+            hash,
+            filters,
+            cacher,
+            sender,
+            n_samples: n_total_samples,
+            max_size,
+        })
+    }
+    pub fn close(&mut self) -> Result<()> {
+        self.sender.try_send(None).unwrap_or_default();
+        if let Some(thread) = self.cacher.take() {
+            thread
+                .join()
+                .map_err(|e| DfDatasetError::ThreadJoinError(format!("{:?}", e)))??;
+        }
+        if let Some(file) = self.file.take() {
+            Arc::try_unwrap(file).unwrap().close()?
+        }
+        Ok(())
+    }
+    pub fn cache_sample(&self, key: u64, sample: &Sample<Complex32>, queue: bool) -> Result<()> {
+        if Self::is_full(self.file.as_ref().unwrap(), self.max_size) {
+            return Ok(());
+        }
+        if queue {
+            return Ok(self.sender.send(Some((key, sample.clone())))?);
+        } else {
+            Self::cache_sample_impl(self.file.as_ref().unwrap(), key, sample, &self.filters)?;
+        };
+        Ok(())
+    }
+    fn get_attr<T: H5Type>(l: &Location, name: &str) -> Result<Option<T>> {
+        match l.attr(name) {
+            Ok(a) => Ok(Some(a.read_scalar()?)),
+            Err(_) => Ok(None),
+        }
+    }
+    fn set_attr<T: H5Type>(l: &Location, name: &str, val: &T) -> Result<()> {
+        let attr = l.new_attr::<T>().create(name).unwrap_or_else(|_| l.attr(name).unwrap());
+        attr.write_scalar(val)?;
+        Ok(())
+    }
+    fn create_ds<'d, A, T, D>(
+        grp: &Group,
+        name: &str,
+        data: A,
+        filters: &[filters::Filter],
+    ) -> Result<()>
+    where
+        A: Into<ArrayView<'d, T, D>>,
+        T: H5Type,
+        D: ndarray::Dimension,
+    {
+        grp.new_dataset_builder().set_filters(filters).with_data(data).create(name)?;
+        Ok(())
+    }
+    fn store_cplx(
+        grp: &Group,
+        name: &str,
+        data: ArrayViewD<Complex32>,
+        filters: &[filters::Filter],
+    ) -> Result<()> {
+        let complex = data.split_complex();
+        let name_re = name.to_owned() + "_re";
+        let name_im = name.to_owned() + "_im";
+        Self::create_ds(
+            grp,
+            &name_re,
+            complex.re.as_standard_layout().view(),
+            filters,
+        )?;
+        Self::create_ds(
+            grp,
+            &name_im,
+            complex.im.as_standard_layout().view(),
+            filters,
+        )?;
+        Ok(())
+    }
+    fn load_cplx(grp: &Group, name: &str) -> Result<ArrayD<Complex32>> {
+        let name_re = name.to_owned() + "_re";
+        let name_im = name.to_owned() + "_im";
+        let ds_re = grp.dataset(&*name_re)?;
+        let ds_im = grp.dataset(&*name_im)?;
+        let mut out = ArrayD::zeros(ds_re.shape());
+        let view = out.view_mut().split_complex();
+        ds_re.read_dyn()?.move_into(view.re);
+        ds_im.read_dyn()?.move_into(view.im);
+        Ok(out)
+    }
+    fn is_full(file: &File, max_size: u64) -> bool {
+        if max_size == 0 {
+            return false;
+        }
+        file.size() > max_size
+    }
+    fn should_cache(
+        file: &File,
+        max_size: Option<u64>,
+        total_samples: Option<u64>,
+    ) -> Result<bool> {
+        let max_size = max_size
+            .unwrap_or_else(|| Self::get_attr(file, "max_size").unwrap().unwrap_or_default());
+        if Self::is_full(file, max_size) {
+            return Ok(false);
+        }
+        let n_samples = file.len();
+        if n_samples == 0 {
+            return Ok(true);
+        }
+        let total_samples = total_samples
+            .unwrap_or_else(|| Self::get_attr(file, "total_samples").unwrap().unwrap_or_default());
+        let avg_size = file.size() as f32 / n_samples as f32;
+        let total_estimate = total_samples as f32 * avg_size;
+        if total_estimate < max_size as f32 {
+            Ok(true)
+        } else {
+            // Since we can't store all samples, let's just store a fraction by uniformly
+            // sampling.
+            let p = max_size as f32 / total_estimate;
+            Ok(p >= thread_rng()?.gen_range(0f32..1f32))
+        }
+    }
+    fn cache_sample_impl(
+        file: &File,
+        key: u64,
+        sample: &Sample<Complex32>,
+        filters: &[filters::Filter],
+    ) -> Result<()> {
+        if !Self::should_cache(file, None, None)? {
+            return Ok(());
+        }
+        let grp = file.create_group(&key.to_string())?;
+        Self::store_cplx(&grp, "speech", sample.speech.view(), filters)?;
+        Self::store_cplx(&grp, "noise", sample.noise.view(), filters)?;
+        Self::store_cplx(&grp, "noisy", sample.noisy.view(), filters)?;
+        if let Some(feat_erb) = sample.feat_erb.as_ref() {
+            Self::create_ds(&grp, "feat_erb", feat_erb.view(), filters)?;
+        }
+        if let Some(feat_spec) = sample.feat_spec.as_ref() {
+            Self::store_cplx(&grp, "feat_spec", feat_spec.view(), filters)?;
+        }
+        Self::set_attr(&grp, "max_freq", &sample.max_freq)?;
+        Self::set_attr(&grp, "snr", &sample.snr)?;
+        Self::set_attr(&grp, "gain", &sample.gain)?;
+        Self::set_attr(&grp, "idx", &sample.idx)?;
+        Ok(())
+    }
+    fn load_sample(&self, key: u64) -> Result<Option<Sample<Complex32>>> {
+        let grp = match self.file.as_ref().unwrap().group(&key.to_string()) {
+            Ok(g) => g,
+            Err(_) => return Ok(None),
+        };
+        Ok(Some(Sample {
+            speech: Self::load_cplx(&grp, "speech")?,
+            noise: Self::load_cplx(&grp, "noise")?,
+            noisy: Self::load_cplx(&grp, "noisy")?,
+            feat_erb: Some(grp.dataset("feat_erb")?.read_dyn()?),
+            feat_spec: Some(Self::load_cplx(&grp, "feat_spec")?),
+            max_freq: Self::get_attr(&grp, "max_freq")?.unwrap(),
+            snr: Self::get_attr(&grp, "snr")?.unwrap(),
+            gain: Self::get_attr(&grp, "gain")?.unwrap(),
+            idx: Self::get_attr(&grp, "idx")?.unwrap(),
+        }))
+    }
+}
+
 #[derive(Debug)]
 pub struct Hdf5Dataset {
     file: File,
-    dstype: DsType,
-    sr: Option<usize>,
-    codec: Option<Codec>,
+    pub dstype: DsType,
+    pub sr: Option<usize>,
+    pub codec: Option<Codec>,
     max_freq: Option<usize>,
     dtype: Option<DType>,
 }
 
+#[derive(Hash, Debug, Clone)]
+pub struct Hdf5CacheHash {
+    seed: u64,
+    ds_hashes: Vec<(String, u64)>, // dataset filename, u64 hash via calculate_hash(&DatasetModified)
+}
+
 fn get_dstype(file: &File) -> Option<DsType> {
     for g in file.member_names().unwrap_or_default() {
         match g.to_lowercase().as_str() {
@@ -870,7 +1378,7 @@ fn get_dstype(file: &File) -> Option<DsType> {
 }
 
 impl Hdf5Dataset {
-    fn new(path: &str) -> Result<Self> {
+    pub fn new(path: &str) -> Result<Self> {
         let file = File::open(path).map_err(move |e: hdf5::Error| -> DfDatasetError {
             DfDatasetError::Hdf5ErrorDetail {
                 source: e,
@@ -919,8 +1427,11 @@ impl Hdf5Dataset {
     fn name(&self) -> String {
         self.file.filename()
     }
+    pub fn ds_type(&self) -> String {
+        self.dstype.to_string().to_lowercase()
+    }
     fn group(&self) -> Result<hdf5::Group> {
-        Ok(self.file.group(&self.dstype.to_string().to_lowercase())?)
+        Ok(self.file.group(&self.ds_type())?)
     }
     pub fn len(&self) -> usize {
         self.group().unwrap().len() as usize
@@ -966,11 +1477,23 @@ impl Hdf5Dataset {
     /// * `ds`: `hdf5::Dataset` containing a vorbis (ogg) encoded audio sample.
     fn sample_len_vorbis(&self, ds: hdf5::Dataset) -> Result<usize> {
         let mut rdr = OggPacketReader::new(ds.as_byte_reader()?);
-        // Seek to almost end to get the last ogg package
+        // Seek almost to end to get the last ogg package
         rdr.seek_bytes(std::io::SeekFrom::End(-4096))?;
+        let mut pkg = rdr.read_packet();
+        if pkg.is_err() {
+            // Maybe seek a little further or start entirely from the beginning.
+            rdr.seek_bytes(std::io::SeekFrom::End(-8192))?;
+            pkg = rdr.read_packet();
+            if pkg.is_err() {
+                rdr.seek_absgp(None, 0)?;
+                pkg = rdr.read_packet();
+            };
+        }
+        // Also check if there are some packges left
         let mut absgp = 0;
-        while let Some(pkg) = rdr.read_packet()? {
-            absgp = pkg.absgp_page();
+        while let Some(p) = pkg? {
+            absgp = p.absgp_page();
+            pkg = rdr.read_packet();
         }
         Ok(absgp as usize)
     }
@@ -1071,22 +1594,12 @@ impl Hdf5Dataset {
             ds: format!("{:?}", self.file),
         })
     }
-    /// Read a Flac encoded sample from an `hdf5::Dataset`.
-    ///
-    /// Arguments:
-    ///
-    /// * `key`: String idendifier to load the dataset.
-    /// * `channel`: Optional channel. `-1` will load a random channel, `None` will return all channels.
-    /// * `r`: Optional range in samples (time axis). `None` will return all samples.
     #[cfg(feature = "flac")]
-    fn read_flac(
+    fn _read_flac<R: std::io::Read>(
         &self,
         key: &str,
-        channel: Option<isize>,
-        r: Option<Range<usize>>,
+        mut reader: claxon::FlacReader<R>,
     ) -> Result<Array2<f32>> {
-        let ds = self.group()?.dataset(key)?;
-        let mut reader = claxon::FlacReader::new(ds.as_byte_reader()?)?;
         let info = reader.streaminfo();
         assert_eq!(
             info.bits_per_sample, 16,
@@ -1098,7 +1611,19 @@ impl Hdf5Dataset {
         let mut out: Array2<f32> = Array2::zeros((ch, samples));
         let mut block = claxon::Block::empty();
         let mut idx = 0;
-        while let Some(next) = frame_reader.read_next_or_eof(block.into_buffer())? {
+        loop {
+            let next = match frame_reader.read_next_or_eof(block.into_buffer()) {
+                Ok(Some(n)) => n,
+                Ok(None) => break,
+                Err(e) => {
+                    eprintln!("Error decoding flac dataset {} {:?}", key, e);
+                    if e.to_string().contains("CRC") {
+                        break;
+                    } else {
+                        return Err(e.into());
+                    }
+                }
+            };
             let numel = (next.len() / next.channels()) as usize;
             debug_assert_eq!(ch, next.channels() as usize);
             for i in 0..ch {
@@ -1112,7 +1637,39 @@ impl Hdf5Dataset {
             idx += numel;
             block = next
         }
-
+        Ok(out)
+    }
+    #[cfg(feature = "flac")]
+    fn read_flac_byte_reader(&self, key: &str) -> Result<Array2<f32>> {
+        let ds = self.group()?.dataset(key)?;
+        let reader = claxon::FlacReader::new(ds.as_byte_reader()?)?;
+        self._read_flac(key, reader)
+    }
+    #[cfg(feature = "flac")]
+    fn read_flac_ds(&self, key: &str) -> Result<Array2<f32>> {
+        let ds = self.group()?.dataset(key)?;
+        let encoded = ds.read_1d()?;
+        let reader = claxon::FlacReader::new(encoded.as_slice().unwrap())?;
+        self._read_flac(key, reader)
+    }
+    /// Read a Flac encoded sample from an `hdf5::Dataset`.
+    ///
+    /// Arguments:
+    ///
+    /// * `key`: String idendifier to load the dataset.
+    /// * `channel`: Optional channel. `-1` will load a random channel, `None` will return all channels.
+    /// * `r`: Optional range in samples (time axis). `None` will return all samples.
+    #[cfg(feature = "flac")]
+    fn read_flac(
+        &self,
+        key: &str,
+        channel: Option<isize>,
+        r: Option<Range<usize>>,
+    ) -> Result<Array2<f32>> {
+        let out = match self.read_flac_byte_reader(key) {
+            Ok(x) => x,
+            Err(_) => self.read_flac_ds(key)?,
+        };
         let mut out = self.match_ch(out, 0, channel)?;
         if let Some(r) = r {
             out.slice_axis_inplace(Axis(1), Slice::from(r));
@@ -1154,22 +1711,43 @@ impl Hdf5Dataset {
         } else {
             24000 // start with 0.5s if sr=48000
         };
+        let mut pck = match srr.read_dec_packet_itl() {
+            Ok(pck) => pck,
+            Err(e) => {
+                if r.is_some() {
+                    // Undo the seek and just start from the beginning
+                    srr = OggStreamReader::new(ds.as_byte_reader()?).unwrap();
+                    srr.read_dec_packet_itl().unwrap()
+                } else {
+                    return Err(e.into());
+                }
+            }
+        };
         let mut out: Vec<i16> = Vec::with_capacity(len);
-        while let Some(mut pck) = srr.read_dec_packet_itl()? {
+        while let Some(mut p) = pck {
+            out.append(&mut p);
             if let (Some(r), Some(pos)) = (r.as_ref(), srr.get_last_absgp().map(|p| p as usize)) {
                 if pos >= r.end {
-                    pck.truncate(pck.len() - (pos - r.end) * ch);
-                    out.append(&mut pck);
+                    // In some rare curcumstances the decoding start is behind r.start.
+                    // Nor sure if this is an error on my end or on lewton's. Anyways we just
+                    // return a slightly delayed slice, therefore the max operation.
+                    out.truncate((out.len() - (pos - r.end) * ch).max((r.end - r.start) * ch));
+                    debug_assert!(out.len() >= r.end - r.start);
                     break;
                 }
             }
-            out.append(&mut pck);
+            pck = srr.read_dec_packet_itl()?;
         }
         let out = if let Some(r) = r {
             // We already have a coarse range. The start may contain more samples from its
-            // corresponging ogg page. The end is already exact. Thus, truncate the beginning.
-            let start_pos = out.len() - (r.end - r.start) * ch;
-            ArrayView2::from_shape((len, ch), &out[start_pos..])?.to_owned()
+            // corresponding ogg page. The end is already exact. Thus, truncate the beginning.
+            // let start_pos = out.len() - (r.end - r.start) * ch;
+            let start_pos = out.len().saturating_sub((r.end - r.start) * ch);
+            // Due to some rare bug, the slice will end up shorter than `len`. I guess similar to
+            // above, the first decoded sample starts after `r.begin`. Threrfore we needed
+            // `saturating_sub` and recalulate the length via `slice.len()`.
+            let slice = &out[start_pos..];
+            ArrayView2::from_shape((slice.len() / ch, ch), slice)?.to_owned()
         } else {
             Array2::from_shape_vec((out.len() / ch, ch), out)?
         };
@@ -1259,8 +1837,6 @@ fn combine_noises(
 /// * `noise` - A noise signal of shape `[C, N]`. Will be modified in place.
 /// * `snr_db` - Signal to noise ratio in decibel used for mixing.
 /// * `gain_db` - Gain to apply to the clean signal in decibel before mixing.
-/// * `atten_db` - Target attenuation limit in decibel. The resulting clean target will contain
-///                `atten_db` less noise compared to the noisy output.
 /// * `noise_resample`: Optional resample parameters which will be used to apply a low-pass via
 ///                     resampling to the noise signal. This may be used to make sure a speech
 ///                     signal with a lower sampling rate will also be mixed with noise having the
@@ -1271,7 +1847,6 @@ fn mix_audio_signal(
     mut noise: Array2<f32>,
     snr_db: f32,
     gain_db: f32,
-    atten_db: Option<f32>,
     noise_resample: Option<LpParam>,
 ) -> Result<(Signal, Signal, Signal)> {
     let len = clean.len_of(Axis(1));
@@ -1287,13 +1862,6 @@ fn mix_audio_signal(
     let clean_mix = clean_rev.map(|c| &c * g).unwrap_or_else(|| clean_out.clone());
     // For energy calculation use clean speech to also consider direct-to-reverberant ratio
     noise *= mix_f(clean_out.view(), noise.view(), snr_db);
-    if let Some(atten_db) = atten_db {
-        // Create a mixture with a higher SNR as target signal
-        let k_target = 1. / mix_f(clean_out.view(), noise.view(), snr_db + atten_db);
-        clean_out *= k_target;
-        // clean_mix *= k_target;
-        clean_out += &noise;
-    }
     let mut mixture = clean_mix + &noise;
     // Guard against clipping
     let max = &([&clean_out, &noise, &mixture].iter().map(|x| find_max_abs(x.iter())))
@@ -1308,6 +1876,47 @@ fn mix_audio_signal(
     Ok((clean_out, noise, mixture))
 }
 
+#[derive(Clone, Debug)]
+pub enum LogLevel {
+    Debug,
+    Info,
+    Warning,
+}
+impl fmt::Display for LogLevel {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Debug => write!(f, "DEBUG"),
+            Self::Info => write!(f, "INFO"),
+            Self::Warning => write!(f, "WARNING"),
+        }
+    }
+}
+impl From<LogLevel> for String {
+    fn from(l: LogLevel) -> String {
+        String::from(&l)
+    }
+}
+impl From<&LogLevel> for String {
+    fn from(l: &LogLevel) -> String {
+        match l {
+            LogLevel::Debug => "DEBUG".to_string(),
+            LogLevel::Info => "INFO".to_string(),
+            LogLevel::Warning => "WARNING".to_string(),
+        }
+    }
+}
+
+impl From<&str> for LogLevel {
+    fn from(loglevel: &str) -> Self {
+        match loglevel.to_lowercase().as_str() {
+            "debug" => LogLevel::Debug,
+            "info" => LogLevel::Info,
+            "warning" => LogLevel::Warning,
+            s => panic!("LogLevel '{}' does not exist.", s),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::BTreeSet;
@@ -1506,7 +2115,7 @@ mod tests {
             wav_utils::write_wav_arr2(filename, samples_raw.view(), hdf5.sr.unwrap() as u32)
                 .unwrap();
             let dsn =
-                &str::replace(ds, "../assets/noise", "").replace("_", "").replace(".hdf5", "");
+                &str::replace(ds, "../assets/noise", "").replace('_', "").replace(".hdf5", "");
             let filename = &str::replace(
                 basen,
                 ".wav",
@@ -1528,55 +2137,90 @@ mod tests {
         let noise = arr1(rng_uniform(n, -0.1, 0.1)?.as_slice()).into_shape([1, n])?;
         let gains = [-6., 0., 6.];
         let snrs = [-10., -5., 0., 5., 10., 20., 40.];
-        let atten_limits = [None, Some(20.), Some(10.), Some(3.)];
         let atol = 1e-4;
         for clean_rev in [None, Some(clean.clone())] {
             for gain in gains {
                 for snr in snrs {
-                    for attn in atten_limits {
-                        let (c, n, m) = mix_audio_signal(
-                            clean.clone(),
-                            clean_rev.clone(),
-                            noise.clone(),
-                            snr,
-                            gain,
-                            attn,
-                            None,
-                        )?;
-                        if attn.is_none() {
-                            assert_eq!(&c + &n, m);
-                        }
-                        dbg!(clean_rev.is_some(), gain, snr, attn);
-                        // Input SNR of mixture
-                        let snr_inp_m = calc_snr_xv(m.iter(), n.iter());
-                        assert!(
-                            (snr_inp_m - snr).abs() < atol,
-                            "Input SNR does not match: {}, {}",
-                            snr_inp_m,
-                            snr
-                        );
-                        // Target SNR between noise and target (clean) speech. Clean speech may
-                        // contain noise due to the attenuation limit, effectively resulting in an
-                        // `attn` higher SNR.
-                        let snr_target_c = if attn.is_none() {
-                            calc_snr(c.iter(), n.iter())
-                        } else {
-                            // With enabled attenuation limiting, the target signal `c` contains
-                            // some noise so that its SNR is `attn` higher then the input mixture.
-                            calc_snr_xv(c.iter(), n.iter())
-                        };
-                        assert!(
-                            (snr_target_c - (snr + attn.unwrap_or(0.))).abs() < atol,
-                            "Target SNR does not match: {}, {}",
-                            snr_target_c,
-                            snr + attn.unwrap_or(0.),
-                        );
-                        // Test the SNR difference between input and target
-                        assert!((snr_inp_m + attn.unwrap_or(0.) - snr_target_c).abs() < atol);
-                    }
+                    let (c, n, m) = mix_audio_signal(
+                        clean.clone(),
+                        clean_rev.clone(),
+                        noise.clone(),
+                        snr,
+                        gain,
+                        None,
+                    )?;
+                    assert_eq!(&c + &n, m);
+                    dbg!(clean_rev.is_some(), gain, snr);
+                    // Input SNR of mixture
+                    let snr_inp_m = calc_snr_xv(m.iter(), n.iter());
+                    assert!(
+                        (snr_inp_m - snr).abs() < atol,
+                        "Input SNR does not match: {}, {}",
+                        snr_inp_m,
+                        snr
+                    );
+                    // Target SNR between noise and target (clean) speech.
+                    let snr_target_c = calc_snr(c.iter(), n.iter());
+                    assert!(
+                        (snr_target_c - snr).abs() < atol,
+                        "Target SNR does not match: {}, {}",
+                        snr_target_c,
+                        snr,
+                    );
+                    // Test the SNR difference between input and target
+                    assert!((snr_inp_m - snr_target_c).abs() < atol);
+                }
+            }
+        }
+        Ok(())
+    }
+    #[test]
+    pub fn test_cached_valid_dataset() -> Result<()> {
+        seed_from_u64(42);
+        let fft_size = 960;
+        let hop_size = Some(480);
+        let nb_erb = Some(32);
+        let nb_spec = Some(32);
+        let norm_alpha = None;
+        let sr = 48000;
+        let ds_dir = "../assets/";
+        for file in fs::read_dir(Path::new("../assets/"))? {
+            let file = file?.path();
+            if file.is_file()
+                && file.file_name().unwrap().to_str().unwrap().starts_with("valid_cache_")
+            {
+                println!("Removing existing cache '{:?}'", file);
+                fs::remove_file(file)?;
+            }
+        }
+        let mut cfg = DatasetConfigJson::open("../assets/dataset.cfg")?;
+        for c in cfg.valid.iter_mut() {
+            c.1 = 3.0; // Set sampling factor
+        }
+        let builder = DatasetBuilder::new(ds_dir, sr)
+            .df_params(fft_size, hop_size, nb_erb, nb_spec, norm_alpha)
+            .max_len(1.);
+        let mut val_ds = builder
+            .cache_valid_dataset()
+            .dataset(cfg.split_config(Split::Valid))
+            .build_fft_dataset()?;
+        let ds_len = val_ds.len();
+        dbg!(ds_len);
+        for seed in [42, 43] {
+            dbg!(seed);
+            for _ in 0..2 {
+                val_ds.set_seed(seed);
+                if val_ds.need_generate_keys() {
+                    dbg!();
+                    val_ds.generate_keys()?
+                }
+                for idx in 0..ds_len {
+                    let sample = val_ds.get_sample(idx, Some(seed))?;
+                    dbg!(idx, sample);
                 }
             }
         }
+
         Ok(())
     }
 }
diff --git a/libDF/src/util.rs b/libDF/src/util.rs
index 12dc4c0d7..0308f23d3 100644
--- a/libDF/src/util.rs
+++ b/libDF/src/util.rs
@@ -18,7 +18,7 @@ type Result<T> = std::result::Result<T, UtilsError>;
 pub enum UtilsError {
     #[error("NaN detected")]
     NaN,
-    #[error("Random seed is not initalized using seed_from_u64(x)")]
+    #[error("Random seed is not initialized using seed_from_u64(x)")]
     SeedNotInitialized,
 }
 
diff --git a/pyDF-data/Cargo.toml b/pyDF-data/Cargo.toml
index 0a89e0410..da77b7761 100644
--- a/pyDF-data/Cargo.toml
+++ b/pyDF-data/Cargo.toml
@@ -12,5 +12,6 @@ path = "src/lib.rs"
 [dependencies]
 deep_filter = { features = ["dataset", "vorbis", "flac"], path = "../libDF" }
 pyo3 = { version = "^0.16", default_features = true, features = ["extension-module"]}
+crossbeam-channel = "^0.5"
 numpy = "^0.16"
 ndarray = "^0.15"
diff --git a/pyDF-data/libdfdata/torch_dataloader.py b/pyDF-data/libdfdata/torch_dataloader.py
index e520a6b26..c955d0163 100644
--- a/pyDF-data/libdfdata/torch_dataloader.py
+++ b/pyDF-data/libdfdata/torch_dataloader.py
@@ -1,8 +1,9 @@
 import atexit
 import queue
 import threading
+import time
 import warnings
-from typing import Iterator, Optional, Tuple
+from typing import Iterator, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -20,28 +21,23 @@ def __init__(self, b: Tuple[np.ndarray, ...]):
         # safe to assume writable.
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", UserWarning)
-            if len(b) == 10:  # FftDataloader
-                speech, noise, noisy, erb, spec, lengths, max_freq, snr, gain, atten = b
-                if erb.size <= 1:
-                    self.feat_erb = None
-                if spec.size <= 1:
-                    self.feat_spec = None
-                else:
-                    self.feat_erb = torch.from_numpy(erb)
-                    self.feat_spec = torch.from_numpy(spec)
-            else:
-                speech, noise, noisy, lengths, max_freq, snr, gain, atten = b
+            assert len(b) == 10
+            speech, noise, noisy, erb, spec, lengths, max_freq, snr, gain, timings = b
+            if erb.size <= 1:
                 self.feat_erb = None
+            if spec.size <= 1:
                 self.feat_spec = None
+            else:
+                self.feat_erb = torch.from_numpy(erb)
+                self.feat_spec = torch.from_numpy(spec)
             self.speech = torch.from_numpy(speech)
             self.noise = torch.from_numpy(noise)
             self.noisy = torch.from_numpy(noisy)
             self.lengths = torch.from_numpy(lengths.astype(np.int64)).long()
             self.snr = torch.from_numpy(snr)
             self.gain = torch.from_numpy(gain)
-            self.atten = torch.from_numpy(atten).long()
-            self.atten[self.atten == 0] = 1000
             self.max_freq = torch.from_numpy(max_freq.astype(np.int32)).long()
+            self.timings = torch.from_numpy(timings.astype(np.float32))
 
     def pin_memory(self):
         self.speech = self.speech.pin_memory()
@@ -52,16 +48,15 @@ def pin_memory(self):
             self.feat_spec = self.feat_spec.pin_memory()
         if self.max_freq is not None:
             self.max_freq = self.max_freq.pin_memory()
-        self.atten.pin_memory()
         return self
 
     def __repr__(self):
         bs = len(self.lengths)
         s = f"Batch of size {bs}:\n"
-        for i in range(bs):
-            s += f"    length: {self.lengths[i]}\n"
-            s += f"    snr: {self.snr[i]}\n"
-            s += f"    gain: {self.gain[i]}\n"
+        snrs = "".join(f"{s}," for s in self.snr)[:-1]
+        gains = "".join(f"{g}," for g in self.gain)[:-1]
+        s += f"    SNRs: {snrs}\n"
+        s += f"    Gains: {gains}\n"
         return s
 
 
@@ -73,7 +68,7 @@ def __init__(
         sr: int,
         batch_size: int,
         max_len_s: Optional[float] = 10.0,
-        prefetch=4,
+        prefetch=8,
         num_workers=None,
         pin_memory=True,
         drop_last=False,  # Drop the last batch if it contains fewer samples then batch_size
@@ -83,16 +78,20 @@ def __init__(
         nb_spec: int = None,  # Number of complex spectrogram bins
         norm_alpha: float = None,  # Exponential normalization decay for erb_feat and spec_feat
         batch_size_eval: Optional[int] = None,  # Different batch size for evaluation
-        p_atten_lim: Optional[float] = None,  # Limit attenuation by providing a noisy target
         p_reverb: Optional[float] = None,  # Percentage of reverberant speech/noise samples
         overfit=False,  # Overfit on one epoch
+        cache_valid=False,  # Cache validiation dataset via a hdf5 dataset
         seed=0,
         min_nb_erb_freqs: int = None,  # Minimum number of frequency bins per ERB band
+        log_timings: bool = False,
+        global_sampling_factor: float = None,  # Additional over/undersampling of all datasets
+        snrs=None,  # Signal to noise ratios (SNRs) to generate. Defaults to [-5,0,5,10,20,40] dB
+        gains=None,  # Additional gains applied to speech. Defaults to [-6,0,6] dB
     ):
         self.fft_size = fft_size
         self.batch_size = batch_size
         self.batch_size_eval = batch_size_eval
-        prefetch_loader = prefetch * batch_size * num_workers
+        prefetch_loader = batch_size * (num_workers or 1)
         logger.info(f"Initializing dataloader with data directory {ds_dir}")
         assert self.fft_size is not None, "No fft_size provided"
         self.loader = _FdDataLoader(
@@ -108,14 +107,19 @@ def __init__(
             nb_erb=nb_erb,
             nb_spec=nb_spec,
             norm_alpha=norm_alpha,
-            p_atten_lim=p_atten_lim,
             p_reverb=p_reverb,
             prefetch=prefetch_loader,
             drop_last=drop_last,
             overfit=overfit,
+            cache_valid=cache_valid,
             seed=seed,
             min_nb_erb_freqs=min_nb_erb_freqs,
+            global_sampling_factor=global_sampling_factor,
+            snrs=snrs,
+            gains=gains,
         )
+        for message in self.loader.get_log_messages():
+            logger.log(message[0], "Dataloader: " + message[1])
         self.prefetch = prefetch
         self.pin_memory = pin_memory if torch.cuda.is_available() else False
         self.idx = 0
@@ -123,8 +127,14 @@ def __init__(
         self.pin_memory_thread_done_event: Optional[threading.Event] = None
         self.pin_memory_thread: Optional[threading.Thread] = None
         self.data_queue = None
+        self.log_timings = log_timings
+        self.timings_py: List[float] = []
+        self.timings_rs: List[torch.Tensor] = []
         atexit.register(self.loader.cleanup)
 
+    def set_batch_size(self, batch_size: int, split: str):
+        self.loader.set_batch_size(batch_size, split)
+
     def cleanup_pin_memory_thread(self):
         if self.pin_memory:
             # Check if still running from previous epoch
@@ -197,9 +207,16 @@ def get(*__args, **__kwargs) -> Tuple[int, Batch]:
         return q
 
     def _get_batch(self) -> Batch:
+        if self.log_timings:
+            t0 = time.time()
         _, batch = self.data_queue.get()
         if isinstance(batch, ExceptionWrapper):
             batch.reraise()
+        if self.log_timings:
+            self.timings_py.append(time.time() - t0)
+            self.timings_rs.append(batch.timings)
+        for message in self.loader.get_log_messages():
+            logger.log(message[0], "Dataloader: " + message[1])
         return batch
 
     def iter_epoch(self, split: str, seed: int) -> Iterator[Batch]:
@@ -207,6 +224,11 @@ def iter_epoch(self, split: str, seed: int) -> Iterator[Batch]:
         # Initializes workers. This needs to be done before pin_memory thread is
         # started via setup_data_queue().
         self.loader.start_epoch(split, seed)
+        for message in self.loader.get_log_messages():
+            logger.log(message[0], "Dataloader: " + message[1])
+        if self.log_timings:
+            self.timings_py = []
+            self.timings_rs = []
         # Initialize data out queue (maybe incl. pin memory thread)
         self.setup_data_queue()
         try:
@@ -225,4 +247,19 @@ def iter_epoch(self, split: str, seed: int) -> Iterator[Batch]:
         except StopIteration:
             if self.pin_memory_thread_done_event is not None:
                 self.pin_memory_thread_done_event.set()
+            if self.log_timings:
+                logger.info("Avg batch loading time (py): {:.4f}s", np.mean(self.timings_py))
+                # Last batch may be shorter, just skip it
+                if len(self.timings_rs[-1]) != len(self.timings_rs[0]):
+                    self.timings_rs = self.timings_rs[:-1]
+                timings_rs = torch.stack(self.timings_rs)
+                timings_samples = timings_rs[:, :-1]
+                timings_batch = timings_rs[:, -1]  # last if for whole batch
+                logger.info("Avg batch loading time (rs): {:.4f}s", timings_batch.mean().item())
+                logger.info(
+                    "Min/Avg/Max sample loading time (rs): {:.4f}/{:.4f}/{:.4f}s",
+                    timings_samples.min().item(),
+                    timings_samples.mean().item(),
+                    timings_samples.max().item(),
+                )
             return
diff --git a/pyDF-data/pyproject.toml b/pyDF-data/pyproject.toml
index 32641c39d..d4631ebfe 100644
--- a/pyDF-data/pyproject.toml
+++ b/pyDF-data/pyproject.toml
@@ -17,7 +17,7 @@ strip = true
 
 [tool.black]
 line-length = 100
-target-version = ["py37", "py38", "py39"]
+target-version = ["py37", "py38", "py39", "py310"]
 include = '\.pyi?$'
 
 [tool.isort]
diff --git a/pyDF-data/src/lib.rs b/pyDF-data/src/lib.rs
index 3ea19d0ef..ec049e8a9 100644
--- a/pyDF-data/src/lib.rs
+++ b/pyDF-data/src/lib.rs
@@ -1,6 +1,14 @@
+use std::path::{Path, PathBuf};
+use std::thread;
+use std::time::Instant;
+
+use crossbeam_channel::{unbounded, Receiver, TryRecvError};
 use df::augmentations::seed_from_u64;
 use df::dataloader::{DataLoader, DfDataloaderError};
-use df::dataset::{DatasetBuilder, DatasetConfigJson, Datasets, DfDatasetError, Split::*};
+use df::dataset::{
+    DatasetBuilder, DatasetConfigCacheJson, DatasetConfigJson, Datasets, DfDatasetError,
+    FftDataset, Hdf5Cfg, LogLevel, Split,
+};
 use df::Complex32;
 use ndarray::{ArrayD, ShapeError};
 use numpy::{IntoPyArray, PyArray1, PyArray4};
@@ -18,6 +26,7 @@ struct _FdDataLoader {
     loader: DataLoader,
     finished: bool,
     cur_id: isize,
+    logger: Receiver<(LogLevel, String)>,
 }
 
 // TODO: Does not work due to pyo3 restrictions; instead return tuples
@@ -59,7 +68,7 @@ type FdBatch<'py> = (
     &'py PyArray1<usize>,     // max_freq
     &'py PyArray1<i8>,        // snr
     &'py PyArray1<i8>,        // gain
-    &'py PyArray1<u8>,        // attenuation limit
+    &'py PyArray1<f32>,       // Timings until each sample and the overall batch was ready
 );
 
 #[pymethods]
@@ -67,6 +76,7 @@ impl _FdDataLoader {
     #[allow(clippy::too_many_arguments)]
     #[new]
     fn new(
+        py: Python,
         ds_dir: &str,
         config_path: &str,
         sr: usize,
@@ -80,15 +90,19 @@ impl _FdDataLoader {
         norm_alpha: Option<f32>,
         num_threads: Option<usize>,
         prefetch: Option<usize>,
-        p_atten_lim: Option<f32>,
         p_reverb: Option<f32>,
         drop_last: Option<bool>,
         overfit: Option<bool>,
+        cache_valid: Option<bool>,
         seed: Option<u64>,
         min_nb_erb_freqs: Option<usize>,
+        global_sampling_factor: Option<f32>,
+        snrs: Option<Vec<i8>>,
+        gains: Option<Vec<i8>>,
     ) -> PyResult<Self> {
+        let (log_sender, log_receiver) = unbounded();
         seed_from_u64(42);
-        let cfg = match DatasetConfigJson::open(config_path) {
+        let mut cfg = match DatasetConfigJson::open(config_path) {
             Err(e) => {
                 return Err(PyRuntimeError::new_err(format!(
                     "DF dataset config not found at '{}' ({:?})",
@@ -97,14 +111,13 @@ impl _FdDataLoader {
             }
             Ok(cfg) => cfg,
         };
+        load_cache(config_path, &mut cfg);
         let mut ds_builder = DatasetBuilder::new(ds_dir, sr)
             .df_params(fft_size, hop_size, nb_erb, nb_spec, norm_alpha);
+        py.check_signals()?;
         if let Some(max_len_s) = max_len_s {
             ds_builder = ds_builder.max_len(max_len_s)
         }
-        if let Some(p_atten_lim) = p_atten_lim {
-            ds_builder = ds_builder.prob_atten_lim(p_atten_lim)
-        }
         if let Some(seed) = seed {
             ds_builder = ds_builder.seed(seed)
         }
@@ -114,22 +127,49 @@ impl _FdDataLoader {
         if let Some(nb_freqs) = min_nb_erb_freqs {
             ds_builder = ds_builder.min_nb_erb_freqs(nb_freqs)
         }
-        let valid_ds = ds_builder
-            .clone()
-            .dataset(cfg.split_config(Valid))
-            .build_fft_dataset()
-            .to_py_err()?;
-        let test_ds = ds_builder
-            .clone()
-            .dataset(cfg.split_config(Test))
-            .build_fft_dataset()
-            .to_py_err()?;
-        ds_builder = ds_builder.p_sample_full_speech(1.0);
-        let train_ds = ds_builder
-            .clone()
-            .dataset(cfg.split_config(Train))
-            .build_fft_dataset()
-            .to_py_err()?;
+        if let Some(f) = global_sampling_factor {
+            ds_builder = ds_builder.global_sample_factor(f)
+        }
+        if let Some(snrs) = snrs {
+            ds_builder = ds_builder.snrs(snrs);
+        }
+        if let Some(gains) = gains {
+            ds_builder = ds_builder.gains(gains);
+        }
+        if let Some(num_threads) = num_threads {
+            ds_builder = ds_builder.num_threads(num_threads);
+        }
+        let valid_handle = {
+            let valid_cfg = cfg.split_config(Split::Valid);
+            let valid_ds_builder = ds_builder.clone().add_logger(log_sender.clone());
+            let valid_ds_builder = if cache_valid.unwrap_or(false) {
+                valid_ds_builder.cache_valid_dataset()
+            } else {
+                valid_ds_builder
+            };
+            thread::spawn(|| valid_ds_builder.dataset(valid_cfg).build_fft_dataset())
+        };
+        let test_handle = {
+            let test_cfg = cfg.split_config(Split::Test);
+            let test_ds_builder = ds_builder.clone().add_logger(log_sender.clone());
+            thread::spawn(|| test_ds_builder.dataset(test_cfg).build_fft_dataset())
+        };
+        ds_builder = ds_builder.p_sample_full_speech(1.0).add_logger(log_sender);
+        let train_handle = {
+            let train_cfg = cfg.split_config(Split::Train);
+            thread::spawn(|| ds_builder.dataset(train_cfg).build_fft_dataset())
+        };
+        let msg = "Unable to join dataset builder thread";
+        let valid_ds = valid_handle.join().expect(msg).to_py_err()?;
+        update_keys(ds_dir, &mut cfg.valid, &valid_ds);
+        py.check_signals()?;
+        let test_ds = test_handle.join().expect(msg).to_py_err()?;
+        update_keys(ds_dir, &mut cfg.test, &test_ds);
+        py.check_signals()?;
+        let train_ds = train_handle.join().expect(msg).to_py_err()?;
+        update_keys(ds_dir, &mut cfg.train, &train_ds);
+        write_cache(config_path, &cfg);
+        py.check_signals()?;
         let ds = Datasets {
             train: train_ds,
             valid: valid_ds,
@@ -145,17 +185,19 @@ impl _FdDataLoader {
         if let Some(bs_eval) = batch_size_eval {
             dl_builder = dl_builder.batch_size_eval(bs_eval);
         }
-        if let Some(drop_last) = drop_last {
-            dl_builder = dl_builder.drop_last(drop_last);
+        if drop_last.unwrap_or(false) {
+            dl_builder = dl_builder.drop_last();
         }
-        if let Some(overfit) = overfit {
-            dl_builder = dl_builder.overfit(overfit);
+        if overfit.unwrap_or(false) {
+            dl_builder = dl_builder.overfit();
         }
+        py.check_signals()?;
         let loader = dl_builder.build().to_py_err()?;
         Ok(_FdDataLoader {
             loader,
             finished: false,
             cur_id: -1,
+            logger: log_receiver,
         })
     }
 
@@ -169,6 +211,7 @@ impl _FdDataLoader {
     }
 
     fn get_batch<'py>(&'py mut self, py: Python<'py>) -> PyResult<FdBatch<'py>> {
+        let t0 = Instant::now();
         if self.finished {
             return Err(PyStopIteration::new_err("Epoch finished"));
         }
@@ -190,7 +233,7 @@ impl _FdDataLoader {
                     batch.max_freq.into_pyarray(py),
                     batch.snr.into_pyarray(py),
                     batch.gain.into_pyarray(py),
-                    batch.atten.into_pyarray(py),
+                    push_ret(batch.timings, (Instant::now() - t0).as_secs_f32()).into_pyarray(py),
                 ))
             }
             None => {
@@ -212,6 +255,77 @@ impl _FdDataLoader {
     fn dataset_len(&self, split: &str) -> usize {
         self.loader.dataset_len(split)
     }
+
+    fn set_batch_size(&mut self, batch_size: usize, split: &str) {
+        self.loader.set_batch_size(batch_size, split)
+    }
+
+    fn get_log_messages(&mut self) -> Vec<(String, String)> {
+        let mut messages = Vec::new();
+        loop {
+            match self.logger.try_recv() {
+                Ok(m) => messages.push((m.0.into(), m.1)),
+                Err(TryRecvError::Empty) => break,
+                Err(TryRecvError::Disconnected) => {
+                    eprintln!("Dataloader logger disconnected unexpectetly!");
+                    break;
+                }
+            }
+        }
+        messages
+    }
+}
+fn cache_path(cfg_path: &str) -> PathBuf {
+    let mut p = Path::new(cfg_path).to_path_buf();
+    let cache_file_name = p.file_stem().unwrap().to_str().unwrap().to_owned();
+    p.set_file_name(".cache_".to_owned() + &cache_file_name);
+    p.set_extension("cfg");
+    p
+}
+fn load_cache(cfg_path: &str, cfg: &mut DatasetConfigJson) {
+    let cache_path = cache_path(cfg_path);
+    if !cache_path.is_file() {
+        return;
+    }
+    println!(
+        "Loading HDF5 key cache from {}",
+        cache_path.to_str().unwrap_or_default()
+    );
+    match DatasetConfigCacheJson::open(cache_path.to_str().unwrap()) {
+        Err(e) => eprintln!("Could not load dataset keys cache: {}", e),
+        Ok(cache) => {
+            cfg.set_keys(Split::Train, cache.keys()).expect("Could not set cached keys");
+            cfg.set_keys(Split::Valid, cache.keys()).expect("Could not set cached keys");
+            cfg.set_keys(Split::Test, cache.keys()).expect("Could not set cached keys");
+        }
+    }
+}
+fn write_cache(cfg_path: &str, cfg: &DatasetConfigJson) {
+    let cache_path = cache_path(cfg_path);
+    let mut cache = Vec::new();
+    cache.extend(cfg.train.iter().map(|x| x.keys_unchecked().unwrap().clone()));
+    cache.extend(cfg.valid.iter().map(|x| x.keys_unchecked().unwrap().clone()));
+    cache.extend(cfg.test.iter().map(|x| x.keys_unchecked().unwrap().clone()));
+    let cache = DatasetConfigCacheJson::new(cache);
+    cache.write(cache_path.to_str().unwrap()).expect("Failed to write cache.");
+}
+/// Upates HDF5 keys from the dataset to cfgs
+fn update_keys(ds_dir: &str, cfgs: &mut [Hdf5Cfg], ds: &FftDataset) {
+    for hdf5cfg in cfgs.iter_mut() {
+        let ds_path = ds_dir.to_owned() + "/" + hdf5cfg.filename();
+        let cfg = ds.get_hdf5cfg(hdf5cfg.filename()).expect("Could not get hdf5cfg");
+        let hash = cfg
+            .hash()
+            .unwrap_or_else(|| cfg.hash_from_ds_path(&ds_path).expect("Could not calculate hash"));
+        if let Some(ds_keys) = cfg.load_keys(hash).expect("Could not load Hdf5Keys.") {
+            hdf5cfg.set_keys(ds_keys.clone()).expect("Could not update keys");
+        }
+    }
+}
+
+fn push_ret<T>(mut a: Vec<T>, b: T) -> Vec<T> {
+    a.push(b);
+    a
 }
 
 fn sort<A, T>(mut array: A) -> A
diff --git a/pyDF/libdf.pyi b/pyDF/libdf.pyi
index f0ddc8425..72d291ca6 100644
--- a/pyDF/libdf.pyi
+++ b/pyDF/libdf.pyi
@@ -42,6 +42,7 @@ class DF:
         ...
     def erb_widths(self) -> ndarray: ...
     def fft_window(self) -> ndarray: ...
+    def sr(self) -> int: ...
     def reset(self) -> None: ...
 
 def erb(
diff --git a/pyDF/pyproject.toml b/pyDF/pyproject.toml
index f5d9e6360..29fa5404b 100644
--- a/pyDF/pyproject.toml
+++ b/pyDF/pyproject.toml
@@ -17,7 +17,7 @@ strip = true
 
 [tool.black]
 line-length = 100
-target-version = ["py37", "py38", "py39"]
+target-version = ["py37", "py38", "py39", "py310"]
 include = '\.pyi?$'
 
 [tool.isort]
diff --git a/pyDF/src/lib.rs b/pyDF/src/lib.rs
index 69eed36a1..704450938 100644
--- a/pyDF/src/lib.rs
+++ b/pyDF/src/lib.rs
@@ -108,6 +108,10 @@ impl DF {
         Ok(self.state.window.clone().into_pyarray(py))
     }
 
+    fn sr(&self) -> usize {
+        self.state.sr
+    }
+
     fn reset(&mut self) {
         self.state.reset();
     }
diff --git a/pyproject.toml b/pyproject.toml
index c0bf8d61e..30e35c23c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.black]
 line-length = 100
-target-version = ["py37", "py38", "py39"]
+target-version = ["py37", "py38", "py39", "py310"]
 include = '\.pyi?$'
 
 [tool.isort]
diff --git a/scripts/copy_datadir.py b/scripts/copy_datadir.py
new file mode 100755
index 000000000..55f800af0
--- /dev/null
+++ b/scripts/copy_datadir.py
@@ -0,0 +1,206 @@
+import argparse
+import concurrent.futures
+import json
+import os
+import shutil
+import subprocess
+import warnings
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from time import sleep
+from typing import DefaultDict, Dict, List, Optional, Tuple
+
+import h5py
+
+TIMESTAMP_FORMAT = "%Y%m%d%H%M"
+timestamp = datetime.now().strftime("%Y%m%d%H%M")
+
+
+@dataclass
+class DsConfig:
+    sampling_rate: int
+    sampling_factor: int
+    max_freq: int
+
+
+def ln(src, tgt):
+    """Link a file via ln"""
+    return subprocess.check_call(["ln", "-s", src, tgt])
+
+
+def du(path):
+    """disk usage in human readable format (e.g. '2,1')"""
+    return float(
+        subprocess.check_output(["du", "-shD", "--block-size=1G", path]).split()[0].decode("utf-8")
+    )
+
+
+def cp(src, tgt, verbose=0):
+    """Copy a file via rsync"""
+    info = []
+    if verbose == 1:
+        info = ["--info=name,stats"]
+    elif verbose > 1:
+        info = ["--info=name,stats2"]
+    try:
+        subprocess.check_call(["rsync", "-aL", *info, src, tgt])
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to copy file {src}: {e.returncode}")
+        if os.path.exists(tgt):
+            os.remove(tgt)
+        print("Linking instead")
+        ln(src, tgt)
+
+
+def has_locks(directory: str, lock: Optional[str] = None, wait_write_lock: bool = False):
+    lock_f = os.path.join(directory, ".lock")
+    have_read_locks = False
+    have_write_locks = False
+    # We can have a write lock allowing exclusive access as well as multiple parallel read locks
+    if os.path.isfile(lock_f):
+        tries = 1
+        while any(
+            line.strip().endswith(".write")
+            and (lock is not None and not line.strip().startswith(lock))
+            for line in open(lock_f)
+        ):
+            have_write_locks = True
+            if not wait_write_lock:
+                break
+            # Wait until the current write lock is released
+            warnings.warn(f"<copy_datadir.py>: Could not lock directory {directory}")
+            sleep(tries)
+            tries *= 2
+            if tries > 2**11:  # 2**11 ~ 34 minutes
+                break
+        have_read_locks = False
+        cur_timestamp = datetime.strptime(timestamp, TIMESTAMP_FORMAT)
+        for line in open(lock_f):
+            line = line.strip()
+            if lock is not None and line.startswith(lock):
+                continue  # Same lock should not be relevant
+            try:
+                lock_timestamp = line.split(".")[1]
+                lock_timestamp = datetime.strptime(lock_timestamp, TIMESTAMP_FORMAT)
+            except Exception as e:
+                print(e)
+                continue
+            if cur_timestamp - lock_timestamp < timedelta(days=1):
+                print("Found existing lock", line.strip())
+                have_read_locks = True
+                break
+    return have_read_locks, have_write_locks
+
+
+def copy_datasets(
+    src_dir: str, target_dir: str, cfg_path: str, max_gb: float, lock: Optional[str] = None
+):
+    print(f"Copying datasets from {src_dir} to {target_dir}", flush=True)
+    os.makedirs(target_dir, exist_ok=True)
+    lock_f = os.path.join(target_dir, ".lock")
+    have_read_locks, have_write_locks = has_locks(target_dir, lock, wait_write_lock=True)
+    # Lock the target dir for writing
+    open(lock_f, "a+").write(f"\n{lock}.{timestamp}.write")
+    cfg = json.load(open(cfg_path))
+    os.makedirs(target_dir, exist_ok=True)
+    executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)
+    futures = {}
+    cur_gb = du(target_dir)
+    print(f"Current target dir size: {cur_gb}")
+    print(f"Copying max {max_gb} GB")
+    # Start with train since it will be accessed most of the time
+    for split in ("train", "valid", "test"):
+        # Get all Datasets and sort by dataset type and sampling factor
+        # {DsType: (file_name, sampling_factor)}, DsType is one of (speech, noise, rir)
+        datasets: Dict[str, List[Tuple[str, float]]] = DefaultDict(list)
+        for entry in cfg[split]:
+            fn = entry[0]
+            dstype = list(h5py.File(os.path.join(src_dir, fn)).keys())[0]
+            datasets[dstype].append((fn, float(entry[1])))
+        for dstype in datasets.keys():
+            datasets[dstype] = sorted(datasets[dstype], key=lambda e: e[1], reverse=True)
+        # We will mix multiple noises with a single speech signal, thus start with noise datasets,
+        # since they will be accessed more often.
+        for dstype in ("noise", "speech", "rir"):
+            for fn, _ in datasets[dstype]:
+                fn_src = os.path.join(src_dir, fn)
+                fn_tgt = os.path.join(target_dir, fn)
+                new_gb = du(fn_src)
+                if cur_gb + new_gb > max_gb or "test" in fn_src.lower():
+                    # If too large or a test dataset, link instead
+                    if not os.path.exists(fn_tgt):
+                        print("linking", fn_src, flush=True)
+                        ln(fn_src, fn_tgt)
+                    elif (
+                        not have_read_locks
+                        and os.path.isfile(fn_tgt)
+                        and not os.path.islink(fn_tgt)
+                    ):
+                        # Just run rsync again to make sure the file is not corrupt
+                        print("checking", fn_tgt, flush=True)
+                        futures[executor.submit(cp, fn_src, fn_tgt)] = fn_tgt
+                else:
+                    if os.path.islink(fn_tgt) and not have_read_locks:
+                        os.remove(fn_tgt)  # Only remove if no other process has a readlock
+                    elif have_read_locks and os.path.isfile(fn_tgt):
+                        continue
+                    print("copying", fn_src, flush=True)
+                    cur_gb += new_gb
+                    futures[executor.submit(cp, fn_src, fn_tgt)] = fn_tgt
+    for future in concurrent.futures.as_completed(futures):
+        print("Completed", futures[future], flush=True)
+
+    if lock is not None:
+        remove_lock(target_dir, lock, lock + "." + timestamp + ".read")
+
+
+def remove_lock(target_dir: str, lock: str, new_lock: Optional[str] = None):
+    lock_f = os.path.join(target_dir, ".lock")
+    with open(lock_f, "r+") as f:
+        lines = []
+        for line in f.readlines():
+            line = line.strip()
+            if line == "" or line.startswith(lock):
+                continue
+            if new_lock is not None and line.startswith(new_lock):
+                continue
+            lines.append(line + "\n")
+        if new_lock is not None:
+            lines.append(new_lock + "\n")
+        f.seek(0)
+        f.writelines(lines)
+        f.truncate()
+
+
+def cleanup(target_dir: str, prev_lock: Optional[str] = None):
+    have_read_locks, have_write_locks = has_locks(target_dir, prev_lock, wait_write_lock=True)
+    if have_read_locks or have_write_locks:
+        print("Could not cleanup due to existing locks.")
+        return
+    print(f"Cleaning directory {target_dir}")
+    shutil.rmtree(target_dir)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest="subparser_name")
+    cp_parser = subparsers.add_parser("copy-datasets", aliases=["cp"])
+    cp_parser.add_argument("src_dir")
+    cp_parser.add_argument("target_dir")
+    cp_parser.add_argument("data_cfg")
+    cp_parser.add_argument("--max-gb", type=float, default=100)
+    cp_parser.add_argument("--lock", "-l", type=str, default=None)
+    cleanup_parser = subparsers.add_parser("cleanup")
+    cleanup_parser.add_argument("target_dir")
+    cleanup_parser.add_argument("--lock", type=str, default=None)
+    args = parser.parse_args()
+    if args.subparser_name is None:
+        parser.print_help()
+        exit(1)
+    if args.subparser_name in ("cp", "copy-datasets"):
+        copy_datasets(args.src_dir, args.target_dir, args.data_cfg, args.max_gb, args.lock)
+    else:
+        if args.lock is not None:
+            remove_lock(args.target_dir, args.lock + ".read")
+            remove_lock(args.target_dir, args.lock + ".write")
+        cleanup(args.target_dir, args.lock)
diff --git a/scripts/sbatch_train.sh b/scripts/sbatch_train.sh
index 98d93c719..fc8c73881 100644
--- a/scripts/sbatch_train.sh
+++ b/scripts/sbatch_train.sh
@@ -9,14 +9,14 @@
 #SBATCH --mail-type=ALL
 # Time limit format: "hours:minutes:seconds"
 #SBATCH --time=24:00:00
-# Send the SIGUSR1 signal 2 h before time limit
-#SBATCH --signal=B:USR1@7200
+# Send the SIGUSR1 signal 4 h before time limit
+#SBATCH --signal=B:USR1@14400
 
 set -e
 
 # Cluster directory containing data, project code, logging, summaries,
 # checkpoitns, etc.
-export CLUSTER=/cluster/$USER
+export CLUSTER=/net/cluster/$USER
 cd "$CLUSTER"
 
 # Workaround for our cluster
@@ -33,13 +33,14 @@ DATA_CFG=${DATA_CFG:-$DATA_DIR/datasets.cfg} # Default dataset configuration
 DATA_CFG=$(readlink -f "$DATA_CFG")
 PYTORCH_JIT=${PYTORCH_JIT:-1}                # Set to 0 to disable pytorch JIT compilation
 COPY_DATA=${COPY_DATA:-1}                    # Copy data
+COPY_MAX_GB=${COPY_MAX_GB:-50}               # Max amount to copy hdf5 datasets, rest will be linked
 DEBUG=${DEBUG:-0}                            # Debug mode passed to the python train script
 EXCLUDE=${EXCLUDE:-lme[49,170,171]}          # Slurm nodes to exclude
 
 if [ "$DEBUG" -eq 1 ]; then
   DEBUG="--debug"
 else
-  DEBUG=
+  DEBUG="--no-debug"
 fi
 
 echo "Started sbatch script at $(date) in $(pwd)"
@@ -98,7 +99,7 @@ if [ "$PROJECT_BRANCH_CUR" != "$PROJECT_BRANCH" ]; then
   git -C "$PROJECT_HOME" stash save "$stash"
   git -C "$PROJECT_HOME" checkout "$PROJECT_BRANCH"
   stash_idx=$(git -C "$PROJECT_HOME" stash list | grep "$stash" | cut -d: -f1)
-  if [ ! -z "$stash_idx" -a "$stash_idx" != " " ]; then
+  if [ -n "$stash_idx" ] && [ "$stash_idx" != " " ]; then
     # Try to apply current stash; If not possible just proceed.
     if ! git -C "$PROJECT_HOME" stash pop "$stash_idx"; then
       echo "Could not apply stash to branch $PROJECT_BRANCH"
@@ -117,22 +118,10 @@ setup_env "$CLUSTER" "$PROJECT_HOME" "$MODEL_NAME"
 if [[ -d /scratch ]] && [[ $COPY_DATA -eq 1 ]]; then
   test -d "/scratch/$USER" || mkdir "/scratch/$USER"
   NEW_DATA_DIR=/scratch/"$USER"/"$PROJECT_NAME"
-  echo "Seting up data dir in $NEW_DATA_DIR"
+  echo "Setting up data dir in $NEW_DATA_DIR"
   mkdir -p "$NEW_DATA_DIR"
-  # Check if another process is currently copying
-  while true; do
-    LOCK=$(rg lock "$NEW_DATA_DIR"/data_lock | sed "s/-lock//g")
-    if [[ -n $LOCK ]]; then
-      echo "Scratch dir is currently locked by: $LOCK"
-      sleep 30s
-    else
-      break
-    fi
-  done
-  echo "$MODEL_NAME-lock" >> "$NEW_DATA_DIR"/data_lock  # lock scratch dir
-  MAX_GB=150 "$PROJECT_HOME"/scripts/copy_datadir.sh "$DATA_DIR" "$DATA_CFG" "$NEW_DATA_DIR"
-  # release copy lock; remaining $MODEL_NAME indicates we are using it
-  sed -i "s/$MODEL_NAME-lock/$MODEL_NAME/g"  "$NEW_DATA_DIR"/data_lock
+  python3 "$PROJECT_HOME"/scripts/copy_datadir.py cp "$DATA_DIR" "$NEW_DATA_DIR" "$DATA_CFG" \
+    --lock "$MODEL_NAME" --max-gb "$COPY_MAX_GB"
   DATA_DIR="$NEW_DATA_DIR"
 fi
 
@@ -141,6 +130,12 @@ fi
 # Therefore, we send a SIGUSR1 to the python training script which needs to
 # write a continue file if it wants to get restarted. After training scripts
 # returns, check if there is a continue file and resubmit.
+function _cleanup_scratch {
+  echo "Checking if need to cleanup scratch: $NEW_DATA_DIR"
+  if [[ -d /scratch ]] && [[ $COPY_DATA -eq 1 ]]; then
+    python3 "$PROJECT_HOME"/scripts/copy_datadir.py cleanup "$NEW_DATA_DIR" --lock "$MODEL_NAME"
+  fi
+}
 function _at_exit {
   conda deactivate
   # Check for return code if training was completed
@@ -154,19 +149,10 @@ function _at_exit {
       $PROJECT_HOME/scripts/sbatch_train.sh $BASE_DIR $PROJECT_HOME $PROJECT_BRANCH"
     exit 0
   fi
-  echo "Checking if need to cleanup scratch: $NEW_DATA_DIR"
-  if [[ -d /scratch ]] && [[ $COPY_DATA -eq 1 ]]; then
-    cat "$NEW_DATA_DIR"/data_lock
-    # Remove own lock
-    sed -i /"$MODEL_NAME"/d "$NEW_DATA_DIR"/data_lock
-    if ! wc -l "$NEW_DATA_DIR"/data_lock; then
-      # No other locks found. Cleanup space.
-      echo "cleaning up data dir $NEW_DATA_DIR"
-      rm -r "$NEW_DATA_DIR"
-    fi
-  fi
+  _cleanup_scratch
 }
 trap _at_exit EXIT
+trap _cleanup_scratch ERR
 
 function _usr1 {
   echo "Caught SIGUSR1 signal!"
@@ -181,10 +167,11 @@ cd "$PROJECT_HOME"/DeepFilterNet/df/
 printf "\n***Starting training***\n\n"
 
 PYTHONPATH="$PROJECT_HOME/DeepFilterNet/" python train.py \
+  --host-batchsize-config $CLUSTER/host_batchsize.ini \
   "$DATA_CFG" \
   "$DATA_DIR" \
   "$BASE_DIR" \
-  $DEBUG &
+  "$DEBUG" &
 
 trainprocess=$!
 echo "Started trainprocess: $trainprocess"
diff --git a/scripts/set_batch_size.py b/scripts/set_batch_size.py
new file mode 100644
index 000000000..8b09fa415
--- /dev/null
+++ b/scripts/set_batch_size.py
@@ -0,0 +1,104 @@
+import os
+import sys
+from configparser import ConfigParser
+from socket import gethostname
+from typing import List, Tuple
+
+
+def update_batch_size(
+    host_key: str,
+    config_parser: ConfigParser,
+    host_bs_parser: ConfigParser,
+    config_key: str,
+    batchconfig_key: str,
+) -> Tuple[bool, bool]:
+    current_bs = None
+    if config_parser.has_section("train"):
+        current_bs = config_parser.get("train", config_key, fallback=None)
+    batchconfig_bs = host_bs_parser.get(host_key, batchconfig_key, fallback=None)
+    config_changed = False
+    host_bs_changed = False
+    if batchconfig_bs is not None:
+        # Overwrite training config
+        if current_bs is not None and batchconfig_bs != current_bs:
+            print(
+                f"Found host specific {batchconfig_key} ({batchconfig_bs}) for host {host_key}. "
+                "Updating config."
+            )
+            if not config_parser.has_section("train"):
+                config_parser.add_section("train")
+            config_parser.set("train", config_key, batchconfig_bs)
+            config_changed = True
+    elif current_bs is not None:
+        # No host_key specific batch config found. Store current batch size for current host
+        print(
+            f"Host specific {batchconfig_key} not found for host {host_key}. "
+            "Updating host batch size config."
+        )
+        if not host_bs_parser.has_section(host_key):
+            host_bs_parser.add_section(host_key)
+        host_bs_parser.set(host_key, batchconfig_key, current_bs)
+        host_bs_changed = True
+    return config_changed, host_bs_changed
+
+
+def cast_bool(value):
+    value = str(value).lower()
+    if value in {"true", "yes", "y", "on", "1"}:
+        return True  # type: ignore
+    elif value in {"false", "no", "n", "off", "0"}:
+        return False  # type: ignore
+
+
+def main(config_path, host_bs_config, host_key=None):
+    """Sets host specific batch sizes in the specified config."""
+    if not os.path.isfile(config_path):
+        print(f"Config not found at path {config_path}.")
+        exit(2)
+    if not os.path.isfile(host_bs_config):
+        print(f"Host specific batch size config not found at path {host_bs_config}.")
+        print("Creating default.")
+        open(host_bs_config, "w").close()
+    if host_key is None:
+        host_key = gethostname()
+    config_parser = ConfigParser()
+    host_bs_parser = ConfigParser()
+    with open(config_path, "r") as f:
+        config_parser.read_file(f)
+    with open(host_bs_config, "r") as f:
+        host_bs_parser.read_file(f)
+    # For eval, don't distinguish between autocasting
+    changed: List[Tuple[bool, bool]] = []
+    changed.append(
+        update_batch_size(
+            host_key, config_parser, host_bs_parser, "batch_size_eval", "batch_size_eval"
+        )
+    )
+    autocast_enabled = cast_bool(config_parser.get("train", "train_autocast", fallback=False))
+    if autocast_enabled:
+        changed.append(
+            update_batch_size(
+                host_key, config_parser, host_bs_parser, "batch_size", "batch_size_autocast_train"
+            )
+        )
+    else:
+        changed.append(
+            update_batch_size(
+                host_key, config_parser, host_bs_parser, "batch_size", "batch_size_train"
+            )
+        )
+    # check whether the training config was changed
+    if any(c[0] for c in changed):
+        with open(config_path, "w") as f:
+            config_parser.write(f)
+    # check whether the host batch size config was changed
+    if any(c[1] for c in changed):
+        with open(host_bs_config, "w") as f:
+            host_bs_parser.write(f)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 3:
+        print("Usage: set_batch_size <config-path> <host-bs-config-path>")
+        exit(1)
+    main(sys.argv[1], sys.argv[2])
diff --git a/scripts/setup_env.sh b/scripts/setup_env.sh
index 2f266d686..8faf985f2 100644
--- a/scripts/setup_env.sh
+++ b/scripts/setup_env.sh
@@ -3,6 +3,7 @@
 pytorch_v="1.10"
 cuda_version=11.1
 PYTHON_V=3.9
+MINICONDA_DIR=${MINICONDA_DIR:-"$HOME/miniconda"}
 
 PYTORCH_NIGHTLY=${PYTORCH_NIGHTLY:-0}
 if [ "$PYTORCH_NIGHTLY" -eq 1 ]; then
@@ -14,56 +15,69 @@ else
 fi
 INSTALL_LIBDF=${INSTALL_LIBDF:-1}
 INSTALL_PYDEPS=${INSTALL_PYDEPS:-0}
+INSTALL_PYTESTDEPS=${INSTALL_PYTESTDEPS:-0}
 SUFFIX=""
 
+check_install_conda() {
+  # Check if miniconda is already installed
+  if ! [ -f "$MINICONDA_DIR"/bin/python$PYTHON_V ]; then
+    if [ -d "$MINICONDA_DIR" ]; then
+      echo "Miniconda directory already exists, but python not found."
+      echo "Please cleanup miniconda dir first for an automatic installation."
+      exit 1
+    fi
+    echo "Miniconda not found. Installing..."
+    wget -nv https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh \
+      -O "$CLUSTER"/miniconda.sh
+    bash "$CLUSTER"/miniconda.sh -b -p "$MINICONDA_DIR"
+    rm "$CLUSTER"/miniconda.sh
+  fi
+}
+
+setup_conda() {
+  # Setup miniconda if found; this isn't done lazily, so expect some bash starup delay
+  __conda_setup=$("$MINICONDA_DIR/bin/conda" 'shell.bash' 'hook')
+  if [ $? -eq 0 ]; then
+    eval "$__conda_setup"
+  else
+    if [ -f "$MINICONDA_DIR/etc/profile.d/conda.sh" ]; then
+      . "$MINICONDA_DIR/etc/profile.d/conda.sh"
+    else
+      export PATH="$MINICONDA_DIR/bin:$PATH"
+    fi
+  fi
+  unset __conda_setup
+}
+
 setup_env() {
   CLUSTER=$1
-  MINICONDA_DIR="$CLUSTER/miniconda"
   PROJECT_HOME=$2
-  if [ ! -z "$3" ]; then
+  if [ -n "$3" ]; then
     SUFFIX="-$3"
   fi
+  echo "Running on CUDA: $cuda_version"
 
   env="df-$pytorch_v$nightly-cuda$cuda_version$SUFFIX"
   env_path="$MINICONDA_DIR/envs/$env"
   if ! [ -f "$env_path/bin/python$PYTHON_V" ]; then
-    # Check if miniconda is already installed
     echo "Virtualenv $env not found. Installing..."
-    if ! [ -f "$MINICONDA_DIR"/bin/python$PYTHON_V ]; then
-      if [ -d "$MINICONDA_DIR" ]; then
-        echo "Miniconda directory already exists, but python not found."
-        echo "Please cleanup miniconda dir first for an automatic installation."
-        exit 1
-      fi
-      echo "Miniconda not found. Installing..."
-      wget -nv https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh \
-        -O "$CLUSTER"/miniconda.sh
-      bash "$CLUSTER"/miniconda.sh -b -p "$MINICONDA_DIR"
-      rm "$CLUSTER"/miniconda.sh
-    fi
-    export PATH="$MINICONDA_DIR/bin:$PATH"
-
-    echo "Running on CUDA: $cuda_version"
+    check_install_conda
+    setup_conda
 
     echo "Installing conda env to $env_path"
     conda create -y -q -p "$env_path" \
       python=$PYTHON_V \
       "$pytorch_v_arg" \
       torchaudio \
-      maturin \
-      patchelf \
-      poetry \
       cudatoolkit=$cuda_version -c "pytorch$nightly" -c conda-forge
-    source activate "$env"
+    conda activate "$env"
 
     INSTALL_LIBDF=1
     INSTALL_PYDEPS=1
   else
     # Only need to activate the existing env
-    if ! [ -x "$(command -v conda)" ]; then
-      export PATH="$MINICONDA_DIR/bin:$PATH"
-    fi
-    source activate "$env"
+    setup_conda
+    conda activate "$env"
 
     echo "Running on env: $CONDA_DEFAULT_ENV"
   fi
@@ -72,7 +86,7 @@ setup_env() {
     cd "$PROJECT_HOME"/ || exit 10
     rustup default stable
     rustup update stable
-    cargo build --release
+    pip install -U maturin
     maturin develop --release -m "$PROJECT_HOME"/pyDF/Cargo.toml
     maturin develop --release -m "$PROJECT_HOME"/pyDF-data/Cargo.toml
   fi
@@ -80,6 +94,12 @@ setup_env() {
     echo "Installing requirements"
     echo pip install -r "$PROJECT_HOME"/DeepFilterNet/requirements.txt
     pip install -r "$PROJECT_HOME"/DeepFilterNet/requirements.txt
+    pip install h5py
+  fi
+  if [ $INSTALL_PYTESTDEPS -eq 1 ]; then
+    echo "Installing test requirements"
+    echo pip install -r "$PROJECT_HOME"/DeepFilterNet/requirements.txt
+    pip install -r "$PROJECT_HOME"/DeepFilterNet/requirements_test.txt
   fi
 }
 


Test patch:
diff --git a/.github/workflows/test_df.yml b/.github/workflows/test_df.yml
index 3bb4b3f31..fc1ab7503 100644
--- a/.github/workflows/test_df.yml
+++ b/.github/workflows/test_df.yml
@@ -27,6 +27,8 @@ jobs:
         include:
           - os: "windows-latest"
             extra: "-E soundfile"  # Soundfile backend for torchaudio
+          - os: "ubuntu-latest"
+            extra: "-E semetrics"  # Aditionally test against original matlab composite metrics
     defaults:
       run:
         shell: bash
diff --git a/DeepFilterNet/df/scripts/test_df.py b/DeepFilterNet/df/scripts/test_df.py
index 32c7b261a..2156a7573 100755
--- a/DeepFilterNet/df/scripts/test_df.py
+++ b/DeepFilterNet/df/scripts/test_df.py
@@ -1,68 +1,146 @@
 #!/usr/bin/env python
 
 import os
+import unittest
+from typing import Dict, List, Union
 
 import numpy as np
 import torch
 from loguru import logger
 
 import df
-from df.deepfilternet import ModelParams
-from df.enhance import enhance, init_df, load_audio
-from df.scripts.test_voicebank_demand import HAS_OCTAVE, composite, si_sdr_speechmetrics, stoi
+from df.enhance import DF, enhance, init_df, load_audio
+from df.evaluation_utils import composite, si_sdr_speechmetrics, stoi
 
 __a_tol = 1e-4
 
 
-def try_eval_composite(clean, enhanced, sr):
-    if not HAS_OCTAVE:
-        logger.warning("Octave not found. Skipping.")
-        return
+def eval_composite(clean, enhanced, sr, m_target: List[float]):  # type: ignore
     logger.info("Computing composite metrics")
+    try:
+        m_enh_octave = torch.as_tensor(
+            composite(clean.squeeze(0).numpy(), enhanced.squeeze(0).numpy(), sr, use_octave=True)
+        ).to(torch.float32)
+    except (OSError, ImportError, ModuleNotFoundError):
+        m_enh_octave = None
+        logger.warning("No octave available")
     m_enh = torch.as_tensor(
         composite(clean.squeeze(0).numpy(), enhanced.squeeze(0).numpy(), sr)
     ).to(torch.float32)
     logger.info(f"Got {m_enh}")
-    m_target = torch.as_tensor(
-        [2.30616855621338, 3.832779407501221, 2.362725973129273, 3.05537247657776, -2.7911112308502]
-    )
+    m_target: torch.Tensor = torch.as_tensor(m_target)
     assert torch.isclose(
         m_enh, m_target, atol=__a_tol
     ).all(), f"Metric output not close. Expected {m_target}, got {m_enh}, diff: {m_target-m_enh}"
+    if m_enh_octave is not None:
+        assert torch.isclose(
+            m_enh_octave, m_target, atol=__a_tol
+        ).all(), (
+            f"Metric output not close. Expected {m_target}, got {m_enh}, diff: {m_target-m_enh}"
+        )
 
 
-def eval_pystoi(clean, enhanced, sr):
+def eval_pystoi(clean, enhanced, sr, m_target: float):
     logger.info("Computing STOI")
     m_enh = stoi(clean.squeeze(0), enhanced.squeeze(0), sr)
-    m_target = 0.9689496585281197
-    logger.info(f"Got {m_enh:.4f}")
+    logger.info(f"Got {m_enh}")
     assert np.isclose(
         [m_enh], [m_target], atol=__a_tol
     ), f"Metric output not close. Expected {m_target}, got {m_enh}, diff: {m_target-m_enh}"
 
 
-def eval_sdr(clean, enhanced):
+def eval_sdr(clean, enhanced, m_target: float):
     logger.info("Computing SI-SDR")
     m_enh = si_sdr_speechmetrics(clean.numpy(), enhanced.numpy())
-    m_target = 18.88543128967285
-    logger.info(f"Got {m_enh:.4f}")
+    logger.info(f"Got {m_enh}")
     assert np.isclose(
         [m_enh], [m_target]
     ), f"Metric output not close. Expected {m_target}, got {m_enh}, diff: {m_target-m_enh}"
 
 
+TARGET_METRICS = {
+    "DeepFilterNet": {
+        "composite": [
+            2.30728650093078,
+            3.83064246177673,
+            2.36408281326293,
+            3.05453467369079,
+            -2.7631254196166,
+        ],
+        "stoi": 0.9689496585281197,
+        "sdr": 18.88543128967285,
+    },
+    "DeepFilterNet2a": {
+        "composite": [
+            2.86751246452332,
+            4.03339815139771,
+            2.56429362297058,
+            3.41470885276794,
+            -2.79574084281921,
+        ],
+        "stoi": 0.9707452525900906,
+        "sdr": 13.40160727500915,
+    },
+    "DeepFilterNet2b": {
+        "composite": [
+            2.87229919433594,
+            4.15724086761475,
+            2.62931561470032,
+            3.48965477943420,
+            -2.28056311607361,
+        ],
+        "stoi": 0.9733591821902137,
+        "sdr": 13.59861135482788,
+    },
+}
+
+
+def _load_model(df_dir: str, model_n: str):
+    model_base_dir = os.path.join(df_dir, "pretrained_models", model_n)
+    model, df_state, _ = init_df(model_base_dir, config_allow_defaults=True)
+    logger.info(f"Loaded model {model_n}")
+    return model, df_state
+
+
+class TestDfModels(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        torch.set_printoptions(precision=14, linewidth=120)
+        cls.df_dir = os.path.abspath(os.path.join(os.path.dirname(df.__file__), os.pardir))
+        cls.models = {m: _load_model(cls.df_dir, m) for m in TARGET_METRICS.keys()}
+        return cls
+
+    def _test_model(
+        self,
+        model: torch.nn.Module,
+        df_state: DF,
+        target_metrics: Dict[str, Union[float, List[float]]],
+    ):
+        sr = df_state.sr()
+        logger.info("Loading audios")
+        noisy, _ = load_audio(
+            os.path.join(self.df_dir, os.path.pardir, "assets", "noisy_snr0.wav"), sr
+        )
+        clean, _ = load_audio(
+            os.path.join(self.df_dir, os.path.pardir, "assets", "clean_freesound_33711.wav"), sr
+        )
+        enhanced = enhance(model, df_state, noisy, pad=True)
+        eval_composite(clean, enhanced, sr, target_metrics["composite"])  # type: ignore
+        eval_pystoi(clean, enhanced, sr, m_target=target_metrics["stoi"])  # type: ignore
+        eval_sdr(clean, enhanced, m_target=target_metrics["sdr"])  # type: ignore
+
+    def test_deepfilternet(self):
+        model = "DeepFilterNet"
+        self._test_model(*self.models[model], target_metrics=TARGET_METRICS[model])
+
+    def test_deepfilternet2a(self):
+        model = "DeepFilterNet2a"
+        self._test_model(*self.models[model], target_metrics=TARGET_METRICS[model])
+
+    def test_deepfilternet2(self):
+        model = "DeepFilterNet2b"
+        self._test_model(*self.models[model], target_metrics=TARGET_METRICS[model])
+
+
 if __name__ == "__main__":
-    df_dir = os.path.abspath(os.path.join(os.path.dirname(df.__file__), os.pardir))
-    model_base_dir = os.path.join(df_dir, "pretrained_models", "DeepFilterNet")
-    model, df_state, _ = init_df(model_base_dir)
-    sr = ModelParams().sr
-    logger.info("Loading audios")
-    noisy, _ = load_audio(os.path.join(df_dir, os.path.pardir, "assets", "noisy_snr0.wav"), sr)
-    clean, _ = load_audio(
-        os.path.join(df_dir, os.path.pardir, "assets", "clean_freesound_33711.wav"), sr
-    )
-    logger.info("Running model")
-    enhanced = enhance(model, df_state, noisy, pad=True)
-    try_eval_composite(clean, enhanced, sr)
-    eval_pystoi(clean, enhanced, sr)
-    eval_sdr(clean, enhanced)
+    unittest.main()
diff --git a/DeepFilterNet/df/scripts/test_dns_2020.py b/DeepFilterNet/df/scripts/test_dns_2020.py
new file mode 100644
index 000000000..f8a5368c5
--- /dev/null
+++ b/DeepFilterNet/df/scripts/test_dns_2020.py
@@ -0,0 +1,78 @@
+import glob
+import os
+
+from loguru import logger
+
+from df.enhance import init_df, save_audio, setup_df_argument_parser
+from df.evaluation_utils import evaluation_loop
+from df.model import ModelParams
+
+
+def main(args):
+    model, df_state, suffix = init_df(
+        args.model_base_dir,
+        post_filter=args.pf,
+        log_level=args.log_level,
+        log_file="dns2020-test.log",
+        config_allow_defaults=True,
+        epoch=args.epoch,
+    )
+    assert os.path.isdir(args.dataset_dir)
+    sr = ModelParams().sr
+    datasets = [os.path.join(args.dataset_dir, "no_reverb")]
+    if args.reverb:
+        datasets.append(os.path.join(args.dataset_dir, "with_reverb"))
+
+    save_audio_callback = None  # type: ignore
+    if args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+        def save_audio_callback(cleanfn: str, enh):
+            save_audio(
+                os.path.basename(cleanfn), enh, sr, output_dir=args.output_dir, suffix=suffix
+            )
+
+    for ds_dir in datasets:
+        logger.info(f"Evaluating dataset {os.path.basename(ds_dir)}")
+        assert os.path.isdir(ds_dir), ds_dir
+        noisy_dir = os.path.join(ds_dir, "noisy")
+        clean_dir = os.path.join(ds_dir, "clean")
+        assert os.path.isdir(noisy_dir) and os.path.isdir(clean_dir)
+        noisy_files = sorted(glob.iglob(noisy_dir + "/*.wav"), key=extract_fileid)
+        clean_files = sorted(glob.glob(clean_dir + "/*.wav"), key=extract_fileid)
+
+        metrics = evaluation_loop(
+            df_state,
+            model,
+            clean_files,
+            noisy_files,
+            n_workers=args.metric_workers,
+            save_audio_callback=save_audio_callback,
+        )
+        for k, v in metrics.items():
+            logger.info(f"{k}: {v}")
+
+
+def extract_fileid(fn: str) -> int:
+    return int(os.path.splitext(fn)[0].split("_")[-1])
+
+
+if __name__ == "__main__":
+    parser = setup_df_argument_parser()
+    parser.add_argument(
+        "dataset_dir",
+        type=str,
+        help="DNS 2020 Test set directory. Must contain 'no_reverb/clean' and 'no_reverb/noisy'.",
+    )
+    parser.add_argument(
+        "--metric-workers",
+        "-w",
+        type=int,
+        default=4,
+        help="Number of worker processes for metric calculation.",
+    )
+    parser.add_argument(
+        "--with-reverb", action="store_true", help="Also test on the reverb dataset", dest="reverb"
+    )
+    args = parser.parse_args()
+    main(args)
diff --git a/DeepFilterNet/df/scripts/test_voicebank_demand.py b/DeepFilterNet/df/scripts/test_voicebank_demand.py
index 634e296e3..0033c4e18 100644
--- a/DeepFilterNet/df/scripts/test_voicebank_demand.py
+++ b/DeepFilterNet/df/scripts/test_voicebank_demand.py
@@ -1,52 +1,11 @@
 import glob
 import os
-import sys
-import tempfile
 
-import numpy as np
-import pystoi
-import torch
 from loguru import logger
 
-from df.enhance import df_features, init_df, load_audio, save_audio, setup_df_argument_parser
+from df.enhance import init_df, save_audio, setup_df_argument_parser
+from df.evaluation_utils import evaluation_loop
 from df.model import ModelParams
-from df.modules import get_device
-from df.utils import as_complex, resample
-
-HAS_OCTAVE = True
-try:
-    import semetrics
-except OSError or ImportError:
-    HAS_OCTAVE = False
-
-try:
-    from tqdm import tqdm
-except ImportError:
-
-    def tqdm(iterable, desc="Progress", total=None, fallback_estimate=1000):
-        # tqdm not available using fallback
-        e = ""
-        try:
-            L = iterable.__len__()
-        except AttributeError:
-            L = total or None
-
-        print("{}:  {: >2d}".format(desc, 0), end="")
-        for k, i in enumerate(iterable):
-            yield i
-            if L is not None:
-                p = (k + 1) / L
-                e = "" if k < (L - 1) else "\n"
-            else:
-                # Use an exponentially decaying function
-                p = 1 - np.exp(-k / fallback_estimate)
-            print("\b\b\b\b {: >2d}%".format(int(100 * p)), end=e)
-            sys.stdout.flush()
-        if L is None:
-            print()
-
-
-__resample_method = "sinc_fast"
 
 
 def main(args):
@@ -54,131 +13,32 @@ def main(args):
         args.model_base_dir,
         post_filter=args.pf,
         log_level=args.log_level,
+        log_file="voicebank-test.log",
         config_allow_defaults=True,
+        epoch=args.epoch,
     )
     assert os.path.isdir(args.dataset_dir)
     sr = ModelParams().sr
     noisy_dir = os.path.join(args.dataset_dir, "noisy_testset_wav")
     clean_dir = os.path.join(args.dataset_dir, "clean_testset_wav")
-    enh_dir = None
     assert os.path.isdir(noisy_dir) and os.path.isdir(clean_dir)
-    enh_stoi = []
-    noisy_stoi = []
-    enh_sisdr = []
-    noisy_sisdr = []
-    enh_comp = []
-    noisy_comp = []
-    noisy_files = glob.glob(noisy_dir + "/*wav")
-    clean_files = glob.glob(clean_dir + "/*wav")
-    for noisyfn, cleanfn in tqdm(zip(noisy_files, clean_files), total=len(noisy_files)):
-        noisy, _ = load_audio(noisyfn, sr)
-        clean, _ = load_audio(cleanfn, sr)
-        enh = enhance(model, df_state, noisy)[0]
-        clean = df_state.synthesis(df_state.analysis(clean.numpy()))[0]
-        noisy = df_state.synthesis(df_state.analysis(noisy.numpy()))[0]
-        enh_stoi.append(stoi(clean, enh, sr))
-        noisy_stoi.append(stoi(clean, noisy, sr))
-        enh_sisdr.append(si_sdr_speechmetrics(clean, enh))
-        noisy_sisdr.append(si_sdr_speechmetrics(clean, noisy))
-        noisy_comp.append(composite(clean, noisy, sr))
-        enh_comp.append(composite(clean, enh, sr))
-        if args.log_level.upper() == "DEBUG":
-            print(cleanfn, enh_stoi[-1], enh_comp[-1], enh_sisdr[-1])
-        enh = torch.as_tensor(enh).to(torch.float32).view(1, -1)
-        if enh_dir is not None:
-            save_audio(
-                os.path.basename(cleanfn),
-                enh,
-                sr,
-                output_dir=enh_dir,
-                suffix=f"{suffix}_{enh_comp[-1][0]:.3f}",
-            )
-    logger.info(f"noisy stoi: {np.mean(noisy_stoi)}")
-    logger.info(f"enhanced stoi: {np.mean(enh_stoi)}")
-    noisy_comp = np.stack(noisy_comp)
-    enh_comp = np.stack(enh_comp)
-    noisy_pesq = np.mean(noisy_comp[:, 0])
-    noisy_csig = np.mean(noisy_comp[:, 1])
-    noisy_cbak = np.mean(noisy_comp[:, 2])
-    noisy_covl = np.mean(noisy_comp[:, 3])
-    noisy_ssnr = np.mean(noisy_comp[:, 4])
-    enh_pesq = np.mean(enh_comp[:, 0])
-    enh_csig = np.mean(enh_comp[:, 1])
-    enh_cbak = np.mean(enh_comp[:, 2])
-    enh_covl = np.mean(enh_comp[:, 3])
-    enh_ssnr = np.mean(enh_comp[:, 4])
-    logger.info(f"noisy pesq: {np.mean(noisy_pesq)}")
-    logger.info(f"enhanced pesq: {np.mean(enh_pesq)}")
-    logger.info(f"noisy csig: {np.mean(noisy_csig)}")
-    logger.info(f"enhanced csig: {np.mean(enh_csig)}")
-    logger.info(f"noisy cbak: {np.mean(noisy_cbak)}")
-    logger.info(f"enhanced cbak: {np.mean(enh_cbak)}")
-    logger.info(f"noisy covl: {np.mean(noisy_covl)}")
-    logger.info(f"enhanced covl: {np.mean(enh_covl)}")
-    logger.info(f"noisy ssnr: {np.mean(noisy_ssnr)}")
-    logger.info(f"enhanced ssnr: {np.mean(enh_ssnr)}")
-    logger.info(f"noisy sisdr: {np.mean(noisy_sisdr)}")
-    logger.info(f"enhanced sisdr: {np.mean(enh_sisdr)}")
-
-
-def stoi(clean, degraded, sr, extended=False):
-    assert len(clean.shape) == 1
-    if sr != 10000:
-        clean = resample(torch.as_tensor(clean), sr, 10000, method=__resample_method).numpy()
-        degraded = resample(torch.as_tensor(degraded), sr, 10000, method=__resample_method).numpy()
-        sr = 10000
-    return pystoi.stoi(clean, degraded, sr, extended=extended)
-
-
-def composite(clean: np.ndarray, degraded: np.ndarray, sr: int) -> np.ndarray:
-    """Compute pesq, csig, cbak, covl, ssnr"""
-    assert HAS_OCTAVE
-    assert len(clean.shape) == 1
-    if sr != 16000:
-        clean = resample(torch.as_tensor(clean), sr, 16000, method=__resample_method).numpy()
-        degraded = resample(torch.as_tensor(degraded), sr, 16000, method=__resample_method).numpy()
-        sr = 16000
-    cf = tempfile.NamedTemporaryFile(suffix=".wav")
-    save_audio(cf.name, clean, sr)
-    nf = tempfile.NamedTemporaryFile(suffix=".wav")
-    save_audio(nf.name, degraded, sr)
-    c = semetrics.composite(cf.name, nf.name)
-    cf.close()
-    nf.close()
-    return np.asarray(c)
-
-
-@torch.no_grad()
-def enhance(model, df_state, audio):
-    model.eval()
-    if hasattr(model, "reset_h0"):
-        model.reset_h0(batch_size=1, device=get_device())
-    spec, erb_feat, spec_feat = df_features(audio, df_state, get_device())
-    spec = model(spec, erb_feat, spec_feat)[0]
-    return df_state.synthesis(as_complex(spec.squeeze(0)).cpu().numpy())
-
-
-def si_sdr_speechmetrics(reference: np.ndarray, estimate: np.ndarray):
-    """This implementation is adopted from https://github.com/aliutkus/speechmetrics/blob/dde303e/speechmetrics/relative/sisdr.py"""
-    # as provided by @Jonathan-LeRoux and slightly adapted for the case of just one reference
-    # and one estimate.
-    # see original code here: https://github.com/sigsep/bsseval/issues/3#issuecomment-494995846
-    reference = reference.reshape(-1, 1)
-    estimate = estimate.reshape(-1, 1)
-    eps = np.finfo(reference.dtype).eps
-    Rss = np.dot(reference.T, reference)
-
-    # get the scaling factor for clean sources
-    a = (eps + np.dot(reference.T, estimate)) / (Rss + eps)
-
-    e_true = a * reference
-    e_res = estimate - e_true
-
-    Sss = (e_true**2).sum()
-    Snn = (e_res**2).sum()
-
-    sisdr = 10 * np.log10((eps + Sss) / (eps + Snn))
-    return sisdr
+    clean_files = glob.glob(clean_dir + "/*.wav")
+    noisy_files = glob.glob(noisy_dir + "/*.wav")
+
+    def save_audio_callback(cleanfn: str, enh):
+        save_audio(os.path.basename(cleanfn), enh, sr, output_dir=args.output_dir, suffix=suffix)
+
+    metrics = evaluation_loop(
+        df_state,
+        model,
+        clean_files,
+        noisy_files,
+        n_workers=args.metric_workers,
+        save_audio_callback=save_audio_callback if args.output_dir is not None else None,
+        metric_list=["stoi", "sisdr", "composite"],
+    )
+    for k, v in metrics.items():
+        logger.info(f"{k}: {v}")
 
 
 if __name__ == "__main__":
@@ -186,7 +46,14 @@ def si_sdr_speechmetrics(reference: np.ndarray, estimate: np.ndarray):
     parser.add_argument(
         "dataset_dir",
         type=str,
-        help="Voicebank Demand Test set directory. Must contain 'noisy_testset_wav' and 'clean_testset_wav'",
+        help="Voicebank Demand Test set directory. Must contain 'noisy_testset_wav' and 'clean_testset_wav'.",
+    )
+    parser.add_argument(
+        "--metric-workers",
+        "-w",
+        type=int,
+        default=4,
+        help="Number of worker processes for metric calculation.",
     )
     args = parser.parse_args()
     main(args)
diff --git a/DeepFilterNet/requirements_test.txt b/DeepFilterNet/requirements_test.txt
new file mode 100644
index 000000000..68fcc5943
--- /dev/null
+++ b/DeepFilterNet/requirements_test.txt
@@ -0,0 +1,3 @@
+pystoi
+pesq
+scipy
diff --git a/scripts/sbatch_test_voicebank.sh b/scripts/sbatch_test_voicebank.sh
new file mode 100644
index 000000000..51d1efdf3
--- /dev/null
+++ b/scripts/sbatch_test_voicebank.sh
@@ -0,0 +1,113 @@
+#!/bin/bash
+#SBATCH --job-name=df
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=1
+#SBATCH --mem=8000
+#SBATCH --gres=gpu:1
+#SBATCH -o /cluster/%u/logs/%j.out
+#SBATCH -e /cluster/%u/logs/%j.out
+#SBATCH --mail-type=ALL
+# Time limit format: "hours:minutes:seconds"
+#SBATCH --time=04:00:00
+
+set -e
+
+# Cluster directory containing data, project code, logging, summaries,
+# checkpoitns, etc.
+export CLUSTER=/net/cluster/$USER
+cd "$CLUSTER"
+
+# Workaround for our cluster
+export WORKON_HOME="/cluster/$USER/.cache"
+export XDG_CACHE_DIR="/cluster/$USER/.cache"
+export PYTHONUSERBASE="/cluster/$USER/.python_packages"
+
+PROJECT_NAME=DeepFilterNet
+DATA_DIR=${DATA_DIR:-$CLUSTER/Data/voicebank-demand-test}     # Set to the directory containing the HDF5s
+PYTORCH_JIT=${PYTORCH_JIT:-1}                # Set to 0 to disable pytorch JIT compilation
+
+echo "Started sbatch script at $(date) in $(pwd)"
+
+echo "Found cuda devices: $CUDA_VISIBLE_DEVICES"
+nvidia-smi -L || echo "nvidia-smi not found"
+echo "Running on host: $(hostname)"
+
+# Check base dir file
+if [[ -z $1 ]]; then
+  echo >&2 "No model base directory provided!"
+  exit
+fi
+if [[ ! -d $1 ]]; then
+  echo >&2 "Model base directory not found at $1!"
+  exit
+fi
+
+BASE_DIR=$(readlink -f "$1")
+
+echo "Got base_dir: $BASE_DIR"
+MODEL_NAME=$(basename "$BASE_DIR")
+
+# Git project setup.
+# Creates a separate code directory so that changes of the source code don't
+# have any impact on the automatic resubmission process. Furthermore, a specific
+# branch or commit can be specified that will be tried to checkout. By default,
+# the currently active branch is used.
+PROJECT_BRANCH=${BRANCH:-$PROJECT_BRANCH_CUR}
+if [[ -n $2 ]]; then
+  if [[ ! -d $2 ]]; then
+    echo >&2 "Project home not found at $2!"
+    exit
+  fi
+  PROJECT_HOME=$2
+else
+  PROJECT_ClUSTER_HOME=$CLUSTER/$PROJECT_NAME/
+  PROJECT_HOME=$CLUSTER/sbatch-$PROJECT_NAME/$MODEL_NAME/
+  mkdir -p "$PROJECT_HOME"
+  echo "Copying repo to $PROJECT_HOME"
+  cd "$PROJECT_ClUSTER_HOME"
+  rsync -avq --include .git \
+    --exclude-from="$(git -C "$PROJECT_ClUSTER_HOME" ls-files --exclude-standard -oi --directory >.git/ignores.tmp && echo .git/ignores.tmp)" \
+    "$PROJECT_ClUSTER_HOME" "$PROJECT_HOME" --delete
+fi
+if [ -n "$3" ]; then
+  # Checkout specified branch from previous job
+  PROJECT_BRANCH_CUR=$3
+else
+  # Use current branch of project on /cluster
+  PROJECT_BRANCH_CUR=$(git -C "$PROJECT_ClUSTER_HOME" rev-parse --abbrev-ref HEAD)
+fi
+echo "Running on branch $PROJECT_BRANCH in dir $PROJECT_HOME"
+if [ "$PROJECT_BRANCH_CUR" != "$PROJECT_BRANCH" ]; then
+  stash="stash_$SLURM_JOB_ID"
+  git -C "$PROJECT_HOME" stash save "$stash"
+  git -C "$PROJECT_HOME" checkout "$PROJECT_BRANCH"
+  stash_idx=$(git -C "$PROJECT_HOME" stash list | grep "$stash" | cut -d: -f1)
+  if [ -n "$stash_idx" ] && [ "$stash_idx" != " " ]; then
+    # Try to apply current stash; If not possible just proceed.
+    if ! git -C "$PROJECT_HOME" stash pop "$stash_idx"; then
+      echo "Could not apply stash to branch $PROJECT_BRANCH"
+      git -C "$PROJECT_HOME" checkout -f
+    fi
+  fi
+fi
+
+# Setup conda environment.
+# This installs miniconda environment if not existing, pytorch with cuda
+# integration and pip packages in requirements.txt
+. "$PROJECT_HOME"/scripts/setup_env.sh --source-only
+setup_env "$CLUSTER" "$PROJECT_HOME" "$MODEL_NAME"
+
+cd "$PROJECT_HOME"/DeepFilterNet/df/
+
+# Start training
+printf "\n***Starting training***\n\n"
+
+PYTHONPATH="$PROJECT_HOME/DeepFilterNet/" python scripts/test_voicebank_demand.py \
+  -m "$BASE_DIR" \
+  "$DATA_DIR"
+
+trainprocess=$!
+echo "Started trainprocess: $trainprocess"
+
+wait $trainprocess
+echo "Training stopped"