用于扫描文档分类的AutoMM

Open In Colab Open In SageMaker Studio Lab

组织中的纸质文件是信息的重要来源,无论行业如何。 处理纸质文件是一件令人头疼的事情,因为它们可能占据大量空间,随着时间的推移容易磨损或褪色,并且难以跟踪。 因此,通过扫描仪、相机等将纸质文件数字化的趋势日益增长。 然而,数字化并不一定带来自动化,识别、分类和分析数字文档仍然可能是一个劳动密集型的过程。 例如,将数字书籍分类为不同的类型,并将扫描的收据分类为公用事业交通保险租金用品等,如果手动完成,既耗时又繁琐。 随着新的人工智能技术的发展,自动化数字文档处理变得更加容易和有效。 可以说,人工智能是现代数字文档处理系统的基石。

在本教程中,我们将展示如何使用Autogluon Multimodal通过几行代码构建一个扫描文档分类器。让我们开始吧!

获取文档数据集

现在让我们下载一个扫描文档数据集。 这个数据集是RVL-CDIP的一个样本,原始数据集包含16个类别的400,000张灰度图像,每个类别有25,000张图像。 在这里,我们采样了大约100个文档和三类文档,包括预算(标记为0)、电子邮件(标记为1)和表格(标记为2)。

import warnings
warnings.filterwarnings('ignore')

import os
import pandas as pd
from autogluon.core.utils.loaders import load_zip

download_dir = './ag_automm_tutorial_doc_classifier'
zip_file = "https://automl-mm-bench.s3.amazonaws.com/doc_classification/rvl_cdip_sample.zip"
load_zip.unzip(zip_file, unzip_dir=download_dir)
Downloading ./ag_automm_tutorial_doc_classifier/file.zip from https://automl-mm-bench.s3.amazonaws.com/doc_classification/rvl_cdip_sample.zip...
0%|          | 0.00/7.95M [00:00<?, ?iB/s]
100%|██████████| 7.95M/7.95M [00:00<00:00, 104MiB/s]

我们加载下面的训练和测试数据。

dataset_path = os.path.join(download_dir, "rvl_cdip_sample")
rvl_cdip_data = pd.read_csv(f"{dataset_path}/rvl_cdip_train_data.csv")
train_data = rvl_cdip_data.sample(frac=0.8, random_state=200)
test_data = rvl_cdip_data.drop(train_data.index)

我们需要扩展文档路径以在训练中加载它们。

from autogluon.multimodal.utils.misc import path_expander

DOC_PATH_COL = "doc_path"

train_data[DOC_PATH_COL] = train_data[DOC_PATH_COL].apply(lambda ele: path_expander(ele, base_folder=download_dir))
test_data[DOC_PATH_COL] = test_data[DOC_PATH_COL].apply(lambda ele: path_expander(ele, base_folder=download_dir))
print(test_data.head())
doc_path  label
1   /home/ci/autogluon/docs/tutorials/multimodal/d...      0
6   /home/ci/autogluon/docs/tutorials/multimodal/d...      0
7   /home/ci/autogluon/docs/tutorials/multimodal/d...      0
11  /home/ci/autogluon/docs/tutorials/multimodal/d...      0
14  /home/ci/autogluon/docs/tutorials/multimodal/d...      0
/home/ci/opt/venv/lib/python3.11/site-packages/mmengine/optim/optimizer/zero_optimizer.py:11: DeprecationWarning: `TorchScript` support for functional optimizers is deprecated and will be removed in a future PyTorch release. Consider using the `torch.compile` optimizer instead.
  from torch.distributed.optim import \

让我们展示其中一个文档。 如您所见,这是一个预算文档,包含账号、账户名称、预算资金、支出等。

from IPython.display import Image, display

example_image = train_data.iloc[0][DOC_PATH_COL]
pil_img = Image(filename=example_image, width=500)
display(pil_img)

使用AutoMM构建扫描文档分类器

您可以使用我们的MultiModalPredictor构建一个扫描文档分类器。 您所需要做的就是创建一个预测器,并使用上述训练数据集进行拟合。 在底层,AutoMM将自动识别手写或打印的文本,并利用识别出的文本、布局信息以及视觉特征进行文档分类。 模型定制也非常简单,您可以使用model.document_transformer.checkpoint_name超参数指定基础模型,AutoMM支持文档基础模型,如layoutlmv3layoutlmv2layoutlm-baselayoutxlm等, 以及纯文本模型,如bertdeberta,仅举几例。

在这里,label 是包含要预测的目标变量的列的名称,例如,在我们的示例中它是“label”。 为了演示目的,我们将训练时间限制设置为120秒。

from autogluon.multimodal import MultiModalPredictor

