02. 使用 Auto Estimator 训练图像分类

本教程介绍了使用 GluonCV auto estimator 通过自定义超参数训练图像分类器的基本步骤。

使用默认配置进行训练

from gluoncv.auto.estimators import ImageClassificationEstimator

在本教程中，我们使用了一个小型样本数据集

train, _, test = ImageClassificationEstimator.Dataset.from_folders(
    'https://autogluon.s3.amazonaws.com/datasets/shopee-iet.zip')
train, val, _ = train.random_split(val_size=0.1, test_size=0)

输出

data/
├── test/
└── train/

使用默认配置创建一个 Estimator。我们只更改 GPU 数量以反映硬件限制。

请注意，即使没有可用的 GPU（设置为 {‘gpus’: []}），您仍然可以启动训练，但请做好准备，这会非常缓慢，并且如果数据集不是非常小，甚至不可能完成。

我们建议您使用至少一块拥有超过 6GB 可用 GPU 内存的英伟达 GPU。

classifier = ImageClassificationEstimator(
    {'gpus': [0], 'train': {'batch_size': 16, 'epochs': 2}})

在训练/验证数据上运行 fit

classifier.fit(train, val)

输出

Downloading /root/.mxnet/models/resnet50_v1-cc729d95.zip from https://apache-mxnet.s3-accelerate.dualstack.amazonaws.com/gluon/models/resnet50_v1-cc729d95.zip...

  0%|          | 0/57421 [00:00<?, ?KB/s]
  0%|          | 99/57421 [00:00<01:12, 793.84KB/s]
  1%|          | 510/57421 [00:00<00:25, 2253.13KB/s]
  4%|3         | 2189/57421 [00:00<00:07, 7331.41KB/s]
 14%|#3        | 7815/57421 [00:00<00:02, 24063.87KB/s]
 25%|##4       | 14265/57421 [00:00<00:01, 37334.23KB/s]
 40%|###9      | 22933/57421 [00:00<00:00, 53084.51KB/s]
 53%|#####3    | 30480/57421 [00:00<00:00, 60094.47KB/s]
 67%|######7   | 38483/57421 [00:00<00:00, 65610.12KB/s]
 82%|########1 | 46854/57421 [00:00<00:00, 71114.66KB/s]
 95%|#########5| 54600/57421 [00:01<00:00, 73035.76KB/s]
100%|##########| 57421/57421 [00:01<00:00, 51737.76KB/s]

在测试集上评估最终模型

eval_result = classifier.evaluate(test)
print("Top1/Top5 on test dataset: {}".format(eval_result))

输出

Top1/Top5 on test dataset: (0.8875, 1.0)

保存到磁盘/从磁盘加载以便后续使用

classifier.save('classifier.pkl')
classifier2 = ImageClassificationEstimator.load('classifier.pkl')

在测试图像上运行预测

pred = classifier2.predict(test.iloc[0]['image'])
print('GroundTruth:', test.iloc[0])
print('prediction:', pred)

输出

GroundTruth: image    /root/.gluoncv/datasets/shopee-iet/data/test/B...
label                                                    0
Name: 0, dtype: object
prediction:               class     score  id
0         BabyPants  0.587603   0
1         BabyShirt  0.388625   1
2   womenchiffontop  0.014393   3
3  womencasualshoes  0.009378   2

自定义您的训练配置

您可以修改配置来自定义您的训练，支持的字段有

print(ImageClassificationEstimator._default_cfg)

输出

