Spaces:
Running
Running
| # modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/data_module.py | |
| from pytorch_lightning import LightningDataModule | |
| from AR.data.bucket_sampler import DistributedBucketSampler | |
| from AR.data.dataset import Text2SemanticDataset | |
| from torch.utils.data import DataLoader | |
| class Text2SemanticDataModule(LightningDataModule): | |
| def __init__(self, config, train_semantic_path, train_phoneme_path,dev_semantic_path=None, dev_phoneme_path=None): | |
| super().__init__() | |
| self.config = config | |
| self.train_semantic_path = train_semantic_path | |
| self.train_phoneme_path = train_phoneme_path | |
| self.dev_semantic_path = dev_semantic_path | |
| self.dev_phoneme_path = dev_phoneme_path | |
| self.num_workers = self.config['data']['num_workers'] | |
| def prepare_data(self): | |
| pass | |
| def setup(self, stage=None, output_logs=False): | |
| self._train_dataset = Text2SemanticDataset( | |
| phoneme_path=self.train_phoneme_path, | |
| semantic_path=self.train_semantic_path, | |
| max_sec=self.config['data']['max_sec'], | |
| pad_val=self.config['data']['pad_val']) | |
| self._dev_dataset = self._train_dataset | |
| # self._dev_dataset = Text2SemanticDataset( | |
| # phoneme_path=self.dev_phoneme_path, | |
| # semantic_path=self.dev_semantic_path, | |
| # max_sample=self.config['data']['max_eval_sample'], | |
| # max_sec=self.config['data']['max_sec'], | |
| # pad_val=self.config['data']['pad_val']) | |
| def train_dataloader(self): | |
| batch_size = self.config['train']['batch_size'] | |
| sampler = DistributedBucketSampler( | |
| self._train_dataset, batch_size=batch_size) | |
| return DataLoader( | |
| self._train_dataset, | |
| batch_size=batch_size, | |
| sampler=sampler, | |
| collate_fn=self._train_dataset.collate, | |
| num_workers=self.num_workers, | |
| persistent_workers=True, | |
| prefetch_factor=16 | |
| ) | |
| def val_dataloader(self): | |
| return DataLoader( | |
| self._dev_dataset, | |
| batch_size=1, | |
| shuffle=False, | |
| collate_fn=self._train_dataset.collate, | |
| num_workers=max(self.num_workers,12), | |
| persistent_workers=True, | |
| prefetch_factor=16 | |
| ) | |
| # 这个会使用到嘛？ | |
| def test_dataloader(self): | |
| return DataLoader( | |
| self._dev_dataset, | |
| batch_size=1, | |
| shuffle=False, | |
| collate_fn=self._train_dataset.collate) | |