predictor = MultiModalPredictor(label="label")
predictor.fit(
    train_data=train_data,
    hyperparameters={"model.document_transformer.checkpoint_name":"microsoft/layoutlm-base-uncased",
    "optimization.top_k_average_method":"best",
    },
    time_limit=120,
)
No path specified. Models will be saved in: "AutogluonModels/ag-20241127_094620"
=================== System Info ===================
AutoGluon Version:  1.2b20241127
Python Version:     3.11.9
Operating System:   Linux
Platform Machine:   x86_64
Platform Version:   #1 SMP Tue Sep 24 10:00:37 UTC 2024
CPU Count:          8
Pytorch Version:    2.5.1+cu124
CUDA Version:       12.4
Memory Avail:       28.42 GB / 30.95 GB (91.8%)
Disk Space Avail:   187.75 GB / 255.99 GB (73.3%)
===================================================
AutoGluon infers your prediction problem is: 'multiclass' (because dtype of label-column == int, but few unique label-values observed).
3 unique label values:  [0, 1, 2]
If 'multiclass' is not the correct problem_type, please manually specify the problem_type parameter during Predictor init (You may specify problem_type as one of: ['binary', 'multiclass', 'regression', 'quantile'])
AutoMM starts to create your model. ✨✨✨

To track the learning progress, you can open a terminal and launch Tensorboard:
    ```shell
    # Assume you have installed tensorboard
    tensorboard --logdir /home/ci/autogluon/docs/tutorials/multimodal/document_prediction/AutogluonModels/ag-20241127_094620
    ```
INFO: Seed set to 0
GPU Count: 1
GPU Count to be Used: 1
GPU 0 Name: Tesla T4
GPU 0 Memory: 0.43GB/15.0GB (Used/Total)
INFO: Using 16bit Automatic Mixed Precision (AMP)
INFO: GPU available: True (cuda), used: True
INFO: TPU available: False, using: 0 TPU cores
INFO: HPU available: False, using: 0 HPUs
INFO: LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
INFO: 
  | Name              | Type                | Params | Mode 
------------------------------------------------------------------
0 | model             | DocumentTransformer | 112 M  | train
1 | validation_metric | MulticlassAccuracy  | 0      | train
2 | loss_func         | CrossEntropyLoss    | 0      | train
------------------------------------------------------------------
112 M     Trainable params
0         Non-trainable params
112 M     Total params
450.521   Total estimated model params size (MB)
4         Modules in train mode
232       Modules in eval mode
INFO: Epoch 0, global step 1: 'val_accuracy' reached 0.43750 (best 0.43750), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/document_prediction/AutogluonModels/ag-20241127_094620/epoch=0-step=1.ckpt' as top 3
INFO: Epoch 1, global step 2: 'val_accuracy' reached 0.56250 (best 0.56250), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/document_prediction/AutogluonModels/ag-20241127_094620/epoch=1-step=2.ckpt' as top 3
INFO: Epoch 2, global step 3: 'val_accuracy' reached 0.68750 (best 0.68750), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/document_prediction/AutogluonModels/ag-20241127_094620/epoch=2-step=3.ckpt' as top 3
INFO: Epoch 3, global step 4: 'val_accuracy' reached 1.00000 (best 1.00000), saving model to '/home/ci/autogluon/docs/tutorials/multimodal/document_prediction/AutogluonModels/ag-20241127_094620/epoch=3-step=4.ckpt' as top 3
AutoMM has created your model. 🎉🎉🎉

To load the model, use the code below:
    ```python
    from autogluon.multimodal import MultiModalPredictor
    predictor = MultiModalPredictor.load("/home/ci/autogluon/docs/tutorials/multimodal/document_prediction/AutogluonModels/ag-20241127_094620")
    ```

If you are not satisfied with the model, try to increase the training time, 
adjust the hyperparameters (https://auto.gluon.ai/stable/tutorials/multimodal/advanced_topics/customization.html),
or post issues on GitHub (https://github.com/autogluon/autogluon/issues).
<autogluon.multimodal.predictor.MultiModalPredictor at 0x7f6238283950>

在测试数据集上评估

你可以在测试数据集上评估分类器,看看它的表现如何:

scores = predictor.evaluate(test_data, metrics=["accuracy"])
print('The test acc: %.3f' % scores["accuracy"])
The test acc: 0.900

预测新文档

给定一个示例文档,让我们首先将其可视化,

doc_path = test_data.iloc[1][DOC_PATH_COL]
from IPython.display import Image, display
pil_img = Image(filename=doc_path, width=500)
display(pil_img)

我们可以轻松地使用最终模型来预测标签,

predictions = predictor.predict({DOC_PATH_COL: [doc_path]})
print(predictions)
[0]

上述输出显示,训练好的模型正确地将给定文档分类到预算类别中。

如果需要所有类别的概率,你可以调用 predict_proba:

proba = predictor.predict_proba({DOC_PATH_COL: [doc_path]})
print(proba)
[[0.9326989  0.048893   0.01840808]]

提取嵌入

从模型学习的整个文档中提取表示也非常有用。 我们提供 extract_embedding 函数,允许预测器返回 N 维文档特征,其中 N 取决于模型。

feature = predictor.extract_embedding({DOC_PATH_COL: [doc_path]})
print(feature[0].shape)
(768,)

其他示例

你可以访问AutoMM 示例来探索更多关于 AutoMM 的示例。

自定义

要学习如何自定义AutoMM,请参考自定义AutoMM