跳到内容

图像分类

这是一个在 MNIST 手写数字数据集上训练图像分类模型的完整示例。

这些交互式笔记本遵循本示例的步骤

  • Ludwig CLI: 使用 Ludwig CLI 进行 MNIST 图像分类
  • Ludwig Python API: 使用 Ludwig API 进行 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)
这将创建两个 pandas DataFrame。train_df 包含 60K 训练图像的文件路径信息。test_df 包含 10K 测试图像的相同信息。

描述
image_path 图像的文件路径字符串
label 0 到 9 的单个数字,表示图像中显示的数字

带有 label 的图像样本。 mnist 样本图像

训练

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

LudwigModel

# 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 命令

ludwig train \
  --dataset mnist_dataset.csv \
  --config config.yaml

train() 方法

# Trains the model. This cell might take a few minutes.
train_stats, preprocessed_data, output_directory = model.train(dataset=train_df)

评估

评估训练好的模型。

ludwig evaluate 命令

ludwig evaluate --model_path results/experiment_run/model \
                 --dataset mnist_dataset.csv \
                 --split test \
                 --output_directory test_results

evaluate() 方法

# 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,
)

confusion matrix and entropy

显示学习曲线图。

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')

confusion learning curves

预测

从测试数据集生成预测结果。

ludwig predict 命令

ludwig predict --model_path results/experiment_run/model \
                --dataset mnist_dataset.csv \
                --split test \
                --output_directory predictions

predict() 方法

predictions, output_directory = model.predict(test_df)

显示真实标签("label") 和预测标签("pred") 的测试图像样本。 mnist 样本预测结果