-
Notifications
You must be signed in to change notification settings - Fork 1
/
train_decisiontree.py
75 lines (61 loc) · 2.7 KB
/
train_decisiontree.py
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
63
64
65
66
67
68
69
70
71
72
73
74
75
# this script builds a simple decision tree classifier
import os
import pandas as pd
import neptune
import pickle
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import f1_score, accuracy_score, average_precision_score
from sklearn.tree import DecisionTreeClassifier
from sklearn.pipeline import Pipeline
from data import pecarn
if __name__ == "__main__":
# bring the pecarn data in
df = pecarn.clean(pecarn.load(fromCsv=True))
X = df.drop(columns='PosIntFinal')
y = pecarn.preprocess(df[['PosIntFinal']])
# configure training and test datasets
X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.75, test_size=0.25, stratify=y, random_state=1234)
# create the pipeline/classifier
clf = DecisionTreeClassifier()
pipeline = Pipeline(steps=[
('data.pecarn.preprocess', pecarn.make_preprocess_pipeline()),
('decisiontreeclassifier', clf)
])
# save the feature names for use later in prediction
pipeline.input_features = X[0:0]
# neptune initialization - NEPTUNE_API_TOKEN and NEPTUNE_PROJECT environment variables must be set
neptune.init()
# create a neptune experiment to log to
with neptune.create_experiment(name='sklearn.DecisionTreeClassifier',
params=clf.get_params(),
upload_source_files=[__file__,'src/data/pecarn/*.py'],
send_hardware_metrics=False) as exp:
# train the classifier
pipeline.fit(X_train, y_train)
# calculate scores on train set
y_train_pred = pipeline.predict(X_train)
train_scores = {
'accuracy': accuracy_score(y_train, y_train_pred),
'f1': f1_score(y_train, y_train_pred),
'f1_weighted': f1_score(y_train, y_train_pred, average='weighted'),
'avg_precision': average_precision_score(y_train, y_train_pred)
}
# calculate scores on test set
y_pred = pipeline.predict(X_test)
test_scores = {
'accuracy': accuracy_score(y_test, y_pred),
'f1': f1_score(y_test, y_pred),
'f1_weighted': f1_score(y_test, y_pred, average='weighted'),
'avg_precision': average_precision_score(y_test, y_pred)
}
# log the scores
for metric_name, score in train_scores.items():
exp.send_metric('train_' + metric_name, score)
for metric_name, score in test_scores.items():
exp.send_metric('test_' + metric_name, score)
# log the model
pickle_file = pickle_file = exp.name + '.pkl'
with open(pickle_file, "wb") as f:
pickle.dump(pipeline, f)
exp.log_artifact(pickle_file)
os.remove(pickle_file)