How to train a classification model with scikit-learn
Training a classification model is the heart of Machine Learning: it is the step where the algorithm learns from examples so it can later predict the category of new data. With the scikit-learn library, in Python, you can train a first model in just a few lines. This guide walks the full path, from loading the data to your first prediction and its evaluation.
Prerequisites
- Python 3 installed on your computer.
- The scikit-learn library installed. If you do not have it yet, run
pip install scikit-learnin the terminal. - Basic Python knowledge: variables and how to run a script.
- You do not need to download anything: we use a dataset that already ships with scikit-learn.
Step 1: Load the data
For this example we use the Iris dataset, a classic for learning classification. It contains the measurements of 150 flowers split into three species, and the model's goal will be to predict the species from those measurements.
from sklearn.datasets import load_iris
iris = load_iris()
X = iris.data # as medidas (4 colunas por flor)
y = iris.target # a especie: 0, 1 ou 2
print(X.shape) # (150, 4)
print(iris.target_names) # ['setosa' 'versicolor' 'virginica']
We store the features (the measurements) in X and the correct answer (the species) in y. In Machine Learning it is a convention to call the input data X and what we want to predict y.
Step 2: Split into training and test sets
Before training, we separate the data into two parts: one for the model to learn from (training) and another to check whether it learned well (test). This lets us measure performance on data the model has never seen, which is essential to trust the result.
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
We set aside 20% of the data for testing with test_size=0.2. The random_state=42 makes the split always the same when you run the code again, and stratify=y keeps the same proportion of species in each part.
Step 3: Choose and train the model
We will use logistic regression, a simple and very popular model for classification. Training means calling the fit() method, which fits the model to the training data, that is, it finds the pattern linking the measurements to the species.
from sklearn.linear_model import LogisticRegression
modelo = LogisticRegression(max_iter=1000)
modelo.fit(X_train, y_train)
That is all: one line creates the model and another trains it. We use max_iter=1000 to give the algorithm enough iterations to converge and to avoid a convergence warning. After fit(), the model has already learned and is ready to predict.
Step 4: Make predictions
With the model trained, we ask it to predict the species of the flowers in the test set, using the predict() method.
previsoes = modelo.predict(X_test)
print(previsoes[:5]) # ex.: [1 0 2 1 1]
Each number is the species the model predicted for a flower. By comparing these predictions with the real answers (y_test), we can tell whether the model is good.
Check the result
The most direct way to evaluate is to measure accuracy: the percentage of correct predictions on the test set.
from sklearn.metrics import accuracy_score
exatidao = accuracy_score(y_test, previsoes)
print(f"Exatidao: {exatidao:.2f}")
In this example you should get a value around 0.97, meaning the model gets about 97% of the test flowers right (the exact value may vary slightly). A high number close to 1 indicates that your first model is working.
Tip: accuracy alone is not always enough. When the classes are imbalanced, it is worth also looking at the confusion matrix and at metrics such as precision and recall.
Conclusion
You have just trained a classification model end to end: you loaded the data, split it, trained with fit() and evaluated the predictions. From here you can try other models, such as DecisionTreeClassifier or RandomForestClassifier, by changing just one line of code, or deepen the evaluation with a confusion matrix. Here is a challenge: can you improve the accuracy using a different model on the same data?