ImageClassificationCfg(img_cls=ImageClassification(model=resnet50_v1, use_pretrained=True, use_gn=False, batch_norm=False, use_se=False, last_gamma=False), train=TrainCfg(pretrained_base=True, batch_size=128, epochs=10, lr=0.1, lr_decay=0.1, lr_decay_period=0, lr_decay_epoch=40, 60, lr_mode=step, warmup_lr=0.0, warmup_epochs=0, num_training_samples=1281167, num_workers=4, wd=0.0001, momentum=0.9, teacher=None, hard_weight=0.5, dtype=float32, input_size=224, crop_ratio=0.875, use_rec=False, rec_train=~/.mxnet/datasets/imagenet/rec/train.rec, rec_train_idx=~/.mxnet/datasets/imagenet/rec/train.idx, rec_val=~/.mxnet/datasets/imagenet/rec/val.rec, rec_val_idx=~/.mxnet/datasets/imagenet/rec/val.idx, data_dir=~/.mxnet/datasets/imagenet, mixup=False, no_wd=False, label_smoothing=False, temperature=20, resume_epoch=0, mixup_alpha=0.2, mixup_off_epoch=0, log_interval=50, mode=, start_epoch=0, transfer_lr_mult=0.01, output_lr_mult=0.1, early_stop_patience=-1, early_stop_min_delta=0.001, early_stop_baseline=0.0, early_stop_max_value=1.0), valid=ValidCfg(batch_size=128, num_workers=4), gpus=(0,))

例如，我们可以更改学习率和批量大小

new_classifier = ImageClassificationEstimator({'gpus': [0], 'train': {'batch_size': 16, 'lr': 0.01}})

修改单个超参数更自然的方式是编辑自动保存在 self._logdir 中的 yaml 文件，这里我们仅展示一个如何在 python 中复制/编辑修改后的配置文件并将其加载回 estimator 的示例

您可以直接使用文本编辑器编辑 yaml 文件

import shutil
import os
config_name = 'config.yaml'
shutil.copyfile(os.path.join(classifier2._logdir, config_name), os.path.join('.', config_name))
cfg = open(config_name).read()
print(cfg)

输出

# ImageClassificationCfg
gpus: !!python/tuple
- 0
img_cls:
  batch_norm: false
  last_gamma: false
  model: resnet50_v1
  use_gn: false
  use_pretrained: true
  use_se: false
train:
  batch_size: 16
  crop_ratio: 0.875
  data_dir: ~/.mxnet/datasets/imagenet
  dtype: float32
  early_stop_baseline: 0.0
  early_stop_max_value: 1.0
  early_stop_min_delta: 0.001
  early_stop_patience: -1
  epochs: 10
  hard_weight: 0.5
  input_size: 224
  label_smoothing: false
  log_interval: 50
  lr: 0.01
  lr_decay: 0.1
  lr_decay_epoch: 40, 60
  lr_decay_period: 0
  lr_mode: step
  mixup: false
  mixup_alpha: 0.2
  mixup_off_epoch: 0
  mode: ''
  momentum: 0.9
  no_wd: false
  num_training_samples: 1281167
  num_workers: 4
  output_lr_mult: 0.1
  pretrained_base: true
  rec_train: ~/.mxnet/datasets/imagenet/rec/train.rec
  rec_train_idx: ~/.mxnet/datasets/imagenet/rec/train.idx
  rec_val: ~/.mxnet/datasets/imagenet/rec/val.rec
  rec_val_idx: ~/.mxnet/datasets/imagenet/rec/val.idx
  resume_epoch: 0
  start_epoch: 0
  teacher: null
  temperature: 20
  transfer_lr_mult: 0.01
  use_rec: false
  warmup_epochs: 0
  warmup_lr: 0.0
  wd: 0.0001
valid:
  batch_size: 128
  num_workers: 4

让我们将网络从 resnet50_v1 修改为 resnet18_v1b

import fileinput
with fileinput.FileInput(config_name,
                         inplace = True, backup ='.bak') as f:
    for line in f:
        if 'resnet50_v1' in line:
            new_line = line.replace('resnet50_v1', 'resnet18_v1b')
            print(new_line, end='')
        else:
            print(line, end='')

新的分类器现在应该反映我们刚刚编辑的新配置

new_classifier2 = ImageClassificationEstimator(config_name)