Here's a solution which is:
- simple (< 10 lines code)
- fast (besides one
for loop, pure NumPy)
- no external dependencies other than NumPy
- is very cheap to generate new balanced random samples (just call
np.random.sample()). Useful for generating different shuffled & balanced samples between training epochs
def stratified_random_sample_weights(labels):
sample_weights = np.zeros(num_samples)
for class_i in range(n_classes):
class_indices = np.where(labels[:, class_i]==1) # find indices where class_i is 1
class_indices = np.squeeze(class_indices) # get rid of extra dim
num_samples_class_i = len(class_indices)
assert num_samples_class_i > 0, f"No samples found for class index {class_i}"
sample_weights[class_indices] = 1.0/num_samples_class_i # note: samples with no classes present will get weight=0
return sample_weights / sample_weights.sum() # sum(weights) == 1
Then, you use re-use these weights over and over to generate balanced indices with np.random.sample():
sample_weights = stratified_random_sample_weights(labels)
chosen_indices = np.random.choice(list(range(num_samples)), size=sample_size, replace=True, p=sample_weights)
Full example:
# generate data
from sklearn.preprocessing import OneHotEncoder
num_samples = 10000
n_classes = 10
ground_truth_class_weights = np.logspace(1,3,num=n_classes,base=10,dtype=float) # exponentially growing
ground_truth_class_weights /= ground_truth_class_weights.sum() # sum to 1
labels = np.random.choice(list(range(n_classes)), size=num_samples, p=ground_truth_class_weights)
labels = labels.reshape(-1, 1) # turn each element into a list
labels = OneHotEncoder(sparse=False).fit_transform(labels)
print(f"original counts: {labels.sum(0)}")
# [ 38. 76. 127. 191. 282. 556. 865. 1475. 2357. 4033.]
sample_weights = stratified_random_sample_weights(labels)
sample_size = 1000
chosen_indices = np.random.choice(list(range(num_samples)), size=sample_size, replace=True, p=sample_weights)
print(f"rebalanced counts: {labels[chosen_indices].sum(0)}")
# [104. 107. 88. 107. 94. 118. 92. 99. 100. 91.]
cross_validationand specificallyK-Fold- EdChum