(config)
| 60 | |
| 61 | |
| 62 | def prepare_training_data(config): |
| 63 | # Only use the blending dataset class in training |
| 64 | if 'dataset_type' in config and config['dataset_type'] == 'blend': |
| 65 | if config['model_name'] == 'sbi': |
| 66 | train_set = SBIDataset(config, mode='train') |
| 67 | else: |
| 68 | raise NotImplementedError( |
| 69 | 'Only sbi is currently supported for blending dataset' |
| 70 | ) |
| 71 | elif 'dataset_type' in config and config['dataset_type'] == 'pair': |
| 72 | train_set = pairDataset(config, mode='train') # Only use the pair dataset class in training |
| 73 | else: |
| 74 | train_set = DeepfakeAbstractBaseDataset( |
| 75 | config=config, |
| 76 | mode='train', |
| 77 | ) |
| 78 | |
| 79 | if config['ddp']: |
| 80 | sampler = DistributedSampler(train_set) |
| 81 | train_data_loader = \ |
| 82 | torch.utils.data.DataLoader( |
| 83 | dataset=train_set, |
| 84 | batch_size=config['train_batchSize'], |
| 85 | num_workers=int(config['workers']), |
| 86 | collate_fn=train_set.collate_fn, |
| 87 | sampler=sampler |
| 88 | ) |
| 89 | else: |
| 90 | train_data_loader = \ |
| 91 | torch.utils.data.DataLoader( |
| 92 | dataset=train_set, |
| 93 | batch_size=config['train_batchSize'], |
| 94 | shuffle=True, |
| 95 | num_workers=int(config['workers']), |
| 96 | collate_fn=train_set.collate_fn, |
| 97 | ) |
| 98 | return train_data_loader |
| 99 | |
| 100 | |
| 101 | def prepare_testing_data(config): |
no test coverage detected