图像分类
这是一个在 MNIST 手写数字数据集上训练图像分类模型的完整示例。
这些交互式笔记本遵循本示例的步骤
下载 MNIST 数据集¶
MNIST 是一个手写数字灰度图像的集合。这个集合包含 60,000 张图像用于训练,10,000 张图像用于测试模型性能。每张图像都是 28 X 28 像素的灰度图像。
ludwig datasets download mnist
这个命令会在当前目录创建一个数据集文件 mnist_dataset.csv
。此外,还会创建包含图像的 training/
和 testing/
目录。
数据集中的列如下:
列 | 描述 |
---|---|
image_path | 图像的文件路径字符串 |
label | 0 到 9 的单个数字,表示图像中显示的数字 |
split | 整数值,表示训练样本 (0) 或测试样本 (2) |
from ludwig.datasets import mnist
# Loads the dataset as a pandas.DataFrame
train_df, test_df, _ = mnist.load(split=True)
train_df
包含 60K 训练图像的文件路径信息。test_df
包含 10K 测试图像的相同信息。
列 | 描述 |
---|---|
image_path | 图像的文件路径字符串 |
label | 0 到 9 的单个数字,表示图像中显示的数字 |
带有 label
的图像样本。
训练¶
Ludwig 配置文件描述了机器学习任务。本示例仅使用了 Ludwig 提供选项的一小部分。有关所有详细信息,请参阅配置部分。
首先定义 input_features
。对于图像特征,配置指定了用于编码图像的神经网络架构类型。本示例中,编码器是一个两层的堆叠卷积神经网络 (Stacked Convolutional Neural Network),后接一个带有 dropout 正则化的全连接层。
接下来定义 output_features
。在这种情况下,只有一个输出特征,称为 label
。这是一个类别特征 (categorical feature),表示图像代表的数字,即 0, 1, 2, ..., 9。
配置文件的最后一部分描述了 trainer
如何运行的选项。在本示例中,trainer
将处理训练数据进行 5 个 epoch。
# config.yaml
input_features:
- name: image_path
type: image
encoder:
type: stacked_cnn
conv_layers:
- num_filters: 32
filter_size: 3
pool_size: 2
pool_stride: 2
- num_filters: 64
filter_size: 3
pool_size: 2
pool_stride: 2
dropout: 0.4
fc_layers:
- output_size: 128
dropout: 0.4
output_features:
- name: label
type: category
trainer:
epochs: 5
# create Ludwig configuration dictionary
config = {
'input_features': [
{
'name': 'image_path',
'type': 'image',
'preprocessing': {'num_processes': 4},
'encoder': {
'stacked_cnn',
'conv_layers': [
{'num_filters': 32, 'filter_size': 3, 'pool_size': 2, 'pool_stride': 2},
{'num_filters': 64, 'filter_size': 3, 'pool_size': 2, 'pool_stride': 2, 'dropout': 0.4}
],
'fc_layers': [{'output_size': 128, 'dropout': 0.4}]
}
}
],
'output_features': [{'name': 'label', 'type': 'category'}],
'trainer': {'epochs': 5}
}
# Constructs Ludwig model from config dictionary
model = LudwigModel(config, logging_level=logging.INFO)
训练模型。
ludwig train \
--dataset mnist_dataset.csv \
--config config.yaml
# Trains the model. This cell might take a few minutes.
train_stats, preprocessed_data, output_directory = model.train(dataset=train_df)
评估¶
评估训练好的模型。
ludwig evaluate --model_path results/experiment_run/model \
--dataset mnist_dataset.csv \
--split test \
--output_directory test_results
# Generates predictions and performance statistics for the test set.
test_stats, predictions, output_directory = model.evaluate(
test_df,
collect_predictions=True,
collect_overall_stats=True
)
可视化指标¶
显示混淆矩阵和类别熵图。
ludwig visualize confusion_matrix
命令
ludwig visualize --visualization confusion_matrix \
--ground_truth_metadata results/experiment_run/model/training_set_metadata.json \
--test_statistics test_results/test_statistics.json \
--output_directory visualizations \
--file_format png
visualize.confusion_matrix()
函数
# Visualizes confusion matrix, which gives an overview of classifier performance
# for each class.
from ludwig.visualize import confusion_matrix
confusion_matrix(
[test_stats],
model.training_set_metadata,
'label',
top_n_classes=[5],
model_names=[''],
normalize=True,
)
显示学习曲线图。
ludwig visualize learning_curves
命令
ludwig visualize --visualization learning_curves \
--ground_truth_metadata results/experiment_run/model/training_set_metadata.json \
--training_statistics results/experiment_run/training_statistics.json \
--file_format png \
--output_directory visualizations
visualize.learning_curves()
函数
# Visualizes learning curves, which show how performance metrics changed over
# time during training.
from ludwig.visualize import learning_curves
learning_curves(train_stats, output_feature_name='label')
预测¶
从测试数据集生成预测结果。
ludwig predict --model_path results/experiment_run/model \
--dataset mnist_dataset.csv \
--split test \
--output_directory predictions
predictions, output_directory = model.predict(test_df)
显示真实标签("label") 和预测标签("pred") 的测试图像样本。