zenon/netsec-lab
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Stratify test split (+0.25 points)

Browse Source
This commit is contained in:
Tobias Eidelpes 2021-06-05 18:51:24 +02:00
parent 5e1ae3289b
commit ee039552d2
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
competition/random_forest.py
View File

@ -20,7 +20,7 @@ df.drop(['flowStartMilliseconds'], 1, inplace=True)
X = np.array(df.drop(columns=['sublabel']))
y = np.array(df['sublabel'])
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y)
clf = RandomForestClassifier(n_estimators=50, n_jobs=-1, criterion='gini', random_state=0)
clf.fit(X_train, y_train)