30 lines
853 B
Plaintext
30 lines
853 B
Plaintext
from sklearn.datasets import load_breast_cancer
|
|
from sklearn.tree import DecisionTreeClassifier
|
|
from sklearn.model_selection import train_test_split
|
|
from sklearn.metrics import confusion_matrix
|
|
from sklearn.tree import export_graphviz
|
|
|
|
from IPython.display import Image
|
|
from pydot import graph_from_dot_data
|
|
import pandas as pd
|
|
import numpy as np
|
|
|
|
cancer = load_breast_cancer()
|
|
X = pd.DataFrame(cancer.data, columns=cancer.feature_names)
|
|
y = pd.Categorical.from_codes(cancer.target, cancer.target_names)
|
|
y = pd.get_dummies(y)
|
|
|
|
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=1)
|
|
tree_clf = DecisionTreeClassifier(max_depth=5)
|
|
tree_clf.fit(X_train, y_train)
|
|
|
|
export_graphviz(
|
|
tree_clf,
|
|
out_file="cancer.dot",
|
|
feature_names=cancer.feature_names,
|
|
class_names=cancer.target_names,
|
|
rounded=True,
|
|
filled=True
|
|
)
|
|
|