| import torch | |
| import numpy as np | |
| from torch.utils.data.sampler import WeightedRandomSampler | |
| from .datasets import RealFakeDataset | |
| def get_bal_sampler(dataset): | |
| targets = [] | |
| for d in dataset.datasets: | |
| targets.extend(d.targets) | |
| ratio = np.bincount(targets) | |
| w = 1. / torch.tensor(ratio, dtype=torch.float) | |
| sample_weights = w[targets] | |
| sampler = WeightedRandomSampler(weights=sample_weights, | |
| num_samples=len(sample_weights)) | |
| return sampler | |
| def create_dataloader(opt, preprocess=None): | |
| shuffle = not opt.serial_batches if (opt.isTrain and not opt.class_bal) else False | |
| dataset = RealFakeDataset(opt) | |
| print(len(dataset)) | |
| if '2b' in opt.arch: | |
| dataset.transform = preprocess | |
| sampler = get_bal_sampler(dataset) if opt.class_bal else None | |
| data_loader = torch.utils.data.DataLoader(dataset, | |
| batch_size=opt.batch_size, | |
| shuffle=shuffle, | |
| sampler=sampler, | |
| num_workers=int(opt.num_threads)) | |
| return data_loader | |