diff --git a/src/syncnet/config.py b/src/syncnet/config.py index 895f095..5f921ec 100644 --- a/src/syncnet/config.py +++ b/src/syncnet/config.py @@ -18,7 +18,7 @@ class Config(BaseModel): Attributes: seed: Random seed for reproducibility across runs. - test_split: Fraction of the dataset to use for validation/testing. + val_split: Fraction of the dataset to use for validation/testing. batch_size: Number of samples per batch during training. max_epochs: Maximum number of training epochs. early_stopping_patience: Number of epochs to wait before early stopping. @@ -43,8 +43,8 @@ class Config(BaseModel): seed: int = Field(default=42, description="Random seed for reproducibility.") # Data - test_split: float = Field( - default=0.05, description="Proportion of data for val/testing." + val_split: float = Field( + default=0.05, description="Proportion of data for validation." ) batch_size: int = Field(default=4, description="Batch size.") diff --git a/src/syncnet/datamodule.py b/src/syncnet/datamodule.py index 1c0d276..286d541 100644 --- a/src/syncnet/datamodule.py +++ b/src/syncnet/datamodule.py @@ -62,7 +62,10 @@ def __init__(self, dataset: Dataset, config: Config, num_workers: int = 4) -> No self.config = config self.num_workers = num_workers self.processor = PeAudioVideoProcessor.from_pretrained(config.base_model) - self.resample = torchaudio.transforms.Resample(16000, 48000) + self.resample = torchaudio.transforms.Resample( + 16000, + self.processor.feature_extractor.sampling_rate, # type: ignore[attr-defined] + ) self.resize = torchvision.transforms.Resize( (config.frame_height, config.frame_width) ) @@ -79,7 +82,7 @@ def setup(self, stage: str | None = None) -> None: """ if stage == "fit" or stage is None: dataset_size = len(self.dataset) # type: ignore - train_size = int(0.8 * dataset_size) + train_size = int((1.0 - self.config.val_split) * dataset_size) val_size = dataset_size - train_size self.train_dataset, self.val_dataset = torch.utils.data.random_split( self.dataset, [train_size, val_size] @@ -156,13 +159,13 @@ def pad_collate_fn( audio = audio.mean(dim=0, keepdim=True) # Convert to mono if stereo # Randomly create negative samples by mixing audio/video - mix = random.random() > 0.5 + mix = random.random() > self.config.negative_fraction video_segment, audio_segment = self.sample_random_segment( video, audio, num_frames=self.config.num_frames, - fps=25, - sample_rate=16000, + fps=metadata.get("video_fps", 25), + sample_rate=metadata.get("audio_fps", 16000), mix=mix, ) if audio_segment.shape[1] < 3200: @@ -177,7 +180,7 @@ def pad_collate_fn( audio=audio_segment.squeeze(0), return_tensors="pt", padding=False, - sampling_rate=48000, + sampling_rate=self.processor.feature_extractor.sampling_rate, # type: ignore[attr-defined] ) video_segments.append(input_values["pixel_values_videos"][0]) audio_segments.append(input_values["input_values"][0])