Spaces:
Sleeping
Sleeping
File size: 2,211 Bytes
0788e19 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import torch
import numpy as np
from torch.utils.data.sampler import WeightedRandomSampler
from .datasets import dataset_folder
import os
# 获取数据集，要根据不同的需求进行修改
'''
def get_dataset(opt):
dset_lst = []
for cls in opt.classes:
root = opt.dataroot + '/' + cls
dset = dataset_folder(opt, root)
dset_lst.append(dset)
return torch.utils.data.ConcatDataset(dset_lst)
'''
def get_dataset(opt):
classes = os.listdir(opt.dataroot) if len(opt.classes) == 0 else opt.classes
# 如果类文件夹中没有 '0_real' 和 '1_fake'，会对每个类单独加载数据集，并将它们合并返回
# 如果类文件夹中包含 '0_real' 和 '1_fake'，则直接对整个根目录加载数据集
if '0_real' not in classes or '1_fake' not in classes:
dset_lst = []
for cls in classes:
root = opt.dataroot + '/' + cls
dset = dataset_folder(opt, root)
dset_lst.append(dset)
return torch.utils.data.ConcatDataset(dset_lst)
return dataset_folder(opt, opt.dataroot)
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):
# 是否打乱数据集
shuffle = not opt.serial_batches if (opt.isTrain and not opt.class_bal) else False
dataset = get_dataset(opt)
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),
drop_last=opt.drop_last)
return data_loader
|