45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.colors import ListedColormap
|
|
from sklearn.model_selection import train_test_split
|
|
from sklearn.ensemble import RandomForestClassifier
|
|
from sklearn import datasets
|
|
|
|
# 载入红酒数据
|
|
wine = datasets.load_wine()
|
|
|
|
# 只选取前两个特征
|
|
x = wine.data[:, :2]
|
|
y = wine.target
|
|
|
|
# 拆分训练集和数据集
|
|
x_train, x_test, y_train, y_test = train_test_split(x, y)
|
|
|
|
forest = RandomForestClassifier(n_estimators=6, random_state=3)
|
|
|
|
# 拟合数据
|
|
forest.fit(x_train, y_train)
|
|
|
|
# 绘制图形
|
|
cmap_light = ListedColormap(['#FFAAAA', '#AAFFAA', '#AAAAFF'])
|
|
cmap_bold = ListedColormap(['#FF0000', '#00FF00', '#0000FF'])
|
|
|
|
x_min, x_max = x[:, 0].min() - 1, x[:, 0].max() + 1
|
|
y_min, y_max = x[:, 1].min() - 1, x[:, 1].max() + 1
|
|
|
|
xx, yy = np.meshgrid(np.arange(x_min, x_max, .02), np.arange(y_min, y_max, .02))
|
|
|
|
# 生成随机森林
|
|
z = forest.predict(np.c_[xx.ravel(), yy.ravel()]).reshape(xx.shape)
|
|
|
|
plt.figure()
|
|
plt.pcolormesh(xx, yy, z, cmap=cmap_light)
|
|
|
|
# 绘制训练数据和测试数据的散点图
|
|
plt.scatter(x_train[:, 0], x_train[:, 1], c=y_train, cmap=cmap_bold, edgecolor='k', s=20)
|
|
plt.scatter(x_test[:, 0], x_test[:, 1], c=y_test, cmap=cmap_bold, edgecolor='k', s=50, marker='x')
|
|
|
|
plt.xlim(xx.min(), xx.max())
|
|
plt.ylim(yy.min(), yy.max())
|
|
plt.title("Random Forest Classification")
|
|
plt.show() |