我想以保持类比率的方式拆分数据集。 假设以下内容：
Total Image: 500,Class A 300; Class B 200
----------
Train: 300 (Class A 180,Class B 120)
----------
Valid: 100 (Class A 60,Class B 40)
----------
Test:100 (Class A 60,Class B 40)
在我的总列车中，保持有效和测试的比率，但不保留班级比率
我想以在Pytorch中维护类比率的方式拆分数据集
•
问答
zyf12345678 回答：我想以在Pytorch中维护类比率的方式拆分数据集
您可能正在寻找sklearn的test_train_split，它可以通过stratify参数来保持班级比例。
在您的情况下，您将必须拆分两次，一次拆分出测试数据，然后再次拆分验证和训练。
from sklearn.model_selection import train_test_split
# Create a dataset with your requirements
X = np.random.randn(500)
a = [0] * 300
b = [1] * 200
y = a + b
# Take 100 elements for the test set
X_train,X_test,y_train,y_test = train_test_split(X,y,test_size=0.2,stratify=y)
# Take 100 elements (from the remaining 400) for the validation set
X_train,X_val,y_val = train_test_split(X_train,test_size=0.25,stratify=y_train)
print(len(y_train)) # 300
print(len(y_val)) # 100
print(len(y_test)) # 100
print(sum(y_train)) # 120 (120 class B,180 class A)
print(sum(y_val)) # 40 (40 class B,60 class A)
print(sum(y_test)) # 40 (40 class B,60 class A)