51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
"""
|
|
Example: Creating a Custom Dataset
|
|
|
|
This example shows how to add a new dataset following the project architecture.
|
|
"""
|
|
|
|
from torch.utils.data import Dataset
|
|
from typing import Dict
|
|
import torch
|
|
from src.data_module.dataset import register_dataset
|
|
|
|
|
|
@register_dataset("time_series")
|
|
class TimeSeriesDataset(Dataset):
|
|
"""
|
|
Time series dataset for sequence modeling.
|
|
|
|
Args:
|
|
sequences: List of time series sequences
|
|
seq_length: Fixed sequence length (pad or truncate if needed)
|
|
"""
|
|
|
|
def __init__(self, sequences: list, seq_length: int = 100):
|
|
self.sequences = sequences
|
|
self.seq_length = seq_length
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.sequences)
|
|
|
|
def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
|
|
sequence = self.sequences[i]
|
|
|
|
# Pad or truncate to fixed length
|
|
if len(sequence) < self.seq_length:
|
|
padding = torch.zeros(self.seq_length - len(sequence))
|
|
sequence = torch.cat([sequence, padding])
|
|
else:
|
|
sequence = sequence[:self.seq_length]
|
|
|
|
return {
|
|
"input": sequence,
|
|
"label": sequence, # For autoencoder, etc.
|
|
"length": torch.tensor(min(len(self.sequences[i]), self.seq_length))
|
|
}
|
|
|
|
|
|
# Usage in training:
|
|
# from src.data_module.dataset import DatasetFactory
|
|
# dataset = DatasetFactory("time_series")(sequences=training_data, seq_length=128)
|
|
# dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
|