flamo 0.2.0__tar.gz → 0.2.1__tar.gz

{flamo-0.2.0 → flamo-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flamo
-Version: 0.2.0
+Version: 0.2.1
 Summary: An Open-Source Library for Frequency-Domain Differentiable Audio Processing
 Project-URL: Homepage, https://github.com/gdalsanto/flamo
 Project-URL: Issues, https://github.com/gdalsanto/flamo/issues

{flamo-0.2.0 → flamo-0.2.1}/examples/e4_recursion_nn.py

@@ -84,6 +84,7 @@ class LossyBiquad(dsp.Biquad):
         fs: int = 48000,
         requires_grad: bool = False,
         alias_decay_db: float = 0.0,
+        dtype: torch.dtype = torch.float32,
         device: Optional[str] = None,
     ):
         super().__init__(
@@ -95,6 +96,7 @@ class LossyBiquad(dsp.Biquad):
             requires_grad=requires_grad,
             alias_decay_db=alias_decay_db,
             device=device,
+            dtype=dtype,
         )
 
     def get_map(self):
@@ -193,6 +195,7 @@ class nnComb(nn.Module):
             requires_grad=True,
             alias_decay_db=args.alias_decay_db,
             device=args.device,
+            dtype=args.dtype,
         )
 
         delays = dsp.parallelDelay(
@@ -203,6 +206,7 @@ class nnComb(nn.Module):
             requires_grad=False,
             alias_decay_db=args.alias_decay_db,
             device=args.device,
+            dtype=args.dtype,
         )
         delays.assign_value(delays.sample2s(delay_lengths))
         # To condition the biquad filter, we will need to pass the parameters estimated

{flamo-0.2.0 → flamo-0.2.1}/flamo/optimize/dataset.py

@@ -20,7 +20,7 @@ class Dataset(torch.utils.data.Dataset):
             - **input** (torch.Tensor, optional): The input data tensor. Default: torch.randn(100, 100).
             - **target** (torch.Tensor, optional): The target data tensor. Default: torch.randn(100, 100).
             - **expand** (int): The first shape dimention of the input and target tensor after expansion. Default: 1. This coincides with the length on the dataset.
-            - **device** (str, optional): The device to store the tensors on. Default: 'cpu'.
+            - **device** (torch.device, optional): The device to store the tensors on. Defaults to torch default device.
     """
 
     def __init__(
@@ -28,7 +28,7 @@ class Dataset(torch.utils.data.Dataset):
         input: torch.Tensor = torch.randn(1, 1),
         target: torch.Tensor = torch.randn(1, 1),
         expand: int = 1,
-        device: str = "cpu",
+        device: torch.device = torch.get_default_device(),
         dtype: torch.dtype = torch.float32,
     ):
         self.input = input.to(device).to(dtype)
@@ -59,7 +59,7 @@ class DatasetColorless(Dataset):
             - **input_shape** (tuple): The shape of the input data.
             - **target_shape** (tuple): The shape of the target data.
             - **expand** (int, optional): The number of times to expand the dataset. Defaults to 1000.
-            - **device** (str, optional): The device to use for computation. Defaults to 'cpu'.
+            - **device** (torch.device, optional): The device to use for computation. Defaults to torch default device.
 
     For details on the colorless optimization, see `<https://arxiv.org/abs/2402.11216v2>`_.
     """
@@ -69,7 +69,7 @@ class DatasetColorless(Dataset):
         input_shape: tuple,
         target_shape: tuple,
         expand: int = 1000,
-        device: str = "cpu",
+        device: torch.device = torch.get_default_device(),
         dtype: torch.dtype = torch.float32,
     ):
         input = torch.zeros(input_shape, device=device, dtype=dtype)
@@ -104,13 +104,18 @@ def get_dataloader(
     return dataloader
 
 
-def split_dataset(dataset: torch.utils.data.Dataset, split: float):
+def split_dataset(
+    dataset: torch.utils.data.Dataset,
+    split: float,
+    device: torch.device = torch.get_default_device()
+):
     r"""
     Splits the given dataset into a training set and a validation set based on the specified split ratio.
 
         **Arguments**:
             - **dataset** (torch.utils.data.Dataset): The dataset to be split.
             - **split** (float): The ratio of the training set size to the total dataset size.
+            - **device** (torch.device, optional): The device to use for the random split generator. Defaults to torch default device.
 
         **Returns**:
             tuple: A tuple containing the training set and the validation set.
@@ -119,7 +124,11 @@ def split_dataset(dataset: torch.utils.data.Dataset, split: float):
     train_set_size = int(len(dataset) * split)
     valid_set_size = len(dataset) - train_set_size
 
-    train_set, valid_set = data.random_split(dataset, [train_set_size, valid_set_size])
+    train_set, valid_set = data.random_split(
+        dataset,
+        [train_set_size, valid_set_size],
+        generator=torch.Generator(device=device)
+    )
 
     return train_set, valid_set
 
@@ -129,6 +138,7 @@ def load_dataset(
     batch_size: int = 2000,
     split: float = 0.8,
     shuffle: bool = True,
+    device: torch.device = torch.get_default_device()
 ):
     r"""
     Load and split a dataset into training and validation sets.
@@ -138,12 +148,13 @@ def load_dataset(
             - **batch_size** (int, optional): The batch size for the data loaders. Defaults to 2000.
             - **split** (float, optional): The ratio to split the dataset into training and validation sets. Defaults to 0.8.
             - **shuffle** (bool, optional): Whether to shuffle the data. Defaults to True.
+            - **device** (torch.device, optional): The device to use for the random split generator. Defaults to torch default device.
 
         **Returns**:
             tuple: A tuple containing the training and validation data loaders.
     """
 
-    train_set, valid_set = split_dataset(dataset, split)
+    train_set, valid_set = split_dataset(dataset, split, device)
 
     train_loader = get_dataloader(
         train_set,

{flamo-0.2.0 → flamo-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "flamo"
-version = "0.2.0"
+version = "0.2.1"
 authors = [
   { name="Gloria Dal Santo", email="gloria.dalsanto@aalto.fi" },
   { name="Gian Marco De Bortoli", email="gian.debortoli@aalto.fi"},