wehub-resource-sync
77 lines
3.2 KiB
Python
77 lines
3.2 KiB
Python
from collections import defaultdict
|
|
|
|
import numpy as np
|
|
from torch.utils.data import Sampler, RandomSampler
|
|
|
|
|
|
class PKSampler(Sampler):
|
|
"""随机取一批数据，保证每个类别的数量都是相同的。
|
|
|
|
Args:
|
|
sampler (Dataset): 数据的Sampler
|
|
batch_size (int): batch size
|
|
sample_per_id (int): 每个类别的样本数量
|
|
shuffle (bool, optional): 是否随机打乱数据
|
|
drop_last (bool, optional): 是否丢掉最后一个batch
|
|
"""
|
|
|
|
def __init__(self,
|
|
sampler,
|
|
batch_size,
|
|
sample_per_id,
|
|
shuffle=True,
|
|
drop_last=True):
|
|
super().__init__(sampler)
|
|
self.sampler = sampler
|
|
self.batch_size = batch_size
|
|
self.shuffle = shuffle
|
|
self.drop_last = drop_last
|
|
assert batch_size % sample_per_id == 0, f"batch_size({batch_size})必须是sample_per_id({sample_per_id})的整数倍"
|
|
self.sample_per_id = sample_per_id
|
|
self.label_dict = defaultdict(list)
|
|
# 处理分布式和单卡，获取所有标签
|
|
if isinstance(sampler, RandomSampler):
|
|
labels = self.sampler.data_source.labels
|
|
else:
|
|
labels = sampler.dataset.labels
|
|
for idx, label in enumerate(labels):
|
|
self.label_dict[label].append(idx)
|
|
self.label_list = list(self.label_dict)
|
|
assert len(self.label_list) * self.sample_per_id >= self.batch_size, \
|
|
f"batch_size({self.batch_size})必须大于等于label_list({len(self.label_list)})*sample_per_id({self.sample_per_id})"
|
|
self.epoch = 0
|
|
self.prob_list = np.array([1 / len(self.label_list)] * len(self.label_list))
|
|
diff = np.abs(sum(self.prob_list) - 1)
|
|
if diff > 0.00000001:
|
|
self.prob_list[-1] = 1 - sum(self.prob_list[:-1])
|
|
|
|
def __len__(self):
|
|
if self.drop_last:
|
|
return len(self.sampler) // self.batch_size
|
|
else:
|
|
return (len(self.sampler) + self.batch_size - 1) // self.batch_size
|
|
|
|
def __iter__(self):
|
|
if self.shuffle:
|
|
np.random.RandomState(self.epoch).shuffle(self.label_list)
|
|
np.random.RandomState(self.epoch).shuffle(self.prob_list)
|
|
self.epoch += 1
|
|
|
|
label_per_batch = self.batch_size // self.sample_per_id
|
|
for _ in range(len(self)):
|
|
batch_index = []
|
|
# 从标签列表中随机选择指定数量的标签，概率根据概率列表进行
|
|
batch_label_list = np.random.choice(self.label_list, size=label_per_batch, replace=False, p=self.prob_list)
|
|
for label_i in batch_label_list:
|
|
label_i_indexes = self.label_dict[label_i]
|
|
# 从当前标签的索引列表中随机选择指定数量的样本，如果样本数量不足则允许重复选择
|
|
batch_index.extend(
|
|
np.random.choice(label_i_indexes, size=self.sample_per_id,
|
|
replace=not self.sample_per_id <= len(label_i_indexes)))
|
|
# 再次随机打乱
|
|
if self.shuffle:
|
|
np.random.shuffle(batch_index)
|
|
if not self.drop_last or len(batch_index) == self.batch_size:
|
|
yield batch_index
|
|
self.epoch += 1
|